add async writer for dynamo

robhowley · robhowley · commit 7991146830f7 · 2024-10-24T09:20:39.000-04:00
Signed-off-by: Rob Howley &lt;howley.robert@gmail.com&gt;
diff --git a/sdk/python/feast/infra/online_stores/dynamodb.py b/sdk/python/feast/infra/online_stores/dynamodb.py
@@ -26,6 +26,7 @@
 from feast.infra.online_stores.helpers import compute_entity_id
 from feast.infra.online_stores.online_store import OnlineStore
 from feast.infra.supported_async_methods import SupportedAsyncMethods
+from feast.infra.utils.aws_utils import dynamo_write_items_async
 from feast.protos.feast.core.DynamoDBTable_pb2 import (
     DynamoDBTable as DynamoDBTableProto,
 )
@@ -103,7 +104,7 @@ async def close(self):
 
     @property
     def async_supported(self) -> SupportedAsyncMethods:
-        return SupportedAsyncMethods(read=True)
+        return SupportedAsyncMethods(read=True, write=True)
 
     def update(
         self,
@@ -238,6 +239,42 @@ def online_write_batch(
         )
         self._write_batch_non_duplicates(table_instance, data, progress, config)
 
+    async def online_write_batch_async(
+        self,
+        config: RepoConfig,
+        table: FeatureView,
+        data: List[
+            Tuple[EntityKeyProto, Dict[str, ValueProto], datetime, Optional[datetime]]
+        ],
+        progress: Optional[Callable[[int], Any]],
+    ) -> None:
+        """
+        Writes a batch of feature rows to the online store asynchronously.
+
+        If a tz-naive timestamp is passed to this method, it is assumed to be UTC.
+
+        Args:
+            config: The config for the current feature store.
+            table: Feature view to which these feature rows correspond.
+            data: A list of quadruplets containing feature data. Each quadruplet contains an entity
+                key, a dict containing feature values, an event timestamp for the row, and the created
+                timestamp for the row if it exists.
+            progress: Function to be called once a batch of rows is written to the online store, used
+                to show progress.
+        """
+        online_config = config.online_store
+        assert isinstance(online_config, DynamoDBOnlineStoreConfig)
+
+        table_name = _get_table_name(online_config, config, table)
+        items = [
+            _to_write_item(config, entity_key, features, timestamp)
+            for entity_key, features, timestamp, _ in data
+        ]
+        client = _get_aiodynamodb_client(
+            online_config.region, config.online_store.max_pool_connections
+        )
+        await dynamo_write_items_async(client, table_name, items)
+
     def online_read(
         self,
         config: RepoConfig,
@@ -419,19 +456,8 @@ def _write_batch_non_duplicates(
         """Deduplicate write batch request items on ``entity_id`` primary key."""
         with table_instance.batch_writer(overwrite_by_pkeys=["entity_id"]) as batch:
             for entity_key, features, timestamp, created_ts in data:
-                entity_id = compute_entity_id(
-                    entity_key,
-                    entity_key_serialization_version=config.entity_key_serialization_version,
-                )
                 batch.put_item(
-                    Item={
-                        "entity_id": entity_id,  # PartitionKey
-                        "event_ts": str(utils.make_tzaware(timestamp)),
-                        "values": {
-                            k: v.SerializeToString()
-                            for k, v in features.items()  # Serialized Features
-                        },
-                    }
+                    Item=_to_write_item(config, entity_key, features, timestamp)
                 )
                 if progress:
                     progress(1)
@@ -675,3 +701,18 @@ def _get_dynamodb_resource(self, region: str, endpoint_url: Optional[str] = None
                 region, endpoint_url
             )
         return self._dynamodb_resource
+
+
+def _to_write_item(config, entity_key, features, timestamp):
+    entity_id = compute_entity_id(
+        entity_key,
+        entity_key_serialization_version=config.entity_key_serialization_version,
+    )
+    return {
+        "entity_id": entity_id,  # PartitionKey
+        "event_ts": str(utils.make_tzaware(timestamp)),
+        "values": {
+            k: v.SerializeToString()
+            for k, v in features.items()  # Serialized Features
+        },
+    }
diff --git a/sdk/python/feast/infra/utils/aws_utils.py b/sdk/python/feast/infra/utils/aws_utils.py
@@ -1,4 +1,6 @@
+import asyncio
 import contextlib
+import itertools
 import os
 import tempfile
 import uuid
@@ -10,6 +12,7 @@
 import pyarrow as pa
 import pyarrow.parquet as pq
 from tenacity import (
+    AsyncRetrying,
     retry,
     retry_if_exception_type,
     stop_after_attempt,
@@ -1076,3 +1079,54 @@ def upload_arrow_table_to_athena(
         # Clean up S3 temporary data
         # for file_path in uploaded_files:
         #     s3_resource.Object(bucket, file_path).delete()
+
+
+class DynamoUnprocessedWriteItems(Exception):
+    pass
+
+
+async def dynamo_write_items_async(
+    dynamo_client, table_name: str, items: list[Any]
+) -> None:
+    DYNAMO_MAX_WRITE_BATCH_SIZE = 25
+
+    async def _do_write(items):
+        item_iter = iter(items)
+        item_batches = []
+        while True:
+            item_batch = [
+                item
+                for item in itertools.islice(item_iter, DYNAMO_MAX_WRITE_BATCH_SIZE)
+            ]
+            if not item_batch:
+                break
+
+            item_batches.append(item_batch)
+
+        return await asyncio.gather(
+            *[
+                dynamo_client.batch_write_item(
+                    RequestItems={table_name: item_batch},
+                )
+                for item_batch in item_batches
+            ]
+        )
+
+    put_items = [{"PutRequest": {"Item": item}} for item in items]
+
+    retries = AsyncRetrying(
+        retry=retry_if_exception_type(DynamoUnprocessedWriteItems),
+        wait=wait_exponential(multiplier=1, max=4),
+        reraise=True,
+    )
+
+    for attempt in retries:
+        with attempt:
+            response_batches = await _do_write(put_items)
+
+            put_items = []
+            for response in response_batches:
+                put_items.extend(response["UnprocessedItems"])
+
+            if put_items:
+                raise DynamoUnprocessedWriteItems()
diff --git a/sdk/python/tests/integration/online_store/test_push_features_to_online_store.py b/sdk/python/tests/integration/online_store/test_push_features_to_online_store.py
@@ -8,29 +8,51 @@
 from tests.integration.feature_repos.universal.entities import location
 
 
-@pytest.mark.integration
-@pytest.mark.universal_online_stores
-def test_push_features_and_read(environment, universal_data_sources):
+@pytest.fixture
+def store(environment, universal_data_sources):
     store = environment.feature_store
     _, _, data_sources = universal_data_sources
     feature_views = construct_universal_feature_views(data_sources)
     location_fv = feature_views.pushed_locations
     store.apply([location(), location_fv])
+    return store
+
 
+def _ingest_df():
     data = {
         "location_id": [1],
         "temperature": [4],
         "event_timestamp": [pd.Timestamp(_utc_now()).round("ms")],
         "created": [pd.Timestamp(_utc_now()).round("ms")],
     }
-    df_ingest = pd.DataFrame(data)
+    return pd.DataFrame(data)
 
-    store.push("location_stats_push_source", df_ingest)
+
+def assert_response(online_resp):
+    online_resp_dict = online_resp.to_dict()
+    assert online_resp_dict["location_id"] == [1]
+    assert online_resp_dict["temperature"] == [4]
+
+
+@pytest.mark.integration
+@pytest.mark.universal_online_stores
+def test_push_features_and_read(store):
+    store.push("location_stats_push_source", _ingest_df())
 
     online_resp = store.get_online_features(
         features=["pushable_location_stats:temperature"],
         entity_rows=[{"location_id": 1}],
     )
-    online_resp_dict = online_resp.to_dict()
-    assert online_resp_dict["location_id"] == [1]
-    assert online_resp_dict["temperature"] == [4]
+    assert_response(online_resp)
+
+
+@pytest.mark.integration
+@pytest.mark.universal_online_stores(only=["dynamodb"])
+async def test_push_features_and_read_async(store):
+    await store.push_async("location_stats_push_source", _ingest_df())
+
+    online_resp = await store.get_online_features_async(
+        features=["pushable_location_stats:temperature"],
+        entity_rows=[{"location_id": 1}],
+    )
+    assert_response(online_resp)
