perf: serialize dataset push payloads once per chunk

Author: vdusek

src/apify/storage_clients/_apify/_dataset_client.py

@@ -1,13 +1,13 @@
 from __future__ import annotations
 
 import asyncio
+import json
 from logging import getLogger
 from typing import TYPE_CHECKING
 
 from typing_extensions import override
 
 from crawlee._utils.byte_size import ByteSize
-from crawlee._utils.file import json_dumps
 from crawlee.storage_clients._base import DatasetClient
 from crawlee.storage_clients.models import DatasetItemsListPage, DatasetMetadata
 
@@ -138,18 +138,16 @@ async def drop(self) -> None:
 
     @override
     async def push_data(self, data: Sequence[Mapping[str, JsonSerializable]] | Mapping[str, JsonSerializable]) -> None:
-        async def payloads_generator(items: Sequence[Mapping[str, JsonSerializable]]) -> AsyncIterator[str]:
-            for index, item in enumerate(items):
-                yield await self._check_and_serialize(item, index)
-
         async with self._charge_lock(), self._lock:
             items = data if self._is_sequence_of_items(data) else [data]
             if not items:
                 return
             limit = self._compute_limit_for_push(len(items))
             items = items[:limit]
 
-            async for chunk in self._chunk_by_size(payloads_generator(items)):
+            offset = 0
+            while offset < len(items):
+                chunk, offset = await asyncio.to_thread(self._serialize_chunk, items, offset)
                 await self._api_client.push_items(items=chunk)
 
             await self._charge_for_items(count_items=limit)
@@ -213,58 +211,43 @@ async def iterate_items(
             yield item
 
     @classmethod
-    async def _check_and_serialize(cls, item: Mapping[str, JsonSerializable], index: int | None = None) -> str:
-        """Serialize a given item to JSON, checks its serializability and size against a limit.
+    def _serialize_chunk(cls, items: Sequence[Mapping[str, JsonSerializable]], offset: int) -> tuple[str, int]:
+        """Serialize items starting at `offset` into one JSON array staying within the payload size limit.
+
+        The array holds as many consecutive items as fit below `_EFFECTIVE_LIMIT_SIZE`, always at least one.
+        The result is compact JSON - it goes on the wire, so indentation and separator padding would be
+        pure overhead. This is CPU-bound and blocking; call it via `asyncio.to_thread`.
 
         Args:
-            item: The item to serialize.
-            index: Index of the item, used for error context.
+            items: The items to serialize.
+            offset: Index of the first item to serialize.
 
         Returns:
-            Serialized JSON string.
+            The JSON array string and the index of the first item that did not fit into it.
 
         Raises:
-            ValueError: If item is not JSON serializable or exceeds size limit.
+            ValueError: If an item is not JSON serializable or on its own exceeds the size limit.
         """
-        s = ' ' if index is None else f' at index {index} '
-
-        try:
-            payload = await json_dumps(item)
-        except Exception as exc:
-            raise ValueError(f'Data item{s}is not serializable to JSON.') from exc
-
-        payload_size = ByteSize(len(payload.encode('utf-8')))
-        if payload_size > cls._EFFECTIVE_LIMIT_SIZE:
-            raise ValueError(f'Data item{s}is too large (size: {payload_size}, limit: {cls._EFFECTIVE_LIMIT_SIZE})')
-
-        return payload
-
-    async def _chunk_by_size(self, items: AsyncIterator[str]) -> AsyncIterator[str]:
-        """Yield chunks of JSON arrays composed of input strings, respecting a size limit.
+        payloads: list[str] = []
+        chunk_size = ByteSize(2)  # Add 2 bytes for [] wrapper.
 
-        Groups an iterable of JSON string payloads into larger JSON arrays, ensuring the total size
-        of each array does not exceed `EFFECTIVE_LIMIT_SIZE`. Each output is a JSON array string that
-        contains as many payloads as possible without breaching the size threshold, maintaining the
-        order of the original payloads. Assumes individual items are below the size limit.
+        for index in range(offset, len(items)):
+            try:
+                payload = json.dumps(items[index], ensure_ascii=False, separators=(',', ':'), default=str)
+            except Exception as exc:
+                raise ValueError(f'Data item at index {index} is not serializable to JSON.') from exc
 
-        Args:
-            items: Iterable of JSON string payloads.
-
-        Yields:
-            Strings representing JSON arrays of payloads, each staying within the size limit.
-        """
-        last_chunk_size = ByteSize(2)  # Add 2 bytes for [] wrapper.
-        current_chunk = []
-
-        async for payload in items:
             payload_size = ByteSize(len(payload.encode('utf-8')))
+            if payload_size > cls._EFFECTIVE_LIMIT_SIZE:
+                raise ValueError(
+                    f'Data item at index {index} is too large '
+                    f'(size: {payload_size}, limit: {cls._EFFECTIVE_LIMIT_SIZE})'
+                )
+
+            if payloads and chunk_size + payload_size > cls._EFFECTIVE_LIMIT_SIZE:
+                return f'[{",".join(payloads)}]', index
 
-            if last_chunk_size + payload_size <= self._EFFECTIVE_LIMIT_SIZE:
-                current_chunk.append(payload)
-                last_chunk_size += payload_size + ByteSize(1)  # Add 1 byte for ',' separator.
-            else:
-                yield f'[{",".join(current_chunk)}]'
-                current_chunk = [payload]
-                last_chunk_size = payload_size + ByteSize(2)  # Add 2 bytes for [] wrapper.
+            payloads.append(payload)
+            chunk_size += payload_size + ByteSize(1)  # Add 1 byte for ',' separator.
 
-        yield f'[{",".join(current_chunk)}]'
+        return f'[{",".join(payloads)}]', len(items)

tests/unit/storage_clients/test_apify_dataset_client.py

@@ -1,10 +1,13 @@
 from __future__ import annotations
 
 import asyncio
-from unittest.mock import AsyncMock
+import json
+from unittest.mock import AsyncMock, Mock
 
 import pytest
 
+from crawlee._utils.byte_size import ByteSize
+
 from apify.storage_clients._apify._dataset_client import ApifyDatasetClient
 
 
@@ -31,3 +34,58 @@ async def test_drop_calls_api_delete() -> None:
     client, api_client = _make_dataset_client()
     await client.drop()
     api_client.delete.assert_awaited_once()
+
+
+async def test_push_data_sends_compact_json() -> None:
+    """Pushed payloads carry no indentation or separator padding."""
+    client, api_client = _make_dataset_client()
+
+    await client.push_data([{'id': 1, 'name': 'first'}, {'id': 2, 'name': 'second'}])
+
+    chunk = api_client.push_items.await_args.kwargs['items']
+    assert chunk == '[{"id":1,"name":"first"},{"id":2,"name":"second"}]'
+
+
+async def test_push_data_serializes_in_a_single_thread_hop_per_chunk(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Serialization is offloaded once per pushed chunk rather than once per item."""
+    to_thread = Mock(wraps=asyncio.to_thread)
+    monkeypatch.setattr(asyncio, 'to_thread', to_thread)
+    client, api_client = _make_dataset_client()
+
+    await client.push_data([{'id': i} for i in range(500)])
+
+    api_client.push_items.assert_awaited_once()
+    assert to_thread.call_count == 1
+
+
+async def test_push_data_splits_items_into_chunks_within_the_size_limit(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Items are pushed in several chunks, each staying within the payload size limit."""
+    monkeypatch.setattr(ApifyDatasetClient, '_EFFECTIVE_LIMIT_SIZE', ByteSize(100))
+    client, api_client = _make_dataset_client()
+    items = [{'value': 'x' * 30} for _ in range(5)]
+
+    await client.push_data(items)
+
+    chunks = [call.kwargs['items'] for call in api_client.push_items.await_args_list]
+    assert len(chunks) > 1
+    assert all(len(chunk.encode('utf-8')) <= 100 for chunk in chunks)
+    assert [item for chunk in chunks for item in json.loads(chunk)] == items
+
+
+async def test_push_data_rejects_an_oversized_item(monkeypatch: pytest.MonkeyPatch) -> None:
+    """An item exceeding the payload size limit raises with its index."""
+    monkeypatch.setattr(ApifyDatasetClient, '_EFFECTIVE_LIMIT_SIZE', ByteSize(100))
+    client, _ = _make_dataset_client()
+
+    with pytest.raises(ValueError, match='at index 1 is too large'):
+        await client.push_data([{'id': 1}, {'value': 'x' * 200}])
+
+
+async def test_push_data_rejects_a_non_serializable_item() -> None:
+    """An item that cannot be serialized to JSON raises with its index."""
+    client, _ = _make_dataset_client()
+    circular: dict = {}
+    circular['self'] = circular
+
+    with pytest.raises(ValueError, match='at index 0 is not serializable'):
+        await client.push_data(circular)