Skip to content

Commit ecad170

Browse files
franciscojavierarceofranciscojavierarceo
authored
feat: Make transform on writes configurable for ingestion (#5283)
* feat: Make transform on writes configurable for batch ingestion Signed-off-by: Francisco Javier Arceo <[email protected]> * update test and fix bug to support skipping transformation on writes for ODFV Signed-off-by: Francisco Javier Arceo <[email protected]> --------- Signed-off-by: Francisco Javier Arceo <[email protected]>
1 parent 0ee3a1e commit ecad170

3 files changed

Lines changed: 63 additions & 2 deletions

File tree

‎sdk/python/feast/feature_server.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,13 +61,15 @@ class WriteToFeatureStoreRequest(BaseModel):
6161
feature_view_name: str
6262
df: dict
6363
allow_registry_cache: bool = True
64+
transform_on_write: bool = True
6465

6566

6667
class PushFeaturesRequest(BaseModel):
6768
push_source_name: str
6869
df: dict
6970
allow_registry_cache: bool = True
7071
to: str = "online"
72+
transform_on_write: bool = True
7173

7274

7375
class MaterializeRequest(BaseModel):
@@ -302,6 +304,7 @@ async def push(request: PushFeaturesRequest) -> None:
302304
df=df,
303305
allow_registry_cache=request.allow_registry_cache,
304306
to=to,
307+
transform_on_write=request.transform_on_write,
305308
)
306309

307310
should_push_async = (
@@ -336,6 +339,7 @@ def write_to_online_store(request: WriteToFeatureStoreRequest) -> None:
336339
feature_view_name=feature_view_name,
337340
df=df,
338341
allow_registry_cache=allow_registry_cache,
342+
transform_on_write=request.transform_on_write,
339343
)
340344

341345
@app.get("/health")

‎sdk/python/feast/feature_store.py‎

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1464,6 +1464,7 @@ def push(
14641464
df: pd.DataFrame,
14651465
allow_registry_cache: bool = True,
14661466
to: PushMode = PushMode.ONLINE,
1467+
transform_on_write: bool = True,
14671468
):
14681469
"""
14691470
Push features to a push source. This updates all the feature views that have the push source as stream source.
@@ -1473,13 +1474,17 @@ def push(
14731474
df: The data being pushed.
14741475
allow_registry_cache: Whether to allow cached versions of the registry.
14751476
to: Whether to push to online or offline store. Defaults to online store only.
1477+
transform_on_write: Whether to transform the data before pushing.
14761478
"""
14771479
for fv in self._fvs_for_push_source_or_raise(
14781480
push_source_name, allow_registry_cache
14791481
):
14801482
if to == PushMode.ONLINE or to == PushMode.ONLINE_AND_OFFLINE:
14811483
self.write_to_online_store(
1482-
fv.name, df, allow_registry_cache=allow_registry_cache
1484+
fv.name,
1485+
df,
1486+
allow_registry_cache=allow_registry_cache,
1487+
transform_on_write=transform_on_write,
14831488
)
14841489
if to == PushMode.OFFLINE or to == PushMode.ONLINE_AND_OFFLINE:
14851490
self.write_to_offline_store(
@@ -1521,6 +1526,7 @@ def _get_feature_view_and_df_for_online_write(
15211526
df: Optional[pd.DataFrame] = None,
15221527
inputs: Optional[Union[Dict[str, List[Any]], pd.DataFrame]] = None,
15231528
allow_registry_cache: bool = True,
1529+
transform_on_write: bool = True,
15241530
):
15251531
feature_view_dict = {
15261532
fv_proto.name: fv_proto
@@ -1553,6 +1559,7 @@ def _get_feature_view_and_df_for_online_write(
15531559
if (
15541560
isinstance(feature_view, OnDemandFeatureView)
15551561
and feature_view.write_to_online_store
1562+
and transform_on_write
15561563
):
15571564
if (
15581565
feature_view.mode == "python"
@@ -1638,6 +1645,7 @@ def write_to_online_store(
16381645
df: Optional[pd.DataFrame] = None,
16391646
inputs: Optional[Union[Dict[str, List[Any]], pd.DataFrame]] = None,
16401647
allow_registry_cache: bool = True,
1648+
transform_on_write: bool = True,
16411649
):
16421650
"""
16431651
Persists a dataframe to the online store.
@@ -1647,13 +1655,15 @@ def write_to_online_store(
16471655
df: The dataframe to be persisted.
16481656
inputs: Optional the dictionary object to be written
16491657
allow_registry_cache (optional): Whether to allow retrieving feature views from a cached registry.
1658+
transform_on_write (optional): Whether to transform the data before pushing.
16501659
"""
16511660

16521661
feature_view, df = self._get_feature_view_and_df_for_online_write(
16531662
feature_view_name=feature_view_name,
16541663
df=df,
16551664
inputs=inputs,
16561665
allow_registry_cache=allow_registry_cache,
1666+
transform_on_write=transform_on_write,
16571667
)
16581668
provider = self._get_provider()
16591669
provider.ingest_df(feature_view, df)

‎sdk/python/tests/unit/test_on_demand_python_transformation.py‎

Lines changed: 48 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -874,7 +874,8 @@ def test_stored_writes(self):
874874
assert driver_stats_fv.entity_columns == []
875875

876876
ODFV_STRING_CONSTANT = "guaranteed constant"
877-
ODFV_OTHER_STRING_CONSTANT = "somethign else"
877+
ODFV_OTHER_STRING_CONSTANT = "something else"
878+
ODFV_UNTRANSFORMED_STRING_CONSTANT = "also something else"
878879

879880
@on_demand_feature_view(
880881
entities=[driver],
@@ -1069,6 +1070,52 @@ def python_stored_writes_feature_view(
10691070
assert online_odfv_python_response["string_constant"] != [
10701071
ODFV_OTHER_STRING_CONSTANT
10711072
]
1073+
odfv_entity_rows_to_write_no_transform = [
1074+
{
1075+
"driver_id": 1003,
1076+
"counter": 10,
1077+
"conv_rate": 0.25,
1078+
"acc_rate": 0.50,
1079+
"input_datetime": current_datetime,
1080+
"string_constant": ODFV_UNTRANSFORMED_STRING_CONSTANT,
1081+
}
1082+
]
1083+
odfv_entity_rows_to_read_no_transform = [
1084+
{
1085+
"driver_id": 1003,
1086+
"conv_rate_plus_acc": 7, # note how this is not the correct value and would be calculate on demand
1087+
"conv_rate": 0.25,
1088+
"acc_rate": 0.50,
1089+
"counter": 0,
1090+
"input_datetime": current_datetime,
1091+
"string_constant": ODFV_UNTRANSFORMED_STRING_CONSTANT,
1092+
}
1093+
]
1094+
print("storing ODFV features")
1095+
self.store.write_to_online_store(
1096+
feature_view_name="python_stored_writes_feature_view",
1097+
df=odfv_entity_rows_to_write_no_transform,
1098+
transform_on_write=False,
1099+
)
1100+
online_odfv_python_response_no_transform = self.store.get_online_features(
1101+
entity_rows=odfv_entity_rows_to_read_no_transform,
1102+
features=[
1103+
"python_stored_writes_feature_view:conv_rate_plus_acc",
1104+
"python_stored_writes_feature_view:current_datetime",
1105+
"python_stored_writes_feature_view:counter",
1106+
"python_stored_writes_feature_view:input_datetime",
1107+
"python_stored_writes_feature_view:string_constant",
1108+
],
1109+
).to_dict()
1110+
# note these are approximately correct by
1111+
assert online_odfv_python_response_no_transform == {
1112+
"driver_id": [1003],
1113+
"counter": [10],
1114+
"conv_rate_plus_acc": [None],
1115+
"input_datetime": [current_datetime.replace(microsecond=0)],
1116+
"string_constant": [ODFV_UNTRANSFORMED_STRING_CONSTANT],
1117+
"current_datetime": [None],
1118+
}
10721119

10731120
def test_stored_writes_with_explode(self):
10741121
with tempfile.TemporaryDirectory() as data_dir:

0 commit comments

Comments
 (0)