[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.