collect_list / collect_set have no GroupsAccumulator, and ArrayAggGroupsAccumulator merges per row
- Dominant language
- Rust
- Stars
- 9.3k
- Forks
- 2.4k
- Avg merge
- 3d 7h
- Merged PRs (30d)
- 344
Description
### 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.
Contributor guide
Research direction
Start with datafusion_spark::function::aggregate::collect and the ArrayAggGroupsAccumulator and DistinctArrayAggAccumulator implementations in datafusion-functions-aggregate. Compare them with native/spark-expr/src/agg_funcs/collect.rs from Comet, including its unit tests and criterion benchmark. Done means grouped collect_list and collect_set support the required NULL and distinct semantics, merge paths avoid per-row work, and tests and benchmarks cover the listed cardinalities.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust
- Domain
- data-engineering, performance
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 42/100