You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by we...@apache.org on 2026/08/12 03:29:32 UTC
(spark) branch master updated: [SPARK-58696][SQL] Fix CTE pushdown discarding injected filters
This is an automated email from the ASF dual-hosted git repository.
cloud-fan pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 12a34234c4e6 [SPARK-58696][SQL] Fix CTE pushdown discarding injected filters
12a34234c4e6 is described below
commit 12a34234c4e6ff150f2b1d2c0ea39e0617ea0fea
Author: James Xu <xu...@gmail.com>
AuthorDate: Wed Aug 12 11:29:15 2026 +0800
[SPARK-58696][SQL] Fix CTE pushdown discarding injected filters
### What changes were proposed in this pull request?
Fix `PushdownPredicatesAndPruneColumnsForCTEDef` discarding filters that other optimizer rules injected into a CTE definition between this rule's applications.
The rule previously rebuilt a CTE definition from the frozen pre-push-down snapshot stored in `CTERelationDef.originalPlanWithPredicates` whenever newly gathered reference predicates re-armed its guard. With this PR the rule rebuilds from the **current** child instead: it removes only the push-down filter it placed in the previous pass (located by a descent that mirrors the predicate push-down rules' own translations through `Project`/`Aggregate` aliases — using the same `AliasHelper` [...]
A cross-reference note is added to `PushPredicateThroughNonJoin` so future changes to push-through cases update the mirror.
### Why are the changes needed?
Take the following SQL as an example:
```
CREATE OR REPLACE TEMP VIEW t1 AS
SELECT id AS a FROM RANGE(10) UNION ALL SELECT CAST(NULL AS BIGINT);
CREATE OR REPLACE TEMP VIEW t2 AS SELECT id AS a FROM RANGE(10);
EXPLAIN EXTENDED
WITH c AS (
SELECT t1.a AS a, t2.a AS b, rand(0) AS r
FROM t1 JOIN t2 ON t1.a = t2.a
)
SELECT a, b FROM c WHERE a = 5
UNION ALL
SELECT a, b FROM c WHERE a = 7;
```
The plan snippet for `c` is the following:
```
:- Filter (isnotnull(a#43L) AND (a#43L = 5))
: +- Project [a#52L AS a#43L, b#63L AS b#44L]
: +- Exchange RoundRobinPartitioning(200), REPARTITION_BY_NUM, [plan_id=120]
: +- Project [a#52L, a#56L AS b#63L]
: +- BroadcastHashJoin [a#52L], [a#56L], Inner, BuildRight, false
: :- Project [id#61L AS a#52L]
: : +- Filter ((id#61L = 5) OR (id#61L = 7))
: : +- Range (0, 10, step=1, splits=12)
: +- BroadcastExchange HashedRelationBroadcastMode(List(input[0, bigint, false]),false), [plan_id=117]
: +- Project [id#62L AS a#56L]
: +- Range (0, 10, step=1, splits=12)
```
The key here is that for the CTE `c`, the reference site have predicate for table t1(e.g. `a=5`), so it gets the predicates(`Filter ((id#61L = 5) OR (id#61L = 7))` for the first Range). But no explicit predicate for t2(The second Range), so in real production query, this means a full table scan, very bad performance.
With the help of https://github.com/apache/spark/pull/56499 we can now infer `t2.a=5` from the reference site: `a=5`, but currently the predicate is stripped by PushdownPredicatesAndPruneColumnsForCTEDef.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
- New `PushdownPredicatesForCTEDefStalenessSuite` (6 tests): the staleness regression test **fails on unmodified master**; further tests cover idempotency under foreign mutation and previous-push removal through `Project`/`Join`/`Union`/`Aggregate`/`Window`.
- `CTEInlineSuite*` + `CTEHintSuite`: 63/63 pass — the plan shapes asserted by existing CTE push-down tests (including "combined predicate") are unchanged.
- Full `sql/catalyst` optimizer package: 1401/1401 pass.
- `catalyst/scalastyle` and `catalyst/Test/scalastyle`: clean.
### Was this patch authored or generated with the assistance of AI tools?
Yes. An AI coding agent (Claude) assisted with root-cause analysis, reproduction, fix design, and test authoring. All changes were reviewed, directed, and validated by the author; full test suites were run locally as listed above.
Closes #57901 from xumingming/cte-pushdown-stale-snapshot-fix.
Authored-by: James Xu <xu...@gmail.com>
Signed-off-by: Wenchen Fan <we...@databricks.com>
---
.../spark/sql/catalyst/optimizer/Optimizer.scala | 4 +
...ushdownPredicatesAndPruneColumnsForCTEDef.scala | 127 ++++++-
.../sql/catalyst/plans/logical/cteOperators.scala | 14 +-
...PushdownPredicatesForCTEDefStalenessSuite.scala | 366 +++++++++++++++++++++
4 files changed, 493 insertions(+), 18 deletions(-)
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala
index 802e86374f50..ea78b8be49be 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala
@@ -2194,6 +2194,10 @@ object PushDownPredicates extends Rule[LogicalPlan] {
* 2) the predicate is deterministic and the operator will not change any of rows.
* 3) We don't add double evaluation OR double evaluation would be cheap OR we're configured to.
*
+ * Note: if a new push-through case is added here, or the translation applied to pushed
+ * conditions changes (e.g. how aliases are substituted), also update
+ * `removePushedDownFilter` in [[PushdownPredicatesAndPruneColumnsForCTEDef]], which mirrors
+ * this rule's cases to locate and remove filters previously pushed into CTE definitions.
*/
object PushPredicateThroughNonJoin extends Rule[LogicalPlan] with PredicateHelper {
def apply(plan: LogicalPlan): LogicalPlan = plan transform applyLocally
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/PushdownPredicatesAndPruneColumnsForCTEDef.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/PushdownPredicatesAndPruneColumnsForCTEDef.scala
index 2339596920b0..7a6827cefcb3 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/PushdownPredicatesAndPruneColumnsForCTEDef.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/PushdownPredicatesAndPruneColumnsForCTEDef.scala
@@ -19,7 +19,8 @@ package org.apache.spark.sql.catalyst.optimizer
import scala.collection.mutable
-import org.apache.spark.sql.catalyst.expressions.{And, Attribute, AttributeSet, Expression, Literal, Or, SubqueryExpression}
+import org.apache.spark.sql.catalyst.expressions.{And, Attribute, AttributeSet}
+import org.apache.spark.sql.catalyst.expressions.{Expression, Literal, Or, PredicateHelper, SubqueryExpression}
import org.apache.spark.sql.catalyst.planning.PhysicalOperation
import org.apache.spark.sql.catalyst.plans.logical._
import org.apache.spark.sql.catalyst.rules.Rule
@@ -30,7 +31,7 @@ import org.apache.spark.util.collection.Utils
* Infer predicates and column pruning for [[CTERelationDef]] from its reference points, and push
* the disjunctive predicates as well as the union of attributes down the CTE plan.
*/
-object PushdownPredicatesAndPruneColumnsForCTEDef extends Rule[LogicalPlan] {
+object PushdownPredicatesAndPruneColumnsForCTEDef extends Rule[LogicalPlan] with PredicateHelper {
// CTE_id - (CTE_definition, precedence, predicates_to_push_down, attributes_to_prune)
private type CTEMap = mutable.HashMap[Long, (CTERelationDef, Int, Seq[Expression], AttributeSet)]
@@ -113,32 +114,52 @@ object PushdownPredicatesAndPruneColumnsForCTEDef extends Rule[LogicalPlan] {
* In order to guarantee idempotency, we keep the predicates (if any) being pushed down by the
* last iteration of this rule in a temporary field of `CTERelationDef`, so that on the current
* iteration, we only push down predicates for a CTE def if there exists any new predicate that
- * has not been pushed before. Also, since part of a new predicate might overlap with some
- * existing predicate and it can be hard to extract only the non-overlapping part, we also keep
- * the original CTE definition plan without any predicate push-down in that temporary field so
- * that when we do a new predicate push-down, we can construct a new plan with all latest
- * predicates over the original plan without having to figure out the exact predicate difference.
+ * has not been pushed before. When such a new predicate push-down happens, the CTE definition
+ * is rebuilt from its CURRENT child: the push-down filter this rule placed in the previous
+ * iteration is removed (wherever it sits) and the result is wrapped with the latest combined
+ * predicate. This preserves any change other rules made to the CTE definition's child in
+ * between (e.g. filters injected by `InferFiltersFromConstraints`, which runs in the `Once`
+ * batch sandwiched between the two fixedPoint batches containing this rule). If the previous
+ * push-down can no longer be found (another rule rewrote or merged it with other filters),
+ * the current child is used as-is: re-pushing the combined predicate is redundant but always
+ * semantics-preserving, since the disjunction of the reference predicates is valid for every
+ * row of the CTE definition.
*/
private def pushdownPredicatesAndAttributes(
plan: LogicalPlan,
cteMap: CTEMap): LogicalPlan = plan.transformWithSubqueries {
case cteDef @ CTERelationDef(child, id, originalPlanWithPredicates, _, _, _) =>
val (_, _, newPreds, newAttrSet) = cteMap(id)
- val originalPlan = originalPlanWithPredicates.map(_._1).getOrElse(child)
val preds = originalPlanWithPredicates.map(_._2).getOrElse(Seq.empty)
if (!isTruePredicate(newPreds) &&
newPreds.exists(newPred => !preds.exists(_.semanticEquals(newPred)))) {
+ val basePlan = originalPlanWithPredicates match {
+ case Some((_, prevPreds)) if prevPreds.nonEmpty =>
+ // Remove the push-down filter this rule placed in the previous iteration. It is
+ // usually the top-level node of the child, but rules sharing the fixedPoint
+ // batches with this rule (e.g. `PushDownPredicates`) may have moved it deeper,
+ // possibly across attribute-renaming projections - hence the comparison is done
+ // on canonicalized conditions. If the previous push-down can no longer be found
+ // (another rule rewrote or merged it with other filters), the current child is
+ // used as-is: re-pushing the combined predicate is redundant but always
+ // semantics-preserving, since the disjunction of the reference predicates is
+ // valid for every row of the CTE definition.
+ removePushedDownFilter(child, prevPreds.reduce(Or))
+ case _ => child
+ }
val newCombinedPred = newPreds.reduce(Or)
- val newChild = if (needsPruning(originalPlan, newAttrSet)) {
- Project(newAttrSet.toSeq, originalPlan)
+ val newChild = if (needsPruning(basePlan, newAttrSet)) {
+ Project(newAttrSet.toSeq, basePlan)
} else {
- originalPlan
+ basePlan
}
cteDef.copy(child = Filter(newCombinedPred, newChild),
- originalPlanWithPredicates = Some((originalPlan, newPreds)))
+ // The plan component of `originalPlanWithPredicates` is recorded but never read
+ // back: on the next iteration only the pushed predicates are consulted (see the
+ // `basePlan` computation above).
+ originalPlanWithPredicates = Some((basePlan, newPreds)))
} else if (needsPruning(cteDef.child, newAttrSet)) {
- cteDef.copy(child = Project(newAttrSet.toSeq, cteDef.child),
- originalPlanWithPredicates = Some((originalPlan, preds)))
+ cteDef.copy(child = Project(newAttrSet.toSeq, cteDef.child))
} else {
cteDef
}
@@ -156,6 +177,84 @@ object PushdownPredicatesAndPruneColumnsForCTEDef extends Rule[LogicalPlan] {
}
}
+ /**
+ * Removes the previous push-down filter (identified by its condition, `predicate`) from
+ * `plan`, wherever predicate push-down rules sharing the fixedPoint batches with this rule
+ * (e.g. `PushDownPredicates`) may have moved it. The descent mirrors the cases of
+ * `PushPredicateThroughNonJoin` and `PushPredicateThroughJoin`, translating `predicate` back
+ * the same way they translate the pushed condition (for `Project` and `Aggregate` it uses
+ * the very same `AliasHelper` utilities, so the two cannot drift apart): through projection
+ * and grouping-key aliases, positionally into each branch (`Union`), and unchanged
+ * through operators that pass the referenced attributes verbatim (`Join`, `Window`, and
+ * output-preserving unary nodes like `Filter`, `Sort`, `Repartition`). Descent stops at
+ * operators that remap attributes in other ways (e.g. `Generate`, `Expand`): if the filter
+ * cannot be located, the input plan is returned unchanged, and the caller re-pushes on top,
+ * which is redundant but always semantics-preserving.
+ *
+ * Ancestors of the removed filter are rebuilt with `withNewChildren` rather than direct
+ * case-class copies so that `TreeNode` tags (e.g. `Project.hiddenOutputTag`, which the
+ * analyzer sets on the projection above a natural/USING join) survive the rebuild.
+ */
+ private def removePushedDownFilter(plan: LogicalPlan, predicate: Expression): LogicalPlan = {
+ def remove(current: LogicalPlan, target: Expression): (LogicalPlan, Boolean) = current match {
+ case Filter(cond, inner) if cond.canonicalized == target.canonicalized =>
+ (inner, true)
+ case p: Project =>
+ // Mirror PushPredicateThroughNonJoin: translate the target through the projection's
+ // aliases with the same helper it uses to move the filter below the projection.
+ val translated = replaceAlias(target, getAliasMap(p))
+ val (newChild, removed) = remove(p.child, translated)
+ if (removed) (p.withNewChildren(Seq(newChild)), true) else (p, false)
+ case j: Join =>
+ val (newLeft, removedFromLeft) = remove(j.left, target)
+ if (removedFromLeft) {
+ (j.withNewChildren(Seq(newLeft, j.right)), true)
+ } else {
+ val (newRight, removedFromRight) = remove(j.right, target)
+ if (removedFromRight) (j.withNewChildren(Seq(j.left, newRight)), true) else (j, false)
+ }
+ case u: Union =>
+ // PushDownPredicates copies the filter into every branch, mapping the union output
+ // attributes to each branch's output positionally; remove it from every branch where
+ // it is found. The ExprId-to-output-index map is built once for all branches.
+ val outputIndexByExprId = u.output.map(_.exprId).zipWithIndex.toMap
+ var removedAny = false
+ val newChildren = u.children.map { branch =>
+ val branchTarget = target.transform {
+ case a: Attribute if outputIndexByExprId.contains(a.exprId) =>
+ branch.output(outputIndexByExprId(a.exprId))
+ }
+ val (newBranch, removed) = remove(branch, branchTarget)
+ if (removed) {
+ removedAny = true
+ newBranch
+ } else {
+ branch
+ }
+ }
+ if (removedAny) (u.withNewChildren(newChildren), true) else (u, false)
+ case agg: Aggregate =>
+ // Mirror PushPredicateThroughNonJoin: translate the target through the grouping
+ // aliases with the same helper it uses to move grouping-key filters below the
+ // aggregate (filters referencing aggregate expressions stay up, so they are always
+ // found above this node).
+ val translated = replaceAlias(target, getAliasMap(agg))
+ val (newChild, removed) = remove(agg.child, translated)
+ if (removed) (agg.withNewChildren(Seq(newChild)), true) else (agg, false)
+ case w: Window =>
+ // PushDownPredicates pushes filters referencing only partition columns below the
+ // window unchanged (partition columns are input attributes, so no translation).
+ val (newChild, removed) = remove(w.child, target)
+ if (removed) (w.withNewChildren(Seq(newChild)), true) else (w, false)
+ case other if other.children.length == 1 &&
+ other.outputSet == other.children.head.outputSet =>
+ val (newChild, removed) = remove(other.children.head, target)
+ if (removed) (other.withNewChildren(Seq(newChild)), true) else (other, false)
+ case _ => (current, false)
+ }
+ remove(plan, predicate)._1
+ }
+
private def isTruePredicate(predicates: Seq[Expression]): Boolean = {
predicates.length == 1 && predicates.head == Literal.TrueLiteral
}
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/cteOperators.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/cteOperators.scala
index d7a2f0e76e37..dd9a14c4ce33 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/cteOperators.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/cteOperators.scala
@@ -99,10 +99,16 @@ case class UnionLoopRef(
* A wrapper for CTE definition plan with a unique ID.
* @param child The CTE definition query plan.
* @param id The unique ID for this CTE definition.
- * @param originalPlanWithPredicates The original query plan before predicate pushdown and the
- * predicates that have been pushed down into `child`. This is
- * a temporary field used by optimization rules for CTE predicate
- * pushdown to help ensure rule idempotency.
+ * @param originalPlanWithPredicates The base plan of the last predicate pushdown and the
+ * predicates that have been pushed down into `child`. The
+ * base plan (the definition's child at the time of the last
+ * pushdown, with the filter that pushdown inserted removed)
+ * is recorded but not read back: the rule rebuilds from the
+ * definition's current child and only consults the
+ * predicates, both to detect newly appeared predicates and
+ * to locate the previous pushdown for removal. This is a
+ * temporary field used by optimization rules for CTE
+ * predicate pushdown to help ensure rule idempotency.
* @param underSubquery If true, it means we don't need to add a shuffle for this CTE relation as
* subquery reuse will be applied to reuse CTE relation output.
* @param maxDepth The maximal depth of a recursion in a recursive CTE.
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/PushdownPredicatesForCTEDefStalenessSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/PushdownPredicatesForCTEDefStalenessSuite.scala
new file mode 100644
index 000000000000..b696ec307dee
--- /dev/null
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/PushdownPredicatesForCTEDefStalenessSuite.scala
@@ -0,0 +1,366 @@
+/*
+ * 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 org.apache.spark.sql.catalyst.expressions.{Alias, And, Attribute, AttributeReference}
+import org.apache.spark.sql.catalyst.expressions.{Ascending, CurrentRow, RowFrame, RowNumber}
+import org.apache.spark.sql.catalyst.expressions.{EqualTo, Expression, GreaterThan}
+import org.apache.spark.sql.catalyst.expressions.{IsNotNull, Literal}
+import org.apache.spark.sql.catalyst.expressions.{SortOrder, SpecifiedWindowFrame, UnboundedPreceding}
+import org.apache.spark.sql.catalyst.expressions.{WindowExpression, WindowSpecDefinition}
+import org.apache.spark.sql.catalyst.plans.{Inner, PlanTest}
+import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, CTERelationDef, CTERelationRef}
+import org.apache.spark.sql.catalyst.plans.logical.{Filter, Join, JoinHint, LocalRelation}
+import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, Project, Union, Window, WithCTE}
+import org.apache.spark.sql.catalyst.trees.TreeNodeTag
+import org.apache.spark.sql.types.IntegerType
+
+/**
+ * Regression test for a staleness defect in
+ * [[PushdownPredicatesAndPruneColumnsForCTEDef]]: on its second pass the rule rebuilt a
+ * CTE definition from the plan snapshot stored in `originalPlanWithPredicates`, which was
+ * taken during its first pass. Any change made to the CTE definition's child by other rules
+ * in between (e.g. filters injected by `InferFiltersFromConstraints`, which runs in the
+ * "Infer Filters" batch sandwiched between the two fixed-point batches that both contain
+ * this rule) was silently discarded by the rebuild.
+ *
+ * The inter-pass mutations are produced by applying the real rules
+ * ([[InferFiltersFromConstraints]] and [[PushPredicateThroughNonJoin]]) between two real
+ * applications of the rule under test, so the suite keeps exercising the actual cross-rule
+ * interaction as those rules evolve.
+ */
+class PushdownPredicatesForCTEDefStalenessSuite extends PlanTest {
+
+ test("CTE def rebuild must not discard filters injected between rule passes") {
+ val t1a = AttributeReference("a", IntegerType, nullable = true)()
+ val t2b = AttributeReference("b", IntegerType, nullable = true)()
+ val t1 = LocalRelation(t1a)
+ val t2 = LocalRelation(t2b)
+ val join = Join(t1, t2, Inner, Some(EqualTo(t1a, t2b)), JoinHint.NONE)
+
+ val cteId = 0L
+ val cteDef = CTERelationDef(join, cteId)
+ val plan = withTwoRefs(cteDef)(
+ out => EqualTo(out(0), Literal(5)),
+ out => EqualTo(out(0), Literal(7)))
+
+ // Pass 1 (batch "Operator Optimization before Inferring Filters"): the rule pushes the
+ // combined reference predicates into the CTE definition and records them in
+ // originalPlanWithPredicates.
+ val afterPass1 = PushdownPredicatesAndPruneColumnsForCTEDef.apply(plan)
+ assert(theOnlyDef(afterPass1).originalPlanWithPredicates.isDefined)
+
+ // Run the real "Infer Filters" batch rule between the two passes (it is sandwiched
+ // between the two fixed-point batches that both contain this rule). It strengthens the
+ // pushed filter itself (propagating a = 5 | a = 7 through the join condition onto b),
+ // injects isnotnull filters below the join, and enriches the reference sites with
+ // isnotnull, which re-arms the rule's guard on the next pass.
+ val afterInfer = InferFiltersFromConstraints.apply(afterPass1)
+ val injectedConds = theOnlyDef(afterInfer).child.collect { case f: Filter => f.condition }
+ val pass1FilterCount = theOnlyDef(afterPass1).child.collect { case f: Filter => f }.length
+ assert(injectedConds.length > pass1FilterCount,
+ "test setup failed: InferFiltersFromConstraints did not inject filters into the " +
+ "CTE definition")
+
+ // Pass 2 (batch "Operator Optimization after Inferring Filters"): the rule sees the
+ // enriched reference predicates and rebuilds the definition. Every filter present
+ // after the Infer Filters batch must survive the rebuild; before the fix the rebuild
+ // used the stale first-pass snapshot and discarded all of them.
+ val defAfterPass2 = theOnlyDef(PushdownPredicatesAndPruneColumnsForCTEDef.apply(afterInfer))
+ injectedConds.foreach { cond =>
+ assert(hasFilterOn(defAfterPass2.child, cond),
+ s"PushdownPredicatesAndPruneColumnsForCTEDef discarded a filter that " +
+ s"InferFiltersFromConstraints injected between its two passes: $cond")
+ }
+ }
+
+ test("rule is idempotent when no new predicates appear between passes") {
+ val t1a = AttributeReference("a", IntegerType, nullable = true)()
+ val t2b = AttributeReference("b", IntegerType, nullable = true)()
+ val t1 = LocalRelation(t1a)
+ val t2 = LocalRelation(t2b)
+ val join = Join(t1, t2, Inner, Some(EqualTo(t1a, t2b)), JoinHint.NONE)
+
+ val cteId = 0L
+ val cteDef = CTERelationDef(join, cteId)
+ val plan = withTwoRefs(cteDef)(
+ out => EqualTo(out(0), Literal(5)),
+ out => EqualTo(out(0), Literal(7)))
+
+ // First application pushes the combined reference predicates and records them.
+ val once = PushdownPredicatesAndPruneColumnsForCTEDef.apply(plan)
+ // A second application with no new reference predicates must be a no-op:
+ // no re-push, no stacked filters.
+ val twice = PushdownPredicatesAndPruneColumnsForCTEDef.apply(once)
+ comparePlans(once, twice)
+
+ // Even after a foreign mutation inside the CTE definition (e.g. a filter injected
+ // by InferFiltersFromConstraints), as long as no NEW reference-site predicate
+ // appears, the rule must leave the plan untouched. (The mutation is hand-written
+ // because no real rule mutates only the definition without also enriching the
+ // reference sites, which is exactly what would re-arm the guard.)
+ val mutated = once.transform {
+ case d @ CTERelationDef(Filter(cond, j: Join), `cteId`, Some(_), _, _, _) =>
+ d.copy(child = Filter(cond, j.copy(right = Filter(GreaterThan(t2b, Literal(0)), j.right))))
+ }
+ val afterMutation = PushdownPredicatesAndPruneColumnsForCTEDef.apply(mutated)
+ comparePlans(mutated, afterMutation)
+ }
+
+ test("rebuild removes the previous push-down even after it was pushed deeper") {
+ val t1a = AttributeReference("a", IntegerType, nullable = true)()
+ val t2b = AttributeReference("b", IntegerType, nullable = true)()
+ val t1 = LocalRelation(t1a)
+ val t2 = LocalRelation(t2b)
+ val join = Join(t1, t2, Inner, Some(EqualTo(t1a, t2b)), JoinHint.NONE)
+
+ // The CTE definition ends in a renaming projection, like a view selecting aliased columns.
+ val pa = Alias(t1a, "pa")()
+ val pb = Alias(t2b, "pb")()
+ val project = Project(Seq(pa, pb), join)
+
+ val cteId = 0L
+ val cteDef = CTERelationDef(project, cteId)
+ val plan = withTwoRefs(cteDef)(
+ out => EqualTo(out(0), Literal(5)),
+ out => EqualTo(out(0), Literal(7)))
+
+ // Pass 1: pushes the combined predicate on top of the definition's projection.
+ val afterPass1 = PushdownPredicatesAndPruneColumnsForCTEDef.apply(plan)
+
+ // Consume the pushed filter with the real push-down rule that shares the fixedPoint
+ // batches with the rule under test: it moves below the renaming projection with the
+ // attributes rewritten to the projection's input.
+ val consumed = addIsNotNullToRef(PushPredicateThroughNonJoin.apply(afterPass1),
+ cteId, Literal(7))
+ assert(theOnlyDef(consumed).child match {
+ case Project(_, Filter(_, _: Join)) => true
+ case _ => false
+ }, "test setup failed: push-down did not move the filter below the projection")
+
+ // Pass 2: the rule must remove its previous push-down from wherever push-down left
+ // it; otherwise the rebuilt definition carries the old and new combined predicates
+ // as redundant stacked filters.
+ assertSingleEnrichedTopFilter(consumed,
+ "previous push-down was not removed before re-pushing")
+ }
+
+ test("rebuild removes the previous push-down pushed into union branches") {
+ val l1 = AttributeReference("a", IntegerType, nullable = true)()
+ val l2 = AttributeReference("b", IntegerType, nullable = true)()
+ val r1 = AttributeReference("a", IntegerType, nullable = true)()
+ val r2 = AttributeReference("b", IntegerType, nullable = true)()
+ // The union output shares exprIds with the first branch; the def output is the union output.
+ val union = Union(Seq(LocalRelation(l1, l2), LocalRelation(r1, r2)))
+
+ val cteId = 0L
+ val cteDef = CTERelationDef(union, cteId)
+ val plan = withTwoRefs(cteDef)(
+ out => And(EqualTo(out(0), Literal(5)), GreaterThan(out(1), Literal(0))),
+ out => And(EqualTo(out(0), Literal(7)), GreaterThan(out(1), Literal(0))))
+
+ val afterPass1 = PushdownPredicatesAndPruneColumnsForCTEDef.apply(plan)
+
+ // The real push-down rule copies the pushed filter into every union branch,
+ // translating the union output attributes to each branch's output positionally.
+ val consumed = addIsNotNullToRef(PushPredicateThroughNonJoin.apply(afterPass1),
+ cteId, Literal(7))
+ assert(theOnlyDef(consumed).child match {
+ case u: Union => u.children.forall(_.isInstanceOf[Filter])
+ case _ => false
+ }, "test setup failed: push-down did not copy the filter into every union branch")
+
+ assertSingleEnrichedTopFilter(consumed,
+ "previous push-down was not removed from the union branches")
+ }
+
+ test("rebuild removes the previous push-down pushed below an aggregate") {
+ val a = AttributeReference("a", IntegerType, nullable = true)()
+ val b = AttributeReference("b", IntegerType, nullable = true)()
+ val pa = Alias(a, "pa")()
+ // Grouping-only aggregate (distinct), so the definition has a single output column.
+ val aggregate = Aggregate(Seq(pa), Seq(pa), LocalRelation(a, b))
+
+ val cteId = 0L
+ val cteDef = CTERelationDef(aggregate, cteId)
+ val plan = withTwoRefs(cteDef)(
+ out => EqualTo(out(0), Literal(5)),
+ out => EqualTo(out(0), Literal(7)))
+
+ val afterPass1 = PushdownPredicatesAndPruneColumnsForCTEDef.apply(plan)
+
+ // The real push-down rule moves the filter below the aggregate, translating the
+ // grouping alias back to its child attribute.
+ val consumed = addIsNotNullToRef(PushPredicateThroughNonJoin.apply(afterPass1),
+ cteId, Literal(7))
+ assert(theOnlyDef(consumed).child match {
+ case Aggregate(_, _, _: Filter, _) => true
+ case _ => false
+ }, "test setup failed: push-down did not move the filter below the aggregate")
+
+ assertSingleEnrichedTopFilter(consumed,
+ "previous push-down was not removed below the aggregate")
+ }
+
+ test("rebuild removes the previous push-down pushed below a window") {
+ val a = AttributeReference("a", IntegerType, nullable = true)()
+ val b = AttributeReference("b", IntegerType, nullable = true)()
+ val rn = Alias(
+ WindowExpression(RowNumber(), WindowSpecDefinition(Seq(a), Seq(SortOrder(a, Ascending)),
+ SpecifiedWindowFrame(RowFrame, UnboundedPreceding, CurrentRow))),
+ "rn")()
+ val window = Window(Seq(rn), Seq(a), Seq(SortOrder(a, Ascending)), LocalRelation(a, b))
+
+ val cteId = 0L
+ val cteDef = CTERelationDef(window, cteId)
+ val plan = withTwoRefs(cteDef)(
+ out => EqualTo(out(0), Literal(5)),
+ out => EqualTo(out(0), Literal(7)))
+
+ val afterPass1 = PushdownPredicatesAndPruneColumnsForCTEDef.apply(plan)
+
+ // The real push-down rule moves the filter below the window unchanged (the predicate
+ // references only the partition column, which is an input attribute).
+ val consumed = addIsNotNullToRef(PushPredicateThroughNonJoin.apply(afterPass1),
+ cteId, Literal(7))
+ assert(theOnlyDef(consumed).child match {
+ case w: Window => w.child.isInstanceOf[Filter]
+ case _ => false
+ }, "test setup failed: push-down did not move the filter below the window")
+
+ assertSingleEnrichedTopFilter(consumed,
+ "previous push-down was not removed below the window")
+ }
+
+ test("rebuild preserves TreeNode tags on the nodes it rebuilds") {
+ val t1a = AttributeReference("a", IntegerType, nullable = true)()
+ val t2b = AttributeReference("b", IntegerType, nullable = true)()
+ val t1 = LocalRelation(t1a)
+ val t2 = LocalRelation(t2b)
+ val join = Join(t1, t2, Inner, Some(EqualTo(t1a, t2b)), JoinHint.NONE)
+ // Identity projection mimicking the analyzer-inserted projection above a natural or
+ // USING join, which carries the hidden join-key columns in Project.hiddenOutputTag.
+ val project = Project(Seq(t1a, t2b), join)
+
+ val cteId = 0L
+ val cteDef = CTERelationDef(project, cteId)
+ val plan = withTwoRefs(cteDef)(
+ out => EqualTo(out(0), Literal(5)),
+ out => EqualTo(out(0), Literal(7)))
+
+ val afterPass1 = PushdownPredicatesAndPruneColumnsForCTEDef.apply(plan)
+
+ // The real push-down rules move the pushed filter below the projection and then into
+ // the join's left branch (the predicate references only the left side of the inner
+ // join), so the next rebuild has to rebuild both the projection and the join.
+ val consumed = addIsNotNullToRef(
+ PushPredicateThroughJoin.apply(PushPredicateThroughNonJoin.apply(afterPass1)),
+ cteId, Literal(7))
+ assert(theOnlyDef(consumed).child match {
+ case Project(_, Join(_: Filter, _, _, _, _)) => true
+ case _ => false
+ }, "test setup failed: push-down did not move the filter into the join branch")
+
+ // Tag the nodes on the removal path in place, so the tags are present when the rule
+ // under test rebuilds them. (The intermediate push-down rules rebuild nodes with
+ // plain copies of their own; their tag handling is out of scope here.)
+ val taggedProject = theOnlyDef(consumed).child.asInstanceOf[Project]
+ val taggedJoin = taggedProject.child.asInstanceOf[Join]
+ taggedProject.setTagValue(Project.hiddenOutputTag, Seq(t2b))
+ taggedJoin.setTagValue(testTag, "preserved")
+
+ // Pass 2: removing the previous push-down from the join branch rebuilds both the
+ // join and the projection. The rebuild must carry their tags over.
+ val rebuilt = theOnlyDef(PushdownPredicatesAndPruneColumnsForCTEDef.apply(consumed)).child
+ val rebuiltProject = rebuilt.collect { case p: Project => p }.head
+ val rebuiltJoin = rebuilt.collect { case j: Join => j }.head
+ assert(rebuiltProject.getTagValue(Project.hiddenOutputTag).contains(Seq(t2b)),
+ s"rebuild dropped Project.hiddenOutputTag: $rebuilt")
+ assert(rebuiltJoin.getTagValue(testTag).contains("preserved"),
+ s"rebuild dropped tags on the join: $rebuilt")
+ }
+
+ /** A tag with no consumer, used to verify that rebuilds preserve arbitrary tags. */
+ private val testTag = TreeNodeTag[String]("cte_pushdown_test_tag")
+
+ /**
+ * Applies the rule under test to `plan` (pass 2) and asserts that the rebuilt CTE
+ * definition carries exactly one filter: the freshly re-pushed combined predicate
+ * (recognizable by the IsNotNull enrichment). Any leftover copy of the previous
+ * push-down shows up as an additional filter and fails the test.
+ */
+ private def assertSingleEnrichedTopFilter(plan: LogicalPlan, message: String): Unit = {
+ val cteDef = theOnlyDef(PushdownPredicatesAndPruneColumnsForCTEDef.apply(plan))
+ val filters = cteDef.child.collect { case f: Filter => f }
+ assert(filters.length == 1, s"$message: ${cteDef.child}")
+ assert(filters.head.condition.find(_.isInstanceOf[IsNotNull]).isDefined,
+ s"expected the enriched combined predicate on top of the def: ${cteDef.child}")
+ }
+
+ /**
+ * Builds one reference to the given CTE definition: `Project -> Filter -> CTERelationRef`
+ * with fresh output attribute instances (`CTERelationRef` is a `MultiInstanceRelation`,
+ * so every reference must own fresh exprIds). The project list covers all definition
+ * output columns so that column pruning never kicks in. `pred` receives the fresh
+ * output attributes and returns the reference-site predicate.
+ */
+ private def mkRef(cteId: Long, defOutput: Seq[Attribute])(
+ pred: Seq[Attribute] => Expression): Project = {
+ val out = defOutput.map(_.newInstance())
+ Project(out, Filter(pred(out), CTERelationRef(cteId, true, out, false)))
+ }
+
+ /**
+ * Builds `WithCTE(Union(refs), Seq(cteDef))` with two reference sites carrying the given
+ * predicates, like a view referenced from multiple call sites.
+ */
+ private def withTwoRefs(cteDef: CTERelationDef)(
+ pred1: Seq[Attribute] => Expression,
+ pred2: Seq[Attribute] => Expression): LogicalPlan = {
+ val refs = Seq(pred1, pred2).map(p => mkRef(cteDef.id, cteDef.output)(p))
+ WithCTE(Union(refs), Seq(cteDef))
+ }
+
+ /** The single CTE definition in the given plan. */
+ private def theOnlyDef(plan: LogicalPlan): CTERelationDef = {
+ plan.collect { case d: CTERelationDef => d }.head
+ }
+
+ /**
+ * Simulates the reference-site enrichment `InferFiltersFromConstraints` performs between
+ * the rule's two passes: conjoins `IsNotNull` on the column the site predicate compares
+ * to `marker` (e.g. the second reference's `Literal(7)`), which re-arms the rule's
+ * rebuild on the next pass. The removal tests use this targeted mutation instead of the
+ * real rule because the real rule would also rewrite the pushed filter inside the
+ * definition, defeating the exact-match removal those tests isolate.
+ */
+ private def addIsNotNullToRef(
+ plan: LogicalPlan, cteId: Long, marker: Literal): LogicalPlan = {
+ plan.transform {
+ case f @ Filter(cond, ref: CTERelationRef) if ref.cteId == cteId =>
+ cond.collectFirst { case EqualTo(a: Attribute, m: Literal) if m == marker => a } match {
+ case Some(a) => Filter(And(cond, IsNotNull(a)), ref)
+ case _ => f
+ }
+ }
+ }
+
+ private def hasFilterOn(plan: LogicalPlan, condition: Expression): Boolean = {
+ plan.collect { case f: Filter => f.condition }.exists(_.semanticEquals(condition))
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org