Skip to content

feat: add native distinct-combined collect_list aggregate support - #4727

Open
peterxcli wants to merge 21 commits into
apache:mainfrom
peterxcli:multi-stage-distinct-combined-collect-list-collect-set
Open

feat: add native distinct-combined collect_list aggregate support#4727
peterxcli wants to merge 21 commits into
apache:mainfrom
peterxcli:multi-stage-distinct-combined-collect-list-collect-set

Conversation

@peterxcli

@peterxcli peterxcli commented Jun 25, 2026

Copy link
Copy Markdown
Member

Which issue does this PR close?

Closes #4724.
Closes #4766.

Rationale for this change

collect_list / array_agg and collect_set use Spark TypedImperativeAggregate buffers that Spark declares as serialized BinaryType, while Comet’s native implementations keep the real aggregate state as an Arrow/Spark ArrayType.

The existing schema adjustment only handled simple two-stage Partial -> Final collect aggregates. Spark’s distinct-aggregate rewrite can introduce multi-stage plans with PartialMerge stages, for example:

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

Without correcting the intermediate buffer schema for these stages, a fully-native pipeline can fail when native list state is treated as Spark binary state. This change makes the native array state round-trip through Partial, PartialMerge, and mixed {Partial, PartialMerge} stages so collect_list / collect_set can run fully native in distinct-combined aggregate plans.

What changes are included in this PR?

  • Adds native collect_list / array_agg aggregate support:

    • Adds CollectList to the aggregate expression proto.
    • Adds JVM serde for Spark CollectList.
    • Registers CollectList -> CometCollectList.
    • Wires native planning to datafusion_spark::function::aggregate::collect::SparkCollectList.
  • Extends collect aggregate native-state schema adjustment:

    • Updates CometObjectHashAggregateExec.adjustOutputForNativeState to handle both CollectList and CollectSet.
    • Applies the rewrite to intermediate Partial, PartialMerge, and mixed {Partial, PartialMerge} stages, not only pure Partial.
    • Rewrites Spark’s serialized BinaryType buffer attributes to Comet’s native ArrayType(elementType, containsNull = true) state type.
  • Adds regression coverage for fully-native distinct-combined collect aggregates.

  • Updates expression support docs and aggregate audit notes for collect_list / array_agg.

How are these changes tested?

Ran:

cargo fmt --manifest-path native/Cargo.toml --all
cargo check --manifest-path native/Cargo.toml -p datafusion-comet
make core
./mvnw test -Dtest=none -Dsuites="org.apache.comet.exec.CometAggregateSuite collect_list/collect_set combined with distinct aggregate runs fully native" -Dscalastyle.skip=true -Pspark-3.5
./mvnw spotless:check -Dscalastyle.skip=true -Pspark-3.5

explain the sql with:

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

was:

ObjectHashAggregate(keys=[x#16], functions=[collect_list(z#18, 0, 0), count(distinct y#17)])
+- Exchange hashpartitioning(x#16, 2), ENSURE_REQUIREMENTS, [plan_id=84]
   +- ObjectHashAggregate(keys=[x#16], functions=[merge_collect_list(z#18, 0, 0), partial_count(distinct y#17)])
      +- ObjectHashAggregate(keys=[x#16, y#17], functions=[merge_collect_list(z#18, 0, 0)])
         +- Exchange hashpartitioning(x#16, y#17, 2), ENSURE_REQUIREMENTS, [plan_id=81]
            +- ObjectHashAggregate(keys=[x#16, y#17], functions=[partial_collect_list(z#18, 0, 0)])
               +- CometNativeColumnarToRow
                  +- CometNativeScan parquet [x#16,y#17,z#18] Batched: true, DataFilters: [], Format: Parquet, Location: InMemoryFileIndex(1 paths)[file:/var/folders/6h/_n7pxl3n7knbwg5lwlrkzr840000gn/T/comet_distinct_c..., PartitionFilters: [], PushedFilters: [], ReadSchema: struct<x:int,y:int,z:string>

after:

== Physical Plan ==
CometNativeColumnarToRow
+- CometHashAggregate [x#16, buf#24, count#27L], [Final], [x#16], [collect_list(z#18, 0, 0), count(distinct y#17)]
   +- CometExchange hashpartitioning(x#16, 2), ENSURE_REQUIREMENTS, CometNativeShuffle, [plan_id=69]
      +- CometHashAggregate [x#16, y#17, buf#24], [PartialMerge, Partial], [x#16], [merge_collect_list(z#18, 0, 0), partial_count(distinct y#17)]
         +- CometHashAggregate [x#16, y#17, buf#24], [PartialMerge], [x#16, y#17], [merge_collect_list(z#18, 0, 0)]
            +- CometExchange hashpartitioning(x#16, y#17, 2), ENSURE_REQUIREMENTS, CometNativeShuffle, [plan_id=66]
               +- CometHashAggregate [x#16, y#17, z#18], [Partial], [x#16, y#17], [partial_collect_list(z#18, 0, 0)]
                  +- CometNativeScan parquet [x#16,y#17,z#18] Batched: true, DataFilters: [], Format: Parquet, Location: InMemoryFileIndex(1 paths)[file:/var/folders/6h/_n7pxl3n7knbwg5lwlrkzr840000gn/T/comet_distinct_c..., PartitionFilters: [], PushedFilters: [], ReadSchema: struct<x:int,y:int,z:string>

* binary). However, the native Comet aggregate produces the actual state type (e.g.,
* ArrayType(elementType) for CollectSet). This method corrects the output schema to match the
* native state types so the shuffle exchange schema is consistent with the actual data.
* For intermediate aggregates containing TypedImperativeAggregate functions (like CollectSet or

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@peterxcli
peterxcli marked this pull request as draft June 25, 2026 17:54
@peterxcli
peterxcli marked this pull request as ready for review June 27, 2026 11:03
@peterxcli peterxcli changed the title feat: Add native collect_list aggregate support feat: Add native distinct-combined collect_list aggregate support Jun 27, 2026
@peterxcli peterxcli changed the title feat: Add native distinct-combined collect_list aggregate support feat: add native distinct-combined collect_list aggregate support Jun 27, 2026
@peterxcli
peterxcli marked this pull request as draft June 28, 2026 18:38
@peterxcli

Copy link
Copy Markdown
Member Author

Noticed that I didn't solve the root cause, we need a native UnsafeRow / UnsafeArrayData bytes decoder. addressing...

@peterxcli
peterxcli marked this pull request as ready for review July 1, 2026 09:31
@peterxcli
peterxcli requested a review from hsiang-c July 1, 2026 12:43
@peterxcli
peterxcli force-pushed the multi-stage-distinct-combined-collect-list-collect-set branch from fead19b to 4e4d29d Compare July 8, 2026 17:18
@peterxcli

Copy link
Copy Markdown
Member Author

@andygrove would you mind taking a look when you have a chance? This is related to #4720. Thanks!

…t-combined-collect-list-collect-set

# Conflicts:
#	dev/diffs/3.4.3.diff
#	dev/diffs/3.5.8.diff
#	dev/diffs/4.0.2.diff
#	dev/diffs/4.1.2.diff
#	docs/source/contributor-guide/expression-audits/agg_funcs.md
#	native/core/src/execution/planner.rs
#	native/proto/src/proto/expr.proto
#	spark/src/main/scala/org/apache/comet/serde/aggregates.scala
#	spark/src/test/resources/sql-tests/expressions/aggregate/collect_list.sql
#	spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

2 participants