| FazBrowse GitHub Viewer | Trending | | Home |
| Tools: [Download Repo ZIP] [Original HTTPS Page] |
1 parent 3425783 commit a42dc85
2 files changed
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -3806,20 +3806,19 @@ def write_to_offline_store( | |||
| 3806 | 3806 | Fails if the dataframe columns do not match the columns of the batch data source. Optionally | |
| 3807 | 3807 | reorders the columns of the dataframe to match. | |
| 3808 | 3808 | """ | |
| 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 | + ) | ||
| 3823 | 3822 | ||
| 3824 | 3823 | provider = self._get_provider() | |
| 3825 | 3824 | # Get columns of the batch source and the input dataframe. | |
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -1,9 +1,12 @@ | |||
| 1 | 1 | from dataclasses import dataclass | |
| 2 | 2 | from typing import Dict, List | |
| 3 | + from unittest.mock import MagicMock, patch | ||
| 3 | 4 | ||
| 5 | + import pandas as pd | ||
| 4 | 6 | import pytest | |
| 5 | 7 | ||
| 6 | 8 | from feast import utils | |
| 9 | + from feast.feature_store import FeatureStore | ||
| 7 | 10 | from feast.protos.feast.types.Value_pb2 import Value | |
| 8 | 11 | ||
| 9 | 12 | ||
@@ -129,3 +132,52 @@ def test_get_unique_entities_missing_all_join_keys_error(): | |||
| 129 | 132 | in error_message | |
| 130 | 133 | ) | |
| 131 | 134 | 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() | ||
| Back | FazBrowse Home | New Git URL |
0 commit comments