Exemple #1
0
    def test_list_feature_sets(self, mocked_client, mocker):
        mocker.patch.object(
            mocked_client,
            "_core_service_stub",
            return_value=Core.CoreServiceStub(grpc.insecure_channel("")),
        )

        feature_set_1_proto = FeatureSetProto(
            spec=FeatureSetSpecProto(
                project="test",
                name="driver_car",
                max_age=Duration(seconds=3600),
                labels={"key1": "val1", "key2": "val2"},
                features=[
                    FeatureSpecProto(
                        name="feature_1", value_type=ValueProto.ValueType.FLOAT
                    )
                ],
            )
        )
        feature_set_2_proto = FeatureSetProto(
            spec=FeatureSetSpecProto(
                project="test",
                name="driver_ride",
                max_age=Duration(seconds=3600),
                labels={"key1": "val1"},
                features=[
                    FeatureSpecProto(
                        name="feature_1", value_type=ValueProto.ValueType.FLOAT
                    )
                ],
            )
        )

        mocker.patch.object(
            mocked_client._core_service_stub,
            "ListFeatureSets",
            return_value=ListFeatureSetsResponse(
                feature_sets=[feature_set_1_proto, feature_set_2_proto]
            ),
        )

        feature_sets = mocked_client.list_feature_sets(labels={"key1": "val1"})
        assert len(feature_sets) == 2

        feature_set = feature_sets[0]
        assert (
            feature_set.name == "driver_car"
            and "key1" in feature_set.labels
            and feature_set.labels["key1"] == "val1"
            and "key2" in feature_set.labels
            and feature_set.labels["key2"] == "val2"
            and feature_set.fields["feature_1"].name == "feature_1"
            and feature_set.fields["feature_1"].dtype == ValueType.FLOAT
            and len(feature_set.features) == 1
        )
Exemple #2
0
 def from_dict(cls, fs_dict):
     if ("kind" not in fs_dict) and (fs_dict["kind"].strip() != "feature_set"):
         raise Exception(f"Resource kind is not a feature set {str(fs_dict)}")
     feature_set_proto = json_format.ParseDict(
         fs_dict, FeatureSetSpecProto(), ignore_unknown_fields=True
     )
     return cls.from_proto(feature_set_proto)
Exemple #3
0
    def test_get_feature_set(self, mocked_client, mocker):
        mocked_client._core_service_stub = Core.CoreServiceStub(
            grpc.insecure_channel("")
        )

        from google.protobuf.duration_pb2 import Duration

        mocker.patch.object(
            mocked_client._core_service_stub,
            "GetFeatureSet",
            return_value=GetFeatureSetResponse(
                feature_set=FeatureSetProto(
                    spec=FeatureSetSpecProto(
                        name="my_feature_set",
                        max_age=Duration(seconds=3600),
                        labels={"key1": "val1", "key2": "val2"},
                        features=[
                            FeatureSpecProto(
                                name="my_feature_1",
                                value_type=ValueProto.ValueType.FLOAT,
                            ),
                            FeatureSpecProto(
                                name="my_feature_2",
                                value_type=ValueProto.ValueType.FLOAT,
                            ),
                        ],
                        entities=[
                            EntitySpecProto(
                                name="my_entity_1",
                                value_type=ValueProto.ValueType.INT64,
                            )
                        ],
                        source=Source(
                            type=SourceType.KAFKA,
                            kafka_source_config=KafkaSourceConfig(
                                bootstrap_servers="localhost:9092", topic="topic"
                            ),
                        ),
                    ),
                    meta=FeatureSetMetaProto(),
                )
            ),
        )
        mocked_client.set_project("my_project")
        feature_set = mocked_client.get_feature_set("my_feature_set")

        assert (
            feature_set.name == "my_feature_set"
            and "key1" in feature_set.labels
            and feature_set.labels["key1"] == "val1"
            and "key2" in feature_set.labels
            and feature_set.labels["key2"] == "val2"
            and feature_set.fields["my_feature_1"].name == "my_feature_1"
            and feature_set.fields["my_feature_1"].dtype == ValueType.FLOAT
            and feature_set.fields["my_entity_1"].name == "my_entity_1"
            and feature_set.fields["my_entity_1"].dtype == ValueType.INT64
            and len(feature_set.features) == 2
            and len(feature_set.entities) == 1
        )
Exemple #4
0
    def test_list_ingest_jobs(self, mocked_client, mocker):
        mocker.patch.object(
            mocked_client,
            "_core_service_stub",
            return_value=Core.CoreServiceStub(grpc.insecure_channel("")),
        )

        feature_set_proto = FeatureSetProto(
            spec=FeatureSetSpecProto(
                project="test", name="driver", max_age=Duration(seconds=3600),
            )
        )

        mocker.patch.object(
            mocked_client._core_service_stub,
            "ListIngestionJobs",
            return_value=ListIngestionJobsResponse(
                jobs=[
                    IngestJobProto(
                        id="kafka-to-redis",
                        external_id="job-2222",
                        status=IngestionJobStatus.RUNNING,
                        feature_sets=[feature_set_proto],
                        source=Source(
                            type=SourceType.KAFKA,
                            kafka_source_config=KafkaSourceConfig(
                                bootstrap_servers="localhost:9092", topic="topic"
                            ),
                        ),
                        store=Store(name="redis"),
                    )
                ]
            ),
        )

        # list ingestion jobs by target feature set reference
        ingest_jobs = mocked_client.list_ingest_jobs(
            feature_set_ref=FeatureSetRef.from_feature_set(
                FeatureSet.from_proto(feature_set_proto)
            )
        )
        assert len(ingest_jobs) >= 1

        ingest_job = ingest_jobs[0]
        assert (
            ingest_job.status == IngestionJobStatus.RUNNING
            and ingest_job.id == "kafka-to-redis"
            and ingest_job.external_id == "job-2222"
            and ingest_job.feature_sets[0].name == "driver"
            and ingest_job.source.source_type == "Kafka"
        )
Exemple #5
0
 def to_proto(self) -> FeatureSetSpecProto:
     return FeatureSetSpecProto(
         name=self.name,
         version=self.version,
         max_age=self.max_age,
         source=self.source.to_proto() if self.source is not None else None,
         features=[
             field.to_proto()
             for field in self._fields.values()
             if type(field) == Feature
         ],
         entities=[
             field.to_proto()
             for field in self._fields.values()
             if type(field) == Entity
         ],
     )
Exemple #6
0
    def to_proto(self) -> FeatureSetSpecProto:
        """
        Converts a feature set object to its protobuf representation

        Returns:
            FeatureSetSpec protobuf
        """

        return FeatureSetSpecProto(
            name=self.name,
            version=self.version,
            max_age=self.max_age,
            source=self.source.to_proto() if self.source is not None else None,
            features=[
                field.to_proto() for field in self._fields.values()
                if type(field) == Feature
            ],
            entities=[
                field.to_proto() for field in self._fields.values()
                if type(field) == Entity
            ],
        )