feat: add BatchNormalizer: re-chunk data source batches by rows and bytes (target_batch_size_bytes)#23390
Draft
adriangb wants to merge 11 commits into
Draft
feat: add BatchNormalizer: re-chunk data source batches by rows and bytes (target_batch_size_bytes)#23390adriangb wants to merge 11 commits into
adriangb wants to merge 11 commits into
Conversation
Adds a sans-IO BatchNormalizer plus a BatchNormalizerStream adapter that re-chunks a stream of RecordBatches towards a target row count AND byte size. Per input batch it either passes through (zero copy, wide acceptance band), coalesces small batches (flushing on rows or bytes, whichever first), splits oversized batches into compact ~target-sized copies, or compacts batches that pin far more memory than they use. Splitting/compaction copies via take + view-array GC: slice/concat are zero-copy for a single input and would keep the oversized parent buffers alive. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01UcPTREZVLXsSZDRCae33gm
…ource batches When set, DataSourceExec wraps its output in BatchNormalizerStream instead of the row-only BatchSplitStream, re-chunking scan output towards both the row target (batch_size) and the new byte target. Default is unset (row-based splitting only, unchanged behavior). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01UcPTREZVLXsSZDRCae33gm
Contributor
Author
|
run benchmarks |
github-actions
Bot
added
documentation
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing batch-normalizer (9e3c032) to c2e3473 (merge-base) diff using: tpcds File an issue against this benchmark runner |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing batch-normalizer (9e3c032) to c2e3473 (merge-base) diff using: tpch File an issue against this benchmark runner |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing batch-normalizer (9e3c032) to c2e3473 (merge-base) diff using: clickbench_partitioned File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: CPU Details (lscpu)Details
Resource Usagetpch — base (merge-base)
tpch — branch
File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: CPU Details (lscpu)Details
Resource Usagetpcds — base (merge-base)
tpcds — branch
File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: CPU Details (lscpu)Details
Resource Usageclickbench_partitioned — base (merge-base)
clickbench_partitioned — branch
File an issue against this benchmark runner |
|
Thank you for opening this pull request! Reviewer note: cargo-semver-checks reported the current version number is not SemVer-compatible with the changes in this pull request (compared against the base branch). Details |
Existing suites (TPC-H, ClickBench) decode to ~1MiB batches and never exercise the oversized-batch regime that target_batch_size_bytes addresses. This benchmark generates incompressible multi-KiB string values (default: 64k rows x 16KiB over 8 files, so each file decodes to a single ~134MiB batch within the 8192-row limit) and runs 5 queries stressing scan+filter, sort (small key / large key), aggregation, and hash join. Designed for A/B runs of target_batch_size_bytes, optionally under --memory-limit. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01UcPTREZVLXsSZDRCae33gm
Contributor
Author
|
| scenario | off | on |
|---|---|---|
| unlimited memory, Q1-Q5 | baseline | +5-10% (worst case: every batch is split, ~1GiB copied per query) |
--memory-limit 4G Q2 sort, n=4 |
819ms | 630ms (-23%) |
--memory-limit 4G Q2 sort, n=12 |
427ms | 360ms (-16%) |
--memory-limit 2G Q2/Q3 sorts |
FAIL in ExternalSorter insert (cannot reserve 2x134MB against fair share) |
FAIL, but much later: merge phase |
64KiB values (536MB decoded batches), --memory-limit 2G: off fails instantly in ExternalSorter; on sorts + spills fine and then fails in the merge with a telling profile:
ExternalSorterMerge[0]#2(can spill: false) consumed 1106.2 MB
ExternalSorterMerge[1]#4(can spill: false) consumed 878.9 MB
Failed to allocate additional 639.1 MB for SortPreservingMergeExec[0]
A 639MB single allocation despite 16MiB input chunks: the sorter re-chunks its own output at batch_size rows (8192 x 64KiB = 536MB), reconstructing oversized batches downstream of the normalizer. Input-side normalization visibly helps (spill-heavy sorts get 16-23% faster, and the failure point moves from "cannot buffer one batch" to the merge), but completing these queries under tight limits needs byte-aware output chunking in sort/merge — the operator-emit-path follow-up called out in the PR description.
🤖 Generated with Claude Code
The sorter and streaming merge re-chunk their output at batch_size rows only, so wide rows rebuild oversized batches downstream of a normalized scan: sorting 536MB decoded batches under a 2GB limit failed in the merge phase with ~1.1GB ExternalSorterMerge reservations (sized by the largest spilled batch) and a single 639MB SortPreservingMergeExec grow. When datafusion.execution.target_batch_size_bytes is set: - IncrementalSortIterator (sorted-output chunking, also used by the aggregate spill path) reduces its chunk row count so chunks are about the byte target - BatchBuilder tracks the estimated bytes of in-progress rows and the streaming merge emits when the row count OR the byte target is reached, plumbed through StreamingMergeBuilder/MultiLevelMergeBuilder and read from session config by SortExec, SortPreservingMergeExec and GroupedHashAggregateStream Unset (default) keeps row-only chunking, byte-for-byte. With this, the large_values 64KiB benchmark (536MB decoded batches, 2GB memory limit) goes from ResourcesExhausted to completing: spill file batches shrink to the byte target, cutting merge reservations from ~1.1GB to ~128MB per partition. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01UcPTREZVLXsSZDRCae33gm
…ill output Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01UcPTREZVLXsSZDRCae33gm
Contributor
Author
Update: byte-aware sort/merge output chunking (commit 88bdade)Following up on the finding above (the sorter re-chunks its own output at Result on the
The mechanism is visible in the memory profiles: spill-file batches now respect the byte target, so per-partition merge reservations (sized from the largest spilled batch) drop from ~1.1GB to ~128MB. Remaining boundary: at 12 concurrent sort partitions in a pool ~2x the data size, the aggregate of per-partition merge reservations plus the top-level SPM buffering can still exhaust the fair pool — that is a genuine concurrent-capacity limit (and FairSpillPool interplay) rather than byte-blindness, left for follow-up. 🤖 Generated with Claude Code |
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01UcPTREZVLXsSZDRCae33gm
Contributor
Author
|
run benchmarks |
Contributor
Author
|
run benchmark tpch tpch_mem clickbench_partitioned tpcds external_aggr baseline: |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing batch-normalizer (047ea98) to c2e3473 (merge-base) diff using: clickbench_partitioned File an issue against this benchmark runner |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing batch-normalizer (047ea98) to 047ea98 diff using: tpch_mem File an issue against this benchmark runner |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing batch-normalizer (047ea98) to 047ea98 diff using: clickbench_partitioned File an issue against this benchmark runner |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing batch-normalizer (047ea98) to c2e3473 (merge-base) diff using: tpcds File an issue against this benchmark runner |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing batch-normalizer (047ea98) to c2e3473 (merge-base) diff using: tpch File an issue against this benchmark runner |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing batch-normalizer (2cf5402) to 2cf5402 diff using: tpch File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: CPU Details (lscpu)Details
Resource Usagetpch — base (merge-base)
tpch — branch
File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: CPU Details (lscpu)Details
Resource Usageclickbench_partitioned — base (merge-base)
clickbench_partitioned — branch
File an issue against this benchmark runner |
Contributor
Author
ClickBench A/B regression: diagnosed and fixed (commit 3a03ec5)The earlier on/off A/B showed clickbench_partitioned +5.1% with the flag on (17 queries slower, up to 1.2x on string-heavy scans). Fix: the absolute-waste floor now scales to Fresh GKE A/B on the fixed head (
The former regressors (Q13/Q17/Q20/Q21/Q22/Q33/Q34) are all back to no-change. Residuals are Q23 at 1.07x (a wide-row 🤖 Generated with Claude Code |
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01UcPTREZVLXsSZDRCae33gm
Adds datafusion.execution.adaptive_target_batch_size (default true): when target_batch_size_bytes is unset and the memory pool reports a finite limit, the byte target is derived as pool_size / target_partitions / 16, capped at 16MiB. The /16 margin covers per-batch multipliers in the sort/merge machinery (2x sort buffering, ~4x spill-merge reservation per stream, SPM buffering all streams); the ceiling is where tight-pool merges stay reliably stable while larger batches add nothing. Adaptive targeting deliberately stays off when there is no memory limit (nothing to protect: identical code path to today) and when the derived target falls below 1MiB (a pool that small cannot be saved by re-chunking, and carefully crafted small-pool OOM scenarios keep their exact behavior). An explicit target_batch_size_bytes always wins. With no configuration at all, the large_values benchmark under a 2GB limit goes from failing its sort queries to completing them at both 4 and 12 partitions, including the 2GiB-dataset variant. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01UcPTREZVLXsSZDRCae33gm
Contributor
Author
Adaptive byte target: derived from the memory limit (commit 775c059)Per review discussion, the byte target can now derive itself instead of requiring manual configuration. New option
Acceptance on
Trade-off note: with a finite limit, already-passing queries whose batches exceed the derived threshold pay the ~5-10% compaction tax (a scan-only query doesn't know whether a downstream sort needs protection). Full workspace test suite passes, including the memory-limit OOM-assertion tests. 🤖 Generated with Claude Code |
Contributor
Author
|
run benchmarks |
Contributor
Author
|
run benchmark tpch sort_tpch clickbench_partitioned external_aggr env: baseline: |
|
Benchmark for this request failed. Last 20 lines of output: Click to expandFile an issue against this benchmark runner |
3 similar comments
|
Benchmark for this request failed. Last 20 lines of output: Click to expandFile an issue against this benchmark runner |
|
Benchmark for this request failed. Last 20 lines of output: Click to expandFile an issue against this benchmark runner |
|
Benchmark for this request failed. Last 20 lines of output: Click to expandFile an issue against this benchmark runner |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing batch-normalizer (775c059) to c2e3473 (merge-base) diff using: tpcds File an issue against this benchmark runner |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing batch-normalizer (775c059) to c2e3473 (merge-base) diff using: clickbench_partitioned File an issue against this benchmark runner |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing batch-normalizer (775c059) to 775c059 diff using: sort_tpch File an issue against this benchmark runner |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing batch-normalizer (775c059) to 775c059 diff using: clickbench_partitioned File an issue against this benchmark runner |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing batch-normalizer (775c059) to c2e3473 (merge-base) diff using: tpch File an issue against this benchmark runner |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing batch-normalizer (775c059) to 775c059 diff using: tpch File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: CPU Details (lscpu)Details
Resource Usagetpch — base (merge-base)
tpch — branch
File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: CPU Details (lscpu)Details
Resource Usagesort_tpch — base (merge-base)
sort_tpch — branch
File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: CPU Details (lscpu)Details
Resource Usagetpch — base (merge-base)
tpch — branch
File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: CPU Details (lscpu)Details
Resource Usagetpcds — base (merge-base)
tpcds — branch
File an issue against this benchmark runner |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing batch-normalizer (775c059) to 775c059 diff using: external_aggr File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: CPU Details (lscpu)Details
Resource Usageclickbench_partitioned — base (merge-base)
clickbench_partitioned — branch
File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: CPU Details (lscpu)Details
Resource Usageclickbench_partitioned — base (merge-base)
clickbench_partitioned — branch
File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: CPU Details (lscpu)Details
Resource Usageexternal_aggr — base (merge-base)
external_aggr — branch
File an issue against this benchmark runner |
Contributor
Author
Final benchmark picture (head 775c059, adaptive targeting on by default)Two runs prove the two halves of the design contract: 1. Default run (merge-base vs this branch, no memory limit). With an infinite pool the adaptive path derives no target and the normalizer wrapper is never installed, so this is a same-code-path regression check:
2. Same-SHA A/B under
Summary: with no memory limit the feature is provably inert; under a memory limit on standard (small-batch) data it costs ~0.7-1.5% total; and for the workloads it exists for ( 🤖 Generated with Claude Code |
Arrow's take and interleave copy view-array views but share the underlying data buffers, so byte-targeted sorted chunks and merged output batches pinned every source batch they drew from: sorting 1GB of 16KiB Utf8View values produced 4MB-logical batches retaining 128MB each (SortExec output_bytes=31.9GB), and downstream reservations (BatchBuilder, SortPreservingMergeExec) count retained size -- so the byte target made memory accounting worse, failing ORDER BY <string> at a 4GB limit where row-only chunking passed. Run gc_view_arrays on IncrementalSortIterator chunks and BatchBuilder output when a byte target is set (the GC is heuristic-gated and no-ops on non-view or non-wasteful batches; row-only mode is untouched). large_values Q3 (ORDER BY val) at 4G: fails -> completes, and adaptive is now on par with row-only where both pass; the 64KiB spill sorts at 2G drop from ~2.3s to ~1.1s. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01UcPTREZVLXsSZDRCae33gm
Contributor
Author
|
| scenario | off | auto (default) |
|---|---|---|
| unlimited, all 5 queries x2 rounds | baseline | identical within ~1% (provably same code path) |
| 2G n=4 | Q2/Q3 FAIL | all 5 pass (Q2 488ms, Q3 506ms); passing queries +3-6% |
| 2G n=12 | Q2/Q3 FAIL | all 5 pass (Q2 367ms, Q3 500ms — the former Q3 flake is gone post-fix) |
| 4G n=12 | all pass | all pass; +10-39% on this worst-case dataset (adaptive still splits 128MB batches while off fits comfortably — the insurance premium when a finite limit is set; adaptive_target_batch_size=false opts out) |
| 64KiB values (2GiB dataset), 2G, n=4 / n=12 | Q2 FAIL | Q2 passes both (~1.1s, down from ~2.3s before the view-GC fix) |
🤖 Generated with Claude Code
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Which issue does this PR close?
Rationale for this change
Operators bound their output batches by row count (
datafusion.execution.batch_size) only. With wide rows (large string values, wide schemas) a batch within the row limit can still be arbitrarily large in bytes — multi-GB batches break memory accounting assumptions and cause OOMs. Conversely, streams of many tiny batches waste per-batch overhead.The existing
BatchSplitStreamonDataSourceExecsplits by rows only, and splitting by zero-copyslicecannot release memory anyway: a slice of a 4GB batch still pins all 4GB of buffers (arrow'sconcatalso short-circuits single inputs to a zero-copy slice).This PR adds a batch normalizer: a stream wrapper that re-chunks data source output towards a target row count and byte size, opt-in via a new config option.
What changes are included in this PR?
BatchNormalizer/BatchNormalizerStream(datafusion/physical-plan/src/batch_normalizer.rs). Per input batch, one of four actions:batch_size / 2or logical bytes ≥target / 2, up to2xthe byte target. The wide acceptance band means near-target batches are never copied.BatchCoalescer) and flushed when the buffer reaches the row target or the byte target, whichever first.take_record_batch+ view-array GC), produced incrementally (one chunk per poll) so peak memory is the parent plus ~one chunk. Copying is load-bearing: zero-copy slices would keep the oversized parent alive.datafusion.execution.target_batch_size_bytes(defaultNone= unchanged behavior). When set,DataSourceExecwraps its output inBatchNormalizerStreaminstead of the row-onlyBatchSplitStream.batches_passed_through,batches_coalesced,batches_split,batches_compacted.Planned follow-ups (intentionally not in this PR): memory-pool reservation for the normalizer's buffer with a spill fallback (following the
RepartitionExecOutputChannel/SpillPoolpattern), and byte-aware bounding for operator emit paths (join/aggregate output amplification is not addressable at the scan).Local benchmark observations (laptop, noisy — GKE runs requested)
Interleaved off/on A/B (2 rounds, per-query minimums, same binary, 16MiB target via
DATAFUSION_EXECUTION_TARGET_BATCH_SIZE_BYTES):batches_passed_through, 0 copies), so the residual delta is per-batch size measurement overhead and/or laptop noise (off-vs-off noise floor spans 0.66-1.22x per query).Laptop numbers are close to the noise floor; treating the dedicated benchmark runner as the source of truth.
Are these changes tested?
Yes — written test-first (TDD):
SessionContextscans.information_schema.sltregenerated.Are there any user-facing changes?
New opt-in config option
datafusion.execution.target_batch_size_bytes(documented inconfigs.md). Default behavior is unchanged. No breaking API changes.🤖 Generated with Claude Code
https://claude.ai/code/session_01UcPTREZVLXsSZDRCae33gm