def test_feature_table_import_export_yaml(self, batch_source):
stream_source = KafkaSource(
field_mapping={
"ride_distance": "ride_distance",
"ride_duration": "ride_duration",
},
bootstrap_servers="localhost:9094",
message_format=ProtoFormat(class_path="class.path"),
topic="test_topic",
event_timestamp_column="ts_col",
)
test_feature_table = FeatureTable(
name="car_driver",
features=[
Feature(name="ride_distance", dtype=ValueType.FLOAT),
Feature(name="ride_duration", dtype=ValueType.STRING),
],
entities=["car_driver_entity"],
labels={"team": "matchmaking"},
batch_source=batch_source,
stream_source=stream_source,
)
# Create a string YAML representation of the feature table
string_yaml = test_feature_table.to_yaml()
# Create a new feature table object from the YAML string
actual_feature_table_from_string = FeatureTable.from_yaml(string_yaml)
# Ensure equality is upheld to original feature table
assert test_feature_table == actual_feature_table_from_string
def test_feature_table_import_export_yaml(self):
batch_source = DataSource(
type=SourceType(1).name,
field_mapping={
"ride_distance": "ride_distance",
"ride_duration": "ride_duration",
},
options=FileOptions(file_format="avro", file_url="data/test.avro"),
timestamp_column="ts_col",
date_partition_column="date_partition_col",
)
stream_source = DataSource(
type=SourceType(3).name,
field_mapping={
"ride_distance": "ride_distance",
"ride_duration": "ride_duration",
},
options=KafkaOptions(
bootstrap_servers="localhost:9094",
class_path="random/path/to/class",
topic="test_topic",
),
timestamp_column="ts_col",
)
test_feature_table = FeatureTable(
name="car_driver",
features=[
FeatureV2(name="ride_distance",
dtype=ValueType.FLOAT).to_proto(),
FeatureV2(name="ride_duration",
dtype=ValueType.STRING).to_proto(),
],
entities=["car_driver_entity"],
labels={"team": "matchmaking"},
batch_source=batch_source.to_proto(),
stream_source=stream_source.to_proto(),
)
# Create a string YAML representation of the feature table
string_yaml = test_feature_table.to_yaml()
# Create a new feature table object from the YAML string
actual_feature_table_from_string = FeatureTable.from_yaml(string_yaml)
# Ensure equality is upheld to original feature table
assert test_feature_table == actual_feature_table_from_string