From 424f5a96640bc80504f9f29225fb13f1fd93ccf1 Mon Sep 17 00:00:00 2001 From: Erik Bogado Date: Sun, 13 Sep 2026 12:10:05 -0300 Subject: [PATCH 1/3] test: reproduce struct collection spill failure Force final aggregation to spill so nested nullability mismatches cannot hide behind shuffle-only spills. Refs #5239 --- .../comet/exec/CometAggregateSuite.scala | 37 +++++++++++++++++++ 1 file changed, 37 insertions(+) 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 4b3e7f4705..646aebeb10 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala @@ -185,6 +185,43 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { classOf[LocalTableScanExec]) } + test("collect_list and collect_set over non-nullable nested fields survive spilling") { + withTempPath { path => + // Keep every group in one input and final partition so aggregation outgrows its memory pool. + 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") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.SHUFFLE_PARTITIONS.key -> "1", + CometConf.COMET_BATCH_SIZE.key -> "1024", + CometConf.COMET_OFFHEAP_MEMORY_POOL_FRACTION.key -> "0.002", + CometConf.COMET_RESPECT_DATAFUSION_CONFIGS.key -> "true", + "spark.comet.datafusion.execution.sort_spill_reservation_bytes" -> "65536") { + // The literal and coalesce create required fields after the Parquet scan, which + // otherwise widens field nullability. Sorting removes collect order differences. + val (_, cometPlan) = checkSparkAnswerAndOperator( + sql(""" + SELECT _1, sort_array(collect_list(s)), sort_array(collect_set(s)) + FROM ( + SELECT _1, named_struct('flag', true, 'id', coalesce(_2, 0), 'value', _3) AS s + FROM tbl + ) GROUP BY _1"""), + Seq(classOf[CometHashAggregateExec])) + val aggregates = collect(cometPlan) { case agg: CometHashAggregateExec => agg } + assert( + aggregates.map(_.metrics("spill_count").value).sum > 0L, + s"Expected the native collection aggregate to spill:\n$cometPlan") + } + } + } + } + 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 aab4e339e5dd036cdc33befd8ead84d522973d14 Mon Sep 17 00:00:00 2001 From: Erik Bogado Date: Sun, 13 Sep 2026 15:25:45 -0300 Subject: [PATCH 2/3] fix: align nested collection buffer nullability Match native collection state so final aggregation can spill structs with required fields. Closes #5239 --- .../main/scala/org/apache/spark/sql/comet/operators.scala | 5 +++-- 1 file changed, 3 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 dfb362595a..f0987fa46b 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 @@ -1903,11 +1903,12 @@ trait CometBaseAggregate { val bufferAttrs = aggFunc.aggBufferAttributes aggFunc match { case cs: CollectSet => - val elementType = cs.children.head.dataType + // Match the native planner's collect argument normalization, including nested fields. + val elementType = cs.children.head.dataType.asNullable val nativeStateType = ArrayType(elementType, containsNull = true) output(bufferIdx) = output(bufferIdx).withDataType(nativeStateType) case cl: CollectList => - val elementType = cl.children.head.dataType + val elementType = cl.children.head.dataType.asNullable val nativeStateType = ArrayType(elementType, containsNull = true) output(bufferIdx) = output(bufferIdx).withDataType(nativeStateType) case _: Percentile => From 6aa8ca8563818400c285640168082e893916b20b Mon Sep 17 00:00:00 2001 From: Erik Bogado Date: Sun, 13 Sep 2026 16:03:34 -0300 Subject: [PATCH 3/3] test: guard recursive collection nullability Flat struct coverage misses nullability below array elements. --- .../comet/exec/CometAggregateSuite.scala | 27 ++++++++++--------- 1 file changed, 15 insertions(+), 12 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 a7c4b8f97f..1c461ca985 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala @@ -205,18 +205,21 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { "spark.comet.datafusion.execution.sort_spill_reservation_bytes" -> "65536") { // The literal and coalesce create required fields after the Parquet scan, which // otherwise widens field nullability. Sorting removes collect order differences. - val (_, cometPlan) = checkSparkAnswerAndOperator( - sql(""" - SELECT _1, sort_array(collect_list(s)), sort_array(collect_set(s)) - FROM ( - SELECT _1, named_struct('flag', true, 'id', coalesce(_2, 0), 'value', _3) AS s - FROM tbl - ) GROUP BY _1"""), - Seq(classOf[CometHashAggregateExec])) - val aggregates = collect(cometPlan) { case agg: CometHashAggregateExec => agg } - assert( - aggregates.map(_.metrics("spill_count").value).sum > 0L, - s"Expected the native collection aggregate to spill:\n$cometPlan") + val struct = "named_struct('flag', true, 'id', coalesce(_2, 0), 'value', _3)" + Seq(struct, s"array($struct)").foreach { value => + val (_, cometPlan) = checkSparkAnswerAndOperator( + 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"""), + Seq(classOf[CometHashAggregateExec])) + val aggregates = collect(cometPlan) { case agg: CometHashAggregateExec => agg } + assert( + aggregates.map(_.metrics("spill_count").value).sum > 0L, + s"Expected the native collection aggregate to spill:\n$cometPlan") + } } } }