perf: use plain int byte arithmetic in the dataset chunking loop

Author: vdusek

src/apify/storage_clients/_apify/_dataset_client.py

@@ -214,9 +214,8 @@ async def iterate_items(
     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`.
+        The array holds as many consecutive items as fit within `_EFFECTIVE_LIMIT_SIZE`, always at least one. Output
+        is compact JSON - it goes straight on the wire. This is CPU-bound and blocking; call it via `asyncio.to_thread`.
 
         Args:
             items: The items to serialize.
@@ -228,26 +227,27 @@ def _serialize_chunk(cls, items: Sequence[Mapping[str, JsonSerializable]], offse
         Raises:
             ValueError: If an item is not JSON serializable or on its own exceeds the size limit.
         """
+        limit = cls._EFFECTIVE_LIMIT_SIZE.bytes
         payloads: list[str] = []
-        chunk_size = ByteSize(2)  # Add 2 bytes for [] wrapper.
+        chunk_size = 2  # Add 2 bytes for [] wrapper.
 
         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
 
-            payload_size = ByteSize(len(payload.encode('utf-8')))
-            if payload_size > cls._EFFECTIVE_LIMIT_SIZE:
+            payload_size = len(payload.encode('utf-8'))
+            if payload_size > limit:
                 raise ValueError(
                     f'Data item at index {index} is too large '
-                    f'(size: {payload_size}, limit: {cls._EFFECTIVE_LIMIT_SIZE})'
+                    f'(size: {ByteSize(payload_size)}, limit: {cls._EFFECTIVE_LIMIT_SIZE})'
                 )
 
-            if payloads and chunk_size + payload_size > cls._EFFECTIVE_LIMIT_SIZE:
+            if payloads and chunk_size + payload_size > limit:
                 return f'[{",".join(payloads)}]', index
 
             payloads.append(payload)
-            chunk_size += payload_size + ByteSize(1)  # Add 1 byte for ',' separator.
+            chunk_size += payload_size + 1  # Add 1 byte for ',' separator.
 
         return f'[{",".join(payloads)}]', len(items)

tests/unit/storage_clients/test_apify_dataset_client.py

@@ -48,14 +48,29 @@ async def test_push_data_sends_compact_json() -> None:
 
 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."""
+    monkeypatch.setattr(ApifyDatasetClient, '_EFFECTIVE_LIMIT_SIZE', ByteSize(200))
     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
+    assert api_client.push_items.await_count > 1
+    assert to_thread.call_count == api_client.push_items.await_count
+
+
+async def test_push_data_makes_progress_when_an_item_fills_a_whole_chunk(monkeypatch: pytest.MonkeyPatch) -> None:
+    """An item that fits the limit only without the array wrapper still yields one chunk per item."""
+    items = [{'value': 'x' * 30} for _ in range(3)]
+    payloads = [json.dumps(item, ensure_ascii=False, separators=(',', ':')) for item in items]
+    monkeypatch.setattr(ApifyDatasetClient, '_EFFECTIVE_LIMIT_SIZE', ByteSize(len(payloads[0].encode('utf-8'))))
+    client, api_client = _make_dataset_client()
+
+    async with asyncio.timeout(5):
+        await client.push_data(items)
+
+    chunks = [call.kwargs['items'] for call in api_client.push_items.await_args_list]
+    assert chunks == [f'[{payload}]' for payload in payloads]
 
 
 async def test_push_data_splits_items_into_chunks_within_the_size_limit(monkeypatch: pytest.MonkeyPatch) -> None: