[SPARK-34581][SQL] Don't optimize out grouping expressions from aggregate expressions without aggregate function #32396
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,79 @@ | ||
| /* | ||
| * 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.{Alias, Expression, NamedExpression} | ||
| import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression | ||
| import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, LogicalPlan, Project} | ||
| import org.apache.spark.sql.catalyst.rules.Rule | ||
|
|
||
| /** | ||
| * This rule ensures that [[Aggregate]] nodes doesn't contain complex grouping expressions in the | ||
| * optimization phase. | ||
| * | ||
| * Complex grouping expressions are pulled out to a [[Project]] node under [[Aggregate]] and are | ||
| * referenced in both grouping expressions and aggregate expressions without aggregate functions. | ||
| * These references ensure that optimization rules don't change the aggregate expressions to invalid | ||
| * ones that no longer refer to any grouping expressions and also simplify the expression | ||
| * transformations on the node (need to transform the expression only once). | ||
| * | ||
| * For example, in the following query Spark shouldn't optimize the aggregate expression | ||
| * `Not(IsNull(c))` to `IsNotNull(c)` as the grouping expression is `IsNull(c)`: | ||
| * SELECT not(c IS NULL) | ||
| * FROM t | ||
| * GROUP BY c IS NULL | ||
| * Instead, the aggregate expression references a `_groupingexpression` attribute: | ||
| * Aggregate [_groupingexpression#233], [NOT _groupingexpression#233 AS (NOT (c IS NULL))#230] | ||
| * +- Project [isnull(c#219) AS _groupingexpression#233] | ||
| * +- LocalRelation [c#219] | ||
| */ | ||
| object PullOutGroupingExpressions extends Rule[LogicalPlan] { | ||
| override def apply(plan: LogicalPlan): LogicalPlan = { | ||
| plan transform { | ||
| case a: Aggregate if a.resolved => | ||
| val complexGroupingExpressionMap = mutable.LinkedHashMap.empty[Expression, NamedExpression] | ||
| val newGroupingExpressions = a.groupingExpressions.map { | ||
| case e if !e.foldable && e.children.nonEmpty => | ||
| complexGroupingExpressionMap | ||
| .getOrElseUpdate(e.canonicalized, Alias(e, s"_groupingexpression")()) | ||
| .toAttribute | ||
| case o => o | ||
| } | ||
| if (complexGroupingExpressionMap.nonEmpty) { | ||
| def replaceComplexGroupingExpressions(e: Expression): Expression = { | ||
| e match { | ||
| case _ if AggregateExpression.isAggregate(e) => e | ||
| case _ if e.foldable => e | ||
| case _ if complexGroupingExpressionMap.contains(e.canonicalized) => | ||
| complexGroupingExpressionMap.get(e.canonicalized).map(_.toAttribute).getOrElse(e) | ||
| case _ => e.mapChildren(replaceComplexGroupingExpressions) | ||
| } | ||
| } | ||
|
|
||
| val newAggregateExpressions = a.aggregateExpressions | ||
| .map(replaceComplexGroupingExpressions(_).asInstanceOf[NamedExpression]) | ||
|
There was a problem hiding this comment. Can it be done with a.transformExpressions similar to this one: There was a problem hiding this comment. We need to do a manual tree traversal if we want to stop recursion earlier, e.g. There was a problem hiding this comment. Does the following one work? a.transformExpressionsWithPruning(e => !(AggregateExpression.isAggregate(e) || e.fordable)) { There was a problem hiding this comment. Hmm, yes, this could work with some explicit casting. There was a problem hiding this comment. But this would traverse on There was a problem hiding this comment.
You're right. I think the following would behave the same as the manual recursion: a.aggregateExpressions.map(_.transformWithPruning(e => !(AggregateExpression.isAggregate(e) || e.fordable))({ There was a problem hiding this comment. Anyway, it's just my small preference -- it seems neater to use framework functions if it works. Feel free to merge whatever you feel comfortable with. There was a problem hiding this comment. I think I'm leaving this PR as it is now. But tested that peter-toth@ed374fe could work, just I need to cast There was a problem hiding this comment. Thanks a lot for exploring this, Peter! I'll think more about such use cases. |
||
| val newChild = Project(a.child.output ++ complexGroupingExpressionMap.values, a.child) | ||
| Aggregate(newGroupingExpressions, newAggregateExpressions, newChild) | ||
| } else { | ||
| a | ||
| } | ||
| } | ||
| } | ||
| } | ||
Uh oh!
There was an error while loading. Please reload this page.