andygrove opened a new pull request, #24384: URL: https://github.com/apache/datafusion/pull/24384
## Which issue does this PR close? Closes #24383. ## Rationale for this change `UnnestExec` emitted exactly one output batch per input batch, however many rows the unnesting produced, and never consulted `datafusion.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? `UnnestStream` now consumes each input batch in chunks rather than whole. `find_longest_length` already computes how many output rows each input row expands into, with all three `NullHandling` modes accounted for. Prefix-summing it gives *exact* chunk boundaries for depth-1 unnesting, so each `build_batch` call produces at most `batch_size` rows 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: - 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 have been unnested Struct-only unnesting doesn't change the row count, so it's already bounded by the input batch size. `batch_size` is read from `TaskContext` in `execute()` and stored on the stream rather than on `UnnestExec`, which keeps it out of the exhaustive-destructure proto round-trip in `try_to_proto`/`try_from_proto`. No serialization changes are needed. ### Tradeoffs worth reviewer attention **`batch_size` is 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_length` per input batch.** `build_batch` recomputes 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 into `build_batch` would 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.rs` cover: the basic batch_size guarantee, output smaller than `batch_size`, a single row exceeding `batch_size`, `PreserveAndExpandEmpty` and `Drop` null 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=4` must 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.slt` coverage runs queries at `batch_size = 3` and 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 failed - full sqllogictest suite — 502 files, 0 mismatches - `cargo fmt --check` and `cargo clippy -p datafusion-physical-plan --all-targets -- -D warnings` clean ## Are there any user-facing changes? `UnnestExec` now produces more, smaller output batches, bounded by `datafusion.execution.batch_size`. Query results are unchanged, including row order. Plans and `EXPLAIN` output are unchanged. Anything reading `UnnestExec` metrics will see `output_batches` no longer equal to `input_batches`. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
