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..d923442b4720f --- /dev/null +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteCollationAggregate.scala @@ -0,0 +1,108 @@ +/* + * 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 + +/** + * 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., 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 = CollationKey.injectCollationKey(ge) + 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, + 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 => + val newOther = replaceGroupingKeyReferences(other) + newOther match { + case ne: NamedExpression => ne + case expr => Alias(expr, expr.prettyName)() + } + } + + a.copy( + groupingExpressions = newGroupingExpressions, + aggregateExpressions = newAggregateExpressions) + } else { + a + } + } + } + } +} 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)"),