feat: Add support for DynamoDB online_read in batches

sdk/python/feast/infra/online_stores/dynamodb.py

@@ -188,6 +188,7 @@ def online_read(
         table: FeatureView,
         entity_keys: List[EntityKeyProto],
         requested_features: Optional[List[str]] = None,
+        batch_size: int = 20,
     ) -> List[Tuple[Optional[datetime], Optional[Dict[str, ValueProto]]]]:
         """
         Retrieve feature values from the online DynamoDB store.
@@ -199,26 +200,59 @@ def online_read(
             config: The RepoConfig for the current FeatureStore.
             table: Feast FeatureView.
             entity_keys: a list of entity keys that should be read from the FeatureStore.
+            batch_size: the number of items to send in a batch_get_item request to DynamoDB.
+                DynamoDB record size limit is 400kb and can retrieve 16MB per call, it is recommended
+                to set batch_size value less than 40 to avoid ``UnprocessedKeys`` and
+                ``ValidationException`` errors.
         """
+        if batch_size > 40:
+            raise ValueError(
+                f"batch_size value must be less than 40, input value is {batch_size}"
+            )
         online_config = config.online_store
         assert isinstance(online_config, DynamoDBOnlineStoreConfig)
         dynamodb_resource = self._get_dynamodb_resource(online_config.region)
+        table_instance = dynamodb_resource.Table(_get_table_name(config, table))
 
         result: List[Tuple[Optional[datetime], Optional[Dict[str, ValueProto]]]] = []
-        for entity_key in entity_keys:
-            table_instance = dynamodb_resource.Table(_get_table_name(config, table))
-            entity_id = compute_entity_id(entity_key)
+        entity_ids = [compute_entity_id(entity_key) for entity_key in entity_keys]
+
+        len_entity_ids = len(entity_ids)
+        # Iterate until the end_index is the value len_entity_ids
+        iters = (
+            len_entity_ids // batch_size + 1
+            if len_entity_ids % batch_size > 0
+            else len_entity_ids // batch_size
+        )
+        for i in range(iters):
+            start_index = min(i * batch_size, len_entity_ids)
+            end_index = min(i * batch_size + batch_size, len_entity_ids)
+
+            batch_entity_ids = {
+                table_instance.name: {
+                    "Keys": [
+                        {"entity_id": entity_id}
+                        for entity_id in entity_ids[start_index:end_index]
+                    ]
+                }
+            }
+
             with tracing_span(name="remote_call"):
-                response = table_instance.get_item(Key={"entity_id": entity_id})
-            value = response.get("Item")
-
-            if value is not None:
-                res = {}
-                for feature_name, value_bin in value["values"].items():
-                    val = ValueProto()
-                    val.ParseFromString(value_bin.value)
-                    res[feature_name] = val
-                result.append((datetime.fromisoformat(value["event_ts"]), res))
+                response = dynamodb_resource.batch_get_item(
+                    RequestItems=batch_entity_ids
+                )
+
+            response = response.get("Responses")
+            table_responses = response.get(table_instance.name)
+
+            if table_responses:
+                for tbl_res in table_responses:
+                    res = {}
+                    for feature_name, value_bin in tbl_res["values"].items():
+                        val = ValueProto()
+                        val.ParseFromString(value_bin.value)
+                        res[feature_name] = val
+                    result.append((datetime.fromisoformat(tbl_res["event_ts"]), res))
             else:
                 result.append((None, None))
         return result

sdk/python/tests/unit/online_store/test_dynamodb_online_store.py

@@ -0,0 +1,57 @@
+from dataclasses import dataclass
+
+import pytest
+from moto import mock_dynamodb2
+
+from feast.infra.offline_stores.file import FileOfflineStoreConfig
+from feast.infra.online_stores.dynamodb import (
+    DynamoDBOnlineStore,
+    DynamoDBOnlineStoreConfig,
+)
+from feast.repo_config import RepoConfig
+from tests.utils.online_store_utils import (
+    _create_n_customer_test_samples,
+    _create_test_table,
+    _insert_data_test_table,
+)
+
+REGISTRY = "s3://test_registry/registry.db"
+PROJECT = "test_aws"
+PROVIDER = "aws"
+TABLE_NAME = "dynamodb_online_store"
+REGION = "us-west-2"
+
+
+@dataclass
+class MockFeatureView:
+    name: str
+
+
+@pytest.fixture
+def repo_config():
+    return RepoConfig(
+        registry=REGISTRY,
+        project=PROJECT,
+        provider=PROVIDER,
+        online_store=DynamoDBOnlineStoreConfig(region=REGION),
+        offline_store=FileOfflineStoreConfig(),
+    )
+
+
+@mock_dynamodb2
+@pytest.mark.parametrize("n_samples", [5, 50, 100])
+def test_online_read(repo_config, n_samples):
+    """Test DynamoDBOnlineStore online_read method."""
+    _create_test_table(PROJECT, f"{TABLE_NAME}_{n_samples}", REGION)
+    data = _create_n_customer_test_samples(n=n_samples)
+    _insert_data_test_table(data, PROJECT, f"{TABLE_NAME}_{n_samples}", REGION)
+
+    entity_keys, features = zip(*data)
+    dynamodb_store = DynamoDBOnlineStore()
+    returned_items = dynamodb_store.online_read(
+        config=repo_config,
+        table=MockFeatureView(name=f"{TABLE_NAME}_{n_samples}"),
+        entity_keys=entity_keys,
+    )
+    assert len(returned_items) == len(data)
+    assert [item[1] for item in returned_items] == list(features)

sdk/python/tests/utils/online_store_utils.py

@@ -0,0 +1,54 @@
+from datetime import datetime
+
+import boto3
+
+from feast import utils
+from feast.infra.online_stores.helpers import compute_entity_id
+from feast.protos.feast.types.EntityKey_pb2 import EntityKey as EntityKeyProto
+from feast.protos.feast.types.Value_pb2 import Value as ValueProto
+
+
+def _create_n_customer_test_samples(n=10):
+    return [
+        (
+            EntityKeyProto(
+                join_keys=["customer"], entity_values=[ValueProto(string_val=str(i))]
+            ),
+            {
+                "avg_orders_day": ValueProto(float_val=1.0),
+                "name": ValueProto(string_val="John"),
+                "age": ValueProto(int64_val=3),
+            },
+        )
+        for i in range(n)
+    ]
+
+
+def _create_test_table(project, tbl_name, region):
+    client = boto3.client("dynamodb", region_name=region)
+    client.create_table(
+        TableName=f"{project}.{tbl_name}",
+        KeySchema=[{"AttributeName": "entity_id", "KeyType": "HASH"}],
+        AttributeDefinitions=[{"AttributeName": "entity_id", "AttributeType": "S"}],
+        BillingMode="PAY_PER_REQUEST",
+    )
+
+
+def _delete_test_table(project, tbl_name, region):
+    client = boto3.client("dynamodb", region_name=region)
+    client.delete_table(TableName=f"{project}.{tbl_name}")
+
+
+def _insert_data_test_table(data, project, tbl_name, region):
+    dynamodb_resource = boto3.resource("dynamodb", region_name=region)
+    table_instance = dynamodb_resource.Table(f"{project}.{tbl_name}")
+    for entity_key, features in data:
+        entity_id = compute_entity_id(entity_key)
+        with table_instance.batch_writer() as batch:
+            batch.put_item(
+                Item={
+                    "entity_id": entity_id,
+                    "event_ts": str(utils.make_tzaware(datetime.utcnow())),
+                    "values": {k: v.SerializeToString() for k, v in features.items()},
+                }
+            )