[SPARK-32750][SQL] Support whole-stage codegen for SortAggregate with grouping keys

[ulysses-you]

What changes were proposed in this pull request?

This PR adds whole-stage code-gen support for SortAggregateExec when it has grouping
keys. Previously only the no-grouping-keys path was code-gen'd; with keys it fell back to
the interpreted SortBasedAggregationIterator.

The implementation:

• AggregateCodegenSupport is refactored to share the aggregation-buffer creation
  (createAggBufVars) and buffer-update (generateAggBufferUpdateCode) logic between the
  no-keys path and the new sort-based with-keys path.
• SortAggregateExec implements doProduceWithKeys / doConsumeWithKeys. Since the input
  is sorted by the grouping keys, a group's result is emitted as soon as the next group
  starts (detected by comparing the binary representation of the grouping key). The produce
  loop is resumable across processNext invocations via shouldStop(), so a completed
  group's output row is not overwritten by the reused output buffer.
• The with-keys path only supports binary-stable grouping-key types (guarded in
  supportCodegen), because group boundaries are detected via UnsafeRow.equals.
• A new internal config spark.sql.codegen.aggregate.sortAggregate.withKeys.enabled
  (default true) gates this path; it takes effect only when
  spark.sql.codegen.aggregate.sortAggregate.enabled is enabled.

Why are the changes needed?

Sort aggregate with grouping keys was the only remaining aggregate path without whole-stage
code-gen, forcing it through the slower interpreted iterator. SortAggregateBenchmark shows
a consistent 1.2~1.3x speedup for grouped aggregates on the code-gen path.

Does this PR introduce any user-facing change?

No. This is an internal code-gen optimization; results are unchanged. The new config is
internal().

How was this patch tested?

• New tests in WholeStageCodegenSuite covering: multiple/numeric/string/decimal grouping
  keys, null keys and values, single-row groups, a single all-rows group, downstream limit
  (resumable shouldStop path), empty input, single partition, split aggregate functions,
  FILTER clauses, HAVING-style filters, and the config gate. Each test asserts a code-gen'd
  SortAggregateExec in the plan and checks the result matches the interpreted (code-gen
  disabled) result.
• Added SortAggregateBenchmark with results committed for JDK 17/21/25.

The generated code of a simple query: select id , count(*) from t1 group by id:

/* 001 */ public Object generate(Object[] references) {
/* 002 */   return new GeneratedIteratorForCodegenStage1(references);
/* 003 */ }
/* 004 */
/* 005 */ // codegenStageId=1
/* 006 */ final class GeneratedIteratorForCodegenStage1 extends org.apache.spark.sql.execution.BufferedRowIterator {
/* 007 */   private Object[] references;
/* 008 */   private scala.collection.Iterator[] inputs;
/* 009 */   private boolean sortAgg_bufIsNull_0;
/* 010 */   private long sortAgg_bufValue_0;
/* 011 */   private UnsafeRow sortAgg_currentGroupingKey_0;
/* 012 */   private boolean sortAgg_initGroup_0;
/* 013 */   private boolean sortAgg_noMoreInputTerm_0;
/* 014 */   private boolean sort_needToSort_0;
/* 015 */   private org.apache.spark.sql.execution.UnsafeExternalRowSorter sort_sorter_0;
/* 016 */   private org.apache.spark.executor.TaskMetrics sort_metrics_0;
/* 017 */   private scala.collection.Iterator<UnsafeRow> sort_sortedIter_0;
/* 018 */   private int columnartorow_batchIdx_0;
/* 019 */   private org.apache.spark.sql.vectorized.ColumnarBatch[] columnartorow_mutableStateArray_1 = new org.apache.spark.sql.vectorized.ColumnarBatch[1];
/* 020 */   private org.apache.spark.sql.execution.vectorized.OnHeapColumnVector[] columnartorow_mutableStateArray_2 = new org.apache.spark.sql.execution.vectorized.OnHeapColumnVector[1];
/* 021 */   private scala.collection.Iterator[] columnartorow_mutableStateArray_0 = new scala.collection.Iterator[1];
/* 022 */   private org.apache.spark.sql.catalyst.expressions.codegen.UnsafeRowWriter[] sortAgg_mutableStateArray_0 = new org.apache.spark.sql.catalyst.expressions.codegen.UnsafeRowWriter[3];
/* 023 */
/* 024 */   public GeneratedIteratorForCodegenStage1(Object[] references) {
/* 025 */     this.references = references;
/* 026 */   }
/* 027 */
/* 028 */   public void init(int index, scala.collection.Iterator[] inputs) {
/* 029 */     partitionIndex = index;
/* 030 */     this.inputs = inputs;
/* 031 */
/* 032 */     sortAgg_mutableStateArray_0[0] = new org.apache.spark.sql.catalyst.expressions.codegen.UnsafeRowWriter(2, 0);
/* 033 */     sort_needToSort_0 = true;
/* 034 */     sort_sorter_0 = ((org.apache.spark.sql.execution.SortExec) references[1] /* plan */).createSorter();
/* 035 */     sort_metrics_0 = org.apache.spark.TaskContext.get().taskMetrics();
/* 036 */     columnartorow_mutableStateArray_0[0] = inputs[0];
/* 037 */     sortAgg_mutableStateArray_0[1] = new org.apache.spark.sql.catalyst.expressions.codegen.UnsafeRowWriter(1, 0);
/* 038 */     sortAgg_mutableStateArray_0[2] = new org.apache.spark.sql.catalyst.expressions.codegen.UnsafeRowWriter(1, 0);
/* 039 */
/* 040 */   }
/* 041 */
/* 042 */   private void sortAgg_doConsume_0(InternalRow sort_outputRow_0, long sortAgg_expr_0_0, boolean sortAgg_exprIsNull_0_0) throws java.io.IOException {
/* 043 */     sortAgg_mutableStateArray_0[2].reset();
/* 044 */
/* 045 */     sortAgg_mutableStateArray_0[2].zeroOutNullBytes();
/* 046 */
/* 047 */     sortAgg_mutableStateArray_0[2].writeNullable(0, sortAgg_expr_0_0, sortAgg_exprIsNull_0_0);
/* 048 */     if (!sortAgg_initGroup_0) {
/* 049 */       sortAgg_initGroup_0 = true;
/* 050 */       sortAgg_currentGroupingKey_0 = (sortAgg_mutableStateArray_0[2].getRow()).copy();
/* 051 */       sortAgg_bufIsNull_0 = false;
/* 052 */       sortAgg_bufValue_0 = 0L;
/* 053 */     } else if (!sortAgg_currentGroupingKey_0.equals((sortAgg_mutableStateArray_0[2].getRow()))) {
/* 054 */       sortAgg_doAggregateWithKeysOutput_0();
/* 055 */       sortAgg_currentGroupingKey_0 = (sortAgg_mutableStateArray_0[2].getRow()).copy();
/* 056 */       sortAgg_bufIsNull_0 = false;
/* 057 */       sortAgg_bufValue_0 = 0L;
/* 058 */     }
/* 059 */
/* 060 */     // do aggregate
/* 061 */     // common sub-expressions
/* 062 */
/* 063 */     // evaluate aggregate functions and update aggregation buffers
/* 064 */
/* 065 */     long sortAgg_value_7 = -1L;
/* 066 */
/* 067 */     sortAgg_value_7 = org.apache.spark.sql.catalyst.util.MathUtils.addExact(sortAgg_bufValue_0, 1L, ((org.apache.spark.sql.catalyst.trees.SQLQueryContext) references[7] /* errCtx */));
/* 068 */
/* 069 */     sortAgg_bufIsNull_0 = false;
/* 070 */     sortAgg_bufValue_0 = sortAgg_value_7;
/* 071 */
/* 072 */   }
/* 073 */
/* 074 */   private void wholestagecodegen_doAggregateWithKeys_0() throws java.io.IOException {
/* 075 */     if (sort_needToSort_0) {
/* 076 */       long sort_spillSizeBefore_0 = sort_metrics_0.memoryBytesSpilled();
/* 077 */       sort_addToSorter_0(partitionIndex);
/* 078 */       sort_sortedIter_0 = sort_sorter_0.sort();
/* 079 */       ((org.apache.spark.sql.execution.metric.SQLMetric) references[6] /* sortTime */).add(sort_sorter_0.getSortTimeNanos() / 1000000);
/* 080 */       ((org.apache.spark.sql.execution.metric.SQLMetric) references[4] /* peakMemory */).add(sort_sorter_0.getPeakMemoryUsage());
/* 081 */       ((org.apache.spark.sql.execution.metric.SQLMetric) references[5] /* spillSize */).add(sort_metrics_0.memoryBytesSpilled() - sort_spillSizeBefore_0);
/* 082 */       sort_metrics_0.incPeakExecutionMemory(sort_sorter_0.getPeakMemoryUsage());
/* 083 */       sort_needToSort_0 = false;
/* 084 */     }
/* 085 */
/* 086 */     while ( sort_sortedIter_0.hasNext()) {
/* 087 */       UnsafeRow sort_outputRow_0 = (UnsafeRow)sort_sortedIter_0.next();
/* 088 */
/* 089 */       boolean sort_isNull_0 = sort_outputRow_0.isNullAt(0);
/* 090 */       long sort_value_0 = sort_isNull_0 ?
/* 091 */       -1L : (sort_outputRow_0.getLong(0));
/* 092 */
/* 093 */       sortAgg_doConsume_0(sort_outputRow_0, sort_value_0, sort_isNull_0);
/* 094 */
/* 095 */       if (shouldStop()) return;
/* 096 */     }
/* 097 */
/* 098 */   }
/* 099 */
/* 100 */   private void sort_addToSorter_0(int partitionIndex) throws java.io.IOException {
/* 101 */     if (columnartorow_mutableStateArray_1[0] == null) {
/* 102 */       columnartorow_nextBatch_0();
/* 103 */     }
/* 104 */     while ( columnartorow_mutableStateArray_1[0] != null) {
/* 105 */       int columnartorow_numRows_0 = columnartorow_mutableStateArray_1[0].numRows();
/* 106 */       int columnartorow_localEnd_0 = columnartorow_numRows_0 - columnartorow_batchIdx_0;
/* 107 */       for (int columnartorow_localIdx_0 = 0; columnartorow_localIdx_0 < columnartorow_localEnd_0; columnartorow_localIdx_0++) {
/* 108 */         int columnartorow_rowIdx_0 = columnartorow_batchIdx_0 + columnartorow_localIdx_0;
/* 109 */         boolean columnartorow_isNull_0 = columnartorow_mutableStateArray_2[0].isNullAt(columnartorow_rowIdx_0);
/* 110 */         long columnartorow_value_0 = columnartorow_isNull_0 ? -1L : (columnartorow_mutableStateArray_2[0].getLong(columnartorow_rowIdx_0));
/* 111 */         sortAgg_mutableStateArray_0[1].reset();
/* 112 */
/* 113 */         sortAgg_mutableStateArray_0[1].zeroOutNullBytes();
/* 114 */
/* 115 */         sortAgg_mutableStateArray_0[1].writeNullable(0, columnartorow_value_0, columnartorow_isNull_0);
/* 116 */         sort_sorter_0.insertRow((UnsafeRow)(sortAgg_mutableStateArray_0[1].getRow()));
/* 117 */         // shouldStop check is eliminated
/* 118 */       }
/* 119 */       columnartorow_batchIdx_0 = columnartorow_numRows_0;
/* 120 */       columnartorow_nextBatch_0();
/* 121 */     }
/* 122 */     // clean up resources
/* 123 */     if (columnartorow_mutableStateArray_1[0] != null) {
/* 124 */       columnartorow_mutableStateArray_1[0].close();
/* 125 */     }
/* 126 */
/* 127 */   }
/* 128 */
/* 129 */   protected void processNext() throws java.io.IOException {
/* 130 */     if (!sortAgg_noMoreInputTerm_0) {
/* 131 */       wholestagecodegen_doAggregateWithKeys_0();
/* 132 */       if (!shouldStop()) {
/* 133 */         sortAgg_noMoreInputTerm_0 = true;
/* 134 */         if (sortAgg_initGroup_0) {
/* 135 */           sortAgg_doAggregateWithKeysOutput_0();
/* 136 */         }
/* 137 */       }
/* 138 */     }
/* 139 */   }
/* 140 */
/* 141 */   private void columnartorow_nextBatch_0() throws java.io.IOException {
/* 142 */     columnartorow_mutableStateArray_1[0] = org.apache.spark.sql.execution.ColumnarToRowExec.advanceBatch(
/* 143 */       columnartorow_mutableStateArray_0[0], columnartorow_mutableStateArray_1[0], ((org.apache.spark.sql.execution.metric.SQLMetric) references[3] /* numInputBatches */), ((org.apache.spark.sql.execution.metric.SQLMetric) references[2] /* numOutputRows */));
/* 144 */     if (columnartorow_mutableStateArray_1[0] != null) {
/* 145 */       columnartorow_batchIdx_0 = 0;
/* 146 */       columnartorow_mutableStateArray_2[0] = (org.apache.spark.sql.execution.vectorized.OnHeapColumnVector) columnartorow_mutableStateArray_1[0].column(0);
/* 147 */
/* 148 */     }
/* 149 */   }
/* 150 */
/* 151 */   private void sortAgg_doAggregateWithKeysOutput_0() throws java.io.IOException {
/* 152 */     ((org.apache.spark.sql.execution.metric.SQLMetric) references[0] /* numOutputRows */).add(1);
/* 153 */
/* 154 */     boolean sortAgg_isNull_1 = sortAgg_currentGroupingKey_0.isNullAt(0);
/* 155 */     long sortAgg_value_1 = sortAgg_isNull_1 ?
/* 156 */     -1L : (sortAgg_currentGroupingKey_0.getLong(0));
/* 157 */
/* 158 */     sortAgg_mutableStateArray_0[0].reset();
/* 159 */
/* 160 */     sortAgg_mutableStateArray_0[0].zeroOutNullBytes();
/* 161 */
/* 162 */     sortAgg_mutableStateArray_0[0].writeNullable(0, sortAgg_value_1, sortAgg_isNull_1);
/* 163 */
/* 164 */     sortAgg_mutableStateArray_0[0].write(1, sortAgg_bufValue_0);
/* 165 */     append((sortAgg_mutableStateArray_0[0].getRow()));
/* 166 */
/* 167 */   }
/* 168 */
/* 169 */ }

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

Generated-by: Claude Code

[ulysses-you]

.cc @viirya @dongjoon-hyun @cloud-fan if you have time to take a look, thank you!

[LuciferYang]

Long time no see.

[LuciferYang]

The $doAgg wrapper here doesn't take an int partitionIndex parameter, unlike the no-keys path, HashAggregateExec, and SortExec. The child's produce emits bare partitionIndex references (e.g. addToSorter(partitionIndex) for a Sort child). When the outer generated class exceeds GENERATED_CLASS_SIZE_THRESHOLD (1MB) and addNewFunction spills $doAgg into a private class NestedClass, that bare reference resolves to the protected BufferedRowIterator.partitionIndex field, which the inner (non-subclass) class can't access, so Janino compilation throws IllegalAccessError. That's an Error, not NonFatal, so the codegen fallback in WholeStageCodegenExec won't catch it and the task fails. Suggest matching the other three wrappers: private void $doAgg(int partitionIndex), called as $doAggFuncName(partitionIndex).

[LuciferYang]

This PR removes the only two throw sites for _LEGACY_ERROR_TEMP_3170 in SortAggregateExec, leaving the entry at error-conditions.json:11601 unreferenced. `

[LuciferYang]

supportCodegenWithKeys gates only on isBinaryStable, and Double/Float/Decimal are all binary-stable, so they take this new path. Group boundaries use raw UnsafeRow.equals, so float-key correctness rides on the planner's NormalizeFloatingNumbers folding -0.0/NaN first, but the new tests only cover long/string keys with no float grouping-key correctness assertion (decimal appears only in the benchmark). Suggest adding a float grouping-key case (with

[LuciferYang]

The third branch of generateResultFunctionForKeys (no aggregate functions, grouping-only) is codegen'd, but every new test carries at least one aggregate function, so this branch and its empty bufVars/reInitBufferCode path are uncovered. Suggest adding a groupBy("k").agg() or distinct case.