fix: make UnnestExec respect datafusion.execution.batch_size - #24384
Open
andygrove wants to merge 3 commits into
Open
fix: make UnnestExec respect datafusion.execution.batch_size#24384andygrove wants to merge 3 commits into
andygrove wants to merge 3 commits into
Conversation
`UnnestExec` emitted exactly one output batch per input batch, however many rows the unnesting produced. An 8192-row batch of 100-element lists came back as a single 819,200-row batch, and `batch_size` was never consulted at all. Besides handing downstream operators arbitrarily large batches, this meant peak memory scaled with the input batch size times the list length rather than with `batch_size`. `UnnestStream` now consumes each input batch in chunks. `find_longest_length` already computes how many output rows each input row expands into, including null handling, so prefix-summing it gives exact chunk boundaries for depth-1 unnesting: each build produces at most `batch_size` rows, so the oversized intermediate is never materialized. Two cases can't be chunked on the input side and are handled by slicing the built batch instead: - a single input row whose list is longer than `batch_size`, since one row is never split across output batches - recursive unnest (`depth > 1`), where a row's expansion depends on inner list lengths that only exist once the outer levels are unnested Struct-only unnesting doesn't change the row count, so it's already bounded by the input batch size. Chunk boundaries fall on input-row boundaries, so an output batch can be shorter than `batch_size`; the guarantee is an upper bound, not an exact size.
github-actions
Bot
added
sqllogictest
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #24384 +/- ##
========================================
Coverage 81.19% 81.20%
========================================
Files 1110 1110
Lines 388737 389008 +271
Branches 388737 389008 +271
========================================
+ Hits 315626 315875 +249
- Misses 54527 54532 +5
- Partials 18584 18601 +17 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
…hunking With 10-element lists at batch_size=8 every input row overflows on its own, so the test only re-covered the oversized-build path that test_unnest_stream_single_row_exceeds_batch_size already pins. Use 3-element lists so chunks pack several input rows, and assert the exact batch shape.
andygrove
added a commit
to andygrove/datafusion-comet
that referenced
this pull request
Aug 14, 2026
- Use upstream's public `ListUnnest` instead of the copy in the vendored region. It is the one item there that is not private to `datafusion-physical-plan`, so copying it was unjustified and left Comet exporting a same-named twin of a public DataFusion type. - Correct the deletion trigger. apache#5210 tracks adopting upstream `unnest_outer` (apache/datafusion#22100), not batch size, so closing it is not a signal to delete this fork; the trigger is apache/datafusion#24384. - Stop claiming the vendored region is byte-identical. Comet's rustfmt reflows it (`max_width = 100` vs upstream's 90), so document the audit recipe instead. - Drop the comment claiming everything below it was Comet code, which contradicted the banner further down the same file. - Fix a test that did not test what it was named for: with 10-element arrays at `batch_size = 8` every row overflows, so `respects_batch_size` only re-covered the oversized-build path. Use 3-element arrays so chunks pack several input rows, and assert the exact `[6, 6, 6, 6, 6]` shape. - Collapse the two near-identical null-handling tests onto a shared helper, add `sizes`/`seq` test helpers, and use the `AsArray` idiom already in scope. - Simplify: `Option::filter` over a match with a wildcard arm, `Count` imported rather than fully qualified twice, and `elapsed_compute().timer()` without the needless clone, matching scan.rs and shuffle_scan.rs.
- Adopt the existing `BatchSplitStream` for output slicing rather than hand- rolling `split_off_head` + `pending_output`. `BatchSplitStream` tracks an offset into the current batch instead of re-slicing the tail on every emit, so unnesting a large recursive expansion is now O(m) rather than O(m^2), and the metrics it publishes (`batches_split`, correct `output_bytes` via `RecordBatchMemoryCounter`-safe accounting) are shared with the rest of the codebase (already used by `DataSourceExec` at `datafusion/datasource/src/source.rs:478`). - Thread the per-row lengths that `predict_output_lens` already computes into `build_batch` -> `list_unnest_at_level` so `find_longest_length` runs once per input batch instead of twice (prediction + re-computation per chunk). For a single-column depth-1 unnest at fanout 1 this drops PR-vs-main overhead from 1.18-1.49x to 1.04-1.08x; the residual is the boundary walk. - Replace the `Vec<usize>` copy with the `PrimitiveArray<Int64Type>` that `find_longest_length` already returns, and drop the `usize::try_from(len) .unwrap_or(0)` guard, which cannot fail on any 64-bit target and would silently zero rows on 32-bit rather than error. - `output_lens` is now `Option<PrimitiveArray>` rather than an empty-vec sentinel. That drops the paragraph of doc comment on the field and also removes the incidental collision with a genuinely empty batch. - `next_chunk_rows` uses a single-index slice loop instead of two counters, and asserts that `PendingInput` is never entered with zero remaining rows (previously enforced by an extra loop stage, now by a `debug_assert!`). - Flatten `poll_next_impl`: no more `return Poll::Ready(match ... continue)` where control flow escapes the operand of a `return`, and the exhausted- PendingInput disposal stage is folded into the drain point. - Consolidate four size-axis tests into one table-driven test that asserts exact output batch shapes; the recursive test uses the shared helper extended with a `depth` parameter. - Trim the slt block: batch sizes are invisible to sqllogictest, so it kept only the ordered `WHERE id IN (...)` query (where chunk boundaries could actually perturb visible output) and the recursive `unnest(unnest(...))` case (distinct code path). Dropped three `count(*)` cases that any correct implementation - including the pre-fix code - passes.
Contributor
|
FYI @duongcongtoai I wonder if you have some time to help review this PR as I think you are familar with the UnnestExec code |
Member
Author
|
This is a change in default behavior. I did think about modifying the API so that we could make the new behavior optional, but that would be a breaking API change, and I'd really like to backport this to 55.1.0 so Comet can get the fix. If that doesn't happen then we can temporarily fork this operator in Comet. |
jayzhan211
approved these changes
Aug 16, 2026
jayzhan211
left a comment
jayzhan211
left a comment
Contributor
There was a problem hiding this comment.
The change makes sense to me
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
4 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
UnnestExecemitted exactly one output batch per input batch, however many rows the unnesting produced, and never consulteddatafusion.execution.batch_size. An 8192-row batch of 100-element lists came back as a single 819,200-row batch.Besides handing downstream operators arbitrarily large batches, this meant peak memory scaled with input batch size times list length rather than with
batch_size, because the full expansion of an input batch was materialized at once.What changes are included in this PR?
UnnestStreamnow consumes each input batch in chunks rather than whole.find_longest_lengthalready computes how many output rows each input row expands into, with all threeNullHandlingmodes accounted for. Prefix-summing it gives exact chunk boundaries for depth-1 unnesting, so eachbuild_batchcall produces at mostbatch_sizerows and the oversized intermediate is never materialized. This is the part that bounds memory, not just output size.Two cases can't be chunked on the input side and are handled by slicing the built batch instead:
batch_size, since one row is never split across output batchesdepth > 1), where a row's expansion depends on inner list lengths that only exist once the outer levels have been unnestedStruct-only unnesting doesn't change the row count, so it's already bounded by the input batch size.
batch_sizeis read fromTaskContextinexecute()and stored on the stream rather than onUnnestExec, which keeps it out of the exhaustive-destructure proto round-trip intry_to_proto/try_from_proto. No serialization changes are needed.Tradeoffs worth reviewer attention
batch_sizeis an upper bound, not an exact size. Chunk boundaries fall on input-row boundaries, so a short tail chunk per input batch is expected. Guaranteeing exact sizes would need a coalescer on top, costing a full copy of the data for little gain.One extra
find_longest_lengthper input batch.build_batchrecomputes it per chunk, so the prediction pass is redundant work — roughly one extra pass over the length arrays, not over the data. Threading the precomputed lengths intobuild_batchwould remove it at the cost of a wider internal signature. Happy to do that if preferred.Are these changes tested?
Yes.
Seven new unit tests in
unnest.rscover: the basic batch_size guarantee, output smaller thanbatch_size, a single row exceedingbatch_size,PreserveAndExpandEmptyandDropnull handling under chunking, multiple input batches, and recursive unnest.One test deliberately pins how the limit is met rather than merely that it's met: 3 rows of 3 elements at
batch_size=4must yield[3, 3, 3](input chunked per row), not[4, 4, 1](built whole, then sliced). Those two strategies are indistinguishable by row counts alone but have very different memory profiles, so without this a future refactor could silently regress to build-then-slice.New
unnest.sltcoverage runs queries atbatch_size = 3and again at the default, asserting identical results — chunk boundaries must not affect output. Note that slt can only compare result sets, so it guards correctness under chunking; the batch-shape guarantees live in the unit tests.Full runs, all green:
cargo test -p datafusion-physical-plan— 1695 passed, 0 failedcargo fmt --checkandcargo clippy -p datafusion-physical-plan --all-targets -- -D warningscleanAre there any user-facing changes?
UnnestExecnow produces more, smaller output batches, bounded bydatafusion.execution.batch_size. Query results are unchanged, including row order. Plans andEXPLAINoutput are unchanged.Anything reading
UnnestExecmetrics will seeoutput_batchesno longer equal toinput_batches.