andygrove opened a new pull request, #5803: URL: https://github.com/apache/datafusion-comet/pull/5803
## Which issue does this PR close? Closes #5797. ## Rationale for this change Neither `SparkCollectList` nor `SparkCollectSet` declares `groups_accumulator_supported`, so grouped `collect_list` / `collect_set` run through DataFusion's `GroupsAccumulatorAdapter`. That adapter keeps one boxed `Accumulator` per group and slices every input batch into a per-group `update_batch` call, so its cost scales with the number of groups rather than with the data. On the high-cardinality grouping in #5797 that overhead is most of the runtime: the same grouping with `count(*)` is 2.88x faster than Spark, and adding `collect_list` on top of it drops to 0.49x. `collect_set` paid a second cost on top: one `RowConverter` plus one hash table and one owned encoded `Row` per distinct value **per group**, and a `merge_batch` that walked the state list row by row calling `update_batch` on one-element arrays. ## What changes are included in this PR? `CometCollectList` / `CometCollectSet` in `native/spark-expr/src/agg_funcs/collect.rs`. Both delegate `return_type`, `state_fields`, `accumulator` and `default_value` to `datafusion_spark`'s implementations, so the ungrouped path (a global aggregate, a window frame) is unchanged, and both add a `GroupsAccumulator`: - `CollectListGroupsAccumulator` retains the input arrays rather than copying per group, and records a `(group, row range)` contribution per batch (consecutive rows of the same group coalesce into one range). `evaluate` counting-sorts those contributions into group order and gathers the values in a single pass: `concat` when the runs are long (merging partial states), `interleave` when the rows are scattered (a high-cardinality `update_batch`). A group that collected nothing gets `[]`, matching Spark, rather than a null list. - `CollectSetGroupsAccumulator` encodes each batch once for all of its groups and keeps the distinct values row-encoded in a single arena, deduplicated on insert against an open-addressed index keyed by `(group, encoded value)`. `evaluate` decodes the surviving rows back in one `convert_rows` call. Deduplication still compares `arrow::row` encoded bytes, exactly as `DistinctArrayAggAccumulator` did, so which values collapse together is unchanged. What does change is that `collect_set` now emits each group's values in insertion order instead of hash-table order; Spark does not define the order either way. The planner change is the two `AggregateUDF::new_from_impl` call sites. `coerce_collect_child_nullability` stays for the ungrouped path; the grouped accumulators normalize the arrays they take in themselves, which `test_collect_agg_absorbs_nested_nullability_drift` now also pins. Upstream `datafusion-spark` has the same gap, and DataFusion's own `ArrayAggGroupsAccumulator` has the `merge_batch` weakness the range/`concat` path here avoids. I will open an issue on `apache/datafusion` with these findings so the work can move back upstream later. ## How are these changes tested? New criterion benchmark `native/spark-expr/benches/collect.rs` runs both functions over a two-stage `AggregateExec` for six shapes (int64 / utf8 / struct elements, high and low grouping cardinality, heavily duplicated elements, and 33% nulls), in partial-only and partial+final form. All 24 shapes improve, none regress: | shape | partial | partial + final | |---|---|---| | `collect_list` int64, 16k groups | −97.9% | −97.0% | | `collect_list` utf8, 16k groups | −97.6% | −96.7% | | `collect_list` utf8, 64 groups | −18.0% | −22.9% | | `collect_list` utf8, duplicate elements | −97.7% | −96.8% | | `collect_list` utf8, 33% nulls | −97.4% | −96.2% | | `collect_list` struct, 16k groups | −99.2% | −98.3% | | `collect_set` int64, 16k groups | −89.3% | −86.9% | | `collect_set` utf8, 16k groups | −87.6% | −84.8% | | `collect_set` utf8, 64 groups | −59.8% | −60.5% | | `collect_set` utf8, duplicate elements | −91.4% | −92.7% | | `collect_set` utf8, 33% nulls | −92.5% | −90.2% | | `collect_set` struct, 16k groups | −92.9% | −90.7% | A first cut delegated `collect_list`'s groups accumulator to DataFusion's `array_agg` one; that regressed the low-cardinality merge by 15% because it expands each state list into per-element `interleave` indices where the old accumulator concatenated whole slices. The range representation is what removes that regression. Correctness: - 15 Rust unit tests in `collect.rs` covering per-group input order, dropped nulls, empty lists for groups that collected nothing, filters, merging partial states, `EmitTo::First` compaction and renumbering, `convert_to_state`, both gather paths agreeing, index growth past its initial capacity, and dedup of ints, strings and structs. - `test_collect_agg_absorbs_nested_nullability_drift` extended to assert the grouped state type matches the declared one with and without the argument coercion. - New `CometAggregateSuite` test "grouped collect_list/collect_set over nulls, duplicates and several batches", which forces a small batch size so a group is fed by several batches and several partial states, and includes a group whose every input is NULL. - `CometAggregateSuite` (91), `CometExecSuite` (144), `CometFuzzAggregateSuite` (27) and the native test suites all pass. Scaffolded with the repo's `optimize-comet-expression` skill (baseline-before-change benchmarking and its no-regression gate). -- 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]
