fix: make CometExplodeExec respect batch size

[andygrove]

Which issue does this PR close?

Related to #5210, which tracks the unnest-related workarounds in the explode planner.

Rationale for this change

CometExplodeExec plans onto DataFusion's UnnestExec, which emits exactly one output batch per input batch however many rows the unnesting produces. It never consults datafusion.execution.batch_size — there is no batch_size anywhere in upstream unnest.rs, and UnnestStream calls build_batch once per input batch and returns the whole result.

So an 8192-row batch of 100-element arrays comes back as a single 819,200-row batch, and input batch count always equals output batch count.

Two consequences:

1. Downstream operators receive arbitrarily large batches.
2. Peak memory scales with input batch size times array length rather than with batch_size — the full expansion of an input batch is materialized at once, which is an OOM risk on high-fanout explode.

What changes are included in this PR?

Adds ExplodeExec (native/core/src/execution/operators/explode.rs), a temporary fork of UnnestExec that consumes each input batch in chunks, and points the planner at it instead of UnnestExec.

find_longest_length already computes how many output rows each input row expands into, with null handling accounted for. Prefix-summing it gives exact chunk boundaries for the depth-1 unnesting Comet plans, so each build_batch call produces at most batch_size rows and the oversized intermediate is never materialized. That is what 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 array is longer than batch_size, since one row is never split across output batches
• recursive unnest (depth > 1), which Comet does not plan today, but the fallback keeps the operator correct if that changes

Why a fork

The unnesting kernels (build_batch and everything it calls) are private to datafusion-physical-plan, so they cannot be called from Comet without copying them. They are copied verbatim from DataFusion 54.1.0 (upstream revision cc7565be1ee97ba8fa2f5d6da373c5e38d81bb13) and marked as such, so the eventual deletion is mechanical. All Comet-specific behavior lives in ExplodeExec / ExplodeStream above the copied section.

Note that 54.1.0 predates upstream's NullHandling enum and still uses UnnestOptions::preserve_nulls, which is why the planner still wraps empty arrays with ListEmptyToNullExpr for explode_outer semantics. That is unchanged by this PR.

Deleting this fork

The fix has been submitted upstream:

• UnnestExec ignores datafusion.execution.batch_size datafusion#24383 (issue)
• fix: make UnnestExec respect datafusion.execution.batch_size datafusion#24384 (PR)

Once Comet moves to a DataFusion release carrying that PR, delete this module and go back to UnnestExec. The module docs say so, and #5210 already tracks the related cleanup.

I went with a fork rather than an opt-in upstream API because Comet is pinned to DataFusion 54.1.0, so anything landing in 55 or 56 does nothing here until the upgrade. A fork gets the fix now and deletes cleanly.

Are these changes tested?

Yes — six new unit tests in the new module:

• respects_batch_size — 10 rows x 10 elements at batch_size=8, every output batch bounded, values and order preserved
• chunks_input_rather_than_slicing_output — pins how the limit is met: 3 rows of 3 elements at batch_size=4 must give [3, 3, 3] (input chunked per row), not [4, 4, 1] (built whole, then sliced). These are indistinguishable by row counts alone but have very different memory profiles, so without this a refactor could silently regress to build-then-slice.
• single_row_exceeding_batch_size_is_sliced — the unavoidable case
• chunking_preserves_outer_semantics / chunking_preserves_non_outer_semantics — chunked vs unchunked output compared directly for both preserve_nulls settings
• multiple_input_batches — short tail chunk per input batch boundary

Full native suite: 163 passed, 0 failed. cargo fmt and cargo clippy --all-targets -- -D warnings clean.

JVM-side suite runs still in progress; I'll report results before marking this ready for review.

Are there any user-facing changes?

Explode now produces more, smaller batches, bounded by datafusion.execution.batch_size, and uses substantially less memory on high-fanout arrays. Query results are unchanged.

CometExplodeExec now displays as CometExplodeExec rather than UnnestExec in native plan output.

[sunchao]

Thanks for working on bounded explode batches. I found four issues in stream cleanup, physical-plan properties, input lifetimes, and repeated list-length computation. Details are inline.

[sunchao]

[P1] Release the exhausted child before returning EOF

Could we restore the EOF cleanup from DataFusion's UnnestStream here? The original operator replaces self.input with EmptyRecordBatchStream when the child is exhausted. Without that, a real HashJoin -> Projection -> Explode -> ShuffleWriter plan keeps the partitioned hash join's build table and MemoryReservation alive while shuffle_write() finalizes. ParquetWriterExec similarly retains the input through writer.close(). This can force avoidable spills or out-of-memory failures when those writers need the memory. A child stream with a drop guard or reservation would make a useful regression test.

[andygrove]

Fixed, and thanks for spotting it — this was the fork drifting from the upstream PR rather than a deliberate choice, which is true of all four of these.

ExplodeStream now swaps in EmptyRecordBatchStream when the child is depleted, same as UnnestStream. Took your suggestion on the test too: releases_exhausted_child_at_eof drives a child stream carrying a drop flag and asserts the flag flips on the poll that reports EOF, not when the stream is finally dropped. It fails on the previous code.

Related: the output now goes through DataFusion's public BatchSplitStream instead of the hand-rolled pending_output/split_off_head pair, which is what upstream uses for the same job. That releases its input on EOF as well, so the whole chain unwinds at the same poll.

[sunchao]

[P2] Preserve physical properties for passthrough columns

Could we preserve the projected physical properties from the original UnnestExec::compute_properties? It retains ordering and equivalences for passthrough columns, but starting from empty EquivalenceProperties loses them. For SortMergeJoin -> Explode -> GROUP BY k, a sorted passthrough key k now makes the downstream AggregateExec choose InputOrderMode::Linear instead of Sorted, so groups accumulate or spill instead of streaming. This increases memory consumption and runtime on existing native plans.

[andygrove]

You are right, and my comment justifying the empty properties was wrong on both counts. Unnesting only rewrites the list and struct columns, so whatever the child guarantees about the rest still holds, and hardcoding UnknownPartitioning(1) because Comet happens to plan explode on one partition is not the operator's call to make.

compute_properties now mirrors UnnestExec::compute_properties: build a ProjectionMapping over the non-unnested indices, project the child's equivalences and partitioning through it, and drop only the constraints, since row duplication genuinely does invalidate uniqueness and primary keys. new() is fallible now as a consequence.

preserves_passthrough_orderings pins it — a sorted passthrough key over a DataSourceExec with sort information, asserting the ordering is still there on the ExplodeExec above it. It fails if the properties go back to empty.

[sunchao]

[P2] Drop consumed input before yielding its final output batch

Could we clear pending_input immediately after pending.row_offset consumes its final row, before yielding output? As written, the final batch is returned while the full source RecordBatch, predicted lengths, and any full-batch posexplode positions remain live until the next poll. The upstream implementation checks whether the input is drained and drops pending_input before emitting. Preserving that order avoids holding the entire input during downstream or JVM processing of the final batch.

[andygrove]

Fixed. The loop now computes drained right after advancing row_offset and clears pending_input before the built batch is returned, so the source RecordBatch, the predicted lengths, and the posexplode position array are all released ahead of downstream and JVM processing rather than surviving until the next poll. That also let the remaining_rows() == 0 re-entry branch go, since PendingInput is now only ever live with rows left in it — there is a debug_assert to that effect.

[sunchao]

[P2] Reuse the list lengths already computed for chunking

Could we keep the predicted lengths as an Arrow PrimitiveArray<Int64Type> and pass zero-copy slices into build_batch? predict_output_lens already runs length, cast, is_not_null, and zip over every input row, but converts the result into Vec<usize>. list_unnest_at_level then runs the same kernels again for every chunk. This duplicates work even when no split is needed, and posexplode repeats it across both list columns. The linked upstream fix already reuses the precomputed chunk lengths.

[andygrove]

Done, and it now matches what the upstream PR does. predict_output_lens keeps the result as a PrimitiveArray<Int64Type> rather than collecting into Vec<usize>, PendingInput::chunk_lengths hands out a zero-copy slice of it per chunk, and that slice threads down through build_batch into list_unnest_at_level, which uses it in place of calling find_longest_length again. Reuse is gated on max_recursion == 1, since with recursion the deeper levels depend on arrays that do not exist yet — the same reason the prediction is skipped there in the first place.

This does mean two functions in the vendored region gain a parameter, which cuts against the "do not change it semantically" banner, so I called it out explicitly in the module docs as a deliberate edit that is itself part of apache/datafusion#24384 and therefore disappears with the rest of the fork.