Skip to content

[SPARK-40382][SQL] Group distinct aggregate expressions by semantically equivalent children in RewriteDistinctAggregates #37825

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
bersprockets wants to merge 13 commits into from
Closed

[SPARK-40382][SQL] Group distinct aggregate expressions by semantically equivalent children in RewriteDistinctAggregates #37825

Changes from 1 commit
Commits
Show all changes
13 commits
Select commit Hold shift + click to select a range
a5a6fc0
First attempt
bersprockets Mar 11, 2022
0a109d9
Fix for foldables
bersprockets Mar 11, 2022
38f1f6a
Update
bersprockets Mar 18, 2022
3fa3588
Update test
bersprockets Mar 18, 2022
4a40f91
Simplify
bersprockets Sep 3, 2022
165f558
Rename
bersprockets Sep 4, 2022
27dcffe
Update comments
bersprockets Sep 7, 2022
484ca8e
Replace Symbol usage with $"" in new unit tests
bersprockets Sep 7, 2022
208fe82
Update tests
bersprockets Sep 8, 2022
882cdaa
Update
bersprockets Sep 26, 2022
f53136d
Use ExpressionSet as key for various distinct aggregate child maps
bersprockets Sep 27, 2022
9938252
Handle case of one distinct grouping with superficially different fun…
bersprockets Oct 3, 2022
f7d29df
Update
bersprockets Oct 6, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
Rename
  • Loading branch information
@bersprockets
bersprockets committed Oct 11, 2022
commit 165f558590dc46ffe50ebcbc9fcc9cd81f801c96
33 changes: 22 additions & 11 deletions ...st/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteDistinctAggregates.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -213,10 +213,13 @@ object RewriteDistinctAggregates extends Rule[LogicalPlan] {
case a: Aggregate if mayNeedtoRewrite(a) => rewrite(a)
}

def rewrite(aRaw: Aggregate): Aggregate = {
def rewrite(aOrig: Aggregate): Aggregate = {
// make children of distinct aggregations the same if they are different
// only because of superficial differences.
val a = getSanitizedAggregate(aRaw)
// only because of superficial reasons, e.g.:
// "1 + col1" vs "col1 + 1", both become "1 + col1"
// or
// "col1" vs "Col1", both become "col1"
val a = reduceDistinctAggregateGroups(aOrig)
Copy link
Contributor

@cloud-fan cloud-fan Sep 21, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

shall we just canonicalize the function inputs when group by them? e.g. change e.aggregateFunction.children.filter(!_.foldable).toSet to ExpressionSet(e.aggregateFunction.children.filter(!_.foldable))

Copy link
Contributor Author

@bersprockets bersprockets Sep 24, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks! I am working on it, just working through some small complications.

Copy link
Contributor Author

@bersprockets bersprockets Sep 26, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I made the change to use ExpressionSet and also commented on some of the issues.

I still prefer 'sanitizing' each original function child to use the first semantically equivalent child, in essence creating a new set of "original" children, as it bypasses some complexities (in particular the one where we may lose some of the original children as keys when we group by ExpressionSet).


val aggExpressions = collectAggregateExprs(a)
val distinctAggs = aggExpressions.filter(_.isDistinct)
Expand Down Expand Up @@ -248,6 +251,13 @@ object RewriteDistinctAggregates extends Rule[LogicalPlan] {
}
val groupByAttrs = groupByMap.map(_._2)

def patchAggregateFunctionChildren(
af: AggregateFunction)(
attrs: Expression => Option[Expression]): AggregateFunction = {
val newChildren = af.children.map(c => attrs(c).getOrElse(c))
af.withNewChildren(newChildren).asInstanceOf[AggregateFunction]
}

// Setup unique distinct aggregate children.
val distinctAggChildren = distinctAggGroups.keySet.flatten.toSeq.distinct
Copy link
Contributor

@cloud-fan cloud-fan Sep 27, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

not sure if this is necessary, but it's better to use ExpressionSet(distinctAggGroups.keySet.flatten).toSeq, instead of calling .distinct on Seq[Expression]

val distinctAggChildAttrMap = distinctAggChildren.map(expressionAttributePair)
Expand Down Expand Up @@ -409,14 +419,7 @@ object RewriteDistinctAggregates extends Rule[LogicalPlan] {
}}
}

private def patchAggregateFunctionChildren(
af: AggregateFunction)(
attrs: Expression => Option[Expression]): AggregateFunction = {
val newChildren = af.children.map(c => attrs(c).getOrElse(c))
af.withNewChildren(newChildren).asInstanceOf[AggregateFunction]
}

private def getSanitizedAggregate(a: Aggregate): Aggregate = {
private def reduceDistinctAggregateGroups(a: Aggregate): Aggregate = {
val aggExpressions = collectAggregateExprs(a)
val distinctAggs = aggExpressions.filter(_.isDistinct)

Expand All @@ -436,6 +439,14 @@ object RewriteDistinctAggregates extends Rule[LogicalPlan] {
(e, funcChildren.find(fc => e.semanticEquals(fc)).getOrElse(e))
Copy link
Contributor

@beliefer beliefer Sep 21, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It seems the code lead to find out itself.

}.toMap

val funcChildrenPatched = funcChildren.map { e =>
funcChildrenLookup.getOrElse(e, e)
}

if (funcChildren.distinct.size == funcChildrenPatched.distinct.size) {
return a;
}

val patchedAggExpressions = a.aggregateExpressions.map { e =>
e.transformDown {
case e: Expression =>
Expand Down