def test_source_raise(self):
with pytest.raises(ValueError,
match="source must be a Source instance"):
FeatureSetPipeline(
spark_client=SparkClient(),
source=Mock(
spark_client=SparkClient(),
readers=[
TableReader(
id="source_a",
database="db",
table="table",
),
],
query="select * from source_a",
),
feature_set=Mock(
spec=FeatureSet,
name="feature_set",
entity="entity",
description="description",
keys=[
KeyFeature(
name="user_id",
description="The user's Main ID or device ID",
dtype=DataType.INTEGER,
)
],
timestamp=TimestampFeature(from_column="ts"),
features=[
Feature(
name="listing_page_viewed__rent_per_month",
description="Average of something.",
transformation=SparkFunctionTransform(functions=[
Function(functions.avg, DataType.FLOAT),
Function(functions.stddev_pop, DataType.FLOAT),
], ).with_window(
partition_by="user_id",
order_by=TIMESTAMP_COLUMN,
window_definition=["7 days", "2 weeks"],
mode="fixed_windows",
),
),
],
),
sink=Mock(
spec=Sink,
writers=[HistoricalFeatureStoreWriter(db_config=None)],
),
)
def test_get_schema(self, target_df: DataFrame,
spark_session: SparkSession) -> None:
# arrange
spark_client = SparkClient()
create_temp_view(dataframe=target_df, name="temp_view")
create_db_and_table(
spark=spark_session,
database="test_db",
table="test_table",
view="temp_view",
)
expected_schema = [
{
"col_name": "col1",
"data_type": "string"
},
{
"col_name": "col2",
"data_type": "bigint"
},
]
# act
schema = spark_client.get_schema(table="test_table",
database="test_db")
# assert
assert schema, expected_schema
def test_read(
self,
format: str,
stream: bool,
schema: Optional[StructType],
path: Any,
options: Any,
target_df: DataFrame,
mocked_spark_read: Mock,
) -> None:
# arrange
spark_client = SparkClient()
mocked_spark_read.load.return_value = target_df
spark_client._session = mocked_spark_read
# act
result_df = spark_client.read(format=format,
schema=schema,
stream=stream,
path=path,
**options)
# assert
mocked_spark_read.format.assert_called_once_with(format)
mocked_spark_read.load.assert_called_once_with(path=path, **options)
assert target_df.collect() == result_df.collect()
def test_add_table_partitions(self, mock_spark_sql: Mock):
# arrange
target_command = (f"ALTER TABLE `db`.`table` ADD IF NOT EXISTS "
f"PARTITION ( year = 2020, month = 8, day = 14 ) "
f"PARTITION ( year = 2020, month = 8, day = 15 ) "
f"PARTITION ( year = 2020, month = 8, day = 16 )")
spark_client = SparkClient()
spark_client._session = mock_spark_sql
partitions = [
{
"year": 2020,
"month": 8,
"day": 14
},
{
"year": 2020,
"month": 8,
"day": 15
},
{
"year": 2020,
"month": 8,
"day": 16
},
]
# act
spark_client.add_table_partitions(partitions, "table", "db")
# assert
mock_spark_sql.assert_called_once_with(target_command)
def test_write_interval_mode(
self,
feature_set_dataframe,
historical_feature_set_dataframe,
mocker,
feature_set,
):
# given
spark_client = SparkClient()
spark_client.write_table = mocker.stub("write_table")
spark_client.conn.conf.set("spark.sql.sources.partitionOverwriteMode",
"dynamic")
writer = HistoricalFeatureStoreWriter(interval_mode=True)
# when
writer.write(
feature_set=feature_set,
dataframe=feature_set_dataframe,
spark_client=spark_client,
)
result_df = spark_client.write_table.call_args[1]["dataframe"]
# then
assert_dataframe_equality(historical_feature_set_dataframe, result_df)
assert writer.database == spark_client.write_table.call_args[1][
"database"]
assert feature_set.name == spark_client.write_table.call_args[1][
"table_name"]
assert (writer.PARTITION_BY == spark_client.write_table.call_args[1]
["partition_by"])
def test_csv_file_with_schema_and_header(self):
# given
spark_client = SparkClient()
schema_csv = StructType(
[
StructField("A", LongType()),
StructField("B", DoubleType()),
StructField("C", StringType()),
]
)
file = "tests/unit/butterfree/extract/readers/file-reader-test.csv"
# when
file_reader = FileReader(
id="id",
path=file,
format="csv",
schema=schema_csv,
format_options={"header": True},
)
df = file_reader.consume(spark_client)
# assert
assert schema_csv == df.schema
assert df.columns == ["A", "B", "C"]
for value in range(3):
assert df.first()[value] != ["A", "B", "C"][value]
def test_feature_transform_with_distinct_empty_subset(
self, timestamp_c, feature_set_with_distinct_dataframe):
spark_client = SparkClient()
with pytest.raises(ValueError,
match="The distinct subset param can't be empty."):
AggregatedFeatureSet(
name="name",
entity="entity",
description="description",
features=[
Feature(
name="feature",
description="test",
transformation=AggregatedTransform(functions=[
Function(functions.sum, DataType.INTEGER)
]),
),
],
keys=[
KeyFeature(name="h3",
description="test",
dtype=DataType.STRING)
],
timestamp=timestamp_c,
).with_windows(["3 days"]).with_distinct(
subset=[],
keep="first").construct(feature_set_with_distinct_dataframe,
spark_client,
end_date="2020-01-10")
def test_write_stream(self, mocked_stream_df: Mock) -> None:
# arrange
spark_client = SparkClient()
processing_time = "0 seconds"
output_mode = "update"
checkpoint_path = "s3://path/to/checkpoint"
# act
stream_handler = spark_client.write_stream(
mocked_stream_df,
processing_time,
output_mode,
checkpoint_path,
format_="parquet",
mode="append",
)
# assert
assert isinstance(stream_handler, StreamingQuery)
mocked_stream_df.trigger.assert_called_with(processingTime=processing_time)
mocked_stream_df.outputMode.assert_called_with(output_mode)
mocked_stream_df.option.assert_called_with(
"checkpointLocation", checkpoint_path
)
mocked_stream_df.foreachBatch.assert_called_once()
mocked_stream_df.start.assert_called_once()
def test_feature_set_with_invalid_feature(self, key_id, timestamp_c,
dataframe):
spark_client = SparkClient()
with pytest.raises(ValueError):
AggregatedFeatureSet(
name="name",
entity="entity",
description="description",
features=[
Feature(
name="feature1",
description="test",
transformation=SparkFunctionTransform(functions=[
Function(functions.avg, DataType.FLOAT)
], ).with_window(
partition_by="id",
mode="row_windows",
window_definition=["2 events"],
),
),
],
keys=[key_id],
timestamp=timestamp_c,
).construct(dataframe, spark_client)
def test_feature_transform_with_distinct(
self,
timestamp_c,
feature_set_with_distinct_dataframe,
target_with_distinct_dataframe,
):
spark_client = SparkClient()
fs = (AggregatedFeatureSet(
name="name",
entity="entity",
description="description",
features=[
Feature(
name="feature",
description="test",
transformation=AggregatedTransform(
functions=[Function(functions.sum, DataType.INTEGER)]),
),
],
keys=[
KeyFeature(name="h3",
description="test",
dtype=DataType.STRING)
],
timestamp=timestamp_c,
).with_windows(["3 days"]).with_distinct(subset=["id"], keep="last"))
# assert
output_df = fs.construct(feature_set_with_distinct_dataframe,
spark_client,
end_date="2020-01-10")
assert_dataframe_equality(output_df, target_with_distinct_dataframe)
def test_h3_feature_set(self, h3_input_df, h3_target_df):
spark_client = SparkClient()
feature_set = AggregatedFeatureSet(
name="h3_test",
entity="h3geolocation",
description="Test",
keys=[
KeyFeature(
name="h3_id",
description="The h3 hash ID",
dtype=DataType.DOUBLE,
transformation=H3HashTransform(
h3_resolutions=[6, 7, 8, 9, 10, 11, 12],
lat_column="lat",
lng_column="lng",
).with_stack(),
)
],
timestamp=TimestampFeature(),
features=[
Feature(
name="house_id",
description="Count of house ids over a day.",
transformation=AggregatedTransform(
functions=[Function(F.count, DataType.BIGINT)]),
),
],
).with_windows(definitions=["1 day"])
output_df = feature_set.construct(h3_input_df,
client=spark_client,
end_date="2016-04-14")
assert_dataframe_equality(output_df, h3_target_df)
def test_write_in_debug_and_stream_mode(
self, feature_set, spark_session,
):
# arrange
spark_client = SparkClient()
mocked_stream_df = Mock()
mocked_stream_df.isStreaming = True
mocked_stream_df.writeStream = mocked_stream_df
mocked_stream_df.format.return_value = mocked_stream_df
mocked_stream_df.queryName.return_value = mocked_stream_df
mocked_stream_df.start.return_value = Mock(spec=StreamingQuery)
writer = OnlineFeatureStoreWriter(debug_mode=True)
# act
handler = writer.write(
feature_set=feature_set,
dataframe=mocked_stream_df,
spark_client=spark_client,
)
# assert
mocked_stream_df.format.assert_called_with("memory")
mocked_stream_df.queryName.assert_called_with(
f"online_feature_store__{feature_set.name}"
)
assert isinstance(handler, StreamingQuery)
def test_sink(input_dataframe, feature_set):
# arrange
client = SparkClient()
client.conn.conf.set("spark.sql.sources.partitionOverwriteMode", "dynamic")
feature_set_df = feature_set.construct(input_dataframe, client)
target_latest_df = OnlineFeatureStoreWriter.filter_latest(
feature_set_df, id_columns=[key.name for key in feature_set.keys])
columns_sort = feature_set_df.schema.fieldNames()
# setup historical writer
s3config = Mock()
s3config.mode = "overwrite"
s3config.format_ = "parquet"
s3config.get_options = Mock(
return_value={"path": "test_folder/historical/entity/feature_set"})
s3config.get_path_with_partitions = Mock(
return_value="test_folder/historical/entity/feature_set")
historical_writer = HistoricalFeatureStoreWriter(db_config=s3config,
interval_mode=True)
# setup online writer
# TODO: Change for CassandraConfig when Cassandra for test is ready
online_config = Mock()
online_config.mode = "overwrite"
online_config.format_ = "parquet"
online_config.get_options = Mock(
return_value={"path": "test_folder/online/entity/feature_set"})
online_writer = OnlineFeatureStoreWriter(db_config=online_config)
writers = [historical_writer, online_writer]
sink = Sink(writers)
# act
client.sql("CREATE DATABASE IF NOT EXISTS {}".format(
historical_writer.database))
sink.flush(feature_set, feature_set_df, client)
# get historical results
historical_result_df = client.read(
s3config.format_,
path=s3config.get_path_with_partitions(feature_set.name,
feature_set_df),
)
# get online results
online_result_df = client.read(
online_config.format_, **online_config.get_options(feature_set.name))
# assert
# assert historical results
assert sorted(feature_set_df.select(*columns_sort).collect()) == sorted(
historical_result_df.select(*columns_sort).collect())
# assert online results
assert sorted(target_latest_df.select(*columns_sort).collect()) == sorted(
online_result_df.select(*columns_sort).collect())
# tear down
shutil.rmtree("test_folder")
def test_flush(self, feature_set_dataframe, mocker):
# given
spark_client = SparkClient()
writer = [
HistoricalFeatureStoreWriter(),
OnlineFeatureStoreWriter(),
]
for w in writer:
w.write = mocker.stub("write")
feature_set = mocker.stub("feature_set")
feature_set.entity = "house"
feature_set.name = "test"
# when
sink = Sink(writers=writer)
sink.flush(
dataframe=feature_set_dataframe,
feature_set=feature_set,
spark_client=spark_client,
)
# then
for w in writer:
w.write.assert_called_once()
def test_construct(
self, feature_set_dataframe, fixed_windows_output_feature_set_dataframe
):
# given
spark_client = SparkClient()
# arrange
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"],
),
),
Feature(
name="divided_feature",
description="unit test",
dtype=DataType.FLOAT,
transformation=CustomTransform(
transformer=divide, column1="feature1", column2="feature2",
),
),
],
keys=[
KeyFeature(
name="id",
description="The user's Main ID or device ID",
dtype=DataType.INTEGER,
)
],
timestamp=TimestampFeature(),
)
output_df = (
feature_set.construct(feature_set_dataframe, client=spark_client)
.orderBy(feature_set.timestamp_column)
.select(feature_set.columns)
)
target_df = fixed_windows_output_feature_set_dataframe.orderBy(
feature_set.timestamp_column
).select(feature_set.columns)
# assert
assert_dataframe_equality(output_df, target_df)
def test_write(
self,
feature_set_dataframe,
historical_feature_set_dataframe,
mocker,
feature_set,
):
# given
spark_client = SparkClient()
spark_client.write_table = mocker.stub("write_table")
writer = HistoricalFeatureStoreWriter()
# when
writer.write(
feature_set=feature_set,
dataframe=feature_set_dataframe,
spark_client=spark_client,
)
result_df = spark_client.write_table.call_args[1]["dataframe"]
# then
assert_dataframe_equality(historical_feature_set_dataframe, result_df)
assert (writer.db_config.format_ ==
spark_client.write_table.call_args[1]["format_"])
assert writer.db_config.mode == spark_client.write_table.call_args[1][
"mode"]
assert (writer.PARTITION_BY == spark_client.write_table.call_args[1]
["partition_by"])
def test_write_in_debug_mode_with_interval_mode(
self,
feature_set_dataframe,
historical_feature_set_dataframe,
feature_set,
spark_session,
mocker,
):
# given
spark_client = SparkClient()
spark_client.write_dataframe = mocker.stub("write_dataframe")
spark_client.conn.conf.set("spark.sql.sources.partitionOverwriteMode",
"dynamic")
writer = HistoricalFeatureStoreWriter(debug_mode=True,
interval_mode=True)
# when
writer.write(
feature_set=feature_set,
dataframe=feature_set_dataframe,
spark_client=spark_client,
)
result_df = spark_session.table(
f"historical_feature_store__{feature_set.name}")
# then
assert_dataframe_equality(historical_feature_set_dataframe, result_df)
def test_add_invalid_partitions(self, mock_spark_sql: Mock, partition):
# arrange
spark_client = SparkClient()
spark_client._session = mock_spark_sql
# act and assert
with pytest.raises(ValueError):
spark_client.add_table_partitions(partition, "table", "db")
def __init__(
self,
database: str = None,
) -> None:
self._db_config = MetastoreConfig()
self.database = database or environment.get_variable(
"FEATURE_STORE_HISTORICAL_DATABASE")
super(MetastoreMigration, self).__init__(SparkClient())
def test_read_invalid_params(self, format: Optional[str],
path: Any) -> None:
# arrange
spark_client = SparkClient()
# act and assert
with pytest.raises(ValueError):
spark_client.read(format=format, path=path) # type: ignore
def test_read_table_invalid_params(self, database: str,
table: Optional[int]) -> None:
# arrange
spark_client = SparkClient()
# act and assert
with pytest.raises(ValueError):
spark_client.read_table(table, database) # type: ignore
def test_run_agg_with_end_date(self, spark_session):
test_pipeline = FeatureSetPipeline(
spark_client=SparkClient(),
source=Mock(
spec=Source,
readers=[
TableReader(
id="source_a",
database="db",
table="table",
)
],
query="select * from source_a",
),
feature_set=Mock(
spec=AggregatedFeatureSet,
name="feature_set",
entity="entity",
description="description",
keys=[
KeyFeature(
name="user_id",
description="The user's Main ID or device ID",
dtype=DataType.INTEGER,
)
],
timestamp=TimestampFeature(from_column="ts"),
features=[
Feature(
name="listing_page_viewed__rent_per_month",
description="Average of something.",
transformation=AggregatedTransform(functions=[
Function(functions.avg, DataType.FLOAT),
Function(functions.stddev_pop, DataType.FLOAT),
], ),
),
],
),
sink=Mock(
spec=Sink,
writers=[HistoricalFeatureStoreWriter(db_config=None)],
),
)
# feature_set need to return a real df for streaming validation
sample_df = spark_session.createDataFrame([{
"a": "x",
"b": "y",
"c": "3"
}])
test_pipeline.feature_set.construct.return_value = sample_df
test_pipeline.run(end_date="2016-04-18")
test_pipeline.source.construct.assert_called_once()
test_pipeline.feature_set.construct.assert_called_once()
test_pipeline.sink.flush.assert_called_once()
test_pipeline.sink.validate.assert_called_once()
def test_construct_rolling_windows_with_end_date(
self,
feature_set_dataframe,
rolling_windows_output_feature_set_dataframe_base_date,
):
# given
spark_client = SparkClient()
# arrange
feature_set = AggregatedFeatureSet(
name="feature_set",
entity="entity",
description="description",
features=[
Feature(
name="feature1",
description="test",
transformation=AggregatedTransform(
functions=[
Function(F.avg, DataType.DOUBLE),
Function(F.stddev_pop, DataType.DOUBLE),
],
),
),
Feature(
name="feature2",
description="test",
transformation=AggregatedTransform(
functions=[
Function(F.avg, DataType.DOUBLE),
Function(F.stddev_pop, DataType.DOUBLE),
],
),
),
],
keys=[
KeyFeature(
name="id",
description="The user's Main ID or device ID",
dtype=DataType.INTEGER,
)
],
timestamp=TimestampFeature(),
).with_windows(definitions=["1 day", "1 week"])
# act
output_df = feature_set.construct(
feature_set_dataframe, client=spark_client, end_date="2016-04-18"
).orderBy("timestamp")
target_df = rolling_windows_output_feature_set_dataframe_base_date.orderBy(
feature_set.timestamp_column
).select(feature_set.columns)
# assert
assert_dataframe_equality(output_df, target_df)
def test_conn(self) -> None:
# arrange
spark_client = SparkClient()
# act
start_conn = spark_client._session
# assert
assert start_conn is None
def test_read_invalid_params(
self, format: Optional[str], options: Union[Dict[str, Any], str]
) -> None:
# arrange
spark_client = SparkClient()
# act and assert
with pytest.raises(ValueError):
spark_client.read(format, options) # type: ignore
def test_sql(self, target_df: DataFrame) -> None:
# arrange
spark_client = SparkClient()
create_temp_view(target_df, "test")
# act
result_df = spark_client.sql("select * from test")
# assert
assert result_df.collect() == target_df.collect()
def test_write_dataframe_invalid_params(
self, target_df: DataFrame, format: Optional[str], mode: Union[str, int]
) -> None:
# arrange
spark_client = SparkClient()
# act and assert
with pytest.raises(ValueError):
spark_client.write_dataframe(
dataframe=target_df, format_=format, mode=mode # type: ignore
)
def test_create_temporary_view(self, target_df: DataFrame,
spark_session: SparkSession) -> None:
# arrange
spark_client = SparkClient()
# act
spark_client.create_temporary_view(target_df, "temp_view")
result_df = spark_session.table("temp_view")
# assert
assert_dataframe_equality(target_df, result_df)
def __init__(
self,
source: Source,
feature_set: FeatureSet,
sink: Sink,
spark_client: SparkClient = None,
):
self.source = source
self.feature_set = feature_set
self.sink = sink
self.spark_client = spark_client or SparkClient()
def test_build_table_expression(self, table, database,
target_table_expression):
# arrange
spark_client = SparkClient()
# act
result_table_expression = SparkTableSchemaCompatibilityHook(
spark_client, table, database).table_expression
# assert
assert target_table_expression == result_table_expression