fix(serializers): byte-aware Arrow batch sizing holds the memory bound on skewed frames (LAB-110) - #244
Conversation
…d on skewed frames (LAB-110) _bounded_chunksize estimated bytes-per-row as a uniform average (table.nbytes // table.num_rows) and capped record batches by row count only, so a clustered run of wide rows (e.g. JSON-blob cells) produced batches ~12x over the 8 MiB target — and zstd compresses per-batch, so the compressor working set spiked in exactly the OOM regime the bound exists to prevent (#161). Replace the estimate with _write_bounded_batches: binary-search the largest row count whose slice ACTUALLY fits the byte budget per batch. Table.slice() is zero-copy and .nbytes accounts for slice offsets, so each probe is O(columns); batch bytes are monotone in row count, so the search is valid. A single row wider than the target gets its own batch. Uniform frames still get maximal near-target batches (no fragmentation; 122 MiB frame serializes in ~0.34s incl. zstd). On the issue's 100k tiny rows + 50x2 MiB cells case the worst batch drops from 95.4 MiB (11.9x target) to 7.6 MiB. No wire-format change: Arrow IPC files were already multi-batch; only the internal batch boundaries move. Closes #161 Co-authored-by: multica-agent <github@multica.ai>
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: ASSERTIVE Plan: Pro Run ID: 📒 Files selected for processing (2)
WalkthroughArrow serialization now writes byte-bounded IPC record batches using binary-search row selection. Tests inspect decompressed batches for size limits, fragmentation, oversized-row handling, and DataFrame round trips. ChangesArrow batching
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant ArrowSerializer
participant _write_bounded_batches
participant PyArrowIPCWriter
ArrowSerializer->>_write_bounded_batches: pass Arrow table
_write_bounded_batches->>_write_bounded_batches: binary-search rows within byte budget
_write_bounded_batches->>PyArrowIPCWriter: write selected table slices
Possibly related PRs
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
27Bslash6
changed the title
|
@coderabbitai review |
✅ Action performedReview finished.
|
|
@coderabbitai full review |
✅ Action performedFull review finished. |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
|
@27Bslash6 — automated review gate cleared for LAB-110, ready for your signoff / merge.
Over to you for human review and merge. |
|
Tip For best results, initiate chat on the files or code changes. 🧩 Analysis chain🏁 Script executed: #!/bin/bash
set -euo pipefail
echo '=== Changed files and diff ==='
git diff --unified=80 HEAD^ HEAD -- \
src/cachekit/serializers/arrow_serializer.py \
tests/unit/test_arrow_serializer.py
echo '=== Serializer structure ==='
ast-grep outline src/cachekit/serializers/arrow_serializer.py --items all --view expanded
echo '=== Test structure ==='
ast-grep outline tests/unit/test_arrow_serializer.py --items all --view expandedLength of output: 35111
The binary search correctly chooses the largest byte-bounded slice, always progresses for oversized single rows, and retains schema-only behaviour for empty tables. The new tests directly validate skewed, uniform, and oversize-row batch boundaries as well as round trips. The change is scoped to internal Arrow IPC record-batch boundaries and does not alter the envelope or public wire format. I cannot merge the PR directly. For an explicit CodeRabbit approval, post |
1 participant
Fixes #161 · Closes LAB-110
Problem
_bounded_chunksizeestimated bytes-per-row as a uniform average (table.nbytes // table.num_rows) and capped record batches by row count only. On a skewed frame, a clustered run of wide rows lands in one batch far over the 8 MiB target — and since zstd compresses per-batch, the compressor's working set spikes in exactly the OOM regime the bounded-batching feature exists to prevent.Measured on the issue's case (100k tiny rows + 50 clustered 2 MiB cells): worst batch 95.4 MiB — 11.9x the target.
Fix
Replace the uniform estimate with
_write_bounded_batches: per batch, binary-search the largest row count whose slice actually fits the byte budget.Table.slice()is zero-copy and.nbytesaccounts for slice offsets, so each probe is O(columns), not O(bytes).Post-fix, the same skewed frame's worst batch is 7.6 MiB (under target). Uniform frames still get maximal, near-target batches — no fragmentation — and a 122 MiB uniform frame serializes in ~0.34 s including zstd, so the probe overhead is noise.
Acceptance criteria
TestBoundedBatching)test_uniform_frame_batching_not_fragmentedguards against fragmentation; all 59 existing serializer tests pass)Testing
tests/unit/test_arrow_serializer.py: 62 tests pass (3 new inTestBoundedBatching: skewed frame, uniform no-fragmentation, single-oversized-row)ruff check+ruff formatclean; pre-commit (incl. basedpyright) greentests/performance/+ fixture-teardown noise, reproduced identically on cleanmain; serializer benchmarks passDocs
No user-facing doc documents record-batch sizing internals (checked
docs/, README, feature matrices — the only "batch" mentions are about batching queries). The in-code comments/docstring describing the strategy are updated in this diff.Summary by CodeRabbit
Bug Fixes
Tests