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 da26c91f9c7..7704f037f4c 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 @@ -2492,12 +2492,9 @@ trait CometBaseAggregate[T <: BaseAggregateExec] extends CometOperatorSerde[T] { val aggFunc = aggExpr.aggregateFunction val bufferAttrs = aggFunc.aggBufferAttributes aggFunc match { - case cs: CollectSet => - val elementType = cs.children.head.dataType - val nativeStateType = ArrayType(elementType, containsNull = true) - output(bufferIdx) = output(bufferIdx).withDataType(nativeStateType) - case cl: CollectList => - val elementType = cl.children.head.dataType + case _: CollectSet | _: CollectList => + // Match the native planner's collect argument normalization, including nested fields. + val elementType = aggFunc.children.head.dataType.asNullable val nativeStateType = ArrayType(elementType, containsNull = true) output(bufferIdx) = output(bufferIdx).withDataType(nativeStateType) case _: Percentile => 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 index 9bf0f59ad78..10e64c94d6a 100644 --- a/spark/src/test/resources/sql-tests/expressions/aggregate/collect_list.sql +++ b/spark/src/test/resources/sql-tests/expressions/aggregate/collect_list.sql @@ -369,6 +369,28 @@ INSERT INTO cl_src_struct VALUES query SELECT grp, sort_array(collect_list(s)) FROM cl_src_struct GROUP BY grp ORDER BY grp +-- ============================================================ +-- Struct with required fields built after the Parquet scan, ungrouped so the +-- boxed Accumulator builds the final state (issue #5239). +-- ============================================================ + +statement +CREATE TABLE cl_src_struct_required(a bigint, b bigint, c string) USING parquet + +statement +INSERT INTO cl_src_struct_required VALUES + (1, 10, 'a'), (2, 20, 'b'), (3, 30, 'c'), (4, NULL, 'd'), (5, NULL, NULL) + +query +SELECT sort_array(collect_list(s)) FROM ( + SELECT named_struct('flag', true, 'id', coalesce(b, 0), 'value', c) AS s + FROM cl_src_struct_required) + +query +SELECT sort_array(collect_list(s)) FROM ( + SELECT array(named_struct('flag', true, 'id', coalesce(b, 0), 'value', c)) AS s + FROM cl_src_struct_required) + -- ============================================================ -- Nested array input -- ============================================================ diff --git a/spark/src/test/resources/sql-tests/expressions/aggregate/collect_set.sql b/spark/src/test/resources/sql-tests/expressions/aggregate/collect_set.sql index c2f04e37f45..8d7e6df70f8 100644 --- a/spark/src/test/resources/sql-tests/expressions/aggregate/collect_set.sql +++ b/spark/src/test/resources/sql-tests/expressions/aggregate/collect_set.sql @@ -280,3 +280,25 @@ FROM VALUES (1, 1, X'ABCD'), (1, 2, X'ABCD') AS t(x, y, b) GROUP BY x query SELECT grp, sort_array(collect_set(i)) FROM cs_src_int GROUP BY grp HAVING size(collect_set(i)) > 1 ORDER BY grp + +-- ============================================================ +-- Struct with required fields built after the Parquet scan, ungrouped so the +-- boxed Accumulator builds the final state (issue #5239). +-- ============================================================ + +statement +CREATE TABLE cs_src_struct_required(a bigint, b bigint, c string) USING parquet + +statement +INSERT INTO cs_src_struct_required VALUES + (1, 10, 'a'), (2, 20, 'b'), (3, 30, 'c'), (4, NULL, 'd'), (5, NULL, NULL) + +query +SELECT sort_array(collect_set(s)) FROM ( + SELECT named_struct('flag', true, 'id', coalesce(b, 0), 'value', c) AS s + FROM cs_src_struct_required) + +query +SELECT sort_array(collect_set(s)) FROM ( + SELECT array(named_struct('flag', true, 'id', coalesce(b, 0), 'value', c)) AS s + FROM cs_src_struct_required) 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 47eb0203d9e..5c73433a552 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala @@ -366,6 +366,42 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { classOf[LocalTableScanExec]) } + test("grouped collect_list and collect_set over non-nullable nested fields") { + withTempPath { path => + spark + .createDataFrame((0 until 32768).map { i => + (i % 4096, i % 8192, s"${i % 8192}-" + "x" * 256) + }) + .coalesce(1) + .write + .parquet(path.getAbsolutePath) + withParquetTable(path.getAbsolutePath, "tbl") { + // The literal and coalesce create required fields after the Parquet scan, which + // otherwise widens field nullability. Sorting removes collect order differences. + val struct = "named_struct('flag', true, 'id', coalesce(_2, 0), 'value', _3)" + // This reuses the helper for its AQE-off and single-shuffle-partition settings, which are + // what make the query reproduce the mismatch: Spark's defaults insert a shuffle-boundary + // cast that hides it. The helper's memory settings are not here to force spilling — + // neither stage spills on this fixture at any pool size tried (4 MiB, 2 MiB, 400 KiB), and + // shrinking the pool past the Final state's first allocation fails the task with + // "Additional allocation failed for FinalHashAggregateStream[0]" instead of spilling. The + // assertion is therefore native execution plus the answer. The ungrouped Accumulator path + // is covered by collect_list.sql and collect_set.sql. + withDistinctMergeMemoryPressure { + Seq(struct, s"array($struct)").foreach { value => + val df = sql(s""" + SELECT _1, sort_array(collect_list(s)), sort_array(collect_set(s)) + FROM ( + SELECT _1, $value AS s + FROM tbl + ) GROUP BY _1""") + checkSparkAnswerAndOperator(df, Seq(classOf[CometHashAggregateExec])) + } + } + } + } + } + test("grouped collect_list/collect_set over nulls, duplicates and several batches") { // Grouped collect_list/collect_set are served by a native GroupsAccumulator rather than one // boxed accumulator per group, so the group's identity travels with each row instead of being