[SPARK-58294][SQL] Fix Hive UDAF with two aggregation buffers in Complete mode#57464
[SPARK-58294][SQL] Fix Hive UDAF with two aggregation buffers in Complete mode#57464ulysses-you wants to merge 1 commit into
Conversation
…lete mode A Hive UDAF whose evaluator uses distinct aggregation-buffer classes per mode (e.g. `MockUDAF2`) fails with a `ClassCastException` when planned in `Complete` mode, which `CombineAdjacentAggregation` produces by merging an adjacent partial/final pair. In Complete mode `update` is called but `merge` is not, so `eval` receives a PARTIAL1-mode buffer and hands it directly to the FINAL evaluator's `terminate`, which expects a FINAL-mode buffer. Extract the PARTIAL1 -> FINAL buffer conversion that `merge` already performs into a `toFinalBuffer` helper and apply it in `eval` too, so the Complete-mode path terminates on a FINAL-mode buffer. Extracted from apache#57363 as a standalone correctness fix: the Complete-mode path is reachable today by setting `spark.sql.execution.combineAdjacentAggregation` to `true`, independent of the `replaceHashWithSortAgg` default flip, so it can be reviewed and backported on its own. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
sunchao
left a comment
There was a problem hiding this comment.
Summary
This is a well-scoped correctness fix for mode-aware Hive UDAFs under Spark Complete-mode aggregation. I verified the evaluator lifecycle, existing partial/final and spill behavior, both actual sort-fallback paths, and the exact-commit Hive CI. The implementation fixes the underlying buffer-mode mismatch without changing the normal aggregation path.
Prior state and problem
CombineAdjacentAggregation can replace an adjacent Partial/Final pair with one Complete aggregate. In that plan, HiveUDAFFunction.update creates a PARTIAL1 buffer, but eval terminates through a separately initialized FINAL evaluator. Valid Hive UDAFs that use different aggregation-buffer classes for those modes consequently fail with ClassCastException, even though the original two-stage plan succeeds.
Design approach
Factor the already-established PARTIAL1-to-FINAL conversion out of merge and reuse that same conversion before final evaluation. This fixes the evaluator lifecycle at the Hive integration boundary rather than excluding Hive UDAFs from the physical optimization or changing the optimizer behavior.
Correctness / compatibility analysis
A buffer that has already entered the merge-capable FINAL state is returned unchanged, preserving existing partial/final execution. An update-only buffer is materialized into a fresh FINAL buffer using terminatePartial and final-evaluator merge before terminate. The existing empty-buffer branch is preserved, and the regression confirms that object-hash aggregation is correct both when it remains in memory and when it falls back to sort-based aggregation.
Key design decisions
Reusing the exact conversion that previously lived in merge keeps the existing mode-specific evaluator and object-inspector contract intact. The new regression explicitly verifies that the physical plan contains a single Complete-mode ObjectHashAggregateExec, preventing a passing test from accidentally exercising the original two-stage plan.
Implementation sketch
toFinalBuffer returns merge-ready buffers directly and converts update-only buffers through terminatePartial -> merge. Both merge and eval call the same helper. The regression constructs a fresh DataFrame under each fallback threshold and verifies the actual numTasksFallBacked metric, so both configured paths genuinely execute instead of reusing a memoized execution.
Behavioral changes worth calling out
Mode-sensitive Hive UDAFs now succeed when adjacent aggregation is combined into Complete mode. Existing Partial/Final aggregation, empty-input handling, and the pre-existing SPARK-24935 coverage remain intact. The exact-head full CI has 26 successful jobs and three expected skips; the slow Hive job explicitly ran the new SPARK-58294 regression and reports 2,246 passing tests with zero failures.
Suggested improvements
No changes requested. The implementation and regression coverage are sufficient for this scoped correctness fix.
What changes were proposed in this pull request?
A Hive UDAF whose
GenericUDAFEvaluatorlegally uses different aggregation-buffer classes per mode (one for consuming raw input inPARTIAL1, another for merging partial buffers inFINAL-- asMockUDAF2does in the test) throws aClassCastExceptionwhen the aggregate is planned inCompletemode.Completemode is produced byCombineAdjacentAggregation, which merges an adjacent partial/final pair (no shuffle between them) into a single complete-mode aggregate. InCompletemode,updateis called butmergeis not, so the buffer reachingevalis aPARTIAL1-mode buffer.evalpreviously passedbuffer.bufstraight to theFINALevaluator'sterminate, which expects aFINAL-mode buffer, causing the cast failure.mergealready performs an on-demandPARTIAL1 -> FINALbuffer conversion for the same reason. This PR extracts that logic into atoFinalBufferhelper and applies it inevalas well, so theComplete-mode path terminates on aFINAL-mode buffer.This is extracted from #57363 as a standalone correctness fix. The
Complete-mode path is reachable today by settingspark.sql.execution.combineAdjacentAggregationtotrue(introduced in SPARK-43317), independent of thespark.sql.execution.replaceHashWithSortAggdefault 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
ClassCastExceptionwhenever the optimizer combines its partial/final aggregates intoCompletemode.Does this PR introduce any user-facing change?
Yes. A Hive UDAF that previously threw a
ClassCastExceptionunderspark.sql.execution.combineAdjacentAggregation=truenow evaluates correctly.How was this patch tested?
New test
SPARK-58294: Hive UDAF with two aggregation buffers in Complete modeinHiveUDAFSuite, using the existingMockUDAF2(distinct buffer classes per mode). It enablescombineAdjacentAggregation, asserts the plan is a singleComplete-modeObjectHashAggregateExec, and checks the result under both the sort-based fallback path (OBJECT_AGG_SORT_BASED_FALLBACK_THRESHOLD = 1, assertingnumTasksFallBacked > 0) and the non-fallback path (= 100, asserting== 0).Mutation-tested: reverting the
evalfix (passingbuffer.bufinstead oftoFinalBuffer(buffer).buf) fails this test with the exactClassCastException: MockUDAFBuffer cannot be cast to MockUDAFBuffer2; it passes with the fix. The existingSPARK-24935test (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)