Skip to content

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

Open
ulysses-you wants to merge 1 commit into
apache:masterfrom
ulysses-you:hive-udaf-complete-mode-fix
Open

[SPARK-58294][SQL] Fix Hive UDAF with two aggregation buffers in Complete mode#57464
ulysses-you wants to merge 1 commit into
apache:masterfrom
ulysses-you:hive-udaf-complete-mode-fix

Conversation

@ulysses-you

Copy link
Copy Markdown
Contributor

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)

…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 sunchao left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants