diff --git a/sdk/python/feast/infra/online_stores/online_store.py b/sdk/python/feast/infra/online_stores/online_store.py index cdf06639fe0..69d596dc6bf 100644 --- a/sdk/python/feast/infra/online_stores/online_store.py +++ b/sdk/python/feast/infra/online_stores/online_store.py @@ -237,6 +237,7 @@ def get_online_features( online_features_response, full_feature_names, output_len, + first_idxs, grouped_refs, registry, project, @@ -522,6 +523,7 @@ async def get_online_features_async( online_features_response, full_feature_names, output_len, + first_idxs, grouped_refs, registry, project, @@ -606,12 +608,16 @@ def _try_precomputed_fast_path( online_features_response: Any, full_feature_names: bool, num_rows: int, + entity_row_indices: Sequence[Sequence[int]], grouped_refs: List, registry: BaseRegistry, project: str, ) -> bool: """Build the response from pre-computed vectors. + ``entity_row_indices`` maps each deduplicated entity vector to its original + request rows, preserving input order and repeated entities in the response. + Returns True if the fast path succeeded for ALL entities, False otherwise (caller should fall back to per-FV reads). """ @@ -695,8 +701,10 @@ def _try_precomputed_fast_path( now_secs = _time_mod.time() if any_ttl else 0.0 - for row_idx, vec in enumerate(vectors): - ts_list[row_idx] = vec.precomputed_at + for entity_idx, vec in enumerate(vectors): + destination_row_indices = entity_row_indices[entity_idx] + for row_idx in destination_row_indices: + ts_list[row_idx] = vec.precomputed_at # Build per-FV expiry flags once per entity (not per feature). fv_expired: Optional[Dict[str, bool]] = None @@ -720,14 +728,16 @@ def _try_precomputed_fast_path( stored_values = vec.values for out_idx in range(n_features): src_idx = reorder_map[out_idx] if reorder_map else out_idx - feat_values[out_idx][row_idx] = stored_values[src_idx] + status = PRESENT if fv_expired: feat_fv = feat_fv_names[out_idx] if feat_fv and fv_expired.get(feat_fv, False): - feat_statuses[out_idx][row_idx] = OUTSIDE_MAX_AGE - continue - feat_statuses[out_idx][row_idx] = PRESENT + status = OUTSIDE_MAX_AGE + + for row_idx in destination_row_indices: + feat_values[out_idx][row_idx] = stored_values[src_idx] + feat_statuses[out_idx][row_idx] = status online_features_response.metadata.feature_names.val.extend( expected_feature_names diff --git a/sdk/python/tests/unit/test_precomputed_feature_vectors.py b/sdk/python/tests/unit/test_precomputed_feature_vectors.py index 80b0b83f53f..ec7b4721ef5 100644 --- a/sdk/python/tests/unit/test_precomputed_feature_vectors.py +++ b/sdk/python/tests/unit/test_precomputed_feature_vectors.py @@ -1,5 +1,6 @@ """Tests for pre-computed feature vectors (issue #6185).""" +import asyncio from datetime import datetime, timedelta, timezone from unittest.mock import MagicMock @@ -272,13 +273,15 @@ class TestSchemaMismatchFallback: def _fast_path(self, blobs, expected_names, grouped_refs=None, num_rows=None): from feast.infra.online_stores.online_store import OnlineStore + resolved_num_rows = num_rows or len(blobs) response = GetOnlineFeaturesResponse(results=[]) result = OnlineStore._try_precomputed_fast_path( blobs=blobs, expected_feature_names=expected_names, online_features_response=response, full_feature_names=True, - num_rows=num_rows or len(blobs), + num_rows=resolved_num_rows, + entity_row_indices=tuple([i] for i in range(len(blobs))), grouped_refs=grouped_refs or [], registry=MagicMock(), project="test", @@ -370,16 +373,28 @@ def test_schema_mismatch_in_second_entity_returns_false(self): class TestFastPathMultiEntity: - def _fast_path(self, blobs, expected_names, grouped_refs=None, num_rows=None): + def _fast_path( + self, + blobs, + expected_names, + grouped_refs=None, + num_rows=None, + entity_row_indices=None, + ): from feast.infra.online_stores.online_store import OnlineStore + resolved_num_rows = num_rows or len(blobs) + if entity_row_indices is None: + entity_row_indices = tuple([i] for i in range(len(blobs))) + response = GetOnlineFeaturesResponse(results=[]) result = OnlineStore._try_precomputed_fast_path( blobs=blobs, expected_feature_names=expected_names, online_features_response=response, full_feature_names=True, - num_rows=num_rows or len(blobs), + num_rows=resolved_num_rows, + entity_row_indices=entity_row_indices, grouped_refs=grouped_refs or [], registry=MagicMock(), project="test", @@ -418,6 +433,39 @@ def test_multiple_entities_all_present(self): assert a_values == pytest.approx([1.0, 2.0, 3.0]) assert b_values == pytest.approx([10.0, 20.0, 30.0]) + def test_restores_original_order_and_duplicate_entities(self): + timestamps = [ + _make_timestamp(datetime(2025, 1, day, tzinfo=timezone.utc)) + for day in (1, 2, 3) + ] + vectors = [ + _make_vector( + feature_names=["fv1__a"], + values=[ValueProto(float_val=float(entity_id))], + precomputed_at=timestamps[entity_id - 1], + ) + for entity_id in (1, 2, 3) + ] + + ok, response = self._fast_path( + [vector.SerializeToString() for vector in vectors], + ["fv1__a"], + num_rows=4, + entity_row_indices=([1], [3], [0, 2]), + ) + + assert ok is True + assert [value.float_val for value in response.results[0].values] == [ + 3.0, + 1.0, + 3.0, + 2.0, + ] + assert list(response.results[0].statuses) == [FieldStatus.PRESENT] * 4 + assert [ + timestamp.seconds for timestamp in response.results[0].event_timestamps + ] == [timestamps[index].seconds for index in (2, 0, 2, 1)] + def test_statuses_are_present(self): vec = _make_vector( feature_names=["fv1__a"], @@ -439,6 +487,89 @@ def test_feature_names_in_metadata(self): assert list(response.metadata.feature_names.val) == ["fv1__x", "fv2__y"] +# ═══════════════════════════════════════════════════════════════════════════════ +# Public online retrieval: entity row ordering +# ═══════════════════════════════════════════════════════════════════════════════ + + +@pytest.mark.parametrize("use_async", [False, True], ids=["sync", "async"]) +def test_precomputed_retrieval_preserves_entity_order_and_duplicates(use_async): + from feast import Entity + from feast.infra.online_stores.online_store import OnlineStore + from feast.value_type import ValueType + + entity = Entity( + name="driver", + join_keys=["driver_id"], + value_type=ValueType.INT64, + ) + feature_view = FeatureView( + name="driver_stats", + entities=[entity], + schema=[ + Field(name="driver_id", dtype=Int64), + Field(name="score", dtype=Float32), + ], + source=_make_file_source("driver_stats_src"), + ) + feature_service = FeatureService( + name="driver_service", + features=[feature_view], + precompute_online=True, + ) + + class PrecomputedStore(OnlineStore): + def online_write_batch(self, config, table, data, progress): + pass + + def online_read(self, config, table, entity_keys, requested_features=None): + raise AssertionError("regular online reads should not be used") + + def update(self, *args, **kwargs): + pass + + def teardown(self, *args, **kwargs): + pass + + def read_precomputed_vectors( + self, config, feature_service_name, project, entity_keys + ): + blobs = [] + for entity_key in entity_keys: + entity_id = entity_key.entity_values[0].int64_val + vector = _make_vector( + feature_names=["driver_stats__score"], + values=[ValueProto(float_val=float(entity_id * 10))], + ) + blobs.append(vector.SerializeToString()) + return blobs + + registry = MagicMock() + registry.cached_registry_proto_created = None + registry.enable_online_versioning = False + registry.get_feature_service.return_value = feature_service + registry.get_any_feature_view.return_value = feature_view + registry.list_entities.return_value = [entity] + + request_kwargs = { + "config": MagicMock(), + "features": feature_service, + "entity_rows": {"driver_id": [3, 1, 3, 2]}, + "registry": registry, + "project": "test", + "full_feature_names": True, + } + store = PrecomputedStore() + if use_async: + response = asyncio.run(store.get_online_features_async(**request_kwargs)) + else: + response = store.get_online_features(**request_kwargs) + + result = response.to_dict() + assert result["driver_id"] == [3, 1, 3, 2] + assert result["driver_stats__score"] == [30.0, 10.0, 30.0, 20.0] + + # ═══════════════════════════════════════════════════════════════════════════════ # TTL enforcement in fast path # ═══════════════════════════════════════════════════════════════════════════════ @@ -472,6 +603,7 @@ def _fast_path_with_ttl(self, fv_event_dt, fv_ttl_seconds, fv_name="fv1"): online_features_response=response, full_feature_names=True, num_rows=1, + entity_row_indices=([0],), grouped_refs=grouped_refs, registry=MagicMock(), project="test", @@ -515,6 +647,7 @@ def test_zero_ttl_means_no_expiry(self): online_features_response=response, full_feature_names=True, num_rows=1, + entity_row_indices=([0],), grouped_refs=[(fv, ["f1"])], registry=MagicMock(), project="test", @@ -548,6 +681,7 @@ def test_mixed_ttl_some_expired_some_fresh(self): online_features_response=response, full_feature_names=True, num_rows=1, + entity_row_indices=([0],), grouped_refs=[(fv1, ["f1"]), (fv2, ["f2"])], registry=MagicMock(), project="test",