def get_db_schema(self, feature_set: FeatureSet):
"""Get desired database schema.
Args:
feature_set: object processed with feature set metadata.
Returns:
Desired database schema.
"""
db_schema = self.db_config.translate(feature_set.get_schema())
return db_schema
def test_get_schema(self):
expected_schema = [
{
"column_name": "id",
"type": LongType(),
"primary_key": True
},
{
"column_name": "timestamp",
"type": TimestampType(),
"primary_key": False
},
{
"column_name": "feature1__avg_over_2_minutes_fixed_windows",
"type": FloatType(),
"primary_key": False,
},
{
"column_name": "feature1__avg_over_15_minutes_fixed_windows",
"type": FloatType(),
"primary_key": False,
},
{
"column_name":
"feature1__stddev_pop_over_2_minutes_fixed_windows",
"type": FloatType(),
"primary_key": False,
},
{
"column_name":
"feature1__stddev_pop_over_15_minutes_fixed_windows",
"type": FloatType(),
"primary_key": False,
},
]
feature_set = FeatureSet(
name="feature_set",
entity="entity",
description="description",
features=[
Feature(
name="feature1",
description="test",
transformation=SparkFunctionTransform(functions=[
Function(F.avg, DataType.FLOAT),
Function(F.stddev_pop, DataType.FLOAT),
]).with_window(
partition_by="id",
order_by=TIMESTAMP_COLUMN,
mode="fixed_windows",
window_definition=["2 minutes", "15 minutes"],
),
),
],
keys=[
KeyFeature(
name="id",
description="The user's Main ID or device ID",
dtype=DataType.BIGINT,
)
],
timestamp=TimestampFeature(),
)
schema = feature_set.get_schema()
assert schema == expected_schema