[BUG] Fail closed on incomplete aggregate result mapping

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
   }
 }
 

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.