diff --git a/src/cachekit/serializers/arrow_serializer.py b/src/cachekit/serializers/arrow_serializer.py index 61e89d1..3a70856 100644 --- a/src/cachekit/serializers/arrow_serializer.py +++ b/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 diff --git a/tests/unit/test_arrow_serializer.py b/tests/unit/test_arrow_serializer.py index 6aa3ba9..972df67 100644 --- a/tests/unit/test_arrow_serializer.py +++ b/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."""