fix(dynamodb): split read_items keys into BatchGetItem requests of at most 100

awswrangler/dynamodb/_read.py

@@ -42,6 +42,8 @@
 
 _ItemsListType = List[Dict[str, "TableAttributeValueTypeDef"]]
 
+_MAX_BATCH_GET_ITEM_KEYS = 100
+
 
 def _read_chunked(iterator: Iterator[dict[str, "TableAttributeValueTypeDef"]]) -> Iterator[pd.DataFrame]:
     for item in iterator:
@@ -279,21 +281,30 @@ def _read_batch_items_chunked(
     dynamodb_client = dynamodb_client if dynamodb_client else _utils.client("dynamodb")
     deserializer = TypeDeserializer()
 
-    response = dynamodb_client.batch_get_item(RequestItems={table_name: kwargs})
-    yield [_deserialize_item(d, deserializer) for d in response.get("Responses", {table_name: []}).get(table_name, [])]
-
-    # SEE: handle possible unprocessed keys. As suggested in Boto3 docs,
-    # this approach should involve exponential backoff, but this should be
-    # already managed by AWS SDK itself, as stated
-    # [here](https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/Programming.Errors.html)
-    while response["UnprocessedKeys"]:
-        kwargs["Keys"] = response["UnprocessedKeys"][table_name]["Keys"]
+    # BatchGetItem accepts at most 100 keys per request
+    # SEE: https://docs.aws.amazon.com/amazondynamodb/latest/APIReference/API_BatchGetItem.html
+    keys = kwargs.pop("Keys")
+    for i in range(0, len(keys), _MAX_BATCH_GET_ITEM_KEYS):
+        request = {**kwargs, "Keys": keys[i : i + _MAX_BATCH_GET_ITEM_KEYS]}
 
-        response = dynamodb_client.batch_get_item(RequestItems={table_name: kwargs})
+        response = dynamodb_client.batch_get_item(RequestItems={table_name: request})
         yield [
             _deserialize_item(d, deserializer) for d in response.get("Responses", {table_name: []}).get(table_name, [])
         ]
 
+        # SEE: handle possible unprocessed keys. As suggested in Boto3 docs,
+        # this approach should involve exponential backoff, but this should be
+        # already managed by AWS SDK itself, as stated
+        # [here](https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/Programming.Errors.html)
+        while response["UnprocessedKeys"]:
+            request["Keys"] = response["UnprocessedKeys"][table_name]["Keys"]
+
+            response = dynamodb_client.batch_get_item(RequestItems={table_name: request})
+            yield [
+                _deserialize_item(d, deserializer)
+                for d in response.get("Responses", {table_name: []}).get(table_name, [])
+            ]
+
 
 @_handle_reserved_keyword_error
 def _read_batch_items(

tests/unit/test_moto.py

@@ -1089,6 +1089,25 @@ def test_dynamodb_read_items_missing_single_key(moto_dynamodb_client, moto_dynam
     assert items == []
 
 
+@pytest.mark.parametrize("chunked", [False, True])
+def test_dynamodb_read_items_more_than_100_partition_values(moto_dynamodb_client, moto_dynamodb_table, chunked) -> None:
+    # BatchGetItem accepts at most 100 keys per request
+    items = [{"key": i, "value": f"v{i}"} for i in range(250)]
+    wr.dynamodb.put_items(items=items, table_name=moto_dynamodb_table)
+
+    result = wr.dynamodb.read_items(
+        table_name=moto_dynamodb_table,
+        partition_values=[item["key"] for item in items],
+        as_dataframe=False,
+        chunked=chunked,
+    )
+    if chunked:
+        result = [item for chunk in result for item in chunk]
+
+    assert len(result) == len(items)
+    assert sorted(int(item["key"]) for item in result) == list(range(250))
+
+
 def _capture_timestream_write_records(session: boto3.Session) -> list:
     captured: list = []