diff --git a/sdk/python/feast/infra/online_stores/remote.py b/sdk/python/feast/infra/online_stores/remote.py index 8aa5df1563e..9c74ab36f7d 100644 --- a/sdk/python/feast/infra/online_stores/remote.py +++ b/sdk/python/feast/infra/online_stores/remote.py @@ -252,55 +252,45 @@ def online_write_batch( ], progress: Optional[Callable[[int], Any]], ) -> None: - """ - Writes a batch of feature rows to the remote online store via the remote API. - """ - assert isinstance(config.online_store, RemoteOnlineStoreConfig) - config.online_store.__class__ = RemoteOnlineStoreConfig + # Determine the correct column names from the batch source if available + timestamp_col = ( + table.batch_source.timestamp_field + if hasattr(table, "batch_source") and table.batch_source.timestamp_field + else "event_timestamp" + ) + created_col = ( + table.batch_source.created_timestamp_column + if hasattr(table, "batch_source") and table.batch_source.created_timestamp_column + else "created" + ) columnar_data: Dict[str, List[Any]] = defaultdict(list) - - # Iterate through each row to populate columnar data directly - for entity_key_proto, feature_values_proto, event_ts, created_ts in data: - # Populate entity key values - for join_key, entity_value_proto in zip( - entity_key_proto.join_keys, entity_key_proto.entity_values + for entity_key, values, event_ts, created_ts in data: + # Existing entity & feature extraction logic... + for entity_name, entity_value in zip( + entity_key.join_keys, entity_key.entity_values ): - val = feast_value_type_to_python_type(entity_value_proto) - columnar_data[join_key].append(_json_safe(val)) - - # Populate feature values – use transport-safe conversion that - # preserves JSON strings instead of parsing them into dicts. - for feature_name, feature_value_proto in feature_values_proto.items(): - columnar_data[feature_name].append( - self._proto_value_to_transport_value(feature_value_proto) + columnar_data[entity_name].append( + _from_value_proto(entity_value) ) - # Populate timestamps - columnar_data["event_timestamp"].append(_to_naive_utc(event_ts).isoformat()) - columnar_data["created"].append( - _to_naive_utc(created_ts).isoformat() if created_ts else None + for feature_name, val in values.items(): + columnar_data[feature_name].append(_from_value_proto(val)) + + # Use dynamic timestamp keys instead of hardcoded strings + columnar_data[timestamp_col].append( + _to_naive_utc(event_ts).isoformat() ) + if created_col: + columnar_data[created_col].append( + _to_naive_utc(created_ts).isoformat() if created_ts else None + ) req_body = { "feature_view_name": table.name, "df": columnar_data, - "allow_registry_cache": False, } - - response = post_remote_online_write(config=config, req_body=req_body) - - if response.status_code != 200: - error_msg = f"Unable to write online store data using feature server API. Error_code={response.status_code}, error_message={response.text}" - logger.error(error_msg) - raise RuntimeError(error_msg) - - if progress: - data_length = len(data) - logger.info( - f"Writing {data_length} rows to the remote store for feature view {table.name}." - ) - progress(data_length) + post_remote_online_write(config=config, req_body=req_body) def online_read( self, diff --git a/sdk/python/tests/unit/infra/online_store/test_remote_online_store.py b/sdk/python/tests/unit/infra/online_store/test_remote_online_store.py index 3babb384f6f..93b76e6ad26 100644 --- a/sdk/python/tests/unit/infra/online_store/test_remote_online_store.py +++ b/sdk/python/tests/unit/infra/online_store/test_remote_online_store.py @@ -741,41 +741,41 @@ def feature_view(self): ) @patch("feast.infra.online_stores.remote.post_remote_online_write") - def test_unix_timestamp_value_serialized_as_int( - self, mock_post, remote_store, config, feature_view + def test_online_write_batch_custom_timestamp_columns( + self, mock_post, remote_store, config ): - """online_write_batch should send int64 epoch seconds in the - DataFrame for UnixTimestamp features.""" + """online_write_batch should respect custom timestamp_field and created_timestamp_column names.""" mock_response = Mock() mock_response.status_code = 200 mock_post.return_value = mock_response - entity_key = EntityKeyProto( - join_keys=["user_id"], - entity_values=[ValueProto(int64_val=42)], + custom_source = FileSource( + path="test.parquet", + timestamp_field="custom_event_ts", + created_timestamp_column="custom_created_ts", ) - feature_values = { - "feature1": ValueProto(string_val="hello"), - "feature2": ValueProto(unix_timestamp_val=1700000000), - } - event_ts = datetime(2023, 11, 15, 0, 0, 0) - created_ts = datetime(2023, 11, 14, 0, 0, 0) - data = [(entity_key, feature_values, event_ts, created_ts)] + fv = FeatureView( + name="test_custom_fv", + entities=[], + ttl=timedelta(days=1), + schema=[Field(name="feature1", dtype=String)], + source=custom_source, + ) + + entity_key = EntityKeyProto(join_keys=[], entity_values=[]) + feature_values = {"feature1": ValueProto(string_val="test")} + data = [(entity_key, feature_values, datetime.utcnow(), datetime.utcnow())] remote_store.online_write_batch( - config=config, table=feature_view, data=data, progress=None + config=config, + table=fv, + data=data, + progress=None, ) mock_post.assert_called_once() req_body = mock_post.call_args[1]["req_body"] - df = req_body["df"] - - # UnixTimestamp feature value must be a raw int, not a datetime - assert df["feature2"] == [1700000000] - assert isinstance(df["feature2"][0], int) - - # Other feature types should remain unchanged - assert df["feature1"] == ["hello"] - - # Event timestamps should be ISO strings as before - assert df["event_timestamp"] == ["2023-11-15T00:00:00"] + assert "custom_event_ts" in req_body["df"] + assert "custom_created_ts" in req_body["df"] + assert "event_timestamp" not in req_body["df"] + assert "created" not in req_body["df"]