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)