Skip to content

Commit a42dc85

Browse files
adarshsmntkathole
adarshsm
authored and
ntkathole
committed
fix: Resolve write_to_offline_store feature view with a single registry lookup
write_to_offline_store resolved the feature view with a try/except chain that called get_stream_feature_view, then get_feature_view, then get_label_view in turn. A plain FeatureView -- the common case -- never matches the first lookup, so it was always issued, failed with FeatureViewNotFoundException, and discarded before the lookup that could succeed ran. On a RemoteRegistry each attempt is a gRPC round-trip to the registry server (no client-side cache), so every batch on this per-batch write path paid one or more guaranteed-miss RPCs before the real one. allow_registry_cache=True, the default, does not avoid it on a remote registry. Resolve the feature view with a single registry.get_any_feature_view lookup instead -- the unified accessor added in #4235 for exactly the case where a caller holds only a name. It covers FeatureView, StreamFeatureView, OnDemandFeatureView, and LabelView, so it is a behaviour-preserving replacement for the chain while collapsing up to three registry lookups into one. Fixes #6671. Signed-off-by: adarshsm <24850536+adarshsm@users.noreply.github.com>
1 parent 3425783 commit a42dc85

2 files changed

Lines changed: 65 additions & 14 deletions

File tree

‎sdk/python/feast/feature_store.py‎

Lines changed: 13 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -3806,20 +3806,19 @@ def write_to_offline_store(
38063806
Fails if the dataframe columns do not match the columns of the batch data source. Optionally
38073807
reorders the columns of the dataframe to match.
38083808
"""
3809-
# TODO: restrict this to work with online StreamFeatureViews and validate the FeatureView type
3810-
try:
3811-
feature_view: FeatureView = self.get_stream_feature_view(
3812-
feature_view_name, allow_registry_cache=allow_registry_cache
3813-
)
3814-
except FeatureViewNotFoundException:
3815-
try:
3816-
feature_view = self.get_feature_view(
3817-
feature_view_name, allow_registry_cache=allow_registry_cache
3818-
)
3819-
except FeatureViewNotFoundException:
3820-
feature_view = self.get_label_view( # type: ignore[assignment]
3821-
feature_view_name, allow_registry_cache=allow_registry_cache
3822-
)
3809+
# Resolve the feature view with a single registry lookup regardless of its
3810+
# type. The previous try/except chain tried get_stream_feature_view, then
3811+
# get_feature_view, then get_label_view in turn, so the common plain
3812+
# FeatureView always paid one guaranteed-miss lookup first. On a
3813+
# RemoteRegistry each miss is a wasted gRPC round-trip on this per-batch
3814+
# write path (see #6671).
3815+
# TODO: validate that the resolved feature view type supports offline writes.
3816+
feature_view = cast(
3817+
FeatureView,
3818+
self.registry.get_any_feature_view(
3819+
feature_view_name, self.project, allow_cache=allow_registry_cache
3820+
),
3821+
)
38233822

38243823
provider = self._get_provider()
38253824
# Get columns of the batch source and the input dataframe.

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

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,12 @@
11
from dataclasses import dataclass
22
from typing import Dict, List
3+
from unittest.mock import MagicMock, patch
34

5+
import pandas as pd
46
import pytest
57

68
from feast import utils
9+
from feast.feature_store import FeatureStore
710
from feast.protos.feast.types.Value_pb2 import Value
811

912

@@ -129,3 +132,52 @@ def test_get_unique_entities_missing_all_join_keys_error():
129132
in error_message
130133
)
131134
assert "Provided join_key_values: ['entity_3']" in error_message
135+
136+
137+
def test_write_to_offline_store_resolves_feature_view_with_single_lookup():
138+
"""``write_to_offline_store`` must resolve the feature view with a single
139+
registry lookup via ``get_any_feature_view`` rather than the legacy
140+
per-type try/except chain.
141+
142+
The old chain called ``get_stream_feature_view`` first, so the common plain
143+
``FeatureView`` always paid a guaranteed-miss lookup before the one that
144+
could succeed. On a ``RemoteRegistry`` each miss is a wasted gRPC round-trip
145+
on this per-batch write path (#6671).
146+
"""
147+
store = FeatureStore.__new__(FeatureStore)
148+
149+
feature_view = MagicMock()
150+
feature_view.name = "driver_hourly_stats"
151+
feature_view.batch_source = MagicMock()
152+
153+
registry = MagicMock()
154+
registry.get_any_feature_view.return_value = feature_view
155+
store._registry = registry
156+
157+
current_project = MagicMock()
158+
current_project.get.return_value = "test_project"
159+
store._current_project = current_project
160+
store.config = MagicMock()
161+
162+
provider = MagicMock()
163+
provider.get_table_column_names_and_types_from_data_source.return_value = [
164+
("driver_id", "INT64"),
165+
]
166+
167+
df = pd.DataFrame({"driver_id": [1, 2, 3]})
168+
169+
with patch.object(FeatureStore, "_get_provider", return_value=provider):
170+
store.write_to_offline_store("driver_hourly_stats", df, reorder_columns=False)
171+
172+
# A single unified lookup, with the default allow_registry_cache=True
173+
# forwarded as allow_cache.
174+
registry.get_any_feature_view.assert_called_once_with(
175+
"driver_hourly_stats", "test_project", allow_cache=True
176+
)
177+
# None of the legacy per-type getters are used, so there is no
178+
# guaranteed-miss lookup before the real one.
179+
registry.get_stream_feature_view.assert_not_called()
180+
registry.get_feature_view.assert_not_called()
181+
registry.get_label_view.assert_not_called()
182+
183+
provider.ingest_df_to_offline_store.assert_called_once()

0 commit comments

Comments
 (0)