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.
mbutrovich
left a comment
mbutrovich
left a comment
There was a problem hiding this comment.
Thanks @andygrove, the delegation to SparkCollectList and the nullability coercion are clean and the SQL coverage is excellent. My main comments are on the two multi-stage fallback guards. The one in operators.scala looks unreachable now that the general mixed-execution guard has merged, so I think it can go. The one in CometExecRule is the opposite, it's catching a case the general path misses, so it should stay with a comment explaining why. The rest are minor.
| } | ||
| } | ||
|
|
||
| // CollectList/CollectSet declare their buffer as BinaryType in Spark but produce a native |
There was a problem hiding this comment.
Isn't this block unreachable? The missingCometProducer guard at line 1632 already returns None for this case. A PartialMerge stage makes consumesBuffers true, and CollectList/CollectSet are always in aggsNotSupportingMixedExecution since they don't override supportsMixedPartialFinal. Looks like it predates the merge that brought that guard over from main. Can we drop it?
| } | ||
| } | ||
|
|
||
| // CollectList/CollectSet round-trip an ArrayType buffer that Spark declares as BinaryType. |
There was a problem hiding this comment.
Could you add a comment on why this is separate from the tagging block just above? They look like duplicates. The difference is that the block above only tags when the Final can't convert, and canAggregateBeConverted skips the child-native check, so an all-native collect_list Final converts and slips past it. This block is what catches that all-native distinct case, which is the Parquet half of the new suite test.
|
|
||
| /** | ||
| * Returns true if any aggregate function produces a native intermediate buffer whose Arrow type | ||
| * (e.g. ArrayType for CollectList/CollectSet) differs from the BinaryType that Spark declares |
There was a problem hiding this comment.
The comment describes this as matching any agg with a native ArrayType buffer, but it only matches the collect functions. Percentile has the same buffer shape (see adjustOutputForNativeState) and isn't matched. It's fine in practice since percentile with distinct already passes, but the comment reads broader than the code. Worth narrowing the comment, or inlining the match since it's only used in one place now.
| * (`ignoreNulls = false` via `RESPECT NULLS`), so this shim reports the actual value. The serde | ||
| * falls back to Spark when it is `false`, since the native path always drops nulls. | ||
| */ | ||
| object CometCollectShim { |
There was a problem hiding this comment.
These four files (3.4 through 4.1) are identical and only 4.2 differs. Could the 3.4/3.5 pair share the existing spark-3.x dir the way CometTypeShim does, so this is one fewer copy to keep in sync?
| CREATE TABLE cl_src_nulls(val int, grp string) USING parquet | ||
|
|
||
| statement | ||
| INSERT INTO cl_src_nulls VALUES |
There was a problem hiding this comment.
Could we add a collect_list(x) FILTER (WHERE ...) case, and something that exercises a map input? The map path hits the Map arm in make_all_fields_nullable that nothing covers yet. sort_array won't order a map, so it'd need spark_answer_only or a size check.
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.