You are viewing a plain text version of this content. The canonical link for it is here.
Posted to reviews@spark.apache.org by "peter-toth (via GitHub)" <gi...@apache.org> on 2025/11/02 12:24:55 UTC
[PR] [SPARK-54136][SQL] Extract plan merging logic from `MergeScalarSubqueries` to `PlanMerger` [spark]
peter-toth opened a new pull request, #52835:
URL: https://github.com/apache/spark/pull/52835
### What changes were proposed in this pull request?
This PR extracts the plan merging logic from `MergeScalarSubqueries` to `PlanMerger` so as to other rules can reuse it.
While the plan merging logic is extracted without modification to `PlanMerger`, `MergeScalarSubqueries` required a significant adjustment. This is because [SPARK-40618](https://issues.apache.org/jira/browse/SPARK-40618) / https://github.com/apache/spark/pull/38093 added subquery reference tracking so as to avoid trying to merge a subquery to any of its nested subqueries. This kind of reference trancking doesn't work well with a general `PlanMerger` so this PR modifies `MergeScalarSubqueries` to use a separate `PlanMerger`s by each subquery level.
### Why are the changes needed?
To be able to reuse plan merging logic.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
Existing UTs.
### Was this patch authored or co-authored using generative AI tooling?
Yes, Claude gave me suggestions to improve documentation.
--
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-54136][SQL] Extract plan merging logic from `MergeScalarSubqueries` to `PlanMerger` [spark]
Posted by "peter-toth (via GitHub)" <gi...@apache.org>.
peter-toth commented on code in PR #52835:
URL: https://github.com/apache/spark/pull/52835#discussion_r2485989939
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/MergeScalarSubqueriesSuite.scala:
##########
@@ -568,14 +568,14 @@ class MergeScalarSubqueriesSuite extends PlanTest {
val subquery5 = ScalarSubquery(testRelation.select((Symbol("a") + 2).as("a_plus2_2")))
val subquery6 = ScalarSubquery(testRelation.select(Symbol("b").as("b_2")))
val originalQuery = testRelation
- .select(
Review Comment:
This minor test chnage in the order of subqueries is needed because the modified `MergeScalarSubqueries` encounters them in different order. So we either change the expected query to expect `a_plus1_2` and `a_plus2_2` or reorder the subqueries in the query like here. I applied the latter.
--
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-54136][SQL] Extract plan merging logic from `MergeScalarSubqueries` to `PlanMerger` [spark]
Posted by "peter-toth (via GitHub)" <gi...@apache.org>.
peter-toth commented on code in PR #52835:
URL: https://github.com/apache/spark/pull/52835#discussion_r2485847706
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/MergeScalarSubqueries.scala:
##########
@@ -117,297 +116,126 @@ object MergeScalarSubqueries extends Rule[LogicalPlan] {
}
}
- /**
- * An item in the cache of merged scalar subqueries.
- *
- * @param plan The plan of a merged scalar subquery.
- * @param merged A flag to identify if this item is the result of merging subqueries.
- * Please note that `attributes.size == 1` doesn't always mean that the plan is not
- * merged as there can be subqueries that are different ([[checkIdenticalPlans]] is
- * false) due to an extra [[Project]] node in one of them. In that case
- * `attributes.size` remains 1 after merging, but the merged flag becomes true.
- * @param references A set of subquery indexes in the cache to track all (including transitive)
- * nested subqueries.
- */
- case class Header(
- plan: LogicalPlan,
- merged: Boolean,
- references: Set[Int])
-
private def extractCommonScalarSubqueries(plan: LogicalPlan) = {
- val cache = ArrayBuffer.empty[Header]
- val planWithReferences = insertReferences(plan, cache)
- cache.zipWithIndex.foreach { case (header, i) =>
- cache(i) = cache(i).copy(plan =
- if (header.merged) {
+ // Collect `ScalarSubquery` plans by level into `PlanMerger`s and insert references in place of
+ // `ScalarSubquery`s.
+ val planMergers = ArrayBuffer.empty[PlanMerger]
+ val planWithReferences = insertReferences(plan, planMergers)._1
+
+ // Traverse level by level and convert merged plans to `CTERelationDef`s and keep non-merged
+ // ones. While traversing replace references in plans back to `CTERelationRef`s or to original
+ // `ScalarSubquery`s. This is safe as a subquery plan at a level can reference only lower level
+ // other subqueries.
+ val subqueryPlansByLevel = ArrayBuffer.empty[IndexedSeq[LogicalPlan]]
+ planMergers.foreach { planMerger =>
+ val mergedPlans = planMerger.mergedPlans()
+ subqueryPlansByLevel += mergedPlans.map { mergedPlan =>
+ val planWithoutReferences = if (subqueryPlansByLevel.isEmpty) {
+ // Level 0 plans can't contain references
+ mergedPlan.plan
+ } else {
+ removeReferences(mergedPlan.plan, subqueryPlansByLevel)
+ }
+ if (mergedPlan.merged) {
CTERelationDef(
- createProject(header.plan.output, removeReferences(header.plan, cache)),
+ Project(
+ Seq(Alias(
+ CreateNamedStruct(
+ planWithoutReferences.output.flatMap(a => Seq(Literal(a.name), a))),
+ "mergedValue")()),
+ planWithoutReferences),
underSubquery = true)
} else {
- removeReferences(header.plan, cache)
- })
+ planWithoutReferences
+ }
+ }
}
- val newPlan = removeReferences(planWithReferences, cache)
- val subqueryCTEs = cache.filter(_.merged).map(_.plan.asInstanceOf[CTERelationDef])
+
+ // Replace references back to `CTERelationRef`s or to original `ScalarSubquery`s in the main
+ // plan.
+ val newPlan = removeReferences(planWithReferences, subqueryPlansByLevel)
+
+ // Add `CTERelationDef`s to the plan.
+ val subqueryCTEs = subqueryPlansByLevel.flatMap(_.collect { case cte: CTERelationDef => cte })
if (subqueryCTEs.nonEmpty) {
WithCTE(newPlan, subqueryCTEs.toSeq)
} else {
newPlan
}
}
- // First traversal builds up the cache and inserts `ScalarSubqueryReference`s to the plan.
- private def insertReferences(plan: LogicalPlan, cache: ArrayBuffer[Header]): LogicalPlan = {
- plan.transformUpWithSubqueries {
- case n => n.transformExpressionsUpWithPruning(_.containsAnyPattern(SCALAR_SUBQUERY)) {
- // The subquery could contain a hint that is not propagated once we cache it, but as a
- // non-correlated scalar subquery won't be turned into a Join the loss of hints is fine.
+ // First traversal inserts `ScalarSubqueryReference`s to the plan and tries to merge subquery
+ // plans by each level.
+ private def insertReferences(
+ plan: LogicalPlan,
+ planMergers: ArrayBuffer[PlanMerger]): (LogicalPlan, Int) = {
+ // The level of a subquery plan is maximum level of its inner subqueries + 1 or 0 if it has no
+ // inner subqueries.
+ var maxLevel = 0
+ val planWithReferences =
+ plan.transformAllExpressionsWithPruning(_.containsAnyPattern(SCALAR_SUBQUERY)) {
case s: ScalarSubquery if !s.isCorrelated && s.deterministic =>
- val (subqueryIndex, headerIndex) = cacheSubquery(s.plan, cache)
- ScalarSubqueryReference(subqueryIndex, headerIndex, s.dataType, s.exprId)
- }
- }
- }
-
- // Caching returns the index of the subquery in the cache and the index of scalar member in the
- // "Header".
- private def cacheSubquery(plan: LogicalPlan, cache: ArrayBuffer[Header]): (Int, Int) = {
Review Comment:
Briefly, `cacheSubquery()` became `PlanMerger.merge()` with some modifications.
As `PlanMerger.merge()` is kind of an internal API now, which can be used by other rules. It returns `MergeResult` instead of the tuple after this refactor.
The `Header` objects, which contined the information about the merged plans became `MergedPlan`s and can be accessed as `PlanMerger.mergedPlans()` by the rules.
The main difference between `cacheSubquery()` and the new `PlanMerger.merge()` is that I removed `Header.references` and the related safeguards like the `if !references.contains(subqueryIndex)` below.
We filled up `Header.references` and with nested subquery indices so as to avoid trying to merge a subquery with other subqueries that are nested into itself. E.g.
```
SELECT (
SELECT avg(a) FROM t WHERE c = ( -- subquery 1
SELECT sum(b) FROM t2) -- subquery 2 nested into subquery 1
)
)
```
Here it doesn't make sense to try merging the plan of `subquery 1` and `subquery 2` as the former contains the latter. The problem of trying to merge the 2 was not just it doesn't make sense, but in certain weird cases (see the example in [SPARK-40618](https://issues.apache.org/jira/browse/SPARK-40618)) the merge was successful and resulted in invalid plans.
But `PlanMerger` should be a general plan merging tool so tracking the nested subqueries in it is not the best way to deal with the problem (in some future rules it will be used to merge other than subqueries). This is why the refactored `MergeScalarSubqueries` uses a different approach. We have one `PlanMerger` object by each subquery level. E.g. `subquery 2` is a level 0 (leaf) subquery, while `subquery 1` is a level 1 subquery (because it has an inner level 0 subquery). So the 2 will be added to different `PlanMerger`s. This way a level n subquery is tried to merge with other level n subqueries only.
--
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-54136][SQL] Extract plan merging logic from `MergeScalarSubqueries` to `PlanMerger` [spark]
Posted by "cloud-fan (via GitHub)" <gi...@apache.org>.
cloud-fan commented on code in PR #52835:
URL: https://github.com/apache/spark/pull/52835#discussion_r2500787430
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/PlanMerger.scala:
##########
@@ -0,0 +1,303 @@
+/*
+ * 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.ArrayBuffer
+
+import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, AttributeMap, Expression, NamedExpression}
+import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression
+import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Filter, Join, LogicalPlan, Project}
+
+/**
+ * Result of attempting to merge a plan via [[PlanMerger.merge]].
+ *
+ * @param mergedPlan The resulting plan, either:
+ * - An existing cached plan (if identical match found)
+ * - A newly merged plan combining the input with a cached plan
+ * - The original input plan (if no merge was possible)
+ * @param mergedPlanIndex The index of this plan in the PlanMerger's cache.
+ * @param merged Whether the plan was merged with an existing cached plan (true) or
+ * is a new entry (false).
+ * @param outputMap Maps attributes from the input plan to corresponding attributes in
+ * `mergedPlan`. Used to rewrite expressions referencing the original plan
+ * to reference the merged plan instead.
+ */
+case class MergeResult(
+ mergedPlan: LogicalPlan,
Review Comment:
shall we use `mergedPlan: MergedPlan` to save the `merged: Boolean` parameter in this class?
--
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-54136][SQL] Extract plan merging logic from `MergeScalarSubqueries` to `PlanMerger` [spark]
Posted by "peter-toth (via GitHub)" <gi...@apache.org>.
peter-toth commented on PR #52835:
URL: https://github.com/apache/spark/pull/52835#issuecomment-3479841185
cc @beliefer as well.
The goal here is to extract the merging logic and reuse it in [SPARK-44571](https://issues.apache.org/jira/browse/SPARK-44571) / https://github.com/apache/spark/pull/42223 and [SPARK-43025](https://issues.apache.org/jira/browse/SPARK-43025) / https://github.com/apache/spark/pull/40661.
--
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-54136][SQL] Extract plan merging logic from `MergeScalarSubqueries` to `PlanMerger` [spark]
Posted by "peter-toth (via GitHub)" <gi...@apache.org>.
peter-toth commented on PR #52835:
URL: https://github.com/apache/spark/pull/52835#issuecomment-3505052975
Thank you for the review @cloud-fan and @dongjoon-hyun, really appreciated.
Merged to `master` (4.2.0).
--
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-54136][SQL] Extract plan merging logic from `MergeScalarSubqueries` to `PlanMerger` [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #52835:
URL: https://github.com/apache/spark/pull/52835#discussion_r2484972241
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/MergeScalarSubqueries.scala:
##########
@@ -117,297 +116,126 @@ object MergeScalarSubqueries extends Rule[LogicalPlan] {
}
}
- /**
- * An item in the cache of merged scalar subqueries.
- *
- * @param plan The plan of a merged scalar subquery.
- * @param merged A flag to identify if this item is the result of merging subqueries.
- * Please note that `attributes.size == 1` doesn't always mean that the plan is not
- * merged as there can be subqueries that are different ([[checkIdenticalPlans]] is
- * false) due to an extra [[Project]] node in one of them. In that case
- * `attributes.size` remains 1 after merging, but the merged flag becomes true.
- * @param references A set of subquery indexes in the cache to track all (including transitive)
- * nested subqueries.
- */
- case class Header(
- plan: LogicalPlan,
- merged: Boolean,
- references: Set[Int])
-
private def extractCommonScalarSubqueries(plan: LogicalPlan) = {
- val cache = ArrayBuffer.empty[Header]
- val planWithReferences = insertReferences(plan, cache)
- cache.zipWithIndex.foreach { case (header, i) =>
- cache(i) = cache(i).copy(plan =
- if (header.merged) {
+ // Collect `ScalarSubquery` plans by level into `PlanMerger`s and insert references in place of
+ // `ScalarSubquery`s.
+ val planMergers = ArrayBuffer.empty[PlanMerger]
+ val planWithReferences = insertReferences(plan, planMergers)._1
+
+ // Traverse level by level and convert merged plans to `CTERelationDef`s and keep non-merged
+ // ones. While traversing replace references in plans back to `CTERelationRef`s or to original
+ // `ScalarSubquery`s. This is safe as a subquery plan at a level can reference only lower level
+ // other subqueries.
+ val subqueryPlansByLevel = ArrayBuffer.empty[IndexedSeq[LogicalPlan]]
+ planMergers.foreach { planMerger =>
+ val mergedPlans = planMerger.mergedPlans()
+ subqueryPlansByLevel += mergedPlans.map { mergedPlan =>
+ val planWithoutReferences = if (subqueryPlansByLevel.isEmpty) {
+ // Level 0 plans can't contain references
+ mergedPlan.plan
+ } else {
+ removeReferences(mergedPlan.plan, subqueryPlansByLevel)
+ }
+ if (mergedPlan.merged) {
CTERelationDef(
- createProject(header.plan.output, removeReferences(header.plan, cache)),
+ Project(
Review Comment:
May I ask why this PR proposes to remove the existing `createProject` method and use `Project(...)` directly here? Is this better in a way?
--
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-54136][SQL] Extract plan merging logic from `MergeScalarSubqueries` to `PlanMerger` [spark]
Posted by "peter-toth (via GitHub)" <gi...@apache.org>.
peter-toth commented on code in PR #52835:
URL: https://github.com/apache/spark/pull/52835#discussion_r2485847706
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/MergeScalarSubqueries.scala:
##########
@@ -117,297 +116,126 @@ object MergeScalarSubqueries extends Rule[LogicalPlan] {
}
}
- /**
- * An item in the cache of merged scalar subqueries.
- *
- * @param plan The plan of a merged scalar subquery.
- * @param merged A flag to identify if this item is the result of merging subqueries.
- * Please note that `attributes.size == 1` doesn't always mean that the plan is not
- * merged as there can be subqueries that are different ([[checkIdenticalPlans]] is
- * false) due to an extra [[Project]] node in one of them. In that case
- * `attributes.size` remains 1 after merging, but the merged flag becomes true.
- * @param references A set of subquery indexes in the cache to track all (including transitive)
- * nested subqueries.
- */
- case class Header(
- plan: LogicalPlan,
- merged: Boolean,
- references: Set[Int])
-
private def extractCommonScalarSubqueries(plan: LogicalPlan) = {
- val cache = ArrayBuffer.empty[Header]
- val planWithReferences = insertReferences(plan, cache)
- cache.zipWithIndex.foreach { case (header, i) =>
- cache(i) = cache(i).copy(plan =
- if (header.merged) {
+ // Collect `ScalarSubquery` plans by level into `PlanMerger`s and insert references in place of
+ // `ScalarSubquery`s.
+ val planMergers = ArrayBuffer.empty[PlanMerger]
+ val planWithReferences = insertReferences(plan, planMergers)._1
+
+ // Traverse level by level and convert merged plans to `CTERelationDef`s and keep non-merged
+ // ones. While traversing replace references in plans back to `CTERelationRef`s or to original
+ // `ScalarSubquery`s. This is safe as a subquery plan at a level can reference only lower level
+ // other subqueries.
+ val subqueryPlansByLevel = ArrayBuffer.empty[IndexedSeq[LogicalPlan]]
+ planMergers.foreach { planMerger =>
+ val mergedPlans = planMerger.mergedPlans()
+ subqueryPlansByLevel += mergedPlans.map { mergedPlan =>
+ val planWithoutReferences = if (subqueryPlansByLevel.isEmpty) {
+ // Level 0 plans can't contain references
+ mergedPlan.plan
+ } else {
+ removeReferences(mergedPlan.plan, subqueryPlansByLevel)
+ }
+ if (mergedPlan.merged) {
CTERelationDef(
- createProject(header.plan.output, removeReferences(header.plan, cache)),
+ Project(
+ Seq(Alias(
+ CreateNamedStruct(
+ planWithoutReferences.output.flatMap(a => Seq(Literal(a.name), a))),
+ "mergedValue")()),
+ planWithoutReferences),
underSubquery = true)
} else {
- removeReferences(header.plan, cache)
- })
+ planWithoutReferences
+ }
+ }
}
- val newPlan = removeReferences(planWithReferences, cache)
- val subqueryCTEs = cache.filter(_.merged).map(_.plan.asInstanceOf[CTERelationDef])
+
+ // Replace references back to `CTERelationRef`s or to original `ScalarSubquery`s in the main
+ // plan.
+ val newPlan = removeReferences(planWithReferences, subqueryPlansByLevel)
+
+ // Add `CTERelationDef`s to the plan.
+ val subqueryCTEs = subqueryPlansByLevel.flatMap(_.collect { case cte: CTERelationDef => cte })
if (subqueryCTEs.nonEmpty) {
WithCTE(newPlan, subqueryCTEs.toSeq)
} else {
newPlan
}
}
- // First traversal builds up the cache and inserts `ScalarSubqueryReference`s to the plan.
- private def insertReferences(plan: LogicalPlan, cache: ArrayBuffer[Header]): LogicalPlan = {
- plan.transformUpWithSubqueries {
- case n => n.transformExpressionsUpWithPruning(_.containsAnyPattern(SCALAR_SUBQUERY)) {
- // The subquery could contain a hint that is not propagated once we cache it, but as a
- // non-correlated scalar subquery won't be turned into a Join the loss of hints is fine.
+ // First traversal inserts `ScalarSubqueryReference`s to the plan and tries to merge subquery
+ // plans by each level.
+ private def insertReferences(
+ plan: LogicalPlan,
+ planMergers: ArrayBuffer[PlanMerger]): (LogicalPlan, Int) = {
+ // The level of a subquery plan is maximum level of its inner subqueries + 1 or 0 if it has no
+ // inner subqueries.
+ var maxLevel = 0
+ val planWithReferences =
+ plan.transformAllExpressionsWithPruning(_.containsAnyPattern(SCALAR_SUBQUERY)) {
case s: ScalarSubquery if !s.isCorrelated && s.deterministic =>
- val (subqueryIndex, headerIndex) = cacheSubquery(s.plan, cache)
- ScalarSubqueryReference(subqueryIndex, headerIndex, s.dataType, s.exprId)
- }
- }
- }
-
- // Caching returns the index of the subquery in the cache and the index of scalar member in the
- // "Header".
- private def cacheSubquery(plan: LogicalPlan, cache: ArrayBuffer[Header]): (Int, Int) = {
Review Comment:
Briefly, `cacheSubquery()` became `PlanMerger.merge()` with some modifications.
As `PlanMerger.merge()` is kind of an internal API now, which can be used by other rules. It returns `MergeResult` instead of the tuple after this refactor.
The `Header` objects, which contined the information about the merged plans became `MergedPlan`s and can be accessed as `PlanMerger.mergedPlans()` by the rules.
The main difference between `cacheSubquery()` and the new `PlanMerger.merge()` is that I removed `Header.references` and the related safeguards like the `if !references.contains(subqueryIndex)` below.
Previously, we filled up `Header.references` and with nested subquery indices so as to avoid trying to merge a subquery with other subqueries that are nested into it.
E.g.
```
SELECT (
SELECT avg(a) FROM t WHERE c = ( -- subquery 1
SELECT sum(b) FROM t2) -- subquery 2 nested into subquery 1
)
)
```
Here it doesn't make sense to try merging the plan of `subquery 1` and `subquery 2` as the former contains the latter. The problem of trying to merge the 2 was not just it doesn't make sense, but in certain weird cases (see the example in [SPARK-40618](https://issues.apache.org/jira/browse/SPARK-40618)) the merge was successful and resulted in invalid plans.
But because `PlanMerger` should be a general plan merging tool and in some future rules it will be used to merge other than subqueries, tracking the nested subqueries in it is not the best way to deal with the problem.
This is why the refactored `MergeScalarSubqueries` uses a different approach. We have one `PlanMerger` object by each subquery level. E.g. `subquery 2` is a level 0 (leaf) subquery, while `subquery 1` is a level 1 subquery because it has an inner, level 0 subquery. So the 2 will be added to different `PlanMerger`s. This way a level n subquery is tried to merge with other level n subqueries only.
--
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-54136][SQL] Extract plan merging logic from `MergeScalarSubqueries` to `PlanMerger` [spark]
Posted by "peter-toth (via GitHub)" <gi...@apache.org>.
peter-toth commented on code in PR #52835:
URL: https://github.com/apache/spark/pull/52835#discussion_r2503100575
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/PlanMerger.scala:
##########
@@ -0,0 +1,303 @@
+/*
+ * 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.ArrayBuffer
+
+import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, AttributeMap, Expression, NamedExpression}
+import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression
+import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Filter, Join, LogicalPlan, Project}
+
+/**
+ * Result of attempting to merge a plan via [[PlanMerger.merge]].
+ *
+ * @param mergedPlan The resulting plan, either:
+ * - An existing cached plan (if identical match found)
+ * - A newly merged plan combining the input with a cached plan
+ * - The original input plan (if no merge was possible)
+ * @param mergedPlanIndex The index of this plan in the PlanMerger's cache.
+ * @param merged Whether the plan was merged with an existing cached plan (true) or
+ * is a new entry (false).
+ * @param outputMap Maps attributes from the input plan to corresponding attributes in
+ * `mergedPlan`. Used to rewrite expressions referencing the original plan
+ * to reference the merged plan instead.
+ */
+case class MergeResult(
+ mergedPlan: LogicalPlan,
Review Comment:
Thanks, fixed in https://github.com/apache/spark/pull/52835/commits/f6899a0f28e893dab4c2f99b9c82f2c229fb69d9.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/PlanMerger.scala:
##########
@@ -0,0 +1,303 @@
+/*
+ * 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.ArrayBuffer
+
+import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, AttributeMap, Expression, NamedExpression}
+import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression
+import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Filter, Join, LogicalPlan, Project}
+
+/**
+ * Result of attempting to merge a plan via [[PlanMerger.merge]].
+ *
+ * @param mergedPlan The resulting plan, either:
+ * - An existing cached plan (if identical match found)
+ * - A newly merged plan combining the input with a cached plan
+ * - The original input plan (if no merge was possible)
+ * @param mergedPlanIndex The index of this plan in the PlanMerger's cache.
+ * @param merged Whether the plan was merged with an existing cached plan (true) or
+ * is a new entry (false).
+ * @param outputMap Maps attributes from the input plan to corresponding attributes in
+ * `mergedPlan`. Used to rewrite expressions referencing the original plan
+ * to reference the merged plan instead.
+ */
+case class MergeResult(
+ mergedPlan: LogicalPlan,
+ mergedPlanIndex: Int,
+ merged: Boolean,
+ outputMap: AttributeMap[Attribute])
+
+/**
+ * Represents a plan in the PlanMerger's cache.
+ *
+ * @param plan The logical plan, which may have been merged from multiple original plans.
+ * @param merged Whether this plan is the result of merging two or more plans (true), or
+ * is an original unmerged plan (false). Merged plans typically require special
+ * handling such as wrapping in CTEs.
+ */
+case class MergedPlan(plan: LogicalPlan, merged: Boolean)
+
+/**
+ * A stateful utility for merging identical or similar logical plans to enable query plan reuse.
+ *
+ * `PlanMerger` maintains a cache of previously seen plans and attempts to either:
+ * 1. Reuse an identical plan already in the cache
+ * 2. Merge a new plan with a cached plan by combining their outputs
+ *
+ * The merging process preserves semantic equivalence while combining outputs from multiple
+ * plans into a single plan. This is primarily used by [[MergeScalarSubqueries]] to deduplicate
+ * scalar subquery execution.
+ *
+ * Supported plan types for merging:
+ * - [[Project]]: Merges project lists
+ * - [[Aggregate]]: Merges aggregate expressions with identical grouping
+ * - [[Filter]]: Requires identical filter conditions
+ * - [[Join]]: Requires identical join type, hints, and conditions
+ *
+ * @example
+ * {{{
+ * val merger = PlanMerger()
+ * val result1 = merger.merge(plan1) // Adds plan1 to cache
+ * val result2 = merger.merge(plan2) // Merges with plan1 if compatible
+ * // result2.merged == true if plans were merged
+ * // result2.outputMap maps plan2's attributes to the merged plan's attributes
+ * }}}
+ */
+class PlanMerger {
+ val cache = ArrayBuffer.empty[MergedPlan]
+
+ /**
+ * Attempts to merge the given plan with cached plans, or adds it to the cache.
+ *
+ * The method tries the following in order:
+ * 1. Check if an identical plan exists in cache (using canonicalized comparison)
+ * 2. Try to merge with each cached plan using [[tryMergePlans]]
+ * 3. If no merge is possible, add as a new cache entry
+ *
+ * @param plan The logical plan to merge or cache.
+ * @return A [[MergeResult]] containing:
+ * - The merged/cached plan to use
+ * - Its index in the cache
+ * - Whether it was merged with an existing plan
+ * - An attribute mapping for rewriting expressions
+ */
+ def merge(plan: LogicalPlan): MergeResult = {
+ cache.zipWithIndex.collectFirst(Function.unlift {
+ case (mp, i) =>
+ checkIdenticalPlans(plan, mp.plan).map { outputMap =>
+ MergeResult(mp.plan, i, true, outputMap)
+ }.orElse {
+ tryMergePlans(plan, mp.plan).map {
+ case (mergedPlan, outputMap) =>
+ cache(i) = MergedPlan(mergedPlan, true)
+ MergeResult(mergedPlan, i, true, outputMap)
+ }
+ }
+ case _ => None
+ }).getOrElse {
+ cache += MergedPlan(plan, false)
+ val outputMap = AttributeMap(plan.output.map(a => a -> a))
+ MergeResult(plan, cache.length - 1, false, outputMap)
+ }
+ }
+
+ /**
+ * Returns all plans currently in the cache as an immutable indexed sequence.
+ *
+ * @return An indexed sequence of [[MergedPlan]]s in cache order. The index of each plan
+ * corresponds to the `mergedPlanIndex` returned by [[merge]].
+ */
+ def mergedPlans(): IndexedSeq[MergedPlan] = cache.toIndexedSeq
+
+ // If 2 plans are identical return the attribute mapping from the new to the cached version.
+ private def checkIdenticalPlans(
+ newPlan: LogicalPlan,
+ cachedPlan: LogicalPlan): Option[AttributeMap[Attribute]] = {
+ if (newPlan.canonicalized == cachedPlan.canonicalized) {
+ Some(AttributeMap(newPlan.output.zip(cachedPlan.output)))
+ } else {
+ None
+ }
+ }
+
+ /**
+ * Recursively attempts to merge two plans by traversing their tree structures.
+ *
+ * Two plans can be merged if:
+ * - They are identical (canonicalized forms match), OR
+ * - They have compatible root nodes with mergeable children
+ *
+ * Supported merge patterns:
+ * - Project nodes: Combines project lists from both plans
+ * - Aggregate nodes: Combines aggregate expressions if grouping is identical and both
+ * support the same aggregate implementation (hash/object-hash/sort-based)
+ * - Filter nodes: Only if filter conditions are identical
+ * - Join nodes: Only if join type, hints, and conditions are identical
+ *
+ * @param newPlan The plan to merge into the cached plan.
+ * @param cachedPlan The cached plan to merge with.
+ * @return Some((mergedPlan, outputMap)) if merge succeeds, where:
+ * - mergedPlan is the combined plan
+ * - outputMap maps newPlan's attributes to mergedPlan's attributes
+ * Returns None if plans cannot be merged.
+ */
+ private def tryMergePlans(
+ newPlan: LogicalPlan,
+ cachedPlan: LogicalPlan): Option[(LogicalPlan, AttributeMap[Attribute])] = {
+ checkIdenticalPlans(newPlan, cachedPlan).map(cachedPlan -> _).orElse(
+ (newPlan, cachedPlan) match {
+ case (np: Project, cp: Project) =>
+ tryMergePlans(np.child, cp.child).map { case (mergedChild, outputMap) =>
+ val (mergedProjectList, newOutputMap) =
+ mergeNamedExpressions(np.projectList, outputMap, cp.projectList)
+ val mergedPlan = Project(mergedProjectList, mergedChild)
+ mergedPlan -> newOutputMap
+ }
+ case (np, cp: Project) =>
+ tryMergePlans(np, cp.child).map { case (mergedChild, outputMap) =>
+ val (mergedProjectList, newOutputMap) =
+ mergeNamedExpressions(np.output, outputMap, cp.projectList)
+ val mergedPlan = Project(mergedProjectList, mergedChild)
+ mergedPlan -> newOutputMap
+ }
+ case (np: Project, cp) =>
+ tryMergePlans(np.child, cp).map { case (mergedChild, outputMap) =>
+ val (mergedProjectList, newOutputMap) =
+ mergeNamedExpressions(np.projectList, outputMap, cp.output)
+ val mergedPlan = Project(mergedProjectList, mergedChild)
+ mergedPlan -> newOutputMap
+ }
+ case (np: Aggregate, cp: Aggregate) if supportedAggregateMerge(np, cp) =>
+ tryMergePlans(np.child, cp.child).flatMap { case (mergedChild, outputMap) =>
+ val mappedNewGroupingExpression =
+ np.groupingExpressions.map(mapAttributes(_, outputMap))
+ // Order of grouping expression does matter as merging different grouping orders can
+ // introduce "extra" shuffles/sorts that might not present in all of the original
+ // subqueries.
+ if (mappedNewGroupingExpression.map(_.canonicalized) ==
+ cp.groupingExpressions.map(_.canonicalized)) {
+ val (mergedAggregateExpressions, newOutputMap) =
+ mergeNamedExpressions(np.aggregateExpressions, outputMap, cp.aggregateExpressions)
+ val mergedPlan =
+ Aggregate(cp.groupingExpressions, mergedAggregateExpressions, mergedChild)
+ Some(mergedPlan -> newOutputMap)
+ } else {
+ None
+ }
+ }
+
+ case (np: Filter, cp: Filter) =>
+ tryMergePlans(np.child, cp.child).flatMap { case (mergedChild, outputMap) =>
+ val mappedNewCondition = mapAttributes(np.condition, outputMap)
+ // Comparing the canonicalized form is required to ignore different forms of the same
+ // expression.
+ if (mappedNewCondition.canonicalized == cp.condition.canonicalized) {
+ val mergedPlan = cp.withNewChildren(Seq(mergedChild))
+ Some(mergedPlan -> outputMap)
+ } else {
+ None
+ }
+ }
+
+ case (np: Join, cp: Join) if np.joinType == cp.joinType && np.hint == cp.hint =>
+ tryMergePlans(np.left, cp.left).flatMap { case (mergedLeft, leftOutputMap) =>
+ tryMergePlans(np.right, cp.right).flatMap { case (mergedRight, rightOutputMap) =>
+ val outputMap = leftOutputMap ++ rightOutputMap
+ val mappedNewCondition = np.condition.map(mapAttributes(_, outputMap))
+ // Comparing the canonicalized form is required to ignore different forms of the same
+ // expression and `AttributeReference.qualifier`s in `cp.condition`.
+ if (mappedNewCondition.map(_.canonicalized) == cp.condition.map(_.canonicalized)) {
+ val mergedPlan = cp.withNewChildren(Seq(mergedLeft, mergedRight))
+ Some(mergedPlan -> outputMap)
+ } else {
+ None
+ }
+ }
+ }
+
+ // Otherwise merging is not possible.
+ case _ => None
+ })
+ }
+
+ private def mapAttributes[T <: Expression](expr: T, outputMap: AttributeMap[Attribute]) = {
+ expr.transform {
+ case a: Attribute => outputMap.getOrElse(a, a)
+ }.asInstanceOf[T]
+ }
+
+ // Applies `outputMap` attribute mapping on attributes of `newExpressions` and merges them into
+ // `cachedExpressions`. Returns the merged expressions and the attribute mapping from the new to
+ // the merged version that can be propagated up during merging nodes.
+ private def mergeNamedExpressions(
+ newExpressions: Seq[NamedExpression],
+ outputMap: AttributeMap[Attribute],
+ cachedExpressions: Seq[NamedExpression]) = {
+ val mergedExpressions = ArrayBuffer[NamedExpression](cachedExpressions: _*)
+ val newOutputMap = AttributeMap(newExpressions.map { ne =>
+ val mapped = mapAttributes(ne, outputMap)
+ val withoutAlias = mapped match {
+ case Alias(child, _) => child
+ case e => e
+ }
+ ne.toAttribute -> mergedExpressions.find {
+ case Alias(child, _) => child semanticEquals withoutAlias
+ case e => e semanticEquals withoutAlias
+ }.getOrElse {
+ mergedExpressions += mapped
+ mapped
+ }.toAttribute
+ })
+ (mergedExpressions.toSeq, newOutputMap)
+ }
+
+ // Only allow aggregates of the same implementation because merging different implementations
+ // could cause performance regression.
+ private def supportedAggregateMerge(newPlan: Aggregate, cachedPlan: Aggregate) = {
+ val aggregateExpressionsSeq = Seq(newPlan, cachedPlan).map { plan =>
+ plan.aggregateExpressions.flatMap(_.collect {
+ case a: AggregateExpression => a
+ })
+ }
+ val groupByExpressionSeq = Seq(newPlan, cachedPlan).map(_.groupingExpressions)
+
+ val Seq(newPlanSupportsHashAggregate, cachedPlanSupportsHashAggregate) =
+ aggregateExpressionsSeq.zip(groupByExpressionSeq).map {
+ case (aggregateExpressions, groupByExpressions) =>
+ Aggregate.supportsHashAggregate(
+ aggregateExpressions.flatMap(
+ _.aggregateFunction.aggBufferAttributes), groupByExpressions)
+ }
+
+ newPlanSupportsHashAggregate && cachedPlanSupportsHashAggregate ||
+ newPlanSupportsHashAggregate == cachedPlanSupportsHashAggregate && {
+ val Seq(newPlanSupportsObjectHashAggregate, cachedPlanSupportsObjectHashAggregate) =
+ aggregateExpressionsSeq.zip(groupByExpressionSeq).map {
+ case (aggregateExpressions, groupByExpressions) =>
+ Aggregate.supportsObjectHashAggregate(aggregateExpressions, groupByExpressions)
+ }
+ newPlanSupportsObjectHashAggregate && cachedPlanSupportsObjectHashAggregate ||
+ newPlanSupportsObjectHashAggregate == cachedPlanSupportsObjectHashAggregate
+ }
+ }
+}
+
+object PlanMerger {
+ def apply(): PlanMerger = new PlanMerger
Review Comment:
Done in https://github.com/apache/spark/pull/52835/commits/f6899a0f28e893dab4c2f99b9c82f2c229fb69d9.
--
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-54136][SQL] Extract plan merging logic from `MergeScalarSubqueries` to `PlanMerger` [spark]
Posted by "peter-toth (via GitHub)" <gi...@apache.org>.
peter-toth commented on code in PR #52835:
URL: https://github.com/apache/spark/pull/52835#discussion_r2501976269
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/PlanMerger.scala:
##########
@@ -0,0 +1,303 @@
+/*
+ * 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.ArrayBuffer
+
+import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, AttributeMap, Expression, NamedExpression}
+import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression
+import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Filter, Join, LogicalPlan, Project}
+
+/**
+ * Result of attempting to merge a plan via [[PlanMerger.merge]].
+ *
+ * @param mergedPlan The resulting plan, either:
+ * - An existing cached plan (if identical match found)
+ * - A newly merged plan combining the input with a cached plan
+ * - The original input plan (if no merge was possible)
+ * @param mergedPlanIndex The index of this plan in the PlanMerger's cache.
+ * @param merged Whether the plan was merged with an existing cached plan (true) or
+ * is a new entry (false).
+ * @param outputMap Maps attributes from the input plan to corresponding attributes in
+ * `mergedPlan`. Used to rewrite expressions referencing the original plan
+ * to reference the merged plan instead.
+ */
+case class MergeResult(
+ mergedPlan: LogicalPlan,
Review Comment:
Good idea! Let me 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-54136][SQL] Extract plan merging logic from `MergeScalarSubqueries` to `PlanMerger` [spark]
Posted by "peter-toth (via GitHub)" <gi...@apache.org>.
peter-toth commented on code in PR #52835:
URL: https://github.com/apache/spark/pull/52835#discussion_r2485734149
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/MergeScalarSubqueries.scala:
##########
@@ -117,297 +116,126 @@ object MergeScalarSubqueries extends Rule[LogicalPlan] {
}
}
- /**
- * An item in the cache of merged scalar subqueries.
- *
- * @param plan The plan of a merged scalar subquery.
- * @param merged A flag to identify if this item is the result of merging subqueries.
- * Please note that `attributes.size == 1` doesn't always mean that the plan is not
- * merged as there can be subqueries that are different ([[checkIdenticalPlans]] is
- * false) due to an extra [[Project]] node in one of them. In that case
- * `attributes.size` remains 1 after merging, but the merged flag becomes true.
- * @param references A set of subquery indexes in the cache to track all (including transitive)
- * nested subqueries.
- */
- case class Header(
- plan: LogicalPlan,
- merged: Boolean,
- references: Set[Int])
-
private def extractCommonScalarSubqueries(plan: LogicalPlan) = {
- val cache = ArrayBuffer.empty[Header]
- val planWithReferences = insertReferences(plan, cache)
- cache.zipWithIndex.foreach { case (header, i) =>
- cache(i) = cache(i).copy(plan =
- if (header.merged) {
+ // Collect `ScalarSubquery` plans by level into `PlanMerger`s and insert references in place of
+ // `ScalarSubquery`s.
+ val planMergers = ArrayBuffer.empty[PlanMerger]
+ val planWithReferences = insertReferences(plan, planMergers)._1
+
+ // Traverse level by level and convert merged plans to `CTERelationDef`s and keep non-merged
+ // ones. While traversing replace references in plans back to `CTERelationRef`s or to original
+ // `ScalarSubquery`s. This is safe as a subquery plan at a level can reference only lower level
+ // other subqueries.
+ val subqueryPlansByLevel = ArrayBuffer.empty[IndexedSeq[LogicalPlan]]
+ planMergers.foreach { planMerger =>
+ val mergedPlans = planMerger.mergedPlans()
+ subqueryPlansByLevel += mergedPlans.map { mergedPlan =>
+ val planWithoutReferences = if (subqueryPlansByLevel.isEmpty) {
+ // Level 0 plans can't contain references
+ mergedPlan.plan
+ } else {
+ removeReferences(mergedPlan.plan, subqueryPlansByLevel)
+ }
+ if (mergedPlan.merged) {
CTERelationDef(
- createProject(header.plan.output, removeReferences(header.plan, cache)),
+ Project(
Review Comment:
In the current version of the rule `createProject()` is a small helper function and it is used only here, so I felt it doesn't need to be a separate function.
--
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-54136][SQL] Extract plan merging logic from `MergeScalarSubqueries` to `PlanMerger` [spark]
Posted by "peter-toth (via GitHub)" <gi...@apache.org>.
peter-toth commented on code in PR #52835:
URL: https://github.com/apache/spark/pull/52835#discussion_r2485976169
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/MergeScalarSubqueries.scala:
##########
@@ -117,297 +116,126 @@ object MergeScalarSubqueries extends Rule[LogicalPlan] {
}
}
- /**
- * An item in the cache of merged scalar subqueries.
- *
- * @param plan The plan of a merged scalar subquery.
- * @param merged A flag to identify if this item is the result of merging subqueries.
- * Please note that `attributes.size == 1` doesn't always mean that the plan is not
- * merged as there can be subqueries that are different ([[checkIdenticalPlans]] is
- * false) due to an extra [[Project]] node in one of them. In that case
- * `attributes.size` remains 1 after merging, but the merged flag becomes true.
- * @param references A set of subquery indexes in the cache to track all (including transitive)
- * nested subqueries.
- */
- case class Header(
- plan: LogicalPlan,
- merged: Boolean,
- references: Set[Int])
-
private def extractCommonScalarSubqueries(plan: LogicalPlan) = {
- val cache = ArrayBuffer.empty[Header]
- val planWithReferences = insertReferences(plan, cache)
- cache.zipWithIndex.foreach { case (header, i) =>
- cache(i) = cache(i).copy(plan =
- if (header.merged) {
+ // Collect `ScalarSubquery` plans by level into `PlanMerger`s and insert references in place of
+ // `ScalarSubquery`s.
+ val planMergers = ArrayBuffer.empty[PlanMerger]
+ val planWithReferences = insertReferences(plan, planMergers)._1
+
+ // Traverse level by level and convert merged plans to `CTERelationDef`s and keep non-merged
+ // ones. While traversing replace references in plans back to `CTERelationRef`s or to original
+ // `ScalarSubquery`s. This is safe as a subquery plan at a level can reference only lower level
+ // other subqueries.
+ val subqueryPlansByLevel = ArrayBuffer.empty[IndexedSeq[LogicalPlan]]
+ planMergers.foreach { planMerger =>
+ val mergedPlans = planMerger.mergedPlans()
+ subqueryPlansByLevel += mergedPlans.map { mergedPlan =>
+ val planWithoutReferences = if (subqueryPlansByLevel.isEmpty) {
+ // Level 0 plans can't contain references
+ mergedPlan.plan
+ } else {
+ removeReferences(mergedPlan.plan, subqueryPlansByLevel)
+ }
+ if (mergedPlan.merged) {
CTERelationDef(
- createProject(header.plan.output, removeReferences(header.plan, cache)),
+ Project(
+ Seq(Alias(
+ CreateNamedStruct(
+ planWithoutReferences.output.flatMap(a => Seq(Literal(a.name), a))),
+ "mergedValue")()),
+ planWithoutReferences),
underSubquery = true)
} else {
- removeReferences(header.plan, cache)
- })
+ planWithoutReferences
+ }
+ }
}
- val newPlan = removeReferences(planWithReferences, cache)
- val subqueryCTEs = cache.filter(_.merged).map(_.plan.asInstanceOf[CTERelationDef])
+
+ // Replace references back to `CTERelationRef`s or to original `ScalarSubquery`s in the main
+ // plan.
+ val newPlan = removeReferences(planWithReferences, subqueryPlansByLevel)
+
+ // Add `CTERelationDef`s to the plan.
+ val subqueryCTEs = subqueryPlansByLevel.flatMap(_.collect { case cte: CTERelationDef => cte })
if (subqueryCTEs.nonEmpty) {
WithCTE(newPlan, subqueryCTEs.toSeq)
} else {
newPlan
}
}
- // First traversal builds up the cache and inserts `ScalarSubqueryReference`s to the plan.
- private def insertReferences(plan: LogicalPlan, cache: ArrayBuffer[Header]): LogicalPlan = {
- plan.transformUpWithSubqueries {
- case n => n.transformExpressionsUpWithPruning(_.containsAnyPattern(SCALAR_SUBQUERY)) {
- // The subquery could contain a hint that is not propagated once we cache it, but as a
- // non-correlated scalar subquery won't be turned into a Join the loss of hints is fine.
+ // First traversal inserts `ScalarSubqueryReference`s to the plan and tries to merge subquery
+ // plans by each level.
+ private def insertReferences(
+ plan: LogicalPlan,
+ planMergers: ArrayBuffer[PlanMerger]): (LogicalPlan, Int) = {
+ // The level of a subquery plan is maximum level of its inner subqueries + 1 or 0 if it has no
+ // inner subqueries.
+ var maxLevel = 0
+ val planWithReferences =
+ plan.transformAllExpressionsWithPruning(_.containsAnyPattern(SCALAR_SUBQUERY)) {
case s: ScalarSubquery if !s.isCorrelated && s.deterministic =>
- val (subqueryIndex, headerIndex) = cacheSubquery(s.plan, cache)
- ScalarSubqueryReference(subqueryIndex, headerIndex, s.dataType, s.exprId)
- }
- }
- }
-
- // Caching returns the index of the subquery in the cache and the index of scalar member in the
- // "Header".
- private def cacheSubquery(plan: LogicalPlan, cache: ArrayBuffer[Header]): (Int, Int) = {
- val output = plan.output.head
- val references = mutable.HashSet.empty[Int]
- plan.transformAllExpressionsWithPruning(_.containsAnyPattern(SCALAR_SUBQUERY_REFERENCE)) {
- case ssr: ScalarSubqueryReference =>
- references += ssr.subqueryIndex
- references ++= cache(ssr.subqueryIndex).references
- ssr
- }
-
- cache.zipWithIndex.collectFirst(Function.unlift {
- case (header, subqueryIndex) if !references.contains(subqueryIndex) =>
- checkIdenticalPlans(plan, header.plan).map { outputMap =>
- val mappedOutput = mapAttributes(output, outputMap)
- val headerIndex = header.plan.output.indexWhere(_.exprId == mappedOutput.exprId)
- subqueryIndex -> headerIndex
- }.orElse{
- tryMergePlans(plan, header.plan).map {
- case (mergedPlan, outputMap) =>
- val mappedOutput = mapAttributes(output, outputMap)
- val headerIndex = mergedPlan.output.indexWhere(_.exprId == mappedOutput.exprId)
- cache(subqueryIndex) = Header(mergedPlan, true, header.references ++ references)
- subqueryIndex -> headerIndex
- }
- }
- case _ => None
- }).getOrElse {
- cache += Header(plan, false, references.toSet)
- cache.length - 1 -> 0
- }
- }
-
- // If 2 plans are identical return the attribute mapping from the new to the cached version.
- private def checkIdenticalPlans(
- newPlan: LogicalPlan,
- cachedPlan: LogicalPlan): Option[AttributeMap[Attribute]] = {
- if (newPlan.canonicalized == cachedPlan.canonicalized) {
- Some(AttributeMap(newPlan.output.zip(cachedPlan.output)))
- } else {
- None
- }
- }
-
- // Recursively traverse down and try merging 2 plans. If merge is possible then return the merged
- // plan with the attribute mapping from the new to the merged version.
- // Please note that merging arbitrary plans can be complicated, the current version supports only
- // some of the most important nodes.
- private def tryMergePlans(
- newPlan: LogicalPlan,
- cachedPlan: LogicalPlan): Option[(LogicalPlan, AttributeMap[Attribute])] = {
- checkIdenticalPlans(newPlan, cachedPlan).map(cachedPlan -> _).orElse(
- (newPlan, cachedPlan) match {
- case (np: Project, cp: Project) =>
- tryMergePlans(np.child, cp.child).map { case (mergedChild, outputMap) =>
- val (mergedProjectList, newOutputMap) =
- mergeNamedExpressions(np.projectList, outputMap, cp.projectList)
- val mergedPlan = Project(mergedProjectList, mergedChild)
- mergedPlan -> newOutputMap
- }
- case (np, cp: Project) =>
- tryMergePlans(np, cp.child).map { case (mergedChild, outputMap) =>
- val (mergedProjectList, newOutputMap) =
- mergeNamedExpressions(np.output, outputMap, cp.projectList)
- val mergedPlan = Project(mergedProjectList, mergedChild)
- mergedPlan -> newOutputMap
- }
- case (np: Project, cp) =>
- tryMergePlans(np.child, cp).map { case (mergedChild, outputMap) =>
- val (mergedProjectList, newOutputMap) =
- mergeNamedExpressions(np.projectList, outputMap, cp.output)
- val mergedPlan = Project(mergedProjectList, mergedChild)
- mergedPlan -> newOutputMap
- }
- case (np: Aggregate, cp: Aggregate) if supportedAggregateMerge(np, cp) =>
- tryMergePlans(np.child, cp.child).flatMap { case (mergedChild, outputMap) =>
- val mappedNewGroupingExpression =
- np.groupingExpressions.map(mapAttributes(_, outputMap))
- // Order of grouping expression does matter as merging different grouping orders can
- // introduce "extra" shuffles/sorts that might not present in all of the original
- // subqueries.
- if (mappedNewGroupingExpression.map(_.canonicalized) ==
- cp.groupingExpressions.map(_.canonicalized)) {
- val (mergedAggregateExpressions, newOutputMap) =
- mergeNamedExpressions(np.aggregateExpressions, outputMap, cp.aggregateExpressions)
- val mergedPlan =
- Aggregate(cp.groupingExpressions, mergedAggregateExpressions, mergedChild)
- Some(mergedPlan -> newOutputMap)
- } else {
- None
- }
- }
-
- case (np: Filter, cp: Filter) =>
- tryMergePlans(np.child, cp.child).flatMap { case (mergedChild, outputMap) =>
- val mappedNewCondition = mapAttributes(np.condition, outputMap)
- // Comparing the canonicalized form is required to ignore different forms of the same
- // expression.
- if (mappedNewCondition.canonicalized == cp.condition.canonicalized) {
- val mergedPlan = cp.withNewChildren(Seq(mergedChild))
- Some(mergedPlan -> outputMap)
- } else {
- None
- }
- }
-
- case (np: Join, cp: Join) if np.joinType == cp.joinType && np.hint == cp.hint =>
- tryMergePlans(np.left, cp.left).flatMap { case (mergedLeft, leftOutputMap) =>
- tryMergePlans(np.right, cp.right).flatMap { case (mergedRight, rightOutputMap) =>
- val outputMap = leftOutputMap ++ rightOutputMap
- val mappedNewCondition = np.condition.map(mapAttributes(_, outputMap))
- // Comparing the canonicalized form is required to ignore different forms of the same
- // expression and `AttributeReference.quailifier`s in `cp.condition`.
- if (mappedNewCondition.map(_.canonicalized) == cp.condition.map(_.canonicalized)) {
- val mergedPlan = cp.withNewChildren(Seq(mergedLeft, mergedRight))
- Some(mergedPlan -> outputMap)
- } else {
- None
- }
- }
- }
-
- // Otherwise merging is not possible.
- case _ => None
- })
- }
-
- private def createProject(attributes: Seq[Attribute], plan: LogicalPlan): Project = {
- Project(
- Seq(Alias(
- CreateNamedStruct(attributes.flatMap(a => Seq(Literal(a.name), a))),
- "mergedValue")()),
- plan)
- }
-
- private def mapAttributes[T <: Expression](expr: T, outputMap: AttributeMap[Attribute]) = {
- expr.transform {
- case a: Attribute => outputMap.getOrElse(a, a)
- }.asInstanceOf[T]
- }
-
- // Applies `outputMap` attribute mapping on attributes of `newExpressions` and merges them into
- // `cachedExpressions`. Returns the merged expressions and the attribute mapping from the new to
- // the merged version that can be propagated up during merging nodes.
- private def mergeNamedExpressions(
- newExpressions: Seq[NamedExpression],
- outputMap: AttributeMap[Attribute],
- cachedExpressions: Seq[NamedExpression]) = {
- val mergedExpressions = ArrayBuffer[NamedExpression](cachedExpressions: _*)
- val newOutputMap = AttributeMap(newExpressions.map { ne =>
- val mapped = mapAttributes(ne, outputMap)
- val withoutAlias = mapped match {
- case Alias(child, _) => child
- case e => e
- }
- ne.toAttribute -> mergedExpressions.find {
- case Alias(child, _) => child semanticEquals withoutAlias
- case e => e semanticEquals withoutAlias
- }.getOrElse {
- mergedExpressions += mapped
- mapped
- }.toAttribute
- })
- (mergedExpressions.toSeq, newOutputMap)
- }
-
- // Only allow aggregates of the same implementation because merging different implementations
- // could cause performance regression.
- private def supportedAggregateMerge(newPlan: Aggregate, cachedPlan: Aggregate) = {
- val aggregateExpressionsSeq = Seq(newPlan, cachedPlan).map { plan =>
- plan.aggregateExpressions.flatMap(_.collect {
- case a: AggregateExpression => a
- })
- }
- val groupByExpressionSeq = Seq(newPlan, cachedPlan).map(_.groupingExpressions)
-
- val Seq(newPlanSupportsHashAggregate, cachedPlanSupportsHashAggregate) =
- aggregateExpressionsSeq.zip(groupByExpressionSeq).map {
- case (aggregateExpressions, groupByExpressions) =>
- Aggregate.supportsHashAggregate(
- aggregateExpressions.flatMap(
- _.aggregateFunction.aggBufferAttributes), groupByExpressions)
- }
-
- newPlanSupportsHashAggregate && cachedPlanSupportsHashAggregate ||
- newPlanSupportsHashAggregate == cachedPlanSupportsHashAggregate && {
- val Seq(newPlanSupportsObjectHashAggregate, cachedPlanSupportsObjectHashAggregate) =
- aggregateExpressionsSeq.zip(groupByExpressionSeq).map {
- case (aggregateExpressions, groupByExpressions) =>
- Aggregate.supportsObjectHashAggregate(aggregateExpressions, groupByExpressions)
- }
- newPlanSupportsObjectHashAggregate && cachedPlanSupportsObjectHashAggregate ||
- newPlanSupportsObjectHashAggregate == cachedPlanSupportsObjectHashAggregate
+ val (planWithReferences, level) = insertReferences(s.plan, planMergers)
+
+ while (level >= planMergers.size) planMergers += PlanMerger()
+ // The subquery could contain a hint that is not propagated once we merge it, but as a
+ // non-correlated scalar subquery won't be turned into a Join the loss of hints is fine.
+ val planMergeResult = planMergers(level).merge(planWithReferences)
+
+ maxLevel = maxLevel.max(level + 1)
+
+ val mergedOutput = planMergeResult.outputMap(planWithReferences.output.head)
+ val headerIndex =
+ planMergeResult.mergedPlan.output.indexWhere(_.exprId == mergedOutput.exprId)
+ ScalarSubqueryReference(
+ level,
+ planMergeResult.mergedPlanIndex,
+ headerIndex,
+ s.dataType,
+ s.exprId)
+ case o => o
}
+ (planWithReferences, maxLevel)
}
// Second traversal replaces `ScalarSubqueryReference`s to either
// `GetStructField(ScalarSubquery(CTERelationRef to the merged plan)` if the plan is merged from
// multiple subqueries or `ScalarSubquery(original plan)` if it isn't.
private def removeReferences(
plan: LogicalPlan,
- cache: ArrayBuffer[Header]) = {
- plan.transformUpWithSubqueries {
- case n =>
- n.transformExpressionsWithPruning(_.containsAnyPattern(SCALAR_SUBQUERY_REFERENCE)) {
- case ssr: ScalarSubqueryReference =>
- val header = cache(ssr.subqueryIndex)
- if (header.merged) {
- val subqueryCTE = header.plan.asInstanceOf[CTERelationDef]
- GetStructField(
- ScalarSubquery(
- CTERelationRef(subqueryCTE.id, _resolved = true, subqueryCTE.output,
- subqueryCTE.isStreaming),
- exprId = ssr.exprId),
- ssr.headerIndex)
- } else {
- ScalarSubquery(header.plan, exprId = ssr.exprId)
- }
+ subqueryPlansByLevel: ArrayBuffer[IndexedSeq[LogicalPlan]]) = {
+ plan.transformAllExpressionsWithPruning(_.containsAnyPattern(SCALAR_SUBQUERY_REFERENCE)) {
+ case ssr: ScalarSubqueryReference =>
+ subqueryPlansByLevel(ssr.level)(ssr.mergedPlanIndex) match {
+ case cte: CTERelationDef =>
+ GetStructField(
+ ScalarSubquery(
+ CTERelationRef(cte.id, _resolved = true, cte.output, cte.isStreaming),
+ exprId = ssr.exprId),
+ ssr.outputIndex)
+ case o => ScalarSubquery(o, exprId = ssr.exprId)
}
}
}
}
/**
- * Temporal reference to a cached subquery.
+ * Temporal reference to a subquery which is added to a `PlanMerger`.
*
- * @param subqueryIndex A subquery index in the cache.
- * @param headerIndex An index in the output of merged subquery.
- * @param dataType The dataType of origin scalar subquery.
+ * @param level The level of the replaced subquery. It defines the `PlanMerger` instance into which
+ * the subquery is merged.
+ * @param mergedPlanIndex The index of the merged plan in the `PlanMerger`.
+ * @param outputIndex The index of the output attribute of the merged plan.
+ * @param dataType The dataType of original scalar subquery.
+ * @param exprId The expression id of the original scalar subquery.
*/
case class ScalarSubqueryReference(
- subqueryIndex: Int,
- headerIndex: Int,
+ level: Int,
+ mergedPlanIndex: Int,
+ outputIndex: Int,
dataType: DataType,
exprId: ExprId) extends LeafExpression with Unevaluable {
override def nullable: Boolean = true
final override val nodePatterns: Seq[TreePattern] = Seq(SCALAR_SUBQUERY_REFERENCE)
-
- override def stringArgs: Iterator[Any] = Iterator(subqueryIndex, headerIndex, dataType, exprId.id)
Review Comment:
`ScalarSubqueryReference` never shows up in any stringified plan. It's just a temporary object that is added and then removed when `MergeScalarSubqueries` runs so I think we don't need 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-54136][SQL] Extract plan merging logic from `MergeScalarSubqueries` to `PlanMerger` [spark]
Posted by "peter-toth (via GitHub)" <gi...@apache.org>.
peter-toth commented on code in PR #52835:
URL: https://github.com/apache/spark/pull/52835#discussion_r2485847706
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/MergeScalarSubqueries.scala:
##########
@@ -117,297 +116,126 @@ object MergeScalarSubqueries extends Rule[LogicalPlan] {
}
}
- /**
- * An item in the cache of merged scalar subqueries.
- *
- * @param plan The plan of a merged scalar subquery.
- * @param merged A flag to identify if this item is the result of merging subqueries.
- * Please note that `attributes.size == 1` doesn't always mean that the plan is not
- * merged as there can be subqueries that are different ([[checkIdenticalPlans]] is
- * false) due to an extra [[Project]] node in one of them. In that case
- * `attributes.size` remains 1 after merging, but the merged flag becomes true.
- * @param references A set of subquery indexes in the cache to track all (including transitive)
- * nested subqueries.
- */
- case class Header(
- plan: LogicalPlan,
- merged: Boolean,
- references: Set[Int])
-
private def extractCommonScalarSubqueries(plan: LogicalPlan) = {
- val cache = ArrayBuffer.empty[Header]
- val planWithReferences = insertReferences(plan, cache)
- cache.zipWithIndex.foreach { case (header, i) =>
- cache(i) = cache(i).copy(plan =
- if (header.merged) {
+ // Collect `ScalarSubquery` plans by level into `PlanMerger`s and insert references in place of
+ // `ScalarSubquery`s.
+ val planMergers = ArrayBuffer.empty[PlanMerger]
+ val planWithReferences = insertReferences(plan, planMergers)._1
+
+ // Traverse level by level and convert merged plans to `CTERelationDef`s and keep non-merged
+ // ones. While traversing replace references in plans back to `CTERelationRef`s or to original
+ // `ScalarSubquery`s. This is safe as a subquery plan at a level can reference only lower level
+ // other subqueries.
+ val subqueryPlansByLevel = ArrayBuffer.empty[IndexedSeq[LogicalPlan]]
+ planMergers.foreach { planMerger =>
+ val mergedPlans = planMerger.mergedPlans()
+ subqueryPlansByLevel += mergedPlans.map { mergedPlan =>
+ val planWithoutReferences = if (subqueryPlansByLevel.isEmpty) {
+ // Level 0 plans can't contain references
+ mergedPlan.plan
+ } else {
+ removeReferences(mergedPlan.plan, subqueryPlansByLevel)
+ }
+ if (mergedPlan.merged) {
CTERelationDef(
- createProject(header.plan.output, removeReferences(header.plan, cache)),
+ Project(
+ Seq(Alias(
+ CreateNamedStruct(
+ planWithoutReferences.output.flatMap(a => Seq(Literal(a.name), a))),
+ "mergedValue")()),
+ planWithoutReferences),
underSubquery = true)
} else {
- removeReferences(header.plan, cache)
- })
+ planWithoutReferences
+ }
+ }
}
- val newPlan = removeReferences(planWithReferences, cache)
- val subqueryCTEs = cache.filter(_.merged).map(_.plan.asInstanceOf[CTERelationDef])
+
+ // Replace references back to `CTERelationRef`s or to original `ScalarSubquery`s in the main
+ // plan.
+ val newPlan = removeReferences(planWithReferences, subqueryPlansByLevel)
+
+ // Add `CTERelationDef`s to the plan.
+ val subqueryCTEs = subqueryPlansByLevel.flatMap(_.collect { case cte: CTERelationDef => cte })
if (subqueryCTEs.nonEmpty) {
WithCTE(newPlan, subqueryCTEs.toSeq)
} else {
newPlan
}
}
- // First traversal builds up the cache and inserts `ScalarSubqueryReference`s to the plan.
- private def insertReferences(plan: LogicalPlan, cache: ArrayBuffer[Header]): LogicalPlan = {
- plan.transformUpWithSubqueries {
- case n => n.transformExpressionsUpWithPruning(_.containsAnyPattern(SCALAR_SUBQUERY)) {
- // The subquery could contain a hint that is not propagated once we cache it, but as a
- // non-correlated scalar subquery won't be turned into a Join the loss of hints is fine.
+ // First traversal inserts `ScalarSubqueryReference`s to the plan and tries to merge subquery
+ // plans by each level.
+ private def insertReferences(
+ plan: LogicalPlan,
+ planMergers: ArrayBuffer[PlanMerger]): (LogicalPlan, Int) = {
+ // The level of a subquery plan is maximum level of its inner subqueries + 1 or 0 if it has no
+ // inner subqueries.
+ var maxLevel = 0
+ val planWithReferences =
+ plan.transformAllExpressionsWithPruning(_.containsAnyPattern(SCALAR_SUBQUERY)) {
case s: ScalarSubquery if !s.isCorrelated && s.deterministic =>
- val (subqueryIndex, headerIndex) = cacheSubquery(s.plan, cache)
- ScalarSubqueryReference(subqueryIndex, headerIndex, s.dataType, s.exprId)
- }
- }
- }
-
- // Caching returns the index of the subquery in the cache and the index of scalar member in the
- // "Header".
- private def cacheSubquery(plan: LogicalPlan, cache: ArrayBuffer[Header]): (Int, Int) = {
Review Comment:
Briefly, `cacheSubquery()` became `PlanMerger.merge()` with some modifications.
As `PlanMerger.merge()` is kind of an internal API now, which can be used by other rules. It returns `MergeResult` instead of the tuple after this refactor.
The `Header` objects, which contined the information about the merged plans became `MergedPlan`s and can be accessed as `PlanMerger.mergedPlans()` by the rules.
The main difference between `cacheSubquery()` and the new `PlanMerger.merge()` is that I removed `Header.references` and the related safeguards like the `if !references.contains(subqueryIndex)` below.
Previously, we filled up `Header.references` and with nested subquery indices so as to avoid trying to merge a subquery with other subqueries that are nested into it.
E.g.
```
SELECT (
SELECT avg(a) FROM t WHERE c = ( -- subquery 1
SELECT sum(b) FROM t2) -- subquery 2 nested into subquery 1
)
)
```
Here it doesn't make sense to try merging the plan of `subquery 1` and `subquery 2` as the former contains the latter. The problem of trying to merge the 2 was not just it doesn't make sense, but in certain weird cases (see the example in [SPARK-40618](https://issues.apache.org/jira/browse/SPARK-40618)) the merge was successful and resulted in invalid plans.
But because `PlanMerger` should be a general plan merging tool and in some future rules it will be used to merge other than subqueries, tracking the nested subqueries in it is not the best way to deal with the problem.
This is why the refactored `MergeScalarSubqueries` uses a different approach. We have one `PlanMerger` object per each subquery level. E.g. `subquery 2` is a level 0 (leaf) subquery, but `subquery 1` is a level 1 subquery because it has an inner, level 0 subquery. So the 2 will be added to different `PlanMerger`s. This way a level n subquery is tried to merge with other level n subqueries only.
--
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-54136][SQL] Extract plan merging logic from `MergeScalarSubqueries` to `PlanMerger` [spark]
Posted by "peter-toth (via GitHub)" <gi...@apache.org>.
peter-toth commented on code in PR #52835:
URL: https://github.com/apache/spark/pull/52835#discussion_r2485708460
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/PlanMerger.scala:
##########
@@ -0,0 +1,303 @@
+/*
+ * 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.ArrayBuffer
+
+import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, AttributeMap, Expression, NamedExpression}
+import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression
+import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Filter, Join, LogicalPlan, Project}
+
+/**
+ * Result of attempting to merge a plan via [[PlanMerger.merge]].
+ *
+ * @param mergedPlan The resulting plan, either:
+ * - An existing cached plan (if identical match found)
+ * - A newly merged plan combining the input with a cached plan
+ * - The original input plan (if no merge was possible)
+ * @param mergedPlanIndex The index of this plan in the PlanMerger's cache.
+ * @param merged Whether the plan was merged with an existing cached plan (true) or
+ * is a new entry (false).
+ * @param outputMap Maps attributes from the input plan to corresponding attributes in
+ * `mergedPlan`. Used to rewrite expressions referencing the original plan
+ * to reference the merged plan instead.
+ */
+case class MergeResult(
+ mergedPlan: LogicalPlan,
+ mergedPlanIndex: Int,
+ merged: Boolean,
+ outputMap: AttributeMap[Attribute])
+
+/**
+ * Represents a plan in the PlanMerger's cache.
+ *
+ * @param plan The logical plan, which may have been merged from multiple original plans.
+ * @param merged Whether this plan is the result of merging two or more plans (true), or
+ * is an original unmerged plan (false). Merged plans typically require special
+ * handling such as wrapping in CTEs.
+ */
+case class MergedPlan(plan: LogicalPlan, merged: Boolean)
+
+/**
+ * A stateful utility for merging identical or similar logical plans to enable query plan reuse.
+ *
+ * `PlanMerger` maintains a cache of previously seen plans and attempts to either:
+ * 1. Reuse an identical plan already in the cache
+ * 2. Merge a new plan with a cached plan by combining their outputs
+ *
+ * The merging process preserves semantic equivalence while combining outputs from multiple
+ * plans into a single plan. This is primarily used by [[MergeScalarSubqueries]] to deduplicate
+ * scalar subquery execution.
+ *
+ * Supported plan types for merging:
+ * - [[Project]]: Merges project lists
+ * - [[Aggregate]]: Merges aggregate expressions with identical grouping
+ * - [[Filter]]: Requires identical filter conditions
+ * - [[Join]]: Requires identical join type, hints, and conditions
+ *
+ * @example
+ * {{{
+ * val merger = PlanMerger()
+ * val result1 = merger.merge(plan1) // Adds plan1 to cache
+ * val result2 = merger.merge(plan2) // Merges with plan1 if compatible
+ * // result2.merged == true if plans were merged
+ * // result2.outputMap maps plan2's attributes to the merged plan's attributes
+ * }}}
+ */
+class PlanMerger {
+ val cache = ArrayBuffer.empty[MergedPlan]
+
+ /**
+ * Attempts to merge the given plan with cached plans, or adds it to the cache.
+ *
+ * The method tries the following in order:
+ * 1. Check if an identical plan exists in cache (using canonicalized comparison)
+ * 2. Try to merge with each cached plan using [[tryMergePlans]]
+ * 3. If no merge is possible, add as a new cache entry
+ *
+ * @param plan The logical plan to merge or cache.
+ * @return A [[MergeResult]] containing:
+ * - The merged/cached plan to use
+ * - Its index in the cache
+ * - Whether it was merged with an existing plan
+ * - An attribute mapping for rewriting expressions
+ */
+ def merge(plan: LogicalPlan): MergeResult = {
+ cache.zipWithIndex.collectFirst(Function.unlift {
+ case (mp, i) =>
+ checkIdenticalPlans(plan, mp.plan).map { outputMap =>
+ MergeResult(mp.plan, i, true, outputMap)
+ }.orElse {
+ tryMergePlans(plan, mp.plan).map {
+ case (mergedPlan, outputMap) =>
+ cache(i) = MergedPlan(mergedPlan, true)
+ MergeResult(mergedPlan, i, true, outputMap)
+ }
+ }
+ case _ => None
+ }).getOrElse {
+ cache += MergedPlan(plan, false)
+ val outputMap = AttributeMap(plan.output.map(a => a -> a))
+ MergeResult(plan, cache.length - 1, false, outputMap)
+ }
+ }
+
+ /**
+ * Returns all plans currently in the cache as an immutable indexed sequence.
+ *
+ * @return An indexed sequence of [[MergedPlan]]s in cache order. The index of each plan
+ * corresponds to the `mergedPlanIndex` returned by [[merge]].
+ */
+ def mergedPlans(): IndexedSeq[MergedPlan] = cache.toIndexedSeq
+
+ // If 2 plans are identical return the attribute mapping from the new to the cached version.
+ private def checkIdenticalPlans(
+ newPlan: LogicalPlan,
+ cachedPlan: LogicalPlan): Option[AttributeMap[Attribute]] = {
+ if (newPlan.canonicalized == cachedPlan.canonicalized) {
+ Some(AttributeMap(newPlan.output.zip(cachedPlan.output)))
+ } else {
+ None
+ }
+ }
+
+ /**
+ * Recursively attempts to merge two plans by traversing their tree structures.
+ *
+ * Two plans can be merged if:
+ * - They are identical (canonicalized forms match), OR
+ * - They have compatible root nodes with mergeable children
+ *
+ * Supported merge patterns:
+ * - Project nodes: Combines project lists from both plans
+ * - Aggregate nodes: Combines aggregate expressions if grouping is identical and both
+ * support the same aggregate implementation (hash/object-hash/sort-based)
+ * - Filter nodes: Only if filter conditions are identical
+ * - Join nodes: Only if join type, hints, and conditions are identical
+ *
+ * @param newPlan The plan to merge into the cached plan.
+ * @param cachedPlan The cached plan to merge with.
+ * @return Some((mergedPlan, outputMap)) if merge succeeds, where:
+ * - mergedPlan is the combined plan
+ * - outputMap maps newPlan's attributes to mergedPlan's attributes
+ * Returns None if plans cannot be merged.
+ */
+ private def tryMergePlans(
Review Comment:
Yes, `tryMergePlans()`, `mapAttributes()`, `mergeNamedExpressions()` and `supportedAggregateMerge()` are just copied without any changes.
--
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-54136][SQL] Extract plan merging logic from `MergeScalarSubqueries` to `PlanMerger` [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #52835:
URL: https://github.com/apache/spark/pull/52835#discussion_r2484971747
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/PlanMerger.scala:
##########
@@ -0,0 +1,303 @@
+/*
+ * 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.ArrayBuffer
+
+import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, AttributeMap, Expression, NamedExpression}
+import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression
+import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Filter, Join, LogicalPlan, Project}
+
+/**
+ * Result of attempting to merge a plan via [[PlanMerger.merge]].
+ *
+ * @param mergedPlan The resulting plan, either:
+ * - An existing cached plan (if identical match found)
+ * - A newly merged plan combining the input with a cached plan
+ * - The original input plan (if no merge was possible)
+ * @param mergedPlanIndex The index of this plan in the PlanMerger's cache.
+ * @param merged Whether the plan was merged with an existing cached plan (true) or
+ * is a new entry (false).
+ * @param outputMap Maps attributes from the input plan to corresponding attributes in
+ * `mergedPlan`. Used to rewrite expressions referencing the original plan
+ * to reference the merged plan instead.
+ */
+case class MergeResult(
+ mergedPlan: LogicalPlan,
+ mergedPlanIndex: Int,
+ merged: Boolean,
+ outputMap: AttributeMap[Attribute])
+
+/**
+ * Represents a plan in the PlanMerger's cache.
+ *
+ * @param plan The logical plan, which may have been merged from multiple original plans.
+ * @param merged Whether this plan is the result of merging two or more plans (true), or
+ * is an original unmerged plan (false). Merged plans typically require special
+ * handling such as wrapping in CTEs.
+ */
+case class MergedPlan(plan: LogicalPlan, merged: Boolean)
+
+/**
+ * A stateful utility for merging identical or similar logical plans to enable query plan reuse.
+ *
+ * `PlanMerger` maintains a cache of previously seen plans and attempts to either:
+ * 1. Reuse an identical plan already in the cache
+ * 2. Merge a new plan with a cached plan by combining their outputs
+ *
+ * The merging process preserves semantic equivalence while combining outputs from multiple
+ * plans into a single plan. This is primarily used by [[MergeScalarSubqueries]] to deduplicate
+ * scalar subquery execution.
+ *
+ * Supported plan types for merging:
+ * - [[Project]]: Merges project lists
+ * - [[Aggregate]]: Merges aggregate expressions with identical grouping
+ * - [[Filter]]: Requires identical filter conditions
+ * - [[Join]]: Requires identical join type, hints, and conditions
+ *
+ * @example
+ * {{{
+ * val merger = PlanMerger()
+ * val result1 = merger.merge(plan1) // Adds plan1 to cache
+ * val result2 = merger.merge(plan2) // Merges with plan1 if compatible
+ * // result2.merged == true if plans were merged
+ * // result2.outputMap maps plan2's attributes to the merged plan's attributes
+ * }}}
+ */
+class PlanMerger {
+ val cache = ArrayBuffer.empty[MergedPlan]
+
+ /**
+ * Attempts to merge the given plan with cached plans, or adds it to the cache.
+ *
+ * The method tries the following in order:
+ * 1. Check if an identical plan exists in cache (using canonicalized comparison)
+ * 2. Try to merge with each cached plan using [[tryMergePlans]]
+ * 3. If no merge is possible, add as a new cache entry
+ *
+ * @param plan The logical plan to merge or cache.
+ * @return A [[MergeResult]] containing:
+ * - The merged/cached plan to use
+ * - Its index in the cache
+ * - Whether it was merged with an existing plan
+ * - An attribute mapping for rewriting expressions
+ */
+ def merge(plan: LogicalPlan): MergeResult = {
+ cache.zipWithIndex.collectFirst(Function.unlift {
+ case (mp, i) =>
+ checkIdenticalPlans(plan, mp.plan).map { outputMap =>
+ MergeResult(mp.plan, i, true, outputMap)
+ }.orElse {
+ tryMergePlans(plan, mp.plan).map {
+ case (mergedPlan, outputMap) =>
+ cache(i) = MergedPlan(mergedPlan, true)
+ MergeResult(mergedPlan, i, true, outputMap)
+ }
+ }
+ case _ => None
+ }).getOrElse {
+ cache += MergedPlan(plan, false)
+ val outputMap = AttributeMap(plan.output.map(a => a -> a))
+ MergeResult(plan, cache.length - 1, false, outputMap)
+ }
+ }
+
+ /**
+ * Returns all plans currently in the cache as an immutable indexed sequence.
+ *
+ * @return An indexed sequence of [[MergedPlan]]s in cache order. The index of each plan
+ * corresponds to the `mergedPlanIndex` returned by [[merge]].
+ */
+ def mergedPlans(): IndexedSeq[MergedPlan] = cache.toIndexedSeq
+
+ // If 2 plans are identical return the attribute mapping from the new to the cached version.
+ private def checkIdenticalPlans(
+ newPlan: LogicalPlan,
+ cachedPlan: LogicalPlan): Option[AttributeMap[Attribute]] = {
+ if (newPlan.canonicalized == cachedPlan.canonicalized) {
+ Some(AttributeMap(newPlan.output.zip(cachedPlan.output)))
+ } else {
+ None
+ }
+ }
+
+ /**
+ * Recursively attempts to merge two plans by traversing their tree structures.
+ *
+ * Two plans can be merged if:
+ * - They are identical (canonicalized forms match), OR
+ * - They have compatible root nodes with mergeable children
+ *
+ * Supported merge patterns:
+ * - Project nodes: Combines project lists from both plans
+ * - Aggregate nodes: Combines aggregate expressions if grouping is identical and both
+ * support the same aggregate implementation (hash/object-hash/sort-based)
+ * - Filter nodes: Only if filter conditions are identical
+ * - Join nodes: Only if join type, hints, and conditions are identical
+ *
+ * @param newPlan The plan to merge into the cached plan.
+ * @param cachedPlan The cached plan to merge with.
+ * @return Some((mergedPlan, outputMap)) if merge succeeds, where:
+ * - mergedPlan is the combined plan
+ * - outputMap maps newPlan's attributes to mergedPlan's attributes
+ * Returns None if plans cannot be merged.
+ */
+ private def tryMergePlans(
Review Comment:
At the first glance, this looks identical in terms of the code. Am I right?
--
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-54136][SQL] Extract plan merging logic from `MergeScalarSubqueries` to `PlanMerger` [spark]
Posted by "peter-toth (via GitHub)" <gi...@apache.org>.
peter-toth commented on PR #52835:
URL: https://github.com/apache/spark/pull/52835#issuecomment-3491097083
If there is no objection I would like to merge this refactor by the end of this week.
--
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-54136][SQL] Extract plan merging logic from `MergeScalarSubqueries` to `PlanMerger` [spark]
Posted by "peter-toth (via GitHub)" <gi...@apache.org>.
peter-toth closed pull request #52835: [SPARK-54136][SQL] Extract plan merging logic from `MergeScalarSubqueries` to `PlanMerger`
URL: https://github.com/apache/spark/pull/52835
--
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-54136][SQL] Extract plan merging logic from `MergeScalarSubqueries` to `PlanMerger` [spark]
Posted by "cloud-fan (via GitHub)" <gi...@apache.org>.
cloud-fan commented on code in PR #52835:
URL: https://github.com/apache/spark/pull/52835#discussion_r2500796562
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/PlanMerger.scala:
##########
@@ -0,0 +1,303 @@
+/*
+ * 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.ArrayBuffer
+
+import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, AttributeMap, Expression, NamedExpression}
+import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression
+import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Filter, Join, LogicalPlan, Project}
+
+/**
+ * Result of attempting to merge a plan via [[PlanMerger.merge]].
+ *
+ * @param mergedPlan The resulting plan, either:
+ * - An existing cached plan (if identical match found)
+ * - A newly merged plan combining the input with a cached plan
+ * - The original input plan (if no merge was possible)
+ * @param mergedPlanIndex The index of this plan in the PlanMerger's cache.
+ * @param merged Whether the plan was merged with an existing cached plan (true) or
+ * is a new entry (false).
+ * @param outputMap Maps attributes from the input plan to corresponding attributes in
+ * `mergedPlan`. Used to rewrite expressions referencing the original plan
+ * to reference the merged plan instead.
+ */
+case class MergeResult(
+ mergedPlan: LogicalPlan,
+ mergedPlanIndex: Int,
+ merged: Boolean,
+ outputMap: AttributeMap[Attribute])
+
+/**
+ * Represents a plan in the PlanMerger's cache.
+ *
+ * @param plan The logical plan, which may have been merged from multiple original plans.
+ * @param merged Whether this plan is the result of merging two or more plans (true), or
+ * is an original unmerged plan (false). Merged plans typically require special
+ * handling such as wrapping in CTEs.
+ */
+case class MergedPlan(plan: LogicalPlan, merged: Boolean)
+
+/**
+ * A stateful utility for merging identical or similar logical plans to enable query plan reuse.
+ *
+ * `PlanMerger` maintains a cache of previously seen plans and attempts to either:
+ * 1. Reuse an identical plan already in the cache
+ * 2. Merge a new plan with a cached plan by combining their outputs
+ *
+ * The merging process preserves semantic equivalence while combining outputs from multiple
+ * plans into a single plan. This is primarily used by [[MergeScalarSubqueries]] to deduplicate
+ * scalar subquery execution.
+ *
+ * Supported plan types for merging:
+ * - [[Project]]: Merges project lists
+ * - [[Aggregate]]: Merges aggregate expressions with identical grouping
+ * - [[Filter]]: Requires identical filter conditions
+ * - [[Join]]: Requires identical join type, hints, and conditions
+ *
+ * @example
+ * {{{
+ * val merger = PlanMerger()
+ * val result1 = merger.merge(plan1) // Adds plan1 to cache
+ * val result2 = merger.merge(plan2) // Merges with plan1 if compatible
+ * // result2.merged == true if plans were merged
+ * // result2.outputMap maps plan2's attributes to the merged plan's attributes
+ * }}}
+ */
+class PlanMerger {
+ val cache = ArrayBuffer.empty[MergedPlan]
+
+ /**
+ * Attempts to merge the given plan with cached plans, or adds it to the cache.
+ *
+ * The method tries the following in order:
+ * 1. Check if an identical plan exists in cache (using canonicalized comparison)
+ * 2. Try to merge with each cached plan using [[tryMergePlans]]
+ * 3. If no merge is possible, add as a new cache entry
+ *
+ * @param plan The logical plan to merge or cache.
+ * @return A [[MergeResult]] containing:
+ * - The merged/cached plan to use
+ * - Its index in the cache
+ * - Whether it was merged with an existing plan
+ * - An attribute mapping for rewriting expressions
+ */
+ def merge(plan: LogicalPlan): MergeResult = {
+ cache.zipWithIndex.collectFirst(Function.unlift {
+ case (mp, i) =>
+ checkIdenticalPlans(plan, mp.plan).map { outputMap =>
+ MergeResult(mp.plan, i, true, outputMap)
+ }.orElse {
+ tryMergePlans(plan, mp.plan).map {
+ case (mergedPlan, outputMap) =>
+ cache(i) = MergedPlan(mergedPlan, true)
+ MergeResult(mergedPlan, i, true, outputMap)
+ }
+ }
+ case _ => None
+ }).getOrElse {
+ cache += MergedPlan(plan, false)
+ val outputMap = AttributeMap(plan.output.map(a => a -> a))
+ MergeResult(plan, cache.length - 1, false, outputMap)
+ }
+ }
+
+ /**
+ * Returns all plans currently in the cache as an immutable indexed sequence.
+ *
+ * @return An indexed sequence of [[MergedPlan]]s in cache order. The index of each plan
+ * corresponds to the `mergedPlanIndex` returned by [[merge]].
+ */
+ def mergedPlans(): IndexedSeq[MergedPlan] = cache.toIndexedSeq
+
+ // If 2 plans are identical return the attribute mapping from the new to the cached version.
+ private def checkIdenticalPlans(
+ newPlan: LogicalPlan,
+ cachedPlan: LogicalPlan): Option[AttributeMap[Attribute]] = {
+ if (newPlan.canonicalized == cachedPlan.canonicalized) {
+ Some(AttributeMap(newPlan.output.zip(cachedPlan.output)))
+ } else {
+ None
+ }
+ }
+
+ /**
+ * Recursively attempts to merge two plans by traversing their tree structures.
+ *
+ * Two plans can be merged if:
+ * - They are identical (canonicalized forms match), OR
+ * - They have compatible root nodes with mergeable children
+ *
+ * Supported merge patterns:
+ * - Project nodes: Combines project lists from both plans
+ * - Aggregate nodes: Combines aggregate expressions if grouping is identical and both
+ * support the same aggregate implementation (hash/object-hash/sort-based)
+ * - Filter nodes: Only if filter conditions are identical
+ * - Join nodes: Only if join type, hints, and conditions are identical
+ *
+ * @param newPlan The plan to merge into the cached plan.
+ * @param cachedPlan The cached plan to merge with.
+ * @return Some((mergedPlan, outputMap)) if merge succeeds, where:
+ * - mergedPlan is the combined plan
+ * - outputMap maps newPlan's attributes to mergedPlan's attributes
+ * Returns None if plans cannot be merged.
+ */
+ private def tryMergePlans(
+ newPlan: LogicalPlan,
+ cachedPlan: LogicalPlan): Option[(LogicalPlan, AttributeMap[Attribute])] = {
+ checkIdenticalPlans(newPlan, cachedPlan).map(cachedPlan -> _).orElse(
+ (newPlan, cachedPlan) match {
+ case (np: Project, cp: Project) =>
+ tryMergePlans(np.child, cp.child).map { case (mergedChild, outputMap) =>
+ val (mergedProjectList, newOutputMap) =
+ mergeNamedExpressions(np.projectList, outputMap, cp.projectList)
+ val mergedPlan = Project(mergedProjectList, mergedChild)
+ mergedPlan -> newOutputMap
+ }
+ case (np, cp: Project) =>
+ tryMergePlans(np, cp.child).map { case (mergedChild, outputMap) =>
+ val (mergedProjectList, newOutputMap) =
+ mergeNamedExpressions(np.output, outputMap, cp.projectList)
+ val mergedPlan = Project(mergedProjectList, mergedChild)
+ mergedPlan -> newOutputMap
+ }
+ case (np: Project, cp) =>
+ tryMergePlans(np.child, cp).map { case (mergedChild, outputMap) =>
+ val (mergedProjectList, newOutputMap) =
+ mergeNamedExpressions(np.projectList, outputMap, cp.output)
+ val mergedPlan = Project(mergedProjectList, mergedChild)
+ mergedPlan -> newOutputMap
+ }
+ case (np: Aggregate, cp: Aggregate) if supportedAggregateMerge(np, cp) =>
+ tryMergePlans(np.child, cp.child).flatMap { case (mergedChild, outputMap) =>
+ val mappedNewGroupingExpression =
+ np.groupingExpressions.map(mapAttributes(_, outputMap))
+ // Order of grouping expression does matter as merging different grouping orders can
+ // introduce "extra" shuffles/sorts that might not present in all of the original
+ // subqueries.
+ if (mappedNewGroupingExpression.map(_.canonicalized) ==
+ cp.groupingExpressions.map(_.canonicalized)) {
+ val (mergedAggregateExpressions, newOutputMap) =
+ mergeNamedExpressions(np.aggregateExpressions, outputMap, cp.aggregateExpressions)
+ val mergedPlan =
+ Aggregate(cp.groupingExpressions, mergedAggregateExpressions, mergedChild)
+ Some(mergedPlan -> newOutputMap)
+ } else {
+ None
+ }
+ }
+
+ case (np: Filter, cp: Filter) =>
+ tryMergePlans(np.child, cp.child).flatMap { case (mergedChild, outputMap) =>
+ val mappedNewCondition = mapAttributes(np.condition, outputMap)
+ // Comparing the canonicalized form is required to ignore different forms of the same
+ // expression.
+ if (mappedNewCondition.canonicalized == cp.condition.canonicalized) {
+ val mergedPlan = cp.withNewChildren(Seq(mergedChild))
+ Some(mergedPlan -> outputMap)
+ } else {
+ None
+ }
+ }
+
+ case (np: Join, cp: Join) if np.joinType == cp.joinType && np.hint == cp.hint =>
+ tryMergePlans(np.left, cp.left).flatMap { case (mergedLeft, leftOutputMap) =>
+ tryMergePlans(np.right, cp.right).flatMap { case (mergedRight, rightOutputMap) =>
+ val outputMap = leftOutputMap ++ rightOutputMap
+ val mappedNewCondition = np.condition.map(mapAttributes(_, outputMap))
+ // Comparing the canonicalized form is required to ignore different forms of the same
+ // expression and `AttributeReference.qualifier`s in `cp.condition`.
+ if (mappedNewCondition.map(_.canonicalized) == cp.condition.map(_.canonicalized)) {
+ val mergedPlan = cp.withNewChildren(Seq(mergedLeft, mergedRight))
+ Some(mergedPlan -> outputMap)
+ } else {
+ None
+ }
+ }
+ }
+
+ // Otherwise merging is not possible.
+ case _ => None
+ })
+ }
+
+ private def mapAttributes[T <: Expression](expr: T, outputMap: AttributeMap[Attribute]) = {
+ expr.transform {
+ case a: Attribute => outputMap.getOrElse(a, a)
+ }.asInstanceOf[T]
+ }
+
+ // Applies `outputMap` attribute mapping on attributes of `newExpressions` and merges them into
+ // `cachedExpressions`. Returns the merged expressions and the attribute mapping from the new to
+ // the merged version that can be propagated up during merging nodes.
+ private def mergeNamedExpressions(
+ newExpressions: Seq[NamedExpression],
+ outputMap: AttributeMap[Attribute],
+ cachedExpressions: Seq[NamedExpression]) = {
+ val mergedExpressions = ArrayBuffer[NamedExpression](cachedExpressions: _*)
+ val newOutputMap = AttributeMap(newExpressions.map { ne =>
+ val mapped = mapAttributes(ne, outputMap)
+ val withoutAlias = mapped match {
+ case Alias(child, _) => child
+ case e => e
+ }
+ ne.toAttribute -> mergedExpressions.find {
+ case Alias(child, _) => child semanticEquals withoutAlias
+ case e => e semanticEquals withoutAlias
+ }.getOrElse {
+ mergedExpressions += mapped
+ mapped
+ }.toAttribute
+ })
+ (mergedExpressions.toSeq, newOutputMap)
+ }
+
+ // Only allow aggregates of the same implementation because merging different implementations
+ // could cause performance regression.
+ private def supportedAggregateMerge(newPlan: Aggregate, cachedPlan: Aggregate) = {
+ val aggregateExpressionsSeq = Seq(newPlan, cachedPlan).map { plan =>
+ plan.aggregateExpressions.flatMap(_.collect {
+ case a: AggregateExpression => a
+ })
+ }
+ val groupByExpressionSeq = Seq(newPlan, cachedPlan).map(_.groupingExpressions)
+
+ val Seq(newPlanSupportsHashAggregate, cachedPlanSupportsHashAggregate) =
+ aggregateExpressionsSeq.zip(groupByExpressionSeq).map {
+ case (aggregateExpressions, groupByExpressions) =>
+ Aggregate.supportsHashAggregate(
+ aggregateExpressions.flatMap(
+ _.aggregateFunction.aggBufferAttributes), groupByExpressions)
+ }
+
+ newPlanSupportsHashAggregate && cachedPlanSupportsHashAggregate ||
+ newPlanSupportsHashAggregate == cachedPlanSupportsHashAggregate && {
+ val Seq(newPlanSupportsObjectHashAggregate, cachedPlanSupportsObjectHashAggregate) =
+ aggregateExpressionsSeq.zip(groupByExpressionSeq).map {
+ case (aggregateExpressions, groupByExpressions) =>
+ Aggregate.supportsObjectHashAggregate(aggregateExpressions, groupByExpressions)
+ }
+ newPlanSupportsObjectHashAggregate && cachedPlanSupportsObjectHashAggregate ||
+ newPlanSupportsObjectHashAggregate == cachedPlanSupportsObjectHashAggregate
+ }
+ }
+}
+
+object PlanMerger {
+ def apply(): PlanMerger = new PlanMerger
Review Comment:
I feel it's better to explicitly write `new PlanMerger()` instead of `PlanMerger()`
--
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-54136][SQL] Extract plan merging logic from `MergeScalarSubqueries` to `PlanMerger` [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #52835:
URL: https://github.com/apache/spark/pull/52835#discussion_r2484974873
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/MergeScalarSubqueries.scala:
##########
@@ -117,297 +116,126 @@ object MergeScalarSubqueries extends Rule[LogicalPlan] {
}
}
- /**
- * An item in the cache of merged scalar subqueries.
- *
- * @param plan The plan of a merged scalar subquery.
- * @param merged A flag to identify if this item is the result of merging subqueries.
- * Please note that `attributes.size == 1` doesn't always mean that the plan is not
- * merged as there can be subqueries that are different ([[checkIdenticalPlans]] is
- * false) due to an extra [[Project]] node in one of them. In that case
- * `attributes.size` remains 1 after merging, but the merged flag becomes true.
- * @param references A set of subquery indexes in the cache to track all (including transitive)
- * nested subqueries.
- */
- case class Header(
- plan: LogicalPlan,
- merged: Boolean,
- references: Set[Int])
-
private def extractCommonScalarSubqueries(plan: LogicalPlan) = {
- val cache = ArrayBuffer.empty[Header]
- val planWithReferences = insertReferences(plan, cache)
- cache.zipWithIndex.foreach { case (header, i) =>
- cache(i) = cache(i).copy(plan =
- if (header.merged) {
+ // Collect `ScalarSubquery` plans by level into `PlanMerger`s and insert references in place of
+ // `ScalarSubquery`s.
+ val planMergers = ArrayBuffer.empty[PlanMerger]
+ val planWithReferences = insertReferences(plan, planMergers)._1
+
+ // Traverse level by level and convert merged plans to `CTERelationDef`s and keep non-merged
+ // ones. While traversing replace references in plans back to `CTERelationRef`s or to original
+ // `ScalarSubquery`s. This is safe as a subquery plan at a level can reference only lower level
+ // other subqueries.
+ val subqueryPlansByLevel = ArrayBuffer.empty[IndexedSeq[LogicalPlan]]
+ planMergers.foreach { planMerger =>
+ val mergedPlans = planMerger.mergedPlans()
+ subqueryPlansByLevel += mergedPlans.map { mergedPlan =>
+ val planWithoutReferences = if (subqueryPlansByLevel.isEmpty) {
+ // Level 0 plans can't contain references
+ mergedPlan.plan
+ } else {
+ removeReferences(mergedPlan.plan, subqueryPlansByLevel)
+ }
+ if (mergedPlan.merged) {
CTERelationDef(
- createProject(header.plan.output, removeReferences(header.plan, cache)),
+ Project(
+ Seq(Alias(
+ CreateNamedStruct(
+ planWithoutReferences.output.flatMap(a => Seq(Literal(a.name), a))),
+ "mergedValue")()),
+ planWithoutReferences),
underSubquery = true)
} else {
- removeReferences(header.plan, cache)
- })
+ planWithoutReferences
+ }
+ }
}
- val newPlan = removeReferences(planWithReferences, cache)
- val subqueryCTEs = cache.filter(_.merged).map(_.plan.asInstanceOf[CTERelationDef])
+
+ // Replace references back to `CTERelationRef`s or to original `ScalarSubquery`s in the main
+ // plan.
+ val newPlan = removeReferences(planWithReferences, subqueryPlansByLevel)
+
+ // Add `CTERelationDef`s to the plan.
+ val subqueryCTEs = subqueryPlansByLevel.flatMap(_.collect { case cte: CTERelationDef => cte })
if (subqueryCTEs.nonEmpty) {
WithCTE(newPlan, subqueryCTEs.toSeq)
} else {
newPlan
}
}
- // First traversal builds up the cache and inserts `ScalarSubqueryReference`s to the plan.
- private def insertReferences(plan: LogicalPlan, cache: ArrayBuffer[Header]): LogicalPlan = {
- plan.transformUpWithSubqueries {
- case n => n.transformExpressionsUpWithPruning(_.containsAnyPattern(SCALAR_SUBQUERY)) {
- // The subquery could contain a hint that is not propagated once we cache it, but as a
- // non-correlated scalar subquery won't be turned into a Join the loss of hints is fine.
+ // First traversal inserts `ScalarSubqueryReference`s to the plan and tries to merge subquery
+ // plans by each level.
+ private def insertReferences(
+ plan: LogicalPlan,
+ planMergers: ArrayBuffer[PlanMerger]): (LogicalPlan, Int) = {
+ // The level of a subquery plan is maximum level of its inner subqueries + 1 or 0 if it has no
+ // inner subqueries.
+ var maxLevel = 0
+ val planWithReferences =
+ plan.transformAllExpressionsWithPruning(_.containsAnyPattern(SCALAR_SUBQUERY)) {
case s: ScalarSubquery if !s.isCorrelated && s.deterministic =>
- val (subqueryIndex, headerIndex) = cacheSubquery(s.plan, cache)
- ScalarSubqueryReference(subqueryIndex, headerIndex, s.dataType, s.exprId)
- }
- }
- }
-
- // Caching returns the index of the subquery in the cache and the index of scalar member in the
- // "Header".
- private def cacheSubquery(plan: LogicalPlan, cache: ArrayBuffer[Header]): (Int, Int) = {
- val output = plan.output.head
- val references = mutable.HashSet.empty[Int]
- plan.transformAllExpressionsWithPruning(_.containsAnyPattern(SCALAR_SUBQUERY_REFERENCE)) {
- case ssr: ScalarSubqueryReference =>
- references += ssr.subqueryIndex
- references ++= cache(ssr.subqueryIndex).references
- ssr
- }
-
- cache.zipWithIndex.collectFirst(Function.unlift {
- case (header, subqueryIndex) if !references.contains(subqueryIndex) =>
- checkIdenticalPlans(plan, header.plan).map { outputMap =>
- val mappedOutput = mapAttributes(output, outputMap)
- val headerIndex = header.plan.output.indexWhere(_.exprId == mappedOutput.exprId)
- subqueryIndex -> headerIndex
- }.orElse{
- tryMergePlans(plan, header.plan).map {
- case (mergedPlan, outputMap) =>
- val mappedOutput = mapAttributes(output, outputMap)
- val headerIndex = mergedPlan.output.indexWhere(_.exprId == mappedOutput.exprId)
- cache(subqueryIndex) = Header(mergedPlan, true, header.references ++ references)
- subqueryIndex -> headerIndex
- }
- }
- case _ => None
- }).getOrElse {
- cache += Header(plan, false, references.toSet)
- cache.length - 1 -> 0
- }
- }
-
- // If 2 plans are identical return the attribute mapping from the new to the cached version.
- private def checkIdenticalPlans(
- newPlan: LogicalPlan,
- cachedPlan: LogicalPlan): Option[AttributeMap[Attribute]] = {
- if (newPlan.canonicalized == cachedPlan.canonicalized) {
- Some(AttributeMap(newPlan.output.zip(cachedPlan.output)))
- } else {
- None
- }
- }
-
- // Recursively traverse down and try merging 2 plans. If merge is possible then return the merged
- // plan with the attribute mapping from the new to the merged version.
- // Please note that merging arbitrary plans can be complicated, the current version supports only
- // some of the most important nodes.
- private def tryMergePlans(
- newPlan: LogicalPlan,
- cachedPlan: LogicalPlan): Option[(LogicalPlan, AttributeMap[Attribute])] = {
- checkIdenticalPlans(newPlan, cachedPlan).map(cachedPlan -> _).orElse(
- (newPlan, cachedPlan) match {
- case (np: Project, cp: Project) =>
- tryMergePlans(np.child, cp.child).map { case (mergedChild, outputMap) =>
- val (mergedProjectList, newOutputMap) =
- mergeNamedExpressions(np.projectList, outputMap, cp.projectList)
- val mergedPlan = Project(mergedProjectList, mergedChild)
- mergedPlan -> newOutputMap
- }
- case (np, cp: Project) =>
- tryMergePlans(np, cp.child).map { case (mergedChild, outputMap) =>
- val (mergedProjectList, newOutputMap) =
- mergeNamedExpressions(np.output, outputMap, cp.projectList)
- val mergedPlan = Project(mergedProjectList, mergedChild)
- mergedPlan -> newOutputMap
- }
- case (np: Project, cp) =>
- tryMergePlans(np.child, cp).map { case (mergedChild, outputMap) =>
- val (mergedProjectList, newOutputMap) =
- mergeNamedExpressions(np.projectList, outputMap, cp.output)
- val mergedPlan = Project(mergedProjectList, mergedChild)
- mergedPlan -> newOutputMap
- }
- case (np: Aggregate, cp: Aggregate) if supportedAggregateMerge(np, cp) =>
- tryMergePlans(np.child, cp.child).flatMap { case (mergedChild, outputMap) =>
- val mappedNewGroupingExpression =
- np.groupingExpressions.map(mapAttributes(_, outputMap))
- // Order of grouping expression does matter as merging different grouping orders can
- // introduce "extra" shuffles/sorts that might not present in all of the original
- // subqueries.
- if (mappedNewGroupingExpression.map(_.canonicalized) ==
- cp.groupingExpressions.map(_.canonicalized)) {
- val (mergedAggregateExpressions, newOutputMap) =
- mergeNamedExpressions(np.aggregateExpressions, outputMap, cp.aggregateExpressions)
- val mergedPlan =
- Aggregate(cp.groupingExpressions, mergedAggregateExpressions, mergedChild)
- Some(mergedPlan -> newOutputMap)
- } else {
- None
- }
- }
-
- case (np: Filter, cp: Filter) =>
- tryMergePlans(np.child, cp.child).flatMap { case (mergedChild, outputMap) =>
- val mappedNewCondition = mapAttributes(np.condition, outputMap)
- // Comparing the canonicalized form is required to ignore different forms of the same
- // expression.
- if (mappedNewCondition.canonicalized == cp.condition.canonicalized) {
- val mergedPlan = cp.withNewChildren(Seq(mergedChild))
- Some(mergedPlan -> outputMap)
- } else {
- None
- }
- }
-
- case (np: Join, cp: Join) if np.joinType == cp.joinType && np.hint == cp.hint =>
- tryMergePlans(np.left, cp.left).flatMap { case (mergedLeft, leftOutputMap) =>
- tryMergePlans(np.right, cp.right).flatMap { case (mergedRight, rightOutputMap) =>
- val outputMap = leftOutputMap ++ rightOutputMap
- val mappedNewCondition = np.condition.map(mapAttributes(_, outputMap))
- // Comparing the canonicalized form is required to ignore different forms of the same
- // expression and `AttributeReference.quailifier`s in `cp.condition`.
- if (mappedNewCondition.map(_.canonicalized) == cp.condition.map(_.canonicalized)) {
- val mergedPlan = cp.withNewChildren(Seq(mergedLeft, mergedRight))
- Some(mergedPlan -> outputMap)
- } else {
- None
- }
- }
- }
-
- // Otherwise merging is not possible.
- case _ => None
- })
- }
-
- private def createProject(attributes: Seq[Attribute], plan: LogicalPlan): Project = {
- Project(
- Seq(Alias(
- CreateNamedStruct(attributes.flatMap(a => Seq(Literal(a.name), a))),
- "mergedValue")()),
- plan)
- }
-
- private def mapAttributes[T <: Expression](expr: T, outputMap: AttributeMap[Attribute]) = {
- expr.transform {
- case a: Attribute => outputMap.getOrElse(a, a)
- }.asInstanceOf[T]
- }
-
- // Applies `outputMap` attribute mapping on attributes of `newExpressions` and merges them into
- // `cachedExpressions`. Returns the merged expressions and the attribute mapping from the new to
- // the merged version that can be propagated up during merging nodes.
- private def mergeNamedExpressions(
- newExpressions: Seq[NamedExpression],
- outputMap: AttributeMap[Attribute],
- cachedExpressions: Seq[NamedExpression]) = {
- val mergedExpressions = ArrayBuffer[NamedExpression](cachedExpressions: _*)
- val newOutputMap = AttributeMap(newExpressions.map { ne =>
- val mapped = mapAttributes(ne, outputMap)
- val withoutAlias = mapped match {
- case Alias(child, _) => child
- case e => e
- }
- ne.toAttribute -> mergedExpressions.find {
- case Alias(child, _) => child semanticEquals withoutAlias
- case e => e semanticEquals withoutAlias
- }.getOrElse {
- mergedExpressions += mapped
- mapped
- }.toAttribute
- })
- (mergedExpressions.toSeq, newOutputMap)
- }
-
- // Only allow aggregates of the same implementation because merging different implementations
- // could cause performance regression.
- private def supportedAggregateMerge(newPlan: Aggregate, cachedPlan: Aggregate) = {
- val aggregateExpressionsSeq = Seq(newPlan, cachedPlan).map { plan =>
- plan.aggregateExpressions.flatMap(_.collect {
- case a: AggregateExpression => a
- })
- }
- val groupByExpressionSeq = Seq(newPlan, cachedPlan).map(_.groupingExpressions)
-
- val Seq(newPlanSupportsHashAggregate, cachedPlanSupportsHashAggregate) =
- aggregateExpressionsSeq.zip(groupByExpressionSeq).map {
- case (aggregateExpressions, groupByExpressions) =>
- Aggregate.supportsHashAggregate(
- aggregateExpressions.flatMap(
- _.aggregateFunction.aggBufferAttributes), groupByExpressions)
- }
-
- newPlanSupportsHashAggregate && cachedPlanSupportsHashAggregate ||
- newPlanSupportsHashAggregate == cachedPlanSupportsHashAggregate && {
- val Seq(newPlanSupportsObjectHashAggregate, cachedPlanSupportsObjectHashAggregate) =
- aggregateExpressionsSeq.zip(groupByExpressionSeq).map {
- case (aggregateExpressions, groupByExpressions) =>
- Aggregate.supportsObjectHashAggregate(aggregateExpressions, groupByExpressions)
- }
- newPlanSupportsObjectHashAggregate && cachedPlanSupportsObjectHashAggregate ||
- newPlanSupportsObjectHashAggregate == cachedPlanSupportsObjectHashAggregate
+ val (planWithReferences, level) = insertReferences(s.plan, planMergers)
+
+ while (level >= planMergers.size) planMergers += PlanMerger()
+ // The subquery could contain a hint that is not propagated once we merge it, but as a
+ // non-correlated scalar subquery won't be turned into a Join the loss of hints is fine.
+ val planMergeResult = planMergers(level).merge(planWithReferences)
+
+ maxLevel = maxLevel.max(level + 1)
+
+ val mergedOutput = planMergeResult.outputMap(planWithReferences.output.head)
+ val headerIndex =
+ planMergeResult.mergedPlan.output.indexWhere(_.exprId == mergedOutput.exprId)
+ ScalarSubqueryReference(
+ level,
+ planMergeResult.mergedPlanIndex,
+ headerIndex,
+ s.dataType,
+ s.exprId)
+ case o => o
}
+ (planWithReferences, maxLevel)
}
// Second traversal replaces `ScalarSubqueryReference`s to either
// `GetStructField(ScalarSubquery(CTERelationRef to the merged plan)` if the plan is merged from
// multiple subqueries or `ScalarSubquery(original plan)` if it isn't.
private def removeReferences(
plan: LogicalPlan,
- cache: ArrayBuffer[Header]) = {
- plan.transformUpWithSubqueries {
- case n =>
- n.transformExpressionsWithPruning(_.containsAnyPattern(SCALAR_SUBQUERY_REFERENCE)) {
- case ssr: ScalarSubqueryReference =>
- val header = cache(ssr.subqueryIndex)
- if (header.merged) {
- val subqueryCTE = header.plan.asInstanceOf[CTERelationDef]
- GetStructField(
- ScalarSubquery(
- CTERelationRef(subqueryCTE.id, _resolved = true, subqueryCTE.output,
- subqueryCTE.isStreaming),
- exprId = ssr.exprId),
- ssr.headerIndex)
- } else {
- ScalarSubquery(header.plan, exprId = ssr.exprId)
- }
+ subqueryPlansByLevel: ArrayBuffer[IndexedSeq[LogicalPlan]]) = {
+ plan.transformAllExpressionsWithPruning(_.containsAnyPattern(SCALAR_SUBQUERY_REFERENCE)) {
+ case ssr: ScalarSubqueryReference =>
+ subqueryPlansByLevel(ssr.level)(ssr.mergedPlanIndex) match {
+ case cte: CTERelationDef =>
+ GetStructField(
+ ScalarSubquery(
+ CTERelationRef(cte.id, _resolved = true, cte.output, cte.isStreaming),
+ exprId = ssr.exprId),
+ ssr.outputIndex)
+ case o => ScalarSubquery(o, exprId = ssr.exprId)
}
}
}
}
/**
- * Temporal reference to a cached subquery.
+ * Temporal reference to a subquery which is added to a `PlanMerger`.
*
- * @param subqueryIndex A subquery index in the cache.
- * @param headerIndex An index in the output of merged subquery.
- * @param dataType The dataType of origin scalar subquery.
+ * @param level The level of the replaced subquery. It defines the `PlanMerger` instance into which
+ * the subquery is merged.
+ * @param mergedPlanIndex The index of the merged plan in the `PlanMerger`.
+ * @param outputIndex The index of the output attribute of the merged plan.
+ * @param dataType The dataType of original scalar subquery.
+ * @param exprId The expression id of the original scalar subquery.
*/
case class ScalarSubqueryReference(
- subqueryIndex: Int,
- headerIndex: Int,
+ level: Int,
+ mergedPlanIndex: Int,
+ outputIndex: Int,
dataType: DataType,
exprId: ExprId) extends LeafExpression with Unevaluable {
override def nullable: Boolean = true
final override val nodePatterns: Seq[TreePattern] = Seq(SCALAR_SUBQUERY_REFERENCE)
-
- override def stringArgs: Iterator[Any] = Iterator(subqueryIndex, headerIndex, dataType, exprId.id)
Review Comment:
May I ask why we don't need to override `stringArgs` at all after this PR? Is this related to this PR's contribution or this was not used before this PR?
--
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-54136][SQL] Extract plan merging logic from `MergeScalarSubqueries` to `PlanMerger` [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #52835:
URL: https://github.com/apache/spark/pull/52835#discussion_r2484975684
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/MergeScalarSubqueries.scala:
##########
@@ -117,297 +116,126 @@ object MergeScalarSubqueries extends Rule[LogicalPlan] {
}
}
- /**
- * An item in the cache of merged scalar subqueries.
- *
- * @param plan The plan of a merged scalar subquery.
- * @param merged A flag to identify if this item is the result of merging subqueries.
- * Please note that `attributes.size == 1` doesn't always mean that the plan is not
- * merged as there can be subqueries that are different ([[checkIdenticalPlans]] is
- * false) due to an extra [[Project]] node in one of them. In that case
- * `attributes.size` remains 1 after merging, but the merged flag becomes true.
- * @param references A set of subquery indexes in the cache to track all (including transitive)
- * nested subqueries.
- */
- case class Header(
- plan: LogicalPlan,
- merged: Boolean,
- references: Set[Int])
-
private def extractCommonScalarSubqueries(plan: LogicalPlan) = {
- val cache = ArrayBuffer.empty[Header]
- val planWithReferences = insertReferences(plan, cache)
- cache.zipWithIndex.foreach { case (header, i) =>
- cache(i) = cache(i).copy(plan =
- if (header.merged) {
+ // Collect `ScalarSubquery` plans by level into `PlanMerger`s and insert references in place of
+ // `ScalarSubquery`s.
+ val planMergers = ArrayBuffer.empty[PlanMerger]
+ val planWithReferences = insertReferences(plan, planMergers)._1
+
+ // Traverse level by level and convert merged plans to `CTERelationDef`s and keep non-merged
+ // ones. While traversing replace references in plans back to `CTERelationRef`s or to original
+ // `ScalarSubquery`s. This is safe as a subquery plan at a level can reference only lower level
+ // other subqueries.
+ val subqueryPlansByLevel = ArrayBuffer.empty[IndexedSeq[LogicalPlan]]
+ planMergers.foreach { planMerger =>
+ val mergedPlans = planMerger.mergedPlans()
+ subqueryPlansByLevel += mergedPlans.map { mergedPlan =>
+ val planWithoutReferences = if (subqueryPlansByLevel.isEmpty) {
+ // Level 0 plans can't contain references
+ mergedPlan.plan
+ } else {
+ removeReferences(mergedPlan.plan, subqueryPlansByLevel)
+ }
+ if (mergedPlan.merged) {
CTERelationDef(
- createProject(header.plan.output, removeReferences(header.plan, cache)),
+ Project(
+ Seq(Alias(
+ CreateNamedStruct(
+ planWithoutReferences.output.flatMap(a => Seq(Literal(a.name), a))),
+ "mergedValue")()),
+ planWithoutReferences),
underSubquery = true)
} else {
- removeReferences(header.plan, cache)
- })
+ planWithoutReferences
+ }
+ }
}
- val newPlan = removeReferences(planWithReferences, cache)
- val subqueryCTEs = cache.filter(_.merged).map(_.plan.asInstanceOf[CTERelationDef])
+
+ // Replace references back to `CTERelationRef`s or to original `ScalarSubquery`s in the main
+ // plan.
+ val newPlan = removeReferences(planWithReferences, subqueryPlansByLevel)
+
+ // Add `CTERelationDef`s to the plan.
+ val subqueryCTEs = subqueryPlansByLevel.flatMap(_.collect { case cte: CTERelationDef => cte })
if (subqueryCTEs.nonEmpty) {
WithCTE(newPlan, subqueryCTEs.toSeq)
} else {
newPlan
}
}
- // First traversal builds up the cache and inserts `ScalarSubqueryReference`s to the plan.
- private def insertReferences(plan: LogicalPlan, cache: ArrayBuffer[Header]): LogicalPlan = {
- plan.transformUpWithSubqueries {
- case n => n.transformExpressionsUpWithPruning(_.containsAnyPattern(SCALAR_SUBQUERY)) {
- // The subquery could contain a hint that is not propagated once we cache it, but as a
- // non-correlated scalar subquery won't be turned into a Join the loss of hints is fine.
+ // First traversal inserts `ScalarSubqueryReference`s to the plan and tries to merge subquery
+ // plans by each level.
+ private def insertReferences(
+ plan: LogicalPlan,
+ planMergers: ArrayBuffer[PlanMerger]): (LogicalPlan, Int) = {
+ // The level of a subquery plan is maximum level of its inner subqueries + 1 or 0 if it has no
+ // inner subqueries.
+ var maxLevel = 0
+ val planWithReferences =
+ plan.transformAllExpressionsWithPruning(_.containsAnyPattern(SCALAR_SUBQUERY)) {
case s: ScalarSubquery if !s.isCorrelated && s.deterministic =>
- val (subqueryIndex, headerIndex) = cacheSubquery(s.plan, cache)
- ScalarSubqueryReference(subqueryIndex, headerIndex, s.dataType, s.exprId)
- }
- }
- }
-
- // Caching returns the index of the subquery in the cache and the index of scalar member in the
- // "Header".
- private def cacheSubquery(plan: LogicalPlan, cache: ArrayBuffer[Header]): (Int, Int) = {
Review Comment:
Is this method removed simply?
--
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-54136][SQL] Extract plan merging logic from `MergeScalarSubqueries` to `PlanMerger` [spark]
Posted by "peter-toth (via GitHub)" <gi...@apache.org>.
peter-toth commented on code in PR #52835:
URL: https://github.com/apache/spark/pull/52835#discussion_r2485989939
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/MergeScalarSubqueriesSuite.scala:
##########
@@ -568,14 +568,14 @@ class MergeScalarSubqueriesSuite extends PlanTest {
val subquery5 = ScalarSubquery(testRelation.select((Symbol("a") + 2).as("a_plus2_2")))
val subquery6 = ScalarSubquery(testRelation.select(Symbol("b").as("b_2")))
val originalQuery = testRelation
- .select(
Review Comment:
This minor chnage in the order of subqueries is needed because the modified `MergeScalarSubqueries` encounters them in different order. So we either change the expected query to expect `a_plus1_2` and `a_plus2_2` or reorder the subqueries in the query like here. I applied the latter.
--
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-54136][SQL] Extract plan merging logic from `MergeScalarSubqueries` to `PlanMerger` [spark]
Posted by "peter-toth (via GitHub)" <gi...@apache.org>.
peter-toth commented on code in PR #52835:
URL: https://github.com/apache/spark/pull/52835#discussion_r2485847706
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/MergeScalarSubqueries.scala:
##########
@@ -117,297 +116,126 @@ object MergeScalarSubqueries extends Rule[LogicalPlan] {
}
}
- /**
- * An item in the cache of merged scalar subqueries.
- *
- * @param plan The plan of a merged scalar subquery.
- * @param merged A flag to identify if this item is the result of merging subqueries.
- * Please note that `attributes.size == 1` doesn't always mean that the plan is not
- * merged as there can be subqueries that are different ([[checkIdenticalPlans]] is
- * false) due to an extra [[Project]] node in one of them. In that case
- * `attributes.size` remains 1 after merging, but the merged flag becomes true.
- * @param references A set of subquery indexes in the cache to track all (including transitive)
- * nested subqueries.
- */
- case class Header(
- plan: LogicalPlan,
- merged: Boolean,
- references: Set[Int])
-
private def extractCommonScalarSubqueries(plan: LogicalPlan) = {
- val cache = ArrayBuffer.empty[Header]
- val planWithReferences = insertReferences(plan, cache)
- cache.zipWithIndex.foreach { case (header, i) =>
- cache(i) = cache(i).copy(plan =
- if (header.merged) {
+ // Collect `ScalarSubquery` plans by level into `PlanMerger`s and insert references in place of
+ // `ScalarSubquery`s.
+ val planMergers = ArrayBuffer.empty[PlanMerger]
+ val planWithReferences = insertReferences(plan, planMergers)._1
+
+ // Traverse level by level and convert merged plans to `CTERelationDef`s and keep non-merged
+ // ones. While traversing replace references in plans back to `CTERelationRef`s or to original
+ // `ScalarSubquery`s. This is safe as a subquery plan at a level can reference only lower level
+ // other subqueries.
+ val subqueryPlansByLevel = ArrayBuffer.empty[IndexedSeq[LogicalPlan]]
+ planMergers.foreach { planMerger =>
+ val mergedPlans = planMerger.mergedPlans()
+ subqueryPlansByLevel += mergedPlans.map { mergedPlan =>
+ val planWithoutReferences = if (subqueryPlansByLevel.isEmpty) {
+ // Level 0 plans can't contain references
+ mergedPlan.plan
+ } else {
+ removeReferences(mergedPlan.plan, subqueryPlansByLevel)
+ }
+ if (mergedPlan.merged) {
CTERelationDef(
- createProject(header.plan.output, removeReferences(header.plan, cache)),
+ Project(
+ Seq(Alias(
+ CreateNamedStruct(
+ planWithoutReferences.output.flatMap(a => Seq(Literal(a.name), a))),
+ "mergedValue")()),
+ planWithoutReferences),
underSubquery = true)
} else {
- removeReferences(header.plan, cache)
- })
+ planWithoutReferences
+ }
+ }
}
- val newPlan = removeReferences(planWithReferences, cache)
- val subqueryCTEs = cache.filter(_.merged).map(_.plan.asInstanceOf[CTERelationDef])
+
+ // Replace references back to `CTERelationRef`s or to original `ScalarSubquery`s in the main
+ // plan.
+ val newPlan = removeReferences(planWithReferences, subqueryPlansByLevel)
+
+ // Add `CTERelationDef`s to the plan.
+ val subqueryCTEs = subqueryPlansByLevel.flatMap(_.collect { case cte: CTERelationDef => cte })
if (subqueryCTEs.nonEmpty) {
WithCTE(newPlan, subqueryCTEs.toSeq)
} else {
newPlan
}
}
- // First traversal builds up the cache and inserts `ScalarSubqueryReference`s to the plan.
- private def insertReferences(plan: LogicalPlan, cache: ArrayBuffer[Header]): LogicalPlan = {
- plan.transformUpWithSubqueries {
- case n => n.transformExpressionsUpWithPruning(_.containsAnyPattern(SCALAR_SUBQUERY)) {
- // The subquery could contain a hint that is not propagated once we cache it, but as a
- // non-correlated scalar subquery won't be turned into a Join the loss of hints is fine.
+ // First traversal inserts `ScalarSubqueryReference`s to the plan and tries to merge subquery
+ // plans by each level.
+ private def insertReferences(
+ plan: LogicalPlan,
+ planMergers: ArrayBuffer[PlanMerger]): (LogicalPlan, Int) = {
+ // The level of a subquery plan is maximum level of its inner subqueries + 1 or 0 if it has no
+ // inner subqueries.
+ var maxLevel = 0
+ val planWithReferences =
+ plan.transformAllExpressionsWithPruning(_.containsAnyPattern(SCALAR_SUBQUERY)) {
case s: ScalarSubquery if !s.isCorrelated && s.deterministic =>
- val (subqueryIndex, headerIndex) = cacheSubquery(s.plan, cache)
- ScalarSubqueryReference(subqueryIndex, headerIndex, s.dataType, s.exprId)
- }
- }
- }
-
- // Caching returns the index of the subquery in the cache and the index of scalar member in the
- // "Header".
- private def cacheSubquery(plan: LogicalPlan, cache: ArrayBuffer[Header]): (Int, Int) = {
Review Comment:
Briefly, `cacheSubquery()` became `PlanMerger.merge()` with some modifications.
As `PlanMerger.merge()` is kind of an internal API now, which can be used by other rules. It returns `MergeResult` instead of the tuple after this refactor.
The `Header` objects, which contined the information about the merged plans became `MergedPlan`s and can be accessed as `PlanMerger.mergedPlans()` by the rules.
The main difference between `cacheSubquery()` and the new `PlanMerger.merge()` is that I removed `Header.references` and the related safeguards like the `if !references.contains(subqueryIndex)` below.
Previously, we filled up `Header.references` and with nested subquery indices so as to avoid trying to merge a subquery with other subqueries that are nested into it.
E.g.
```
SELECT (
SELECT avg(a) FROM t WHERE c = ( -- subquery 1
SELECT sum(b) FROM t2) -- subquery 2 nested into subquery 1
)
)
```
Here it doesn't make sense to try merging the plan of `subquery 1` and `subquery 2` as the former contains the latter. The problem of trying to merge the 2 was not just it doesn't make sense, but in certain weird cases (see the example in [SPARK-40618](https://issues.apache.org/jira/browse/SPARK-40618)) the merge was successful and resulted in invalid plans.
But because `PlanMerger` should be a general plan merging tool and in some future rules it will be used to merge other than subqueries, tracking the nested subqueries in it is not the best way to deal with the problem.
This is why the refactored `MergeScalarSubqueries` uses a different approach. We have one `PlanMerger` object per each subquery level. E.g. `subquery 2` is a level 0 (leaf) subquery, while `subquery 1` is a level 1 subquery because it has an inner, level 0 subquery. So the 2 will be added to different `PlanMerger`s. This way a level n subquery is tried to merge with other level n subqueries only.
--
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