From 72cd19892c2895670c45c555409cb4b16f4d5a5d Mon Sep 17 00:00:00 2001 From: Erik Bogado Date: Thu, 22 Oct 2026 21:37:41 -0300 Subject: [PATCH 1/2] fix: align nested collect buffer nullability collect_list and collect_set declare their intermediate buffer as array, and the native accumulator emits the all-nullable form of the argument. A struct or array argument built after the Parquet scan can still carry required nested fields, so the declared buffer and the emitted state disagreed on nested nullability and the group aggregate failed with "column types must match schema types" when it materialized that state. Apply Spark's recursive asNullable to the element type for both functions, matching the native planner's normalization, and keep one branch that mirrors it. Cover the ungrouped Accumulator path in the collect_list and collect_set SQL fixtures, which assert native execution by default, and drop the grouped test's spill assertion: neither stage spills on that fixture, and its old sum-over-all-aggregates assertion does not hold on current main. --- .../apache/spark/sql/comet/operators.scala | 9 ++--- .../expressions/aggregate/collect_list.sql | 22 ++++++++++++ .../expressions/aggregate/collect_set.sql | 22 ++++++++++++ .../comet/exec/CometAggregateSuite.scala | 36 +++++++++++++++++++ 4 files changed, 83 insertions(+), 6 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 ee5c5358e73..52bca909f4c 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..ffe9f99698f 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)" + // The four settings below are load-bearing for the plan that reproduces the mismatch: + // without them Spark keeps AQE and the default shuffle partition count, and this query + // passes even with the strict buffer declaration. They 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 From 2add60f845522b01a743da2fcc2edda91e8b5155 Mon Sep 17 00:00:00 2001 From: Erik Bogado Date: Mon, 26 Oct 2026 10:46:24 -0300 Subject: [PATCH 2/2] test(aggregate): clarify why the collect regression reuses the helper --- .../org/apache/comet/exec/CometAggregateSuite.scala | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) 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 ffe9f99698f..5c73433a552 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala @@ -379,14 +379,14 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { // 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)" - // The four settings below are load-bearing for the plan that reproduces the mismatch: - // without them Spark keeps AQE and the default shuffle partition count, and this query - // passes even with the strict buffer declaration. They are not here to force spilling — + // 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. + // "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"""