diff --git a/sql/hive/src/main/scala/org/apache/spark/sql/hive/hiveUDFs.scala b/sql/hive/src/main/scala/org/apache/spark/sql/hive/hiveUDFs.scala index 129c5e2cc053b..e5db4d9c7384a 100644 --- a/sql/hive/src/main/scala/org/apache/spark/sql/hive/hiveUDFs.scala +++ b/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 } )) } diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/HiveUDAFSuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/HiveUDAFSuite.scala index 16ae456c41bf5..45ec91ddeff27 100644 --- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/HiveUDAFSuite.scala +++ b/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) {