LuciferYang commented on code in PR #12756: URL: https://github.com/apache/gluten/pull/12756#discussion_r3822503878
########## backends-velox/src/main/scala/org/apache/gluten/extension/RewriteSelfJoinInequalityToAggregate.scala: ########## @@ -0,0 +1,722 @@ +/* + * 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.gluten.extension + +import org.apache.gluten.config.VeloxConfig + +import org.apache.spark.internal.Logging +import org.apache.spark.sql.SparkSession +import org.apache.spark.sql.catalyst.expressions._ +import org.apache.spark.sql.catalyst.expressions.RowOrdering +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.types.LongType + +/** + * Rewrites self-join with inequality into GROUP BY + HAVING COUNT(DISTINCT) > 1. + * + * Targets three patterns; all require an existence-only context (LeftSemi/LeftAnti join, or + * InSubquery/Exists expression) so that row-count multiplicity from the self-join cross-product + * does not affect semantics. + * + * - Pattern A' (InSubquery/Exists primary): InSubquery/Exists whose subquery top-level join is a + * direct self-join. The primary path for TPC-DS Q95. + * - Pattern A2 (nested): InSubquery/Exists whose subquery contains an outer InnerJoin that has a + * self-join child. Only the self-join child is replaced with Aggregate; the outer join is + * preserved. + * - Pattern A (LeftSemi/LeftAnti): LeftSemi/LeftAnti whose right child is an Inner self-join + * (possibly wrapped in Project). Matches semi/anti joins that already exist in the input -- + * e.g. from an explicit `LEFT SEMI JOIN` clause. Note: this rule is injected via + * `injectOptimizerRule`, which places it in the operator-optimization batch that runs BEFORE + * `RewritePredicateSubquery`; A is NOT a post-subquery-rewrite fallback for A'. + * + * Correlated subqueries (outer references / joinCond in ListQuery/Exists) are fail-closed at the + * entry expression, since our ExprId canonicalization does not remap those predicates. + * + * All three share: + * - [[buildAggregateHavingDistinctGt1]] to construct `Filter(cnt > 1, Aggregate)` + * - [[canonicalizeWrapper]] to rebuild a wrapping Project so every equi-key reference points to + * the sjLeft-side attribute, with **fresh exprIds** (Spark's SPARK-21835 style -- no reuse of + * original exprIds), returning an old->new attribute remap for downstream rewrite. + * + * Controlled by `spark.gluten.sql.rewrite.selfJoinInequality` (default false, opt-in). + */ +case class RewriteSelfJoinInequalityToAggregate(spark: SparkSession) + extends Rule[LogicalPlan] + with PredicateHelper + with Logging { + + private val CountDistinctAliasName = "_gluten_rw_selfjoin_cnt_distinct" + + override def apply(plan: LogicalPlan): LogicalPlan = { + if (!VeloxConfig.get.enableRewriteSelfJoinInequality) { + logDebug("RewriteSelfJoinInequalityToAggregate: disabled via config, skipping") + return plan + } + + // Pattern A: rewrite LeftSemi/LeftAnti whose right child is an Inner self-join. + val afterOps = plan.transformUp { + case j: Join + if (j.joinType == LeftSemi || j.joinType == LeftAnti) && + j.condition.isDefined && + isInnerJoinShape(j.right) => + tryRewriteSemiWithSelfJoinChild(j).getOrElse(j) + case other => other + } + + // Pattern A' / A2: rewrite subquery plans embedded in InSubquery/Exists. + // Type-based matching (`x: T`) + named-argument copy keeps this portable across + // Spark 3.3/3.4/3.5/4.x where ListQuery/Exists case-class arity has drifted. + // + // Correlated subquery fail-closed: `SubqueryExpression.children.nonEmpty` iff the + // subquery has outer references / correlated join conditions. These predicates + // reference attributes INSIDE the subquery plan by ExprId; our canonicalizeWrapper + // rewrites those ExprIds without remapping the correlated predicates, which would + // leave dangling references after `RewritePredicateSubquery` folds them back into + // the semi-join condition. Target workload (TPC-DS Q95) is uncorrelated, so bail + // on any correlated candidate rather than growing the remap surface. + val rewritten = afterOps.transformAllExpressions { + 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 + } + case ex: Exists if ex.children.isEmpty => Review Comment: The `case ex: Exists` at line 100 is unreachable. An uncorrelated `Exists` is replaced by `IsNotNull(ScalarSubquery(Limit 1, Project(1, plan)))` by `RewriteNonCorrelatedExists` in `Batch("Finish Analysis")`, far ahead of the operator-optimization batch where injected rules live, so they never see it. A correlated `Exists` is blocked by the `children.isEmpty` guard on the same line. The branch is therefore dead code on 3.3 through 4.1, checked version by version, and the `Exists` mentioned in the config doc, in `docs/velox-configuration.md` and in the class comment never applies. The earlier comment about the literal projection describes the same symptom, but that layer is never reached. Real EXISTS support means matching the rewritten `ScalarSubquery(Limit 1, ...)` shape, and `Limit` is itself rejected by `isRowBagRepeatable`. If that is out of scope, drop the branch and remove `Exists` from the config and docs wording. ########## backends-velox/src/main/scala/org/apache/gluten/extension/RewriteSelfJoinInequalityToAggregate.scala: ########## @@ -0,0 +1,722 @@ +/* + * 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.gluten.extension + +import org.apache.gluten.config.VeloxConfig + +import org.apache.spark.internal.Logging +import org.apache.spark.sql.SparkSession +import org.apache.spark.sql.catalyst.expressions._ +import org.apache.spark.sql.catalyst.expressions.RowOrdering +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.types.LongType + +/** + * Rewrites self-join with inequality into GROUP BY + HAVING COUNT(DISTINCT) > 1. + * + * Targets three patterns; all require an existence-only context (LeftSemi/LeftAnti join, or + * InSubquery/Exists expression) so that row-count multiplicity from the self-join cross-product + * does not affect semantics. + * + * - Pattern A' (InSubquery/Exists primary): InSubquery/Exists whose subquery top-level join is a + * direct self-join. The primary path for TPC-DS Q95. + * - Pattern A2 (nested): InSubquery/Exists whose subquery contains an outer InnerJoin that has a + * self-join child. Only the self-join child is replaced with Aggregate; the outer join is + * preserved. + * - Pattern A (LeftSemi/LeftAnti): LeftSemi/LeftAnti whose right child is an Inner self-join + * (possibly wrapped in Project). Matches semi/anti joins that already exist in the input -- + * e.g. from an explicit `LEFT SEMI JOIN` clause. Note: this rule is injected via + * `injectOptimizerRule`, which places it in the operator-optimization batch that runs BEFORE + * `RewritePredicateSubquery`; A is NOT a post-subquery-rewrite fallback for A'. + * + * Correlated subqueries (outer references / joinCond in ListQuery/Exists) are fail-closed at the + * entry expression, since our ExprId canonicalization does not remap those predicates. + * + * All three share: + * - [[buildAggregateHavingDistinctGt1]] to construct `Filter(cnt > 1, Aggregate)` + * - [[canonicalizeWrapper]] to rebuild a wrapping Project so every equi-key reference points to + * the sjLeft-side attribute, with **fresh exprIds** (Spark's SPARK-21835 style -- no reuse of + * original exprIds), returning an old->new attribute remap for downstream rewrite. + * + * Controlled by `spark.gluten.sql.rewrite.selfJoinInequality` (default false, opt-in). + */ +case class RewriteSelfJoinInequalityToAggregate(spark: SparkSession) + extends Rule[LogicalPlan] + with PredicateHelper + with Logging { + + private val CountDistinctAliasName = "_gluten_rw_selfjoin_cnt_distinct" + + override def apply(plan: LogicalPlan): LogicalPlan = { + if (!VeloxConfig.get.enableRewriteSelfJoinInequality) { + logDebug("RewriteSelfJoinInequalityToAggregate: disabled via config, skipping") + return plan + } + + // Pattern A: rewrite LeftSemi/LeftAnti whose right child is an Inner self-join. + val afterOps = plan.transformUp { + case j: Join + if (j.joinType == LeftSemi || j.joinType == LeftAnti) && + j.condition.isDefined && + isInnerJoinShape(j.right) => + tryRewriteSemiWithSelfJoinChild(j).getOrElse(j) + case other => other + } + + // Pattern A' / A2: rewrite subquery plans embedded in InSubquery/Exists. + // Type-based matching (`x: T`) + named-argument copy keeps this portable across + // Spark 3.3/3.4/3.5/4.x where ListQuery/Exists case-class arity has drifted. + // + // Correlated subquery fail-closed: `SubqueryExpression.children.nonEmpty` iff the + // subquery has outer references / correlated join conditions. These predicates + // reference attributes INSIDE the subquery plan by ExprId; our canonicalizeWrapper + // rewrites those ExprIds without remapping the correlated predicates, which would + // leave dangling references after `RewritePredicateSubquery` folds them back into + // the semi-join condition. Target workload (TPC-DS Q95) is uncorrelated, so bail + // on any correlated candidate rather than growing the remap surface. + val rewritten = afterOps.transformAllExpressions { + 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 + } + case ex: Exists if ex.children.isEmpty => + rewriteSubqueryPlan(ex.plan) match { + case Some(newSub) => ex.copy(plan = newSub) + case None => ex + } + } + if (!(rewritten eq plan)) { + logDebug( + "RewriteSelfJoinInequalityToAggregate: rewrote self-join to " + + "GROUP BY + HAVING COUNT(DISTINCT) > 1") + } + rewritten + } + + // ============================================================================ + // Shared helpers + // ============================================================================ + + private def isInnerJoinShape(plan: LogicalPlan): Boolean = plan match { + case Project(_, j: Join) if j.joinType == Inner && j.condition.isDefined => true + case j: Join if j.joinType == Inner && j.condition.isDefined => true + case _ => false + } + + /** + * Build `Filter(cnt > 1, Aggregate(equiKeys, [equiKeys, cnt_alias], Filter(IsNotNull(equiKeys), + * child)))`. Returns the Filter node whose output is `equiKeys ++ [count_alias_attr]`. + * + * The extra `IsNotNull(equiKeys)` filter is essential to preserve the original equi-join's NULL + * semantics. Under SQL 3VL, `left.k = right.k` never matches when either side is NULL, so the + * original self-join drops rows with NULL equi-keys. Aggregate, in contrast, groups NULL keys + * together into a single "NULL group" -- if that group has >= 2 distinct non-null neq values, + * COUNT(DISTINCT) > 1 fires and injects NULL into the subquery output. That leaked NULL then + * turns `NOT IN` into a spurious empty result (Spark's null-aware anti-join uses + * `Or(equi, IsNull(equi))` which any NULL sub-row satisfies) and can flip EXISTS/IN outcomes. The + * neq column needs no such filter: `COUNT(DISTINCT col)` already ignores NULL. + */ + private def buildAggregateHavingDistinctGt1( + equiKeys: Seq[Attribute], + neqCol: Attribute, + child: LogicalPlan): LogicalPlan = { + val countExpr = AggregateExpression( + Count(Seq(neqCol)), + mode = Complete, + isDistinct = true, + filter = None, + NamedExpression.newExprId) + val countAlias = Alias(countExpr, CountDistinctAliasName)() + // Seq[Attribute] is a Seq[NamedExpression] via covariance; no cast needed. + val aggExprs: Seq[NamedExpression] = equiKeys :+ countAlias + val nonNullChild = equiKeys + .map(a => IsNotNull(a): Expression) + .reduceOption(And) + .map(Filter(_, child)) + .getOrElse(child) + val agg = Aggregate(equiKeys, aggExprs, nonNullChild) + Filter(GreaterThan(countAlias.toAttribute, Literal(1L, LongType)), agg) + } + + /** + * Canonicalize a Project so every equi-key reference points at the sjLeft-side attribute (both + * sides of a valid self-join share names, so this substitution is semantically safe). Uses + * **fresh exprIds** (no reuse of original wrapper output exprIds) -- the same technique Spark's + * own `dedupSubqueryOnSelfJoin` uses when it needs to change subquery output. + * + * Returns the rebuilt Project and a map `oldWrapperOutputExprId -> newWrapperOutputAttr`, so + * downstream references (outer join condition, top-level Project) can be updated consistently. + * + * `equiPairs` provides the definitive ExprId-based lookup: `equiPair (l, r)` binds + * `l.exprId -> l` (identity) and `r.exprId -> l` (sjRight -> sjLeft). Attribute identity in + * Catalyst is ExprId, not name; two columns can share a name with distinct ExprIds. Name-based + * lookup would silently drop such entries via `.toMap`. + * + * Fails (returns None) when a projectList entry is neither an equi-key Attribute (by ExprId) nor + * `Alias(equi-key Attribute, _)`. Fail-closed. + */ + private def canonicalizeWrapper( + projectList: Seq[NamedExpression], + equiPairs: Seq[(Attribute, Attribute)], + newChild: LogicalPlan): Option[(Project, Map[ExprId, Attribute])] = { + // ExprId-based canonical map: any equi-key attribute (either side) -> sjLeft attribute. + val exprIdToLeft: Map[ExprId, Attribute] = + equiPairs.flatMap { case (l, r) => Seq(l.exprId -> l, r.exprId -> l) }.toMap + val oldOutput: Seq[Attribute] = projectList.map(_.toAttribute) + val mapped: Seq[Option[NamedExpression]] = projectList.map { + case a: Attribute if exprIdToLeft.contains(a.exprId) => + // Wrap every rewritten output slot in a fresh Alias. + // + // When a wrapper reprojects BOTH sides of the same equi pair (e.g. + // `SELECT s1.k, s2.k FROM T s1 JOIN T s2 ON s1.k = s2.k AND s1.v <> s2.v`), + // both entries collapse to the same sjLeft Attribute after the self-join is + // rewritten. Duplicate output ExprIds are not illegal in Spark (`SELECT a, a` + // is a valid Project), but fresh Aliases give each output slot an independent + // identity, which keeps the `oldOutput -> newOutput` remap 1-to-1 and lets + // downstream references (outer join condition, top-level Project) be updated + // unambiguously via ExprId. + // + // The fresh ExprId is on the Alias ITSELF; the referenced child keeps its + // original ExprId. Spark's logical-plan integrity checks reject reusing a + // referenced ExprId as the Alias's own ExprId, not duplication across slots. + Some(Alias(exprIdToLeft(a.exprId), a.name)(): NamedExpression) + case al @ Alias(a: Attribute, _) if exprIdToLeft.contains(a.exprId) => + // Fresh exprId; do NOT reuse `al.exprId`. Reusing another expression's exprId + // is the pattern that Spark 3.3 flags via structural-integrity checks. + Some(Alias(exprIdToLeft(a.exprId), al.name)(): NamedExpression) + case _ => None + } + if (mapped.exists(_.isEmpty)) { + None + } else { + val newProjectList = mapped.flatten + val newWrapper = Project(newProjectList, newChild) + val newOutput = newWrapper.output + val remap: Map[ExprId, Attribute] = + oldOutput.zip(newOutput).map { case (o, n) => o.exprId -> n }.toMap + Some((newWrapper, remap)) + } + } + + /** + * Replace equi-key attribute references inside a NamedExpression according to `remap`, while + * preserving the NamedExpression shape. + * + * `Expression.transformUp` returns `Expression`, not `NamedExpression`. We avoid a blanket + * `asInstanceOf[NamedExpression]` by handling the two shapes that can appear in a Project's + * `projectList` explicitly: a bare Attribute (whose top-level may itself be replaced) and an + * Alias (which stays an Alias while its child is transformed). Anything else in a projectList -- + * e.g. computed expressions we don't own -- is passed through unchanged. + */ + private def remapNamedExpressionAttributes( + ne: NamedExpression, + remap: Map[ExprId, Attribute]): NamedExpression = ne match { + case a: Attribute if remap.contains(a.exprId) => remap(a.exprId) + case a: Attribute => a + case al: Alias => + val newChild = al.child.transformUp { + case a: Attribute if remap.contains(a.exprId) => remap(a.exprId) + } + if (newChild eq al.child) al + else Alias(newChild, al.name)(al.exprId, al.qualifier, al.explicitMetadata) + case other => other + } + + // ============================================================================ + // Pattern A' / A2 dispatch (subquery plans of InSubquery / Exists) + // ============================================================================ + + private def rewriteSubqueryPlan(plan: LogicalPlan): Option[LogicalPlan] = { + // Candidate-level nondeterminism guard: reject if ANY node in the whole subquery plan + // is non-repeatable (Rand, LIMIT-without-ORDER-BY, Sample, Offset, streaming). This + // catches nondeterminism that has been hoisted above the self-join by an earlier + // optimizer rule -- the per-side `isSameBaseRelation` check alone would miss it because + // both innerLeft/innerRight can look deterministic after such a hoist. + if (!isRepeatablePlan(plan)) return None + + 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 + } + + 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] = { + val innerLeft = innerJoin.left + val innerRight = innerJoin.right + val innerCond = innerJoin.condition.get + + val parsed = parseSelfJoinCondition(innerCond, innerLeft, innerRight) + if (parsed.isEmpty) return None + // parseSelfJoinCondition guarantees Seq[(Attribute, Attribute)] and distinct equi-key names. + val (equiPairs, neqPairs) = parsed.get + + val innerLeftEquiAttrs: Seq[Attribute] = equiPairs.map(_._1) + val innerLeftNeqAttr: Attribute = neqPairs.head._1 + val filtered = buildAggregateHavingDistinctGt1(innerLeftEquiAttrs, innerLeftNeqAttr, innerLeft) + + // Fail-closed on bare-Join subqueries: without a wrapping Project the subquery output + // is the full self-join output (both sides' columns). Replacing that with + // `Project(equiKeys, filtered)` shrinks the output; if the enclosing InSubquery + // referenced a non-equi column by position, `values.zip(sub.output).map(EqualTo.tupled)` + // inside RewritePredicateSubquery would build an incorrect semi condition. Q95's + // subqueries all have an explicit Project wrapper, so this branch does not affect it. + projectListOpt match { + case None => + None + case Some(pl) => + canonicalizeWrapper(pl, equiPairs, filtered).map { + case (newWrapper, _) => + logDebug( + s"Pattern A' - equiKeys=[${innerLeftEquiAttrs.map(_.name).mkString(",")}]" + + s", neqCol=${innerLeftNeqAttr.name}" + + s", outCols=[${newWrapper.projectList.map(_.name).mkString(",")}]") + 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 sjRight = selfJoin.right + val sjCond = selfJoin.condition.get + if (!isSameBaseRelation(sjLeft, sjRight)) return None + + val parsed = parseSelfJoinCondition(sjCond, sjLeft, sjRight) + if (parsed.isEmpty) return None + // parseSelfJoinCondition guarantees Seq[(Attribute, Attribute)] and distinct equi-key names. + 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 + // 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 + + // Outer join condition may reference only equi-key attrs from the self-join side. + val outerCondRefs = outerCond.references.filter(selfJoinOutputSet.contains) + if (!outerCondRefs.forall(a => allEquiExprIds.contains(a.exprId))) return None + + // Top-level subquery Project may reference only equi-key attrs from the self-join side. + 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 = buildAggregateHavingDistinctGt1(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 neither a wrapper Project around the self-join nor a top-level + // subquery Project, the outer join currently exposes every self-join column, and + // replacing the self-join with `Project(equiKeys, filtered)` would shrink the outer + // join's right-hand output arity. RewritePredicateSubquery's positional zip + // (`values.zip(sub.output).map(EqualTo.tupled)`) would then bind semi predicates to + // the wrong attributes -- silently dropping components of a tuple IN/EXISTS. A + // top-level Project (`projectListOpt`) is what would let the arity be preserved + // by the top-level rewrite loop; without one, refuse to rewrite. + return None + case None => + // No wrapper Project but there IS a top-level subquery Project: shrinking the outer + // join's self-join-side output is safe because the top-level Project is rewritten + // consistently via `outputRemap` below and the top-level rewrite loop ensures + // subquery output arity matches what the enclosing InSubquery/Exists expects. + // Outer references may point at sjRight equi-attributes; remap them to sjLeft + // (same names 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) + } + + // Rewrite outer join condition to use new wrapper output attributes. + 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)) + } + + // Rewrite top-level Project references. + val result = projectListOpt match { + case Some(pl) => + val newPl = pl.map(ne => remapNamedExpressionAttributes(ne, outputRemap)) + Project(newPl, newOuterJoin) + case None => newOuterJoin + } + + logDebug( + s"Pattern A2 - equiKeys=[${sjLeftEquiAttrs.map(_.name).mkString(",")}]" + + s", neqCol=${sjLeftNeqAttr.name}") + Some(result) + } + + 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 + } + if (!isSameBaseRelation(join.left, join.right)) return None + val parsed = parseSelfJoinCondition(join.condition.get, join.left, join.right) + if (parsed.isEmpty) return None + Some(join) + } + + // ============================================================================ + // Pattern A : LeftSemi/LeftAnti whose right child is an Inner self-join + // ============================================================================ + + private def tryRewriteSemiWithSelfJoinChild(original: Join): Option[LogicalPlan] = { + // Candidate-level nondeterminism guard on the entire right subtree. Same rationale as + // in [[rewriteSubqueryPlan]]: catches Rand/Limit/Sample/Offset hoisted between the semi + // join and the inner self-join by an earlier optimizer rule. + if (!isRepeatablePlan(original.right)) return None + + val left = original.left + val right = original.right + val semiCondition = original.condition.get + + val (innerJoin, wrapper): (Join, Option[Project]) = right match { + case p @ Project(_, j: Join) if j.joinType == Inner && j.condition.isDefined => (j, Some(p)) + case j: Join if j.joinType == Inner && j.condition.isDefined => (j, None) + case _ => return None + } + + val innerLeft = innerJoin.left + val innerRight = innerJoin.right + val innerCond = innerJoin.condition.get + if (!isSameBaseRelation(innerLeft, innerRight)) return None + + val parsed = parseSelfJoinCondition(innerCond, innerLeft, innerRight) + if (parsed.isEmpty) return None + // parseSelfJoinCondition guarantees Seq[(Attribute, Attribute)] and distinct equi-key names. + val (innerEquiPairs, innerNeqPairs) = parsed.get + + // All semi predicates must be pure equi-join. + val semiPreds = splitConjunctivePredicates(semiCondition) + val leftOutputSet = left.outputSet + val rightOutputSet = right.outputSet + val semiEquiPairs = semiPreds.collect { + case EqualTo(l: Attribute, r: Attribute) + if leftOutputSet.contains(l) && rightOutputSet.contains(r) => + (l, r) + case EqualTo(r: Attribute, l: Attribute) + if leftOutputSet.contains(l) && rightOutputSet.contains(r) => + (l, r) + } + if (semiEquiPairs.size != semiPreds.size) return None + if (semiEquiPairs.isEmpty) return None + + val innerLeftEquiAttrs: Seq[Attribute] = innerEquiPairs.map(_._1) + val innerRightEquiAttrs: Seq[Attribute] = innerEquiPairs.map(_._2) + val innerEquiLeftAttrIds = innerLeftEquiAttrs.map(_.exprId).toSet + val innerEquiRightAttrIds = innerRightEquiAttrs.map(_.exprId).toSet + val innerEquiAllIds = innerEquiLeftAttrIds ++ innerEquiRightAttrIds + val innerNeqAttr: Attribute = innerNeqPairs.head._1 + + // Semi right-side keys must derive from inner equi-key attributes, possibly via wrapper alias. + val rightKeyIds = semiEquiPairs.map(_._2.exprId).toSet + val validSemiKeys = rightKeyIds.forall { + id => + innerEquiAllIds.contains(id) || wrapper.exists { + p => + p.projectList.exists { + case al @ Alias(a: Attribute, _) => + al.exprId == id && innerEquiAllIds.contains(a.exprId) + case a: Attribute => + a.exprId == id && innerEquiAllIds.contains(a.exprId) + case _ => false + } + } + } + if (!validSemiKeys) return None + + val filtered = buildAggregateHavingDistinctGt1(innerLeftEquiAttrs, innerNeqAttr, innerLeft) + + // Replace the entire right subtree with a Project of just the equi keys. + val projectedKeys = Project(innerLeftEquiAttrs, filtered) + + // Build an ExprId-keyed remap from every attribute reachable via the old right side to + // the corresponding sjLeft equi-attribute: + // (a) direct innerLeft equi attr -> identity (by ExprId) + // (b) direct innerRight equi attr -> paired sjLeft attr (looked up via equiPairs, NOT name) + // (c) wrapper `Alias(equiAttr, name)` output -> paired sjLeft attr (via inner Attribute's + // ExprId, NOT the alias name) + // Never use column name as identity: Catalyst allows same-name attributes with distinct + // ExprIds (e.g. `SELECT a AS k, b AS k`) and `.toMap` by name would silently drop one. + val exprIdToInnerLeft: Map[ExprId, Attribute] = + innerEquiPairs.flatMap { + case (l, r) => Seq(l.exprId -> l, r.exprId -> l) + }.toMap + val wrapperRemap: Map[ExprId, Attribute] = wrapper.map { + p => + p.projectList.flatMap { + case al @ Alias(a: Attribute, _) if exprIdToInnerLeft.contains(a.exprId) => + Some(al.exprId -> exprIdToInnerLeft(a.exprId)) + case _ => None + }.toMap + }.getOrElse(Map.empty) + val oldToNewMap: Map[ExprId, Attribute] = exprIdToInnerLeft ++ wrapperRemap + + // Fail-closed: refuse to rewrite if any attribute in the old right side is unresolvable. + val unresolved = semiCondition.collect { + case a: Attribute if rightOutputSet.contains(a) && !oldToNewMap.contains(a.exprId) => a + } + if (unresolved.nonEmpty) return None + + val newSemiCondition = semiCondition.transformUp { + case a: Attribute if oldToNewMap.contains(a.exprId) && rightOutputSet.contains(a) => + oldToNewMap(a.exprId) + } + + logDebug( + s"Pattern A (${original.joinType}) - equiKeys=[" + + innerLeftEquiAttrs.map(_.name).mkString(",") + + s"], neqCol=${innerNeqAttr.name}") + Some(original.copy(right = projectedKeys, condition = Some(newSemiCondition))) + } + + // ============================================================================ + // parseSelfJoinCondition + isSameBaseRelation + // ============================================================================ + + /** + * Parse a join condition into equi-pairs and inequality-pairs. Accepts only: + * - `EqualTo(attr, attr)` where the two attrs come from opposite sides, + * - `Not(EqualTo(attr, attr))` -- same side rule, + * - `IsNotNull(attr)` where the attr is one of the join columns. + * Anything else in the condition disqualifies the whole rewrite (fail-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 predicates on join columns are safe to drop -- they're redundant with + // the join semantics or auto-added by InferFiltersFromConstraints. IsNotNull on other + // columns changes semantics if we drop it; 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 + + // Only rewrite single-inequality case. Multi-column inequality + // (col_a<>col_a OR col_b<>col_b) + // is NOT equivalent to COUNT(DISTINCT single_col) > 1. + if (neqPairs.size != 1) return None + + // COUNT(DISTINCT neqCol) requires a hashable/orderable data type. Use Spark's public + // `RowOrdering.isOrderable` (which delegates to `OrderUtils.isOrderable`) to reject + // UDTs / Maps / and any complex type whose order-or-hash isn't defined. `AtomicType` + // would be a simpler predicate, but `AtomicType` is `protected[sql]` and thus not + // usable from Gluten's package. Fail-closed here so a query that would otherwise + // run does not crash in Catalyst's CheckAnalysis after our rule fires. + if (!neqPairs.forall { case (l, _) => RowOrdering.isOrderable(l.dataType) }) return None + + // Self-join invariant: equi and neq columns share names between the two sides. + val equiValid = equiPairs.forall { case (l, r) => l.name == r.name } Review Comment: `isSameBaseRelation` proves "both sides are the same relation" with `left.canonicalized == right.canonicalized`, but `parseSelfJoinCondition` proves "the two ends are the same column" with `l.name == r.name`. Canonicalization erases names: `QueryPlan.doCanonicalize` rewrites every `Alias` to `Alias(normalizedChild, "")`. So two projections that differ only by swapped column aliases canonicalize equal, and the name check then wires the equi pair and the neq pair to different underlying columns. With the two sides' aliases swapped (`a AS k, b AS v` against `a AS v, b AS k`), `s1.k = s2.k` really compares `t1.a` with `t2.b`, so IN returns extra rows and LeftAnti drops rows it should keep. Make the column-identity check positional too, each side looked up by its own end's ExprId: require `leftPlan.output.indexWhere(_.exprId == l.exprId)` to be `>= 0` and to equal `rightPlan.output.indexWhere(_.exprId == r.exprId)`. Do not use `indexOf` here: `AttributeReference.equals` also compares name and qualifier, and the attributes in the condition still carry their `SubqueryAlias` qualifier, so both sides would return `-1` and the rule would stop firing altogether, q95 included. The equi-key dedup at 638-639 needs the same treatment; `semiEquiPairs` does not, since its two sides are different relations. ########## backends-velox/src/test/scala/org/apache/gluten/extension/RewriteSelfJoinInequalityToAggregateSuite.scala: ########## @@ -0,0 +1,446 @@ +/* + * 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.gluten.extension + +import org.apache.gluten.execution.WholeStageTransformerSuite + +import org.apache.spark.SparkConf +import org.apache.spark.sql.Row +import org.apache.spark.sql.catalyst.expressions.Alias +import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan +import org.apache.spark.sql.internal.SQLConf + +/** + * Correctness tests for [[RewriteSelfJoinInequalityToAggregate]]. + * + * Assertions center on **result equivalence** between `rewrite=true` and `rewrite=false`. That is + * the direct check of the rule's semantic contract: the rewrite must not change what the query + * returns. Plan-shape assertions (e.g. "an Aggregate node exists") are avoided because downstream + * optimizer rules differ across Spark versions -- e.g. Spark 4.0's `RewritePredicateSubquery` and + * constant folding collapse LocalRelation-based EXISTS bodies so aggressively that our rule may + * never see the original self-join shape yet the final result is still correct. + * + * Where a plan-level signal is useful, we look for the alias name our rule injects + * (`_gluten_rw_selfjoin_cnt_distinct`) as a soft indicator that the rule fired. Its absence is not + * treated as a failure -- an equivalent result via a different path is still a pass. + */ +class RewriteSelfJoinInequalityToAggregateSuite extends WholeStageTransformerSuite { + + override protected val resourcePath: String = "/tpch-data-parquet" + override protected val fileFormat: String = "parquet" + + override protected def sparkConf: SparkConf = super.sparkConf + .set("spark.gluten.sql.rewrite.selfJoinInequality", "true") + .set(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key, "-1") + + /** Signature alias produced by the rewrite; presence => rule definitely fired. */ + private val CountDistinctAlias = "_gluten_rw_selfjoin_cnt_distinct" + + private def ruleFired(plan: LogicalPlan): Boolean = + plan.exists { + p => + p.expressions.exists(_.exists { + case a: Alias if a.name == CountDistinctAlias => true + case _ => false + }) + } + + /** Run `sql` twice, first with rewrite ON then OFF, and return the two result row sets. */ + private def runBoth(sql: String): (Set[Row], Set[Row]) = { + var on: Set[Row] = null + var off: Set[Row] = null + withSQLConf("spark.gluten.sql.rewrite.selfJoinInequality" -> "true") { + on = spark.sql(sql).collect().toSet + } + withSQLConf("spark.gluten.sql.rewrite.selfJoinInequality" -> "false") { + off = spark.sql(sql).collect().toSet + } + (on, off) + } + + private def setupTable(): Unit = { + // k=1: distinct v={10,20} -> matches (has 2 non-null distinct) + // k=2: distinct v={30} -> no match (only 1) + // k=3: distinct v={40,50,60} -> matches + // k=4: v={70, NULL} -> no match (only 1 non-null) + // k=5: v={NULL, NULL} -> no match (0 non-null) + // k=6: v={80, 90, NULL} -> matches + spark.sql( + """CREATE OR REPLACE TEMP VIEW T AS SELECT * FROM VALUES + | (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)) + |AS T(k, v)""".stripMargin) + } + + // ==================== Positive: rewrite is semantically equivalent ==================== + + test("Pattern A': EXISTS subquery with bare self-join produces equivalent results") { + setupTable() + val sql = + """SELECT k FROM T ws1 WHERE EXISTS ( + | SELECT 1 FROM T s WHERE s.k = ws1.k AND s.v <> ws1.v)""".stripMargin + val (on, off) = runBoth(sql) + assert(on == off, s"rewrite ON $on != OFF $off") + // Ground truth: only k in {1,3,6} have >=2 non-null distinct v. + assert(on == Set(Row(1), Row(3), Row(6)), s"expected {1,3,6}, got $on") + } + + test("Pattern A': InSubquery with bare self-join produces equivalent results") { + setupTable() + val sql = + """SELECT k FROM T ws1 WHERE k IN ( + | SELECT s1.k FROM T s1, T s2 + | WHERE s1.k = s2.k AND s1.v <> s2.v)""".stripMargin + val (on, off) = runBoth(sql) + assert(on == off, s"rewrite ON $on != OFF $off") + assert(on == Set(Row(1), Row(3), Row(6))) + } + + test("Pattern A2: self-join nested inside outer InnerJoin produces equivalent results") { + setupTable() + // Only k in {1,3,6} qualify from the self-join side; the outer InnerJoin with D + // (values {1,3,6}) intersects, so the final answer is again {1,3,6}. + 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, T s2 + | WHERE s1.k = s2.k AND s1.v <> s2.v) sj + | WHERE d.k = sj.k)""".stripMargin + 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))) + } + + // ==================== Semantic parity on NULL / 3VL ==================== + + test("NULL / 3VL: rows with only-NULL or single-non-null inequality column are excluded") { + setupTable() + val sql = + """SELECT k FROM T ws1 WHERE EXISTS ( + | SELECT 1 FROM T s WHERE s.k = ws1.k AND s.v <> ws1.v)""".stripMargin + val (on, off) = runBoth(sql) + // k=4 (v={70,NULL}) fails: <> with NULL is UNKNOWN -> filtered by WHERE. + // k=5 (v={NULL,NULL}) fails: every <> is UNKNOWN. + 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") + } + + // ==================== Negative: rewrite must produce equivalent results (or bail) ========== + + test("Plain InnerJoin at top level: results unchanged (rewrite must not touch it)") { + setupTable() + val sql = + """SELECT ws1.k FROM T ws1 JOIN T ws2 + |ON ws1.k = ws2.k AND ws1.v <> ws2.v""".stripMargin + // Row-multiplicity matters here; using count() to catch any drop or dup. + var onCount: Long = -1L + var offCount: Long = -1L + withSQLConf("spark.gluten.sql.rewrite.selfJoinInequality" -> "true") { + onCount = spark.sql(sql).count() + } + withSQLConf("spark.gluten.sql.rewrite.selfJoinInequality" -> "false") { + offCount = spark.sql(sql).count() + } + assert( + onCount == offCount, + s"plain InnerJoin row-count differs: rewrite=$onCount vs baseline=$offCount") + } + + test("IS DISTINCT FROM: NULL-safe inequality preserves original semantics") { + setupTable() + // IS DISTINCT FROM treats NULL as distinguishable (NULL IS DISTINCT FROM x = TRUE, + // NULL IS DISTINCT FROM NULL = FALSE). Our rewrite must NOT fold this into + // COUNT(DISTINCT), because COUNT(DISTINCT) ignores NULL. + val sql = + """SELECT k FROM T ws1 WHERE EXISTS ( + | SELECT 1 FROM T s WHERE s.k = ws1.k + | AND (s.v IS DISTINCT FROM ws1.v))""".stripMargin + val (on, off) = runBoth(sql) + assert(on == off, s"IS DISTINCT FROM semantics diverge: ON=$on OFF=$off") + // Sanity check: k=4 has (70, NULL) -- pair (v=70, v=NULL) IS DISTINCT FROM => TRUE + // so k=4 must be included (unlike the plain-neq case above where it's excluded). + assert(on.contains(Row(4)), s"k=4 should be in IS DISTINCT FROM result: $on") + // And rule fire signal must be absent: this is a rejection path. + val plan = spark.sql(sql).queryExecution.optimizedPlan + assert(!ruleFired(plan), s"rule must not fire on IS DISTINCT FROM:\n$plan") Review Comment: The negative tests are all written as correlated subqueries, but `apply` only handles `InSubquery`/`Exists` with `children.isEmpty`, and a correlated subquery's `outerAttrs`/`joinCond` are non-empty by the time an injected rule runs, so the expression comes back untouched. The `IS DISTINCT FROM`, `IsNotNull`-on-non-join-column, multi-column-inequality, `rand()` and equi/neq-overlap tests, together with lines 95, 137 and 284, are eight queries whose subquery holds a single table and no Join node at all, so even dropping the correlation would not reach the branch each test is named after. Only the IN test at line 106 and the A2 test at line 117 trigger the rewrite, and only lines 386 and 407 are effective negative tests: line 426 uses `rand(41)` on one side and `rand(42)` on the other, so the candidate-level `plan.deterministic` is what rejects it first, and `isSameBaseRelation` catches it too, so deleting the determinism check keeps the test green. Replace `apply` with the identity function and all 19 still pass, because none of them asserts that the rule fired. Please rewrite the negative tests as uncorrelated subqueries with the self-join inside and a constant `WHERE`, and add `assert(ruleFired(plan))` to the positive ones. `IsNotNull(equiKeys)` additionally needs a fixture with a NULL equi key plus a `NOT IN` query. ########## backends-velox/src/main/scala/org/apache/gluten/extension/RewriteSelfJoinInequalityToAggregate.scala: ########## @@ -0,0 +1,722 @@ +/* + * 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.gluten.extension + +import org.apache.gluten.config.VeloxConfig + +import org.apache.spark.internal.Logging +import org.apache.spark.sql.SparkSession +import org.apache.spark.sql.catalyst.expressions._ +import org.apache.spark.sql.catalyst.expressions.RowOrdering +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.types.LongType + +/** + * Rewrites self-join with inequality into GROUP BY + HAVING COUNT(DISTINCT) > 1. + * + * Targets three patterns; all require an existence-only context (LeftSemi/LeftAnti join, or + * InSubquery/Exists expression) so that row-count multiplicity from the self-join cross-product + * does not affect semantics. + * + * - Pattern A' (InSubquery/Exists primary): InSubquery/Exists whose subquery top-level join is a + * direct self-join. The primary path for TPC-DS Q95. + * - Pattern A2 (nested): InSubquery/Exists whose subquery contains an outer InnerJoin that has a + * self-join child. Only the self-join child is replaced with Aggregate; the outer join is + * preserved. + * - Pattern A (LeftSemi/LeftAnti): LeftSemi/LeftAnti whose right child is an Inner self-join + * (possibly wrapped in Project). Matches semi/anti joins that already exist in the input -- + * e.g. from an explicit `LEFT SEMI JOIN` clause. Note: this rule is injected via + * `injectOptimizerRule`, which places it in the operator-optimization batch that runs BEFORE + * `RewritePredicateSubquery`; A is NOT a post-subquery-rewrite fallback for A'. + * + * Correlated subqueries (outer references / joinCond in ListQuery/Exists) are fail-closed at the + * entry expression, since our ExprId canonicalization does not remap those predicates. + * + * All three share: + * - [[buildAggregateHavingDistinctGt1]] to construct `Filter(cnt > 1, Aggregate)` + * - [[canonicalizeWrapper]] to rebuild a wrapping Project so every equi-key reference points to + * the sjLeft-side attribute, with **fresh exprIds** (Spark's SPARK-21835 style -- no reuse of + * original exprIds), returning an old->new attribute remap for downstream rewrite. + * + * Controlled by `spark.gluten.sql.rewrite.selfJoinInequality` (default false, opt-in). + */ +case class RewriteSelfJoinInequalityToAggregate(spark: SparkSession) + extends Rule[LogicalPlan] + with PredicateHelper + with Logging { + + private val CountDistinctAliasName = "_gluten_rw_selfjoin_cnt_distinct" + + override def apply(plan: LogicalPlan): LogicalPlan = { + if (!VeloxConfig.get.enableRewriteSelfJoinInequality) { + logDebug("RewriteSelfJoinInequalityToAggregate: disabled via config, skipping") + return plan + } + + // Pattern A: rewrite LeftSemi/LeftAnti whose right child is an Inner self-join. + val afterOps = plan.transformUp { + case j: Join + if (j.joinType == LeftSemi || j.joinType == LeftAnti) && + j.condition.isDefined && + isInnerJoinShape(j.right) => + tryRewriteSemiWithSelfJoinChild(j).getOrElse(j) + case other => other + } + + // Pattern A' / A2: rewrite subquery plans embedded in InSubquery/Exists. + // Type-based matching (`x: T`) + named-argument copy keeps this portable across + // Spark 3.3/3.4/3.5/4.x where ListQuery/Exists case-class arity has drifted. + // + // Correlated subquery fail-closed: `SubqueryExpression.children.nonEmpty` iff the + // subquery has outer references / correlated join conditions. These predicates + // reference attributes INSIDE the subquery plan by ExprId; our canonicalizeWrapper + // rewrites those ExprIds without remapping the correlated predicates, which would + // leave dangling references after `RewritePredicateSubquery` folds them back into + // the semi-join condition. Target workload (TPC-DS Q95) is uncorrelated, so bail + // on any correlated candidate rather than growing the remap surface. + val rewritten = afterOps.transformAllExpressions { + 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 + } + case ex: Exists if ex.children.isEmpty => + rewriteSubqueryPlan(ex.plan) match { + case Some(newSub) => ex.copy(plan = newSub) + case None => ex + } + } + if (!(rewritten eq plan)) { + logDebug( + "RewriteSelfJoinInequalityToAggregate: rewrote self-join to " + + "GROUP BY + HAVING COUNT(DISTINCT) > 1") + } + rewritten + } + + // ============================================================================ + // Shared helpers + // ============================================================================ + + private def isInnerJoinShape(plan: LogicalPlan): Boolean = plan match { + case Project(_, j: Join) if j.joinType == Inner && j.condition.isDefined => true + case j: Join if j.joinType == Inner && j.condition.isDefined => true + case _ => false + } + + /** + * Build `Filter(cnt > 1, Aggregate(equiKeys, [equiKeys, cnt_alias], Filter(IsNotNull(equiKeys), + * child)))`. Returns the Filter node whose output is `equiKeys ++ [count_alias_attr]`. + * + * The extra `IsNotNull(equiKeys)` filter is essential to preserve the original equi-join's NULL + * semantics. Under SQL 3VL, `left.k = right.k` never matches when either side is NULL, so the + * original self-join drops rows with NULL equi-keys. Aggregate, in contrast, groups NULL keys + * together into a single "NULL group" -- if that group has >= 2 distinct non-null neq values, + * COUNT(DISTINCT) > 1 fires and injects NULL into the subquery output. That leaked NULL then + * turns `NOT IN` into a spurious empty result (Spark's null-aware anti-join uses + * `Or(equi, IsNull(equi))` which any NULL sub-row satisfies) and can flip EXISTS/IN outcomes. The + * neq column needs no such filter: `COUNT(DISTINCT col)` already ignores NULL. + */ + private def buildAggregateHavingDistinctGt1( + equiKeys: Seq[Attribute], + neqCol: Attribute, + child: LogicalPlan): LogicalPlan = { + val countExpr = AggregateExpression( + Count(Seq(neqCol)), + mode = Complete, + isDistinct = true, + filter = None, + NamedExpression.newExprId) + val countAlias = Alias(countExpr, CountDistinctAliasName)() + // Seq[Attribute] is a Seq[NamedExpression] via covariance; no cast needed. + val aggExprs: Seq[NamedExpression] = equiKeys :+ countAlias + val nonNullChild = equiKeys + .map(a => IsNotNull(a): Expression) + .reduceOption(And) + .map(Filter(_, child)) + .getOrElse(child) + val agg = Aggregate(equiKeys, aggExprs, nonNullChild) + Filter(GreaterThan(countAlias.toAttribute, Literal(1L, LongType)), agg) + } + + /** + * Canonicalize a Project so every equi-key reference points at the sjLeft-side attribute (both + * sides of a valid self-join share names, so this substitution is semantically safe). Uses + * **fresh exprIds** (no reuse of original wrapper output exprIds) -- the same technique Spark's + * own `dedupSubqueryOnSelfJoin` uses when it needs to change subquery output. + * + * Returns the rebuilt Project and a map `oldWrapperOutputExprId -> newWrapperOutputAttr`, so + * downstream references (outer join condition, top-level Project) can be updated consistently. + * + * `equiPairs` provides the definitive ExprId-based lookup: `equiPair (l, r)` binds + * `l.exprId -> l` (identity) and `r.exprId -> l` (sjRight -> sjLeft). Attribute identity in + * Catalyst is ExprId, not name; two columns can share a name with distinct ExprIds. Name-based + * lookup would silently drop such entries via `.toMap`. + * + * Fails (returns None) when a projectList entry is neither an equi-key Attribute (by ExprId) nor + * `Alias(equi-key Attribute, _)`. Fail-closed. + */ + private def canonicalizeWrapper( + projectList: Seq[NamedExpression], + equiPairs: Seq[(Attribute, Attribute)], + newChild: LogicalPlan): Option[(Project, Map[ExprId, Attribute])] = { + // ExprId-based canonical map: any equi-key attribute (either side) -> sjLeft attribute. + val exprIdToLeft: Map[ExprId, Attribute] = + equiPairs.flatMap { case (l, r) => Seq(l.exprId -> l, r.exprId -> l) }.toMap + val oldOutput: Seq[Attribute] = projectList.map(_.toAttribute) + val mapped: Seq[Option[NamedExpression]] = projectList.map { + case a: Attribute if exprIdToLeft.contains(a.exprId) => + // Wrap every rewritten output slot in a fresh Alias. + // + // When a wrapper reprojects BOTH sides of the same equi pair (e.g. + // `SELECT s1.k, s2.k FROM T s1 JOIN T s2 ON s1.k = s2.k AND s1.v <> s2.v`), + // both entries collapse to the same sjLeft Attribute after the self-join is + // rewritten. Duplicate output ExprIds are not illegal in Spark (`SELECT a, a` + // is a valid Project), but fresh Aliases give each output slot an independent + // identity, which keeps the `oldOutput -> newOutput` remap 1-to-1 and lets + // downstream references (outer join condition, top-level Project) be updated + // unambiguously via ExprId. + // + // The fresh ExprId is on the Alias ITSELF; the referenced child keeps its + // original ExprId. Spark's logical-plan integrity checks reject reusing a + // referenced ExprId as the Alias's own ExprId, not duplication across slots. + Some(Alias(exprIdToLeft(a.exprId), a.name)(): NamedExpression) + case al @ Alias(a: Attribute, _) if exprIdToLeft.contains(a.exprId) => + // Fresh exprId; do NOT reuse `al.exprId`. Reusing another expression's exprId + // is the pattern that Spark 3.3 flags via structural-integrity checks. + Some(Alias(exprIdToLeft(a.exprId), al.name)(): NamedExpression) + case _ => None + } + if (mapped.exists(_.isEmpty)) { + None + } else { + val newProjectList = mapped.flatten + val newWrapper = Project(newProjectList, newChild) + val newOutput = newWrapper.output + val remap: Map[ExprId, Attribute] = + oldOutput.zip(newOutput).map { case (o, n) => o.exprId -> n }.toMap + Some((newWrapper, remap)) + } + } + + /** + * Replace equi-key attribute references inside a NamedExpression according to `remap`, while + * preserving the NamedExpression shape. + * + * `Expression.transformUp` returns `Expression`, not `NamedExpression`. We avoid a blanket + * `asInstanceOf[NamedExpression]` by handling the two shapes that can appear in a Project's + * `projectList` explicitly: a bare Attribute (whose top-level may itself be replaced) and an + * Alias (which stays an Alias while its child is transformed). Anything else in a projectList -- + * e.g. computed expressions we don't own -- is passed through unchanged. + */ + private def remapNamedExpressionAttributes( + ne: NamedExpression, + remap: Map[ExprId, Attribute]): NamedExpression = ne match { + case a: Attribute if remap.contains(a.exprId) => remap(a.exprId) + case a: Attribute => a + case al: Alias => + val newChild = al.child.transformUp { + case a: Attribute if remap.contains(a.exprId) => remap(a.exprId) + } + if (newChild eq al.child) al + else Alias(newChild, al.name)(al.exprId, al.qualifier, al.explicitMetadata) + case other => other + } + + // ============================================================================ + // Pattern A' / A2 dispatch (subquery plans of InSubquery / Exists) + // ============================================================================ + + private def rewriteSubqueryPlan(plan: LogicalPlan): Option[LogicalPlan] = { + // Candidate-level nondeterminism guard: reject if ANY node in the whole subquery plan + // is non-repeatable (Rand, LIMIT-without-ORDER-BY, Sample, Offset, streaming). This + // catches nondeterminism that has been hoisted above the self-join by an earlier + // optimizer rule -- the per-side `isSameBaseRelation` check alone would miss it because + // both innerLeft/innerRight can look deterministic after such a hoist. + if (!isRepeatablePlan(plan)) return None + + 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 + } + + 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] = { + val innerLeft = innerJoin.left + val innerRight = innerJoin.right + val innerCond = innerJoin.condition.get + + val parsed = parseSelfJoinCondition(innerCond, innerLeft, innerRight) + if (parsed.isEmpty) return None + // parseSelfJoinCondition guarantees Seq[(Attribute, Attribute)] and distinct equi-key names. + val (equiPairs, neqPairs) = parsed.get + + val innerLeftEquiAttrs: Seq[Attribute] = equiPairs.map(_._1) + val innerLeftNeqAttr: Attribute = neqPairs.head._1 + val filtered = buildAggregateHavingDistinctGt1(innerLeftEquiAttrs, innerLeftNeqAttr, innerLeft) + + // Fail-closed on bare-Join subqueries: without a wrapping Project the subquery output + // is the full self-join output (both sides' columns). Replacing that with + // `Project(equiKeys, filtered)` shrinks the output; if the enclosing InSubquery + // referenced a non-equi column by position, `values.zip(sub.output).map(EqualTo.tupled)` + // inside RewritePredicateSubquery would build an incorrect semi condition. Q95's + // subqueries all have an explicit Project wrapper, so this branch does not affect it. + projectListOpt match { + case None => + None + case Some(pl) => + canonicalizeWrapper(pl, equiPairs, filtered).map { + case (newWrapper, _) => + logDebug( + s"Pattern A' - equiKeys=[${innerLeftEquiAttrs.map(_.name).mkString(",")}]" + + s", neqCol=${innerLeftNeqAttr.name}" + + s", outCols=[${newWrapper.projectList.map(_.name).mkString(",")}]") + 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 sjRight = selfJoin.right + val sjCond = selfJoin.condition.get + if (!isSameBaseRelation(sjLeft, sjRight)) return None + + val parsed = parseSelfJoinCondition(sjCond, sjLeft, sjRight) + if (parsed.isEmpty) return None + // parseSelfJoinCondition guarantees Seq[(Attribute, Attribute)] and distinct equi-key names. + 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 + // 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 + + // Outer join condition may reference only equi-key attrs from the self-join side. + val outerCondRefs = outerCond.references.filter(selfJoinOutputSet.contains) + if (!outerCondRefs.forall(a => allEquiExprIds.contains(a.exprId))) return None + + // Top-level subquery Project may reference only equi-key attrs from the self-join side. + 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 = buildAggregateHavingDistinctGt1(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 neither a wrapper Project around the self-join nor a top-level + // subquery Project, the outer join currently exposes every self-join column, and + // replacing the self-join with `Project(equiKeys, filtered)` would shrink the outer + // join's right-hand output arity. RewritePredicateSubquery's positional zip + // (`values.zip(sub.output).map(EqualTo.tupled)`) would then bind semi predicates to + // the wrong attributes -- silently dropping components of a tuple IN/EXISTS. A + // top-level Project (`projectListOpt`) is what would let the arity be preserved + // by the top-level rewrite loop; without one, refuse to rewrite. + return None + case None => + // No wrapper Project but there IS a top-level subquery Project: shrinking the outer + // join's self-join-side output is safe because the top-level Project is rewritten + // consistently via `outputRemap` below and the top-level rewrite loop ensures + // subquery output arity matches what the enclosing InSubquery/Exists expects. + // Outer references may point at sjRight equi-attributes; remap them to sjLeft + // (same names 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) + } + + // Rewrite outer join condition to use new wrapper output attributes. + 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)) + } + + // Rewrite top-level Project references. + val result = projectListOpt match { + case Some(pl) => + val newPl = pl.map(ne => remapNamedExpressionAttributes(ne, outputRemap)) + Project(newPl, newOuterJoin) + case None => newOuterJoin + } + + logDebug( + s"Pattern A2 - equiKeys=[${sjLeftEquiAttrs.map(_.name).mkString(",")}]" + + s", neqCol=${sjLeftNeqAttr.name}") + Some(result) + } + + 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 + } + if (!isSameBaseRelation(join.left, join.right)) return None + val parsed = parseSelfJoinCondition(join.condition.get, join.left, join.right) + if (parsed.isEmpty) return None + Some(join) + } + + // ============================================================================ + // Pattern A : LeftSemi/LeftAnti whose right child is an Inner self-join + // ============================================================================ + + private def tryRewriteSemiWithSelfJoinChild(original: Join): Option[LogicalPlan] = { Review Comment: `tryRewriteSemiWithSelfJoinChild` plus `validSemiKeys`, `wrapperRemap` and `newSemiCondition` is about 110 lines, and it is the only path in this PR that rewrites a join condition in the main plan rather than inside a subquery. The suite contains zero `LEFT SEMI`/`LEFT ANTI` queries, so none of it has ever run. The only way in is an explicitly written `LEFT SEMI JOIN`/`LEFT ANTI JOIN`: `semiEquiPairs` accepts only `EqualTo`, while `ReplaceIntersectWithSemiJoin`/`ReplaceExceptWithAntiJoin` emit `EqualNullSafe`, so `INTERSECT`/`EXCEPT DISTINCT` never get in, and the semi joins that IN/EXISTS turn into only appear in the `RewriteSubquery` batch, after this rule has run. Please add explicit semi/anti tests, including the aliased shape that goes through the wrapper Project. Alternatively ship A' and A2 only for now and drop Pattern A together with that part of the comment. -- 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]
