[SPARK-58294][SQL] Fix Hive UDAF with two aggregation buffers in Complete mode

sql/hive/src/main/scala/org/apache/spark/sql/hive/hiveUDFs.scala

@@ -454,14 +454,7 @@ private[hive] case class HiveUDAFFunction(
     // implementation can't mix the `update` and `merge` calls during its life cycle. To work
     // around it, here we create a fresh buffer with final evaluator, and merge the existing buffer
     // to it, and replace the existing buffer with it.
-    val mergeableBuf = if (!nonNullBuffer.canDoMerge) {
-      val newBuf = finalHiveEvaluator.evaluator.getNewAggregationBuffer
-      finalHiveEvaluator.evaluator.merge(
-        newBuf, partial1HiveEvaluator.evaluator.terminatePartial(nonNullBuffer.buf))
-      HiveUDAFBuffer(newBuf, true)
-    } else {
-      nonNullBuffer
-    }
+    val mergeableBuf = toFinalBuffer(nonNullBuffer)
 
     // The 2nd argument of the Hive `GenericUDAFEvaluator.merge()` method is an input aggregation
     // buffer in the 3rd format mentioned in the ScalaDoc of this class. Originally, Hive converts
@@ -472,12 +465,30 @@ private[hive] case class HiveUDAFFunction(
     mergeableBuf
   }
 
+  // Convert a buffer to a FINAL-mode buffer if it is a PARTIAL1-mode buffer created by `update`
+  // (or deserialized). Some Hive UDAFs legally use different buffer classes per mode, so a
+  // PARTIAL1 buffer cannot be passed directly to the FINAL evaluator's `terminate`. This mirrors
+  // the on-demand conversion `merge` performs, and is also needed by `eval` for the Complete-mode
+  // path (enabled by `CombineAdjacentAggregation`), where `update` is called but `merge` is not.
+  private def toFinalBuffer(buffer: HiveUDAFBuffer): HiveUDAFBuffer = {
+    if (!buffer.canDoMerge) {
+      val newBuf = finalHiveEvaluator.evaluator.getNewAggregationBuffer
+      finalHiveEvaluator.evaluator.merge(
+        newBuf, partial1HiveEvaluator.evaluator.terminatePartial(buffer.buf))
+      HiveUDAFBuffer(newBuf, true)
+    } else {
+      buffer
+    }
+  }
+
   override def eval(buffer: HiveUDAFBuffer): Any = {
     resultUnwrapper(finalHiveEvaluator.evaluator.terminate(
       if (buffer == null) {
         finalHiveEvaluator.evaluator.getNewAggregationBuffer
       } else {
-        buffer.buf
+        // A buffer that only went through `update` (Complete mode) is a PARTIAL1-mode buffer; turn
+        // it into a FINAL-mode buffer before terminating so mode-aware UDAFs see the right class.
+        toFinalBuffer(buffer).buf
       }
     ))
   }

sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/HiveUDAFSuite.scala

@@ -29,8 +29,9 @@ import org.apache.hadoop.hive.serde2.typeinfo.TypeInfo
 import test.org.apache.spark.sql.MyDoubleAvg
 
 import org.apache.spark.SPARK_DOC_ROOT
-import org.apache.spark.sql.{AnalysisException, QueryTest, Row}
+import org.apache.spark.sql.{AnalysisException, DataFrame, QueryTest, Row}
 import org.apache.spark.sql.catalyst.expressions.Cast._
+import org.apache.spark.sql.catalyst.expressions.aggregate.Complete
 import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
 import org.apache.spark.sql.execution.aggregate.ObjectHashAggregateExec
 import org.apache.spark.sql.hive.test.TestHiveSingleton
@@ -42,6 +43,14 @@ class HiveUDAFSuite extends QueryTest
   with TestHiveSingleton with AdaptiveSparkPlanHelper {
   import testImplicits._
 
+  // Sums the `numTasksFallBacked` metric across every `ObjectHashAggregateExec` in the executed
+  // plan. The DataFrame must be executed (e.g. via `checkAnswer`) before this is read.
+  private def numFallbackTasks(df: DataFrame): Long = {
+    collect(df.queryExecution.executedPlan) {
+      case agg: ObjectHashAggregateExec => agg.metrics("numTasksFallBacked").value
+    }.sum
+  }
+
   protected override def beforeAll(): Unit = {
     super.beforeAll()
     sql(s"CREATE TEMPORARY FUNCTION mock AS '${classOf[MockUDAF].getName}'")
@@ -128,6 +137,52 @@ class HiveUDAFSuite extends QueryTest
     }
   }
 
+  test("SPARK-58294: Hive UDAF with two aggregation buffers in Complete mode") {
+    withTempView("v") {
+      spark.range(100).createTempView("v")
+      // With `CombineAdjacentAggregation` enabled and a single input partition, the adjacent
+      // partial/final pair is merged into a single `Complete`-mode `ObjectHashAggregateExec`.
+      // `MockUDAF2` deliberately uses distinct aggregation-buffer classes per mode (one for
+      // consuming original input via PARTIAL1, another for merging partial buffers via FINAL).
+      // The Hive UDAF wrapper must convert the PARTIAL1 buffer to a FINAL buffer on `eval` so the
+      // Complete-mode path (which calls `update` but not `merge`) terminates correctly instead of
+      // handing a PARTIAL1 buffer to the FINAL evaluator.
+      val query = "SELECT id % 2, mock2(id) FROM v GROUP BY id % 2"
+      withSQLConf(SQLConf.COMBINE_ADJACENT_AGGREGATION_ENABLED.key -> "true") {
+        val aggs = collect(sql(query).queryExecution.executedPlan) {
+          case agg: ObjectHashAggregateExec => agg
+        }
+        // Combined into a single `Complete`-mode aggregate.
+        assert(aggs.length == 1)
+        assert(aggs.head.aggregateExpressions.forall(_.mode == Complete))
+
+        // Rebuild the query inside each `withSQLConf` block: `SparkPlan.executeRDD` memoizes the
+        // RDD, so reusing one DataFrame across thresholds would replay the first execution instead
+        // of re-planning under the new threshold, exercising only one of the two paths.
+        // Threshold 1 forces every task with more than one group to fall back to sort-based
+        // aggregation (there are two groups: `id % 2` in {0, 1}).
+        withSQLConf(SQLConf.OBJECT_AGG_SORT_BASED_FALLBACK_THRESHOLD.key -> "1") {
+          val df = sql(query)
+          checkAnswer(df, Seq(
+            Row(0, Row(50, 0)),
+            Row(1, Row(50, 0))
+          ))
+          assert(numFallbackTasks(df) > 0, "threshold 1 must trigger the sort-based fallback path")
+        }
+
+        // Threshold 100 is above the two groups, so no task falls back.
+        withSQLConf(SQLConf.OBJECT_AGG_SORT_BASED_FALLBACK_THRESHOLD.key -> "100") {
+          val df = sql(query)
+          checkAnswer(df, Seq(
+            Row(0, Row(50, 0)),
+            Row(1, Row(50, 0))
+          ))
+          assert(numFallbackTasks(df) == 0, "threshold 100 must not trigger the fallback path")
+        }
+      }
+    }
+  }
+
   test("call JAVA UDAF") {
     withTempView("temp") {
       withUserDefinedFunction("myDoubleAvg" -> false) {