LuciferYang commented on code in PR #58424: URL: https://github.com/apache/spark/pull/58424#discussion_r3998997760
########## 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 => Review Comment: `rewriteSubqueryPlan` accepts at most one `Project` above the self-join (:170), while Q95's subquery plus the CTE's own projection is two, and the inner one carries two columns that are not equi keys. It matches today because `OptimizeSubqueries` runs the whole optimizer recursively over each subquery plan before the operator batches, so `ColumnPruning` has dropped those columns and `CollapseProject` has merged the two projections by the time this rule looks. Nothing covers that chain. The new suite never mentions `q95`, `ws_wh` or `web_sales`, and the config is set only inside that suite. If the conditions in `ColumnPruning` or `CollapseProject` shift, the subquery arrives with two projections, the rule returns None at :174, the benefit disappears and all 38 tests still pass. A Q95-shaped test would be the cheapest fix: the CTE plus both IN subqueries with the config on, asserting the min/max signature aliases appear for each. Those two subqueries cover Pattern A' and Pattern A2 respectively. ########## 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( Review Comment: The comment above `buildAggregateHavingMultipleDistinct` argues NULL semantics carefully but leaves out two preconditions. First, the rewrite preserves the distinct set, not the row bag: the original self-join emits one row per qualifying pair, the aggregate one row per key. That is sound only because the rule fires exclusively under `InSubquery` (:60), which `RewritePredicateSubquery` turns into a semi/anti join where multiplicity is unobservable. Second, an output column name can change. `canonicalizeWrapper` preserves names; the branch that does not is the A2 path with no wrapper `Project` on the self-join side (:294-299), which rewrites the top-level `Project`'s references to right-side attributes into the left-side ones (:147), so the name changes whenever the two aliases differ. That is harmless only because `InSubquery` binds positionally and compares types with `equalsStructurally`, which ignores field names. Two sentences in that comment would cover it. Whoever reuses this rule outside `InSubquery`, or attaches a consumer that reads columns by name, breaks one of the two silently. ########## 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 { Review Comment: The comment above the allowlist gives `(NaN, signed zero)` as the reason for excluding Float/Double, which reads as a claim that comparison equality and grouping equality diverge for floats. In Spark they agree: `genEqual` emits `(isNaN && isNaN) || c1 == c2` for both types (`CodeGenerator.scala:670`), `SQLOrderingUtil` matches it, and aggregate grouping keys go through `NormalizeFloatingNumbers.normalize` during physical planning, as the comment at `SparkStrategies.scala:596` explains. So `{0.0, -0.0}` and `{NaN, NaN}` produce no row either way, and `{1.0, NaN}` produces one either way. Excluding floats is fine in itself, it only loses rewrites. What is worth fixing is the wording: stating a conservative choice as a correctness requirement means nobody revisits it. The suite is more candid than the comment here, since the test at :987 calls itself defensive. ########## 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 Review Comment: The string branch of the type allowlist gates on `supportsBinaryEquality` (:456), but the rewrite replaces `<>` with `MIN`/`MAX`, so what it depends on is the ordering, and codegen's own byte fast path for equality keys off `supportsBinaryOrdering` (`CodeGenerator.scala:673`). The invariant in `CollationFactory` runs one way only: `:153-158` says binary ordering implies binary equality, not the reverse. Today the two fields are computed from the identical expression at `:208-209`, so they agree. If a collation ever supports binary equality with a non-binary ordering, say case-sensitive equality with case-insensitive ordering, a group `{'a', 'A'}` yields a row from the original `<>` join. After the rewrite `Least`/`Greatest` keep the value already accumulated when the ordering ties, so `MIN` and `MAX` both return the same value, `EqualTo(min, max)` is true, the group is dropped and the `IN` loses a match. Gating on `supportsBinaryOrdering` fixes it with no behavior change today. ########## 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") { Review Comment: Three groups of negative tests never reach the guard their name points at. In `Inequality column overlapping an equi-key is rejected` (:749), the two mutually negating conjuncts are folded away by `BooleanSimplification` while `OptimizeSubqueries` optimizes the subquery, so what reaches the rule no longer contains a self-join shape and it returns at :174, never reaching :425-427. In `Aggregate (first)` (:819) and `Window (row_number)` (:839), `first` and `row_number` are already outside the expression allowlist, so deleting the operator allowlist at :486-503 leaves both green. `Plain InnerJoin` (:533) and `LeftOuter join` (:705) contain no `InSubquery` at all, so the IN_SUBQUERY pruning means the rule never inspects that join. Meanwhile the guard with real correctness weight has no negative test: a `LEFT JOIN` self-join returns every left key, so `MIN <> MAX` is the wrong answer, yet relaxing `joinType == Inner` at :170 to admit LeftOuter keeps all 38 tests green. Worth making each group reach the guard it names, plus an Inner-only case. To reach :427, :749 needs the two negating conjuncts separated, since `BooleanSimplification` only matches an adjacent pair: a second equi key in between, as in `ON s1.k = s2.k AND s1.v = s2.v AND s1.k <> s2.k`. -- 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]
