LuciferYang commented on code in PR #58424: URL: https://github.com/apache/spark/pull/58424#discussion_r3939100190
########## sql/core/src/main/scala/org/apache/spark/sql/execution/RewriteSelfJoinInequalityToAggregate.scala: ########## @@ -0,0 +1,746 @@ +/* + * 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.execution + +import org.apache.spark.sql.catalyst.expressions._ +import org.apache.spark.sql.catalyst.expressions.aggregate._ +import org.apache.spark.sql.catalyst.plans._ +import org.apache.spark.sql.catalyst.plans.logical._ +import org.apache.spark.sql.catalyst.rules.Rule +import org.apache.spark.sql.catalyst.trees.TreePattern.IN_SUBQUERY +import org.apache.spark.sql.catalyst.util.CharVarcharUtils +import org.apache.spark.sql.execution.datasources.{HadoopFsRelation, LogicalRelation} +import org.apache.spark.sql.execution.datasources.parquet.ParquetFileFormat +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types._ + +/** + * Rewrites a self-join with an inequality into GROUP BY + HAVING COUNT(DISTINCT) > 1. + * + * Targets the two uncorrelated InSubquery shapes exercised by TPC-DS Q95: + * + * - Pattern A': the subquery top-level InnerJoin is a direct self-join. + * - Pattern A2: the subquery contains an outer InnerJoin with a self-join child; only the + * self-join child is replaced with Aggregate and the outer join is preserved. + * + * Both patterns require an existence-only membership context so row-count multiplicity from the + * original self-join cross-product does not affect semantics. Correlated InSubquery expressions are + * intentionally fail-closed because the ExprId remapping performed here does not rewrite correlated + * predicates. + * + * This rule runs in `extendedOperatorOptimizationRules`, which is part of the operator optimization + * batch and therefore executes before `RewritePredicateSubquery` turns the predicate subquery into + * a semi/anti/existence join. It only observes the uncorrelated `InSubquery` shape at that phase, + * so there is no separate LeftSemi/LeftAnti ("Pattern A") or `Exists` handling. + * + * Both patterns share: + * - [[buildAggregateHavingDistinctGt1]] to construct `Filter(cnt > 1, Aggregate)` + * - [[canonicalizeWrapper]] to rebuild a wrapping Project so every equi-key reference points to + * the sjLeft-side attribute, with **fresh exprIds** (Spark's SPARK-21835 style -- no reuse of + * original exprIds), returning an old->new attribute remap for downstream rewrite. + * + * Controlled by `spark.sql.optimizer.rewriteSelfJoinInequalityToAggregate.enabled` + * (default false, opt-in). + */ +object RewriteSelfJoinInequalityToAggregate extends Rule[LogicalPlan] with PredicateHelper { Review Comment: Compared to the precedent of RewriteNonCorrelatedExists (11 lines, always-on), this rule is three hand-curated whitelists plus a private config covering a single TPC-DS shape, and each whitelist silently narrows as Spark adds new expressions and types. I agree with opt-in (see my comment on the cost gate), but the PR description should state the promotion or removal criterion for this flag: what measured win on real workloads graduates it, and what has to be solved before it can default on. Otherwise it stays experimental, and no one is willing to touch or delete it. ########## sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala: ########## @@ -669,6 +669,18 @@ object SQLConf { "for using switch statements in InSet must be non-negative and less than or equal to 600") .createWithDefault(400) + val REWRITE_SELF_JOIN_INEQUALITY_TO_AGGREGATE_ENABLED = + buildConf("spark.sql.optimizer.rewriteSelfJoinInequalityToAggregate.enabled") + .internal() + .doc("When true, rewrites a supported existence-only inequality self-join inside an " + Review Comment: The conf doc's "proven to be the same repeatable relation" reads like a general capability, but isRowBagRepeatable only admits stock ParquetFileFormat (exact getClass match), Range and LocalRelation. ORC/CSV/JSON, DataSourceV2 and cached tables are all silently excluded: users enable the flag, EXPLAIN does not change, and nothing tells them why. The boundary comes from validation scope and deserves to be stated. Name the supported sources in the conf doc and the PR description, and consider a DEBUG log saying why the whitelist rejected the plan when the shape matches (that state only becomes observable once the shape match is moved ahead of the audit, per my other comment); the rule runs in fixedPoint batches, so log once per subquery rather than on every rejection. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/RewriteSelfJoinInequalityToAggregate.scala: ########## @@ -0,0 +1,746 @@ +/* + * 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.execution + +import org.apache.spark.sql.catalyst.expressions._ +import org.apache.spark.sql.catalyst.expressions.aggregate._ +import org.apache.spark.sql.catalyst.plans._ +import org.apache.spark.sql.catalyst.plans.logical._ +import org.apache.spark.sql.catalyst.rules.Rule +import org.apache.spark.sql.catalyst.trees.TreePattern.IN_SUBQUERY +import org.apache.spark.sql.catalyst.util.CharVarcharUtils +import org.apache.spark.sql.execution.datasources.{HadoopFsRelation, LogicalRelation} +import org.apache.spark.sql.execution.datasources.parquet.ParquetFileFormat +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types._ + +/** + * Rewrites a self-join with an inequality into GROUP BY + HAVING COUNT(DISTINCT) > 1. + * + * Targets the two uncorrelated InSubquery shapes exercised by TPC-DS Q95: + * + * - Pattern A': the subquery top-level InnerJoin is a direct self-join. + * - Pattern A2: the subquery contains an outer InnerJoin with a self-join child; only the + * self-join child is replaced with Aggregate and the outer join is preserved. + * + * Both patterns require an existence-only membership context so row-count multiplicity from the + * original self-join cross-product does not affect semantics. Correlated InSubquery expressions are + * intentionally fail-closed because the ExprId remapping performed here does not rewrite correlated + * predicates. + * + * This rule runs in `extendedOperatorOptimizationRules`, which is part of the operator optimization + * batch and therefore executes before `RewritePredicateSubquery` turns the predicate subquery into + * a semi/anti/existence join. It only observes the uncorrelated `InSubquery` shape at that phase, + * so there is no separate LeftSemi/LeftAnti ("Pattern A") or `Exists` handling. + * + * Both patterns share: + * - [[buildAggregateHavingDistinctGt1]] to construct `Filter(cnt > 1, Aggregate)` + * - [[canonicalizeWrapper]] to rebuild a wrapping Project so every equi-key reference points to + * the sjLeft-side attribute, with **fresh exprIds** (Spark's SPARK-21835 style -- no reuse of + * original exprIds), returning an old->new attribute remap for downstream rewrite. + * + * Controlled by `spark.sql.optimizer.rewriteSelfJoinInequalityToAggregate.enabled` + * (default false, opt-in). + */ +object RewriteSelfJoinInequalityToAggregate extends Rule[LogicalPlan] with PredicateHelper { + + private val CountDistinctAliasName = "_rewrite_selfjoin_inequality_cnt_distinct" + + override def apply(plan: LogicalPlan): LogicalPlan = { + if (!conf.getConf(SQLConf.REWRITE_SELF_JOIN_INEQUALITY_TO_AGGREGATE_ENABLED)) { + return plan + } + + // Pattern A' / A2: rewrite uncorrelated InSubquery plans. + // Correlated subqueries carry outer references / correlated join conditions in + // `SubqueryExpression.children`; fail closed because this rule does not remap them. + val rewritten = plan.transformAllExpressionsWithPruning(_.containsPattern(IN_SUBQUERY)) { + case in @ InSubquery(_, lq: ListQuery) if lq.children.isEmpty => + rewriteSubqueryPlan(lq.plan) match { + case Some(newSub) => in.copy(query = lq.copy(plan = newSub)) + case None => in + } + } + if (!(rewritten eq plan)) { + logDebug( + "RewriteSelfJoinInequalityToAggregate: rewrote self-join to " + + "GROUP BY + HAVING COUNT(DISTINCT) > 1") + } + rewritten + } + // ============================================================================ + // Shared helpers + // ============================================================================ + + /** + * Build `Filter(cnt > 1, Aggregate(equiKeys, [equiKeys, cnt_alias], Filter(IsNotNull(equiKeys), + * child)))`. Returns the Filter node whose output is `equiKeys ++ [count_alias_attr]`. + * + * The extra `IsNotNull(equiKeys)` filter is essential to preserve the original equi-join's NULL + * semantics. Under SQL 3VL, `left.k = right.k` never matches when either side is NULL, so the + * original self-join drops rows with NULL equi-keys. Aggregate, in contrast, groups NULL keys + * together into a single "NULL group" -- if that group has >= 2 distinct non-null neq values, + * COUNT(DISTINCT) > 1 fires and injects NULL into the subquery output. That leaked NULL then + * turns `NOT IN` into a spurious empty result (Spark's null-aware anti-join uses + * `Or(equi, IsNull(equi))` which any NULL sub-row satisfies) and can flip IN/NOT IN outcomes. The + * neq column needs no such filter: `COUNT(DISTINCT col)` already ignores NULL. + */ + private def buildAggregateHavingDistinctGt1( + equiKeys: Seq[Attribute], + neqCol: Attribute, + child: LogicalPlan): LogicalPlan = { + val countExpr = AggregateExpression( + Count(Seq(neqCol)), + mode = Complete, + isDistinct = true, + filter = None, + NamedExpression.newExprId) + val countAlias = Alias(countExpr, CountDistinctAliasName)() + // Seq[Attribute] is a Seq[NamedExpression] via covariance; no cast needed. + val aggExprs: Seq[NamedExpression] = equiKeys :+ countAlias + val nonNullChild = equiKeys + .map(a => IsNotNull(a): Expression) + .reduceOption(And) + .map(Filter(_, child)) + .getOrElse(child) + val agg = Aggregate(equiKeys, aggExprs, nonNullChild) + Filter(GreaterThan(countAlias.toAttribute, Literal(1L, LongType)), agg) + } + + /** + * Canonicalize a Project so every equi-key reference points at the sjLeft-side attribute. + * [[parseSelfJoinCondition]] has already verified that each pair refers to the same output + * position on the two structurally identical self-join sides. Uses **fresh exprIds** (no reuse of + * original wrapper output exprIds) -- the same technique Spark's own `dedupSubqueryOnSelfJoin` + * uses when it needs to change subquery output. + * + * Returns the rebuilt Project and a map `oldWrapperOutputExprId -> newWrapperOutputAttr`, so + * downstream references (outer join condition, top-level Project) can be updated consistently. + * + * `equiPairs` provides the definitive ExprId-based lookup: `equiPair (l, r)` binds + * `l.exprId -> l` (identity) and `r.exprId -> l` (sjRight -> sjLeft). Attribute identity in + * Catalyst is ExprId, not name; two columns can share a name with distinct ExprIds. Name-based + * lookup would silently drop such entries via `.toMap`. + * + * Fails (returns None) when a projectList entry is neither an equi-key Attribute (by ExprId) nor + * `Alias(equi-key Attribute, _)`. Fail-closed. + */ + private def canonicalizeWrapper( + projectList: Seq[NamedExpression], + equiPairs: Seq[(Attribute, Attribute)], + newChild: LogicalPlan): Option[(Project, Map[ExprId, Attribute])] = { + // ExprId-based canonical map: any equi-key attribute (either side) -> sjLeft attribute. + val exprIdToLeft: Map[ExprId, Attribute] = + equiPairs.flatMap { case (l, r) => Seq(l.exprId -> l, r.exprId -> l) }.toMap + val oldOutput: Seq[Attribute] = projectList.map(_.toAttribute) + val mapped: Seq[Option[NamedExpression]] = projectList.map { + case a: Attribute if exprIdToLeft.contains(a.exprId) => + // Wrap every rewritten output slot in a fresh Alias. + // + // When a wrapper reprojects BOTH sides of the same equi pair (e.g. + // `SELECT s1.k, s2.k FROM T s1 JOIN T s2 ON s1.k = s2.k AND s1.v <> s2.v`), + // both entries collapse to the same sjLeft Attribute after the self-join is + // rewritten. Duplicate output ExprIds are not illegal in Spark (`SELECT a, a` + // is a valid Project), but fresh Aliases give each output slot an independent + // identity, which keeps the `oldOutput -> newOutput` remap 1-to-1 and lets + // downstream references (outer join condition, top-level Project) be updated + // unambiguously via ExprId. + // + // The fresh ExprId is on the Alias ITSELF; the referenced child keeps its + // original ExprId. Spark's logical-plan integrity checks reject reusing a + // referenced ExprId as the Alias's own ExprId, not duplication across slots. + Some(Alias(exprIdToLeft(a.exprId), a.name)(): NamedExpression) + case al @ Alias(a: Attribute, _) if exprIdToLeft.contains(a.exprId) => + // Fresh exprId; do NOT reuse `al.exprId`. Reusing another expression's exprId + // is the pattern that Spark 3.3 flags via structural-integrity checks. + Some(Alias(exprIdToLeft(a.exprId), al.name)(): NamedExpression) + case _ => None + } + if (mapped.exists(_.isEmpty)) { + None + } else { + val newProjectList = mapped.flatten + val newWrapper = Project(newProjectList, newChild) + val newOutput = newWrapper.output + val remap: Map[ExprId, Attribute] = + oldOutput.zip(newOutput).map { case (o, n) => o.exprId -> n }.toMap + Some((newWrapper, remap)) + } + } + + /** + * Replace equi-key attribute references inside a NamedExpression according to `remap`, while + * preserving the NamedExpression shape. + * + * `Expression.transformUp` returns `Expression`, not `NamedExpression`. We avoid a blanket + * `asInstanceOf[NamedExpression]` by handling the two shapes that can appear in a Project's + * `projectList` explicitly: a bare Attribute (whose top-level may itself be replaced) and an + * Alias (which stays an Alias while its child is transformed). Any other NamedExpression shape we + * do not rewrite is left as-is ONLY if it does not reference a replaced self-join output; + * otherwise it would carry a stale ExprId, so returns None to fail the whole rewrite closed. + */ + private def remapNamedExpressionAttributes( + ne: NamedExpression, + remap: Map[ExprId, Attribute]): Option[NamedExpression] = ne match { + case a: Attribute if remap.contains(a.exprId) => Some(remap(a.exprId)) + case a: Attribute => Some(a) + case al: Alias => + val newChild = al.child.transformUp { + case a: Attribute if remap.contains(a.exprId) => remap(a.exprId) + } + Some( + if (newChild eq al.child) al + else Alias(newChild, al.name)(al.exprId, al.qualifier, al.explicitMetadata)) + case other if other.references.exists(a => remap.contains(a.exprId)) => + // Fail-closed: a NamedExpression we do not rewrite (neither a bare Attribute nor an Alias) + // that still references a replaced self-join output would be left with a dangling ExprId. + // Refuse the rewrite rather than emit a plan with a stale reference. + None + case other => Some(other) + } + + // ============================================================================ + // Pattern A' / A2 dispatch (subquery plans of InSubquery) + // ============================================================================ + + private def rewriteSubqueryPlan(plan: LogicalPlan): Option[LogicalPlan] = { Review Comment: The rewrite has no cost gate: it fires on plan shape alone. A small dimension table's self-join currently plans as a broadcast hash join with no shuffle exchange; rewriting it to COUNT(DISTINCT) adds a hashpartitioning(equiKeys) exchange plus a two-phase aggregate per qualifying subquery, which the Q95-style high-multiplicity win does not offset. With the flag enabled globally, queries over small tables regress across the board. Once the shape and self-join guards pass, skip the rewrite when sjLeft.stats.sizeInBytes is at or below autoBroadcastJoinThreshold. Without ANALYZE that number is the raw file size, the same one JoinSelection's static broadcast check uses, and both sides are the same relation so one check suffices. ########## sql/core/src/test/scala/org/apache/spark/sql/execution/RewriteSelfJoinInequalityToAggregateSuite.scala: ########## @@ -0,0 +1,1037 @@ +/* + * 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.execution + +import org.apache.spark.SparkThrowable +import org.apache.spark.sql.{QueryTest, Row} +import org.apache.spark.sql.catalyst.expressions.Alias +import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.test.SharedSparkSession +import org.apache.spark.sql.types.{IntegerType, StructField, StructType} + +/** + * Correctness tests for [[RewriteSelfJoinInequalityToAggregate]]. + * + * Positive A' / A2 cases assert both result equivalence and that the rewrite actually fired. + * + * `assert(!ruleFired(plan))` on its own only proves the rewrite did not happen -- not that it was + * the guard under test that stopped it. A fixture whose two self-join sides are not structurally + * identical is rejected by `isSameBaseRelation` before any predicate is even parsed, and such a + * test passes while covering nothing. So six important rejection paths -- the predicate parser, the + * single-inequality requirement, output-position identity, the nondeterminism guard, the + * leaf-source allowlist (LogicalRDD vs Parquet), and the expression-type allowlist (`abs(v)` vs + * `v + 1`) -- are tested as single-variable pairs: the same fixture and the same query shape, one + * control query that must fire and one variant that changes only the feature under test and must + * not. A firing control does not pin the rejection to a particular line, but it does rule out an + * unrelated fixture mismatch as the reason its partner was rejected. The row-bag whitelist + * (Aggregate, Window) stays a plain negative: dropping the operator would change the query shape + * rather than one feature. + * + * Self-joined fixtures are real tables, not temp views over VALUES. Spark deduplicates a self-join + * over a [[org.apache.spark.sql.catalyst.analysis.MultiInstanceRelation]] via `newInstance()`, + * which refreshes one side's ExprIds without inserting a rename-only Project, so both sides stay + * structurally identical. A temp view over VALUES cannot, and Spark renames one side with a Project + * instead, which would make `isSameBaseRelation` false for every self-join below. `range()` needs + * no such treatment -- Range is a MultiInstanceRelation already. + */ +class RewriteSelfJoinInequalityToAggregateSuite extends QueryTest with SharedSparkSession { + + private val rewriteConf = SQLConf.REWRITE_SELF_JOIN_INEQUALITY_TO_AGGREGATE_ENABLED.key + + /** Signature alias produced by the rewrite; presence => rule definitely fired. */ + private val CountDistinctAlias = "_rewrite_selfjoin_inequality_cnt_distinct" + + private def ruleFired(plan: LogicalPlan): Boolean = + plan.exists { + p => + p.expressions.exists(_.exists { + case a: Alias if a.name == CountDistinctAlias => true + case _ => false + }) + } + + private def assertRuleFired(sql: String): Unit = { + withSQLConf(rewriteConf -> "true") { + val plan = spark.sql(sql).queryExecution.optimizedPlan + assert(ruleFired(plan), s"self-join inequality rewrite should fire:\n$plan") + } + } + + private def assertRuleNotFired(sql: String): Unit = { + withSQLConf(rewriteConf -> "true") { + val plan = spark.sql(sql).queryExecution.optimizedPlan + assert(!ruleFired(plan), s"self-join inequality rewrite must not fire:\n$plan") + } + } + + /** + * A real table, so that a self-join of it dedups into two structurally identical sides. See the + * class comment for why a temp view over VALUES cannot be used for a self-joined fixture. + */ + private def createTable(name: String, schema: String, values: String): Unit = { + spark.sql(s"DROP TABLE IF EXISTS $name") + spark.sql(s"CREATE TABLE $name($schema) USING parquet") + spark.sql(s"INSERT INTO $name SELECT * FROM VALUES $values") + } + + /** Run `sql` twice, first with rewrite ON then OFF, and return the two result row sets. */ + private def runBoth(sql: String): (Set[Row], Set[Row]) = { + var on: Set[Row] = null + var off: Set[Row] = null + withSQLConf(rewriteConf -> "true") { + on = spark.sql(sql).collect().toSet + } + withSQLConf(rewriteConf -> "false") { + off = spark.sql(sql).collect().toSet + } + (on, off) + } + + private def setupTable(): Unit = { + // k=1: distinct v={10,20} -> matches (has 2 non-null distinct) + // k=2: distinct v={30} -> no match (only 1) + // k=3: distinct v={40,50,60} -> matches + // k=4: v={70, NULL} -> no match (only 1 non-null) + // k=5: v={NULL, NULL} -> no match (0 non-null) + // k=6: v={80, 90, NULL} -> matches + // k=7: v={100,100} -> no match: duplicate-only. Proves DISTINCT is required; + // a plain COUNT(v) > 1 would wrongly match this group. + createTable( + "T", + "k INT, v INT", + """ (1, 10), (1, 10), (1, 20), + | (2, 30), + | (3, 40), (3, 50), (3, 60), + | (4, 70), (4, CAST(NULL AS INT)), + | (5, CAST(NULL AS INT)), (5, CAST(NULL AS INT)), + | (6, 80), (6, 90), (6, CAST(NULL AS INT)), + | (7, 100), (7, 100)""".stripMargin + ) + } + + // ==================== Positive: rewrite fires and is semantically equivalent =============== + + test("Pattern A': direct InSubquery self-join is rewritten") { + setupTable() + val sql = + """SELECT k FROM T outer_t WHERE k IN ( + | SELECT s1.k FROM T s1 JOIN T s2 + | ON s1.k = s2.k AND s1.v <> s2.v)""".stripMargin + + assertRuleFired(sql) + val (on, off) = runBoth(sql) + assert(on == off, s"rewrite ON $on != OFF $off") + assert(on == Set(Row(1), Row(3), Row(6)), s"expected {1,3,6}, got $on") + } + + test("Pattern A2: nested self-join is rewritten") { + setupTable() + spark.sql( + """CREATE OR REPLACE TEMP VIEW D AS SELECT * FROM VALUES + | (1), (3), (6) AS D(k)""".stripMargin) + val sql = + """SELECT k FROM T outer_t WHERE k IN ( + | SELECT d.k + | FROM D d, (SELECT s1.k FROM T s1 JOIN T s2 + | ON s1.k = s2.k AND s1.v <> s2.v) sj + | WHERE d.k = sj.k)""".stripMargin + + assertRuleFired(sql) + val (on, off) = runBoth(sql) + assert(on == off, s"Pattern A2 rewrite ON $on != OFF $off") + assert(on == Set(Row(1), Row(3), Row(6))) + } + + test("Pattern A2: self-join on the LEFT of the outer join is rewritten") { + // Mirror of the Pattern A2 test above. There the self-join is the RIGHT child of the outer join + // (`selfJoinOnRight = true`); here it is the LEFT child (`selfJoinOnRight = false`). The rule + // has an explicit branch for each side, so both are covered. + setupTable() + spark.sql( + """CREATE OR REPLACE TEMP VIEW D AS SELECT * FROM VALUES + | (1), (3), (6) AS D(k)""".stripMargin) + val sql = + """SELECT k FROM T outer_t WHERE k IN ( + | SELECT d.k + | FROM (SELECT s1.k FROM T s1 JOIN T s2 + | ON s1.k = s2.k AND s1.v <> s2.v) sj, D d + | WHERE sj.k = d.k)""".stripMargin + + assertRuleFired(sql) + val (on, off) = runBoth(sql) + assert(on == off, s"Pattern A2 (self-join on left) rewrite ON $on != OFF $off") + assert(on == Set(Row(1), Row(3), Row(6))) + } + + test("Pattern A2: nondeterminism in the outer join condition must not be rewritten") { + // Both self-join sides are still repeatable here, so the per-side `isSameBaseRelation` check + // would pass; the `rand()` conjunct lives on the outer join ABOVE the self-join. Only the + // candidate-level `isRepeatablePlan` walk over the whole subquery catches it, so the rule must + // fail closed. This is the case the candidate-level guard exists for. + setupTable() + spark.sql( + """CREATE OR REPLACE TEMP VIEW D AS SELECT * FROM VALUES + | (1), (3), (6) AS D(k)""".stripMargin) + val sql = + """SELECT k FROM T outer_t WHERE k IN ( + | SELECT d.k + | FROM D d, (SELECT s1.k FROM T s1 JOIN T s2 + | ON s1.k = s2.k AND s1.v <> s2.v) sj + | WHERE d.k = sj.k AND rand() < 0.5)""".stripMargin + + assertRuleNotFired(sql) + } + + test("Pattern A': multi-equi tuple IN with sjRight key remap is rewritten") { + // Exercises the multi-equi-key path: two equi keys (k1, k2) drive the GROUP BY, and the tuple + // IN projects `s1.k1, s2.k2` -- so the second output column comes from the RIGHT self-join side + // and must be remapped to its sjLeft counterpart by `canonicalizeWrapper`. This one case covers + // multiple equi keys, tuple IN output arity, the two injected IsNotNull(equiKey) filters, and + // the sjRight-attribute remap at once. + createTable( + "TM", + "k1 INT, k2 INT, v INT", + """ (1, 1, 10), (1, 1, 20), + | (1, 2, 30), (1, 2, 30), + | (2, 1, 40), (2, 1, 50), + | (CAST(NULL AS INT), 1, 60), (CAST(NULL AS INT), 1, 70), + | (3, CAST(NULL AS INT), 80), (3, CAST(NULL AS INT), 90)""".stripMargin) + val sql = + """SELECT k1, k2 FROM TM outer_t WHERE (k1, k2) IN ( + | SELECT s1.k1, s2.k2 FROM TM s1 JOIN TM s2 + | ON s1.k1 = s2.k1 AND s1.k2 = s2.k2 AND s1.v <> s2.v)""".stripMargin + + assertRuleFired(sql) + val (on, off) = runBoth(sql) + assert(on == off, s"multi-equi tuple IN rewrite ON $on != OFF $off") + // (1,1): distinct v={10,20} -> matches; (1,2): v={30} -> no; (2,1): v={40,50} -> matches; + // (NULL,1) and (3,NULL): NULL equi key filtered out by the injected IsNotNull. -> {(1,1),(2,1)} + assert(on == Set(Row(1, 1), Row(2, 1)), s"expected {(1,1),(2,1)}, got $on") + } + + test("NULL / 3VL on inequality column is preserved") { + setupTable() + val sql = + """SELECT k FROM T outer_t WHERE k IN ( + | SELECT s1.k FROM T s1 JOIN T s2 + | ON s1.k = s2.k AND s1.v <> s2.v)""".stripMargin + + assertRuleFired(sql) + val (on, off) = runBoth(sql) + // k=4 (v={70,NULL}) and k=5 (v={NULL,NULL}) do not satisfy plain SQL <>. + assert(on == Set(Row(1), Row(3), Row(6)), s"expected {1,3,6}, got $on") + assert(off == on, s"NULL/3VL semantics diverge between rewrite ON and OFF: $on vs $off") + } + + test("NULL equi-key is filtered before aggregation for NOT IN") { + createTable( + "TN", + "k INT, v INT", + """ (CAST(NULL AS INT), 10), + | (CAST(NULL AS INT), 20), + | (1, 10), (1, 20), + | (2, 30)""".stripMargin + ) + spark.sql( + """CREATE OR REPLACE TEMP VIEW OuterKeys AS SELECT * FROM VALUES + | (1), (2), (3) AS OuterKeys(k)""".stripMargin) + val sql = + """SELECT k FROM OuterKeys o WHERE k NOT IN ( + | SELECT s1.k FROM TN s1 JOIN TN s2 + | ON s1.k = s2.k AND s1.v <> s2.v)""".stripMargin + + assertRuleFired(sql) + val (on, off) = runBoth(sql) + assert(on == off, s"NULL equi-key NOT IN semantics diverge: ON=$on OFF=$off") + assert(on == Set(Row(2), Row(3)), s"expected {2,3}, got $on") + } + + test("Swapped aliases must not be treated as the same self-join columns") { + // The guard under test is `sameOutputPosition`. Both queries alias the same two base columns + // to the names `k` and `v` on both sides, so a rule that compares attribute names would fire + // on both; only the output ordinal tells them apart. The control fires, which is what makes + // the negative case evidence that the ordinal check -- not a structural mismatch -- rejected + // the swapped one. + createTable("AliasBase", "a INT, b INT", " (1, 10), (1, 20), (2, 30)") + + val alignedSql = + """SELECT a FROM AliasBase outer_t WHERE a IN ( + | SELECT s1.k + | FROM (SELECT a AS k, b AS v FROM AliasBase) s1 + | JOIN (SELECT a AS k, b AS v FROM AliasBase) s2 + | ON s1.k = s2.k AND s1.v <> s2.v)""".stripMargin + assertRuleFired(alignedSql) + val (alignedOn, alignedOff) = runBoth(alignedSql) + assert( + alignedOn == alignedOff, + s"aligned-alias control diverges: ON=$alignedOn OFF=$alignedOff") + assert(alignedOn == Set(Row(1)), s"aligned-alias control expected {1}, got $alignedOn") + + // s1.k is `a` (output position 0) but s2.k is `b` (output position 1): same name, different + // column. Rewriting this would count distinct `b` per `a`, which is a different query. + val swappedSql = + """SELECT a FROM AliasBase outer_t WHERE a IN ( + | SELECT s1.k + | FROM (SELECT a AS k, b AS v FROM AliasBase) s1 + | JOIN (SELECT a AS v, b AS k FROM AliasBase) s2 + | ON s1.k = s2.k AND s1.v <> s2.v)""".stripMargin + assertRuleNotFired(swappedSql) + val (on, off) = runBoth(swappedSql) + assert(on == off, s"swapped-alias semantics diverge: ON=$on OFF=$off") + assert(on.isEmpty, s"swapped-alias baseline should be empty, got $on") + } + + test("Two different relations with the same schema must not be treated as a self-join") { + // The guard under test is `isSameBaseRelation`: it must reject a join between two DIFFERENT + // base tables even when they share a schema and column names. Distinct Parquet tables + // canonicalize to distinct `rootPaths`, so `left.canonicalized == right.canonicalized` is + // false and the rewrite must not fire. This pins a real correctness boundary, not just a + // missed optimization: rewriting `TLeft JOIN TRight` as COUNT(DISTINCT) over TLeft alone would + // drop TRight's rows and change the answer, so removing the guard would make ON diverge from + // OFF here. + createTable("TLeft", "k INT, v INT", " (1, 10), (1, 10), (2, 30)") + createTable("TRight", "k INT, v INT", " (1, 20), (1, 20), (2, 30)") + val sql = + """SELECT k FROM TLeft outer_t WHERE k IN ( + | SELECT s1.k FROM TLeft s1 JOIN TRight s2 + | ON s1.k = s2.k AND s1.v <> s2.v)""".stripMargin + assertRuleNotFired(sql) + val (on, off) = runBoth(sql) + assert(on == off, s"different-relation join semantics diverge: ON=$on OFF=$off") + // k=1: TLeft v={10} vs TRight v={20} -> 10<>20 true -> qualifies; k=2: 30<>30 false -> no. + assert(on == Set(Row(1)), s"expected {1}, got $on") + } + + // ==================== Positive: rewritten plan is structurally the aggregate ================ + // + // Result parity (ON == OFF) proves the two queries return the same rows; it does not prove the + // rewrite produced the specific GROUP BY + HAVING COUNT(DISTINCT) > 1 shape rather than, say, + // leaving the self-join and happening to agree. These controls compare the rewritten plan against + // a hand-written aggregate SQL that spells out the intended shape, INCLUDING two `IS NOT NULL` + // filters: on the equi-key (which the rewrite injects to preserve equi-join NULL semantics) and + // on the neq column. `v IS NOT NULL` is semantically redundant for COUNT(DISTINCT v), which + // already ignores NULL, but the original `s1.v <> s2.v` lets InferFiltersFromConstraints derive + // `isnotnull(v)` and push it below the aggregate, so the equivalent SQL must include it to match + // the shape the optimizer actually produces. The comparison itself is `compareCanonicalizedPlans` + // (see its doc for why canonicalizing first, and disabling checkAnalysis, is required here). + + private def optimizedPlanWith(sql: String, rewrite: Boolean): LogicalPlan = + withSQLConf(rewriteConf -> rewrite.toString) { + spark.sql(sql).queryExecution.optimizedPlan + } + + /** + * Assert two optimized plans are structurally equal. Canonicalize first: the rewrite creates + * fresh aliases, whose names and exprIds are cosmetic for this structural check, while the + * Join-vs-Aggregate shape difference this asserts on survives canonicalization. `checkAnalysis` + * is disabled because both plans are already analyzed and optimized, and a canonicalized plan is + * not re-analyzable (its HAVING references a zeroed exprId), which the default would reject + * before any comparison. + */ + private def compareCanonicalizedPlans(actual: LogicalPlan, expected: LogicalPlan): Unit = + comparePlans(actual.canonicalized, expected.canonicalized, checkAnalysis = false) + + test("Pattern A' rewritten plan is structurally the equivalent aggregate") { + setupTable() + val selfJoinSql = + """SELECT k FROM T outer_t WHERE k IN ( + | SELECT s1.k FROM T s1 JOIN T s2 + | ON s1.k = s2.k AND s1.v <> s2.v)""".stripMargin + val aggregateSql = + """SELECT k FROM T outer_t WHERE k IN ( + | SELECT k FROM T WHERE k IS NOT NULL AND v IS NOT NULL + | GROUP BY k HAVING count(DISTINCT v) > 1)""".stripMargin + + val actual = optimizedPlanWith(selfJoinSql, rewrite = true) + val expected = optimizedPlanWith(aggregateSql, rewrite = false) + assert(ruleFired(actual), s"precondition: rewrite should fire:\n$actual") + assert(!ruleFired(expected), s"precondition: expected plan is the hand-written aggregate") + compareCanonicalizedPlans(actual, expected) + } + + test("Pattern A2 rewritten plan is structurally the equivalent aggregate") { + setupTable() + spark.sql( + """CREATE OR REPLACE TEMP VIEW D AS SELECT * FROM VALUES + | (1), (3), (6) AS D(k)""".stripMargin) + val selfJoinSql = + """SELECT k FROM T outer_t WHERE k IN ( + | SELECT d.k + | FROM D d, (SELECT s1.k FROM T s1 JOIN T s2 + | ON s1.k = s2.k AND s1.v <> s2.v) sj + | WHERE d.k = sj.k)""".stripMargin + val aggregateSql = + """SELECT k FROM T outer_t WHERE k IN ( + | SELECT d.k + | FROM D d, (SELECT k FROM T WHERE k IS NOT NULL AND v IS NOT NULL + | GROUP BY k HAVING count(DISTINCT v) > 1) sj + | WHERE d.k = sj.k)""".stripMargin + + val actual = optimizedPlanWith(selfJoinSql, rewrite = true) + val expected = optimizedPlanWith(aggregateSql, rewrite = false) + assert(ruleFired(actual), s"precondition: rewrite should fire:\n$actual") + assert(!ruleFired(expected), s"precondition: expected plan is the hand-written aggregate") + compareCanonicalizedPlans(actual, expected) + } + + // ==================== Negative: rewrite must produce equivalent results (or bail) ========== + + test("Plain InnerJoin at top level: results unchanged (rewrite must not touch it)") { + setupTable() + val sql = + """SELECT ws1.k FROM T ws1 JOIN T ws2 + |ON ws1.k = ws2.k AND ws1.v <> ws2.v""".stripMargin + // Row-multiplicity matters here; using count() to catch any drop or dup. + var onCount: Long = -1L + var offCount: Long = -1L + withSQLConf(rewriteConf -> "true") { + onCount = spark.sql(sql).count() + } + withSQLConf(rewriteConf -> "false") { + offCount = spark.sql(sql).count() + } + assert( + onCount == offCount, + s"plain InnerJoin row-count differs: rewrite=$onCount vs baseline=$offCount") + assertRuleNotFired(sql) + } + + test("IS DISTINCT FROM is rejected by the self-join condition parser") { + setupTable() + val sql = + """SELECT k FROM T outer_t WHERE k IN ( + | SELECT s1.k FROM T s1 JOIN T s2 + | ON s1.k = s2.k AND s1.v IS DISTINCT FROM s2.v)""".stripMargin + val (on, off) = runBoth(sql) + assert(on == off, s"IS DISTINCT FROM semantics diverge: ON=$on OFF=$off") + // Assert the full result, not just contains(4): unlike `<>`, `IS DISTINCT FROM` treats NULL as + // a value, so k=4 (v={70,NULL}) qualifies alongside k=1, k=3 and k=6. + assert(on == Set(Row(1), Row(3), Row(4), Row(6)), s"expected {1,3,4,6}, got $on") + assertRuleNotFired(sql) + } + + test("IsNotNull on a non-join column is rejected") { + // The guard under test is the predicate parser: it accepts IsNotNull only on a column the + // join condition already references, because such a predicate is implied by the equi-key or + // the inequality and can be dropped, while IsNotNull(w) filters rows the aggregate would + // otherwise count. The control is the same query without that one conjunct. + createTable( + "T3", + "k INT, v INT, w INT", + """ (1, 10, 100), (1, 20, 200), + | (2, 30, CAST(NULL AS INT)), (2, 40, CAST(NULL AS INT))""".stripMargin) + + val controlSql = + """SELECT k FROM T3 outer_t WHERE k IN ( + | SELECT s1.k FROM T3 s1 JOIN T3 s2 + | ON s1.k = s2.k AND s1.v <> s2.v)""".stripMargin + assertRuleFired(controlSql) + val (controlOn, controlOff) = runBoth(controlSql) + assert(controlOn == controlOff, s"T3 control diverges: ON=$controlOn OFF=$controlOff") + assert(controlOn == Set(Row(1), Row(2)), s"T3 control expected {1,2}, got $controlOn") + + val sql = + """SELECT k FROM T3 outer_t WHERE k IN ( + | SELECT s1.k FROM T3 s1 JOIN T3 s2 + | ON s1.k = s2.k AND s1.v <> s2.v AND s1.w IS NOT NULL)""".stripMargin + assertRuleNotFired(sql) + val (on, off) = runBoth(sql) + assert(on == off, s"IsNotNull(non-join-col) semantics diverge: ON=$on OFF=$off") + assert(on == Set(Row(1)), s"expected {1}, got $on") + } + + test("Multiple inequality columns are rejected") { + // The guard under test is `neqPairs.size != 1`. Two inequalities need "at least two rows + // differing in v AND in w", which no count-distinct over a single column can express. The + // control is the same query with only the first inequality. + createTable( + "T2", + "k INT, v INT, w INT", + """ (1, 10, 100), (1, 20, 200), + | (2, 30, 300), + | (3, 40, 100), (3, 50, 100)""".stripMargin) + + val controlSql = + """SELECT k FROM T2 outer_t WHERE k IN ( + | SELECT s1.k FROM T2 s1 JOIN T2 s2 + | ON s1.k = s2.k AND s1.v <> s2.v)""".stripMargin + assertRuleFired(controlSql) + val (controlOn, controlOff) = runBoth(controlSql) + assert(controlOn == controlOff, s"T2 control diverges: ON=$controlOn OFF=$controlOff") + assert(controlOn == Set(Row(1), Row(3)), s"T2 control expected {1,3}, got $controlOn") + + val sql = + """SELECT k FROM T2 outer_t WHERE k IN ( + | SELECT s1.k FROM T2 s1 JOIN T2 s2 + | ON s1.k = s2.k AND s1.v <> s2.v AND s1.w <> s2.w)""".stripMargin + assertRuleNotFired(sql) + val (on, off) = runBoth(sql) + assert(on == off, s"multi-column neq semantics diverge: ON=$on OFF=$off") + assert(on == Set(Row(1)), s"expected {1}, got $on") + } + + test("LeftOuter join is outside existence context: results unchanged") { + setupTable() + val sql = + """SELECT ws1.k FROM T ws1 LEFT OUTER JOIN T ws2 + |ON ws1.k = ws2.k AND ws1.v <> ws2.v""".stripMargin + // Row multiplicity matters here. + var onCount: Long = -1L + var offCount: Long = -1L + withSQLConf(rewriteConf -> "true") { + onCount = spark.sql(sql).count() + } + withSQLConf(rewriteConf -> "false") { + offCount = spark.sql(sql).count() + } + assert(onCount == offCount, s"LeftOuter row-count differs: $onCount vs $offCount") + assertRuleNotFired(sql) + } + + test("Inequality column overlapping an equi-key is rejected") { + setupTable() + val sql = + """SELECT k FROM T outer_t WHERE k IN ( + | SELECT s1.k FROM T s1 JOIN T s2 + | ON s1.k = s2.k AND s1.k <> s2.k)""".stripMargin + val (on, off) = runBoth(sql) + assert(on == off, s"unsatisfiable predicate diverges: ON=$on OFF=$off") + assert(on.isEmpty, s"unsatisfiable predicate should produce empty set, got $on") + assertRuleNotFired(sql) + } + + test("Config gate: rewrite disabled leaves a valid A' candidate untouched") { + setupTable() + val sql = + """SELECT k FROM T outer_t WHERE k IN ( + | SELECT s1.k FROM T s1 JOIN T s2 + | ON s1.k = s2.k AND s1.v <> s2.v)""".stripMargin + withSQLConf(rewriteConf -> "false") { + val plan = spark.sql(sql).queryExecution.optimizedPlan + assert(!ruleFired(plan), s"config off must not fire rewrite:\n$plan") + val res = spark.sql(sql).collect().toSet + assert(res == Set(Row(1), Row(3), Row(6)), s"config off correctness broken: $res") + } + } + + // ==================== Correlated subquery: rule must fail-closed ==================== + + private def setupOuterT(): Unit = { + spark.sql( + """CREATE OR REPLACE TEMP VIEW OuterT AS SELECT * FROM VALUES + | (1), (3), (6) AS OuterT(k)""".stripMargin) + } + + test("Correlated InSubquery is fail-closed") { + setupTable() Review Comment: Three small additions: a case with the InSubquery in the SELECT list or a CASE WHEN (the rule scans all expressions, so it fires there too); a direct test that re-applying the rule leaves the plan unchanged (it runs in two fixedPoint batches; idempotence is currently covered only indirectly by the structural tests); and a precondition assert in the correlated test that with the rewrite off the plan still holds an InSubquery with outer references, so a future pullup change cannot make this negative test vacuous. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/RewriteSelfJoinInequalityToAggregate.scala: ########## @@ -0,0 +1,746 @@ +/* + * 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.execution + +import org.apache.spark.sql.catalyst.expressions._ +import org.apache.spark.sql.catalyst.expressions.aggregate._ +import org.apache.spark.sql.catalyst.plans._ +import org.apache.spark.sql.catalyst.plans.logical._ +import org.apache.spark.sql.catalyst.rules.Rule +import org.apache.spark.sql.catalyst.trees.TreePattern.IN_SUBQUERY +import org.apache.spark.sql.catalyst.util.CharVarcharUtils +import org.apache.spark.sql.execution.datasources.{HadoopFsRelation, LogicalRelation} +import org.apache.spark.sql.execution.datasources.parquet.ParquetFileFormat +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types._ + +/** + * Rewrites a self-join with an inequality into GROUP BY + HAVING COUNT(DISTINCT) > 1. + * + * Targets the two uncorrelated InSubquery shapes exercised by TPC-DS Q95: + * + * - Pattern A': the subquery top-level InnerJoin is a direct self-join. + * - Pattern A2: the subquery contains an outer InnerJoin with a self-join child; only the + * self-join child is replaced with Aggregate and the outer join is preserved. + * + * Both patterns require an existence-only membership context so row-count multiplicity from the + * original self-join cross-product does not affect semantics. Correlated InSubquery expressions are + * intentionally fail-closed because the ExprId remapping performed here does not rewrite correlated + * predicates. + * + * This rule runs in `extendedOperatorOptimizationRules`, which is part of the operator optimization + * batch and therefore executes before `RewritePredicateSubquery` turns the predicate subquery into + * a semi/anti/existence join. It only observes the uncorrelated `InSubquery` shape at that phase, + * so there is no separate LeftSemi/LeftAnti ("Pattern A") or `Exists` handling. + * + * Both patterns share: + * - [[buildAggregateHavingDistinctGt1]] to construct `Filter(cnt > 1, Aggregate)` + * - [[canonicalizeWrapper]] to rebuild a wrapping Project so every equi-key reference points to + * the sjLeft-side attribute, with **fresh exprIds** (Spark's SPARK-21835 style -- no reuse of + * original exprIds), returning an old->new attribute remap for downstream rewrite. + * + * Controlled by `spark.sql.optimizer.rewriteSelfJoinInequalityToAggregate.enabled` + * (default false, opt-in). + */ +object RewriteSelfJoinInequalityToAggregate extends Rule[LogicalPlan] with PredicateHelper { + + private val CountDistinctAliasName = "_rewrite_selfjoin_inequality_cnt_distinct" + + override def apply(plan: LogicalPlan): LogicalPlan = { + if (!conf.getConf(SQLConf.REWRITE_SELF_JOIN_INEQUALITY_TO_AGGREGATE_ENABLED)) { + return plan + } + + // Pattern A' / A2: rewrite uncorrelated InSubquery plans. + // Correlated subqueries carry outer references / correlated join conditions in + // `SubqueryExpression.children`; fail closed because this rule does not remap them. + val rewritten = plan.transformAllExpressionsWithPruning(_.containsPattern(IN_SUBQUERY)) { + case in @ InSubquery(_, lq: ListQuery) if lq.children.isEmpty => + rewriteSubqueryPlan(lq.plan) match { + case Some(newSub) => in.copy(query = lq.copy(plan = newSub)) + case None => in + } + } + if (!(rewritten eq plan)) { + logDebug( + "RewriteSelfJoinInequalityToAggregate: rewrote self-join to " + + "GROUP BY + HAVING COUNT(DISTINCT) > 1") + } + rewritten + } + // ============================================================================ + // Shared helpers + // ============================================================================ + + /** + * Build `Filter(cnt > 1, Aggregate(equiKeys, [equiKeys, cnt_alias], Filter(IsNotNull(equiKeys), + * child)))`. Returns the Filter node whose output is `equiKeys ++ [count_alias_attr]`. + * + * The extra `IsNotNull(equiKeys)` filter is essential to preserve the original equi-join's NULL + * semantics. Under SQL 3VL, `left.k = right.k` never matches when either side is NULL, so the + * original self-join drops rows with NULL equi-keys. Aggregate, in contrast, groups NULL keys + * together into a single "NULL group" -- if that group has >= 2 distinct non-null neq values, + * COUNT(DISTINCT) > 1 fires and injects NULL into the subquery output. That leaked NULL then + * turns `NOT IN` into a spurious empty result (Spark's null-aware anti-join uses + * `Or(equi, IsNull(equi))` which any NULL sub-row satisfies) and can flip IN/NOT IN outcomes. The + * neq column needs no such filter: `COUNT(DISTINCT col)` already ignores NULL. + */ + private def buildAggregateHavingDistinctGt1( + equiKeys: Seq[Attribute], + neqCol: Attribute, + child: LogicalPlan): LogicalPlan = { + val countExpr = AggregateExpression( + Count(Seq(neqCol)), + mode = Complete, + isDistinct = true, + filter = None, + NamedExpression.newExprId) + val countAlias = Alias(countExpr, CountDistinctAliasName)() + // Seq[Attribute] is a Seq[NamedExpression] via covariance; no cast needed. + val aggExprs: Seq[NamedExpression] = equiKeys :+ countAlias + val nonNullChild = equiKeys + .map(a => IsNotNull(a): Expression) + .reduceOption(And) + .map(Filter(_, child)) + .getOrElse(child) + val agg = Aggregate(equiKeys, aggExprs, nonNullChild) + Filter(GreaterThan(countAlias.toAttribute, Literal(1L, LongType)), agg) + } + + /** + * Canonicalize a Project so every equi-key reference points at the sjLeft-side attribute. + * [[parseSelfJoinCondition]] has already verified that each pair refers to the same output + * position on the two structurally identical self-join sides. Uses **fresh exprIds** (no reuse of + * original wrapper output exprIds) -- the same technique Spark's own `dedupSubqueryOnSelfJoin` + * uses when it needs to change subquery output. + * + * Returns the rebuilt Project and a map `oldWrapperOutputExprId -> newWrapperOutputAttr`, so + * downstream references (outer join condition, top-level Project) can be updated consistently. + * + * `equiPairs` provides the definitive ExprId-based lookup: `equiPair (l, r)` binds + * `l.exprId -> l` (identity) and `r.exprId -> l` (sjRight -> sjLeft). Attribute identity in + * Catalyst is ExprId, not name; two columns can share a name with distinct ExprIds. Name-based + * lookup would silently drop such entries via `.toMap`. + * + * Fails (returns None) when a projectList entry is neither an equi-key Attribute (by ExprId) nor + * `Alias(equi-key Attribute, _)`. Fail-closed. + */ + private def canonicalizeWrapper( + projectList: Seq[NamedExpression], + equiPairs: Seq[(Attribute, Attribute)], + newChild: LogicalPlan): Option[(Project, Map[ExprId, Attribute])] = { + // ExprId-based canonical map: any equi-key attribute (either side) -> sjLeft attribute. + val exprIdToLeft: Map[ExprId, Attribute] = + equiPairs.flatMap { case (l, r) => Seq(l.exprId -> l, r.exprId -> l) }.toMap + val oldOutput: Seq[Attribute] = projectList.map(_.toAttribute) + val mapped: Seq[Option[NamedExpression]] = projectList.map { + case a: Attribute if exprIdToLeft.contains(a.exprId) => + // Wrap every rewritten output slot in a fresh Alias. + // + // When a wrapper reprojects BOTH sides of the same equi pair (e.g. + // `SELECT s1.k, s2.k FROM T s1 JOIN T s2 ON s1.k = s2.k AND s1.v <> s2.v`), + // both entries collapse to the same sjLeft Attribute after the self-join is + // rewritten. Duplicate output ExprIds are not illegal in Spark (`SELECT a, a` + // is a valid Project), but fresh Aliases give each output slot an independent + // identity, which keeps the `oldOutput -> newOutput` remap 1-to-1 and lets + // downstream references (outer join condition, top-level Project) be updated + // unambiguously via ExprId. + // + // The fresh ExprId is on the Alias ITSELF; the referenced child keeps its + // original ExprId. Spark's logical-plan integrity checks reject reusing a + // referenced ExprId as the Alias's own ExprId, not duplication across slots. + Some(Alias(exprIdToLeft(a.exprId), a.name)(): NamedExpression) + case al @ Alias(a: Attribute, _) if exprIdToLeft.contains(a.exprId) => + // Fresh exprId; do NOT reuse `al.exprId`. Reusing another expression's exprId + // is the pattern that Spark 3.3 flags via structural-integrity checks. + Some(Alias(exprIdToLeft(a.exprId), al.name)(): NamedExpression) + case _ => None + } + if (mapped.exists(_.isEmpty)) { + None + } else { + val newProjectList = mapped.flatten + val newWrapper = Project(newProjectList, newChild) + val newOutput = newWrapper.output + val remap: Map[ExprId, Attribute] = + oldOutput.zip(newOutput).map { case (o, n) => o.exprId -> n }.toMap + Some((newWrapper, remap)) + } + } + + /** + * Replace equi-key attribute references inside a NamedExpression according to `remap`, while + * preserving the NamedExpression shape. + * + * `Expression.transformUp` returns `Expression`, not `NamedExpression`. We avoid a blanket + * `asInstanceOf[NamedExpression]` by handling the two shapes that can appear in a Project's + * `projectList` explicitly: a bare Attribute (whose top-level may itself be replaced) and an + * Alias (which stays an Alias while its child is transformed). Any other NamedExpression shape we + * do not rewrite is left as-is ONLY if it does not reference a replaced self-join output; + * otherwise it would carry a stale ExprId, so returns None to fail the whole rewrite closed. + */ + private def remapNamedExpressionAttributes( + ne: NamedExpression, + remap: Map[ExprId, Attribute]): Option[NamedExpression] = ne match { + case a: Attribute if remap.contains(a.exprId) => Some(remap(a.exprId)) + case a: Attribute => Some(a) + case al: Alias => + val newChild = al.child.transformUp { + case a: Attribute if remap.contains(a.exprId) => remap(a.exprId) + } + Some( + if (newChild eq al.child) al + else Alias(newChild, al.name)(al.exprId, al.qualifier, al.explicitMetadata)) + case other if other.references.exists(a => remap.contains(a.exprId)) => + // Fail-closed: a NamedExpression we do not rewrite (neither a bare Attribute nor an Alias) + // that still references a replaced self-join output would be left with a dangling ExprId. + // Refuse the rewrite rather than emit a plan with a stale reference. + None + case other => Some(other) + } + + // ============================================================================ + // Pattern A' / A2 dispatch (subquery plans of InSubquery) + // ============================================================================ + + private def rewriteSubqueryPlan(plan: LogicalPlan): Option[LogicalPlan] = { + // Candidate-level nondeterminism guard: reject if ANY node in the whole subquery plan + // is non-repeatable (Rand, LIMIT-without-ORDER-BY, Sample, Offset, streaming). This catches + // nondeterminism that lives ABOVE the self-join rather than on either side -- e.g. a Pattern + // A2 outer join whose condition is `d.k = sj.k AND rand() < 0.5`. Both self-join sides stay + // repeatable there, so the per-side `isSameBaseRelation` check would pass, yet the enclosing + // subquery is not repeatable. + if (!isRepeatablePlan(plan)) return None Review Comment: rewriteSubqueryPlan runs the full isRepeatablePlan audit before the top-level shape match, so the common `IN (SELECT y FROM t)` (no Join at the top) still pays subqueriesAll and two uncached whole-tree whitelist walks, the second of which runs the per-expression whitelist recursion, and the rule sits in two fixedPoint batches. Move the Project/Join shape match ahead of isRepeatablePlan: semantics are unchanged (both guards still apply) and non-candidate subqueries are left with a handful of pattern matches. -- 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]
