def test_feature_set_ingest_success(self, dataframe, client):
driver_fs = FeatureSet("driver-feature-set")
driver_fs.add(Feature(name="feature_1", dtype=ValueType.FLOAT))
driver_fs.add(Feature(name="feature_2", dtype=ValueType.STRING))
driver_fs.add(Feature(name="feature_3", dtype=ValueType.INT64))
driver_fs.add(Entity(name="entity_id", dtype=ValueType.INT64))
driver_fs.source = KafkaSource(topic="feature-topic",
brokers="127.0.0.1")
driver_fs._message_producer = MagicMock()
driver_fs._message_producer.send = MagicMock()
# Register with Feast core
client.apply(driver_fs)
# Ingest data into Feast
driver_fs.ingest(dataframe=dataframe)
# Make sure message producer is called
driver_fs._message_producer.send.assert_called()
def test_feature_set_types_success(self, client, dataframe):
all_types_fs = FeatureSet(
name="all_types",
entities=[Entity(name="user_id", dtype=ValueType.INT64)],
features=[
Feature(name="float_feature", dtype=ValueType.FLOAT),
Feature(name="int64_feature", dtype=ValueType.INT64),
Feature(name="int32_feature", dtype=ValueType.INT32),
Feature(name="string_feature", dtype=ValueType.STRING),
Feature(name="bytes_feature", dtype=ValueType.BYTES),
Feature(name="bool_feature", dtype=ValueType.BOOL),
Feature(name="double_feature", dtype=ValueType.DOUBLE),
Feature(name="float_list_feature", dtype=ValueType.FLOAT_LIST),
Feature(name="int64_list_feature", dtype=ValueType.INT64_LIST),
Feature(name="int32_list_feature", dtype=ValueType.INT32_LIST),
Feature(name="string_list_feature",
dtype=ValueType.STRING_LIST),
Feature(name="bytes_list_feature", dtype=ValueType.BYTES_LIST),
Feature(name="bool_list_feature", dtype=ValueType.BOOL_LIST),
Feature(name="double_list_feature",
dtype=ValueType.DOUBLE_LIST),
],
max_age=Duration(seconds=3600),
)
all_types_fs.source = KafkaSource(topic="feature-topic",
brokers="127.0.0.1")
all_types_fs._message_producer = MagicMock()
all_types_fs._message_producer.send = MagicMock()
# Register with Feast core
client.apply(all_types_fs)
# Ingest data into Feast
all_types_fs.ingest(dataframe=dataframe)
# Make sure message producer is called
all_types_fs._message_producer.send.assert_called()
def test_feature_set_ingest_failure(self, client, dataframe, exception):
with pytest.raises(exception):
# Create feature set
driver_fs = FeatureSet("driver-feature-set")
driver_fs.source = KafkaSource(topic="feature-topic",
brokers="fake.broker.com")
driver_fs._message_producer = MagicMock()
driver_fs._message_producer.send = MagicMock()
# Update based on dataset
driver_fs.update_from_dataset(
dataframe,
column_mapping={
"entity_id": Entity(name="entity", dtype=ValueType.INT64)
},
)
# Register with Feast core
client.apply(driver_fs)
# Ingest data into Feast
driver_fs.ingest(dataframe=dataframe)
