apache / apache/datafusion

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

Open
#25,118 0 comments 0 reactions 0 assignees View on GitHub
enhancement performance spark
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.