andygrove opened a new pull request, #5667: URL: https://github.com/apache/datafusion-comet/pull/5667
## Which issue does this PR close? N/A — no existing issue. Happy to file one if reviewers would rather track it that way. > [!NOTE] > **Stacked on #5381.** The first commit here is that PR's benchmark, which is what > measures the end-to-end effect of the rest. Review the last two commits; this needs a > rebase once #5381 lands. The diff against `main` will shrink to the native changes then. ## Rationale for this change `ExplodeExec` is a fork of DataFusion's `UnnestExec`, and it inherited kernels written to be generic over every list type DataFusion supports. Comet plans exactly one shape: a `List`, depth 1, one column for `explode` and two for `posexplode`. Three costs in the general path are avoidable at that shape. The largest is in `unnest_list_array`. Unnesting a single list column pads nothing, so the take indices it builds — one `i64` per output element — are the contiguous run `offsets.first()..offsets.last()`, and the gather through them reads the child straight through in order. Rows of a `ListArray` are adjacent by construction, since row `i` is `[offsets[i], offsets[i + 1])`, so the answer is already sitting in the child as one slice. The operator was allocating an index buffer the size of its output and copying every element to reproduce it. Two smaller ones. `predict_output_lens` derives per-row lengths through `find_longest_length`, which chains `length`, `cast`, `is_not_null` and `zip` to stay generic — four allocating passes to subtract adjacent offsets, once per input batch, because `length` returns `Int32` for `List` and NULL rows need substituting. And `create_take_indices` appended the repeat indices one element at a time through a builder that has no validity to track. ## What changes are included in this PR? `unnest_list_array` returns a slice of the child when the gather would be a contiguous run. Comet takes this for plain `explode`, and for both columns of `posexplode`, whose position array is built with the same per-row lengths and null mask. `explode_outer` still gathers: a NULL or empty row is padded with a NULL that no slice of the child contains. The guard is not just arithmetic. Checking that the offset span equals the capacity is not sufficient, because Arrow permits a NULL list slot to span elements — the gather skips those and a slice would include them — and such a row can cancel out padding elsewhere and leave the totals agreeing with a run that is not the one to take. `populated_null_row_falls_back_even_when_the_totals_agree` pins that case. The view types are rejected outright, since their per-row offsets are independent and need not be ordered. Worth noting for the guard's reach: it accepts NULL rows whose range is empty, which is what builders and the Parquet reader emit. That matters because `InferFiltersFromGenerate` puts `size(arr) > 0 AND arr IS NOT NULL` below every non-outer generator over a column, so a plain `explode` in production is usually handed an array column with the NULL rows already filtered out — but when it is not, the rows still contribute nothing and the run still holds. `predict_output_lens` computes lengths in one pass for the single-`List` case and keeps `find_longest_length` as the fallback. `create_take_indices` fills the buffer per run. The module header said a change to the forked region "either belongs upstream or does not belong at all". This PR deliberately reverses that: the specializations are for the shape Comet plans, not upstream's, and the header now says so and says what retiring the fork would cost. ### Trade-off Output batches now alias the input's child buffer rather than owning a compacted copy, and `ListArray::slice` leaves `values` whole, so the alias is to the child of the whole input batch, not of the chunk. All chunks from one input batch share that buffer and between them fill it, so there is no amplification when they are all retained; a downstream operator that keeps only some of them pins all of it. That is bounded by one input batch's expansion, which `pending_input` already holds materialized, and `BatchSplitStream` above it already hands out slices of the input. ## How are these changes tested? Eight new unit tests in `explode.rs` cover the fast path — child aliasing, a non-zero offset base from slicing, empty and dropped-NULL rows, the padded fallback, the populated-NULL-row case above, and view input — plus equivalence of the fused length computation against `find_longest_length` across NULL handling, empty batches, and slicing. The existing chunking tests already assert that chunked output matches unchunked output exactly, which is what would catch a fast path that fired when it should not have. `CometGenerateExecSuite` (37 tests) and `CometExecSuite` (144 tests) pass. ### Measurements `cargo bench --bench explode`, added in the second commit, on an Apple M3 Max: | Case | Change | | ---- | ------ | | `explode_fan_out/2` | −74% | | `explode_fan_out/10` | −82% | | `explode_fan_out/100` | −90% | | `explode_element_type/bigint` | −82% | | `explode_element_type/string` | −91% | | `explode_element_type/struct` | −91% | | `explode_carried_columns/0` | −82% | | `explode_carried_columns/3` | −60% | | `explode_outer_with_nulls/bigint` | −13% | | `explode_outer_with_nulls/string` | −8% | The outer cases keep gathering, so they get only the length and index changes. The carried-column case improves least of the sliceable ones because replicating the passthrough columns is a real gather that this does not touch. `CometExplodeBenchmark` from #5381, same machine, `local[1]`, Spark 4.1 / Scala 2.13. These are whole-query times including the Parquet scan and the counting aggregate, so they are the diluted view (Best ms): | Case | Spark before / after | Comet before / after | Relative before / after | | ---- | -------------------- | -------------------- | ----------------------- | | fan-out 2 | 42 / 41 | 25 / 23 | 1.7X / 1.8X | | fan-out 10 | 59 / 58 | 34 / 29 | 1.7X / 2.0X | | fan-out 100 | 239 / 231 | 180 / 136 | 1.3X / 1.7X | | posexplode | 55 / 54 | 40 / 34 | 1.4X / 1.6X | | explode_outer | 55 / 55 | 40 / 36 | 1.4X / 1.5X | | posexplode_outer | 56 / 57 | 48 / 45 | 1.2X / 1.3X | | element bigint | 54 / 55 | 34 / 28 | 1.6X / 2.0X | | element string | 81 / 82 | 53 / 43 | 1.5X / 1.9X | | element struct | 91 / 92 | 77 / 67 | 1.2X / 1.4X | | explode alone | 71 / 72 | 44 / 38 | 1.6X / 1.9X | | plus 3 carried columns | 75 / 73 | 53 / 49 | 1.4X / 1.5X | The Spark arm is unchanged across every case, which is the control: the two runs are comparable and the movement is Comet's. -- 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]
