From 0a82a426b790e1d7a7faa1c89b80264c465f6915 Mon Sep 17 00:00:00 2001 From: Shinoaki14 Date: Fri, 21 Aug 2026 12:03:25 +0700 Subject: [PATCH 1/3] [SPARK-57743][SQL] Support hash-based aggregation for collated grouping keys --- .../sql/catalyst/optimizer/Optimizer.scala | 3 +- .../optimizer/RewriteCollationAggregate.scala | 145 ++++++++++++++++++ .../plans/logical/basicLogicalOperators.scala | 4 +- .../apache/spark/sql/internal/SQLConf.scala | 10 ++ .../collation/CollationAggregationSuite.scala | 71 ++++++--- 5 files changed, 209 insertions(+), 24 deletions(-) create mode 100644 sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteCollationAggregate.scala diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala index 77ffd7e29da03..d3b96f097288f 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala @@ -367,7 +367,8 @@ abstract class Optimizer(catalogManager: CatalogManager) RewriteNearestByJoin, EvalInlineTables, ReplaceTranspose, - RewriteCollationJoin + RewriteCollationJoin, + RewriteCollationAggregate ) override def apply(plan: LogicalPlan): LogicalPlan = { diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteCollationAggregate.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteCollationAggregate.scala new file mode 100644 index 0000000000000..33a28448dfd17 --- /dev/null +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteCollationAggregate.scala @@ -0,0 +1,145 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.catalyst.optimizer + +import scala.collection.mutable + +import org.apache.spark.sql.catalyst.expressions._ +import org.apache.spark.sql.catalyst.expressions.aggregate._ +import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, LogicalPlan} +import org.apache.spark.sql.catalyst.rules.Rule +import org.apache.spark.sql.catalyst.trees.TreePattern.AGGREGATE +import org.apache.spark.sql.catalyst.util.UnsafeRowUtils +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types._ +import org.apache.spark.util.ArrayImplicits.SparkArrayOps + +/** + * This rule rewrites Aggregate grouping expressions to ensure that non-binary collated strings + * are converted to their binary-stable collation keys (via [[CollationKey]]). + * + * This allows hash-based aggregation (e.g., [[org.apache.spark.sql.execution.aggregate.ObjectHashAggregateExec]]) + * to work properly on data with non-binary collations, avoiding full sorting and spilling. + * + * Any original grouping expression referenced in the aggregate expressions (output) is preserved + * by wrapping it in `First(expr, ignoreNulls = false)`, an arbitrary representative of each + * collation-equal group. + */ +object RewriteCollationAggregate extends Rule[LogicalPlan] { + def apply(plan: LogicalPlan): LogicalPlan = { + if (!conf.collationHashAggregationEnabled) { + plan + } else { + plan.transformWithPruning(_.containsPattern(AGGREGATE)) { + case a @ Aggregate(groupingExpressions, aggregateExpressions, child, _) + if a.resolved && groupingExpressions.exists(e => !UnsafeRowUtils.isBinaryStable(e.dataType)) => + val keyMapping = mutable.LinkedHashMap.empty[Expression, Expression] + val newGroupingExpressions = groupingExpressions.map { ge => + val processed = processExpression(ge, ge.dataType) + if (!processed.fastEquals(ge)) { + keyMapping.put(ge.canonicalized, ge) + processed + } else { + ge + } + } + + if (keyMapping.nonEmpty) { + def replaceGroupingKeyReferences(e: Expression): Expression = { + e match { + case _: AggregateExpression => e + case _ if e.foldable => e + case _ if keyMapping.contains(e.canonicalized) => + val origExpr = keyMapping(e.canonicalized) + First(origExpr, ignoreNulls = false).toAggregateExpression() + case _ => + e.mapChildren(replaceGroupingKeyReferences) + } + } + + 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] + } + + a.copy( + groupingExpressions = newGroupingExpressions, + aggregateExpressions = newAggregateExpressions) + } else { + a + } + } + } + } + + /** + * 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) + + // Recursively process struct fields for non-binary structs. + case StructType(fields) => + processStruct(expr, fields) + + // Recursively process array elements for non-binary arrays. + case ArrayType(et, containsNull) => + processArray(expr, et, containsNull) + + case _ => + expr + } + } + + private def processStruct(str: Expression, fields: Array[StructField]): Expression = { + val struct = CreateNamedStruct(fields.zipWithIndex.flatMap { case (f, i) => + Seq(Literal(f.name), processExpression(GetStructField(str, i, Some(f.name)), f.dataType)) + }.toImmutableArraySeq) + if (str.nullable) { + If(IsNull(str), Literal(null, struct.dataType), struct) + } else { + struct + } + } + + private def processArray(arr: Expression, et: DataType, containsNull: Boolean): Expression = { + val param: NamedExpression = NamedLambdaVariable("a", et, containsNull) + val funcBody: Expression = processExpression(param, et) + if (!funcBody.fastEquals(param)) { + ArrayTransform(arr, LambdaFunction(funcBody, Seq(param))) + } else { + arr + } + } +} diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicLogicalOperators.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicLogicalOperators.scala index fb2dffed92b78..71d9e15941553 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicLogicalOperators.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicLogicalOperators.scala @@ -1337,7 +1337,9 @@ object Aggregate { return false } - aggregateExpressions.map(_.aggregateFunction).exists { + val schema = DataTypeUtils.fromAttributes( + aggregateExpressions.flatMap(_.aggregateFunction.aggBufferAttributes)) + !isAggregateBufferMutable(schema) || aggregateExpressions.map(_.aggregateFunction).exists { case _: TypedImperativeAggregate[_] => true case _ => false } diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala index 0ed4399361d0d..fb4d423b3385a 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala @@ -1387,6 +1387,14 @@ object SQLConf { .booleanConf .createWithDefault(false) + val COLLATION_HASH_AGGREGATION_ENABLED = + buildConf("spark.sql.collation.hashAggregation.enabled") + .doc("When true, allows hash-based aggregation for non-binary collated strings by " + + "rewriting grouping keys to collation keys.") + .version("4.1.0") + .booleanConf + .createWithDefault(true) + val OBJECT_LEVEL_COLLATIONS_ENABLED = buildConf("spark.sql.collation.objectLevel.enabled") .internal() @@ -8844,6 +8852,8 @@ class SQLConf extends Serializable with Logging with SqlApiConf { def allowCollationsInMapKeys: Boolean = getConf(ALLOW_COLLATIONS_IN_MAP_KEYS) + def collationHashAggregationEnabled: Boolean = getConf(COLLATION_HASH_AGGREGATION_ENABLED) + def objectLevelCollationsEnabled: Boolean = getConf(OBJECT_LEVEL_COLLATIONS_ENABLED) def schemaLevelCollationsEnabled: Boolean = getConf(SCHEMA_LEVEL_COLLATIONS_ENABLED) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/collation/CollationAggregationSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/collation/CollationAggregationSuite.scala index b894abe614761..61cdfeab03f27 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/collation/CollationAggregationSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/collation/CollationAggregationSuite.scala @@ -20,36 +20,64 @@ package org.apache.spark.sql.collation import org.apache.spark.sql.Row import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper import org.apache.spark.sql.execution.aggregate.{HashAggregateExec, ObjectHashAggregateExec, SortAggregateExec} +import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession class CollationAggregationSuite extends SharedSparkSession with AdaptiveSparkPlanHelper { - test("group by collated column doesn't work with obj hash aggregate") { + test("hash aggregate on collated grouping key with RewriteCollationAggregate") { val tblName = "grp_by_tbl" withTable(tblName) { sql(s"CREATE TABLE $tblName (c1 STRING COLLATE UTF8_LCASE, c2 INT) USING PARQUET") sql(s"INSERT INTO $tblName VALUES ('hello', 1), ('HELLO', 2), ('HeLlO', 3)") - // Result is correct without forcing object hash aggregate. - checkAnswer( - sql(s"SELECT COUNT(*) FROM $tblName GROUP BY c1"), - Seq(Row(3))) + val df = sql(s"SELECT c1, COUNT(*), SUM(c2) FROM $tblName GROUP BY c1") + val executedPlan = df.queryExecution.executedPlan - withSQLConf("spark.sql.test.forceApplyObjectHashAggregate" -> true.toString) { - checkAnswer( - sql(s"SELECT COUNT(*) FROM $tblName GROUP BY c1"), - Seq(Row(1), Row(1), Row(1))) + assert(collectFirst(executedPlan) { + case _: ObjectHashAggregateExec => true + }.nonEmpty) + assert(collectFirst(executedPlan) { + case _: SortAggregateExec => true + }.isEmpty) - checkAnswer( - sql(s"SELECT COLLECT_LIST(c2) AS c3 FROM $tblName GROUP BY c1 ORDER BY c3"), - Seq(Row(Seq(1)), Row(Seq(2)), Row(Seq(3)))) + val res = df.collect() + assert(res.length == 1) + assert(res(0).getString(0).toLowerCase() == "hello") + assert(res(0).getLong(1) == 3L) + assert(res(0).getLong(2) == 6L) + } + } + + test("disable hash aggregate on collated column via SQLConf") { + val tblName = "grp_by_disabled_tbl" + withTable(tblName) { + sql(s"CREATE TABLE $tblName (c1 STRING COLLATE UTF8_LCASE, c2 INT) USING PARQUET") + sql(s"INSERT INTO $tblName VALUES ('hello', 1), ('HELLO', 2), ('HeLlO', 3)") + + withSQLConf(SQLConf.COLLATION_HASH_AGGREGATION_ENABLED.key -> "false") { + val df = sql(s"SELECT c1, COUNT(*) FROM $tblName GROUP BY c1") + val executedPlan = df.queryExecution.executedPlan + + assert(collectFirst(executedPlan) { + case _: SortAggregateExec => true + }.nonEmpty) + assert(collectFirst(executedPlan) { + case _: ObjectHashAggregateExec => true + case _: HashAggregateExec => true + }.isEmpty) + + val res = df.collect() + assert(res.length == 1) + assert(res(0).getString(0).toLowerCase() == "hello") + assert(res(0).getLong(1) == 3L) } } } - test("imperative aggregate fn does not use objectHashAggregate when group by collated column") { + test("imperative aggregate fn uses objectHashAggregate when group by collated column") { val tblName = "imp_agg" Seq(true, false).foreach { useObjHashAgg => withTable(tblName) { @@ -66,17 +94,16 @@ class CollationAggregationSuite val df = sql(s"SELECT COLLECT_LIST(c2) as list FROM $tblName GROUP BY c1") val executedPlan = df.queryExecution.executedPlan - // Plan should not have any hash aggregate nodes. - collectFirst(executedPlan) { - case _: ObjectHashAggregateExec => fail("ObjectHashAggregateExec should not be used.") - case _: HashAggregateExec => fail("HashAggregateExec should not be used.") + if (useObjHashAgg) { + assert(collectFirst(executedPlan) { + case _: ObjectHashAggregateExec => true + }.nonEmpty) + } else { + assert(collectFirst(executedPlan) { + case _: SortAggregateExec => true + }.nonEmpty) } - // Plan should have a [[SortAggregateExec]] node. - assert(collectFirst(executedPlan) { - case _: SortAggregateExec => true - }.nonEmpty) - checkAnswer( // Sort the values to get deterministic output. df.selectExpr("array_sort(list)"), From 1e031946f63f6bf17fb88ff9d1164ae742c84fb1 Mon Sep 17 00:00:00 2001 From: Shinoaki14 Date: Fri, 21 Aug 2026 12:21:21 +0700 Subject: [PATCH 2/3] [SPARK-58927][SQL] Support hash-based aggregation for collated grouping keys --- .../optimizer/RewriteCollationAggregate.scala | 69 ++++--------------- 1 file changed, 14 insertions(+), 55 deletions(-) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteCollationAggregate.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteCollationAggregate.scala index 33a28448dfd17..a6055a169af09 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteCollationAggregate.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteCollationAggregate.scala @@ -25,9 +25,6 @@ import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, LogicalPlan} import org.apache.spark.sql.catalyst.rules.Rule import org.apache.spark.sql.catalyst.trees.TreePattern.AGGREGATE import org.apache.spark.sql.catalyst.util.UnsafeRowUtils -import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.types._ -import org.apache.spark.util.ArrayImplicits.SparkArrayOps /** * This rule rewrites Aggregate grouping expressions to ensure that non-binary collated strings @@ -50,7 +47,7 @@ object RewriteCollationAggregate extends Rule[LogicalPlan] { if a.resolved && groupingExpressions.exists(e => !UnsafeRowUtils.isBinaryStable(e.dataType)) => val keyMapping = mutable.LinkedHashMap.empty[Expression, Expression] val newGroupingExpressions = groupingExpressions.map { ge => - val processed = processExpression(ge, ge.dataType) + val processed = CollationKey.injectCollationKey(ge) if (!processed.fastEquals(ge)) { keyMapping.put(ge.canonicalized, ge) processed @@ -76,12 +73,23 @@ object RewriteCollationAggregate extends Rule[LogicalPlan] { case a @ Alias(child, name) => val newChild = replaceGroupingKeyReferences(child) if (!newChild.fastEquals(child)) { - Alias(newChild, name)(exprId = a.exprId, explicitMetadata = a.explicitMetadata) + Alias(newChild, name)(exprId = a.exprId, qualifier = a.qualifier, explicitMetadata = a.explicitMetadata) } else { a } + case attr: Attribute => + val newChild = replaceGroupingKeyReferences(attr) + if (!newChild.fastEquals(attr)) { + Alias(newChild, attr.name)(exprId = attr.exprId, qualifier = attr.qualifier) + } else { + attr + } case other => - replaceGroupingKeyReferences(other).asInstanceOf[NamedExpression] + val newOther = replaceGroupingKeyReferences(other) + newOther match { + case ne: NamedExpression => ne + case expr => Alias(expr, expr.prettyName)() + } } a.copy( @@ -93,53 +101,4 @@ object RewriteCollationAggregate extends Rule[LogicalPlan] { } } } - - /** - * 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) - - // Recursively process struct fields for non-binary structs. - case StructType(fields) => - processStruct(expr, fields) - - // Recursively process array elements for non-binary arrays. - case ArrayType(et, containsNull) => - processArray(expr, et, containsNull) - - case _ => - expr - } - } - - private def processStruct(str: Expression, fields: Array[StructField]): Expression = { - val struct = CreateNamedStruct(fields.zipWithIndex.flatMap { case (f, i) => - Seq(Literal(f.name), processExpression(GetStructField(str, i, Some(f.name)), f.dataType)) - }.toImmutableArraySeq) - if (str.nullable) { - If(IsNull(str), Literal(null, struct.dataType), struct) - } else { - struct - } - } - - private def processArray(arr: Expression, et: DataType, containsNull: Boolean): Expression = { - val param: NamedExpression = NamedLambdaVariable("a", et, containsNull) - val funcBody: Expression = processExpression(param, et) - if (!funcBody.fastEquals(param)) { - ArrayTransform(arr, LambdaFunction(funcBody, Seq(param))) - } else { - arr - } - } } From 4d8aafab60159b11a347a3302ee5b7390123cf76 Mon Sep 17 00:00:00 2001 From: Shinoaki14 Date: Fri, 21 Aug 2026 12:56:13 +0700 Subject: [PATCH 3/3] [SPARK-58927][SQL] Support hash-based aggregation on non-binary collated grouping keys --- .../optimizer/RewriteCollationAggregate.scala | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteCollationAggregate.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteCollationAggregate.scala index a6055a169af09..d923442b4720f 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteCollationAggregate.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteCollationAggregate.scala @@ -30,8 +30,8 @@ import org.apache.spark.sql.catalyst.util.UnsafeRowUtils * This rule rewrites Aggregate grouping expressions to ensure that non-binary collated strings * are converted to their binary-stable collation keys (via [[CollationKey]]). * - * This allows hash-based aggregation (e.g., [[org.apache.spark.sql.execution.aggregate.ObjectHashAggregateExec]]) - * to work properly on data with non-binary collations, avoiding full sorting and spilling. + * This allows hash-based aggregation (e.g., ObjectHashAggregateExec) to work properly on data + * with non-binary collations, avoiding full sorting and spilling. * * Any original grouping expression referenced in the aggregate expressions (output) is preserved * by wrapping it in `First(expr, ignoreNulls = false)`, an arbitrary representative of each @@ -44,7 +44,8 @@ object RewriteCollationAggregate extends Rule[LogicalPlan] { } else { plan.transformWithPruning(_.containsPattern(AGGREGATE)) { case a @ Aggregate(groupingExpressions, aggregateExpressions, child, _) - if a.resolved && groupingExpressions.exists(e => !UnsafeRowUtils.isBinaryStable(e.dataType)) => + if a.resolved && + groupingExpressions.exists(e => !UnsafeRowUtils.isBinaryStable(e.dataType)) => val keyMapping = mutable.LinkedHashMap.empty[Expression, Expression] val newGroupingExpressions = groupingExpressions.map { ge => val processed = CollationKey.injectCollationKey(ge) @@ -73,7 +74,10 @@ object RewriteCollationAggregate extends Rule[LogicalPlan] { case a @ Alias(child, name) => val newChild = replaceGroupingKeyReferences(child) if (!newChild.fastEquals(child)) { - Alias(newChild, name)(exprId = a.exprId, qualifier = a.qualifier, explicitMetadata = a.explicitMetadata) + Alias(newChild, name)( + exprId = a.exprId, + qualifier = a.qualifier, + explicitMetadata = a.explicitMetadata) } else { a }