fix(serializers): byte-aware Arrow batch sizing holds the memory bound on skewed frames (LAB-110)

src/cachekit/serializers/arrow_serializer.py

@@ -57,15 +57,29 @@
 _RELEASE_POOL_THRESHOLD = 4 * 1024 * 1024
 
 
-def _bounded_chunksize(table: pa.Table) -> int | None:  # type: ignore[name-defined]
-    """Rows per IPC record-batch so each batch is ~_TARGET_BATCH_BYTES, regardless of width.
-
-    Returns None for empty tables (nothing to chunk). Never returns 0.
+def _write_bounded_batches(writer: pa.ipc.RecordBatchFileWriter, table: pa.Table) -> None:  # type: ignore[name-defined]
+    """Write ``table`` as record batches whose ACTUAL bytes stay <= _TARGET_BATCH_BYTES.
+
+    A uniform rows-per-batch estimate (nbytes // num_rows) breaks on skewed/variable-width
+    frames: a clustered run of wide rows (e.g. JSON-blob cells) can overshoot the target
+    ~12x (#161), spiking the per-batch zstd working set in exactly the OOM regime this
+    bound exists to prevent. Instead, binary-search the largest row count whose slice
+    actually fits the byte budget — Table.slice() is zero-copy and .nbytes accounts for
+    slice offsets, so each probe is O(columns), not O(bytes). Batch bytes are monotone in
+    row count, so the search is valid; a single row wider than the target gets its own
+    batch (rows can't split). Empty tables write no batches (schema-only IPC file).
     """
-    if table.num_rows <= 0:
-        return None
-    bytes_per_row = max(1, table.nbytes // table.num_rows)
-    return max(1, _TARGET_BATCH_BYTES // bytes_per_row)
+    offset, total = 0, table.num_rows
+    while offset < total:
+        lo, hi, rows = 2, total - offset, 1
+        while lo <= hi:
+            mid = (lo + hi) // 2
+            if table.slice(offset, mid).nbytes <= _TARGET_BATCH_BYTES:
+                rows, lo = mid, mid + 1
+            else:
+                hi = mid - 1
+        writer.write_table(table.slice(offset, rows))
+        offset += rows
 
 
 class ArrowSerializer:
@@ -227,14 +241,14 @@ def serialize(self, obj: Any) -> tuple[bytes, SerializationMetadata]:  # type: i
             # Serialize to Arrow IPC. Compression (when enabled) runs per record-batch, so
             # writing in bounded batches keeps the compressor's working set bounded (one big
             # batch makes the codec allocate a full-size working buffer — measured ~3.6x the
-            # payload). Size each batch to ~8 MiB regardless of schema width. compression=None
+            # payload). Batches are sized by actual bytes (<= 8 MiB), not a uniform row-count
+            # estimate, so the bound holds for skewed/wide frames too (#161). compression=None
             # writes uncompressed IPC, which the File backend reads zero-copy via mmap (#171,
             # plaintext, pandas return only).
-            max_chunksize = _bounded_chunksize(table)
             sink = pa.BufferOutputStream()
             write_options = pa.ipc.IpcWriteOptions(compression=self.compression) if self.compression else None
             with pa.ipc.new_file(sink, table.schema, options=write_options) as writer:
-                writer.write_table(table, max_chunksize=max_chunksize)
+                _write_bounded_batches(writer, table)
             del table  # free the Arrow table before materializing the IPC bytes (lowers peak)
 
             # Always integrity-protect: hash over the buffer's memoryview (no copy), then

tests/unit/test_arrow_serializer.py

@@ -462,6 +462,70 @@ def test_auto_falls_back_to_zstd_when_settings_module_unavailable(self, monkeypa
         assert ArrowSerializer._resolve_compression("auto") == "zstd"
 
 
+class TestBoundedBatching:
+    """Record batches must honor the byte budget for SKEWED frames, not just uniform ones.
+
+    The pre-fix estimate (nbytes // num_rows, rows-only cap) let a clustered run of wide
+    rows produce a ~100 MiB batch against the 8 MiB target (#161) — and zstd's working
+    set scales with batch size, so the OOM bound this feature exists for silently broke.
+    """
+
+    # IPC padding/offset-shifting can nudge a read-back batch slightly past the slice's
+    # in-memory nbytes; the guarantee we defend is "no order-of-magnitude overshoot".
+    TOLERANCE = 1.05
+
+    @staticmethod
+    def _read_batches(data: bytes) -> list:
+        """Strip the 8-byte checksum envelope and return decompressed IPC record batches."""
+        reader = pa.ipc.open_file(pa.py_buffer(memoryview(data)[8:]))
+        return [reader.get_batch(i) for i in range(reader.num_record_batches)]
+
+    def test_skewed_frame_batches_stay_bounded(self):
+        """Issue #161's exact case: 100k tiny rows + 50 clustered 2 MiB cells."""
+        from cachekit.serializers.arrow_serializer import _TARGET_BATCH_BYTES
+
+        df = pd.DataFrame({"blob": ["y"] * 100_000 + ["x" * 2_000_000] * 50})
+        serializer = ArrowSerializer()
+        data, meta = serializer.serialize(df)
+
+        batches = self._read_batches(data)
+        worst = max(b.nbytes for b in batches)
+        assert worst <= _TARGET_BATCH_BYTES * self.TOLERANCE, (
+            f"batch of {worst / 2**20:.1f} MiB exceeds the {_TARGET_BATCH_BYTES / 2**20:.0f} MiB "
+            f"target (pre-fix overshoot was ~12.5x)"
+        )
+        pd.testing.assert_frame_equal(serializer.deserialize(data, meta), df)
+
+    def test_uniform_frame_batching_not_fragmented(self):
+        """Common case unchanged: uniform frames still get few, near-target batches."""
+        import math
+
+        from cachekit.serializers.arrow_serializer import _TARGET_BATCH_BYTES
+
+        df = pd.DataFrame({f"c{i}": [float(j) for j in range(500_000)] for i in range(8)})  # ~32 MiB
+        table_nbytes = pa.Table.from_pandas(df, preserve_index=None).nbytes
+        data, meta = ArrowSerializer().serialize(df)
+
+        batches = self._read_batches(data)
+        assert all(b.nbytes <= _TARGET_BATCH_BYTES * self.TOLERANCE for b in batches)
+        # Byte-aware batching must not fragment: batches are maximal, so the count stays
+        # at the minimum the budget allows (uniform estimate gave the same).
+        assert len(batches) <= math.ceil(table_nbytes / _TARGET_BATCH_BYTES) + 1
+        pd.testing.assert_frame_equal(ArrowSerializer().deserialize(data, meta), df)
+
+    def test_single_row_wider_than_target_gets_own_batch(self):
+        """A row that can't fit the budget must still serialize (1-row batch, no hang)."""
+        from cachekit.serializers.arrow_serializer import _TARGET_BATCH_BYTES
+
+        df = pd.DataFrame({"blob": ["small"] * 10 + ["x" * (12 * 2**20)] + ["small"] * 10})
+        serializer = ArrowSerializer()
+        data, meta = serializer.serialize(df)
+
+        for batch in self._read_batches(data):
+            assert batch.nbytes <= _TARGET_BATCH_BYTES * self.TOLERANCE or batch.num_rows == 1
+        pd.testing.assert_frame_equal(serializer.deserialize(data, meta), df)
+
+
 class TestIntegrityAlwaysOn:
     """DATA IS SACRED: corruption is always detected, even with integrity_checking=False."""