Skip to content

[SPARK-54865][CONNECT][SQL] Add foreachWithSubqueriesAndPruning method to QueryPlan #53637

New issue

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

Sign up for GitHub

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

Already on GitHub? Sign in to your account

Closed
heyihong wants to merge 4 commits into from
Closed

[SPARK-54865][CONNECT][SQL] Add foreachWithSubqueriesAndPruning method to QueryPlan #53637

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
d18d8b5
[SPARK-54865] Add foreachWithSubqueriesAndPruning method to QueryPlan
heyihong Dec 29, 2025
ecff57d
foreachWithPruning
heyihong Jan 5, 2026
02cb32d
undo foreachWithPruning
heyihong Jan 7, 2026
9954cae
back to the initial version
heyihong Jan 7, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 14 additions & 0 deletions sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/QueryPlan.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -630,6 +630,20 @@ abstract class QueryPlan[PlanType <: QueryPlan[PlanType]]
foreach(actualFunc)
}

/**
* A variant of [[foreachWithSubqueries]] with pruning support.
* Only traverses nodes that match the given condition.
*/
def foreachWithSubqueriesAndPruning(
cond: TreePatternBits => Boolean)(f: PlanType => Unit): Unit = {
Copy link
Contributor

@cloud-fan cloud-fan Jan 3, 2026

Choose a reason for hiding this comment

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

other pruning methods also have a ruleId parameter, shall we follow?

Copy link
Contributor Author

@heyihong heyihong Jan 5, 2026 •
edited
Loading

Choose a reason for hiding this comment

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

IMHO, ruleId is not applicable here, since it is used for transformation—it represents the transformation rule ID. Also, ruleId is not used by default, as the default value is UnknownRuleId.

If we really need this parameter, we can extend it later (less is more).

if (!cond.apply(this)) {
return
}
f(this)
subqueries.foreach(_.foreachWithSubqueriesAndPruning(cond)(f))
Copy link
Contributor

@cloud-fan cloud-fan Jan 3, 2026

Choose a reason for hiding this comment

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

I think foreachWithSubqueries traverse the children first, not subqueries. Shall we follow it here?

Copy link
Contributor

@cloud-fan cloud-fan Jan 3, 2026

Choose a reason for hiding this comment

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

or it's actually interleaved in foreachWithSubqueries

Copy link
Contributor Author

@heyihong heyihong Jan 5, 2026 •
edited
Loading

Choose a reason for hiding this comment

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

The traversal order of foreachWithSubqueries should also be subqueries first, then children, but its traversal implementation is a bit hard to read and understand, in my opinion.

After taking a closer look at the code, I think it may be a style choice for TreeNode methods to handle plan traversal logic (including withPruning), while QueryPlan simply leverages that to perform additional tasks, such as handling subqueries.

So the correctness may not be affected, but style-wise, it makes more sense to move the pruning logic to TreeNode.

Copy link
Contributor Author

@heyihong heyihong Jan 7, 2026 •
edited
Loading

Choose a reason for hiding this comment

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

Discussed with @cloud-fan offline, while the current convention is to let TreeNode methods to handle plan traversal logic (including withPruning), for foreachWithSubqueriesAndPruning, a clearer implementation should avoid doing so.

children.foreach(_.foreachWithSubqueriesAndPruning(cond)(f))
}

/**
* A variant of `collect`. This method not only apply the given function to all elements in this
* plan, also considering all the plans in its (nested) subqueries.
Expand Down
27 changes: 26 additions & 1 deletion sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/plans/QueryPlanSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,8 @@

package org.apache.spark.sql.catalyst.plans

import scala.collection.mutable.ArrayBuffer

import org.apache.spark.SparkFunSuite
import org.apache.spark.sql.catalyst.TableIdentifier
import org.apache.spark.sql.catalyst.analysis.UnresolvedRelation
Expand All @@ -25,7 +27,7 @@ import org.apache.spark.sql.catalyst.dsl.plans._
import org.apache.spark.sql.catalyst.expressions.{Alias, AttributeReference, Expression, ListQuery, Literal, NamedExpression, Rand}
import org.apache.spark.sql.catalyst.plans.logical.{Filter, LocalRelation, LogicalPlan, Project, Union}
import org.apache.spark.sql.catalyst.rules.Rule
import org.apache.spark.sql.catalyst.trees.{CurrentOrigin, Origin}
import org.apache.spark.sql.catalyst.trees.{CurrentOrigin, Origin, TreePattern}
import org.apache.spark.sql.types.IntegerType

class QueryPlanSuite extends SparkFunSuite {
Expand Down Expand Up @@ -160,4 +162,27 @@ class QueryPlanSuite extends SparkFunSuite {
val planAfterTestRule = testRule(plan)
assert(planAfterTestRule.output(0).nullable)
}

test("SPARK-54865: pruning works correctly in foreachWithSubqueriesAndPruning") {
val a: NamedExpression = AttributeReference("a", IntegerType)()
val plan = Project(
Seq(a),
Filter(
ListQuery(Project(
Seq(a),
UnresolvedRelation(TableIdentifier("t", None))
)),
UnresolvedRelation(TableIdentifier("t", None))
)
)

val visited = ArrayBuffer[LogicalPlan]()
plan.foreachWithSubqueriesAndPruning(_.containsPattern(TreePattern.FILTER)) { p =>
visited += p
}

// Only 2 nodes contain FILTER pattern: outer Project and Filter
assert(visited.size == 2)
assert(visited.forall(_.containsPattern(TreePattern.FILTER)))
}
}
7 changes: 4 additions & 3 deletions ...nect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -242,17 +242,18 @@ class SparkConnectPlanner(
}

if (executeHolderOpt.isDefined) {
plan.transformUpWithSubqueriesAndPruning(_.containsPattern(TreePattern.COLLECT_METRICS)) {
plan.foreachWithSubqueriesAndPruning(_.containsPattern(TreePattern.COLLECT_METRICS)) {
case collectMetrics: CollectMetrics if !collectMetrics.child.isStreaming =>
// TODO this might be too complex for no good reason. It might
// be easier to inspect the plan after it completes.
val observation = session.observationManager.getOrNewObservation(
collectMetrics.name,
collectMetrics.dataframeId)
executeHolder.addObservation(collectMetrics.name, observation)
collectMetrics
case _ =>
}
} else plan
}
plan
}

private def transformRelationPlugin(extension: ProtoAny): LogicalPlan = {
Expand Down
4 changes: 3 additions & 1 deletion sql/core/src/main/scala/org/apache/spark/sql/classic/ObservationManager.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ import java.util.concurrent.ConcurrentHashMap

import org.apache.spark.sql.{Observation, Row}
import org.apache.spark.sql.catalyst.plans.logical.CollectMetrics
import org.apache.spark.sql.catalyst.trees.TreePattern
import org.apache.spark.sql.execution.QueryExecution
import org.apache.spark.sql.util.QueryExecutionListener

Expand Down Expand Up @@ -54,7 +55,8 @@ private[sql] class ObservationManager(session: SparkSession) {

private def tryComplete(qe: QueryExecution): Unit = {
val allMetrics = qe.observedMetrics
qe.logical.foreach {
qe.logical.foreachWithSubqueriesAndPruning(
_.containsPattern(TreePattern.COLLECT_METRICS)) {
case c: CollectMetrics =>
val keyExists = observations.containsKey((c.name, c.dataframeId))
val metrics = allMetrics.get(c.name)
Expand Down