apache / apache/datafusion-comet

Support fully-native multi-stage (distinct-combined) collect_list / collect_set

Open
#4,724 3 comments 0 reactions 1 assignee Claimed by @peterxcli View on GitHub
area:aggregation area:expressions enhancement priority:medium
Dominant language
Scala
Stars
1.3k
Forks
373
Avg merge
2d 4h
Merged PRs (30d)
198

Description

## Describe the bug

`collect_list` / `collect_set` declare their aggregate buffer as `BinaryType` in Spark (serialized `TypedImperativeAggregate` state) but produce a native `ArrayType` (Arrow `List`) state in Comet. Comet bridges this only for the simple two-stage shape: `CometObjectHashAggregateExec.adjustOutputForNativeState` rewrites the buffer column of a **pure-`Partial`** aggregate to `ArrayType`, so `Partial -> Final` runs natively and correctly.

It does **not** handle multi-stage aggregates that contain a `PartialMerge` stage, which Spark introduces for the distinct-aggregate rewrite, e.g.:

```sql
SELECT x, count(DISTINCT y), collect_list(z) FROM t GROUP BY x
```

This plans as `Partial(x,y) -> PartialMerge(x,y) -> [PartialMerge, Partial](x) -> Final(x)`. The intermediate `PartialMerge` outputs are still declared `BinaryType`, so a fully-native pipeline crashes at runtime:

```
CometNativeException: Cast error: Cannot cast LIST to non-list data type Binary
```

(and a related nullability drift `List(non-null T)` vs `List(nullable T)` for nested element types). This is a facet of the broader Arrow-type-drift tracked in #4515.

This affects both `collect_list` (added in #4720) and the already-shipped `collect_set` — `collect_set` has the same latent crash on this shape.

## Current behavior (workaround in #4720)

To avoid the crash, multi-stage `collect_list`/`collect_set` aggregates now fall back to Spark consistently:
- `CometExecRule.tagUnsafePartialAggregates` tags the feeding pure-`Partial` when a `PartialMerge` stage of a `CollectList`/`CollectSet` aggregate is present.
- `CometBaseAggregate.doConvert` falls back a `PartialMerge`/`Final` stage of these functions when no Comet partial produced the buffer (the cross-engine `LocalTableScan` case).

The simple two-stage `collect_list`/`collect_set` cases continue to run natively.

## Expected behavior

Enable fully-native multi-stage execution by correcting the intermediate buffer schema for `PartialMerge` (and multi-mode) stages of `CollectList`/`CollectSet` (extend `adjustOutputForNativeState` beyond pure-`Partial`, and fix the element-nullability drift), so the `ArrayType` buffer round-trips across all stages. Then remove the fallback guards added in #4720.

## Additional context

- Related: #4515 (DataFusion / datafusion-spark functions whose Arrow return type drifts from Spark catalyst's declared type)
- Introduced alongside: #4720 (native `collect_list` / `array_agg`)

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.