You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by ul...@apache.org on 2026/07/24 01:34:48 UTC
(spark) branch branch-4.x updated: [SPARK-58294][SQL] Fix Hive UDAF with two aggregation buffers in Complete mode
This is an automated email from the ASF dual-hosted git repository.
ulysses-you pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new 0d07c87ac108 [SPARK-58294][SQL] Fix Hive UDAF with two aggregation buffers in Complete mode
0d07c87ac108 is described below
commit 0d07c87ac1086fe2524dc806b3a1fe662cec7fa2
Author: Xiduo You <ul...@gmail.com>
AuthorDate: Fri Jul 24 09:33:59 2026 +0800
[SPARK-58294][SQL] Fix Hive UDAF with two aggregation buffers in Complete mode
### What changes were proposed in this pull request?
A Hive UDAF whose `GenericUDAFEvaluator` legally uses different aggregation-buffer classes per mode (one for consuming raw input in `PARTIAL1`, another for merging partial buffers in `FINAL` -- as `MockUDAF2` does in the test) throws a `ClassCastException` when the aggregate is planned in `Complete` mode.
`Complete` mode is produced by `CombineAdjacentAggregation`, which merges an adjacent partial/final pair (no shuffle between them) into a single complete-mode aggregate. In `Complete` mode, `update` is called but `merge` is not, so the buffer reaching `eval` is a `PARTIAL1`-mode buffer. `eval` previously passed `buffer.buf` straight to the `FINAL` evaluator's `terminate`, which expects a `FINAL`-mode buffer, causing the cast failure.
`merge` already performs an on-demand `PARTIAL1 -> FINAL` buffer conversion for the same reason. This PR extracts that logic into a `toFinalBuffer` helper and applies it in `eval` as well, so the `Complete`-mode path terminates on a `FINAL`-mode buffer.
This is extracted from #57363 as a standalone correctness fix. The `Complete`-mode path is reachable today by setting `spark.sql.execution.combineAdjacentAggregation` to `true` (introduced in SPARK-43317), independent of the `spark.sql.execution.replaceHashWithSortAgg` default flip proposed there, so it can be reviewed, merged, and backported on its own.
### Why are the changes needed?
Without the fix, a valid Hive UDAF that uses mode-specific buffer classes fails at runtime with a `ClassCastException` whenever the optimizer combines its partial/final aggregates into `Complete` mode.
### Does this PR introduce _any_ user-facing change?
Yes. A Hive UDAF that previously threw a `ClassCastException` under `spark.sql.execution.combineAdjacentAggregation=true` now evaluates correctly.
### How was this patch tested?
New test `SPARK-58294: Hive UDAF with two aggregation buffers in Complete mode` in `HiveUDAFSuite`, using the existing `MockUDAF2` (distinct buffer classes per mode). It enables `combineAdjacentAggregation`, asserts the plan is a single `Complete`-mode `ObjectHashAggregateExec`, and checks the result under both the sort-based fallback path (`OBJECT_AGG_SORT_BASED_FALLBACK_THRESHOLD = 1`, asserting `numTasksFallBacked > 0`) and the non-fallback path (`= 100`, asserting `== 0`).
Mutation-tested: reverting the `eval` fix (passing `buffer.buf` instead of `toFinalBuffer(buffer).buf`) fails this test with the exact `ClassCastException: MockUDAFBuffer cannot be cast to MockUDAFBuffer2`; it passes with the fix. The existing `SPARK-24935` test (partial/final path) is left unchanged and still passes.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Claude Opus 4.8)
Closes #57464 from ulysses-you/hive-udaf-complete-mode-fix.
Authored-by: Xiduo You <ul...@gmail.com>
Signed-off-by: Xiduo You <ul...@apache.org>
(cherry picked from commit 6dfcc2ccfd7d53f4250e764fb714b5aad1d65d5b)
Signed-off-by: Xiduo You <ul...@apache.org>
---
.../scala/org/apache/spark/sql/hive/hiveUDFs.scala | 29 +++++++----
.../spark/sql/hive/execution/HiveUDAFSuite.scala | 57 +++++++++++++++++++++-
2 files changed, 76 insertions(+), 10 deletions(-)
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 129c5e2cc053..e5db4d9c7384 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 16ae456c41bf..45ec91ddeff2 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) {
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org