cloud-fan commented on code in PR #58424: URL: https://github.com/apache/spark/pull/58424#discussion_r4000488652
########## sql/core/src/main/scala/org/apache/spark/sql/execution/RewriteSelfJoinInequalityToAggregate.scala: ########## @@ -0,0 +1,569 @@ +/* + * 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 supported uncorrelated IN-subquery inequality self-joins into + * `GROUP BY + HAVING MIN(neq) <> MAX(neq)`, avoiding the self-join cross-product. + * + * Supports a direct self-join (Pattern A') and a self-join nested under an outer inner join + * (Pattern A2, where only the self-join child becomes an Aggregate). Unsupported and correlated + * shapes fail closed. + * + * Runs in `extendedOperatorOptimizationRules`, before `RewritePredicateSubquery` turns the + * predicate subquery into a semi/anti/existence join, so it only sees the uncorrelated + * `InSubquery` shape. + * + * Controlled by `spark.sql.optimizer.rewriteSelfJoinInequalityToAggregate.enabled` + * (default false, opt-in). + */ +object RewriteSelfJoinInequalityToAggregate extends Rule[LogicalPlan] with PredicateHelper { + + private val MinNeqAliasName = "_rewrite_selfjoin_inequality_min" + private val MaxNeqAliasName = "_rewrite_selfjoin_inequality_max" + + override def apply(plan: LogicalPlan): LogicalPlan = { + if (!conf.getConf(SQLConf.REWRITE_SELF_JOIN_INEQUALITY_TO_AGGREGATE_ENABLED)) { + return plan + } + + // Fail closed on correlated subqueries: `lq.children` holds the outer references this rule + // does not remap. + 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 + } + } + } + + // ============================================================================ + // Shared helpers + // ============================================================================ + + /** + * Build `Filter(min <> max, Aggregate(equiKeys, child))`, taking MIN and MAX over the neq column. + * `MIN(neqCol) <> MAX(neqCol)` is true exactly when the group holds two or more distinct non-null + * values -- the same test as `COUNT(DISTINCT neqCol) > 1`, but avoids the distinct-dedup + * aggregation stages and supports partial aggregation. + * + * The `IsNotNull(equiKeys)` filter preserves the equi-join's NULL semantics: `=` never matches a + * NULL key, but GROUP BY would fold all NULL keys into one group that can leak NULL into a + * `NOT IN`. The neq column needs no filter -- MIN/MAX ignore NULL, and a group with fewer than + * two non-null values has `min = max` (or both NULL, which makes `<>` NULL), so `<>` is never + * true for it and the group is dropped. + */ + private def buildAggregateHavingMultipleDistinct( + equiKeys: Seq[Attribute], + neqCol: Attribute, + child: LogicalPlan): LogicalPlan = { + val minAlias = Alias(Min(neqCol).toAggregateExpression(), MinNeqAliasName)() + val maxAlias = Alias(Max(neqCol).toAggregateExpression(), MaxNeqAliasName)() + val aggExprs: Seq[NamedExpression] = equiKeys :+ minAlias :+ maxAlias + val nonNullChild = equiKeys + .map(a => IsNotNull(a): Expression) + .reduceOption(And) + .map(Filter(_, child)) + .getOrElse(child) + val agg = Aggregate(equiKeys, aggExprs, nonNullChild) + Filter(Not(EqualTo(minAlias.toAttribute, maxAlias.toAttribute)), agg) + } + + /** + * Rebuild the wrapper Project so every equi-key reference points at the sjLeft attribute with a + * fresh output ExprId, returning `oldOutputExprId -> newOutputAttr` for downstream references + * (outer join condition, top-level Project). Lookup is by ExprId (Catalyst attribute identity), + * not name. Fails closed when an entry is neither an equi-key Attribute nor `Alias(equi-key, _)`. + */ + private def canonicalizeWrapper( + projectList: Seq[NamedExpression], + equiPairs: Seq[(Attribute, Attribute)], + newChild: LogicalPlan): Option[(Project, Map[ExprId, 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) => + // Fresh exprId, but carry over qualifier / metadata so this branch stays consistent with + // the Alias branch and a column keeps its metadata. + Some( + Alias(exprIdToLeft(a.exprId), a.name)( + qualifier = a.qualifier, + explicitMetadata = Some(a.metadata)): NamedExpression) + case al @ Alias(a: Attribute, _) if exprIdToLeft.contains(a.exprId) => + // withNewChild preserves name/qualifier/metadata and exprId; newInstance then re-stamps a + // fresh exprId, so Alias keeps ownership of its own metadata contract instead of us + // re-listing its fields (which drift when Alias gains one). + Some(al.withNewChild(exprIdToLeft(a.exprId)).newInstance()) + case _ => None + } + if (mapped.exists(_.isEmpty)) { + None + } else { + val newProjectList = mapped.flatten + val newWrapper = Project(newProjectList, newChild) + val remap: Map[ExprId, Attribute] = + oldOutput.zip(newWrapper.output).map { case (o, n) => o.exprId -> n }.toMap + Some((newWrapper, remap)) + } + } + + /** + * Replace equi-key references inside a NamedExpression per `remap`, preserving Attribute/Alias + * shape. Any other expression still referencing a replaced output returns None (fail-closed) to + * avoid a dangling ExprId. + */ + 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) + } + // withNewChild preserves the same exprId/qualifier/metadata the manual copy did. + Some(if (newChild eq al.child) al else al.withNewChild(newChild)) + case other if other.references.exists(a => remap.contains(a.exprId)) => + None + case other => Some(other) + } + + // ============================================================================ + // Pattern A' / A2 dispatch (subquery plans of InSubquery) + // ============================================================================ + + private def rewriteSubqueryPlan(plan: LogicalPlan): Option[LogicalPlan] = { + // Match the candidate shape first -- a top-level Inner Join, optionally under one wrapper + // Project. This structural match is cheap, so run it before the whole-subquery + // `isRepeatablePlan` walk and skip that walk entirely for the many subqueries that are not + // even shaped like a self-join. + val (projectListOpt, innerJoin): (Option[Seq[NamedExpression]], Join) = plan match { + case Project(pl, j: Join) if j.joinType == Inner && j.condition.isDefined => + (Some(pl), j) + case j: Join if j.joinType == Inner && j.condition.isDefined => + (None, j) + case _ => return None + } + + // Candidate-level guard: reject if any node in the whole subquery is non-repeatable, catching + // nondeterminism hoisted above the self-join that the per-side `isSameBaseRelation` misses. + if (!isRepeatablePlan(plan)) return None + + if (isSameBaseRelation(innerJoin.left, innerJoin.right)) { + rewriteDirectSelfJoin(projectListOpt, innerJoin) + } else { + rewriteNestedSelfJoin(projectListOpt, innerJoin) + } + } + + // ============================================================================ + // Pattern A' : direct self-join at subquery top level + // ============================================================================ + + private def rewriteDirectSelfJoin( + projectListOpt: Option[Seq[NamedExpression]], + innerJoin: Join): Option[LogicalPlan] = { + // Fail closed on an explicit join hint: it is a directive about the join this rule deletes. + if (!innerJoin.hint.isEmpty) return None + + val innerLeft = innerJoin.left + val innerCond = innerJoin.condition.get + + val parsed = parseSelfJoinCondition(innerCond, innerLeft, innerJoin.right) + if (parsed.isEmpty) return None + val (equiPairs, neqPairs) = parsed.get + + val innerLeftEquiAttrs: Seq[Attribute] = equiPairs.map(_._1) + val innerLeftNeqAttr: Attribute = neqPairs.head._1 + val filtered = + buildAggregateHavingMultipleDistinct(innerLeftEquiAttrs, innerLeftNeqAttr, innerLeft) + + // Fail closed on a bare-Join subquery: with no wrapper Project, replacing the self-join output + // with `Project(equiKeys, filtered)` shrinks arity and RewritePredicateSubquery's positional + // `values.zip(sub.output)` would misbind semi predicates. Q95 subqueries always have a Project. + projectListOpt match { + case None => None + case Some(pl) => + canonicalizeWrapper(pl, equiPairs, filtered).map { case (newWrapper, _) => newWrapper } + } + } + + // ============================================================================ + // Pattern A2 : self-join nested inside another InnerJoin in the subquery + // ============================================================================ + + private def rewriteNestedSelfJoin( + projectListOpt: Option[Seq[NamedExpression]], + outerJoin: Join): Option[LogicalPlan] = { + val outerCond = outerJoin.condition.get + + val (selfJoinSide, selfJoinOnRight) = + tryExtractSelfJoin(outerJoin.right) match { + case Some(_) => (outerJoin.right, true) + case None => + tryExtractSelfJoin(outerJoin.left) match { + case Some(_) => (outerJoin.left, false) + case None => return None + } + } + + val (selfJoinProjectOpt, selfJoin) = selfJoinSide match { + case p @ Project(_, j: Join) if j.joinType == Inner && j.condition.isDefined => + (Some(p), j) + case j: Join if j.joinType == Inner && j.condition.isDefined => + (None, j) + case _ => return None + } + + val sjLeft = selfJoin.left + val sjCond = selfJoin.condition.get + if (!isSameBaseRelation(sjLeft, selfJoin.right)) return None + + val parsed = parseSelfJoinCondition(sjCond, sjLeft, selfJoin.right) + if (parsed.isEmpty) return None + val (equiPairs, neqPairs) = parsed.get + + val sjLeftEquiAttrs: Seq[Attribute] = equiPairs.map(_._1) + val sjLeftNeqAttr: Attribute = neqPairs.head._1 + + val selfJoinOutputSet = selfJoinSide.outputSet + val sjEquiExprIds: Set[ExprId] = + equiPairs.flatMap { case (l, r) => Seq(l.exprId, r.exprId) }.toSet + // A wrapper Project may reproject equi-keys under fresh alias exprIds; include those. + val wrapperEquiExprIds: Set[ExprId] = selfJoinProjectOpt.toSeq.flatMap { p => + p.projectList.flatMap { + case a: Attribute if sjEquiExprIds.contains(a.exprId) => Some(a.exprId) + case al @ Alias(a: Attribute, _) if sjEquiExprIds.contains(a.exprId) => Some(al.exprId) + case _ => None + } + }.toSet + val allEquiExprIds = sjEquiExprIds ++ wrapperEquiExprIds + + // The outer join condition and any top-level Project may reference only equi-key attrs from the + // self-join side (the neq column does not survive the rewrite). + val outerCondRefs = outerCond.references.filter(selfJoinOutputSet.contains) + if (!outerCondRefs.forall(a => allEquiExprIds.contains(a.exprId))) return None + val projectOk = projectListOpt.forall { pl => + val refs = pl.flatMap(_.references).filter(selfJoinOutputSet.contains) + refs.forall(a => allEquiExprIds.contains(a.exprId)) + } + if (!projectOk) return None + + val filtered = buildAggregateHavingMultipleDistinct(sjLeftEquiAttrs, sjLeftNeqAttr, sjLeft) + + val (newSelfJoinSide, outputRemap): (LogicalPlan, Map[ExprId, Attribute]) = + selfJoinProjectOpt match { + case Some(wp) => + canonicalizeWrapper(wp.projectList, equiPairs, filtered) match { + case Some((newWrapper, remap)) => (newWrapper, remap) + case None => return None + } + case None if projectListOpt.isEmpty => + // Fail closed: with no wrapper and no top-level Project, `Project(equiKeys, filtered)` + // shrinks the outer join's arity and RewritePredicateSubquery's positional zip misbinds. + return None + case None => Review Comment: **Non-blocking (P2):** This supported A2 branch changes output arity and remaps right-side ExprIds, but no positive test reaches selfJoinProjectOpt == None with a top-level Project present. All positive A2 cases wrap the self-join, while the bare-Join case has no top Project and intentionally bails out. Please add a bare nested self-join under a top-level Project, require right-key remapping, assert the expected rewrite shape, and compare rewrite-on/off results. ########## sql/core/src/test/scala/org/apache/spark/sql/execution/RewriteSelfJoinInequalityToAggregateSuite.scala: ########## @@ -0,0 +1,1387 @@ +/* + * 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, Attribute, EqualTo, InSubquery, ListQuery, Not} +import org.apache.spark.sql.catalyst.optimizer.ReorderJoin +import org.apache.spark.sql.catalyst.plans.Inner +import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Filter, Join, 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 aliases produced by the rewrite; presence of both => rule definitely fired. */ + private val MinNeqAlias = "_rewrite_selfjoin_inequality_min" + private val MaxNeqAlias = "_rewrite_selfjoin_inequality_max" + + // Descends into subqueries: `QueryPlan.exists` does not, and the rewrite's signature alias lives + // inside the IN-subquery when the rule runs on an analyzed (not-yet-rewritten) plan. + private def hasAlias(plan: LogicalPlan, name: String): Boolean = + plan.collectFirstWithSubqueries { + case p if p.expressions.exists(_.exists { + case a: Alias if a.name == name => true + case _ => false + }) => () + }.isDefined + + /** Require BOTH aliases: a rewrite that emitted MIN but dropped MAX is still a bug. */ + private def ruleFired(plan: LogicalPlan): Boolean = + hasAlias(plan, MinNeqAlias) && hasAlias(plan, MaxNeqAlias) + + 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") + } + } + + /** + * Optimize just the IN-subquery plan (rewrite left at its default-off) and return the result, so + * a test can prove what shape the subquery reaches the rule as -- e.g. a bare Join with no + * wrapper Project -- before asserting the rule declines it. Without this, `assertRuleNotFired` + * alone can pass merely because the fixture never produced the shape the guard means to reject. + */ + private def optimizedInSubqueryPlan(sql: String): LogicalPlan = { + val analyzed = spark.sql(sql).queryExecution.analyzed + val subqueries = analyzed.subqueriesAll + assert( + subqueries.length == 1, + s"expected exactly one subquery in analyzed plan, got ${subqueries.length}:\n$analyzed") + spark.sessionState.optimizer.execute(subqueries.head) + } + + /** + * 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. The rewrite must not change the outer query's + * row multiplicity, so assert ON and OFF agree as MULTISETS (a `.toSet` here would hide a + * duplicated or dropped row) before handing callers the row sets their fixed-value assertions + * compare against. `QueryTest.sameRows` is Spark's own multiset comparison (order-insensitive, + * duplicate-sensitive) and formats the offending rows on mismatch. + */ + private def runBoth(sql: String): (Set[Row], Set[Row]) = { + val on = withSQLConf(rewriteConf -> "true") { + spark.sql(sql).collect().toSeq + } + val off = withSQLConf(rewriteConf -> "false") { + spark.sql(sql).collect().toSeq + } + QueryTest.sameRows(on, off).foreach { error => + fail(s"rewrite changed row multiplicity between ON and OFF:\n$error") + } + (on.toSet, off.toSet) + } + + 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, min(v)==max(v)==100 so min<>max is + // false. 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") { + withTable("T") { + 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("InSubquery in a SELECT-list CASE WHEN is rewritten, not only in a WHERE predicate") { + withTable("T") { + // `apply` rewrites via transformAllExpressionsWithPruning, so an InSubquery anywhere in the + // plan is a candidate -- not just a WHERE filter. A Project is a valid host for an InSubquery + // (see ValidateSubqueryExpression), so pin that the rewrite fires when the same uncorrelated + // self-join subquery sits inside a projected CASE WHEN. Moving it out of WHERE is the only + // change from the Pattern A' control above. + setupTable() + val subquery = + """SELECT s1.k FROM T s1 JOIN T s2 + | ON s1.k = s2.k AND s1.v <> s2.v""".stripMargin + val sql = + s"""SELECT k, CASE WHEN k IN ($subquery) THEN 1 ELSE 0 END AS flag + |FROM T outer_t""".stripMargin + + assertRuleFired(sql) + val (on, off) = runBoth(sql) + assert(on == off, s"SELECT-list InSubquery rewrite ON $on != OFF $off") + } + } + + test("Rewrite is idempotent: a second application on the rewritten plan is a no-op") { + withTable("T") { + // After the rewrite the outer `k IN (...)` is still an InSubquery, now over the aggregate, so + // the rule revisits it on any later pass. Applying the rule to its own output must change + // nothing: the aggregate no longer matches the self-join shape, so the second pass returns + // the plan unchanged. + 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 -> "true") { + val analyzed = spark.sql(sql).queryExecution.analyzed + val once = RewriteSelfJoinInequalityToAggregate(analyzed) + assert(ruleFired(once), s"precondition: first pass should fire:\n$once") + val twice = RewriteSelfJoinInequalityToAggregate(once) + assert(twice == once, s"rule is not idempotent:\n$once\n-- second pass -->\n$twice") + } + } + } + + test("Pattern A2: nested self-join is rewritten") { + withTable("T") { + withTempView("D") { + 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") { + withTable("T") { + withTempView("D") { + // 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") { + withTable("T") { + withTempView("D") { + // 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("Clock-dependent STRING -> TIMESTAMP cast in the subquery is fail-closed") { + withTable("TTs") { + // A time-only STRING cast to TIMESTAMP takes its missing date from the runtime clock + // (LocalDate.now), so two self-join scans that straddle midnight can disagree while the + // rewrite folds them into one evaluation. The neq column here is derived by such a cast, so + // the rewrite must not fire. Isolate the cast as the sole cause with a same-shape control + // whose derived neq column uses a benign INT -> BIGINT cast (which does fire), and cover the + // nested form CAST(CAST(t AS ARRAY<TIMESTAMP>) AS STRING) where the timestamp conversion + // hides inside an array element cast and the final column type is STRING. + createTable( + "TTs", + "k INT, v INT, t STRING", + """ (1, 10, '01:00:00'), (1, 20, '02:00:00'), + | (2, 30, '03:00:00')""".stripMargin) + + def wrapped(neqExpr: String): String = + s"""SELECT k FROM TTs outer_t WHERE k IN ( + | SELECT s1.k FROM + | (SELECT k, $neqExpr AS nc FROM TTs) s1 + | JOIN (SELECT k, $neqExpr AS nc FROM TTs) s2 + | ON s1.k = s2.k AND s1.nc <> s2.nc)""".stripMargin + + // Same wrapper shape with a repeatable cast fires, proving the shape itself is supported. + assertRuleFired(wrapped("CAST(v AS BIGINT)")) + + // Direct and array-nested STRING -> TIMESTAMP casts are both fail-closed. + assertRuleNotFired(wrapped("CAST(t AS TIMESTAMP)")) + assertRuleNotFired(wrapped("CAST(CAST(ARRAY(t) AS ARRAY<TIMESTAMP>) AS STRING)")) + + // STRING -> TIMESTAMP_NTZ stays supported: it does not consult the session time zone and a + // time-only string parses to NULL deterministically rather than borrowing the runtime date, + // so the guard only rejects the clock-dependent LTZ conversion, not all string-to-timestamp. + assertRuleFired(wrapped("CAST(t AS TIMESTAMP_NTZ)")) + } + } + + test("Pattern A': multi-equi tuple IN with sjRight key remap is rewritten") { + withTable("TM") { + // 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") { + withTable("T") { + 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") { + withTable("TN") { + withTempView("OuterKeys") { + 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") { + withTable("AliasBase") { + // 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") { + withTable("TLeft", "TRight") { + // 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 MIN(v) <> MAX(v) 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) does not prove the rewrite produced the GROUP BY + HAVING + // MIN(v) <> MAX(v) shape rather than leaving the self-join and happening to agree. These controls + // assert that shape directly. A full canonicalized-plan comparison against hand-written aggregate + // SQL would be wrong here: InferFiltersFromConstraints adds isnotnull(min)/isnotnull(max) to a + // hand-written HAVING but not to the rule's own filter (created in a later batch), so it would + // fail on redundant null predicates -- and matching it by hardening the rule would be the wrong + // fix. + + private def optimizedPlanWith(sql: String, rewrite: Boolean): LogicalPlan = + withSQLConf(rewriteConf -> rewrite.toString) { + spark.sql(sql).queryExecution.optimizedPlan + } + + // The rewrite shape: an Aggregate emitting both signature aliases, with a Filter on top whose + // condition includes Not(EqualTo(min, max)). Match the two operands by ExprId (Catalyst attribute + // identity), not by name -- the rule itself never trusts names -- so a same-named attribute from + // elsewhere cannot satisfy it. `exists` on the condition, not exact match, because + // InferFiltersFromConstraints may fold redundant isnotnull(min/max) into the same Filter. + private def assertMinMaxRewriteShape(plan: LogicalPlan): Unit = { + val found = plan.collectFirstWithSubqueries { + case Filter(cond, agg: Aggregate) + if { + // Key by alias name so the shape requires exactly one MIN alias AND one MAX alias: two + // same-named aliases collapse to a single map key and fail the keySet check, which a + // bare `size == 2` on exprIds would not catch. + val signatureAttrs = agg.aggregateExpressions.collect { + case a: Alias if a.name == MinNeqAlias || a.name == MaxNeqAlias => + a.name -> a.toAttribute + }.toMap + signatureAttrs.keySet == Set(MinNeqAlias, MaxNeqAlias) && cond.exists { + case Not(EqualTo(l: Attribute, r: Attribute)) => + Set(l.exprId, r.exprId) == signatureAttrs.values.map(_.exprId).toSet + case _ => false + } + } => () + } + assert(found.isDefined, s"expected a MIN(v) <> MAX(v) aggregate rewrite shape:\n$plan") + } + + test("Pattern A' rewritten plan is structurally the equivalent aggregate") { + withTable("T") { + 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 actual = optimizedPlanWith(selfJoinSql, rewrite = true) + assert(ruleFired(actual), s"precondition: rewrite should fire:\n$actual") + assertMinMaxRewriteShape(actual) + } + } + + test("Pattern A2 rewritten plan is structurally the equivalent aggregate") { + withTable("T") { + withTempView("D") { + 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 actual = optimizedPlanWith(selfJoinSql, rewrite = true) + assert(ruleFired(actual), s"precondition: rewrite should fire:\n$actual") + assertMinMaxRewriteShape(actual) + } + } + } + + // ==================== Negative: rewrite must produce equivalent results (or bail) ========== + + test("Plain InnerJoin at top level: results unchanged (rewrite must not touch it)") { + withTable("T") { + 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("Bare-Join subquery (no wrapper Project) fails closed: Pattern A' arity guard") { + withTable("T") { + // The guard under test is the `projectListOpt match { case None => None }` bail in + // `rewriteDirectSelfJoin`. With no wrapper Project the self-join output is `left ++ right`; + // replacing it with `Project(equiKeys, aggregate)` would shrink the arity that + // RewritePredicateSubquery later positionally zips against, misbinding the semi predicates. A + // `SELECT *` over the self-join whose tuple IN references every output column lets + // RemoveNoopOperators strip the identity Project, so the subquery reaches the rule as a bare + // Join -- proven below before asserting the rule declines it. + setupTable() + val sql = + """SELECT k FROM T outer_t WHERE (k, v, k, v) IN ( + | SELECT * FROM T s1 JOIN T s2 + | ON s1.k = s2.k AND s1.v <> s2.v)""".stripMargin + optimizedInSubqueryPlan(sql) match { + case _: Join => + case other => fail(s"expected a bare Join subquery, got:\n$other") + } + assertRuleNotFired(sql) + val (on, off) = runBoth(sql) + assert(on == off, s"bare-Join A' arity guard semantics diverge: ON=$on OFF=$off") + } + } + + test("Bare-Join nested self-join (no Project anywhere) fails closed: Pattern A2 arity guard") { + withTable("T") { + withTempView("D") { + // The guard under test is the `case None if projectListOpt.isEmpty => return None` bail in + // `rewriteNestedSelfJoin`: with no wrapper Project and no top-level Project, changing the + // self-join output arity would misbind the positional semi-predicate zip. `SELECT *` plus a + // tuple IN over every column lets RemoveNoopOperators expose the bare nested self-join. + // ReorderJoin is excluded (a valid config) so the bare nested self-join deterministically + // reaches this rule rather than being reshaped away; the shape is then asserted so the test + // fails loudly if it stops reaching the guard. The two sides use renaming subqueries (ka/va + // vs kb/vb) so `SELECT *` yields distinct names -- a plain `T s1 JOIN T s2` would emit two + // columns named `k` and fail analysis; the renames are children of the self-join, so the + // identity `SELECT *` is still stripped. + 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, v, k, k, v) IN ( + | SELECT * FROM D d JOIN ( + | SELECT * FROM (SELECT k AS ka, v AS va FROM T) s1 + | JOIN (SELECT k AS kb, v AS vb FROM T) s2 + | ON s1.ka = s2.kb AND s1.va <> s2.vb) sj + | ON d.k = sj.ka)""".stripMargin + withSQLConf(SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> ReorderJoin.ruleName) { + // A child that is a `sameResult` Inner Join carrying an inequality is exactly the + // self-join rewriteNestedSelfJoin extracts; asserting it proves execution reaches the + // arity guard. + def isTargetSelfJoin(p: LogicalPlan): Boolean = p match { + case j: Join if j.joinType == Inner => + j.left.sameResult(j.right) && + j.condition.exists(_.exists { case _: Not => true; case _ => false }) + case _ => false + } + optimizedInSubqueryPlan(sql) match { + case j: Join if isTargetSelfJoin(j.left) || isTargetSelfJoin(j.right) => + case other => fail(s"expected a bare nested self-join child, got:\n$other") + } + assertRuleNotFired(sql) + val (on, off) = runBoth(sql) + assert(on == off, s"bare-Join A2 arity guard semantics diverge: ON=$on OFF=$off") + } + } + } + } + + test("IS DISTINCT FROM is rejected by the self-join condition parser") { + withTable("T") { + 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") { + withTable("T3") { + // 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") { + withTable("T2") { + // The guard under test is `neqPairs.size != 1`. Two inequalities need "at least two rows + // differing in v AND in w", which no MIN/MAX 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") { + withTable("T") { + 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("Join hint on the self-join is fail-closed") { + withTable("T") { + // The guard under test is `hint != JoinHint.NONE`. The rewrite deletes the self-join, so a + // hint on it is a directive about a join that would vanish; fail closed instead. Control (no + // hint) fires; the same query with a BROADCAST hint on the inner self-join -- the only change + // -- must not fire. A hint never changes rows, so results are identical either way. + setupTable() + val controlSql = + """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(controlSql) + + val hintedSql = + """SELECT k FROM T outer_t WHERE k IN ( + | SELECT /*+ BROADCAST(s2) */ s1.k FROM T s1 JOIN T s2 + | ON s1.k = s2.k AND s1.v <> s2.v)""".stripMargin + assertRuleNotFired(hintedSql) + val (on, off) = runBoth(hintedSql) + assert(on == off, s"hinted self-join semantics diverge: ON=$on OFF=$off") + assert(on == Set(Row(1), Row(3), Row(6)), s"expected {1,3,6}, got $on") + } + } + + test("Inequality column overlapping an equi-key is rejected") { + withTable("T") { + 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") { + withTable("T") { + 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") { + withTable("T") { + withTempView("OuterT") { + setupTable() + setupOuterT() + val sql = + """SELECT o.k FROM OuterT o WHERE o.k IN ( + | SELECT s1.k FROM T s1 JOIN T s2 + | ON s1.k = s2.k AND s1.v <> s2.v + | WHERE s2.k = o.k)""".stripMargin + // Precondition: the InSubquery is genuinely correlated -- its ListQuery carries outer + // references -- so the not-fired result exercises the `lq.children.isEmpty` fail-closed + // guard in `apply`, not an unrelated shape mismatch that would leave this test vacuous. + val analyzed = spark.sql(sql).queryExecution.analyzed + var sawCorrelated = false + analyzed.foreach { node => + node.expressions.foreach(_.foreach { + case InSubquery(_, lq: ListQuery) => sawCorrelated ||= lq.children.nonEmpty + case _ => + }) + } + assert(sawCorrelated, s"expected a correlated InSubquery in analyzed plan:\n$analyzed") + val (on, off) = runBoth(sql) + assert(on == off, s"correlated IN parity: ON=$on OFF=$off") + assert(on == Set(Row(1), Row(3), Row(6)), s"expected {1,3,6}, got $on") + assertRuleNotFired(sql) + } + } + } + + // ==================== Repeatability whitelist: unknown operators fail-closed ============== + + test("Aggregate (first) inside subquery breaks row-bag repeatability: rule bails out") { + // FIRST() is order-dependent and its aggregate result is not row-bag repeatable across + // two evaluations, yet Catalyst's Expression.deterministic returns true. The whitelist + // in `isRowBagRepeatable` must reject any Aggregate node inside the subquery plan. + // range(...) avoids ConvertToLocalRelation folding the Aggregate away. + val sql = + """SELECT k FROM (SELECT CAST(id AS INT) AS k, CAST(id AS INT) AS v FROM range(100)) t + |WHERE k IN ( + | SELECT s1.k FROM ( + | SELECT CAST(id % 10 AS INT) AS k, first(CAST(id AS INT)) AS v + | FROM range(200) GROUP BY id % 10 + | ) s1 + | JOIN ( + | SELECT CAST(id % 10 AS INT) AS k, first(CAST(id AS INT)) AS v + | FROM range(200) GROUP BY id % 10 + | ) s2 + | ON s1.k = s2.k AND s1.v <> s2.v)""".stripMargin + assertRuleNotFired(sql) + } + + test("Window (row_number) inside subquery breaks row-bag repeatability: rule bails out") { + // ROW_NUMBER over non-total order breaks ties nondeterministically. Whitelist rejects + // any Window node inside the subquery plan. + val sql = + """SELECT k FROM (SELECT CAST(id AS INT) AS k, CAST(id AS INT) AS v FROM range(100)) t + |WHERE k IN ( + | SELECT s1.k FROM ( + | SELECT k, ROW_NUMBER() OVER (PARTITION BY k ORDER BY grp) AS v + | FROM (SELECT CAST(id % 10 AS INT) AS k, CAST(id % 3 AS INT) AS grp FROM range(200)) + | ) s1 + | JOIN ( + | SELECT k, ROW_NUMBER() OVER (PARTITION BY k ORDER BY grp) AS v + | FROM (SELECT CAST(id % 10 AS INT) AS k, CAST(id % 3 AS INT) AS grp FROM range(200)) + | ) s2 + | ON s1.k = s2.k AND s1.v <> s2.v)""".stripMargin + assertRuleNotFired(sql) + } + + test("Nondeterministic self-join input is rejected") { + // The guard under test is `plan.deterministic` inside `isRepeatablePlan`. Both sides use the + // same explicit seed, so the two subplans do have the same canonical shape and the rejection + // cannot come from `isSameBaseRelation`. The control replaces `rand(41) < 0.5` with a + // deterministic filter and nothing else, proving this Range/Filter/Project shape does reach + // the rewrite. + val controlSql = + """SELECT k FROM (SELECT CAST(id AS INT) AS k, CAST(id AS INT) AS v FROM range(100)) t + |WHERE k IN ( + | SELECT s1.k FROM ( + | SELECT CAST(id % 10 AS INT) AS k, CAST(id AS INT) AS v + | FROM range(1000) WHERE id % 2 = 0 + | ) s1 + | JOIN ( + | SELECT CAST(id % 10 AS INT) AS k, CAST(id AS INT) AS v + | FROM range(1000) WHERE id % 2 = 0 + | ) s2 + | ON s1.k = s2.k AND s1.v <> s2.v)""".stripMargin + assertRuleFired(controlSql) + val (controlOn, controlOff) = runBoth(controlSql) + assert(controlOn == controlOff, s"range control diverges: ON=$controlOn OFF=$controlOff") + assert( + controlOn == Set(Row(0), Row(2), Row(4), Row(6), Row(8)), + s"range control expected the even keys, got $controlOn") + + val sql = + """SELECT k FROM (SELECT CAST(id AS INT) AS k, CAST(id AS INT) AS v FROM range(100)) t + |WHERE k IN ( + | SELECT s1.k FROM ( + | SELECT CAST(id % 10 AS INT) AS k, CAST(id AS INT) AS v + | FROM range(1000) WHERE rand(41) < 0.5 + | ) s1 + | JOIN ( + | SELECT CAST(id % 10 AS INT) AS k, CAST(id AS INT) AS v + | FROM range(1000) WHERE rand(41) < 0.5 + | ) s2 + | ON s1.k = s2.k AND s1.v <> s2.v)""".stripMargin + assertRuleNotFired(sql) + } + + test("LogicalRDD leaf is not a trusted repeatable source: rule bails out") { + withTable("RddCtl") { + withTempView("RddT") { + // The guard under test is the leaf allowlist in `isRowBagRepeatable`: a Parquet + // `LogicalRelation` is trusted, but a `LogicalRDD` (createDataFrame over an RDD) wraps an + // arbitrary RDD lineage whose runtime row bag Catalyst cannot prove repeatable, so it must + // fail closed even though `plan.deterministic` is true. Both fixtures are + // MultiInstanceRelation leaves, so each self-join dedups into two structurally identical + // sides without a rename-only Project -- the rejection therefore comes from the leaf + // allowlist, not `isSameBaseRelation`. The Parquet control uses the same schema, data and + // query shape and fires, which is what makes the LogicalRDD negative evidence that the leaf + // allowlist -- not a structural mismatch -- rejected it. + createTable("RddCtl", "k INT, v INT", " (1, 10), (1, 20), (2, 30)") + val controlSql = + """SELECT k FROM RddCtl outer_t WHERE k IN ( + | SELECT s1.k FROM RddCtl s1 JOIN RddCtl s2 + | ON s1.k = s2.k AND s1.v <> s2.v)""".stripMargin + assertRuleFired(controlSql) + val (controlOn, controlOff) = runBoth(controlSql) + assert(controlOn == controlOff, s"Parquet control diverges: ON=$controlOn OFF=$controlOff") + assert(controlOn == Set(Row(1)), s"Parquet control expected {1}, got $controlOn") + + val schema = StructType(Seq(StructField("k", IntegerType), StructField("v", IntegerType))) + val rows = spark.sparkContext.parallelize(Seq(Row(1, 10), Row(1, 20), Row(2, 30))) + spark.createDataFrame(rows, schema).createOrReplaceTempView("RddT") + val sql = + """SELECT k FROM RddT outer_t WHERE k IN ( + | SELECT s1.k FROM RddT s1 JOIN RddT s2 + | ON s1.k = s2.k AND s1.v <> s2.v)""".stripMargin + assertRuleNotFired(sql) + val (on, off) = runBoth(sql) + assert(on == off, s"LogicalRDD semantics diverge: ON=$on OFF=$off") + assert(on == Set(Row(1)), s"expected {1}, got $on") + } + } + } + + test("Non-allowlisted deterministic expression (Abs) fails closed") { + withTable("T") { + // The guard under test is the expression allowlist in `isRepeatableExpression`: it trusts + // expression TYPES, not merely `deterministic`. Abs is deterministic, but is intentionally + // not yet part of the expression allowlist; until its repeatability contract is explicitly + // admitted there, a self-join whose side projects abs(v) fails closed -- a missed + // optimization, not a correctness bug. This test verifies that an unknown-but-deterministic + // expression is not silently let through by `plan.deterministic`. + // + // Control and negative are a single-variable pair: both project one derived column and differ + // ONLY in its expression. The control uses `v + 1` (Add over Attribute + Literal, all + // allowlisted) and fires; wrapping the same column in abs() -- the sole change -- makes it + // not fire, so the rejection is attributable to the expression allowlist rather than a + // structural mismatch. `+ 1` (not `+ 0`) is used so the Add survives arithmetic + // simplification and the control genuinely exercises a compound allowlisted expression. v is + // INT here, so the Add is a plain `Add(v, 1)` with no decimal PromotePrecision / + // CheckOverflow wrappers. + setupTable() + + val controlSql = + """SELECT k FROM T outer_t WHERE k IN ( + | SELECT s1.k + | FROM (SELECT k, v + 1 AS x FROM T) s1 + | JOIN (SELECT k, v + 1 AS x FROM T) s2 + | ON s1.k = s2.k AND s1.x <> s2.x)""".stripMargin + assertRuleFired(controlSql) + val (controlOn, controlOff) = runBoth(controlSql) + assert(controlOn == controlOff, s"Add control diverges: ON=$controlOn OFF=$controlOff") + // v+1 is injective over the (non-null) v values, so distinctness per k is unchanged: {1,3,6}. + assert( + controlOn == Set(Row(1), Row(3), Row(6)), + s"Add control expected {1,3,6}, got $controlOn") + + val sql = + """SELECT k FROM T outer_t WHERE k IN ( + | SELECT s1.k + | FROM (SELECT k, abs(v) AS x FROM T) s1 + | JOIN (SELECT k, abs(v) AS x FROM T) s2 + | ON s1.k = s2.k AND s1.x <> s2.x)""".stripMargin + assertRuleNotFired(sql) + val (on, off) = runBoth(sql) + assert(on == off, s"Abs-projected self-join semantics diverge: ON=$on OFF=$off") + } + } + + // ==================== Data-type safety: comparison vs grouping/MIN-MAX equality ============= + // + // The rewrite turns `<>` into MIN(v) <> MAX(v) and `=` into GROUP BY, so it is only sound on + // types where SQL comparison equality coincides with grouping/MIN-MAX ordering equality. + // `parseSelfJoinCondition` gates BOTH the neq column and every equi-key through + // `isSafeComparisonGroupingType` (a positive allowlist, not `RowOrdering.isOrderable`). These + // tests pin the boundary for the risky types. + + test("Float/Double neq column is rejected (comparison-vs-MIN-MAX contract, defensive)") { + withTable("TFloat") { + // The guard under test is `isSafeComparisonGroupingType` on the NEQ column. Control and + // negative differ only in which column feeds the inequality: `vi` (Int, allowlisted) fires, + // `vd` (Double) does not. Double fails closed defensively: the rewrite depends on comparison + // equality and grouping/MIN-MAX ordering equality agreeing, and for floating point that + // agreement on signed zero and NaN rests on normalization details (NormalizeFloatingNumbers) + // that need not match across Spark versions or native backends. On current Spark, comparison + // semantics and aggregation normalization are aligned for these cases, so the OFF baseline of + // the negative query is {1}; the rule does not fire, so ON matches it. The test pins + // fail-closed behavior, not a divergence in current Spark. + createTable( + "TFloat", + "k INT, vi INT, vd DOUBLE", + """ (1, 10, 1.0), (1, 20, 2.0), + | (2, 30, 0.0), (2, 30, -0.0), + | (3, 40, CAST('NaN' AS DOUBLE)), (3, 50, CAST('NaN' AS DOUBLE))""".stripMargin) + + val controlSql = + """SELECT k FROM TFloat outer_t WHERE k IN ( + | SELECT s1.k FROM TFloat s1 JOIN TFloat s2 + | ON s1.k = s2.k AND s1.vi <> s2.vi)""".stripMargin + assertRuleFired(controlSql) + val (controlOn, controlOff) = runBoth(controlSql) + assert(controlOn == controlOff, s"Int-neq control diverges: ON=$controlOn OFF=$controlOff") + assert(controlOn == Set(Row(1), Row(3)), s"Int-neq control expected {1,3}, got $controlOn") + + val sql = + """SELECT k FROM TFloat outer_t WHERE k IN ( + | SELECT s1.k FROM TFloat s1 JOIN TFloat s2 + | ON s1.k = s2.k AND s1.vd <> s2.vd)""".stripMargin + assertRuleNotFired(sql) + val (on, off) = runBoth(sql) + assert(on == off, s"Double-neq semantics diverge: ON=$on OFF=$off") + // k=1 matches (1.0 <> 2.0); k=2 does not (0.0 = -0.0); k=3 does not (Spark NaN = NaN). + assert(on == Set(Row(1)), s"Double-neq baseline expected {1}, got $on") + } + } + + test("Float/Double equi-key is rejected (defensive fail-closed)") { + withTable("TFloatKey") { + // The guard under test is `isSafeComparisonGroupingType` on the EQUI key. Current Spark + // aligns floating-point comparison semantics with grouping normalization here, but the rule + // does not depend on that implementation contract, so it fails closed. Control and negative + // differ only in the equi-key column: `ki` (Int) fires, `kd` (Double) does not. + createTable( + "TFloatKey", + "kd DOUBLE, ki INT, v INT", + """ (1.0, 1, 10), (1.0, 1, 20), + | (2.0, 2, 30), + | (0.0, 3, 40), (-0.0, 3, 50)""".stripMargin) + + val controlSql = + """SELECT ki FROM TFloatKey outer_t WHERE ki IN ( + | SELECT s1.ki FROM TFloatKey s1 JOIN TFloatKey s2 + | ON s1.ki = s2.ki AND s1.v <> s2.v)""".stripMargin + assertRuleFired(controlSql) + val (controlOn, controlOff) = runBoth(controlSql) + assert(controlOn == controlOff, s"Int-key control diverges: ON=$controlOn OFF=$controlOff") + assert(controlOn == Set(Row(1), Row(3)), s"Int-key control expected {1,3}, got $controlOn") + + val sql = + """SELECT kd FROM TFloatKey outer_t WHERE kd IN ( + | SELECT s1.kd FROM TFloatKey s1 JOIN TFloatKey s2 + | ON s1.kd = s2.kd AND s1.v <> s2.v)""".stripMargin + assertRuleNotFired(sql) + val (on, off) = runBoth(sql) + assert(on == off, s"Double-key semantics diverge: ON=$on OFF=$off") + } + } + + test("Complex-type neq column (array/struct) is rejected wholesale") { + withTable("TCplx") { + // The guard under test is the wholesale rejection of complex types by + // `isSafeComparisonGroupingType`, which also covers any Float/Double nested inside them. + // Control (`v` Int) fires; the ARRAY<DOUBLE> and STRUCT<..DOUBLE> variants -- identical query + // shape, only the neq column changed -- do not. + createTable( + "TCplx", + "k INT, v INT, a ARRAY<DOUBLE>, s STRUCT<x: INT, y: DOUBLE>", + """ (1, 10, ARRAY(1.0), NAMED_STRUCT('x', 1, 'y', 1.0)), + | (1, 20, ARRAY(2.0), NAMED_STRUCT('x', 2, 'y', 2.0)), + | (2, 30, ARRAY(1.0), NAMED_STRUCT('x', 1, 'y', 1.0))""".stripMargin) + + val controlSql = + """SELECT k FROM TCplx outer_t WHERE k IN ( + | SELECT s1.k FROM TCplx s1 JOIN TCplx s2 + | ON s1.k = s2.k AND s1.v <> s2.v)""".stripMargin + assertRuleFired(controlSql) + val (controlOn, controlOff) = runBoth(controlSql) + assert(controlOn == controlOff, s"Int-neq control diverges: ON=$controlOn OFF=$controlOff") + assert(controlOn == Set(Row(1)), s"Int-neq control expected {1}, got $controlOn") + + val arraySql = + """SELECT k FROM TCplx outer_t WHERE k IN ( + | SELECT s1.k FROM TCplx s1 JOIN TCplx s2 + | ON s1.k = s2.k AND s1.a <> s2.a)""".stripMargin + assertRuleNotFired(arraySql) + val (arrayOn, arrayOff) = runBoth(arraySql) + assert(arrayOn == arrayOff, s"array-neq semantics diverge: ON=$arrayOn OFF=$arrayOff") + + val structSql = + """SELECT k FROM TCplx outer_t WHERE k IN ( + | SELECT s1.k FROM TCplx s1 JOIN TCplx s2 + | ON s1.k = s2.k AND s1.s <> s2.s)""".stripMargin + assertRuleNotFired(structSql) + val (structOn, structOff) = runBoth(structSql) + assert(structOn == structOff, s"struct-neq semantics diverge: ON=$structOn OFF=$structOff") + } + } + + test("String neq/equi key: default (UTF8_BINARY) fires, non-binary collation fails closed") { + withTable("TStrBin", "TStrCiNeq", "TStrCiEqui") { + // The guard under test is the StringType branch of `isSafeComparisonGroupingType`: only + // `supportsBinaryEquality` (byte-wise) strings are admitted, because a non-binary collation + // (Spark 4.0+) routes comparison and grouping through different code paths. The control uses + // a default-collation table for BOTH the equi key and the neq column and fires -- proving + // plain strings are not rejected wholesale. + // + // The two negatives collate exactly ONE column each, so each pins the rejection to a specific + // gate: a UTF8_LCASE NEQ column exercises the neq-side check, a UTF8_LCASE EQUI key exercises + // the equi-side check. Collating both at once would leave it ambiguous which gate fired and + // stay green if either were deleted -- mirroring the split Float neq / Float equi coverage. + createTable( + "TStrBin", + "k STRING, v STRING", + """ ('a', 'x'), ('a', 'y'), + | ('b', 'z')""".stripMargin) + val binSql = + """SELECT k FROM TStrBin outer_t WHERE k IN ( + | SELECT s1.k FROM TStrBin s1 JOIN TStrBin s2 + | ON s1.k = s2.k AND s1.v <> s2.v)""".stripMargin + assertRuleFired(binSql) + val (binOn, binOff) = runBoth(binSql) + assert(binOn == binOff, s"binary-string control diverges: ON=$binOn OFF=$binOff") + assert(binOn == Set(Row("a")), s"binary-string control expected {a}, got $binOn") + + // Negative 1: only the NEQ column is non-binary collated -> neq-side type gate rejects. + createTable( + "TStrCiNeq", + "k STRING, v STRING COLLATE UTF8_LCASE", + """ ('a', 'x'), ('a', 'y'), + | ('b', 'z')""".stripMargin) + val ciNeqSql = + """SELECT k FROM TStrCiNeq outer_t WHERE k IN ( + | SELECT s1.k FROM TStrCiNeq s1 JOIN TStrCiNeq s2 + | ON s1.k = s2.k AND s1.v <> s2.v)""".stripMargin + assertRuleNotFired(ciNeqSql) + val (neqOn, neqOff) = runBoth(ciNeqSql) + assert(neqOn == neqOff, s"collated-neq semantics diverge: ON=$neqOn OFF=$neqOff") + + // Negative 2: only the EQUI key is non-binary collated -> equi-side type gate rejects. + createTable( + "TStrCiEqui", + "k STRING COLLATE UTF8_LCASE, v STRING", + """ ('a', 'x'), ('a', 'y'), + | ('b', 'z')""".stripMargin) + val ciEquiSql = + """SELECT k FROM TStrCiEqui outer_t WHERE k IN ( + | SELECT s1.k FROM TStrCiEqui s1 JOIN TStrCiEqui s2 + | ON s1.k = s2.k AND s1.v <> s2.v)""".stripMargin + assertRuleNotFired(ciEquiSql) + val (equiOn, equiOff) = runBoth(ciEquiSql) + assert(equiOn == equiOff, s"collated-equi semantics diverge: ON=$equiOn OFF=$equiOff") + } + } + + test("CHAR/VARCHAR join keys are rejected via declared-type metadata") { + withTable("TStr", "TCharVarchar") { + // The guard under test is isSafeComparisonGroupingAttribute: CHAR/VARCHAR table columns reach + // the optimizer as annotated StringType (CharVarcharUtils records the declared type in the + // attribute metadata), so a dataType-only check would admit them through the StringType + // branch. The rule recovers the declared raw type from the metadata and fails closed. + // Control: a plain STRING table -- identical query shape and data -- fires, proving strings + // are not rejected wholesale; the CHAR(5) and VARCHAR(5) variants, differing only in the + // declared column type, do not. Both the equi key `k` and the neq column `v` are + // CHAR/VARCHAR, so both gate paths are pinned. + createTable( + "TStr", + "k STRING, v STRING", + """ ('a', 'x'), ('a', 'y'), + | ('b', 'z')""".stripMargin) + val stringSql = + """SELECT k FROM TStr outer_t WHERE k IN ( + | SELECT s1.k FROM TStr s1 JOIN TStr s2 + | ON s1.k = s2.k AND s1.v <> s2.v)""".stripMargin + assertRuleFired(stringSql) + val (strOn, strOff) = runBoth(stringSql) + assert(strOn == strOff, s"string control diverges: ON=$strOn OFF=$strOff") + assert(strOn == Set(Row("a")), s"string control expected {a}, got $strOn") + + Seq("CHAR(5)", "VARCHAR(5)").foreach { keyType => + createTable( + "TCharVarchar", + s"k $keyType, v $keyType", + """ ('a', 'x'), ('a', 'y'), + | ('b', 'z')""".stripMargin) + val sql = + """SELECT k FROM TCharVarchar outer_t WHERE k IN ( + | SELECT s1.k FROM TCharVarchar s1 JOIN TCharVarchar s2 + | ON s1.k = s2.k AND s1.v <> s2.v)""".stripMargin + assertRuleNotFired(sql) + val (on, off) = runBoth(sql) + assert(on == off, s"$keyType not-fired query diverges: ON=$on OFF=$off") + } + } + } + + test("ANSI: rewrite preserves observable error behavior (throw-or-succeed parity)") { + withTable("TAnsi") { + // The guard under test is not a rejection but a parity property: the expression allowlist + // admits Cast, which can throw under ANSI. The rewrite evaluates the projected `CAST(s AS + // INT)` once per row inside the Aggregate, while the baseline self-join evaluates it per row + // on each side -- the same set of rows either way -- so a malformed value must make BOTH + // forms behave the same. k=2 holds a non-numeric 's'. + // + // Parity is checked as observable behavior, not just "an exception happened": both succeed + // with equal rows, or both throw with the same error class. A one-sided throw is a blocker. + createTable( + "TAnsi", + "k INT, s STRING", + """ (1, '10'), (1, '20'), + | (2, '30'), (2, 'xyz')""".stripMargin) + val sql = + """SELECT k FROM TAnsi outer_t WHERE k IN ( + | SELECT s1.k + | FROM (SELECT k, CAST(s AS INT) AS x FROM TAnsi) s1 + | JOIN (SELECT k, CAST(s AS INT) AS x FROM TAnsi) s2 + | ON s1.k = s2.k AND s1.x <> s2.x)""".stripMargin + + Seq("false", "true").foreach { ansi => + withSQLConf(SQLConf.ANSI_ENABLED.key -> ansi) { + // The rule still fires at plan level regardless of ANSI (the cast throws only at + // runtime). + assertRuleFired(sql) + val on = runOutcome(sql, rewrite = true) + val off = runOutcome(sql, rewrite = false) + (on, off) match { + case (Right(onRows), Right(offRows)) => + assert(onRows == offRows, + s"ANSI=$ansi both succeeded but diverged: ON=$onRows OFF=$offRows") + // Concrete positive signal under ANSI off: the cast is defined (yields NULL, does not + // throw) for every surviving row, so the rewrite must return the real membership {1}, + // not merely agree with OFF. + if (ansi == "false") { + assert(onRows == Set(Row(1)), s"ANSI=false expected {1}, got $onRows") + } + case (Left(onErr), Left(offErr)) => + assert(onErr == offErr, + s"ANSI=$ansi both threw but different error: ON=$onErr OFF=$offErr") + if (ansi == "true") { + assert(onErr == "CAST_INVALID_INPUT", + s"ANSI=true expected CAST_INVALID_INPUT, got $onErr") + } + case _ => + fail(s"ANSI=$ansi one-sided error behavior: ON=$on OFF=$off") + } + } + } + } + } + + test("ANSI: Remainder neq column preserves error behavior; Divide is type-gated") { + withTable("TAnsiDiv") { + // Cast is not the only allowlisted expression that can throw under ANSI. This pins the two + // arithmetic ops the review asked about: + // - Remainder (`%`) on INT operands stays INT, an allowlisted type, so the rule fires. It + // can raise DIVIDE_BY_ZERO under ANSI, so it exercises the throw-parity property directly, Review Comment: **Nit (P3):** This comment names DIVIDE_BY_ZERO, but the `% 0` assertion below and Spark's Remainder expression use REMAINDER_BY_ZERO. Please update the comment to match the error class the test actually enforces. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/RewriteSelfJoinInequalityToAggregate.scala: ########## @@ -0,0 +1,569 @@ +/* + * 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 supported uncorrelated IN-subquery inequality self-joins into + * `GROUP BY + HAVING MIN(neq) <> MAX(neq)`, avoiding the self-join cross-product. + * + * Supports a direct self-join (Pattern A') and a self-join nested under an outer inner join + * (Pattern A2, where only the self-join child becomes an Aggregate). Unsupported and correlated + * shapes fail closed. + * + * Runs in `extendedOperatorOptimizationRules`, before `RewritePredicateSubquery` turns the + * predicate subquery into a semi/anti/existence join, so it only sees the uncorrelated + * `InSubquery` shape. + * + * Controlled by `spark.sql.optimizer.rewriteSelfJoinInequalityToAggregate.enabled` + * (default false, opt-in). + */ +object RewriteSelfJoinInequalityToAggregate extends Rule[LogicalPlan] with PredicateHelper { + + private val MinNeqAliasName = "_rewrite_selfjoin_inequality_min" + private val MaxNeqAliasName = "_rewrite_selfjoin_inequality_max" + + override def apply(plan: LogicalPlan): LogicalPlan = { + if (!conf.getConf(SQLConf.REWRITE_SELF_JOIN_INEQUALITY_TO_AGGREGATE_ENABLED)) { + return plan + } + + // Fail closed on correlated subqueries: `lq.children` holds the outer references this rule + // does not remap. + 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 + } + } + } + + // ============================================================================ + // Shared helpers + // ============================================================================ + + /** + * Build `Filter(min <> max, Aggregate(equiKeys, child))`, taking MIN and MAX over the neq column. + * `MIN(neqCol) <> MAX(neqCol)` is true exactly when the group holds two or more distinct non-null + * values -- the same test as `COUNT(DISTINCT neqCol) > 1`, but avoids the distinct-dedup + * aggregation stages and supports partial aggregation. + * + * The `IsNotNull(equiKeys)` filter preserves the equi-join's NULL semantics: `=` never matches a + * NULL key, but GROUP BY would fold all NULL keys into one group that can leak NULL into a + * `NOT IN`. The neq column needs no filter -- MIN/MAX ignore NULL, and a group with fewer than + * two non-null values has `min = max` (or both NULL, which makes `<>` NULL), so `<>` is never + * true for it and the group is dropped. + */ + private def buildAggregateHavingMultipleDistinct( + equiKeys: Seq[Attribute], + neqCol: Attribute, + child: LogicalPlan): LogicalPlan = { + val minAlias = Alias(Min(neqCol).toAggregateExpression(), MinNeqAliasName)() + val maxAlias = Alias(Max(neqCol).toAggregateExpression(), MaxNeqAliasName)() + val aggExprs: Seq[NamedExpression] = equiKeys :+ minAlias :+ maxAlias + val nonNullChild = equiKeys + .map(a => IsNotNull(a): Expression) + .reduceOption(And) + .map(Filter(_, child)) + .getOrElse(child) + val agg = Aggregate(equiKeys, aggExprs, nonNullChild) + Filter(Not(EqualTo(minAlias.toAttribute, maxAlias.toAttribute)), agg) + } + + /** + * Rebuild the wrapper Project so every equi-key reference points at the sjLeft attribute with a + * fresh output ExprId, returning `oldOutputExprId -> newOutputAttr` for downstream references + * (outer join condition, top-level Project). Lookup is by ExprId (Catalyst attribute identity), + * not name. Fails closed when an entry is neither an equi-key Attribute nor `Alias(equi-key, _)`. + */ + private def canonicalizeWrapper( + projectList: Seq[NamedExpression], + equiPairs: Seq[(Attribute, Attribute)], + newChild: LogicalPlan): Option[(Project, Map[ExprId, 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) => + // Fresh exprId, but carry over qualifier / metadata so this branch stays consistent with + // the Alias branch and a column keeps its metadata. + Some( + Alias(exprIdToLeft(a.exprId), a.name)( + qualifier = a.qualifier, + explicitMetadata = Some(a.metadata)): NamedExpression) + case al @ Alias(a: Attribute, _) if exprIdToLeft.contains(a.exprId) => + // withNewChild preserves name/qualifier/metadata and exprId; newInstance then re-stamps a + // fresh exprId, so Alias keeps ownership of its own metadata contract instead of us + // re-listing its fields (which drift when Alias gains one). + Some(al.withNewChild(exprIdToLeft(a.exprId)).newInstance()) + case _ => None + } + if (mapped.exists(_.isEmpty)) { + None + } else { + val newProjectList = mapped.flatten + val newWrapper = Project(newProjectList, newChild) + val remap: Map[ExprId, Attribute] = + oldOutput.zip(newWrapper.output).map { case (o, n) => o.exprId -> n }.toMap + Some((newWrapper, remap)) + } + } + + /** + * Replace equi-key references inside a NamedExpression per `remap`, preserving Attribute/Alias + * shape. Any other expression still referencing a replaced output returns None (fail-closed) to + * avoid a dangling ExprId. + */ + 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) + } + // withNewChild preserves the same exprId/qualifier/metadata the manual copy did. + Some(if (newChild eq al.child) al else al.withNewChild(newChild)) + case other if other.references.exists(a => remap.contains(a.exprId)) => + None + case other => Some(other) + } + + // ============================================================================ + // Pattern A' / A2 dispatch (subquery plans of InSubquery) + // ============================================================================ + + private def rewriteSubqueryPlan(plan: LogicalPlan): Option[LogicalPlan] = { + // Match the candidate shape first -- a top-level Inner Join, optionally under one wrapper + // Project. This structural match is cheap, so run it before the whole-subquery + // `isRepeatablePlan` walk and skip that walk entirely for the many subqueries that are not + // even shaped like a self-join. + val (projectListOpt, innerJoin): (Option[Seq[NamedExpression]], Join) = plan match { + case Project(pl, j: Join) if j.joinType == Inner && j.condition.isDefined => + (Some(pl), j) + case j: Join if j.joinType == Inner && j.condition.isDefined => + (None, j) + case _ => return None + } + + // Candidate-level guard: reject if any node in the whole subquery is non-repeatable, catching + // nondeterminism hoisted above the self-join that the per-side `isSameBaseRelation` misses. + if (!isRepeatablePlan(plan)) return None + + if (isSameBaseRelation(innerJoin.left, innerJoin.right)) { + rewriteDirectSelfJoin(projectListOpt, innerJoin) + } else { + rewriteNestedSelfJoin(projectListOpt, innerJoin) + } + } + + // ============================================================================ + // Pattern A' : direct self-join at subquery top level + // ============================================================================ + + private def rewriteDirectSelfJoin( + projectListOpt: Option[Seq[NamedExpression]], + innerJoin: Join): Option[LogicalPlan] = { + // Fail closed on an explicit join hint: it is a directive about the join this rule deletes. + if (!innerJoin.hint.isEmpty) return None + + val innerLeft = innerJoin.left + val innerCond = innerJoin.condition.get + + val parsed = parseSelfJoinCondition(innerCond, innerLeft, innerJoin.right) + if (parsed.isEmpty) return None + val (equiPairs, neqPairs) = parsed.get + + val innerLeftEquiAttrs: Seq[Attribute] = equiPairs.map(_._1) + val innerLeftNeqAttr: Attribute = neqPairs.head._1 + val filtered = + buildAggregateHavingMultipleDistinct(innerLeftEquiAttrs, innerLeftNeqAttr, innerLeft) + + // Fail closed on a bare-Join subquery: with no wrapper Project, replacing the self-join output + // with `Project(equiKeys, filtered)` shrinks arity and RewritePredicateSubquery's positional + // `values.zip(sub.output)` would misbind semi predicates. Q95 subqueries always have a Project. + projectListOpt match { + case None => None + case Some(pl) => + canonicalizeWrapper(pl, equiPairs, filtered).map { case (newWrapper, _) => newWrapper } + } + } + + // ============================================================================ + // Pattern A2 : self-join nested inside another InnerJoin in the subquery + // ============================================================================ + + private def rewriteNestedSelfJoin( + projectListOpt: Option[Seq[NamedExpression]], + outerJoin: Join): Option[LogicalPlan] = { + val outerCond = outerJoin.condition.get + + val (selfJoinSide, selfJoinOnRight) = + tryExtractSelfJoin(outerJoin.right) match { + case Some(_) => (outerJoin.right, true) + case None => + tryExtractSelfJoin(outerJoin.left) match { + case Some(_) => (outerJoin.left, false) + case None => return None + } + } + + val (selfJoinProjectOpt, selfJoin) = selfJoinSide match { + case p @ Project(_, j: Join) if j.joinType == Inner && j.condition.isDefined => + (Some(p), j) + case j: Join if j.joinType == Inner && j.condition.isDefined => + (None, j) + case _ => return None + } + + val sjLeft = selfJoin.left + val sjCond = selfJoin.condition.get + if (!isSameBaseRelation(sjLeft, selfJoin.right)) return None + + val parsed = parseSelfJoinCondition(sjCond, sjLeft, selfJoin.right) + if (parsed.isEmpty) return None + val (equiPairs, neqPairs) = parsed.get + + val sjLeftEquiAttrs: Seq[Attribute] = equiPairs.map(_._1) + val sjLeftNeqAttr: Attribute = neqPairs.head._1 + + val selfJoinOutputSet = selfJoinSide.outputSet + val sjEquiExprIds: Set[ExprId] = + equiPairs.flatMap { case (l, r) => Seq(l.exprId, r.exprId) }.toSet + // A wrapper Project may reproject equi-keys under fresh alias exprIds; include those. + val wrapperEquiExprIds: Set[ExprId] = selfJoinProjectOpt.toSeq.flatMap { p => + p.projectList.flatMap { + case a: Attribute if sjEquiExprIds.contains(a.exprId) => Some(a.exprId) + case al @ Alias(a: Attribute, _) if sjEquiExprIds.contains(a.exprId) => Some(al.exprId) + case _ => None + } + }.toSet + val allEquiExprIds = sjEquiExprIds ++ wrapperEquiExprIds + + // The outer join condition and any top-level Project may reference only equi-key attrs from the + // self-join side (the neq column does not survive the rewrite). + val outerCondRefs = outerCond.references.filter(selfJoinOutputSet.contains) + if (!outerCondRefs.forall(a => allEquiExprIds.contains(a.exprId))) return None + val projectOk = projectListOpt.forall { pl => + val refs = pl.flatMap(_.references).filter(selfJoinOutputSet.contains) + refs.forall(a => allEquiExprIds.contains(a.exprId)) + } + if (!projectOk) return None + + val filtered = buildAggregateHavingMultipleDistinct(sjLeftEquiAttrs, sjLeftNeqAttr, sjLeft) + + val (newSelfJoinSide, outputRemap): (LogicalPlan, Map[ExprId, Attribute]) = + selfJoinProjectOpt match { + case Some(wp) => + canonicalizeWrapper(wp.projectList, equiPairs, filtered) match { + case Some((newWrapper, remap)) => (newWrapper, remap) + case None => return None + } + case None if projectListOpt.isEmpty => + // Fail closed: with no wrapper and no top-level Project, `Project(equiKeys, filtered)` + // shrinks the outer join's arity and RewritePredicateSubquery's positional zip misbinds. + return None + case None => + // Top-level Project preserves arity via `outputRemap`; remap sjRight equi-refs to sjLeft + // (same output position in a valid self-join). + val newP = Project(sjLeftEquiAttrs, filtered) + val remap: Map[ExprId, Attribute] = equiPairs.map { case (l, r) => r.exprId -> l }.toMap + (newP, remap) + } + + val newOuterCond = outerCond.transformUp { + case a: Attribute if outputRemap.contains(a.exprId) => outputRemap(a.exprId) + } + + val newOuterJoin = if (selfJoinOnRight) { + outerJoin.copy(right = newSelfJoinSide, condition = Some(newOuterCond)) + } else { + outerJoin.copy(left = newSelfJoinSide, condition = Some(newOuterCond)) + } + + projectListOpt match { + case Some(pl) => + val remapped = pl.map(ne => remapNamedExpressionAttributes(ne, outputRemap)) + if (remapped.exists(_.isEmpty)) return None + Some(Project(remapped.flatten, newOuterJoin)) + case None => Some(newOuterJoin) + } + } + + private def tryExtractSelfJoin(plan: LogicalPlan): Option[Join] = { + val join = plan match { + case Project(_, j: Join) if j.joinType == Inner && j.condition.isDefined => j + case j: Join if j.joinType == Inner && j.condition.isDefined => j + case _ => return None + } + // A hinted self-join is not an extraction candidate; see `rewriteDirectSelfJoin`. + if (!join.hint.isEmpty) return None + if (!isSameBaseRelation(join.left, join.right)) return None + if (parseSelfJoinCondition(join.condition.get, join.left, join.right).isEmpty) return None + Some(join) + } + + // ============================================================================ + // parseSelfJoinCondition + isSameBaseRelation + // ============================================================================ + + private def outputOrdinal(plan: LogicalPlan, attr: Attribute): Int = + plan.output.indexWhere(_.exprId == attr.exprId) + + private def sameOutputPosition( + leftPlan: LogicalPlan, + rightPlan: LogicalPlan, + leftAttr: Attribute, + rightAttr: Attribute): Boolean = { + val leftPos = outputOrdinal(leftPlan, leftAttr) + val rightPos = outputOrdinal(rightPlan, rightAttr) + leftPos >= 0 && rightPos >= 0 && leftPos == rightPos + } + + /** + * Parse a join condition into equi-pairs and inequality-pairs. Accepts only `EqualTo(attr, attr)` + * and `Not(EqualTo(attr, attr))` across opposite sides, and `IsNotNull(attr)` on a join column; + * anything else fails the whole rewrite closed. + */ + private def parseSelfJoinCondition( + condition: Expression, + leftPlan: LogicalPlan, + rightPlan: LogicalPlan) + : Option[(Seq[(Attribute, Attribute)], Seq[(Attribute, Attribute)])] = { + + val leftOutput = leftPlan.outputSet + val rightOutput = rightPlan.outputSet + val predicates = splitConjunctivePredicates(condition) + + val equiPairs = predicates.collect { + case EqualTo(l: Attribute, r: Attribute) + if leftOutput.contains(l) && rightOutput.contains(r) => + (l, r) + case EqualTo(r: Attribute, l: Attribute) + if leftOutput.contains(l) && rightOutput.contains(r) => + (l, r) + } + + val neqPairs = predicates.collect { + case Not(EqualTo(l: Attribute, r: Attribute)) + if leftOutput.contains(l) && rightOutput.contains(r) => + (l, r) + case Not(EqualTo(r: Attribute, l: Attribute)) + if leftOutput.contains(l) && rightOutput.contains(r) => + (l, r) + } + + // Only IsNotNull on a join column is safe to drop -- redundant with the join or auto-added by + // InferFiltersFromConstraints. IsNotNull on any other column changes semantics; bail out. + val joinAttrIds: Set[ExprId] = + (equiPairs ++ neqPairs).flatMap { case (l, r) => Seq(l.exprId, r.exprId) }.toSet + val isNotNullOnJoinCols = predicates.count { + case IsNotNull(a: Attribute) if joinAttrIds.contains(a.exprId) => true + case _ => false + } + + val totalMatched = equiPairs.size + neqPairs.size + isNotNullOnJoinCols + if (totalMatched != predicates.size) return None + if (equiPairs.isEmpty || neqPairs.isEmpty) return None + + // A single inequality only: MIN/MAX over one column cannot represent multiple neqs. + if (neqPairs.size != 1) return None + + // The rewrite swaps comparison equality (`=`/`<>`) for grouping and MIN/MAX ordering equality, + // so gate every equi-key and the neq column -- both ends of each pair, since canonicalization + // drops the metadata that `isSafeComparisonGroupingAttribute` reads and the two ends may differ + // -- on a positive type allowlist. Fail closed on anything not proven safe. + val keyAttrs = (equiPairs ++ neqPairs).flatMap { case (l, r) => Seq(l, r) } + if (!keyAttrs.forall(isSafeComparisonGroupingAttribute)) return None + + // The rewrite expresses "two or more distinct values" as MIN(neq) <> MAX(neq), so the neq + // column must be orderable. The allowlist above already implies this, but assert Spark's own + // MIN/MAX input contract (RowOrdering.isOrderable, the same check Min/Max run) explicitly, so + // the requirement is visible at the rewrite site. + if (!RowOrdering.isOrderable(neqPairs.head._1.dataType)) return None + + // Resolve each predicate end by ExprId and require matching output ordinals, not name equality + // (canonicalization erases cosmetic Alias names). + val equiValid = + equiPairs.forall { case (l, r) => sameOutputPosition(leftPlan, rightPlan, l, r) } + val neqValid = neqPairs.forall { case (l, r) => sameOutputPosition(leftPlan, rightPlan, l, r) } + if (!equiValid || !neqValid) return None + + // Equi-key output positions must be distinct, so swapped/duplicate aliases cannot collide. + val leftEquiOrdinals = equiPairs.map { case (l, _) => outputOrdinal(leftPlan, l) } + if (leftEquiOrdinals.exists(_ < 0)) return None + if (leftEquiOrdinals.distinct.size != leftEquiOrdinals.size) return None + + // Reject when the neq column overlaps an equi-key column (e.g. `t1.k = t2.k AND t1.k <> t2.k`). + val neqLeftOrdinal = outputOrdinal(leftPlan, neqPairs.head._1) + if (neqLeftOrdinal < 0 || leftEquiOrdinals.contains(neqLeftOrdinal)) return None + Some((equiPairs, neqPairs)) + } + + /** + * Type gate applied to each equi-key and the neq column. CHAR/VARCHAR reach the optimizer as + * StringType with the declared type recorded in the attribute metadata, so recover the raw type + * from metadata (falling back to `dataType`) before running the datatype allowlist -- otherwise + * they would slip through the StringType branch. + */ + private def isSafeComparisonGroupingAttribute(attr: Attribute): Boolean = { + val rawType = CharVarcharUtils.getRawType(attr.metadata).getOrElse(attr.dataType) + isSafeComparisonGroupingType(rawType) + } + + /** + * Positive allowlist of types where comparison equality (`=`/`<>`) provably coincides with + * grouping and MIN/MAX ordering equality, so a key can move into GROUP BY / MIN-MAX. Float/Double + * (NaN, signed zero), CHAR/VARCHAR (declared-type/padding), non-binary collated strings, complex + * types, UDTs / Variant and unknown types fail closed. + */ + private def isSafeComparisonGroupingType(dt: DataType): Boolean = dt match { + case ByteType | ShortType | IntegerType | LongType => true + case _: DecimalType => true + case BooleanType => true + case DateType => true + case TimestampType | TimestampNTZType => true + case BinaryType => true + case _: CharType | _: VarcharType => false + case st: StringType if st.supportsBinaryEquality => true + case _ => false + } + + /** + * Primary safety guard: the rewrite folds two occurrences of one subtree into a single aggregate, + * so a plan qualifies only when its operators, leaves and expressions are all allowlisted as + * repeatable. `plan.deterministic` alone is insufficient -- Aggregate(First), Window row_number + * over a non-total order and Limit/Sample are row-bag nondeterministic yet report deterministic. + * Embedded expression subqueries also fail closed. + */ + private def isRepeatablePlan(plan: LogicalPlan): Boolean = { + plan.deterministic && + !plan.isStreaming && + plan.subqueriesAll.isEmpty && + isRowBagRepeatable(plan) && + hasRepeatableExpressions(plan) + } + + /** + * Operator/leaf allowlist for repeatable row bags; everything unknown fails closed. Kept narrow: + * the target shape needs only a Parquet scan optionally wrapped in Project / Filter / + * SubqueryAlias plus the self-join. Row ordering is irrelevant to the row-bag contract. + * + * The narrowness is deliberate, not a correctness requirement, but it stays intentionally + * conservative: the exact-`ParquetFileFormat` leaf check admits only stock Parquet scans. Other + * file formats (ORC, JSON, CSV) reach the same FileSourceScan but remain rejected until each is + * separately validated, and likewise for inverting the allowlist into a leaf blocklist. This + * keeps the rule fail-closed on any leaf not proven repeatable. + */ + private def isRowBagRepeatable(plan: LogicalPlan): Boolean = !plan.exists { + // Whitelisted operator => false ("does not break repeatability"); negating `exists` then means + // "every operator is whitelisted". + case _: Project => false + case _: Filter => false + case _: SubqueryAlias => false + case _: Join => false + case _: Range => false + case _: LocalRelation => false + case relation: LogicalRelation => + // Trust a Parquet scan only: exact `ParquetFileFormat` (getClass, not isInstanceOf, since it + // is non-final); any other FileFormat is not provably repeatable. + relation.relation match { + case h: HadoopFsRelation if h.fileFormat.getClass == classOf[ParquetFileFormat] => false + case _ => true + } + case _ => true + } + + private def hasRepeatableExpressions(plan: LogicalPlan): Boolean = { + !plan.exists(node => node.expressions.exists(expr => !isRepeatableExpression(expr))) + } + + /** + * Expression allowlist: repeatable only when the root type is allowlisted and all children are, + * so `Add(v, Abs(w))` is rejected. Unknown types fail closed (a missed optimization, not a bug). + * Decimal wrappers such as `PromotePrecision` / `CheckOverflow` are absent and may fail closed. + */ + private def isRepeatableExpression(expr: Expression): Boolean = expr match { + case _: Attribute | _: Literal => + true + // A STRING -> TIMESTAMP_LTZ cast (micro or nanosecond precision) is not repeatable: for a + // time-only string SparkDateTimeUtils fills the missing date from LocalDate.now(zoneId), which + // ComputeCurrentTime does not stabilize on this path, so two scans that straddle midnight can + // produce different timestamps while the rewrite folds them into a single evaluation. The NTZ + // parse is clock-independent (it returns null for a time-only string), so only the LTZ + // directions are rejected. Reject whenever such a conversion appears at any nesting level, + // including inside array/map/struct element casts -- e.g. CAST(CAST(ss AS ARRAY<TIMESTAMP>) AS + // STRING). Each nested Cast node is itself visited here, so a per-node check catches the inner + // conversion. Fail closed. + case c: Cast if castHasClockDependentStringToTimestamp(c.child.dataType, c.dataType) => + false + case _: Alias | _: Cast | _: Add | _: Subtract | _: Multiply | _: Divide | _: Remainder | + _: And | _: Or | _: Not | _: EqualTo | _: EqualNullSafe | _: LessThan | + _: LessThanOrEqual | _: GreaterThan | _: GreaterThanOrEqual | _: IsNull | _: IsNotNull => + expr.children.forall(isRepeatableExpression) + case _ => + false + } + + /** + * True when casting `from` to `to` performs a clock-dependent STRING -> TIMESTAMP_LTZ conversion + * (microsecond or nanosecond precision) at any nesting level: directly, or inside matching + * array / map / struct element casts. Such a cast supplies a time-only string's missing date from + * the runtime clock, so it is not repeatable; see [[isRepeatableExpression]]. The scalar + * String-source decision is delegated to [[Cast.needsTimeZone]], which enumerates exactly the + * zone-dependent (hence LTZ, hence clock-dependent for a bare time) String conversions and stays + * in sync as Spark adds timestamp types; String -> TIMESTAMP_NTZ is clock-independent and absent + * there. + */ + private def castHasClockDependentStringToTimestamp(from: DataType, to: DataType): Boolean = + (from, to) match { + case (s: StringType, t) => Cast.needsTimeZone(s, t) Review Comment: **Non-blocking (P2):** The clock-dependence check follows declared StringType shapes, so VariantType reaches the fallback even though VariantGet can unwrap a runtime string and invoke the same string-to-TIMESTAMP_LTZ parser. Parquet Variant input and TimestampType output both pass the other guards; if the two original scans evaluate on opposite sides of midnight, the rewrite's single evaluation can change IN/NOT IN membership. Please reject Variant casts whose target is or recursively contains TIMESTAMP_LTZ, while retaining a non-LTZ positive control, and cover scalar plus nested Variant targets. **Recommended change:** Extend clock-dependent-cast classification for Variant sources and add focused optimizer regression coverage for scalar and nested Variant-to-LTZ targets. **Why this works:** When the declared source is VariantType, reject the Cast whenever its target is or recursively contains a TIMESTAMP_LTZ type, because a runtime Variant string will be parsed by the ordinary clock-dependent string-to-LTZ path. Preserve the existing precise StringType structural recursion for non-Variant sources and keep Variant-to-TIMESTAMP_NTZ eligible. **Scope:** sql/core/src/main/scala/org/apache/spark/sql/execution, sql/core/src/test/scala/org/apache/spark/sql/execution **Compatibility:** All currently supported repeatable scalar and structurally nested casts, especially TIMESTAMP_NTZ controls, continue to rewrite. **Risks:** Over-broad Variant rejection would exclude repeatable Variant-to-numeric, date, string, or TIMESTAMP_NTZ projections. Checking only scalar TimestampType would leave Array, Map, Struct, or nanosecond LTZ targets exposed. **Constraints:** Do not change the public CAST or VariantGet contract. Treat unknown or newly introduced nested target shapes conservatively. Retain the current positive control for clock-independent TIMESTAMP_NTZ parsing. **Success:** A Variant runtime string cannot reach a timestamp-LTZ target in a plan accepted as repeatable by this rule. Nested Variant targets containing LTZ at any supported array, map, or struct position are rejected. Variant casts whose target contains no LTZ remain eligible when all other guards 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]
