pratham76 commented on code in PR #58656:
URL: https://github.com/apache/spark/pull/58656#discussion_r4274594390
##########
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))
Review Comment:
Addressed on the current HEAD — apologies for the late reply. This also
closes `r4139416496` from the round before, which I had marked addressed when
it was not.
Both halves of this check out, and both are fixed.
**The flattened sequence.** `AttributeSet(values.flatMap(_.references) ++
joinCond.flatMap(_.references))` is worse than it looks: the `flatMap` builds a
`Seq[Attribute]` as long as the total number of references, and then
`AttributeSet.apply(Iterable[Expression])` is
`fromAttributeSets(baseSet.map(_.references))`, so each of those attributes
gets its own single-element `AttributeSet` allocated purely to be folded away.
`effectiveReferences` now merges the child sets directly:
```scala
AttributeSet.fromAttributeSets(values.map(_.references) ++
joinCond.map(_.references))
```
**The repeated reconstruction and the separate scans.** The same candidate
was classified up to three times once by the both-plan rejection, once by the
UNKNOWN filter, once by the rewrite's `canRewrite` each time rebuilding the set
and scanning it twice, once for `subPlan` and once for
`referencedByOuterPlanOnly`. `referencedByOuterPlanOnly` is replaced by
`planReferences`, which answers both questions in a single pass with an early
exit once both plans are known to be referenced, and the result is memoised per
candidate:
```scala
val classified = new java.util.IdentityHashMap[Expression, PlanReferences]()
```
Identity keying is deliberate: the candidates are nodes of the expressions
being classified, so structural comparison would compare the sub-query plans
they hold.
Per candidate this goes from 3 set constructions and up to 6 scans to 1 and
1.
**On the collision ordering**, which is the compatibility half of your
request. The single-pass form probes `subPlan` first, preserving the "a shared
attribute belongs to the sub-query" tie-break that `referencedByOuterPlanOnly`
had. I first wrote in the scaladoc that no shape in the suite reaches that
branch that was wrong, and you caught it in `r4257086406`: the shared-attribute
fixture in `RewriteSubquerySuite` does reach it. The comment has since been
corrected to distinguish analyzer-produced SQL plans, which do not reach it
because `DeduplicateRelations` has already separated the ExprIds, from that
optimizer-level fixture, which builds the overlap directly. It pins the shape
but not the outcome of the ordering, since `dedupSubqueryOnSelfJoin`'s
pre-existing conflict check reports the duplicated attribute first either way.
The other shapes you named — inner-only, outer-only, both-plan, unreferenced
and simplified-away-correlation are covered by the existing cases in the two
`SPARK-59351` tests, and all pass.
--
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]
