andygrove opened a new issue, #25118: URL: https://github.com/apache/datafusion/issues/25118
### Is your feature request related to a problem or challenge? While tracking down a `collect_list` / `collect_set` slowdown in Apache DataFusion Comet ([comet#5797](https://github.com/apache/datafusion-comet/issues/5797)) I found two things in `datafusion-spark` / `datafusion-functions-aggregate` worth fixing upstream. Comet has worked around them locally ([comet#5803](https://github.com/apache/datafusion-comet/pull/5803)) by supplying its own `GroupsAccumulator`s, and would rather go back to the upstream implementations once these are addressed. **1. `SparkCollectList` / `SparkCollectSet` declare no `GroupsAccumulator`.** `datafusion_spark::function::aggregate::collect` implements only `accumulator()`, so grouped aggregation falls to `GroupsAccumulatorAdapter`: one boxed `Accumulator` per group, plus per-batch slicing and dispatch into each. This is the dominant cost at high grouping cardinality. Measured on 2,000,000 rows grouped by a string key with 200,000 distinct values, `collect_list` ran at 0.49x of Spark's own (JVM) implementation while the identical grouping with `count(*)` ran at 2.88x, i.e. the aggregate, not the group-by, was the problem. `ArrayAgg` already has `ArrayAggGroupsAccumulator`, so `SparkCollectList` can largely reuse it. The one semantic difference is that Spark's `collect_list` returns `[]`, not `NULL`, for a group whose inputs were all NULL, whereas `ArrayAggGroupsAccumulator::evaluate` marks such a group's list entry null. `SparkCollectSet` has no upstream grouped equivalent at all, since `groups_accumulator_supported` on `ArrayAgg` excludes the distinct case. **2. `ArrayAggGroupsAccumulator::merge_batch` is weak for low grouping cardinality.** `merge_batch` expands every state list into one `(group_idx, row_idx)` entry per element, and `evaluate` then gathers them with `interleave`. Merging partial states is exactly the case where each contribution is a long contiguous run, so a per-element gather is the wrong shape: `interleave` over N scattered indices is several times more expensive than copying the same N rows as a handful of contiguous slices. Delegating Comet's grouped `collect_list` to `array_agg_udaf().create_groups_accumulator(...)` made a 64-group `collect_list` **15% slower end to end** than the `GroupsAccumulatorAdapter` it replaced (the final stage alone was ~5x slower), while the high-cardinality shapes got 20-100x faster. So the adapter is currently the better choice for low-cardinality merges, which is worth fixing since `merge_batch` is where partial states always arrive. **3. `DistinctArrayAggAccumulator` is per-row in two places** (already noted in the Comet issue, listed here for completeness): `merge_batch` walks the state `ListArray` row by row calling `update_batch(&[val])` on one-element arrays, so every merged row pays a fresh `RowConverter::append` + `create_hashes` + probe setup; and `evaluate` round-trips every distinct element through `ScalarValue::try_from_array` and `ScalarValue::new_list`, which for struct elements is a full recursive `ScalarValue::Struct` materialisation per element. ### Describe the solution you'd like 1. `SparkCollectList`: implement `groups_accumulator_supported` / `create_groups_accumulator`, reusing `ArrayAggGroupsAccumulator` with `ignore_nulls = true` and rewriting the lists it leaves null into empty lists. 2. `SparkCollectSet`: add a grouped distinct accumulator. 3. `ArrayAggGroupsAccumulator`: represent a contribution as a `(group, start, len)` range rather than one entry per row, coalescing consecutive same-group rows in `update_batch` and recording one range per list row in `merge_batch`. On emit, counting-sort the ranges into group order and pick the gather by average run length: `concat` of slices for long runs, `interleave` for scattered rows. 4. `DistinctArrayAggAccumulator`: encode the whole state batch once in `merge_batch` instead of per row, and decode with a single `RowConverter::convert_rows` in `evaluate` instead of per-element `ScalarValue` round trips. ### Describe alternatives you've considered Comet's implementations of (1)-(3) are in [`native/spark-expr/src/agg_funcs/collect.rs`](https://github.com/apache/datafusion-comet/pull/5803) and could be moved upstream more or less as they are. They are Apache-2.0, in this project's style, and carry unit tests plus a criterion benchmark. `CollectSetGroupsAccumulator` there keeps the distinct values row-encoded in one arena, deduplicated on insert against an open-addressed index keyed by `(group, encoded value)`, so a batch is encoded once for all of its groups. Measurements from that PR, criterion, two-stage `AggregateExec` over 131,072 rows, versus the `GroupsAccumulatorAdapter` baseline: | 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` struct, 16k groups | −99.2% | −98.3% | | `collect_set` utf8, 16k groups | −87.6% | −84.8% | | `collect_set` utf8, 64 groups | −59.8% | −60.5% | | `collect_set` struct, 16k groups | −92.9% | −90.7% | ### Additional context Measured against DataFusion 55.0.0 with arrow-rs 59.2.0. Happy to open PRs for any of the four items. -- 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]
