def test_filter_latest_without_id_columns(
self, feature_set_dataframe, cassandra_config
):
# with
writer = OnlineFeatureStoreWriter(cassandra_config)
# then
with pytest.raises(ValueError, match="must provide the unique identifiers"):
_ = writer.filter_latest(feature_set_dataframe, id_columns=[])
# then
with pytest.raises(KeyError, match="not found"):
_ = writer.filter_latest(
feature_set_dataframe.drop("id"), id_columns=["id"]
)
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_filter_latest_without_ts(
self, feature_set_dataframe_without_ts, cassandra_config
):
# with
writer = OnlineFeatureStoreWriter(cassandra_config)
# then
with pytest.raises(KeyError):
_ = writer.filter_latest(
feature_set_dataframe_without_ts, id_columns=["id"]
)
def test_filter_latest(
self, feature_set_dataframe, latest_feature_set_dataframe, cassandra_config
):
# with
writer = OnlineFeatureStoreWriter(cassandra_config)
# when
result_df = writer.filter_latest(feature_set_dataframe, id_columns=["id"])
sort_columns = result_df.schema.fieldNames()
# then
assert sorted(
latest_feature_set_dataframe.select(*sort_columns).collect()
) == sorted(result_df.select(*sort_columns).collect())
def test_write_stream(self, feature_set, has_checkpoint, monkeypatch):
# arrange
spark_client = SparkClient()
spark_client.write_stream = Mock()
spark_client.write_dataframe = Mock()
spark_client.write_stream.return_value = Mock(spec=StreamingQuery)
dataframe = Mock(spec=DataFrame)
dataframe.isStreaming = True
if has_checkpoint:
monkeypatch.setenv("STREAM_CHECKPOINT_PATH", "test")
cassandra_config = CassandraConfig(keyspace="feature_set")
target_checkpoint_path = (
"test/entity/feature_set"
if cassandra_config.stream_checkpoint_path
else None
)
writer = OnlineFeatureStoreWriter(cassandra_config)
writer.filter_latest = Mock()
# act
stream_handler = writer.write(feature_set, dataframe, spark_client)
# assert
assert isinstance(stream_handler, StreamingQuery)
spark_client.write_stream.assert_any_call(
dataframe,
processing_time=cassandra_config.stream_processing_time,
output_mode=cassandra_config.stream_output_mode,
checkpoint_path=target_checkpoint_path,
format_=cassandra_config.format_,
mode=cassandra_config.mode,
**cassandra_config.get_options(table=feature_set.name),
)
writer.filter_latest.assert_not_called()
spark_client.write_dataframe.assert_not_called()