feat: native collect_list / array_agg aggregate#4720
Conversation
Wires Spark's CollectList aggregate to datafusion-spark's SparkCollectList. array_agg, registered as a SQL alias of CollectList in FunctionRegistry, is also covered. Closes apache#2524.
andygrove
force-pushed
the
feat-collect-list
branch
from
bf6b0c1 to
c9c19e3
Compare
…gates
A distinct aggregate combined with collect_list/collect_set produces a
multi-stage plan (Partial -> PartialMerge -> Final). CollectList/CollectSet
declare a BinaryType buffer in Spark but produce a native ArrayType state, so
Comet cannot read a Spark-produced Binary buffer, nor round-trip its own
ArrayType buffer across the intermediate PartialMerge stages. Both led to
native crashes ("could not cast Binary to List" / "cast List to Binary").
Force these multi-stage aggregates to fall back to Spark consistently:
- tag the feeding Partial when a PartialMerge stage of CollectList/CollectSet
is present (CometExecRule.tagUnsafePartialAggregates), and
- fall back a PartialMerge stage whose buffer was produced by a Spark partial
(CometBaseAggregate.doConvert).
Two-stage collect_list/collect_set continue to run natively.
Patch the upstream SPARK-22223 plan-shape test to disable Comet, since native
collect_list removes the ObjectHashAggregateExec it asserts on.
Enabling fully-native multi-stage execution is tracked in apache#4724.
mbutrovich
left a comment
mbutrovich
left a comment
There was a problem hiding this comment.
First pass. Mostly want to discuss Spark 4.2 or if we should punt.
| inputs: Seq[Attribute], | ||
| binding: Boolean, | ||
| conf: SQLConf): Option[ExprOuterClass.AggExpr] = { | ||
| val child = expr.children.head |
There was a problem hiding this comment.
This does not check CollectList.ignoreNulls, and the upstream SparkCollectList hardcodes ignore_nulls = true (datafusion/spark/src/function/aggregate/collect.rs). On Spark 3.4 through 4.1 CollectList has no ignoreNulls field so this is correct. On Spark 4.2 it gains ignoreNulls: Boolean = true and RESPECT NULLS becomes reachable from SQL. FunctionResolution.applyIgnoreNulls resolves it:
case collectList: CollectList => collectList.copy(ignoreNulls = ignoreNulls)So collect_list(x) RESPECT NULLS sets ignoreNulls = false, which keeps null elements. Comet drops them and returns a different result from Spark with no fallback. 4.2 releases soon, so we might want to tackle this now. The field is absent on 3.4-4.1, so reading it needs a version shim (CometExprShim) rather than a direct expr.ignoreNulls reference. CometCollectSet has the same gap on 4.2, so consider covering both in the shim.
CometFirst (line 249) and CometLast (line 285) already thread ignoreNulls through the proto. First/Last carry the field on every supported Spark version, so no shim is needed there, which is the one extra wrinkle for CollectList/CollectSet. Same wiring shape applies once the field is shimmed.
There was a problem hiding this comment.
Addressed in 9e680c4. Added a CometCollectShim.ignoreNulls shim across the per-version source dirs: 3.4 through 4.1 have no such field so it returns true, and the spark-4.2 shim reads agg.ignoreNulls. getSupportLevel on both CometCollectList and CometCollectSet now returns Unsupported when ignoreNulls = false, so collect_list(x) RESPECT NULLS / collect_set(x) RESPECT NULLS on 4.2 falls back to Spark instead of silently dropping nulls.
| * are therefore only safe to run natively when every stage runs in Comet and there are at most | ||
| * two stages (Partial + Final). | ||
| */ | ||
| def hasIncompatibleBufferAgg(aggExprs: Seq[AggregateExpression]): Boolean = { |
There was a problem hiding this comment.
Nice to factor this out and share it between CometExecRule and operators.scala instead of duplicating. Minor naming thought: "incompatible buffer" is a bit generic. The method really means "produces a native ArrayType state where Spark declares BinaryType, so it cannot round-trip a Spark-produced intermediate buffer." Something like hasNativeArrayBufferAgg would carry more of that at the call sites. The doc comment is clear either way, so take it or leave it.
There was a problem hiding this comment.
Renamed to hasNativeArrayBufferAgg in 9e680c4.
| -- ============================================================ | ||
|
|
||
| query | ||
| SELECT grp, sort_array(collect_list(DISTINCT i)) FROM cl_src_int GROUP BY grp ORDER BY grp |
There was a problem hiding this comment.
For the record, no change needed. A lone collect_list(DISTINCT i) does not create a PartialMerge stage on collect_list itself. Spark does the dedup with group-by-only stages and the collect still runs plain Partial then Final, so this query stays fully native and does not touch the hasIncompatibleBufferAgg fallback. The fallback path needs a distinct aggregate combined with the collect, which the new CometAggregateSuite test covers over both LocalTableScan and Parquet sources. Coverage across the two files is good.
There was a problem hiding this comment.
Agreed, no change needed. The lone collect_list(DISTINCT i) path stays fully native, and the aggregate SQL-file tests plus the CometAggregateSuite distinct-plus-collect fallback test still pass after merging the latest apache/main (which reworked the mixed-execution guard).
# Conflicts: # dev/diffs/3.4.3.diff # docs/source/contributor-guide/expression-audits/agg_funcs.md # native/proto/src/proto/expr.proto # spark/src/main/scala/org/apache/comet/serde/aggregates.scala # spark/src/main/scala/org/apache/spark/sql/comet/operators.scala
Spark 4.2 adds an ignoreNulls field to CollectList and CollectSet, and collect_list(x) RESPECT NULLS sets it to false, keeping null elements. The native path delegates to SparkCollectList/SparkCollectSet, which always drop nulls, so it would silently return a different result from Spark. Add a per-version CometCollectShim that reads ignoreNulls (always true on Spark 3.4 through 4.1, where the field is absent) and fall back to Spark in getSupportLevel when it is false. Also rename QueryPlanSerde.hasIncompatibleBufferAgg to hasNativeArrayBufferAgg to describe what it detects: an aggregate whose native ArrayType state cannot round-trip Spark's declared BinaryType buffer.
# 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/user-guide/latest/expressions.md # native/proto/src/proto/expr.proto # spark/src/main/scala/org/apache/comet/serde/aggregates.scala # spark/src/main/scala/org/apache/spark/sql/comet/operators.scala
The SPARK-22223 ObjectHashAggregate test asserts on the executed plan. With Comet enabled, collect_list runs natively as CometHashAggregateExec, so ObjectHashAggregateExec is no longer present. Rather than disabling Comet, update the operator assertion to also accept CometHashAggregateExec, matching the pattern already used elsewhere in the diffs. The exchange assertion already matches ShuffleExchangeLike, which CometShuffleExchangeExec implements, so the single-shuffle check still holds.
Adding CollectList to the comment pushed a line past the column limit during the apache/main merge; reflow to satisfy spotless.
collect_list and collect_set build their result list with all element fields marked nullable, but SparkCollectList/SparkCollectSet derive the return type from the child, preserving non-nullable nested fields. When the child is a nested type with a non-nullable inner field (e.g. a struct field built from non-nullable columns), the declared aggregate output disagrees with the array the accumulator produces, and the grouped native AggregateExec fails validating its output batch with "column types must match schema types". Cast the collect child to the all-nullable variant of its type so the declared and produced types stay consistent.
2 participants
Which issue does this PR close?
Closes #2524.
Rationale for this change
collect_listand its aliasarray_aggare common aggregate functions that previously fell back to Spark, breaking native execution for many real workloads (notably plans that group rows into arrays before further processing). Adding native support keeps these queries on the Comet path.What changes are included in this PR?
This change was scaffolded with the
implement-comet-expressionskill.CollectListmessage toexpr.protoand a newcollectList = 18arm in theAggExproneof.CometCollectListserde inaggregates.scalaand registerclassOf[CollectList] -> CometCollectListinQueryPlanSerde.datafusion_spark::function::aggregate::collect::SparkCollectList, the upstream Spark-compatible accumulator (Spark 3.4 through 4.1 useignore_nulls = truesemantics that match it). No Comet-local Rust function is added.adjustOutputForNativeStatefix inoperators.scalato coverCollectList(same pattern already used forCollectSet: native producesArrayType(elementType)while Spark declares the buffer asBinaryType).collect_listandarray_aggas supported in the user-guide expression page.docs/source/contributor-guide/expression-audits/agg_funcs.mdcovering Spark 3.4.3, 3.5.8, 4.0.1, 4.1.1.How are these changes tested?
spark/src/test/resources/sql-tests/expressions/aggregate/collect_list.sqlexercises ~30 queries across types (boolean, byte/short/int/bigint, float/double including NaN/Inf/-0, string, binary, decimal up to (38,0), date, timestamp, struct, nested array), GROUP BY, NULLs, all-NULL groups, empty tables, single-row, mixed aggregates, multiplecollect_listcolumns, DISTINCT, HAVING, thearray_aggalias, INT/BIGINT boundary values, and an SPARK-17641 null-filter regression. The fixture runs underConfigMatrix: parquet.enable.dictionary=false,true, so each query executes twice.CometSqlFileTestSuitetests pass (./mvnw test -Dsuites="org.apache.comet.CometSqlFileTestSuite expressions/aggregate" -Dtest=none -Pspark-3.5), confirming no regression tocollect_setor other aggregates.cargo clippy --all-targets --workspace -- -D warningsis clean.