[SPARK-58927][SQL] Support hash-based aggregation for collated groupi… - #58196
Open
Shinoaki0145 wants to merge 3 commits into
Open
[SPARK-58927][SQL] Support hash-based aggregation for collated groupi…#58196Shinoaki0145 wants to merge 3 commits into
Shinoaki0145 wants to merge 3 commits into
Conversation
Shinoaki0145
force-pushed
the
SPARK-58927-Support-hash-based-aggregation
branch
from
August 21, 2026 05:14
9766096 to
0a82a42
Compare
There was a problem hiding this comment.
Pull request overview
Adds Catalyst/planner support to enable hash-based aggregation for non-binary-stable collated grouping keys by rewriting grouping expressions to collation keys, with a SQLConf flag and accompanying test coverage.
Changes:
- Introduces
RewriteCollationAggregate(FinishAnalysis) to rewrite non-binary-stable grouping keys usingCollationKey(...)and preserve original grouping outputs viaFirst(...). - Updates
Aggregate.supportsObjectHashAggregateto allowObjectHashAggregateExecwhen grouping keys are binary-stable and the aggregate buffer schema is non-mutable. - Adds
spark.sql.collation.hashAggregation.enabled(defaulttrue) and expandsCollationAggregationSuiteto validate the new planning behavior and the disable switch.
Reviewed changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| sql/core/src/test/scala/org/apache/spark/sql/collation/CollationAggregationSuite.scala | Updates tests to assert ObjectHashAggregateExec is chosen for collated grouping keys and that the feature flag falls back to sort aggregation. |
| sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala | Adds a new SQLConf flag to enable/disable collation hash aggregation rewrites. |
| sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicLogicalOperators.scala | Broadens supportsObjectHashAggregate to permit object-hash aggregation with non-mutable aggregate buffers when keys are binary-stable. |
| sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteCollationAggregate.scala | New optimizer rule implementing the collation-key rewrite for aggregates and preserving original grouping outputs. |
| sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala | Registers RewriteCollationAggregate in the FinishAnalysis rule set. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+75
to
+85
| val newAggregateExpressions = aggregateExpressions.map { | ||
| case a @ Alias(child, name) => | ||
| val newChild = replaceGroupingKeyReferences(child) | ||
| if (!newChild.fastEquals(child)) { | ||
| Alias(newChild, name)(exprId = a.exprId, explicitMetadata = a.explicitMetadata) | ||
| } else { | ||
| a | ||
| } | ||
| case other => | ||
| replaceGroupingKeyReferences(other).asInstanceOf[NamedExpression] | ||
| } |
Comment on lines
+97
to
+111
| /** | ||
| * Recursively process the expression in order to replace non-binary collated strings with their | ||
| * associated collation keys. This is necessary to ensure grouping is evaluated correctly for all | ||
| * types containing non-binary collated strings, including structs and arrays. | ||
| */ | ||
| private def processExpression(expr: Expression, dt: DataType): Expression = { | ||
| dt match { | ||
| // For binary stable expressions, no special handling is needed. | ||
| case _ if UnsafeRowUtils.isBinaryStable(dt) => | ||
| expr | ||
|
|
||
| // Inject CollationKey for non-binary collated strings. | ||
| case _: StringType => | ||
| CollationKey(expr) | ||
|
|
Comment on lines
+1340
to
1345
| val schema = DataTypeUtils.fromAttributes( | ||
| aggregateExpressions.flatMap(_.aggregateFunction.aggBufferAttributes)) | ||
| !isAggregateBufferMutable(schema) || aggregateExpressions.map(_.aggregateFunction).exists { | ||
| case _: TypedImperativeAggregate[_] => true | ||
| case _ => false | ||
| } |
Shinoaki0145
changed the title
…ted grouping keys
uros-b
reviewed
Aug 23, 2026
uros-b
left a comment
uros-b
left a comment
Member
There was a problem hiding this comment.
I think this might be duplicating #58083, @Shinoaki0145 PTAL
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What changes were proposed in this pull request?
This PR mirrors the approach introduced in
RewriteCollationJoin(SPARK-48000) for hash joins and introduces hash-based aggregation support for non-binary collated grouping keys:RewriteCollationAggregateOptimizer Rule:RewriteCollationAggregateregistered under theFinish Analysisbatch.Aggregateoperators with non-binary-stable grouping keys (e.g.UTF8_LCASE, structs/arrays containing non-binary collated strings), it rewrites the grouping expressions usingCollationKey.injectCollationKey(...)so grouping is performed on binary-stable collation keys.First(origExpr, ignoreNulls = false).toAggregateExpression(), properly aliased with the originalexprIdand qualifier.ObjectHashAggregate Planner Update:
Aggregate.supportsObjectHashAggregateto support aggregations with non-mutable buffer schemas (such as those containingFirst(string)from collation rewrites) when grouping expressions are binary-stable, allowing these queries to plan asObjectHashAggregateExecinstead of falling back toSortAggregateExec.Feature Flag:
spark.sql.collation.hashAggregation.enabled(default:true) to allow reverting to the previous sort-based behavior if necessary.Why are the changes needed?
Currently, a
GROUP BY(or any aggregation) on non-binary collated keys (e.g.,UTF8_LCASE) is planned asSortAggregateExecbecauseUnsafeRowUtils.isBinaryStablereturnsfalsefor non-binary collations, preventing hash aggregation.This forces a full-input sort and can cause heavy disk spilling on large datasets. While hash join already solves this by injecting
CollationKeyinRewriteCollationJoin, aggregations were not given equivalent optimization. Rewriting grouping keys to collation keys and routing toObjectHashAggregateExecavoids the mandatory sort and disk spill, significantly improving query performance on collated data.Does this PR introduce any user-facing change?
Yes:
ObjectHashAggregateExecinstead ofSortAggregateExecby default, dramatically improving execution speed without requiring manual casts to binary.ObjectHashAggregateExecwhenspark.sql.execution.useObjectHashAggregateExecis enabled.Users can disable the collation aggregation rewrite optimization and restore the legacy sort-based execution by setting:
How was this patch tested?
org.apache.spark.sql.collation.CollationAggregationSuite:ObjectHashAggregateExecis selected and output results are correct forGROUP BYonUTF8_LCASE.spark.sql.collation.hashAggregation.enabled = falsefalls back toSortAggregateExec.collect_list).Was this patch authored or co-authored using generative AI tooling?
No.