[cloud-fan]

What changes were proposed in this pull request?

This is a followup of #57742.

When adaptive partial aggregation bypasses the partial aggregation, each passed-through row starts from a fresh initial aggregation buffer. The generated code built that buffer by projecting the initial values again for every passed-through row, inline in the per-row consume method of the aggregation.

This PR builds the initial buffer once per task instead, and resets a reusable buffer from it with UnsafeRow.copyFrom for each passed-through row. This is how the interpreted path already does it (TungstenAggregationIterator.nextPassThroughOutput). Both buffers are registered in doProduceWithKeys, so when the aggregate sits over a fused UnionExec, all of its children share one pair of buffers.

Why are the changes needed?

The per-row projection writes one value per aggregation buffer slot, so its code grows with the number of aggregate functions. It lands in the per-row consume method of every eligible aggregation, although it only runs once pass-through is active, and a fused UnionExec generates one copy of it, with its own row writer, per child. With copyFrom the reset is a single call however many buffer slots there are, and the generated code matches the interpreted path.

Does this PR introduce any user-facing change?

No. It only changes the generated code. Adaptive partial aggregation is new in the unreleased 4.4.0 and disabled by default.

How was this patch tested?

A new test in AdaptivePartialAggregationSuite checks the generated code of an aggregate over a fused two-child UnionExec: each child resets the buffer with copyFrom, and both children reset the same pair of buffers.

The existing tests in the suite exercise the pass-through path with whole-stage codegen on and off, the two-level map on and off, and with and without spills, and compare every result against a run with the feature disabled. build/sbt "sql/testOnly *AdaptivePartialAggregationSuite" passes.

Was this patch authored or co-authored using generative AI tooling?

Generated-by: Claude Code (Claude Opus 5.5)

This pull request and its description were written by Isaac.

[cloud-fan]

cc @ulysses-you

[dongjoon-hyun]

Thank you for the fix, @cloud-fan. I reviewed the latest commit, 0938b6f, and left 6 inline comments. Here is a summary, ordered by priority.

Performance

1. HashAggregateExec L925: each bypassed row still re-projects a constant initial buffer, while the interpreted pass-through path resets the same buffer with one copyFrom (TungstenAggregationIterator.scala:466).
2. HashAggregateExec L926: the extracted makeEmptyAggBuffer still grows with the number of buffer slots, because GenerateUnsafeProjection does not split a top-level projection under whole-stage codegen.
3. HashAggregateExec L927: under a fused UnionExec, the consume code of every child adds its own identical copy of the method and of its UnsafeRowWriter.

Tests

4. HashAggregateExec L928: no test pins the new layout, so a change that inlines the projection again would pass every existing test.

Docs and comments

5. HashAggregateExec L921: the comment and the PR description give no numbers for the JIT-limit motivation, which spark.sql.codegen.maxFields already bounds by default.

Cleanup

6. HashAggregateExec L924: makeEmptyAggBufferFunc departs from the *FuncName naming that this file uses for the results of addNewFunction.

None of these changes the query results, and the generated code is otherwise the same as before. The most important is 1, which would also resolve 2 and most of 3; 4-6 are small.

[dongjoon-hyun]

[1] Each bypassed row still re-projects a constant initial buffer.

The initial values do not depend on the input row, so every call to makeEmptyAggBuffer() writes the same bytes again: rowWriter.reset(), zeroOutNullBytes(), and one write or setNullAt per buffer slot. Once pass-through fires, this runs for every remaining row of the task. The interpreted path already resets the same buffer with a single memory copy, passThroughAggregationBuffer.copyFrom(initialAggregationBuffer) (TungstenAggregationIterator.scala:466), and the generated code already calls getEmptyAggregationBuffer() (L302) to build the fast hash map (L623).

Could the generated code do the same, e.g. keep the row from $thisPlan.getEmptyAggregationBuffer() and a reusable .copy() of it (backed by a byte array, as copyFrom requires) as two mutable states, and reset the copy with copyFrom for each bypassed row? Reusing one row is as safe as reusing the writer row today, since handlePassThroughRow copies the buffer before queuing it. The code size would then no longer depend on the number of buffer slots, which also covers item 2.

[cloud-fan]

Done in bbf464a. The initial buffer is now built once per task, and each passed-through row resets a reusable copy with copyFrom, as TungstenAggregationIterator does. makeEmptyAggBuffer is gone.

[dongjoon-hyun]

[2] The extracted method still grows with the number of buffer slots.

GenerateUnsafeProjection splits the field writes of a top-level projection only when ctx.currentVars is null (GenerateUnsafeProjection.scala:188-189, "TODO: support whole stage codegen"). Here currentVars still holds the consume input, so makeEmptyAggBuffer gets every write in one body, on the order of 10-15 bytes of bytecode per slot. With the default spark.sql.codegen.maxFields this stays small (item 5), but when that limit is raised, a partial aggregate with several hundred buffer slots gets a makeEmptyAggBuffer past the 8000-byte limit of HotSpot. That method then stays interpreted while it runs for every bypassed row, so the hot path only moves from one interpreted method to another.

Could item 1 address this as well? If the projection stays, could the comment at L921-923 say that the method itself is not split?

[cloud-fan]

Fixed by the change for item 1: the reset is now a single copyFrom call, so none of the generated code grows with the number of buffer slots.

[dongjoon-hyun]

[3] Under a fused UnionExec, each child adds its own copy of this method.

doConsumeWithKeys runs once per consume, and UnionExec.doProduce calls produce(ctx, this) once per child (basicPhysicalOperators.scala:1566-1570). So groupBy(...).agg(...) over a fused union of N children generates N identical makeEmptyAggBuffer_<i> methods and N UnsafeRowWriter states, all counting toward the class size. The inline projection was duplicated the same way before this PR, but now that it is a method that reads no input, one copy would do.

Minor: could it be generated once in doProduceWithKeys, before doAgg is generated, and kept in a field like adaptiveOutputMapAndFlushFuncName?

[cloud-fan]

Done in bbf464a. Both buffers are registered in doProduceWithKeys, so the children of a fused Union share them. Each child's consume code keeps only its own one-line copyFrom call.

[dongjoon-hyun]

[4] No test pins the new layout.

AdaptivePartialAggregationSuite compares query results only, and the results were the same before this PR, so a later change that inlines the projection again would pass every test. The path where addNewFunction returns a nestedClassInstance.-qualified name, which the inline code could not reach, is not exercised by this suite either.

Minor: could a test check the generated code of a pass-through-eligible aggregate, e.g. that it contains makeEmptyAggBuffer, or compare the largest method size with the helper at WholeStageCodegenSuite.scala:1849?

[cloud-fan]

Added pass-through resets its buffer with copyFrom from buffers registered once. It checks the generated code of an aggregate over a fused two-child Union: each child resets the buffer with copyFrom, and both use the same buffers.

[dongjoon-hyun]

[5] The JIT-limit motivation has no numbers, and the default config bounds it.

CollapseCodegenStages keeps whole-stage codegen only when the output of an operator has at most spark.sql.codegen.maxFields fields, 100 by default (WholeStageCodegenExec.scala:1214), and a partial aggregate outputs its grouping keys plus its buffer slots (AggUtils.scala:174-176). So by default the projection moved here writes fewer than 100 slots, on the order of 1-1.5KB of bytecode. That can still tip a consume method over 8000 bytes, but neither this comment nor the PR description shows a case that does.

Minor: could the PR description give the largest method size before and after for an example aggregation, or say that the gain matters mainly when maxFields is raised?

[cloud-fan]

Updated the PR description. It no longer claims a JIT-limit effect; it describes the change as a constant-size reset that matches the interpreted path, without one copy per Union child.

[dongjoon-hyun]

[6] makeEmptyAggBufferFunc departs from the *FuncName naming of this file.

The value returned by ctx.addNewFunction can be a nestedClassInstance.-qualified name, and this file names such values doAggFuncName (L799) and adaptiveOutputMapAndFlushFuncName (L776), as SampleExec and RangeExec do with initSamplerFuncName and initRangeFuncName.

Nit: could it be makeEmptyAggBufferFuncName?

[cloud-fan]

No longer applies: the generated method is gone, and the two new fields follow the file's *Term naming for mutable states.

[spark-test]

Review summary

This review found no defects. The change only relocates code. makeEmptyAggBuffer runs the same GenerateUnsafeProjection over declFunctions.flatMap(_.initialValues) that the pass-through branch used to inline, and returns the same class-level writer row. The update and copy-before-queue steps that follow therefore see the same buffer as before. The branch calls the method through the name that ctx.addNewFunction returns, so the call still resolves when the method is placed in a nested class. The method body references only the writer state and literal references, not BufferedRowIterator fields. No admitted initial value reads the consume input: built-in initial values are literals, and the typed aggregate replaces the bound references of its serializer with literals. Moving the projection out of the consume method therefore changes no evaluated value, and the largest generated method can only stay the same size or shrink.

AdaptivePartialAggregationSuite exercises the new method. Each whole-stage-codegen cell asserts that rows were bypassed, which happens only on the branch that calls the method, and then compares the results with a run that has the feature disabled.

The open inline suggestions on this commit (resetting the buffer with copyFrom, the unsplit method, one copy per fused UnionExec child, a layout test, numbers for the JIT motivation, and the *FuncName naming) describe the code accurately, and the author may adopt any of them. None of them is a regression introduced here. The per-row re-projection, the unsplit projection, and the per-child copies already existed at the merge base. The remaining items are a missing test for a performance-only layout property, a request for more detail in the description, and a naming nit.

Findings

0 total: 0 P0, 0 P1, 0 P2, 0 P3.

No findings.

Generated by Omnigent on Databricks