From c9c19e3bc55b4629c4b7f237d562d194699c5b80 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Wed, 24 Jun 2026 10:26:03 -0600 Subject: [PATCH 1/6] feat: native collect_list / array_agg aggregate 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 #2524. --- .../expression-audits/agg_funcs.md | 8 + docs/source/user-guide/latest/expressions.md | 4 +- native/core/src/execution/planner.rs | 7 +- native/proto/src/proto/expr.proto | 6 + .../apache/comet/serde/QueryPlanSerde.scala | 1 + .../org/apache/comet/serde/aggregates.scala | 34 +- .../apache/spark/sql/comet/operators.scala | 21 +- .../expressions/aggregate/collect_list.sql | 405 ++++++++++++++++++ 8 files changed, 472 insertions(+), 14 deletions(-) create mode 100644 spark/src/test/resources/sql-tests/expressions/aggregate/collect_list.sql diff --git a/docs/source/contributor-guide/expression-audits/agg_funcs.md b/docs/source/contributor-guide/expression-audits/agg_funcs.md index 8025f31b73..f3d83262b4 100644 --- a/docs/source/contributor-guide/expression-audits/agg_funcs.md +++ b/docs/source/contributor-guide/expression-audits/agg_funcs.md @@ -39,4 +39,12 @@ - Spark 3.5.8 (2026-05-26) - Spark 4.0.1 (2026-05-26) +## collect_list + +- Spark 3.4.3 (audited 2026-06-24): `CollectList` extends `Collect[ArrayBuffer[Any]]`, returns `ArrayType(child.dataType, containsNull = false)`, ignores NULL inputs in `update()` (Hive-compatible semantics), and yields an empty array as `defaultResult`. `nullable = false`. No `checkInputDataTypes` override, so any input type is accepted (including STRUCT, ARRAY, MAP). Registered as both `collect_list` and `array_agg` aliases in `FunctionRegistry`. +- Spark 3.5.8 (audited 2026-06-24): identical to 3.4.3. +- Spark 4.0.1 (audited 2026-06-24): only structural change is adding `with UnaryLike[Expression]` to the case class (no behavior change). +- Spark 4.1.1 (audited 2026-06-24): identical to 4.0.1. +- Comet implementation: native side delegates to `datafusion_spark::function::aggregate::collect::SparkCollectList`, which wraps `ArrayAggAccumulator` with `ignore_nulls = true` and converts a final NULL accumulator state to an empty array (matching Spark's `defaultResult`). The native return type is `List(Field, containsNull = true)`, while Spark uses `containsNull = false`. Because nulls are filtered before insertion, no nulls actually appear in the array, so this is a schema-shape difference only and tests using `checkSparkAnswerAndOperator` accept it (same pattern already in use for `collect_set`). + [Spark Expression Support]: ../../user-guide/latest/expressions.md diff --git a/docs/source/user-guide/latest/expressions.md b/docs/source/user-guide/latest/expressions.md index b854e29d1f..c74d87a6d4 100644 --- a/docs/source/user-guide/latest/expressions.md +++ b/docs/source/user-guide/latest/expressions.md @@ -71,14 +71,14 @@ The tables below list every Spark built-in expression with its current status. | `any` | ✅ | | | `any_value` | ✅ | | | `approx_count_distinct` | 🔜 | tracking [#4098](https://github.com/apache/datafusion-comet/issues/4098) | -| `array_agg` | 🔜 | Array aggregate (related to `collect_list`, [#2524](https://github.com/apache/datafusion-comet/issues/2524)) | +| `array_agg` | ✅ | Alias for `collect_list` | | `avg` | ✅ | Interval types fall back | | `bit_and` | ✅ | | | `bit_or` | ✅ | | | `bit_xor` | ✅ | | | `bool_and` | ✅ | | | `bool_or` | ✅ | | -| `collect_list` | 🔜 | [#2524](https://github.com/apache/datafusion-comet/issues/2524) | +| `collect_list` | ✅ | | | `collect_set` | ✅ | | | `corr` | ✅ | | | `count` | ✅ | | diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index e89f0a8cf4..97800e91c4 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -75,7 +75,7 @@ use datafusion_comet_spark_expr::{ BloomFilterAgg, BloomFilterMightContain, CsvWriteOptions, EvalMode, SparkArraysZipFunc, SparkBloomFilterVersion, SumInteger, ToCsv, }; -use datafusion_spark::function::aggregate::collect::SparkCollectSet; +use datafusion_spark::function::aggregate::collect::{SparkCollectList, SparkCollectSet}; use iceberg::expr::Bind; use crate::execution::operators::ExecutionError::GeneralError; @@ -2558,6 +2558,11 @@ impl PhysicalPlanner { let func = AggregateUDF::new_from_impl(SparkCollectSet::new()); Self::create_aggr_func_expr("collect_set", schema, vec![child], func) } + AggExprStruct::CollectList(expr) => { + let child = self.create_expr(expr.child.as_ref().unwrap(), Arc::clone(&schema))?; + let func = AggregateUDF::new_from_impl(SparkCollectList::new()); + Self::create_aggr_func_expr("collect_list", schema, vec![child], func) + } } } diff --git a/native/proto/src/proto/expr.proto b/native/proto/src/proto/expr.proto index 90e3d87032..e73c07a914 100644 --- a/native/proto/src/proto/expr.proto +++ b/native/proto/src/proto/expr.proto @@ -143,6 +143,7 @@ message AggExpr { Correlation correlation = 15; BloomFilterAgg bloomFilterAgg = 16; CollectSet collectSet = 17; + CollectList collectList = 18; } // Optional filter expression for SQL FILTER (WHERE ...) clause. @@ -267,6 +268,11 @@ message CollectSet { DataType datatype = 2; } +message CollectList { + Expr child = 1; + DataType datatype = 2; +} + enum EvalMode { LEGACY = 0; TRY = 1; diff --git a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala index 143048fb44..8935f7cf22 100644 --- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala +++ b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala @@ -380,6 +380,7 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim { classOf[BitOrAgg] -> CometBitOrAgg, classOf[BitXorAgg] -> CometBitXOrAgg, classOf[BloomFilterAggregate] -> CometBloomFilterAggregate, + classOf[CollectList] -> CometCollectList, classOf[CollectSet] -> CometCollectSet, classOf[Corr] -> CometCorr, classOf[Count] -> CometCount, diff --git a/spark/src/main/scala/org/apache/comet/serde/aggregates.scala b/spark/src/main/scala/org/apache/comet/serde/aggregates.scala index 6221452c86..8696b39ea1 100644 --- a/spark/src/main/scala/org/apache/comet/serde/aggregates.scala +++ b/spark/src/main/scala/org/apache/comet/serde/aggregates.scala @@ -22,7 +22,7 @@ package org.apache.comet.serde import scala.jdk.CollectionConverters._ import org.apache.spark.sql.catalyst.expressions.{Attribute, Literal} -import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, Average, BitAndAgg, BitOrAgg, BitXorAgg, BloomFilterAggregate, CentralMomentAgg, CollectSet, Corr, Count, Covariance, CovPopulation, CovSample, First, Last, Max, Min, StddevPop, StddevSamp, Sum, VariancePop, VarianceSamp} +import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, Average, BitAndAgg, BitOrAgg, BitXorAgg, BloomFilterAggregate, CentralMomentAgg, CollectList, CollectSet, Corr, Count, Covariance, CovPopulation, CovSample, First, Last, Max, Min, StddevPop, StddevSamp, Sum, VariancePop, VarianceSamp} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{ByteType, DecimalType, IntegerType, LongType, ShortType, StringType} @@ -740,6 +740,38 @@ object CometCollectSet extends CometAggregateExpressionSerde[CollectSet] { } } +object CometCollectList extends CometAggregateExpressionSerde[CollectList] { + + override def convert( + aggExpr: AggregateExpression, + expr: CollectList, + inputs: Seq[Attribute], + binding: Boolean, + conf: SQLConf): Option[ExprOuterClass.AggExpr] = { + val child = expr.children.head + val childExpr = exprToProto(child, inputs, binding) + val dataType = serializeDataType(expr.dataType) + + if (childExpr.isDefined && dataType.isDefined) { + val builder = ExprOuterClass.CollectList.newBuilder() + builder.setChild(childExpr.get) + builder.setDatatype(dataType.get) + + Some( + ExprOuterClass.AggExpr + .newBuilder() + .setCollectList(builder) + .build()) + } else if (dataType.isEmpty) { + withFallbackReason(aggExpr, s"datatype ${expr.dataType} is not supported", child) + None + } else { + withFallbackReason(aggExpr, child) + None + } + } +} + object AggSerde { import org.apache.spark.sql.types._ diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala index af9e1df8a3..19a0841650 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala @@ -30,7 +30,7 @@ import org.apache.spark.broadcast.Broadcast import org.apache.spark.rdd.RDD import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions.{Ascending, Attribute, AttributeSet, Expression, ExpressionSet, Generator, NamedExpression, SortOrder} -import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, AggregateMode, CollectSet, Final, First, Last, Partial, PartialMerge} +import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, AggregateMode, CollectList, CollectSet, Final, First, Last, Partial, PartialMerge} import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide} import org.apache.spark.sql.catalyst.plans._ import org.apache.spark.sql.catalyst.plans.physical._ @@ -1892,14 +1892,15 @@ object CometObjectHashAggregateExec } /** - * For Partial mode aggregates containing TypedImperativeAggregate functions (like CollectSet), - * the Spark-side output declares buffer columns as BinaryType (since Spark serializes state to - * 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 Partial mode aggregates containing TypedImperativeAggregate functions (like CollectSet or + * CollectList), the Spark-side output declares buffer columns as BinaryType (since Spark + * serializes state to binary). However, the native Comet aggregate produces the actual state + * type (e.g., ArrayType(elementType) for CollectSet/CollectList). This method corrects the + * output schema to match the native state types so the shuffle exchange schema is consistent + * with the actual data. * - * NOTE: If a new TypedImperativeAggregate function (e.g., CollectList) is added natively, add a - * case branch here mapping it to the native state type. + * NOTE: If a new TypedImperativeAggregate function is added natively, add a case branch here + * mapping it to the native state type. */ private def adjustOutputForNativeState(op: ObjectHashAggregateExec): Seq[Attribute] = { // This adjustment only applies to pure-Partial aggregates (checked below). @@ -1916,8 +1917,8 @@ object CometObjectHashAggregateExec val aggFunc = aggExpr.aggregateFunction val bufferAttrs = aggFunc.aggBufferAttributes aggFunc match { - case cs: CollectSet => - val elementType = cs.children.head.dataType + case _: CollectSet | _: CollectList => + val elementType = aggFunc.children.head.dataType val nativeStateType = ArrayType(elementType, containsNull = true) output(bufferIdx) = output(bufferIdx).withDataType(nativeStateType) case _ => diff --git a/spark/src/test/resources/sql-tests/expressions/aggregate/collect_list.sql b/spark/src/test/resources/sql-tests/expressions/aggregate/collect_list.sql new file mode 100644 index 0000000000..507944419a --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/aggregate/collect_list.sql @@ -0,0 +1,405 @@ +-- Licensed to the Apache Software Foundation (ASF) under one +-- or more contributor license agreements. See the NOTICE file +-- distributed with this work for additional information +-- regarding copyright ownership. The ASF licenses this file +-- to you under the Apache License, Version 2.0 (the +-- "License"); you may not use this file except in compliance +-- with the License. You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, +-- software distributed under the License is distributed on an +-- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +-- KIND, either express or implied. See the License for the +-- specific language governing permissions and limitations +-- under the License. + +-- ConfigMatrix: parquet.enable.dictionary=false,true + +-- collect_list result order is non-deterministic across partitions, so +-- every query wraps the result in sort_array to make comparisons stable. + +-- ============================================================ +-- Setup: tables +-- ============================================================ + +statement +CREATE TABLE cl_src_int(i int, grp string) USING parquet + +statement +INSERT INTO cl_src_int VALUES + (1, 'a'), (2, 'a'), (1, 'a'), (3, 'a'), + (4, 'b'), (4, 'b'), (NULL, 'b'), (5, 'b'), + (NULL, 'c'), (NULL, 'c') + +statement +CREATE TABLE cl_src_nulls(val int, grp string) USING parquet + +statement +INSERT INTO cl_src_nulls VALUES + (NULL, 'a'), (NULL, 'a'), (NULL, 'b'), (1, 'b') + +statement +CREATE TABLE cl_src_empty(val int) USING parquet + +statement +CREATE TABLE cl_src_single(val int) USING parquet + +statement +INSERT INTO cl_src_single VALUES (42) + +statement +CREATE TABLE cl_src_dupes(val int, grp string) USING parquet + +statement +INSERT INTO cl_src_dupes VALUES (7, 'a'), (7, 'a'), (7, 'a'), (8, 'b'), (9, 'b') + +-- ============================================================ +-- Basic: integer (global aggregate, no GROUP BY) — duplicates kept +-- ============================================================ + +query +SELECT sort_array(collect_list(i)) FROM cl_src_int + +-- ============================================================ +-- GROUP BY: integer per group +-- ============================================================ + +query +SELECT grp, sort_array(collect_list(i)) FROM cl_src_int GROUP BY grp ORDER BY grp + +-- ============================================================ +-- NULLs: nulls are dropped; all-NULL group returns empty array +-- ============================================================ + +query +SELECT grp, sort_array(collect_list(val)) FROM cl_src_nulls GROUP BY grp ORDER BY grp + +-- ============================================================ +-- Empty table: returns empty array +-- ============================================================ + +query +SELECT sort_array(collect_list(val)) FROM cl_src_empty + +-- ============================================================ +-- Single value +-- ============================================================ + +query +SELECT sort_array(collect_list(val)) FROM cl_src_single + +-- ============================================================ +-- All duplicates in a group — collect_list keeps repeats +-- ============================================================ + +query +SELECT grp, sort_array(collect_list(val)) FROM cl_src_dupes GROUP BY grp ORDER BY grp + +-- ============================================================ +-- Boolean (with NULLs) +-- ============================================================ + +statement +CREATE TABLE cl_src_bool(v boolean, grp string) USING parquet + +statement +INSERT INTO cl_src_bool VALUES + (true, 'a'), (false, 'a'), (true, 'a'), (NULL, 'a'), + (NULL, 'b'), (true, 'b') + +query +SELECT grp, sort_array(collect_list(v)) FROM cl_src_bool GROUP BY grp ORDER BY grp + +-- ============================================================ +-- Byte / Short (with NULLs) +-- ============================================================ + +statement +CREATE TABLE cl_src_small(b tinyint, s smallint, grp string) USING parquet + +statement +INSERT INTO cl_src_small VALUES + (1, 100, 'a'), (2, 200, 'a'), (1, 100, 'a'), (NULL, NULL, 'a'), + (3, 300, 'b'), (NULL, 300, 'b') + +query +SELECT grp, sort_array(collect_list(b)) FROM cl_src_small GROUP BY grp ORDER BY grp + +query +SELECT grp, sort_array(collect_list(s)) FROM cl_src_small GROUP BY grp ORDER BY grp + +-- ============================================================ +-- Int / BigInt (with NULLs) +-- ============================================================ + +statement +CREATE TABLE cl_src_intbig(i int, bi bigint, grp string) USING parquet + +statement +INSERT INTO cl_src_intbig VALUES + (10, 1000000000000, 'a'), (20, 2000000000000, 'a'), + (10, 1000000000000, 'a'), (NULL, NULL, 'a'), + (30, 3000000000000, 'b'), (30, NULL, 'b') + +query +SELECT grp, sort_array(collect_list(i)) FROM cl_src_intbig GROUP BY grp ORDER BY grp + +query +SELECT grp, sort_array(collect_list(bi)) FROM cl_src_intbig GROUP BY grp ORDER BY grp + +-- ============================================================ +-- Float (with NULLs, NaN, Inf, -Inf, +0, -0) +-- collect_list keeps duplicates verbatim, so floating-point is fine. +-- ============================================================ + +statement +CREATE TABLE cl_src_float(v float, grp string) USING parquet + +statement +INSERT INTO cl_src_float VALUES + (1.5, 'a'), (2.5, 'a'), (1.5, 'a'), (NULL, 'a'), + (CAST('NaN' AS FLOAT), 'b'), (CAST('NaN' AS FLOAT), 'b'), (1.0, 'b'), + (CAST('Infinity' AS FLOAT), 'c'), (CAST('-Infinity' AS FLOAT), 'c'), + (CAST('Infinity' AS FLOAT), 'c'), + (CAST(0.0 AS FLOAT), 'd'), (CAST(-0.0 AS FLOAT), 'd'), (1.0, 'd'), (NULL, 'd') + +query +SELECT grp, sort_array(collect_list(v)) FROM cl_src_float GROUP BY grp ORDER BY grp + +-- ============================================================ +-- Double (with NULLs, NaN, Inf, -Inf, +0, -0) +-- ============================================================ + +statement +CREATE TABLE cl_src_double(v double, grp string) USING parquet + +statement +INSERT INTO cl_src_double VALUES + (1.1, 'a'), (2.2, 'a'), (1.1, 'a'), (NULL, 'a'), + (CAST('NaN' AS DOUBLE), 'b'), (CAST('NaN' AS DOUBLE), 'b'), (1.0, 'b'), + (CAST('Infinity' AS DOUBLE), 'c'), (CAST('-Infinity' AS DOUBLE), 'c'), + (CAST('Infinity' AS DOUBLE), 'c'), + (0.0, 'd'), (-0.0, 'd'), (1.0, 'd'), (NULL, 'd') + +query +SELECT grp, sort_array(collect_list(v)) FROM cl_src_double GROUP BY grp ORDER BY grp + +-- ============================================================ +-- String (with NULLs and empty string) +-- ============================================================ + +statement +CREATE TABLE cl_src_string(v string, grp string) USING parquet + +statement +INSERT INTO cl_src_string VALUES + ('hello', 'a'), ('world', 'a'), ('hello', 'a'), (NULL, 'a'), + ('', 'b'), ('x', 'b'), ('', 'b'), (NULL, 'b') + +query +SELECT grp, sort_array(collect_list(v)) FROM cl_src_string GROUP BY grp ORDER BY grp + +-- ============================================================ +-- Binary (with NULLs) +-- ============================================================ + +statement +CREATE TABLE cl_src_binary(v binary, grp string) USING parquet + +statement +INSERT INTO cl_src_binary VALUES + (X'CAFE', 'a'), (X'BABE', 'a'), (X'CAFE', 'a'), (NULL, 'a'), + (X'', 'b'), (X'FF', 'b'), (NULL, 'b') + +query +SELECT grp, sort_array(collect_list(v)) FROM cl_src_binary GROUP BY grp ORDER BY grp + +-- ============================================================ +-- Decimal (with NULLs) +-- ============================================================ + +statement +CREATE TABLE cl_src_decimal(v decimal(10,2), grp string) USING parquet + +statement +INSERT INTO cl_src_decimal VALUES + (1.50, 'a'), (2.50, 'a'), (1.50, 'a'), (NULL, 'a'), + (0.00, 'b'), (99999999.99, 'b'), (NULL, 'b') + +query +SELECT grp, sort_array(collect_list(v)) FROM cl_src_decimal GROUP BY grp ORDER BY grp + +-- ============================================================ +-- Date (with NULLs) +-- ============================================================ + +statement +CREATE TABLE cl_src_date(v date, grp string) USING parquet + +statement +INSERT INTO cl_src_date VALUES + (DATE '2024-01-01', 'a'), (DATE '2024-06-15', 'a'), (DATE '2024-01-01', 'a'), (NULL, 'a'), + (DATE '1970-01-01', 'b'), (NULL, 'b') + +query +SELECT grp, sort_array(collect_list(v)) FROM cl_src_date GROUP BY grp ORDER BY grp + +-- ============================================================ +-- Timestamp (with NULLs) +-- ============================================================ + +statement +CREATE TABLE cl_src_ts(v timestamp, grp string) USING parquet + +statement +INSERT INTO cl_src_ts VALUES + (TIMESTAMP '2024-01-01 00:00:00', 'a'), (TIMESTAMP '2024-06-15 12:30:00', 'a'), + (TIMESTAMP '2024-01-01 00:00:00', 'a'), (NULL, 'a'), + (TIMESTAMP '1970-01-01 00:00:00', 'b'), (NULL, 'b') + +query +SELECT grp, sort_array(collect_list(v)) FROM cl_src_ts GROUP BY grp ORDER BY grp + +-- ============================================================ +-- Mixed with other aggregates +-- ============================================================ + +query +SELECT grp, sort_array(collect_list(i)), count(*), sum(i) +FROM cl_src_int GROUP BY grp ORDER BY grp + +-- ============================================================ +-- Multiple collect_list in the same query +-- ============================================================ + +statement +CREATE TABLE cl_src_multi(a int, b string, grp string) USING parquet + +statement +INSERT INTO cl_src_multi VALUES + (1, 'x', 'g1'), (2, 'y', 'g1'), (1, 'x', 'g1'), + (3, 'z', 'g2'), (NULL, NULL, 'g2') + +query +SELECT grp, sort_array(collect_list(a)), sort_array(collect_list(b)) +FROM cl_src_multi GROUP BY grp ORDER BY grp + +-- ============================================================ +-- DISTINCT: deduplicates before collecting (different planner path) +-- ============================================================ + +query +SELECT grp, sort_array(collect_list(DISTINCT i)) FROM cl_src_int GROUP BY grp ORDER BY grp + +-- ============================================================ +-- HAVING clause with collect_list +-- ============================================================ + +query +SELECT grp, sort_array(collect_list(i)) +FROM cl_src_int GROUP BY grp HAVING size(collect_list(i)) > 1 ORDER BY grp + +-- ============================================================ +-- Result size matches count of non-null values per group +-- (collect_list ignores NULL inputs, like Hive) +-- ============================================================ + +query +SELECT grp, size(collect_list(val)) FROM cl_src_nulls GROUP BY grp ORDER BY grp + +-- ============================================================ +-- array_agg alias (registered as alias of CollectList in +-- FunctionRegistry: `expression[CollectList]("array_agg")`) +-- ============================================================ + +query +SELECT grp, sort_array(array_agg(i)) FROM cl_src_int GROUP BY grp ORDER BY grp + +-- ============================================================ +-- Struct input (Spark DataFrameAggregateSuite "collect functions structs") +-- ============================================================ + +statement +CREATE TABLE cl_src_struct(s struct, grp string) USING parquet + +statement +INSERT INTO cl_src_struct VALUES + (named_struct('x', 1, 'y', 'a'), 'g1'), + (named_struct('x', 2, 'y', 'b'), 'g1'), + (named_struct('x', 1, 'y', 'a'), 'g1'), + (NULL, 'g1'), + (named_struct('x', 3, 'y', 'c'), 'g2'), + (NULL, 'g2') + +query +SELECT grp, sort_array(collect_list(s)) FROM cl_src_struct GROUP BY grp ORDER BY grp + +-- ============================================================ +-- Nested array input +-- ============================================================ + +statement +CREATE TABLE cl_src_array(a array, grp string) USING parquet + +statement +INSERT INTO cl_src_array VALUES + (array(1, 2), 'g1'), + (array(3, 4, 5), 'g1'), + (NULL, 'g1'), + (array(), 'g2'), + (array(NULL), 'g2'), + (NULL, 'g2') + +query +SELECT grp, sort_array(collect_list(a)) FROM cl_src_array GROUP BY grp ORDER BY grp + +-- ============================================================ +-- DECIMAL boundary precisions +-- ============================================================ + +statement +CREATE TABLE cl_src_decimal38(v decimal(38,0), grp string) USING parquet + +statement +INSERT INTO cl_src_decimal38 VALUES + (CAST('99999999999999999999999999999999999999' AS DECIMAL(38,0)), 'a'), + (CAST('-99999999999999999999999999999999999999' AS DECIMAL(38,0)), 'a'), + (CAST(0 AS DECIMAL(38,0)), 'a'), + (NULL, 'a') + +query +SELECT grp, sort_array(collect_list(v)) FROM cl_src_decimal38 GROUP BY grp ORDER BY grp + +-- ============================================================ +-- INT / BIGINT boundary values +-- ============================================================ + +statement +CREATE TABLE cl_src_bounds(i int, bi bigint) USING parquet + +statement +INSERT INTO cl_src_bounds VALUES + (-2147483648, -9223372036854775808), + (2147483647, 9223372036854775807), + (0, 0), + (NULL, NULL) + +query +SELECT sort_array(collect_list(i)), sort_array(collect_list(bi)) FROM cl_src_bounds + +-- ============================================================ +-- Spark SPARK-17641 regression: collect functions should not +-- collect null values. Verifies the absolute size matches the +-- number of non-null inputs across mixed types. +-- ============================================================ + +statement +CREATE TABLE cl_src_17641(a string, b int) USING parquet + +statement +INSERT INTO cl_src_17641 VALUES ('1', 2), (NULL, 2), ('1', 4) + +query +SELECT sort_array(collect_list(a)), sort_array(collect_list(b)) FROM cl_src_17641 From 6d2aa2b4c5b3b6e14c42c011293851e72b5e654d Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Wed, 24 Jun 2026 14:05:17 -0600 Subject: [PATCH 2/6] fix: fall back collect_list/collect_set in multi-stage distinct aggregates 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 #4724. --- dev/diffs/3.4.3.diff | 20 +++++++++++--- dev/diffs/3.5.8.diff | 18 +++++++++++-- dev/diffs/4.0.2.diff | 18 +++++++++++-- dev/diffs/4.1.2.diff | 18 +++++++++++-- .../apache/comet/rules/CometExecRule.scala | 20 +++++++++++++- .../apache/comet/serde/QueryPlanSerde.scala | 16 +++++++++++ .../apache/spark/sql/comet/operators.scala | 15 +++++++++++ .../comet/exec/CometAggregateSuite.scala | 27 +++++++++++++++++++ 8 files changed, 142 insertions(+), 10 deletions(-) diff --git a/dev/diffs/3.4.3.diff b/dev/diffs/3.4.3.diff index f0cb101532..bc42c80687 100644 --- a/dev/diffs/3.4.3.diff +++ b/dev/diffs/3.4.3.diff @@ -260,7 +260,7 @@ index cf40e944c09..bdd5be4f462 100644 test("A cached table preserves the partitioning and ordering of its cached SparkPlan") { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala -index 1cc09c3d7fc..f031fa45c33 100644 +index 1cc09c3d7fc..2a6b073ca7a 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala @@ -27,7 +27,7 @@ import org.apache.spark.SparkException @@ -272,7 +272,21 @@ index 1cc09c3d7fc..f031fa45c33 100644 import org.apache.spark.sql.expressions.Window import org.apache.spark.sql.functions._ import org.apache.spark.sql.internal.SQLConf -@@ -755,7 +755,7 @@ class DataFrameAggregateSuite extends QueryTest +@@ -733,7 +733,12 @@ class DataFrameAggregateSuite extends QueryTest + } + + test("SPARK-22223: ObjectHashAggregate should not introduce unnecessary shuffle") { +- withSQLConf(SQLConf.USE_OBJECT_HASH_AGG.key -> "true") { ++ withSQLConf( ++ SQLConf.USE_OBJECT_HASH_AGG.key -> "true", ++ // Comet runs collect_list/collect_set natively (CometHashAggregateExec), so the ++ // ObjectHashAggregateExec this test asserts on is no longer present. Disable Comet ++ // to preserve the original plan-shape assertions. ++ "spark.comet.enabled" -> "false") { + val df = Seq(("1", "2", 1), ("1", "2", 2), ("2", "3", 3), ("2", "3", 4)).toDF("a", "b", "c") + .repartition(col("a")) + +@@ -755,7 +760,7 @@ class DataFrameAggregateSuite extends QueryTest assert(objHashAggPlans.nonEmpty) val exchangePlans = collect(aggPlan) { @@ -2977,7 +2991,7 @@ index dd55fcfe42c..d9a3f2df535 100644 spark.internalCreateDataFrame(withoutFilters.execute(), schema) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala b/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala -index ed2e309fa07..0658bfe9e12 100644 +index ed2e309fa07..863868646a8 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala @@ -74,6 +74,20 @@ trait SharedSparkSessionBase diff --git a/dev/diffs/3.5.8.diff b/dev/diffs/3.5.8.diff index 5e875b2c9b..83942632db 100644 --- a/dev/diffs/3.5.8.diff +++ b/dev/diffs/3.5.8.diff @@ -241,7 +241,7 @@ index e5494726695..00937f025c2 100644 test("A cached table preserves the partitioning and ordering of its cached SparkPlan") { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala -index 6f3090d8908..c08a60fb0c2 100644 +index 6f3090d8908..4f2e8970be8 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala @@ -28,7 +28,7 @@ import org.apache.spark.sql.catalyst.plans.logical.Expand @@ -253,7 +253,21 @@ index 6f3090d8908..c08a60fb0c2 100644 import org.apache.spark.sql.expressions.Window import org.apache.spark.sql.functions._ import org.apache.spark.sql.internal.SQLConf -@@ -793,7 +793,7 @@ class DataFrameAggregateSuite extends QueryTest +@@ -771,7 +771,12 @@ class DataFrameAggregateSuite extends QueryTest + } + + test("SPARK-22223: ObjectHashAggregate should not introduce unnecessary shuffle") { +- withSQLConf(SQLConf.USE_OBJECT_HASH_AGG.key -> "true") { ++ withSQLConf( ++ SQLConf.USE_OBJECT_HASH_AGG.key -> "true", ++ // Comet runs collect_list/collect_set natively (CometHashAggregateExec), so the ++ // ObjectHashAggregateExec this test asserts on is no longer present. Disable Comet ++ // to preserve the original plan-shape assertions. ++ "spark.comet.enabled" -> "false") { + val df = Seq(("1", "2", 1), ("1", "2", 2), ("2", "3", 3), ("2", "3", 4)).toDF("a", "b", "c") + .repartition(col("a")) + +@@ -793,7 +798,7 @@ class DataFrameAggregateSuite extends QueryTest assert(objHashAggPlans.nonEmpty) val exchangePlans = collect(aggPlan) { diff --git a/dev/diffs/4.0.2.diff b/dev/diffs/4.0.2.diff index e7fb4b2f79..1e91b55d0f 100644 --- a/dev/diffs/4.0.2.diff +++ b/dev/diffs/4.0.2.diff @@ -378,7 +378,7 @@ index 0f42502f1d9..e9ff802141f 100644 withTempView("t0", "t1", "t2") { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala -index 9db406ff12f..245e4caa319 100644 +index 9db406ff12f..066b03283e2 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala @@ -30,7 +30,7 @@ import org.apache.spark.sql.errors.DataTypeErrors.toSQLId @@ -390,7 +390,21 @@ index 9db406ff12f..245e4caa319 100644 import org.apache.spark.sql.expressions.Window import org.apache.spark.sql.functions._ import org.apache.spark.sql.internal.SQLConf -@@ -855,7 +855,7 @@ class DataFrameAggregateSuite extends QueryTest +@@ -833,7 +833,12 @@ class DataFrameAggregateSuite extends QueryTest + } + + test("SPARK-22223: ObjectHashAggregate should not introduce unnecessary shuffle") { +- withSQLConf(SQLConf.USE_OBJECT_HASH_AGG.key -> "true") { ++ withSQLConf( ++ SQLConf.USE_OBJECT_HASH_AGG.key -> "true", ++ // Comet runs collect_list/collect_set natively (CometHashAggregateExec), so the ++ // ObjectHashAggregateExec this test asserts on is no longer present. Disable Comet ++ // to preserve the original plan-shape assertions. ++ "spark.comet.enabled" -> "false") { + val df = Seq(("1", "2", 1), ("1", "2", 2), ("2", "3", 3), ("2", "3", 4)).toDF("a", "b", "c") + .repartition(col("a")) + +@@ -855,7 +860,7 @@ class DataFrameAggregateSuite extends QueryTest assert(objHashAggPlans.nonEmpty) val exchangePlans = collect(aggPlan) { diff --git a/dev/diffs/4.1.2.diff b/dev/diffs/4.1.2.diff index 402582a5ca..0a6b2bf09a 100644 --- a/dev/diffs/4.1.2.diff +++ b/dev/diffs/4.1.2.diff @@ -392,7 +392,7 @@ index 0d807aeae4d..6d7744e771b 100644 withTempView("t0", "t1", "t2") { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala -index bfe15b33768..55c23a38ccc 100644 +index bfe15b33768..c391a9e1790 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala @@ -31,7 +31,7 @@ import org.apache.spark.sql.errors.DataTypeErrors.toSQLId @@ -404,7 +404,21 @@ index bfe15b33768..55c23a38ccc 100644 import org.apache.spark.sql.expressions.Window import org.apache.spark.sql.functions._ import org.apache.spark.sql.internal.SQLConf -@@ -856,7 +856,7 @@ class DataFrameAggregateSuite extends QueryTest +@@ -834,7 +834,12 @@ class DataFrameAggregateSuite extends QueryTest + } + + test("SPARK-22223: ObjectHashAggregate should not introduce unnecessary shuffle") { +- withSQLConf(SQLConf.USE_OBJECT_HASH_AGG.key -> "true") { ++ withSQLConf( ++ SQLConf.USE_OBJECT_HASH_AGG.key -> "true", ++ // Comet runs collect_list/collect_set natively (CometHashAggregateExec), so the ++ // ObjectHashAggregateExec this test asserts on is no longer present. Disable Comet ++ // to preserve the original plan-shape assertions. ++ "spark.comet.enabled" -> "false") { + val df = Seq(("1", "2", 1), ("1", "2", 2), ("2", "3", 3), ("2", "3", 4)).toDF("a", "b", "c") + .repartition(col("a")) + +@@ -856,7 +861,7 @@ class DataFrameAggregateSuite extends QueryTest assert(objHashAggPlans.nonEmpty) val exchangePlans = collect(aggPlan) { diff --git a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala index 50cd0927b4..54c15f6c5b 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -23,7 +23,7 @@ import scala.collection.mutable.ListBuffer import org.apache.spark.sql.SparkSession import org.apache.spark.sql.catalyst.expressions.{Divide, DoubleLiteral, EqualNullSafe, EqualTo, Expression, FloatLiteral, GreaterThan, GreaterThanOrEqual, KnownFloatingPointNormalized, LessThan, LessThanOrEqual, NamedExpression, Remainder} -import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateMode, Final, Partial} +import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateMode, Final, Partial, PartialMerge} import org.apache.spark.sql.catalyst.optimizer.NormalizeNaNAndZero import org.apache.spark.sql.catalyst.rules.Rule import org.apache.spark.sql.catalyst.trees.TreeNodeTag @@ -836,6 +836,24 @@ case class CometExecRule(session: SparkSession) } } } + + // CollectList/CollectSet round-trip an ArrayType buffer that Spark declares as BinaryType. + // In a multi-stage aggregate with a PartialMerge stage (e.g. Spark's distinct-aggregate + // rewrite), Comet cannot represent that buffer consistently across the intermediate stages + // (issue #4724), so a fully-native pipeline crashes. Force the whole chain to fall back to + // Spark by tagging the feeding pure-Partial; the PartialMerge/Final stages then fall back + // via the buffer-source check in doConvert. + if (agg.aggregateExpressions.exists(_.mode == PartialMerge) && + QueryPlanSerde.hasIncompatibleBufferAgg(agg.aggregateExpressions)) { + findPartialAggInPlan(agg.child).foreach { partial => + if (canAggregateBeConverted(partial, Partial)) { + partial.setTagValue( + CometExecRule.COMET_UNSAFE_PARTIAL, + "Partial aggregate disabled: part of a multi-stage CollectList/CollectSet " + + "aggregate whose intermediate buffer cannot round-trip in Comet (issue #4724)") + } + } + } case _ => } } diff --git a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala index 8935f7cf22..f8d1ad8d01 100644 --- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala +++ b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala @@ -414,6 +414,22 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim { } } + /** + * 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 + * for its serialized TypedImperativeAggregate buffer. Comet cannot interpret Spark's Binary + * buffer for these functions, and cannot yet represent the buffer consistently across the + * intermediate PartialMerge stages of a multi-stage aggregate (issue #4724). Such aggregates + * 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 = { + aggExprs.exists(_.aggregateFunction match { + case _: CollectList | _: CollectSet => true + case _ => false + }) + } + // A unique id for each expression. ~used to look up QueryContext during error creation. private val exprIdCounter = new AtomicLong(0) diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala index 19a0841650..df26d1cc8c 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala @@ -1547,6 +1547,21 @@ trait CometBaseAggregate { return None } + // CollectList/CollectSet declare their buffer as BinaryType in Spark but produce a native + // ArrayType state. A PartialMerge (or multi-mode) stage consuming that buffer can only read it + // when an upstream Comet partial produced it; if the buffer came from a Spark partial (no Comet + // partial in the child subtree), Comet cannot interpret Spark's serialized Binary buffer and + // must fall back. The matching Partial is forced to fall back too (see CometExecRule), so the + // whole multi-stage chain runs consistently in one engine. See issue #4724. + if (modes.contains(PartialMerge) && + QueryPlanSerde.hasIncompatibleBufferAgg(aggregate.aggregateExpressions) && + findCometPartialAgg(aggregate.child).isEmpty) { + withFallbackReason( + aggregate, + "CollectList/CollectSet PartialMerge cannot read a Spark-produced intermediate buffer") + return None + } + // Check if this aggregate has been tagged as unsafe for mixed execution // (Comet partial + Spark final with incompatible intermediate buffers) val unsafeReason = aggregate.getTagValue(CometExecRule.COMET_UNSAFE_PARTIAL) diff --git a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala index ae14c68207..b689a6e2eb 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala @@ -48,6 +48,33 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { override protected def sparkConf: SparkConf = super.sparkConf.set(SQLConf.ANSI_ENABLED.key, "false") + test("collect_list/collect_set combined with distinct aggregate falls back safely") { + // SPARK-17616: a distinct aggregate combined with collect_list/collect_set produces a + // multi-stage plan where the buffer-producing Partial may run in Spark (e.g. over a + // non-native LocalTableScan). Comet cannot read Spark's serialized Binary buffer, so the + // dependent PartialMerge/Final stages must also fall back rather than crash. See issue #4724 + // for enabling the fully-native distinct path. + import org.apache.spark.sql.functions.{collect_list, collect_set, sort_array} + // Non-native source (LocalTableScan): the buffer-producing Partial runs in Spark. + val df = Seq((1, 3, "a"), (1, 2, "b"), (3, 4, "c"), (3, 4, "c"), (3, 5, "d")) + .toDF("x", "y", "z") + checkSparkAnswer( + df.groupBy(col("x")).agg(count_distinct(col("y")), sort_array(collect_list(col("z"))))) + checkSparkAnswer( + df.groupBy(col("x")).agg(count_distinct(col("y")), sort_array(collect_set(col("z"))))) + + // Native source (Parquet): the whole multi-stage distinct chain must still fall back to + // Spark consistently (issue #4724), rather than running a fully-native pipeline that crashes. + withParquetTable( + Seq((1, 3, "a"), (1, 2, "b"), (3, 4, "c"), (3, 4, "c"), (3, 5, "d")), + "t17616") { + for (fn <- Seq("collect_list", "collect_set")) { + checkSparkAnswer( + sql(s"SELECT _1, count(distinct _2), sort_array($fn(_3)) FROM t17616 GROUP BY _1")) + } + } + } + test("min/max floating point with negative zero") { val r = new Random(42) val schema = StructType( From 9e680c42f8118c967530fcbe5b257fbe1fd3e6ad Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Wed, 1 Jul 2026 17:21:23 -0600 Subject: [PATCH 3/6] feat: fall back collect_list/collect_set on Spark 4.2 RESPECT NULLS 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. --- .../expression-audits/agg_funcs.md | 1 + .../apache/comet/rules/CometExecRule.scala | 2 +- .../apache/comet/serde/QueryPlanSerde.scala | 2 +- .../org/apache/comet/serde/aggregates.scala | 32 ++++++++++++++----- .../apache/spark/sql/comet/operators.scala | 2 +- .../apache/comet/shims/CometCollectShim.scala | 32 +++++++++++++++++++ .../apache/comet/shims/CometCollectShim.scala | 32 +++++++++++++++++++ .../apache/comet/shims/CometCollectShim.scala | 32 +++++++++++++++++++ .../apache/comet/shims/CometCollectShim.scala | 32 +++++++++++++++++++ .../apache/comet/shims/CometCollectShim.scala | 32 +++++++++++++++++++ 10 files changed, 188 insertions(+), 11 deletions(-) create mode 100644 spark/src/main/spark-3.4/org/apache/comet/shims/CometCollectShim.scala create mode 100644 spark/src/main/spark-3.5/org/apache/comet/shims/CometCollectShim.scala create mode 100644 spark/src/main/spark-4.0/org/apache/comet/shims/CometCollectShim.scala create mode 100644 spark/src/main/spark-4.1/org/apache/comet/shims/CometCollectShim.scala create mode 100644 spark/src/main/spark-4.2/org/apache/comet/shims/CometCollectShim.scala diff --git a/docs/source/contributor-guide/expression-audits/agg_funcs.md b/docs/source/contributor-guide/expression-audits/agg_funcs.md index 827ea05fe7..deaf1a99dc 100644 --- a/docs/source/contributor-guide/expression-audits/agg_funcs.md +++ b/docs/source/contributor-guide/expression-audits/agg_funcs.md @@ -46,6 +46,7 @@ - Spark 4.0.1 (audited 2026-06-24): only structural change is adding `with UnaryLike[Expression]` to the case class (no behavior change). - Spark 4.1.1 (audited 2026-06-24): identical to 4.0.1. - Comet implementation: native side delegates to `datafusion_spark::function::aggregate::collect::SparkCollectList`, which wraps `ArrayAggAccumulator` with `ignore_nulls = true` and converts a final NULL accumulator state to an empty array (matching Spark's `defaultResult`). The native return type is `List(Field, containsNull = true)`, while Spark uses `containsNull = false`. Because nulls are filtered before insertion, no nulls actually appear in the array, so this is a schema-shape difference only and tests using `checkSparkAnswerAndOperator` accept it (same pattern already in use for `collect_set`). +- Spark 4.2 (preview): `CollectList` and `CollectSet` gain an `ignoreNulls` field (default `true`); `RESPECT NULLS` sets it to `false` and keeps null elements. The native path always drops nulls, so `CometCollectShim` reads the field per Spark version (always `true` on 3.4-4.1) and `CometCollectList` / `CometCollectSet` report `Unsupported` when it is `false`, falling back to Spark. ## median diff --git a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala index b181bced46..137ba4adfe 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -866,7 +866,7 @@ case class CometExecRule(session: SparkSession) // Spark by tagging the feeding pure-Partial; the PartialMerge/Final stages then fall back // via the buffer-source check in doConvert. if (agg.aggregateExpressions.exists(_.mode == PartialMerge) && - QueryPlanSerde.hasIncompatibleBufferAgg(agg.aggregateExpressions)) { + QueryPlanSerde.hasNativeArrayBufferAgg(agg.aggregateExpressions)) { findPartialAggInPlan(agg.child).foreach { partial => if (canAggregateBeConverted(partial, Partial)) { partial.setTagValue( diff --git a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala index f1a8fcba7f..6043b93155 100644 --- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala +++ b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala @@ -434,7 +434,7 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim { * 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 = { + def hasNativeArrayBufferAgg(aggExprs: Seq[AggregateExpression]): Boolean = { aggExprs.exists(_.aggregateFunction match { case _: CollectList | _: CollectSet => true case _ => false diff --git a/spark/src/main/scala/org/apache/comet/serde/aggregates.scala b/spark/src/main/scala/org/apache/comet/serde/aggregates.scala index e9b5984d5a..57f98bf02c 100644 --- a/spark/src/main/scala/org/apache/comet/serde/aggregates.scala +++ b/spark/src/main/scala/org/apache/comet/serde/aggregates.scala @@ -29,7 +29,7 @@ import org.apache.spark.sql.types.{ByteType, DecimalType, DoubleType, IntegerTyp import org.apache.comet.CometConf.COMET_EXEC_STRICT_FLOATING_POINT import org.apache.comet.CometSparkSessionExtensions.{isSpark41Plus, withFallbackReason} import org.apache.comet.serde.QueryPlanSerde.{evalModeToProto, exprToProto, serializeDataType} -import org.apache.comet.shims.CometEvalModeUtil +import org.apache.comet.shims.{CometCollectShim, CometEvalModeUtil} object CometMin extends CometAggregateExpressionSerde[Min] { @@ -784,13 +784,19 @@ object CometCollectSet extends CometAggregateExpressionSerde[CollectSet] { " `spark.comet.expression.CollectSet.allowIncompatible=true` is set.") override def getSupportLevel(expr: CollectSet): SupportLevel = { - SupportLevel - .strictFloatingPointReason( - expr.children.head.dataType, - "collect_set on floating-point types " + - "(Comet deduplicates NaN values while Spark treats each NaN as distinct)") - .map(reason => Incompatible(Some(reason))) - .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)) { + Unsupported(Some("collect_set with RESPECT NULLS (ignoreNulls = false) is not supported")) + } else { + SupportLevel + .strictFloatingPointReason( + expr.children.head.dataType, + "collect_set on floating-point types " + + "(Comet deduplicates NaN values while Spark treats each NaN as distinct)") + .map(reason => Incompatible(Some(reason))) + .getOrElse(Compatible()) + } } override def convert( @@ -825,6 +831,16 @@ object CometCollectSet extends CometAggregateExpressionSerde[CollectSet] { object CometCollectList extends CometAggregateExpressionSerde[CollectList] { + override def getSupportLevel(expr: CollectList): SupportLevel = { + // The native path delegates to SparkCollectList, which always drops null inputs. Spark 4.2 + // adds `RESPECT NULLS` (`ignoreNulls = false`), which keeps nulls, so fall back there. + if (!CometCollectShim.ignoreNulls(expr)) { + Unsupported(Some("collect_list with RESPECT NULLS (ignoreNulls = false) is not supported")) + } else { + Compatible() + } + } + override def convert( aggExpr: AggregateExpression, expr: CollectList, diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala index f76f336414..4115f2c67e 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala @@ -1554,7 +1554,7 @@ trait CometBaseAggregate { // must fall back. The matching Partial is forced to fall back too (see CometExecRule), so the // whole multi-stage chain runs consistently in one engine. See issue #4724. if (modes.contains(PartialMerge) && - QueryPlanSerde.hasIncompatibleBufferAgg(aggregate.aggregateExpressions) && + QueryPlanSerde.hasNativeArrayBufferAgg(aggregate.aggregateExpressions) && findCometPartialAgg(aggregate.child).isEmpty) { withFallbackReason( aggregate, diff --git a/spark/src/main/spark-3.4/org/apache/comet/shims/CometCollectShim.scala b/spark/src/main/spark-3.4/org/apache/comet/shims/CometCollectShim.scala new file mode 100644 index 0000000000..69f33a1b29 --- /dev/null +++ b/spark/src/main/spark-3.4/org/apache/comet/shims/CometCollectShim.scala @@ -0,0 +1,32 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.comet.shims + +import org.apache.spark.sql.catalyst.expressions.aggregate.{CollectList, CollectSet} + +/** + * Shim for the `ignoreNulls` flag on `CollectList` / `CollectSet`. Spark 3.4 through 4.1 have no + * such field: these aggregates always drop null inputs, so this shim reports `true`. Spark 4.2 + * added `ignoreNulls` (settable to `false` via `RESPECT NULLS`), handled by the spark-4.2 shim. + */ +object CometCollectShim { + def ignoreNulls(agg: CollectList): Boolean = true + def ignoreNulls(agg: CollectSet): Boolean = true +} diff --git a/spark/src/main/spark-3.5/org/apache/comet/shims/CometCollectShim.scala b/spark/src/main/spark-3.5/org/apache/comet/shims/CometCollectShim.scala new file mode 100644 index 0000000000..69f33a1b29 --- /dev/null +++ b/spark/src/main/spark-3.5/org/apache/comet/shims/CometCollectShim.scala @@ -0,0 +1,32 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.comet.shims + +import org.apache.spark.sql.catalyst.expressions.aggregate.{CollectList, CollectSet} + +/** + * Shim for the `ignoreNulls` flag on `CollectList` / `CollectSet`. Spark 3.4 through 4.1 have no + * such field: these aggregates always drop null inputs, so this shim reports `true`. Spark 4.2 + * added `ignoreNulls` (settable to `false` via `RESPECT NULLS`), handled by the spark-4.2 shim. + */ +object CometCollectShim { + def ignoreNulls(agg: CollectList): Boolean = true + def ignoreNulls(agg: CollectSet): Boolean = true +} diff --git a/spark/src/main/spark-4.0/org/apache/comet/shims/CometCollectShim.scala b/spark/src/main/spark-4.0/org/apache/comet/shims/CometCollectShim.scala new file mode 100644 index 0000000000..69f33a1b29 --- /dev/null +++ b/spark/src/main/spark-4.0/org/apache/comet/shims/CometCollectShim.scala @@ -0,0 +1,32 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.comet.shims + +import org.apache.spark.sql.catalyst.expressions.aggregate.{CollectList, CollectSet} + +/** + * Shim for the `ignoreNulls` flag on `CollectList` / `CollectSet`. Spark 3.4 through 4.1 have no + * such field: these aggregates always drop null inputs, so this shim reports `true`. Spark 4.2 + * added `ignoreNulls` (settable to `false` via `RESPECT NULLS`), handled by the spark-4.2 shim. + */ +object CometCollectShim { + def ignoreNulls(agg: CollectList): Boolean = true + def ignoreNulls(agg: CollectSet): Boolean = true +} diff --git a/spark/src/main/spark-4.1/org/apache/comet/shims/CometCollectShim.scala b/spark/src/main/spark-4.1/org/apache/comet/shims/CometCollectShim.scala new file mode 100644 index 0000000000..69f33a1b29 --- /dev/null +++ b/spark/src/main/spark-4.1/org/apache/comet/shims/CometCollectShim.scala @@ -0,0 +1,32 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.comet.shims + +import org.apache.spark.sql.catalyst.expressions.aggregate.{CollectList, CollectSet} + +/** + * Shim for the `ignoreNulls` flag on `CollectList` / `CollectSet`. Spark 3.4 through 4.1 have no + * such field: these aggregates always drop null inputs, so this shim reports `true`. Spark 4.2 + * added `ignoreNulls` (settable to `false` via `RESPECT NULLS`), handled by the spark-4.2 shim. + */ +object CometCollectShim { + def ignoreNulls(agg: CollectList): Boolean = true + def ignoreNulls(agg: CollectSet): Boolean = true +} diff --git a/spark/src/main/spark-4.2/org/apache/comet/shims/CometCollectShim.scala b/spark/src/main/spark-4.2/org/apache/comet/shims/CometCollectShim.scala new file mode 100644 index 0000000000..defae750ae --- /dev/null +++ b/spark/src/main/spark-4.2/org/apache/comet/shims/CometCollectShim.scala @@ -0,0 +1,32 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.comet.shims + +import org.apache.spark.sql.catalyst.expressions.aggregate.{CollectList, CollectSet} + +/** + * Shim for the `ignoreNulls` flag on `CollectList` / `CollectSet`. Spark 4.2 added the field + * (`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 { + def ignoreNulls(agg: CollectList): Boolean = agg.ignoreNulls + def ignoreNulls(agg: CollectSet): Boolean = agg.ignoreNulls +} From 589ebddca9891955f276b2ccf9cf9dd76a85bd09 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 9 Jul 2026 15:16:26 -0600 Subject: [PATCH 4/6] test: accept Comet operators in SPARK-22223 instead of disabling Comet 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. --- dev/diffs/3.4.3.diff | 22 +++++++--------------- dev/diffs/3.5.8.diff | 22 +++++++--------------- dev/diffs/4.0.2.diff | 22 +++++++--------------- dev/diffs/4.1.2.diff | 22 +++++++--------------- 4 files changed, 28 insertions(+), 60 deletions(-) diff --git a/dev/diffs/3.4.3.diff b/dev/diffs/3.4.3.diff index 22bf3bf417..133daf681a 100644 --- a/dev/diffs/3.4.3.diff +++ b/dev/diffs/3.4.3.diff @@ -260,7 +260,7 @@ index cf40e944c09..bdd5be4f462 100644 test("A cached table preserves the partitioning and ordering of its cached SparkPlan") { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala -index 1cc09c3d7fc..5451c678ab8 100644 +index 1cc09c3d7fc..9e1e883d450 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala @@ -26,8 +26,9 @@ import org.scalatest.matchers.must.Matchers.the @@ -285,21 +285,13 @@ index 1cc09c3d7fc..5451c678ab8 100644 } // test case for ObjectHashAggregate and SortAggregate -@@ -733,7 +736,12 @@ class DataFrameAggregateSuite extends QueryTest - } - - test("SPARK-22223: ObjectHashAggregate should not introduce unnecessary shuffle") { -- withSQLConf(SQLConf.USE_OBJECT_HASH_AGG.key -> "true") { -+ withSQLConf( -+ SQLConf.USE_OBJECT_HASH_AGG.key -> "true", -+ // Comet runs collect_list/collect_set natively (CometHashAggregateExec), so the -+ // ObjectHashAggregateExec this test asserts on is no longer present. Disable Comet -+ // to preserve the original plan-shape assertions. -+ "spark.comet.enabled" -> "false") { - val df = Seq(("1", "2", 1), ("1", "2", 2), ("2", "3", 3), ("2", "3", 4)).toDF("a", "b", "c") - .repartition(col("a")) +@@ -750,12 +753,12 @@ class DataFrameAggregateSuite extends QueryTest + assert(sortAggPlans.isEmpty) -@@ -755,7 +763,7 @@ class DataFrameAggregateSuite extends QueryTest + val objHashAggPlans = collect(aggPlan) { +- case objHashAgg: ObjectHashAggregateExec => objHashAgg ++ case objHashAgg @ (_: ObjectHashAggregateExec | _: CometHashAggregateExec) => objHashAgg + } assert(objHashAggPlans.nonEmpty) val exchangePlans = collect(aggPlan) { diff --git a/dev/diffs/3.5.8.diff b/dev/diffs/3.5.8.diff index 8b860d794d..428e964bad 100644 --- a/dev/diffs/3.5.8.diff +++ b/dev/diffs/3.5.8.diff @@ -241,7 +241,7 @@ index e5494726695..00937f025c2 100644 test("A cached table preserves the partitioning and ordering of its cached SparkPlan") { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala -index 6f3090d8908..af60512e183 100644 +index 6f3090d8908..4774aad5019 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala @@ -27,8 +27,9 @@ import org.apache.spark.{SparkException, SparkThrowable} @@ -266,21 +266,13 @@ index 6f3090d8908..af60512e183 100644 } // test case for ObjectHashAggregate and SortAggregate -@@ -771,7 +774,12 @@ class DataFrameAggregateSuite extends QueryTest - } - - test("SPARK-22223: ObjectHashAggregate should not introduce unnecessary shuffle") { -- withSQLConf(SQLConf.USE_OBJECT_HASH_AGG.key -> "true") { -+ withSQLConf( -+ SQLConf.USE_OBJECT_HASH_AGG.key -> "true", -+ // Comet runs collect_list/collect_set natively (CometHashAggregateExec), so the -+ // ObjectHashAggregateExec this test asserts on is no longer present. Disable Comet -+ // to preserve the original plan-shape assertions. -+ "spark.comet.enabled" -> "false") { - val df = Seq(("1", "2", 1), ("1", "2", 2), ("2", "3", 3), ("2", "3", 4)).toDF("a", "b", "c") - .repartition(col("a")) +@@ -788,12 +791,12 @@ class DataFrameAggregateSuite extends QueryTest + assert(sortAggPlans.isEmpty) -@@ -793,7 +801,7 @@ class DataFrameAggregateSuite extends QueryTest + val objHashAggPlans = collect(aggPlan) { +- case objHashAgg: ObjectHashAggregateExec => objHashAgg ++ case objHashAgg @ (_: ObjectHashAggregateExec | _: CometHashAggregateExec) => objHashAgg + } assert(objHashAggPlans.nonEmpty) val exchangePlans = collect(aggPlan) { diff --git a/dev/diffs/4.0.2.diff b/dev/diffs/4.0.2.diff index 9107a3df25..8ee55fea1d 100644 --- a/dev/diffs/4.0.2.diff +++ b/dev/diffs/4.0.2.diff @@ -378,7 +378,7 @@ index 0f42502f1d9..e9ff802141f 100644 withTempView("t0", "t1", "t2") { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala -index 9db406ff12f..5c987be87cb 100644 +index 9db406ff12f..b3d55394d25 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala @@ -29,8 +29,9 @@ import org.apache.spark.sql.catalyst.util.AUTO_GENERATED_ALIAS @@ -403,21 +403,13 @@ index 9db406ff12f..5c987be87cb 100644 } // test case for ObjectHashAggregate and SortAggregate -@@ -833,7 +836,12 @@ class DataFrameAggregateSuite extends QueryTest - } - - test("SPARK-22223: ObjectHashAggregate should not introduce unnecessary shuffle") { -- withSQLConf(SQLConf.USE_OBJECT_HASH_AGG.key -> "true") { -+ withSQLConf( -+ SQLConf.USE_OBJECT_HASH_AGG.key -> "true", -+ // Comet runs collect_list/collect_set natively (CometHashAggregateExec), so the -+ // ObjectHashAggregateExec this test asserts on is no longer present. Disable Comet -+ // to preserve the original plan-shape assertions. -+ "spark.comet.enabled" -> "false") { - val df = Seq(("1", "2", 1), ("1", "2", 2), ("2", "3", 3), ("2", "3", 4)).toDF("a", "b", "c") - .repartition(col("a")) +@@ -850,12 +853,12 @@ class DataFrameAggregateSuite extends QueryTest + assert(sortAggPlans.isEmpty) -@@ -855,7 +863,7 @@ class DataFrameAggregateSuite extends QueryTest + val objHashAggPlans = collect(aggPlan) { +- case objHashAgg: ObjectHashAggregateExec => objHashAgg ++ case objHashAgg @ (_: ObjectHashAggregateExec | _: CometHashAggregateExec) => objHashAgg + } assert(objHashAggPlans.nonEmpty) val exchangePlans = collect(aggPlan) { diff --git a/dev/diffs/4.1.2.diff b/dev/diffs/4.1.2.diff index 6cabcbf67c..599e2cd077 100644 --- a/dev/diffs/4.1.2.diff +++ b/dev/diffs/4.1.2.diff @@ -392,7 +392,7 @@ index 0d807aeae4d..6d7744e771b 100644 withTempView("t0", "t1", "t2") { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala -index bfe15b33768..2ccb7f4713c 100644 +index bfe15b33768..13aeb3f6610 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala @@ -30,8 +30,9 @@ import org.apache.spark.sql.catalyst.util.AUTO_GENERATED_ALIAS @@ -417,21 +417,13 @@ index bfe15b33768..2ccb7f4713c 100644 } // test case for ObjectHashAggregate and SortAggregate -@@ -834,7 +837,12 @@ class DataFrameAggregateSuite extends QueryTest - } +@@ -851,12 +854,12 @@ class DataFrameAggregateSuite extends QueryTest + assert(sortAggPlans.isEmpty) - test("SPARK-22223: ObjectHashAggregate should not introduce unnecessary shuffle") { -- withSQLConf(SQLConf.USE_OBJECT_HASH_AGG.key -> "true") { -+ withSQLConf( -+ SQLConf.USE_OBJECT_HASH_AGG.key -> "true", -+ // Comet runs collect_list/collect_set natively (CometHashAggregateExec), so the -+ // ObjectHashAggregateExec this test asserts on is no longer present. Disable Comet -+ // to preserve the original plan-shape assertions. -+ "spark.comet.enabled" -> "false") { - val df = Seq(("1", "2", 1), ("1", "2", 2), ("2", "3", 3), ("2", "3", 4)).toDF("a", "b", "c") - .repartition(col("a")) - -@@ -856,7 +864,7 @@ class DataFrameAggregateSuite extends QueryTest + val objHashAggPlans = collect(aggPlan) { +- case objHashAgg: ObjectHashAggregateExec => objHashAgg ++ case objHashAgg @ (_: ObjectHashAggregateExec | _: CometHashAggregateExec) => objHashAgg + } assert(objHashAggPlans.nonEmpty) val exchangePlans = collect(aggPlan) { From a98c3fcb77e931e31bcc020ecf4d22dfb7e570bb Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 9 Jul 2026 17:02:29 -0600 Subject: [PATCH 5/6] style: reflow adjustOutputForNativeState doc comment Adding CollectList to the comment pushed a line past the column limit during the apache/main merge; reflow to satisfy spotless. --- .../src/main/scala/org/apache/spark/sql/comet/operators.scala | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala index 3d7146484e..5b5a2821a7 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala @@ -1802,8 +1802,8 @@ trait CometBaseAggregate { /** * For partial-like aggregates containing TypedImperativeAggregate functions (like CollectSet, * CollectList, and Percentile), the Spark-side output declares buffer columns as BinaryType - * because Spark serializes state to binary. Native Comet emits the actual state type, so fix the - * exposed output schema before shuffle/exchange code consumes it. + * because Spark serializes state to binary. Native Comet emits the actual state type, so fix + * the exposed output schema before shuffle/exchange code consumes it. * * NOTE: If a new TypedImperativeAggregate function is added natively, add a case branch here * mapping it to the native state type. From 5ed285fcd5124689c4aaffce8b85baee590fa0c4 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 9 Jul 2026 17:02:29 -0600 Subject: [PATCH 6/6] fix: coerce collect_list/collect_set nested field nullability 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. --- native/core/src/execution/planner.rs | 70 +++++++++++++++++++ .../comet/exec/CometAggregateSuite.scala | 32 +++++++++ 2 files changed, 102 insertions(+) diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index c5438b928d..d0cd8b2c1c 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -155,6 +155,53 @@ struct JoinParameters { pub join_type: DFJoinType, } +/// Return a copy of `data_type` with every nested field marked nullable. Map key fields are left +/// non-nullable to preserve Arrow's map invariant. Primitive types are returned unchanged. +fn make_all_fields_nullable(data_type: &DataType) -> DataType { + fn nullable_field(field: &Field, nullable: bool) -> FieldRef { + Arc::new( + Field::new( + field.name(), + make_all_fields_nullable(field.data_type()), + nullable, + ) + .with_metadata(field.metadata().clone()), + ) + } + match data_type { + DataType::Struct(fields) => { + DataType::Struct(fields.iter().map(|f| nullable_field(f, true)).collect()) + } + DataType::List(field) => DataType::List(nullable_field(field, true)), + DataType::LargeList(field) => DataType::LargeList(nullable_field(field, true)), + DataType::FixedSizeList(field, len) => { + DataType::FixedSizeList(nullable_field(field, true), *len) + } + DataType::Map(entries, sorted) => { + // Map entries are a non-nullable struct of {key (non-null), value}. Recurse into the + // key/value types but keep the key field non-nullable per Arrow's map invariant. + match entries.data_type() { + DataType::Struct(kv) => { + let new_kv = kv + .iter() + .enumerate() + .map(|(i, f)| nullable_field(f, i != 0)) + .collect(); + let new_entries = Field::new( + entries.name(), + DataType::Struct(new_kv), + entries.is_nullable(), + ) + .with_metadata(entries.metadata().clone()); + DataType::Map(Arc::new(new_entries), *sorted) + } + _ => data_type.clone(), + } + } + other => other.clone(), + } +} + /// If `expr` evaluates to `Timestamp(_, Some(_))` against `schema`, wrap it in a /// metadata-only cast to `Timestamp(_, None)`. This is required because /// DataFusion's `SortMergeJoinExec` comparator only supports timezone-less @@ -2662,11 +2709,13 @@ impl PhysicalPlanner { } AggExprStruct::CollectSet(expr) => { let child = self.create_expr(expr.child.as_ref().unwrap(), Arc::clone(&schema))?; + let child = Self::coerce_collect_child_nullability(child, &schema)?; let func = AggregateUDF::new_from_impl(SparkCollectSet::new()); Self::create_aggr_func_expr("collect_set", schema, vec![child], func) } AggExprStruct::CollectList(expr) => { let child = self.create_expr(expr.child.as_ref().unwrap(), Arc::clone(&schema))?; + let child = Self::coerce_collect_child_nullability(child, &schema)?; let func = AggregateUDF::new_from_impl(SparkCollectList::new()); Self::create_aggr_func_expr("collect_list", schema, vec![child], func) } @@ -3274,6 +3323,27 @@ impl PhysicalPlanner { Ok(scalar_expr) } + /// `collect_list` / `collect_set` build their result list with all element fields marked + /// nullable, regardless of the input's nullability. However `SparkCollectList::return_type` + /// (and `SparkCollectSet`) derive the element type directly from the child, preserving any + /// non-nullable nested field. When the child is a nested type with a non-nullable inner field + /// (e.g. a struct field), the declared aggregate output disagrees with the array the + /// accumulator actually produces, and DataFusion's grouped `AggregateExec` fails validating + /// its output batch ("column types must match schema types"). Cast the child to the + /// all-nullable variant of its type so the declared and produced types stay consistent. + fn coerce_collect_child_nullability( + child: Arc, + schema: &SchemaRef, + ) -> Result, ExecutionError> { + let child_type = child.data_type(schema.as_ref())?; + let nullable_type = make_all_fields_nullable(&child_type); + if child_type.equals_datatype(&nullable_type) { + Ok(child) + } else { + Ok(Arc::new(CastExpr::new(child, nullable_type, None))) + } + } + fn create_aggr_func_expr( name: &str, schema: SchemaRef, diff --git a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala index 8675d60ab3..3248101a2b 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala @@ -48,6 +48,38 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { override protected def sparkConf: SparkConf = super.sparkConf.set(SQLConf.ANSI_ENABLED.key, "false") + test("collect_list over struct with non-nullable fields") { + // Building a struct from non-nullable columns yields non-nullable struct fields. The native + // collect_list accumulator emits a list whose element fields are all nullable, so the produced + // array must be reconciled with the declared aggregate output type, otherwise the grouped + // native aggregate fails with "column types must match schema types". + import org.apache.spark.sql.functions.{collect_list, expr} + import org.apache.spark.sql.execution.LocalTableScanExec + // One row per (a, b) group keeps each collected list deterministic for the answer comparison. + // repartition inserts a Comet exchange so the aggregate runs natively over a Comet child; the + // in-memory source stays a (non-Comet) LocalTableScan, which we allow via excludedClasses. + val df = Seq(("1", "2", 1), ("2", "3", 3)) + .toDF("a", "b", "c") + .repartition(col("a")) + .withColumn("d", expr("named_struct('a', a, 'b', b, 'c', c)")) + val include = Seq(classOf[CometHashAggregateExec]) + // Single grouped collect_list of a struct with a non-nullable field. + checkSparkAnswerAndOperator( + df.groupBy("a", "b").agg(collect_list("d")), + include, + classOf[LocalTableScanExec]) + // Multi-stage: the array result flows through a projection into a second collect_list + // (the SPARK-22223 plan shape), so the nested struct field nullability must round-trip. + checkSparkAnswerAndOperator( + df.groupBy("a", "b") + .agg(collect_list("d").as("e")) + .withColumn("f", expr("named_struct('b', b, 'e', e)")) + .groupBy("a") + .agg(collect_list("f")), + include, + classOf[LocalTableScanExec]) + } + test("collect_list/collect_set combined with distinct aggregate falls back safely") { // SPARK-17616: a distinct aggregate combined with collect_list/collect_set produces a // multi-stage plan where the buffer-producing Partial may run in Spark (e.g. over a