pratham76 commented on code in PR #58656:
URL: https://github.com/apache/spark/pull/58656#discussion_r4274565628
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/subquery.scala:
##########
@@ -398,62 +406,412 @@ object RewritePredicateSubquery extends
Rule[LogicalPlan] with PredicateHelper {
(newExpr, newPlan)
}
+ /**
+ * Same as [[rewriteExistentialExpr]], but it also returns the newly
introduced attributes, and
+ * it only rewrites the existential sub-queries for which `canRewrite`
returns true. A sub-query
+ * that is not rewritten stays in the returned expression as it is, and is
not descended into:
+ * rewriting an existential sub-query nested in its join condition would
graft an existence join
+ * onto the plan for an `exists` reference that the sub-query left in place
may never evaluate.
+ */
private def rewriteExistentialExprWithAttrs(
exprs: Seq[Expression],
- plan: LogicalPlan): (Option[Expression], LogicalPlan, Seq[Attribute]) = {
+ plan: LogicalPlan,
+ canRewrite: Expression => Boolean = _ => true): (Option[Expression],
LogicalPlan,
+ Seq[Attribute]) = {
var newPlan = plan
val introducedAttrs = ArrayBuffer.empty[Attribute]
- val newExprs = exprs.map { e =>
- e.transformDownWithPruning(_.containsAnyPattern(EXISTS_SUBQUERY,
IN_SUBQUERY)) {
- case Exists(sub, _, _, conditions, subHint) =>
- val exists = AttributeReference("exists", BooleanType, nullable =
false)()
- val existenceJoin = ExistenceJoin(exists)
- val newCondition = conditions.reduceLeftOption(And)
- newPlan =
- buildJoin(newPlan, rewriteDomainJoinsIfPresent(newPlan, sub,
newCondition),
- existenceJoin, newCondition, subHint)
- introducedAttrs += exists
- exists
- case Not(InSubquery(values, ListQuery(sub, _, _, _, conditions,
subHint))) =>
- val exists = AttributeReference("exists", BooleanType, nullable =
false)()
- // Deduplicate conflicting attributes if any.
- val newSub = dedupSubqueryOnSelfJoin(newPlan, sub, Some(values))
- val inConditions = values.zip(newSub.output).map(EqualTo.tupled)
- // To handle a null-aware predicate not-in-subquery in nested
conditions
- // (e.g., `v > 0 OR t1.id NOT IN (SELECT id FROM t2)`), we transform
- // `inCondition` (t1.id=t2.id) into `(inCondition) OR
ISNULL(inCondition)`.
- //
- // For example, `SELECT * FROM t1 WHERE v > 0 OR t1.id NOT IN
(SELECT id FROM t2)`
- // is transformed into a plan below;
- // == Optimized Logical Plan ==
- // Project [id#78, v#79]
- // +- Filter ((v#79 > 0) OR NOT exists#83)
- // +- Join ExistenceJoin(exists#83), ((id#78 = id#80) OR
isnull((id#78 = id#80)))
- // :- Relation[id#78,v#79] parquet
- // +- Relation[id#80] parquet
- val nullAwareJoinConds = inConditions.map(c => Or(c, IsNull(c)))
- val finalJoinCond = (nullAwareJoinConds ++
conditions).reduceLeft(And)
- val joinHint = JoinHint(None, subHint)
- newPlan = Join(newPlan,
- rewriteDomainJoinsIfPresent(newPlan, newSub, Some(finalJoinCond)),
- ExistenceJoin(exists), Some(finalJoinCond), joinHint)
- introducedAttrs += exists
- Not(exists)
- case InSubquery(values, ListQuery(sub, _, _, _, conditions, subHint))
=>
- val exists = AttributeReference("exists", BooleanType, nullable =
false)()
- // Deduplicate conflicting attributes if any.
- val newSub = dedupSubqueryOnSelfJoin(newPlan, sub, Some(values))
- val inConditions = values.zip(newSub.output).map(EqualTo.tupled)
- val newConditions = (inConditions ++
conditions).reduceLeftOption(And)
- val joinHint = JoinHint(None, subHint)
- newPlan = Join(newPlan, rewriteDomainJoinsIfPresent(newPlan, newSub,
newConditions),
- ExistenceJoin(exists), newConditions, joinHint)
- introducedAttrs += exists
- exists
+ // Builds the replacement of a sub-query expression under that
expression's origin and with
+ // its tags, as TreeNode.transformDownWithPruning does for the nodes a
rule replaces. Without
+ // it every node built below, the `exists` attribute and the join
conditions included, would
+ // take the origin of the enclosing operator, and a runtime error raised
from the rewritten
+ // condition would quote that operator rather than the sub-query fragment.
+ def replacing(sq: Expression)(replacement: => Expression): Expression = {
+ val newExpr = CurrentOrigin.withOrigin(sq.origin)(replacement)
+ newExpr.copyTagsFrom(sq)
+ newExpr
+ }
+ def rewrite(expr: Expression): Expression = {
+ if (!expr.containsAnyPattern(EXISTS_SUBQUERY, IN_SUBQUERY)) {
+ return expr
+ }
+ expr match {
+ case sq @ Exists(sub, _, _, conditions, subHint) if canRewrite(sq) =>
+ replacing(sq) {
+ val exists = AttributeReference("exists", BooleanType, nullable =
false)()
+ val existenceJoin = ExistenceJoin(exists)
+ val newCondition = conditions.reduceLeftOption(And)
+ newPlan =
+ buildJoin(newPlan, rewriteDomainJoinsIfPresent(newPlan, sub,
newCondition),
+ existenceJoin, newCondition, subHint)
+ introducedAttrs += exists
+ exists
+ }
+ case sq @ Not(InSubquery(values, ListQuery(sub, _, _, _, conditions,
subHint)))
+ if canRewrite(sq) =>
+ replacing(sq) {
+ val exists = AttributeReference("exists", BooleanType, nullable =
false)()
+ // Deduplicate conflicting attributes if any.
+ val newSub = dedupSubqueryOnSelfJoin(newPlan, sub, Some(values))
+ val inConditions = values.zip(newSub.output).map(EqualTo.tupled)
+ // To handle a null-aware predicate not-in-subquery in nested
conditions
+ // (e.g., `v > 0 OR t1.id NOT IN (SELECT id FROM t2)`), we
transform
+ // `inCondition` (t1.id=t2.id) into `(inCondition) OR
ISNULL(inCondition)`.
+ //
+ // For example, `SELECT * FROM t1 WHERE v > 0 OR t1.id NOT IN
(SELECT id FROM t2)`
+ // is transformed into a plan below;
+ // == Optimized Logical Plan ==
+ // Project [id#78, v#79]
+ // +- Filter ((v#79 > 0) OR NOT exists#83)
+ // +- Join ExistenceJoin(exists#83), ((id#78 = id#80) OR
isnull((id#78 = id#80)))
+ // :- Relation[id#78,v#79] parquet
+ // +- Relation[id#80] parquet
+ val nullAwareJoinConds = inConditions.map(c => Or(c, IsNull(c)))
+ val finalJoinCond = (nullAwareJoinConds ++
conditions).reduceLeft(And)
+ val joinHint = JoinHint(None, subHint)
+ newPlan = Join(newPlan,
+ rewriteDomainJoinsIfPresent(newPlan, newSub,
Some(finalJoinCond)),
+ ExistenceJoin(exists), Some(finalJoinCond), joinHint)
+ introducedAttrs += exists
+ Not(exists)
+ }
+ case sq @ InSubquery(values, ListQuery(sub, _, _, _, conditions,
subHint))
+ if canRewrite(sq) =>
+ replacing(sq) {
+ val exists = AttributeReference("exists", BooleanType, nullable =
false)()
+ // Deduplicate conflicting attributes if any.
+ val newSub = dedupSubqueryOnSelfJoin(newPlan, sub, Some(values))
+ val inConditions = values.zip(newSub.output).map(EqualTo.tupled)
+ val newConditions = (inConditions ++
conditions).reduceLeftOption(And)
+ val joinHint = JoinHint(None, subHint)
+ newPlan = Join(newPlan, rewriteDomainJoinsIfPresent(newPlan,
newSub, newConditions),
+ ExistenceJoin(exists), newConditions, joinHint)
+ introducedAttrs += exists
+ exists
+ }
+ // A sub-query that `canRewrite` declined is left as it is, children
included.
+ case sq @ (_: Exists | Not(_: InSubquery) | _: InSubquery) => sq
+ case other => other.mapChildren(rewrite)
}
}
+ val newExprs = exprs.map(rewrite)
(newExprs.reduceOption(And), newPlan, introducedAttrs.toSeq)
}
+
+ /**
+ * Returns true if `e` references any of the attributes produced by `plan`.
+ */
+ private def referencesPlan(e: Expression, plan: LogicalPlan): Boolean = {
+ // Iterates the references and probes the output set, rather than
materialising the
+ // intersection: AttributeSet.intersect builds a set from its argument, so
the operand order
+ // there would decide whether this walks the references or the whole
output of `plan`. The
+ // membership probes still construct wrappers, so this is cheaper, not
allocation-free.
+ e.references.exists(plan.outputSet.contains)
+ }
+
+ /**
+ * Returns true if `e` effectively references any of the attributes produced
by `plan`, see
+ * [[effectiveReferences]].
+ */
+ private def effectivelyReferencesPlan(e: Expression, plan: LogicalPlan):
Boolean = {
+ effectiveReferences(e).exists(plan.outputSet.contains)
+ }
+
+ /**
+ * Returns true if `attr` is produced by `outerPlan` and not by `subPlan`.
+ *
+ * An attribute that both plans produce is attributed to `subPlan` rather
than to `outerPlan`.
+ * This matters when the two share ExprIds, which a self-join can leave for
the optimizer to
+ * deduplicate (SPARK-21835): the EXISTS arms classify against the raw
sub-query plan, while the
+ * IN arms classify against the plan that `dedupSubqueryOnSelfJoin` has
already separated from
+ * the outer plan. Without this a single shared attribute would satisfy both
sides of the
+ * classification, so a nested sub-query that references only the sub-query
it is nested in
+ * would be taken for one referencing both plans and rejected. A nested
sub-query resolves in
+ * the scope of that sub-query, so a shared attribute is one of its own.
+ */
+ private def referencedByOuterPlanOnly(
+ attr: Attribute,
+ outerPlan: LogicalPlan,
+ subPlan: LogicalPlan): Boolean = {
+ outerPlan.outputSet.contains(attr) && !subPlan.outputSet.contains(attr)
+ }
+
+ /**
+ * The attributes that `e` references and can still evaluate. For an
existential sub-query
+ * expression these are the ones of the values it compares and of the
correlated condition that
+ * [[PullupCorrelatedPredicates]] hoisted into it, rather than the ones of
`references`, which
+ * also include the outer attributes the pull-up retains for idempotency.
Those outlive the
+ * condition that referenced them once BooleanSimplification has eliminated
it, as in
+ * `a IN (SELECT col1 FROM t3 WHERE false AND col1 = c1)`, which keeps `c1`
as an outer
+ * attribute of a sub-query that no longer reads it, and which must still be
rewritten against
+ * the outer plan rather than treated as referencing the sub-query plan.
+ */
+ private def effectiveReferences(e: Expression): AttributeSet = e match {
+ case Exists(_, _, _, joinCond, _) =>
AttributeSet(joinCond.flatMap(_.references))
+ case Not(in: InSubquery) => effectiveReferences(in)
+ case InSubquery(values, ListQuery(_, _, _, _, joinCond, _)) =>
+ AttributeSet(values.flatMap(_.references) ++
joinCond.flatMap(_.references))
+ case _ => e.references
+ }
+
+ /**
+ * Rewrites the existential sub-queries that are nested in the correlated
predicates which
+ * [[PullupCorrelatedPredicates]] hoisted out of a predicate sub-query, and
which therefore end
+ * up in the condition of the semi/anti join that replaces that sub-query.
+ *
+ * A hoisted predicate can carry a nested existential sub-query out of the
sub-query plan,
+ * because a correlated predicate is hoisted as a whole when it is a
disjunction. For example
+ *
+ * SELECT * FROM t1 WHERE EXISTS (
+ * SELECT 1 FROM t2 WHERE t1.a = t2.c1 OR t2.c1 IN (SELECT col1 FROM t3))
+ *
+ * hoists `a = c1 OR c1 IN (SELECT col1 FROM t3)` into the join condition.
Such a nested
+ * sub-query must be rewritten against the plan that produces the attributes
it references:
+ * the one above references `c1`, which is produced by the sub-query plan
and not by the outer
+ * plan, so its existence join has to be built on top of the sub-query plan.
Building it on top
+ * of the outer plan instead yields a join whose condition references an
attribute that neither
+ * of its children can produce (SPARK-59351).
+ *
+ * A nested sub-query that references both plans cannot be rewritten into an
existence join on
+ * either side, since neither alone can evaluate it, so it has to stay in
the join condition,
+ * where both plans are in scope. Leaving it there is only correct while it
is uncorrelated: the
+ * join condition of a correlated one is dropped when it is planned as an
in-subquery filter,
+ * which would silently change the result, so a correlated one is reported
as unsupported
+ * instead. Note that such a sub-query can be
+ * correlated only to the sub-query plan, as being correlated to the outer
plan as well would
+ * require two levels of correlation, which the Analyzer rejects.
+ *
+ * A sub-query left here is further subject to the rewrite of predicate
sub-queries in join
+ * conditions, which rejects one referencing both plans under the default
configuration. It
+ * survives that rewrite when
+ *
`spark.sql.optimizer.decorrelatePredicateSubqueriesInJoinPredicate.enabled` is
disabled, and
+ * also when it is an uncorrelated IN sub-query and
+ *
`spark.sql.optimizer.optimizeUncorrelatedInSubqueriesInJoinCondition.enabled`
is disabled,
+ * since it is then not among the sub-queries that rewrite considers.
+ *
+ * An existence join yields only whether a row matched, so its `exists`
attribute cannot tell
+ * FALSE from unknown, while `IN` is three-valued. The two are
indistinguishable while the value
+ * only feeds a predicate, which is what a hoisted condition normally does,
but not when it
+ * reaches something else, e.g. `(c1 IN (SELECT col1 FROM t3)) <=> false`,
which is FALSE for a
+ * NULL that matches nothing and TRUE for the `exists` attribute. An IN
sub-query whose row
+ * comparison can evaluate to unknown is therefore rejected in that position
rather than
+ * rewritten. NOT IN is rewritten with a null-aware join condition of its
own, which is equally
+ * two-valued, so it is rejected there too.
+ *
+ * Returns the rewritten condition along with the updated outer and
sub-query plans.
+ */
+ private def rewriteExistentialExprInJoinCondition(
+ conditions: Seq[Expression],
+ outerPlan: LogicalPlan,
+ subPlan: LogicalPlan): (Option[Expression], LogicalPlan, LogicalPlan) = {
+ val (subCond, newSubPlan) =
+ rewriteExistentialExprInSubqueryPlan(conditions, outerPlan, subPlan)
+ // The sub-queries that do not reference the sub-query plan are rewritten
against the outer
+ // plan, as they only reference attributes of the outer plan, if any. This
accepts on the
+ // negation of the test above rather than on a positive one, which rests
on the invariant that
+ // `effectiveReferences` sees every attribute the rewrite can put in a
join condition: one it
+ // cannot see would be rewritten here and end up in a condition the outer
plan cannot evaluate.
+ // A positive test is not available without also refusing the sub-queries
whose retained outer
+ // attributes no longer correspond to a condition, which must keep being
rewritten here.
+ val (newCond, newOuterPlan, _) = rewriteExistentialExprWithAttrs(
+ subCond.toSeq, outerPlan, e => !effectivelyReferencesPlan(e, subPlan))
+ (newCond, newOuterPlan, newSubPlan)
+ }
+
+ /**
+ * Rewrites the existential sub-queries in `conditions` that can only be
evaluated by the
+ * sub-query plan, that is those referencing the sub-query plan but not the
outer plan, into
+ * existence joins on top of the sub-query plan. See
+ * [[rewriteExistentialExprInJoinCondition]] for details.
+ *
+ * Returns the rewritten condition along with the updated sub-query plan.
+ */
+ private def rewriteExistentialExprInSubqueryPlan(
+ conditions: Seq[Expression],
+ outerPlan: LogicalPlan,
+ subPlan: LogicalPlan): (Option[Expression], LogicalPlan) = {
+ // Both collectors below match only Exists and InSubquery, and neither
prunes on tree
Review Comment:
Addressed on the current HEAD — apologies for the late reply.
Correct, the comment was wrong about one of the two collectors.
`collectUnknownSensitiveInSubqueries` does open with
```scala
case _ if !expr.containsAnyPattern(EXISTS_SUBQUERY, IN_SUBQUERY) =>
```
so it prunes every subquery-free subtree. It is only the rejection collector
that does not: that one is a plain `TreeNode.collect`, which visits every node.
The comment now says exactly that, and that the guard buys the pruning
collector only its top-level call
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
