You are viewing a plain text version of this content. The canonical link for it is here.
Posted to reviews@spark.apache.org by "kelvinjian-db (via GitHub)" <gi...@apache.org> on 2024/04/12 23:43:40 UTC
[PR] [SPARK-47839][SQL] Fix aggregate bug in RewriteWithExpression [spark]
kelvinjian-db opened a new pull request, #46034:
URL: https://github.com/apache/spark/pull/46034
<!--
Thanks for sending a pull request! Here are some tips for you:
1. If this is your first time, please read our contributor guidelines: https://spark.apache.org/contributing.html
2. Ensure you have added or run the appropriate tests for your PR: https://spark.apache.org/developer-tools.html
3. If the PR is unfinished, add '[WIP]' in your PR title, e.g., '[WIP][SPARK-XXXX] Your PR title ...'.
4. Be sure to keep the PR description updated to reflect all changes.
5. Please write your PR title to summarize what this PR proposes.
6. If possible, provide a concise example to reproduce the issue for a faster review.
7. If you want to add a new configuration, please read the guideline first for naming configurations in
'core/src/main/scala/org/apache/spark/internal/config/ConfigEntry.scala'.
8. If you want to add or modify an error type or message, please read the guideline first in
'common/utils/src/main/resources/error/README.md'.
-->
### What changes were proposed in this pull request?
<!--
Please clarify what changes you are proposing. The purpose of this section is to outline the changes and how this PR fixes the issue.
If possible, please consider writing useful notes for better and faster reviews in your PR. See the examples below.
1. If you refactor some codes with changing classes, showing the class hierarchy will help reviewers.
2. If you fix some SQL features, you can provide some references of other DBMSes.
3. If there is design documentation, please add the link.
4. If there is a discussion in the mailing list, please add the link.
-->
- Fixes a bug where `RewriteWithExpression` can rewrite an `Aggregate` into an invalid one. The fix is done by separating out the "result expressions" from the "aggregate expressions" in the `Aggregate` node, and rewriting them separately.
- Some QOL improvements around `With`:
- Fix aliases created by `With` expression to use the `CommonExpressionId` to avoid duplicate aliases
- Created new `With.create` helper function to make unit tests a little cleaner to write
- Implemented `QueryPlan.transformUpWithSubqueriesAndPruning` that the new logic depends on
### Why are the changes needed?
<!--
Please clarify why the changes are needed. For instance,
1. If you propose a new API, clarify the use case for a new API.
2. If you fix a bug, you can clarify why it is a bug.
-->
See [JIRA ticket](https://issues.apache.org/jira/browse/SPARK-47839) for more details on the bug that this fixes.
### Does this PR introduce _any_ user-facing change?
<!--
Note that it means *any* user-facing change including all aspects such as the documentation fix.
If yes, please clarify the previous behavior and the change this PR proposes - provide the console output, description and/or an example to show the behavior difference if possible.
If possible, please also clarify if this is a user-facing change compared to the released Spark versions or within the unreleased branches such as master.
If no, write 'No'.
-->
No.
### How was this patch tested?
<!--
If tests were added, say they were added here. Please make sure to add some test cases that check the changes thoroughly including negative and positive cases if possible.
If it was tested in a way different from regular unit tests, please clarify how you tested step by step, ideally copy and paste-able, so that other reviewers can test and check, and descendants can verify in the future.
If tests were not added, please describe why they were not added and/or why it was difficult to add.
If benchmark tests were added, please run the benchmarks in GitHub Actions for the consistent environment, and the instructions could accord to: https://spark.apache.org/developer-tools.html#github-workflow-benchmarks.
-->
Added new unit tests.
### Was this patch authored or co-authored using generative AI tooling?
<!--
If generative AI tooling has been used in the process of authoring this patch, please include the
phrase: 'Generated-by: ' followed by the name of the tool and its version.
If no, write 'No'.
Please refer to the [ASF Generative Tooling Guidance](https://www.apache.org/legal/generative-tooling.html) for details.
-->
No.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47839][SQL] Fix aggregate bug in RewriteWithExpression [spark]
Posted by "kelvinjian-db (via GitHub)" <gi...@apache.org>.
kelvinjian-db commented on code in PR #46034:
URL: https://github.com/apache/spark/pull/46034#discussion_r1566230161
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/RewriteWithExpressionSuite.scala:
##########
@@ -29,25 +29,29 @@ import org.apache.spark.sql.types.IntegerType
class RewriteWithExpressionSuite extends PlanTest {
object Optimizer extends RuleExecutor[LogicalPlan] {
- val batches = Batch("Rewrite With expression", Once, RewriteWithExpression) :: Nil
+ val batches = Batch("Rewrite With expression", Once,
+ PullOutGroupingExpressions,
+ RewriteWithExpression) :: Nil
}
private val testRelation = LocalRelation($"a".int, $"b".int)
private val testRelation2 = LocalRelation($"x".int, $"y".int)
test("simple common expression") {
val a = testRelation.output.head
- val commonExprDef = CommonExpressionDef(a)
- val ref = new CommonExpressionRef(commonExprDef)
- val plan = testRelation.select(With(ref + ref, Seq(commonExprDef)).as("col"))
+ val expr = With.create((a, 0)) { case Seq(ref) =>
Review Comment:
~~good point, i think it doesn't, i'll remove these~~ never mind, looks like it does matter since the IDs are used in the aliases due to [this change](https://github.com/apache/spark/pull/46034/files#diff-1eba63e1edc5cae9e89b8e326d7f5797ae5dc9314707268760484a3ed3d57038L96-R117). i originally made that change since it was possible to introduce duplicate alias names using `index`, whereas using the common expression ID was a more natural ID to use in the alias name. however, if the aliases are only used for bookkeeping, i can revert that 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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47839][SQL] Fix aggregate bug in RewriteWithExpression [spark]
Posted by "cloud-fan (via GitHub)" <gi...@apache.org>.
cloud-fan commented on code in PR #46034:
URL: https://github.com/apache/spark/pull/46034#discussion_r1568259745
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteWithExpression.scala:
##########
@@ -21,36 +21,68 @@ import scala.collection.mutable
import org.apache.spark.SparkException
import org.apache.spark.sql.catalyst.expressions._
-import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, PlanHelper, Project}
+import org.apache.spark.sql.catalyst.planning.PhysicalAggregation
+import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, LogicalPlan, PlanHelper, Project}
import org.apache.spark.sql.catalyst.rules.Rule
import org.apache.spark.sql.catalyst.trees.TreePattern.{COMMON_EXPR_REF, WITH_EXPRESSION}
/**
* Rewrites the `With` expressions by adding a `Project` to pre-evaluate the common expressions, or
* just inline them if they are cheap.
*
+ * Since this rule can introduce new `Project` operators, it is advised to run [[CollapseProject]]
+ * after this rule.
+ *
* Note: For now we only use `With` in a few `RuntimeReplaceable` expressions. If we expand its
* usage, we should support aggregate/window functions as well.
*/
object RewriteWithExpression extends Rule[LogicalPlan] {
override def apply(plan: LogicalPlan): LogicalPlan = {
- plan.transformDownWithSubqueriesAndPruning(_.containsPattern(WITH_EXPRESSION)) {
+ plan.transformUpWithSubqueriesAndPruning(_.containsPattern(WITH_EXPRESSION)) {
+ // For aggregates, separate the computation of the aggregations themselves from the final
+ // result by moving the final result computation into a projection above it. This prevents
+ // this rule from producing an invalid Aggregate operator.
+ case p @ PhysicalAggregation(
+ groupingExpressions, aggregateExpressions, resultExpressions, child)
+ if p.expressions.exists(_.containsPattern(WITH_EXPRESSION)) =>
+ // There should not be dangling common expression references in the aggregate expressions.
+ // This can happen if a With is created with an aggregate function in its child.
+ assert(!aggregateExpressions.exists(ae =>
Review Comment:
Shall we do the assert in the constructor of `With`, to fail earlier?
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47839][SQL] Fix aggregate bug in RewriteWithExpression [spark]
Posted by "cloud-fan (via GitHub)" <gi...@apache.org>.
cloud-fan commented on code in PR #46034:
URL: https://github.com/apache/spark/pull/46034#discussion_r1565727826
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteWithExpression.scala:
##########
@@ -21,36 +21,57 @@ import scala.collection.mutable
import org.apache.spark.SparkException
import org.apache.spark.sql.catalyst.expressions._
-import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, PlanHelper, Project}
+import org.apache.spark.sql.catalyst.planning.PhysicalAggregation
+import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, LogicalPlan, PlanHelper, Project}
import org.apache.spark.sql.catalyst.rules.Rule
import org.apache.spark.sql.catalyst.trees.TreePattern.{COMMON_EXPR_REF, WITH_EXPRESSION}
/**
* Rewrites the `With` expressions by adding a `Project` to pre-evaluate the common expressions, or
* just inline them if they are cheap.
*
+ * Since this rule can introduce new `Project` operators, it is advised to run [[CollapseProject]]
+ * after this rule.
+ *
* Note: For now we only use `With` in a few `RuntimeReplaceable` expressions. If we expand its
* usage, we should support aggregate/window functions as well.
*/
object RewriteWithExpression extends Rule[LogicalPlan] {
override def apply(plan: LogicalPlan): LogicalPlan = {
- plan.transformDownWithSubqueriesAndPruning(_.containsPattern(WITH_EXPRESSION)) {
+ plan.transformUpWithSubqueriesAndPruning(_.containsPattern(WITH_EXPRESSION)) {
+ case p @ PhysicalAggregation(
+ groupingExpressions, aggregateExpressions, resultExpressions, child)
+ if p.expressions.exists(_.containsPattern(WITH_EXPRESSION)) =>
+ // For aggregates, separate computation of the aggregations themselves from the final
+ // result by moving the final result computation into a projection above. This prevents
+ // this rule from producing an invalid Aggregate operator.
+ // TODO: the names of these aliases will become outdated after the rewrite
+ val aggExprs = aggregateExpressions.map(ae => Alias(ae, ae.toString)(ae.resultId))
+ // Rewrite the projection and the aggregate separately and then piece them together.
+ val agg = Aggregate(groupingExpressions, groupingExpressions ++ aggExprs, child)
+ val rewrittenAgg = applyInternal(agg)
+ val proj = Project(resultExpressions, rewrittenAgg)
+ applyInternal(proj)
case p if p.expressions.exists(_.containsPattern(WITH_EXPRESSION)) =>
- val inputPlans = p.children.toArray
- var newPlan: LogicalPlan = p.mapExpressions { expr =>
- rewriteWithExprAndInputPlans(expr, inputPlans)
- }
- newPlan = newPlan.withNewChildren(inputPlans.toIndexedSeq)
- // Since we add extra Projects with extra columns to pre-evaluate the common expressions,
- // the current operator may have extra columns if it inherits the output columns from its
- // child, and we need to project away the extra columns to keep the plan schema unchanged.
- assert(p.output.length <= newPlan.output.length)
- if (p.output.length < newPlan.output.length) {
- assert(p.outputSet.subsetOf(newPlan.outputSet))
- Project(p.output, newPlan)
- } else {
- newPlan
- }
+ applyInternal(p)
+ }
+ }
+
+ private def applyInternal(p: LogicalPlan): LogicalPlan = {
+ val inputPlans = p.children.toArray
+ var newPlan: LogicalPlan = p.mapExpressions { expr =>
+ rewriteWithExprAndInputPlans(expr, inputPlans)
+ }
+ newPlan = newPlan.withNewChildren(inputPlans)
Review Comment:
the children do not change, why calling `withNewChildren`?
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47839][SQL] Fix aggregate bug in RewriteWithExpression [spark]
Posted by "kelvinjian-db (via GitHub)" <gi...@apache.org>.
kelvinjian-db commented on PR #46034:
URL: https://github.com/apache/spark/pull/46034#issuecomment-2052701361
cc @cloud-fan @jchen5
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47839][SQL] Fix aggregate bug in RewriteWithExpression [spark]
Posted by "kelvinjian-db (via GitHub)" <gi...@apache.org>.
kelvinjian-db commented on code in PR #46034:
URL: https://github.com/apache/spark/pull/46034#discussion_r1567767274
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/RewriteWithExpressionSuite.scala:
##########
@@ -29,25 +29,29 @@ import org.apache.spark.sql.types.IntegerType
class RewriteWithExpressionSuite extends PlanTest {
object Optimizer extends RuleExecutor[LogicalPlan] {
- val batches = Batch("Rewrite With expression", Once, RewriteWithExpression) :: Nil
+ val batches = Batch("Rewrite With expression", Once,
+ PullOutGroupingExpressions,
+ RewriteWithExpression) :: Nil
}
private val testRelation = LocalRelation($"a".int, $"b".int)
private val testRelation2 = LocalRelation($"x".int, $"y".int)
test("simple common expression") {
val a = testRelation.output.head
- val commonExprDef = CommonExpressionDef(a)
- val ref = new CommonExpressionRef(commonExprDef)
- val plan = testRelation.select(With(ref + ref, Seq(commonExprDef)).as("col"))
+ val expr = With.create((a, 0)) { case Seq(ref) =>
Review Comment:
i ended up changing the tests to manually extract the common expression IDs so we don't need `With.create` anymore
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47839][SQL] Fix aggregate bug in RewriteWithExpression [spark]
Posted by "cloud-fan (via GitHub)" <gi...@apache.org>.
cloud-fan commented on code in PR #46034:
URL: https://github.com/apache/spark/pull/46034#discussion_r1565741611
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/RewriteWithExpressionSuite.scala:
##########
@@ -229,4 +236,85 @@ class RewriteWithExpressionSuite extends PlanTest {
.analyze
)
}
+
+ test("WITH expression in grouping exprs") {
+ val a = testRelation.output.head
+ val expr1 = With.create((a + 1, 0)) { case Seq(ref) =>
+ ref * ref
+ }
+ val expr2 = With.create((a + 1, 1)) { case Seq(ref) =>
+ ref * ref
+ }
+ val expr3 = With.create((a + 1, 2)) { case Seq(ref) =>
+ ref * ref
+ }
+ val plan = testRelation.groupBy(expr1)(
+ (expr2 + 2).as("col1"),
+ count(expr3 - 3).as("col2")
+ )
+ val commonExpr1Name = "_common_expr_0"
+ // Note that _common_expr_1 gets deduplicated by PullOutGroupingExpressions.
+ val commonExpr2Name = "_common_expr_2"
+ val groupingExprName = "_groupingexpression"
+ val countAlias = count(expr3 - 3).toString
+ comparePlans(
+ Optimizer.execute(plan),
+ testRelation
+ .select(testRelation.output :+ (a + 1).as(commonExpr1Name): _*)
+ .select(testRelation.output :+
+ ($"$commonExpr1Name" * $"$commonExpr1Name").as(groupingExprName): _*)
+ .select(testRelation.output ++ Seq($"$groupingExprName", (a + 1).as(commonExpr2Name)): _*)
+ .groupBy($"$groupingExprName")(
+ $"$groupingExprName",
+ count($"$commonExpr2Name" * $"$commonExpr2Name" - 3).as(countAlias)
+ )
+ .select(($"$groupingExprName" + 2).as("col1"), $"`$countAlias`".as("col2"))
+ .analyze
+ )
+ // Running CollapseProject after the rule cleans up the unnecessary projections.
+ comparePlans(
+ CollapseProject(Optimizer.execute(plan)),
+ testRelation
+ .select(testRelation.output :+ (a + 1).as(commonExpr1Name): _*)
+ .select(testRelation.output ++ Seq(
+ ($"$commonExpr1Name" * $"$commonExpr1Name").as(groupingExprName),
+ (a + 1).as(commonExpr2Name)): _*)
+ .groupBy($"$groupingExprName")(
+ ($"$groupingExprName" + 2).as("col1"),
+ count($"$commonExpr2Name" * $"$commonExpr2Name" - 3).as("col2")
+ )
+ .analyze
+ )
+ }
+
+ test("WITH expression in aggregate exprs") {
+ val Seq(a, b) = testRelation.output
+ val expr1 = With.create((a + 1, 0)) { case Seq(ref) =>
+ ref * ref
Review Comment:
can we test aggregate function as the common expression?
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47839][SQL] Fix aggregate bug in RewriteWithExpression [spark]
Posted by "cloud-fan (via GitHub)" <gi...@apache.org>.
cloud-fan commented on PR #46034:
URL: https://github.com/apache/spark/pull/46034#issuecomment-2062852813
thanks, merging to master!
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47839][SQL] Fix aggregate bug in RewriteWithExpression [spark]
Posted by "kelvinjian-db (via GitHub)" <gi...@apache.org>.
kelvinjian-db commented on code in PR #46034:
URL: https://github.com/apache/spark/pull/46034#discussion_r1566227660
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteWithExpression.scala:
##########
@@ -21,36 +21,57 @@ import scala.collection.mutable
import org.apache.spark.SparkException
import org.apache.spark.sql.catalyst.expressions._
-import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, PlanHelper, Project}
+import org.apache.spark.sql.catalyst.planning.PhysicalAggregation
+import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, LogicalPlan, PlanHelper, Project}
import org.apache.spark.sql.catalyst.rules.Rule
import org.apache.spark.sql.catalyst.trees.TreePattern.{COMMON_EXPR_REF, WITH_EXPRESSION}
/**
* Rewrites the `With` expressions by adding a `Project` to pre-evaluate the common expressions, or
* just inline them if they are cheap.
*
+ * Since this rule can introduce new `Project` operators, it is advised to run [[CollapseProject]]
+ * after this rule.
+ *
* Note: For now we only use `With` in a few `RuntimeReplaceable` expressions. If we expand its
* usage, we should support aggregate/window functions as well.
*/
object RewriteWithExpression extends Rule[LogicalPlan] {
override def apply(plan: LogicalPlan): LogicalPlan = {
- plan.transformDownWithSubqueriesAndPruning(_.containsPattern(WITH_EXPRESSION)) {
+ plan.transformUpWithSubqueriesAndPruning(_.containsPattern(WITH_EXPRESSION)) {
+ case p @ PhysicalAggregation(
+ groupingExpressions, aggregateExpressions, resultExpressions, child)
+ if p.expressions.exists(_.containsPattern(WITH_EXPRESSION)) =>
+ // For aggregates, separate computation of the aggregations themselves from the final
+ // result by moving the final result computation into a projection above. This prevents
+ // this rule from producing an invalid Aggregate operator.
+ // TODO: the names of these aliases will become outdated after the rewrite
+ val aggExprs = aggregateExpressions.map(ae => Alias(ae, ae.toString)(ae.resultId))
Review Comment:
good idea, i'll change this
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47839][SQL] Fix aggregate bug in RewriteWithExpression [spark]
Posted by "cloud-fan (via GitHub)" <gi...@apache.org>.
cloud-fan commented on code in PR #46034:
URL: https://github.com/apache/spark/pull/46034#discussion_r1565732326
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/RewriteWithExpressionSuite.scala:
##########
@@ -29,25 +29,29 @@ import org.apache.spark.sql.types.IntegerType
class RewriteWithExpressionSuite extends PlanTest {
object Optimizer extends RuleExecutor[LogicalPlan] {
- val batches = Batch("Rewrite With expression", Once, RewriteWithExpression) :: Nil
+ val batches = Batch("Rewrite With expression", Once,
+ PullOutGroupingExpressions,
+ RewriteWithExpression) :: Nil
}
private val testRelation = LocalRelation($"a".int, $"b".int)
private val testRelation2 = LocalRelation($"x".int, $"y".int)
test("simple common expression") {
val a = testRelation.output.head
- val commonExprDef = CommonExpressionDef(a)
- val ref = new CommonExpressionRef(commonExprDef)
- val plan = testRelation.select(With(ref + ref, Seq(commonExprDef)).as("col"))
+ val expr = With.create((a, 0)) { case Seq(ref) =>
Review Comment:
does the id matter for this test?
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47839][SQL] Fix aggregate bug in RewriteWithExpression [spark]
Posted by "kelvinjian-db (via GitHub)" <gi...@apache.org>.
kelvinjian-db commented on code in PR #46034:
URL: https://github.com/apache/spark/pull/46034#discussion_r1566233143
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/RewriteWithExpressionSuite.scala:
##########
@@ -229,4 +236,85 @@ class RewriteWithExpressionSuite extends PlanTest {
.analyze
)
}
+
+ test("WITH expression in grouping exprs") {
+ val a = testRelation.output.head
+ val expr1 = With.create((a + 1, 0)) { case Seq(ref) =>
+ ref * ref
+ }
+ val expr2 = With.create((a + 1, 1)) { case Seq(ref) =>
+ ref * ref
+ }
+ val expr3 = With.create((a + 1, 2)) { case Seq(ref) =>
+ ref * ref
+ }
+ val plan = testRelation.groupBy(expr1)(
+ (expr2 + 2).as("col1"),
+ count(expr3 - 3).as("col2")
+ )
+ val commonExpr1Name = "_common_expr_0"
+ // Note that _common_expr_1 gets deduplicated by PullOutGroupingExpressions.
+ val commonExpr2Name = "_common_expr_2"
+ val groupingExprName = "_groupingexpression"
+ val countAlias = count(expr3 - 3).toString
+ comparePlans(
+ Optimizer.execute(plan),
+ testRelation
+ .select(testRelation.output :+ (a + 1).as(commonExpr1Name): _*)
+ .select(testRelation.output :+
+ ($"$commonExpr1Name" * $"$commonExpr1Name").as(groupingExprName): _*)
+ .select(testRelation.output ++ Seq($"$groupingExprName", (a + 1).as(commonExpr2Name)): _*)
+ .groupBy($"$groupingExprName")(
+ $"$groupingExprName",
+ count($"$commonExpr2Name" * $"$commonExpr2Name" - 3).as(countAlias)
+ )
+ .select(($"$groupingExprName" + 2).as("col1"), $"`$countAlias`".as("col2"))
+ .analyze
+ )
+ // Running CollapseProject after the rule cleans up the unnecessary projections.
+ comparePlans(
+ CollapseProject(Optimizer.execute(plan)),
+ testRelation
+ .select(testRelation.output :+ (a + 1).as(commonExpr1Name): _*)
+ .select(testRelation.output ++ Seq(
+ ($"$commonExpr1Name" * $"$commonExpr1Name").as(groupingExprName),
+ (a + 1).as(commonExpr2Name)): _*)
+ .groupBy($"$groupingExprName")(
+ ($"$groupingExprName" + 2).as("col1"),
+ count($"$commonExpr2Name" * $"$commonExpr2Name" - 3).as("col2")
+ )
+ .analyze
+ )
+ }
+
+ test("WITH expression in aggregate exprs") {
Review Comment:
doesn't the test above test WITH in both grouping and aggregate expressions? the test here is for testing the motivating example mentioned in https://issues.apache.org/jira/browse/SPARK-47839
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47839][SQL] Fix aggregate bug in RewriteWithExpression [spark]
Posted by "cloud-fan (via GitHub)" <gi...@apache.org>.
cloud-fan closed pull request #46034: [SPARK-47839][SQL] Fix aggregate bug in RewriteWithExpression
URL: https://github.com/apache/spark/pull/46034
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47839][SQL] Fix aggregate bug in RewriteWithExpression [spark]
Posted by "cloud-fan (via GitHub)" <gi...@apache.org>.
cloud-fan commented on PR #46034:
URL: https://github.com/apache/spark/pull/46034#issuecomment-2060437664
The test fails: `org.apache.spark.sql.connect.ProtoToParsedPlanTestSuite`
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47839][SQL] Fix aggregate bug in RewriteWithExpression [spark]
Posted by "cloud-fan (via GitHub)" <gi...@apache.org>.
cloud-fan commented on code in PR #46034:
URL: https://github.com/apache/spark/pull/46034#discussion_r1565739557
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteWithExpression.scala:
##########
@@ -21,36 +21,57 @@ import scala.collection.mutable
import org.apache.spark.SparkException
import org.apache.spark.sql.catalyst.expressions._
-import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, PlanHelper, Project}
+import org.apache.spark.sql.catalyst.planning.PhysicalAggregation
+import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, LogicalPlan, PlanHelper, Project}
import org.apache.spark.sql.catalyst.rules.Rule
import org.apache.spark.sql.catalyst.trees.TreePattern.{COMMON_EXPR_REF, WITH_EXPRESSION}
/**
* Rewrites the `With` expressions by adding a `Project` to pre-evaluate the common expressions, or
* just inline them if they are cheap.
*
+ * Since this rule can introduce new `Project` operators, it is advised to run [[CollapseProject]]
+ * after this rule.
+ *
* Note: For now we only use `With` in a few `RuntimeReplaceable` expressions. If we expand its
* usage, we should support aggregate/window functions as well.
*/
object RewriteWithExpression extends Rule[LogicalPlan] {
override def apply(plan: LogicalPlan): LogicalPlan = {
- plan.transformDownWithSubqueriesAndPruning(_.containsPattern(WITH_EXPRESSION)) {
+ plan.transformUpWithSubqueriesAndPruning(_.containsPattern(WITH_EXPRESSION)) {
+ case p @ PhysicalAggregation(
+ groupingExpressions, aggregateExpressions, resultExpressions, child)
+ if p.expressions.exists(_.containsPattern(WITH_EXPRESSION)) =>
+ // For aggregates, separate computation of the aggregations themselves from the final
+ // result by moving the final result computation into a projection above. This prevents
+ // this rule from producing an invalid Aggregate operator.
+ // TODO: the names of these aliases will become outdated after the rewrite
+ val aggExprs = aggregateExpressions.map(ae => Alias(ae, ae.toString)(ae.resultId))
Review Comment:
The alias name doesn't matter as it's only for internal bookkeeping. `.toString` can be super long if the aggregate function input is a complex expression. Shall we follow `PullOutGroupingExpressions` and use consistent naming? like `_aggregateExpression`?
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47839][SQL] Fix aggregate bug in RewriteWithExpression [spark]
Posted by "cloud-fan (via GitHub)" <gi...@apache.org>.
cloud-fan commented on code in PR #46034:
URL: https://github.com/apache/spark/pull/46034#discussion_r1565740751
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/RewriteWithExpressionSuite.scala:
##########
@@ -229,4 +236,85 @@ class RewriteWithExpressionSuite extends PlanTest {
.analyze
)
}
+
+ test("WITH expression in grouping exprs") {
+ val a = testRelation.output.head
+ val expr1 = With.create((a + 1, 0)) { case Seq(ref) =>
+ ref * ref
+ }
+ val expr2 = With.create((a + 1, 1)) { case Seq(ref) =>
+ ref * ref
+ }
+ val expr3 = With.create((a + 1, 2)) { case Seq(ref) =>
+ ref * ref
+ }
+ val plan = testRelation.groupBy(expr1)(
+ (expr2 + 2).as("col1"),
+ count(expr3 - 3).as("col2")
+ )
+ val commonExpr1Name = "_common_expr_0"
+ // Note that _common_expr_1 gets deduplicated by PullOutGroupingExpressions.
+ val commonExpr2Name = "_common_expr_2"
+ val groupingExprName = "_groupingexpression"
+ val countAlias = count(expr3 - 3).toString
+ comparePlans(
+ Optimizer.execute(plan),
+ testRelation
+ .select(testRelation.output :+ (a + 1).as(commonExpr1Name): _*)
+ .select(testRelation.output :+
+ ($"$commonExpr1Name" * $"$commonExpr1Name").as(groupingExprName): _*)
+ .select(testRelation.output ++ Seq($"$groupingExprName", (a + 1).as(commonExpr2Name)): _*)
+ .groupBy($"$groupingExprName")(
+ $"$groupingExprName",
+ count($"$commonExpr2Name" * $"$commonExpr2Name" - 3).as(countAlias)
+ )
+ .select(($"$groupingExprName" + 2).as("col1"), $"`$countAlias`".as("col2"))
+ .analyze
+ )
+ // Running CollapseProject after the rule cleans up the unnecessary projections.
+ comparePlans(
+ CollapseProject(Optimizer.execute(plan)),
+ testRelation
+ .select(testRelation.output :+ (a + 1).as(commonExpr1Name): _*)
+ .select(testRelation.output ++ Seq(
+ ($"$commonExpr1Name" * $"$commonExpr1Name").as(groupingExprName),
+ (a + 1).as(commonExpr2Name)): _*)
+ .groupBy($"$groupingExprName")(
+ ($"$groupingExprName" + 2).as("col1"),
+ count($"$commonExpr2Name" * $"$commonExpr2Name" - 3).as("col2")
+ )
+ .analyze
+ )
+ }
+
+ test("WITH expression in aggregate exprs") {
Review Comment:
can we merge the two tests and have WITH in both grouping and aggregate expressions?
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47839][SQL] Fix aggregate bug in RewriteWithExpression [spark]
Posted by "kelvinjian-db (via GitHub)" <gi...@apache.org>.
kelvinjian-db commented on code in PR #46034:
URL: https://github.com/apache/spark/pull/46034#discussion_r1566228919
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteWithExpression.scala:
##########
@@ -21,36 +21,57 @@ import scala.collection.mutable
import org.apache.spark.SparkException
import org.apache.spark.sql.catalyst.expressions._
-import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, PlanHelper, Project}
+import org.apache.spark.sql.catalyst.planning.PhysicalAggregation
+import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, LogicalPlan, PlanHelper, Project}
import org.apache.spark.sql.catalyst.rules.Rule
import org.apache.spark.sql.catalyst.trees.TreePattern.{COMMON_EXPR_REF, WITH_EXPRESSION}
/**
* Rewrites the `With` expressions by adding a `Project` to pre-evaluate the common expressions, or
* just inline them if they are cheap.
*
+ * Since this rule can introduce new `Project` operators, it is advised to run [[CollapseProject]]
+ * after this rule.
+ *
* Note: For now we only use `With` in a few `RuntimeReplaceable` expressions. If we expand its
* usage, we should support aggregate/window functions as well.
*/
object RewriteWithExpression extends Rule[LogicalPlan] {
override def apply(plan: LogicalPlan): LogicalPlan = {
- plan.transformDownWithSubqueriesAndPruning(_.containsPattern(WITH_EXPRESSION)) {
+ plan.transformUpWithSubqueriesAndPruning(_.containsPattern(WITH_EXPRESSION)) {
+ case p @ PhysicalAggregation(
+ groupingExpressions, aggregateExpressions, resultExpressions, child)
+ if p.expressions.exists(_.containsPattern(WITH_EXPRESSION)) =>
+ // For aggregates, separate computation of the aggregations themselves from the final
+ // result by moving the final result computation into a projection above. This prevents
+ // this rule from producing an invalid Aggregate operator.
+ // TODO: the names of these aliases will become outdated after the rewrite
+ val aggExprs = aggregateExpressions.map(ae => Alias(ae, ae.toString)(ae.resultId))
+ // Rewrite the projection and the aggregate separately and then piece them together.
+ val agg = Aggregate(groupingExpressions, groupingExpressions ++ aggExprs, child)
+ val rewrittenAgg = applyInternal(agg)
+ val proj = Project(resultExpressions, rewrittenAgg)
+ applyInternal(proj)
case p if p.expressions.exists(_.containsPattern(WITH_EXPRESSION)) =>
- val inputPlans = p.children.toArray
- var newPlan: LogicalPlan = p.mapExpressions { expr =>
- rewriteWithExprAndInputPlans(expr, inputPlans)
- }
- newPlan = newPlan.withNewChildren(inputPlans.toIndexedSeq)
- // Since we add extra Projects with extra columns to pre-evaluate the common expressions,
- // the current operator may have extra columns if it inherits the output columns from its
- // child, and we need to project away the extra columns to keep the plan schema unchanged.
- assert(p.output.length <= newPlan.output.length)
- if (p.output.length < newPlan.output.length) {
- assert(p.outputSet.subsetOf(newPlan.outputSet))
- Project(p.output, newPlan)
- } else {
- newPlan
- }
+ applyInternal(p)
+ }
+ }
+
+ private def applyInternal(p: LogicalPlan): LogicalPlan = {
+ val inputPlans = p.children.toArray
+ var newPlan: LogicalPlan = p.mapExpressions { expr =>
+ rewriteWithExprAndInputPlans(expr, inputPlans)
+ }
+ newPlan = newPlan.withNewChildren(inputPlans)
Review Comment:
my understanding is that we need to call `withNewChildren` because `rewriteWithExprAndInputPlans` puts the new projections in `inputPlans` and we need to update our current plan with them?
(btw this was largely just refactored from the previous code)
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47839][SQL] Fix aggregate bug in RewriteWithExpression [spark]
Posted by "kelvinjian-db (via GitHub)" <gi...@apache.org>.
kelvinjian-db commented on code in PR #46034:
URL: https://github.com/apache/spark/pull/46034#discussion_r1566233683
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/RewriteWithExpressionSuite.scala:
##########
@@ -229,4 +236,85 @@ class RewriteWithExpressionSuite extends PlanTest {
.analyze
)
}
+
+ test("WITH expression in grouping exprs") {
+ val a = testRelation.output.head
+ val expr1 = With.create((a + 1, 0)) { case Seq(ref) =>
+ ref * ref
+ }
+ val expr2 = With.create((a + 1, 1)) { case Seq(ref) =>
+ ref * ref
+ }
+ val expr3 = With.create((a + 1, 2)) { case Seq(ref) =>
+ ref * ref
+ }
+ val plan = testRelation.groupBy(expr1)(
+ (expr2 + 2).as("col1"),
+ count(expr3 - 3).as("col2")
+ )
+ val commonExpr1Name = "_common_expr_0"
+ // Note that _common_expr_1 gets deduplicated by PullOutGroupingExpressions.
+ val commonExpr2Name = "_common_expr_2"
+ val groupingExprName = "_groupingexpression"
+ val countAlias = count(expr3 - 3).toString
+ comparePlans(
+ Optimizer.execute(plan),
+ testRelation
+ .select(testRelation.output :+ (a + 1).as(commonExpr1Name): _*)
+ .select(testRelation.output :+
+ ($"$commonExpr1Name" * $"$commonExpr1Name").as(groupingExprName): _*)
+ .select(testRelation.output ++ Seq($"$groupingExprName", (a + 1).as(commonExpr2Name)): _*)
+ .groupBy($"$groupingExprName")(
+ $"$groupingExprName",
+ count($"$commonExpr2Name" * $"$commonExpr2Name" - 3).as(countAlias)
+ )
+ .select(($"$groupingExprName" + 2).as("col1"), $"`$countAlias`".as("col2"))
+ .analyze
+ )
+ // Running CollapseProject after the rule cleans up the unnecessary projections.
+ comparePlans(
+ CollapseProject(Optimizer.execute(plan)),
+ testRelation
+ .select(testRelation.output :+ (a + 1).as(commonExpr1Name): _*)
+ .select(testRelation.output ++ Seq(
+ ($"$commonExpr1Name" * $"$commonExpr1Name").as(groupingExprName),
+ (a + 1).as(commonExpr2Name)): _*)
+ .groupBy($"$groupingExprName")(
+ ($"$groupingExprName" + 2).as("col1"),
+ count($"$commonExpr2Name" * $"$commonExpr2Name" - 3).as("col2")
+ )
+ .analyze
+ )
+ }
+
+ test("WITH expression in aggregate exprs") {
+ val Seq(a, b) = testRelation.output
+ val expr1 = With.create((a + 1, 0)) { case Seq(ref) =>
+ ref * ref
Review Comment:
yes let me add a test for that
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47839][SQL] Fix aggregate bug in RewriteWithExpression [spark]
Posted by "kelvinjian-db (via GitHub)" <gi...@apache.org>.
kelvinjian-db commented on code in PR #46034:
URL: https://github.com/apache/spark/pull/46034#discussion_r1566230161
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/RewriteWithExpressionSuite.scala:
##########
@@ -29,25 +29,29 @@ import org.apache.spark.sql.types.IntegerType
class RewriteWithExpressionSuite extends PlanTest {
object Optimizer extends RuleExecutor[LogicalPlan] {
- val batches = Batch("Rewrite With expression", Once, RewriteWithExpression) :: Nil
+ val batches = Batch("Rewrite With expression", Once,
+ PullOutGroupingExpressions,
+ RewriteWithExpression) :: Nil
}
private val testRelation = LocalRelation($"a".int, $"b".int)
private val testRelation2 = LocalRelation($"x".int, $"y".int)
test("simple common expression") {
val a = testRelation.output.head
- val commonExprDef = CommonExpressionDef(a)
- val ref = new CommonExpressionRef(commonExprDef)
- val plan = testRelation.select(With(ref + ref, Seq(commonExprDef)).as("col"))
+ val expr = With.create((a, 0)) { case Seq(ref) =>
Review Comment:
good point, i think it doesn't, i'll remove these
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org