Use configured timestamp column names in RemoteOnlineStore.online_write_batch (#6595) - #6773
Open
gouravparmar17 wants to merge 4 commits into
Open
Use configured timestamp column names in RemoteOnlineStore.online_write_batch (#6595)#6773gouravparmar17 wants to merge 4 commits into
gouravparmar17 wants to merge 4 commits into
Conversation
Refactor batch writing logic to use dynamic column names and simplify data extraction.
Test the online_write_batch method for custom timestamp columns in RemoteOnlineStore.
Updated the test to verify custom timestamp column names in online_write_batch.
Contributor
There was a problem hiding this comment.
Pull request overview
This PR updates the Python Remote Online Store client write path to use the FeatureView.batch_source-configured timestamp_field and created_timestamp_column names (instead of hardcoded "event_timestamp" / "created"), so the feature server can reconstruct the DataFrame with the expected schema during /write-to-online-store.
Changes:
- Resolve event/created timestamp column names dynamically from the feature view’s batch source in
RemoteOnlineStore.online_write_batch. - Update unit coverage to validate that custom timestamp column names are used in the outgoing request body.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 5 comments.
| File | Description |
|---|---|
| sdk/python/feast/infra/online_stores/remote.py | Adjusts online_write_batch request DataFrame construction to use configured timestamp column names for remote writes. |
| sdk/python/tests/unit/infra/online_store/test_remote_online_store.py | Updates the remote online write batch test to cover custom timestamp_field / created_timestamp_column behavior. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+268
to
+278
| 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)) |
Comment on lines
289
to
+293
| 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) |
Comment on lines
+293
to
294
| post_remote_online_write(config=config, req_body=req_body) | ||
|
|
Comment on lines
743
to
+781
| @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"] |
Comment on lines
+255
to
+265
| # 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" | ||
| ) |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What this PR does / why we need it:
Updates
RemoteOnlineStore.online_write_batchto dynamically resolvetimestamp_fieldandcreated_timestamp_columnfrom the batch source if configured, instead of hardcoding"event_timestamp"and"created". This enables feature views with custom timestamp column names to write to the remote online store without failing Arrow schema conversion.Which issue(s) this PR fixes:
Fixes #6595
Checks
git commit -s)Testing Strategy
-Unit tests