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.
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.
There was a problem hiding this comment.
ignore nulls is not supported by spark for those functions, Comet shouldn't be handling it as this is misleading
| * 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?
There was a problem hiding this comment.
Dropped. You are right, the general missingCometProducer + aggsNotSupportingMixedExecution guard just above already returns None for the same case, so this block was unreachable.
| } | ||
| } | ||
|
|
||
| // 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.
There was a problem hiding this comment.
Added a comment on the block explaining why it is separate: the tagging block above only fires when the Final cannot convert, but canAggregateBeConverted skips the child-native check, so an all-native distinct collect_list Final converts and slips past. This block catches that case.
|
|
||
| /** | ||
| * 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.
There was a problem hiding this comment.
Narrowed the doc comment to describe what the code actually matches (CollectList/CollectSet) and noted that Percentile has the same buffer shape but is intentionally not matched here since it already passes through the general mixed-execution guard.
| * (`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?
There was a problem hiding this comment.
Moved the identical spark-3.4/spark-3.5 copies to spark-3.x. spark-4.x is used by spark-4.0/4.1/4.2 too, and 4.2 needs the different shim, so I could not consolidate 4.0/4.1 there — happy to add a fresh shared dir if you would prefer that, but it seemed like more build machinery than the reduction was worth.
| 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.
There was a problem hiding this comment.
Added collect_list FILTER (WHERE i > 1) per-group and a map-input case (collect_list(m) on a map<string, int> table with two entries and a null row). The map assertion is a spark_answer_only size-only check since sort_array cannot order a MapType.
…m; add FILTER + map tests operators.scala: drop the CollectList/CollectSet PartialMerge fallback block. It is unreachable now that the general missingCometProducer + aggsNotSupportingMixedExecution guard just above returns None for the same case. CometExecRule.scala: add a comment on the collect-specific tagging block explaining it is separate from the tagging block just above because canAggregateBeConverted skips the child-native check, so an all-native distinct collect chain would otherwise slip past. QueryPlanSerde.scala: narrow the hasNativeArrayBufferAgg doc comment to describe what the code actually matches and note that Percentile has the same shape but is not matched here. Consolidate CometCollectShim: move the identical spark-3.4 and spark-3.5 copies to spark-3.x. Add collect_list FILTER (WHERE ...) test and a map-input test hitting make_all_fields_nullable's Map arm (size-only via spark_answer_only, since sort_array cannot order a MapType).
mbutrovich
left a comment
mbutrovich
left a comment
There was a problem hiding this comment.
Approved pending CI. Thanks @andygrove!
| .getOrElse(Compatible()) | ||
| // The native path always drops null inputs. Spark 4.2 adds `RESPECT NULLS` | ||
| // (`ignoreNulls = false`), which keeps nulls, so fall back there. | ||
| if (!CometCollectShim.ignoreNulls(expr)) { |
There was a problem hiding this comment.
this code is confusing
spark.sql("SELECT collect_list(c) IGNORE NULLS FROM VALUES (1),(NULL) AS t(c);").show(false)
org.apache.spark.sql.AnalysisException: [INVALID_SQL_SYNTAX.FUNCTION_WITH_UNSUPPORTED_SYNTAX] Invalid SQL syntax: The function `collect_list` does not support IGNORE NULLS.
it is always dead code
| * added `ignoreNulls` (settable to `false` via `RESPECT NULLS`), handled by the spark-4.2 shim. | ||
| */ | ||
| object CometCollectShim { | ||
| def ignoreNulls(agg: CollectList): Boolean = true |
There was a problem hiding this comment.
we shouldn't be having that
comphead
left a comment
comphead
left a comment
There was a problem hiding this comment.
Thanks @andygrove I think we should leave IGNORE NULLS handling to Spark itself which would fail
|
This guard is for Spark 4.2 added an case collectList: CollectList => collectList.copy(ignoreNulls = ignoreNulls)
case collectSet: CollectSet => collectSet.copy(ignoreNulls = ignoreNulls)so Your I'll add a 4.2-only SQL test exercising For transparency: I used an LLM to help trace the Spark 4.2 resolution path while working through this. |
Add a Spark 4.2-gated SQL file test asserting that collect_list and collect_set with RESPECT NULLS fall back to Spark with Spark-identical results, and expand the getSupportLevel comments to explain that the fallback branch is only reachable on 4.2+ (a no-op on 3.4 through 4.1).
|
Pushed New SQL file test
This runs in the I also expanded the For transparency: the analysis and these tests were done with LLM assistance. |
3 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.