From 51b75b53ffbcf2eb823ff6b411f1af85f6ae0a55 Mon Sep 17 00:00:00 2001 From: NopeCharon <115794254+NopeCharon@users.noreply.github.com> Date: Wed, 11 Nov 2026 17:19:31 +0000 Subject: [PATCH] feat: Batch DynamoDB online reads across feature views The generic OnlineStore._read_features_per_fv loop reads one feature view after another, so a sync request spanning N feature views costs at least N sequential BatchGetItem round trips, and an async one at least N concurrent requests. Redis and Postgres already override that seam to batch; DynamoDB now does too. Each feature view lives in its own DynamoDB table, and a single BatchGetItem request can read from several tables at once. The override resolves every view's entity ids up front, packs all of the keys into as few requests as possible (at most batch_size keys each, across all of their tables), sends those requests in parallel (a thread pool capped at max_read_workers on the sync path, asyncio.gather on the async path) and then splits the returned items back apart by table and entity id. Each table keeps its own ProjectionExpression and ConsistentRead setting. BatchGetItem rejects a request that repeats a key within a table, so views that map to the same table, such as one feature view requested under two aliases, are read once with the union of their requested features. Following the Postgres implementation, a single-view request keeps the generic path, and requests larger than the new max_batched_result_rows setting (entity rows x feature views, default 2048, 0 disables batching) also fall back to it. Past that size each view already fills requests of its own, so packing saves at most one partially filled request per view while holding every view's items at once. The item parsing in _process_batch_get_response moves into a shared _to_read_row helper so the batched and per-view paths decode items identically. Refs #3259 Signed-off-by: NopeCharon <115794254+NopeCharon@users.noreply.github.com> --- .../online-server-performance-tuning.md | 9 +- docs/reference/online-stores/dynamodb.md | 3 + .../feast/infra/online_stores/dynamodb.py | 332 ++++++++++++++- .../test_dynamodb_batched_read.py | 384 ++++++++++++++++++ 4 files changed, 710 insertions(+), 18 deletions(-) create mode 100644 sdk/python/tests/unit/infra/online_store/test_dynamodb_batched_read.py diff --git a/docs/how-to-guides/online-server-performance-tuning.md b/docs/how-to-guides/online-server-performance-tuning.md index 96ec1168b09..24d3c5d2238 100644 --- a/docs/how-to-guides/online-server-performance-tuning.md +++ b/docs/how-to-guides/online-server-performance-tuning.md @@ -35,6 +35,8 @@ When the server processes a `get_online_features()` call, it groups the requeste **Redis exception:** The Redis online store overrides `_read_features_per_fv()` to batch all `HMGET` commands across every feature view into a **single pipeline execution**. Because all feature views for the same entity share one Redis hash key, the number of Redis round trips is always **1**, regardless of how many feature views the request touches. This means the "fewer feature views" guideline is less critical for Redis than for other stores — but consolidating feature views still reduces serialization and protobuf overhead at the application layer. **PostgreSQL exception:** PostgreSQL stores each feature view in its own table and overrides `_read_features_per_fv()` to combine the per-view reads into a **single `UNION ALL` statement**, so a request touching any number of feature views costs **1** query on both the sync and async paths. This helps most where round trips dominate — a handful of entities against a remote database. Requests larger than `max_batched_result_rows` (entity rows × feature views, default 2048) fall back to one query per view, because at that size the saved round trips no longer offset holding every view's rows at once. A single-view request is never batched, since there is nothing to combine. Set `max_batched_result_rows: 0` in the online store config to turn batching off. As with Redis, consolidating feature views still reduces application-layer overhead. + +**DynamoDB exception:** DynamoDB stores each feature view in its own table and overrides `_read_features_per_fv()` to pack the keys of every feature view into **shared `BatchGetItem` requests**, since one request can read from several tables. A request touching any number of feature views costs as few as **1** request (up to `batch_size` keys per request, sent in parallel) on both the sync and async paths. As with PostgreSQL, requests larger than `max_batched_result_rows` (entity rows × feature views, default 2048) read one view at a time, single-view requests are never batched, and `max_batched_result_rows: 0` turns batching off. {% endhint %} ### Feature services are free (and can be faster) @@ -89,7 +91,7 @@ Requesting just `combined_score` triggers reads from **both** `driver_stats_fv` ## Pre-computed feature vectors -When a `get_online_features()` request touches multiple feature views, the server issues a separate store read per feature view. For services spanning 5–15+ feature views, this fan-out dominates latency — even where the store batches its reads (Redis, PostgreSQL), the protobuf deserialization and response-building overhead grows linearly with the number of views. +When a `get_online_features()` request touches multiple feature views, the server issues a separate store read per feature view. For services spanning 5–15+ feature views, this fan-out dominates latency — even where the store batches its reads (Redis, PostgreSQL, DynamoDB), the protobuf deserialization and response-building overhead grows linearly with the number of views. **Pre-computed feature vectors** solve this by storing all of a feature service's features for each entity as a single serialized blob. At read time, the server fetches one blob per entity instead of N reads per feature view, reducing the operation to O(1). @@ -281,7 +283,7 @@ The online store is the single largest factor in `get_online_features()` latency | Store | Typical p50 latency | Async read | Best for | Key trade-off | | ----- | ------------------- | ---------- | -------- | ------------- | | **Redis / Dragonfly** | < 1 ms | No (threadpool) | Ultra-low latency, high throughput; all FV reads batched into 1 pipeline | Requires in-memory capacity for your dataset | -| **DynamoDB** | 2–5 ms | Yes | Serverless, auto-scaling on AWS | Pay-per-request cost; batch API limits (100 items) | +| **DynamoDB** | 2–5 ms | Yes | Serverless, auto-scaling on AWS; FV reads packed into shared batch requests | Pay-per-request cost; batch API limits (100 items) | | **PostgreSQL** | 3–10 ms | No (threadpool) | Teams with existing Postgres infra; all FV reads batched into 1 query | Connection pooling needed at scale | | **MongoDB** | 2–5 ms | Yes | Flexible schema, async-native | Requires index tuning for large datasets | | **Aerospike** | < 1 ms | No (threadpool) | Ultra-low latency, hybrid memory (RAM + SSD), large datasets | Namespace must be pre-configured on the cluster | @@ -308,7 +310,7 @@ The feature server can read from the online store using either an **async** or * | Store | Async read | Async write | Notes | | ----- | ---------- | ----------- | ----- | -| **DynamoDB** | Yes | Yes | Uses `aiobotocore` for non-blocking I/O | +| **DynamoDB** | Yes | Yes | Uses `aiobotocore` for non-blocking I/O. `_read_features_per_fv` is overridden to pack all feature view reads into shared `BatchGetItem` requests | | **MongoDB** | Yes | Yes | Uses `motor` (async MongoDB driver) | | **PostgreSQL** | Implemented | No | Has `online_read_async` but does not yet advertise via `async_supported`; uses sync/threadpool path. `_read_features_per_fv` is overridden to batch all feature view reads into a single `UNION ALL` query | | **Redis** | Implemented | **Yes** | `online_read_async` and `online_write_batch_async` both implemented; uses sync/threadpool path for `get_online_features` (`_read_features_per_fv` overridden with batched single pipeline) | @@ -343,6 +345,7 @@ Key knobs: - **`batch_size`**: DynamoDB's `BatchGetItem` accepts up to 100 items per request. For 500 entities, this means 5 batches. Keep at 100 unless hitting the 16 MB response limit. - **`max_read_workers`**: Controls parallelism for batch reads. With 10 workers, those 5 batches run concurrently (~10 ms) instead of sequentially (~50 ms). +- **`max_batched_result_rows`**: Requests spanning several feature views pack every view's keys into shared `BatchGetItem` requests, up to this many entity rows × feature views (default 2048). Set to `0` to read one feature view at a time. - **`consistent_reads: false`**: Eventually consistent reads are faster and cheaper. Use `true` only if you need read-after-write consistency. - **`warmup_connections: true`**: Pre-warms the DynamoDB connection pool on server startup by making `max_pool_connections` concurrent lightweight calls (`describe_endpoints`). This avoids a cold-start latency penalty (~20ms) on the first feature requests, including concurrent ones. - **`max_pool_connections`**: Increase for high-throughput deployments to improve HTTP connection reuse to the DynamoDB endpoint. diff --git a/docs/reference/online-stores/dynamodb.md b/docs/reference/online-stores/dynamodb.md index 203cd407c1c..39ef887227a 100644 --- a/docs/reference/online-stores/dynamodb.md +++ b/docs/reference/online-stores/dynamodb.md @@ -49,6 +49,7 @@ online_store: | `table_name_template` | string | `{project}.{table_name}` | Template for table names | | `batch_size` | int | `100` | Number of items per BatchGetItem/BatchWriteItem request (max 100) | | `max_read_workers` | int | `10` | Maximum parallel threads for batch read operations. Higher values improve throughput for large batch reads but increase resource usage | +| `max_batched_result_rows` | int | `2048` | Largest request, in entity rows × feature views, whose feature views are read through shared BatchGetItem requests. Set to `0` to always read one feature view at a time | | `consistent_reads` | bool | `false` | Whether to use strongly consistent reads (higher latency, guaranteed latest data) | | `warmup_connections` | bool | `false` | Whether to pre-warm the async connection pool on startup with `max_pool_connections` concurrent lightweight calls (`describe_endpoints`) | | `tags` | dict | `null` | AWS resource tags added to each table | @@ -63,6 +64,8 @@ online_store: For high-throughput workloads with large entity counts, increase `max_read_workers` (up to 20-30) based on your DynamoDB capacity and network conditions. +**Batched Multi-Feature-View Reads**: Each feature view is stored in its own table, and a single BatchGetItem request can read from several tables. When a request spans several feature views, the keys of every view are packed into the same BatchGetItem requests (up to `batch_size` keys each), which are sent in parallel. A request for 10 feature views and 5 entities therefore costs **1** BatchGetItem call instead of 10, on both the sync and async read paths. Requests larger than `max_batched_result_rows` (entity rows × feature views) read one feature view at a time instead, because at that size each view already fills requests of its own and holding every view's items at once only raises peak memory. A single-view request is never batched. Set `max_batched_result_rows: 0` to turn batching off. + **Batch Size**: Increase `batch_size` up to 100 to reduce the number of API calls. However, larger batches may hit DynamoDB's 16MB response limit for tables with large feature values. **Connection Warmup**: The DynamoDB async client does not establish actual TCP/TLS connections to the AWS endpoint on initialization. The very first feature retrieval request is penalized with a cold-start overhead (~20ms). Setting `warmup_connections: true` opens up to `max_pool_connections` connections during server startup, one per concurrent `describe_endpoints` call. When DynamoDB is reached through a VPC endpoint, these calls need the `dynamodb:DescribeEndpoints` permission; without it the connections are still opened, but a warning is logged. diff --git a/sdk/python/feast/infra/online_stores/dynamodb.py b/sdk/python/feast/infra/online_stores/dynamodb.py index 6c4fd201e88..7c0ea7e9dff 100644 --- a/sdk/python/feast/infra/online_stores/dynamodb.py +++ b/sdk/python/feast/infra/online_stores/dynamodb.py @@ -19,10 +19,21 @@ from collections import OrderedDict, defaultdict from concurrent.futures import ThreadPoolExecutor from datetime import datetime -from typing import Any, Callable, Dict, List, Literal, Optional, Sequence, Tuple, Union +from typing import ( + Any, + Callable, + Dict, + List, + Literal, + NamedTuple, + Optional, + Sequence, + Tuple, + Union, +) from aiobotocore.config import AioConfig -from pydantic import StrictBool, StrictStr +from pydantic import NonNegativeInt, StrictBool, StrictStr from feast import Entity, FeatureView, utils from feast.infra.infra_object import InfraObject @@ -52,6 +63,27 @@ logger = logging.getLogger(__name__) +# Default for ``max_batched_result_rows``: the largest request, measured as entity +# rows x feature views, whose feature views are read through shared BatchGetItem +# requests. Below it, packing the keys of every feature view into the same requests +# replaces one or more requests per view with as few as one in total. Past it, each +# view already fills several requests of its own, so packing saves at most one +# partially filled request per view, while holding every view's items at once makes +# peak memory grow with the number of views. Past the limit the generic path is used +# instead, which reads and releases a single view at a time. +MAX_BATCHED_RESULT_ROWS = 2048 + + +class _BatchedRead(NamedTuple): + """One feature view's share of a batched online read.""" + + table: FeatureView + requested_features: List[str] + table_name: str + entity_ids: List[str] + idxs: Tuple[List[int], ...] + output_len: int + class DynamoDBOnlineStoreConfig(FeastConfigBaseModel): """Online store config for DynamoDB store""" @@ -90,6 +122,11 @@ class DynamoDBOnlineStoreConfig(FeastConfigBaseModel): """Maximum number of parallel threads for batch read operations. Higher values improve throughput for large batch reads but increase resource usage.""" + max_batched_result_rows: NonNegativeInt = MAX_BATCHED_RESULT_ROWS + """Largest request, as entity rows x feature views, whose feature views are read + through shared BatchGetItem requests. Set to 0 to always read one feature view + at a time.""" + max_pool_connections: int = 50 """Max number of connections for async Dynamodb operations. Increase for high-throughput workloads.""" @@ -715,6 +752,271 @@ def to_tbl_resp(raw_client_response): return list(itertools.chain(*result_batches)) + def _read_features_per_fv( + self, + config: RepoConfig, + grouped_refs: List, + join_key_values: Dict, + entity_name_to_join_key_map: Dict, + online_features_response, + full_feature_names: bool, + include_feature_view_version_metadata: bool, + ) -> None: + """Read every requested feature view through shared BatchGetItem requests. + + The generic path reads one feature view after another, with at least one + BatchGetItem request each. Each feature view lives in its own table here, and + a BatchGetItem request can read from several tables at once, so the keys of + every view are packed into the same requests, which are sent in parallel. + """ + if not self._should_batch(config, grouped_refs, join_key_values): + return super()._read_features_per_fv( + config, + grouped_refs, + join_key_values, + entity_name_to_join_key_map, + online_features_response, + full_feature_names, + include_feature_view_version_metadata, + ) + + online_config = config.online_store + assert isinstance(online_config, DynamoDBOnlineStoreConfig) + + reads = self._prepare_batched_reads( + config, grouped_refs, join_key_values, entity_name_to_join_key_map + ) + request_items = self._to_batched_request_items(online_config, reads) + + # boto3 clients are thread-safe, so the requests can share one client. + dynamodb_client = self._get_dynamodb_client( + online_config.region, + online_config.endpoint_url, + online_config.session_based_auth, + ) + + def fetch_batch(batch_request_items: Dict[str, Any]) -> Dict[str, Any]: + return dynamodb_client.batch_get_item(RequestItems=batch_request_items) + + if len(request_items) <= 1: + responses = [fetch_batch(items) for items in request_items] + else: + max_workers = min(len(request_items), online_config.max_read_workers) + with ThreadPoolExecutor(max_workers=max_workers) as executor: + responses = list(executor.map(fetch_batch, request_items)) + + self._populate_from_batched_responses( + reads, + responses, + online_features_response, + full_feature_names, + include_feature_view_version_metadata, + ) + + async def _read_features_per_fv_async( + self, + config: RepoConfig, + grouped_refs: List, + join_key_values: Dict, + entity_name_to_join_key_map: Dict, + online_features_response, + full_feature_names: bool, + include_feature_view_version_metadata: bool, + ) -> None: + """Async version of :meth:`_read_features_per_fv`. + + The generic async path already sends the per-view requests concurrently. + Packing them still replaces at least one request per view with as few as one + in total. + """ + if not self._should_batch(config, grouped_refs, join_key_values): + return await super()._read_features_per_fv_async( + config, + grouped_refs, + join_key_values, + entity_name_to_join_key_map, + online_features_response, + full_feature_names, + include_feature_view_version_metadata, + ) + + online_config = config.online_store + assert isinstance(online_config, DynamoDBOnlineStoreConfig) + + reads = self._prepare_batched_reads( + config, grouped_refs, join_key_values, entity_name_to_join_key_map + ) + request_items = self._to_batched_request_items(online_config, reads) + + client = await self._get_aiodynamodb_client( + online_config.region, + online_config.max_pool_connections, + online_config.keepalive_timeout, + online_config.connect_timeout, + online_config.read_timeout, + online_config.total_max_retry_attempts, + online_config.retry_mode, + online_config.endpoint_url, + ) + responses = await asyncio.gather( + *[ + client.batch_get_item(RequestItems=batch_request_items) + for batch_request_items in request_items + ] + ) + + self._populate_from_batched_responses( + reads, + responses, + online_features_response, + full_feature_names, + include_feature_view_version_metadata, + ) + + @staticmethod + def _should_batch( + config: RepoConfig, grouped_refs: List, join_key_values: Dict + ) -> bool: + """Whether to read this request's feature views through shared requests. + + A single view has nothing to share requests with, so it keeps the generic + path. The size check is estimated from the request shape rather than from the + resolved keys, so that deciding against batching costs nothing. The number of + entity rows in the request is an upper bound on the unique keys any one view + will read. See :data:`MAX_BATCHED_RESULT_ROWS`. + """ + if len(grouped_refs) < 2: + return False + entity_rows = max( + (len(values) for values in join_key_values.values()), default=0 + ) + return ( + entity_rows * len(grouped_refs) + <= config.online_store.max_batched_result_rows + ) + + def _prepare_batched_reads( + self, + config: RepoConfig, + grouped_refs: List, + join_key_values: Dict, + entity_name_to_join_key_map: Dict, + ) -> List[_BatchedRead]: + """Resolve the table and entity ids each feature view needs before reading. + + Feature views can be keyed on different entities, so the entity ids are + computed per view rather than shared. + """ + online_config = config.online_store + reads = [] + for table, requested_features in grouped_refs: + table_entity_values, idxs, output_len = utils._get_unique_entities( + table, join_key_values, entity_name_to_join_key_map + ) + entity_key_protos = utils._get_entity_key_protos(table_entity_values) + reads.append( + _BatchedRead( + table=table, + requested_features=requested_features, + table_name=_get_table_name(online_config, config, table), + entity_ids=self._to_entity_ids(config, entity_key_protos), + idxs=idxs, + output_len=output_len, + ) + ) + return reads + + @staticmethod + def _to_batched_request_items( + online_config: DynamoDBOnlineStoreConfig, reads: List[_BatchedRead] + ) -> List[Dict[str, Any]]: + """Pack the keys of every feature view into as few requests as possible. + + ``batch_size`` caps the keys of one BatchGetItem request across all of its + tables. A request must not repeat a key within a table, so views that share a + table, such as one feature view requested under two aliases, are read once, + projecting the union of their requested features. + """ + features_by_table: Dict[str, Optional[List[str]]] = {} + entity_ids_by_table: Dict[str, Dict[str, None]] = {} + for read in reads: + if read.table_name not in features_by_table: + features_by_table[read.table_name] = read.requested_features or None + entity_ids_by_table[read.table_name] = {} + else: + features = features_by_table[read.table_name] + features_by_table[read.table_name] = ( + list(dict.fromkeys(features + read.requested_features)) + if features and read.requested_features + else None + ) + entity_ids_by_table[read.table_name].update(dict.fromkeys(read.entity_ids)) + + keys = [ + (table_name, entity_id) + for table_name, entity_ids in entity_ids_by_table.items() + for entity_id in entity_ids + ] + batch_size = online_config.batch_size + request_items: List[Dict[str, Any]] = [] + for start in range(0, len(keys), batch_size): + batch_by_table: Dict[str, List[str]] = defaultdict(list) + for table_name, entity_id in keys[start : start + batch_size]: + batch_by_table[table_name].append(entity_id) + batch_request_items: Dict[str, Any] = {} + for table_name, batch in batch_by_table.items(): + batch_request_items.update( + DynamoDBOnlineStore._to_client_batch_get_payload( + online_config, table_name, batch, features_by_table[table_name] + ) + ) + request_items.append(batch_request_items) + return request_items + + def _populate_from_batched_responses( + self, + reads: List[_BatchedRead], + responses: Sequence[Dict[str, Any]], + online_features_response, + full_feature_names: bool, + include_feature_view_version_metadata: bool, + ) -> None: + """Split the responses by table and hand each view its rows in key order. + + BatchGetItem returns items in no particular order, and a request can mix + tables, so the items are indexed by table and entity id first. + """ + if self._type_deserializer is None: + self._type_deserializer = TypeDeserializer() + deserialize = self._type_deserializer.deserialize + + rows_by_table: Dict[ + str, Dict[str, Tuple[Optional[datetime], Optional[Dict[str, ValueProto]]]] + ] = defaultdict(dict) + for response in responses: + for table_name, items in response.get("Responses", {}).items(): + table_rows = rows_by_table[table_name] + for item in items: + table_rows[deserialize(item["entity_id"])] = self._to_read_row( + deserialize(item["event_ts"]), deserialize(item["values"]) + ) + + for read in reads: + table_rows = rows_by_table.get(read.table_name, {}) + utils._populate_response_from_feature_data( + read.requested_features, + [ + table_rows.get(entity_id, (None, None)) + for entity_id in read.entity_ids + ], + read.idxs, + online_features_response, + full_feature_names, + read.table, + read.output_len, + include_feature_view_version_metadata, + ) + def _get_dynamodb_client( self, region: str, @@ -812,22 +1114,22 @@ def _process_batch_get_response( for idx, entity_id in enumerate(batch): tbl_res = response_dict.get(entity_id) if tbl_res is not None: - # Parse feature values - features: Dict[str, ValueProto] = {} - values_data = tbl_res["values"] - for feature_name, value_bin in values_data.items(): - val = ValueProto() - val.ParseFromString(value_bin.value) - features[feature_name] = val - - # Parse timestamp and set result - result[idx] = ( - datetime.fromisoformat(tbl_res["event_ts"]), - features, - ) + result[idx] = self._to_read_row(tbl_res["event_ts"], tbl_res["values"]) return result + @staticmethod + def _to_read_row( + event_ts: str, values: Dict[str, Any] + ) -> Tuple[Optional[datetime], Optional[Dict[str, ValueProto]]]: + """Parse one stored item's timestamp and serialized feature values.""" + features: Dict[str, ValueProto] = {} + for feature_name, value_bin in values.items(): + val = ValueProto() + val.ParseFromString(value_bin.value) + features[feature_name] = val + return datetime.fromisoformat(event_ts), features + @staticmethod def _to_entity_ids(config: RepoConfig, entity_keys: List[EntityKeyProto]): """Convert entity keys to entity IDs.""" diff --git a/sdk/python/tests/unit/infra/online_store/test_dynamodb_batched_read.py b/sdk/python/tests/unit/infra/online_store/test_dynamodb_batched_read.py new file mode 100644 index 00000000000..ca2023107c8 --- /dev/null +++ b/sdk/python/tests/unit/infra/online_store/test_dynamodb_batched_read.py @@ -0,0 +1,384 @@ +"""Unit tests for batching DynamoDB online reads across feature views.""" + +from datetime import timedelta +from typing import List +from unittest.mock import patch + +import pytest +from moto import mock_dynamodb +from pydantic import ValidationError + +from feast import Entity, FeatureView +from feast.field import Field +from feast.infra.offline_stores.dask import DaskOfflineStoreConfig +from feast.infra.online_stores.dynamodb import ( + MAX_BATCHED_RESULT_ROWS, + DynamoDBOnlineStore, + DynamoDBOnlineStoreConfig, +) +from feast.infra.online_stores.online_store import OnlineStore +from feast.protos.feast.serving.ServingService_pb2 import ( + FieldStatus, + GetOnlineFeaturesResponse, +) +from feast.protos.feast.types.EntityKey_pb2 import EntityKey as EntityKeyProto +from feast.protos.feast.types.Value_pb2 import Value as ValueProto +from feast.repo_config import RepoConfig +from feast.types import Int64 +from feast.utils import _utc_now +from feast.value_type import ValueType +from tests.utils.dynamo_table_creator import create_test_table + +PROJECT = "test_aws" +REGION = "us-west-2" + +DRIVER = Entity(name="driver", join_keys=["driver_id"], value_type=ValueType.INT64) +CUSTOMER = Entity( + name="customer", join_keys=["customer_id"], value_type=ValueType.INT64 +) +JOIN_KEY_MAP = {"driver": "driver_id", "customer": "customer_id"} + + +def _feature_view( + name: str, features: List[str], entity: Entity = DRIVER +) -> FeatureView: + return FeatureView( + name=name, + entities=[entity], + ttl=timedelta(days=1), + schema=[Field(name=feature, dtype=Int64) for feature in features], + ) + + +def _config(**online_store) -> RepoConfig: + return RepoConfig( + registry="s3://test_registry/registry.db", + project=PROJECT, + provider="aws", + online_store=DynamoDBOnlineStoreConfig(region=REGION, **online_store), + offline_store=DaskOfflineStoreConfig(), + entity_key_serialization_version=3, + ) + + +def _write(store, config, fv: FeatureView, join_key: str, rows: dict) -> None: + """Create the view's table and write ``{entity value: {feature: value}}``.""" + create_test_table(PROJECT, fv.name, REGION) + store.online_write_batch( + config, + fv, + [ + ( + EntityKeyProto( + join_keys=[join_key], + entity_values=[ValueProto(int64_val=entity_value)], + ), + {name: ValueProto(int64_val=value) for name, value in values.items()}, + _utc_now(), + None, + ) + for entity_value, values in rows.items() + ], + None, + ) + + +@pytest.fixture +def dynamodb(): + # The moto decorator does not await coroutines, so the async tests get the + # mock from a fixture instead. + with mock_dynamodb(): + yield + + +def _drivers(*ids) -> dict: + return {"driver_id": [ValueProto(int64_val=i) for i in ids]} + + +def _args(config, grouped_refs, join_key_values, response): + return dict( + config=config, + grouped_refs=grouped_refs, + join_key_values=join_key_values, + entity_name_to_join_key_map=JOIN_KEY_MAP, + online_features_response=response, + full_feature_names=True, + include_feature_view_version_metadata=False, + ) + + +class _AsyncClient: + """Serves the async client calls from the moto-backed sync client.""" + + def __init__(self, client): + self._client = client + + async def batch_get_item(self, **kwargs): + return self._client.batch_get_item(**kwargs) + + +class _CountingClient: + """Wraps a client and records the RequestItems of every batch_get_item call.""" + + def __init__(self, client, calls: List): + self._client = client + self._calls = calls + + def batch_get_item(self, **kwargs): + self._calls.append(kwargs["RequestItems"]) + return self._client.batch_get_item(**kwargs) + + +def _read_sync(store, config, grouped_refs, join_key_values, generic=False): + response = GetOnlineFeaturesResponse() + read = ( + OnlineStore._read_features_per_fv + if generic + else type(store)._read_features_per_fv + ) + read(store, **_args(config, grouped_refs, join_key_values, response)) + return response + + +async def _read_async(store, config, grouped_refs, join_key_values, generic=False): + response = GetOnlineFeaturesResponse() + read = ( + OnlineStore._read_features_per_fv_async + if generic + else type(store)._read_features_per_fv_async + ) + client = _AsyncClient(store._get_dynamodb_client(REGION)) + + async def get_client(*args, **kwargs): + return client + + with patch.object(store, "_get_aiodynamodb_client", get_client): + await read(store, **_args(config, grouped_refs, join_key_values, response)) + return response + + +async def _read_sync_awaitable(*args, **kwargs): + """Await-compatible wrapper so both overrides share one test body.""" + return _read_sync(*args, **kwargs) + + +READERS = [ + pytest.param(_read_sync_awaitable, id="sync"), + pytest.param(_read_async, id="async"), +] + + +def _values(response: GetOnlineFeaturesResponse) -> dict: + """Map each feature name to its (values, statuses) across entity rows.""" + return { + name: ( + [value.int64_val for value in vector.values], + list(vector.statuses), + ) + for name, vector in zip(response.metadata.feature_names.val, response.results) + } + + +@pytest.mark.usefixtures("dynamodb") +@pytest.mark.parametrize("read", READERS) +async def test_views_read_in_one_request_match_generic_path(read): + """Views sharing an entity id AND a feature name come back from one request. + + Distinct feature names would make a cross-wired table invisible, since the + response is keyed by feature name. + """ + config = _config() + store = DynamoDBOnlineStore() + fv_a = _feature_view("fv_a", ["feat"]) + fv_b = _feature_view("fv_b", ["feat"]) + fv_c = _feature_view("fv_c", ["visits"], entity=CUSTOMER) + _write(store, config, fv_a, "driver_id", {1: {"feat": 11}, 2: {"feat": 12}}) + # Driver 2 only exists in fv_a, so fv_b must report it as not found. + _write(store, config, fv_b, "driver_id", {1: {"feat": 21}}) + _write(store, config, fv_c, "customer_id", {7: {"visits": 70}}) + + grouped_refs = [(fv_a, ["feat"]), (fv_b, ["feat"]), (fv_c, ["visits"])] + join_key_values = { + **_drivers(1, 2, 1), + "customer_id": [ValueProto(int64_val=i) for i in (7, 8, 7)], + } + + calls: List = [] + store._dynamodb_client = _CountingClient(store._get_dynamodb_client(REGION), calls) + batched = await read(store, config, grouped_refs, join_key_values) + assert len(calls) == 1 + assert set(calls[0]) == {"test_aws.fv_a", "test_aws.fv_b", "test_aws.fv_c"} + + generic = await read(store, config, grouped_refs, join_key_values, generic=True) + assert batched == generic + + found, missing = FieldStatus.PRESENT, FieldStatus.NOT_FOUND + assert _values(batched) == { + "fv_a__feat": ([11, 12, 11], [found, found, found]), + "fv_b__feat": ([21, 0, 21], [found, missing, found]), + "fv_c__visits": ([70, 0, 70], [found, missing, found]), + } + + +@pytest.mark.usefixtures("dynamodb") +@pytest.mark.parametrize("read", READERS) +async def test_keys_are_packed_across_views_up_to_batch_size(read): + """Requests are filled with keys of several tables, never above batch_size.""" + config = _config(batch_size=3) + store = DynamoDBOnlineStore() + fv_a = _feature_view("fv_a", ["feat"]) + fv_b = _feature_view("fv_b", ["feat"]) + _write(store, config, fv_a, "driver_id", {i: {"feat": i} for i in range(1, 5)}) + _write(store, config, fv_b, "driver_id", {i: {"feat": -i} for i in range(1, 5)}) + + calls: List = [] + store._dynamodb_client = _CountingClient(store._get_dynamodb_client(REGION), calls) + response = await read( + store, + config, + [(fv_a, ["feat"]), (fv_b, ["feat"])], + _drivers(1, 2, 3, 4), + ) + + # 8 keys at 3 per request; the second request mixes both tables. + assert [ + sum(len(table["Keys"]) for table in request.values()) for request in calls + ] == [3, 3, 2] + assert set(calls[1]) == {"test_aws.fv_a", "test_aws.fv_b"} + assert _values(response) == { + "fv_a__feat": ([1, 2, 3, 4], [FieldStatus.PRESENT] * 4), + "fv_b__feat": ([-1, -2, -3, -4], [FieldStatus.PRESENT] * 4), + } + + +@pytest.mark.usefixtures("dynamodb") +@pytest.mark.parametrize("read", READERS) +async def test_views_sharing_a_table_read_each_key_once(read): + """One view under two aliases maps to one table and must not repeat its keys. + + BatchGetItem rejects a request that lists the same key twice for a table, so + the aliases are read once, projecting the union of their features. + """ + config = _config() + store = DynamoDBOnlineStore() + fv = _feature_view("fv", ["feat_x", "feat_y", "feat_z"]) + _write( + store, + config, + fv, + "driver_id", + {1: {"feat_x": 1, "feat_y": 2, "feat_z": 3}}, + ) + + calls: List = [] + store._dynamodb_client = _CountingClient(store._get_dynamodb_client(REGION), calls) + response = await read( + store, + config, + [ + (fv.with_name("first"), ["feat_x"]), + (fv.with_name("second"), ["feat_y"]), + ], + _drivers(1), + ) + + assert len(calls) == 1 + assert list(calls[0]) == ["test_aws.fv"] + request = calls[0]["test_aws.fv"] + assert len(request["Keys"]) == 1 + assert set(request["ExpressionAttributeNames"].values()) == { + "entity_id", + "event_ts", + "values", + "feat_x", + "feat_y", + } + assert _values(response) == { + "first__feat_x": ([1], [FieldStatus.PRESENT]), + "second__feat_y": ([2], [FieldStatus.PRESENT]), + } + + +@pytest.mark.usefixtures("dynamodb") +@pytest.mark.parametrize("read", READERS) +async def test_requested_features_are_projected_per_table(read): + config = _config() + store = DynamoDBOnlineStore() + fv_a = _feature_view("fv_a", ["kept", "dropped"]) + fv_b = _feature_view("fv_b", ["other"]) + _write(store, config, fv_a, "driver_id", {1: {"kept": 1, "dropped": 2}}) + _write(store, config, fv_b, "driver_id", {1: {"other": 3}}) + + calls: List = [] + store._dynamodb_client = _CountingClient(store._get_dynamodb_client(REGION), calls) + await read(store, config, [(fv_a, ["kept"]), (fv_b, ["other"])], _drivers(1)) + + projected = { + table_name: set(request["ExpressionAttributeNames"].values()) + - {"entity_id", "event_ts", "values"} + for table_name, request in calls[0].items() + } + assert projected == {"test_aws.fv_a": {"kept"}, "test_aws.fv_b": {"other"}} + + +@pytest.mark.parametrize( + "online_store, n_views, n_rows", + [ + pytest.param({}, 1, 1, id="single_view"), + pytest.param({"max_batched_result_rows": 0}, 2, 1, id="disabled"), + pytest.param({"max_batched_result_rows": 5}, 2, 3, id="over_limit"), + ], +) +@pytest.mark.parametrize("read", READERS) +async def test_generic_path_reads_each_view(read, online_store, n_views, n_rows): + """Outside the gate, every view is read by its own online_read call.""" + config = _config(**online_store) + store = DynamoDBOnlineStore() + grouped_refs = [ + (_feature_view(f"fv_{i}", ["feat"]), ["feat"]) for i in range(n_views) + ] + read_tables: List[str] = [] + + def online_read(self, config, table, entity_keys, requested_features=None): + read_tables.append(table.name) + return [(None, None)] * len(entity_keys) + + async def online_read_async(self, *args, **kwargs): + return online_read(self, *args, **kwargs) + + with ( + patch.object(DynamoDBOnlineStore, "online_read", online_read), + patch.object(DynamoDBOnlineStore, "online_read_async", online_read_async), + patch.object( + DynamoDBOnlineStore, + "_to_batched_request_items", + side_effect=AssertionError, + ), + ): + await read(store, config, grouped_refs, _drivers(*range(n_rows))) + + assert read_tables == [f"fv_{i}" for i in range(n_views)] + + +def test_batching_at_the_limit(): + grouped_refs = [(None, []), (None, [])] + config = _config(max_batched_result_rows=4) + assert DynamoDBOnlineStore._should_batch(config, grouped_refs, _drivers(1, 2)) + assert not DynamoDBOnlineStore._should_batch( + config, grouped_refs, _drivers(1, 2, 3) + ) + + +def test_max_batched_result_rows_config(): + assert ( + DynamoDBOnlineStoreConfig(region=REGION).max_batched_result_rows + == MAX_BATCHED_RESULT_ROWS + ) + assert ( + DynamoDBOnlineStoreConfig( + region=REGION, max_batched_result_rows=0 + ).max_batched_result_rows + == 0 + ) + with pytest.raises(ValidationError): + DynamoDBOnlineStoreConfig(region=REGION, max_batched_result_rows=-1)