From c54c01326b3269e830f68ab9d3c318ae005a96c4 Mon Sep 17 00:00:00 2001 From: Zach Puller Date: Fri, 11 Sep 2026 14:45:32 -0700 Subject: [PATCH] [BUG] Fail closed on incomplete aggregate result mapping Signed-off-by: Zach Puller --- .../spark/rapids/GpuAggregateExec.scala | 22 ++++++++++---- .../spark/rapids/HashAggregatesSuite.scala | 29 +++++++++++++++++-- 2 files changed, 43 insertions(+), 8 deletions(-) diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuAggregateExec.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuAggregateExec.scala index 9afdea0b1df..86d6551d574 100644 --- a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuAggregateExec.scala +++ b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuAggregateExec.scala @@ -1424,8 +1424,10 @@ abstract class GpuTypedImperativeSupportedAggregateExecMeta[INPUT <: BaseAggrega (expr.mode == Partial || expr.mode == PartialMerge) } - // overriding data types of Aggregation Buffers if necessary - if (mayNeedAggBufferConversion) overrideAggBufTypes() + // Overriding data types happens before metadata tagging initializes replacement reasons, so + // retain any mapping failure and report it from tagPlanForGpu. + private val aggBufTypeOverrideFailure: Option[String] = + if (mayNeedAggBufferConversion) overrideAggBufTypes() else None override protected lazy val outputTypeMetas: Option[Seq[DataTypeMeta]] = if (mayNeedAggBufferConversion) { @@ -1442,6 +1444,7 @@ abstract class GpuTypedImperativeSupportedAggregateExecMeta[INPUT <: BaseAggrega override def tagPlanForGpu(): Unit = { super.tagPlanForGpu() + aggBufTypeOverrideFailure.foreach(willNotWorkOnGpu) // If a typedImperativeAggregate function run across CPU and GPU (ex: Partial mode on CPU, // Merge mode on GPU), it will lead to a runtime crash. Because aggregation buffers produced @@ -1507,10 +1510,11 @@ abstract class GpuTypedImperativeSupportedAggregateExecMeta[INPUT <: BaseAggrega * At last, we traverse aggregateAttributes and resultExpressions, overriding data type in * RapidsMeta if necessary, in order to ensure TypeChecks tagging exact data types in runtime. */ - private def overrideAggBufTypes(): Unit = { + private def overrideAggBufTypes(): Option[String] = { val desiredAggBufTypes = mutable.HashMap.empty[ExprId, DataType] val desiredInputAggBufTypes = mutable.HashMap.empty[ExprId, DataType] val desiredResultOutputTypes = mutable.HashMap.empty[ExprId, DataType] + var mappingFailure: Option[String] = None // Collects exprId from TypedImperativeAggBufferAttributes, and maps them to the data type // of `TypedImperativeAggExprMeta.aggBufferAttribute`. aggregateExpressions.map(_.childExprs.head).foreach { @@ -1539,9 +1543,14 @@ abstract class GpuTypedImperativeSupportedAggregateExecMeta[INPUT <: BaseAggrega val bufferCount = aggExpr.aggregateFunction.inputAggBufferAttributes.length aggExprMeta.childExprs.head match { case aggMeta: TypedImperativeAggExprMeta[_] => - resultExpressions.lift(resultOffset).foreach { resultMeta => - val resultExpr = resultMeta.wrapped.asInstanceOf[NamedExpression] - desiredResultOutputTypes(resultExpr.exprId) = aggMeta.aggBufferAttribute.dataType + resultExpressions.lift(resultOffset) match { + case Some(resultMeta) => + val resultExpr = resultMeta.wrapped.asInstanceOf[NamedExpression] + desiredResultOutputTypes(resultExpr.exprId) = aggMeta.aggBufferAttribute.dataType + case None => + mappingFailure = Some( + s"Typed imperative aggregate buffer result at offset $resultOffset is " + + s"missing from ${resultExpressions.length} result expressions") } case _ => } @@ -1567,6 +1576,7 @@ abstract class GpuTypedImperativeSupportedAggregateExecMeta[INPUT <: BaseAggrega case _ => } } + mappingFailure } } diff --git a/tests/src/test/scala/com/nvidia/spark/rapids/HashAggregatesSuite.scala b/tests/src/test/scala/com/nvidia/spark/rapids/HashAggregatesSuite.scala index c0a4ed6831b..081e96bcb56 100644 --- a/tests/src/test/scala/com/nvidia/spark/rapids/HashAggregatesSuite.scala +++ b/tests/src/test/scala/com/nvidia/spark/rapids/HashAggregatesSuite.scala @@ -26,10 +26,10 @@ import org.apache.spark.sql.AnalysisException import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, AttributeReference, Expression, ExprId, Literal} -import org.apache.spark.sql.catalyst.expressions.aggregate.Final +import org.apache.spark.sql.catalyst.expressions.aggregate.{Final, Partial} import org.apache.spark.sql.execution.{LeafExecNode, SparkPlan, WholeStageCodegenExec} import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, BroadcastQueryStageExec, QueryStageExec, ShuffleQueryStageExec} -import org.apache.spark.sql.execution.aggregate.SortAggregateExec +import org.apache.spark.sql.execution.aggregate.{ObjectHashAggregateExec, SortAggregateExec} import org.apache.spark.sql.functions._ import org.apache.spark.sql.rapids.ExecutionPlanCaptureCallback import org.apache.spark.sql.rapids.aggregate.{CudfAggregate, GpuAggregateExpression, @@ -41,6 +41,31 @@ class HashAggregatesSuite extends SparkQueryCompareTestSuite { private def floatAggConf: SparkConf = enableCsvConf() .set(RapidsConf.ENABLE_FLOAT_AGG.key, "true") + test("incomplete positional typed aggregate result mapping fails closed") { + withCpuSparkSession({ spark => + val plan = spark.range(10) + .groupBy((col("id") % 2).alias("key")) + .agg(collect_set(col("id")).alias("values")) + .queryExecution.sparkPlan + + val partialAgg = plan.collectFirst { + case agg: ObjectHashAggregateExec + if agg.aggregateExpressions.exists(_.mode == Partial) => agg + }.getOrElse(fail(s"Expected a partial ObjectHashAggregateExec in:\n$plan")) + + // Simulate the inconsistent positional layout from #15808: the typed aggregate still + // needs its shuffle-facing result type remapped, but that result expression is absent. + val malformedAgg = partialAgg.copy( + resultExpressions = partialAgg.resultExpressions.dropRight(1)) + val meta = GpuOverrides.wrapPlan( + malformedAgg, new RapidsConf(Map.empty[String, String]), None) + meta.tagForGpu() + + assert(!meta.canThisBeReplaced, + "An unresolved positional typed aggregate result must prevent GPU conversion") + }, new SparkConf().set("spark.sql.adaptive.enabled", "false")) + } + test("SPARK-55979: GPU aggregate references retain renamed input buffer attributes") { val scanAggBufferAttr = AttributeReference("buf", IntegerType, nullable = true)() // withName preserves exprId while changing the display name, matching the Spark regression.