Copilot commented on code in PR #58196: URL: https://github.com/apache/spark/pull/58196#discussion_r3827432409
########## 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) + Review Comment: This rule duplicates the recursive collation-key injection logic that already exists in `org.apache.spark.sql.catalyst.expressions.CollationKey.injectCollationKey(...)` (including struct/array handling). Reusing that helper would reduce duplication and keep join/aggregate behavior consistent (e.g., any future changes to injection rules would only need to be made in one place). ########## 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] + } Review Comment: `replaceGroupingKeyReferences(other).asInstanceOf[NamedExpression]` can produce a non-`NamedExpression` (e.g., when `other` is a grouping-key `AttributeReference`, the replacement is `First(...).toAggregateExpression()`, which is an `AggregateExpression`). This will throw a `ClassCastException` during optimization/planning for queries like `SELECT c1, count(*) ... GROUP BY c1` (where the grouping key appears directly in `aggregateExpressions`). Wrap replaced grouping-key outputs in an `Alias` (preserving `exprId`/qualifier/metadata) instead of casting. ########## 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 } Review Comment: `supportsObjectHashAggregate` now returns true for *any* aggregate whose buffer schema is non-mutable (even when there are no `TypedImperativeAggregate`s). This broadens planning behavior beyond collated grouping keys (e.g., `GROUP BY int` with `first(string)`/`max(string)` could now plan `ObjectHashAggregateExec` instead of `SortAggregateExec` when `spark.sql.execution.useObjectHashAggregateExec=true`). The PR description's user-facing-change section only mentions collated columns; please either narrow this behavior to the collation rewrite path or update the PR description/tests to reflect the wider change. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
