collect_list / collect_set have no GroupsAccumulator, and ArrayAggGroupsAccumulator merges per row

[andygrove]

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) I found two things in datafusion-spark / datafusion-functions-aggregate worth fixing upstream. Comet has worked around them locally (comet#5803) by supplying its own GroupsAccumulators, 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 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.