pratham76 commented on code in PR #58656:
URL: https://github.com/apache/spark/pull/58656#discussion_r4274278928
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/RewriteSubquerySuite.scala:
##########
@@ -161,4 +163,69 @@ class RewriteSubquerySuite extends PlanTest {
val existsExpected = LocalRelation(existsQuery.analyze.output).analyze
comparePlans(existsOptimized, existsExpected)
}
+
+ test("SPARK-59351: nested subquery of an EXISTS whose plan shares attributes
with the outer " +
+ "plan") {
+ // The EXISTS arms route the nested subquery against the raw subquery
plan, while the IN arms
+ // route it against the plan dedupSubqueryOnSelfJoin has already separated
from the outer
+ // plan. A self-join can leave an ExprId shared between the two plans for
the optimizer
+ // (SPARK-21835), and the analyzer is not involved here, so build that
state directly: `a` is
+ // the same attribute in the outer and in the subquery plan.
+ val a = $"a".int
+ val outer = LocalRelation(a, $"b".int)
+ val sub = LocalRelation(a)
+ val nested = LocalRelation($"col1".int)
+ assert(outer.outputSet.intersect(sub.outputSet).nonEmpty,
+ "the two plans must share an attribute for this test to mean anything")
+
+ // The nested subquery references `a`, which only the subquery plan is
meant to give it, and
+ // is correlated to it, so it must be rewritten against the subquery plan.
Classifying `a` as
+ // a reference to the outer plan as well would report it as referencing
both plans.
+ val nestedInSubquery = InSubquery(
+ Seq(a),
+ ListQuery(nested.select($"col1"), joinCond = Seq(EqualTo($"col1", a))))
+ val query = Filter(
+ Exists(sub, joinCond = Seq(Or(EqualTo($"b", a), nestedInSubquery))),
+ outer)
+
+ // The shared attribute is attributed to the subquery plan, so the nested
subquery is routed
+ // there rather than reported as referencing both plans. What then reports
this state is the
+ // conflict check that dedupSubqueryOnSelfJoin has always run over the
semi join condition,
+ // which names the duplicated attribute: an attribute shared between the
two plans that the
+ // hoisted condition mentions cannot be turned into a semi join whatever
the routing does.
+ val e = intercept[AnalysisException](Optimize.execute(query))
+ assert(e.getCondition == "_LEGACY_ERROR_TEMP_1212",
+ s"expected the pre-existing conflict check to report this, got
${e.getCondition}")
+ }
+
+ test("SPARK-59351: the rewrite keeps the origin and the tags of the
sub-query it replaces") {
+ // The rewrite of a nested existential sub-query builds its replacement by
hand rather than
+ // through TreeNode.transformDownWithPruning, so it has to reproduce what
that does: apply the
+ // rule under the origin of the node being replaced and copy its tags.
Without this a runtime
+ // error raised from the rewritten condition quotes the enclosing operator
instead of the
+ // sub-query fragment, which no result or plan assertion would notice.
+ val outer = LocalRelation($"a".int)
+ val inner = LocalRelation($"x".int)
+ val tag = TreeNodeTag[String]("SPARK-59351")
+ val origin = Origin(line = Some(42), startPosition = Some(7))
+ val a = outer.output.head
+ val x = inner.output.head
+ // Correlated, so that the rule's unary node handler picks the projection
up at all.
+ val exists = CurrentOrigin.withOrigin(origin) {
+ Exists(inner, outerAttrs = Seq(a), joinCond = Seq(EqualTo(a, x)))
+ }
+ exists.setTagValue(tag, "kept")
+
+ // An Exists outside a Filter goes through the unary node handler, which
is the path that
+ // builds the replacement by hand.
+ val rewritten = RewritePredicateSubquery(Project(Seq(Alias(exists,
"e")()), outer))
+ val replacements = rewritten.expressions.flatMap(_.collect {
+ case attr: AttributeReference if attr.name == "exists" => attr
+ })
+ assert(replacements.size == 1, s"expected one exists attribute
in\n$rewritten")
+ assert(replacements.head.origin == origin,
Review Comment:
Confirmed, and your prediction is exact. I added a second test covering the
positive IN and NOT IN arms, then broke **only** those two arms by giving their
replacements the ambient origin instead of the subquery's:
| | existing Exists test | new IN test |
|---|---|---|
| both arms correct | pass | pass |
| only IN arms broken | **pass** | **fail** |
So the old test was blind to it exactly as you said.
The new test asserts on the generated condition nodes rather than the
projected attribute, since that is where they end up:
* **positive IN, uncorrelated** — the condition is the synthesized
`EqualTo`, whose origin must be the subquery's;
* **NOT IN, correlated** — the condition is `And(Or(cmp, IsNull(cmp)),
hoisted)`, and the combined condition, the null-aware wrapper, the comparison
and the null check must all carry the subquery's origin.
It also pins the other direction, which nothing asserted before: the hoisted
condition is built under a *different* origin and must keep it, since the
rewrite reuses that node rather than building it, and clobbering it would
misreport a user-written predicate.
--
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]
