hhr293 commented on code in PR #12756: URL: https://github.com/apache/gluten/pull/12756#discussion_r3795558198
########## backends-velox/src/main/scala/org/apache/gluten/extension/RewriteSelfJoinInequalityToAggregate.scala: ########## @@ -0,0 +1,622 @@ +/* + * 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.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 (post-RewritePredicateSubquery fallback): LeftSemi/LeftAnti whose right child is an + * Inner self-join (possibly wrapped in Project). The outer semi-join condition is pure equi-key + * referencing a column from the self-join output. The self-join condition has equi + inequality on + * the same table. Rewrites the right child to GROUP BY + HAVING. + * + * Pattern A' (pre-RewritePredicateSubquery, primary path): InSubquery(_, ListQuery(sub, ...)) or + * Exists(sub, ...) whose sub is Project(Inner self-join with equi+neq). Rewrites sub to + * Project(equi_keys, Filter(count_distinct>1, Aggregate)). Fires in the Operator Optimization + * batches BEFORE Spark lifts them to LeftSemi/LeftAnti; the outer expression form guarantees + * existence semantics. ScalarSubquery is intentionally NOT matched. + * + * Pattern A2 (nested self-join): InSubquery/Exists whose subquery plan is Project(Join(Inner, + * other_table, self-join)) -- the self-join is a child of another InnerJoin, not the top-level join + * itself. Only the self-join child is replaced with Aggregate; the outer join is preserved. Safe + * because the outer join connects on the equi-key, and the subquery is still consumed as an + * existence set. + * + * Controlled by spark.gluten.sql.rewrite.selfJoinInequality (default false, opt-in until exercised + * more broadly across workloads). + */ +case class RewriteSelfJoinInequalityToAggregate(spark: SparkSession) + extends Rule[LogicalPlan] + with PredicateHelper + with Logging { + + override def apply(plan: LogicalPlan): LogicalPlan = { + if (!VeloxConfig.get.enableRewriteSelfJoinInequality) { + logDebug("RewriteSelfJoinInequalityToAggregate: disabled via config, skipping") + return plan + } + + val afterOps = plan.transformUp { + // Pattern A: LeftSemi/LeftAnti whose right child is an Inner self-join + // (possibly wrapped in Project). + case j: Join + if (j.joinType == LeftSemi || j.joinType == LeftAnti) && + j.condition.isDefined && + isInnerJoinShape(j.right) => + tryRewriteSemiWithSelfJoinChild(j, j.left, j.right, j.joinType, j.condition.get, j.hint) + .getOrElse(j) + + case other => other + } + + // Pattern A': rewrite subquery plans embedded in InSubquery/Exists expressions. + // Fires before Spark's RewritePredicateSubquery (batch pos 26); once the subquery + // plan is rewritten to GROUP BY + HAVING, RewritePredicateSubquery lifts it to + // LeftSemi/LeftAnti in the normal way. + // Use type-based matching (`x: T`) and named-argument copy (`x.copy(plan = ...)`) + // instead of case-class unapply with a fixed parameter list. This keeps the code + // portable across Spark 3.3/3.4/3.5/4.x where the internal ListQuery/Exists + // case classes have added parameters over releases. + val rewritten = afterOps.transformAllExpressions { + case in @ InSubquery(_, lq: ListQuery) => + rewriteSubqueryPlan(lq.plan) match { + case Some(newSub) => in.copy(query = lq.copy(plan = newSub)) + case None => in + } + case ex: Exists => + rewriteSubqueryPlan(ex.plan) match { + case Some(newSub) => ex.copy(plan = newSub) + case None => ex + } + } + if (!(rewritten eq plan)) { + logInfo("RewriteSelfJoinInequalityToAggregate: rewrote self-join to " + + "GROUP BY + HAVING COUNT(DISTINCT) > 1") + } + rewritten + } + + 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 + } + + /** + * Dispatches subquery plan rewriting: tries Pattern A' (direct self-join at top level) first, + * then Pattern A2 (self-join nested as a child of another InnerJoin). + * + * Only called on subquery plans of `InSubquery` / `Exists`, i.e. contexts that consume the output + * as a set of distinct keys. Cardinality of the intermediate is safe to change. + */ + private def rewriteSubqueryPlan(plan: LogicalPlan): Option[LogicalPlan] = { + 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(plan, projectListOpt, innerJoin) + } else { + rewriteNestedSelfJoin(plan, projectListOpt, innerJoin) + } + } + + private def rewriteDirectSelfJoin( + plan: LogicalPlan, + projectListOpt: Option[Seq[NamedExpression]], + innerJoin: Join): Option[LogicalPlan] = { + + // Use field accessors (portable across Spark versions) instead of + // case-class unapply with a fixed parameter list. The caller has already + // pattern-matched innerJoin as Join(_, _, Inner, Some(_), _). + val innerLeft = innerJoin.left + val innerRight = innerJoin.right + val innerCond = innerJoin.condition.get + + val parsed = parseSelfJoinCondition(innerCond, innerLeft, innerRight) + if (parsed.isEmpty) return None + val (equiPairs, neqPairs) = parsed.get + + val innerLeftEquiAttrs = equiPairs.map(_._1).collect { case a: Attribute => a } + if (innerLeftEquiAttrs.size != equiPairs.size) return None + + val innerLeftNeqAttr = neqPairs.head._1 match { + case a: Attribute => a + case _ => return None + } + + val countDistinctExpr = AggregateExpression( + Count(Seq(innerLeftNeqAttr)), + mode = Complete, + isDistinct = true, + filter = None, + NamedExpression.newExprId) + val countAlias = Alias(countDistinctExpr, "_gluten_rw_selfjoin_cnt_distinct")() + + val groupingExprs: Seq[Expression] = innerLeftEquiAttrs + val aggExprs: Seq[NamedExpression] = + innerLeftEquiAttrs.map(_.asInstanceOf[NamedExpression]) :+ countAlias + val aggregate = Aggregate(groupingExprs, aggExprs, innerLeft) + + val filterExpr = GreaterThan(countAlias.toAttribute, Literal(1L, LongType)) + val filtered = Filter(filterExpr, aggregate) + + val nameToInnerLeft: Map[String, Attribute] = innerLeftEquiAttrs.map(a => a.name -> a).toMap + val innerLeftEquiExprIds = innerLeftEquiAttrs.map(_.exprId).toSet + val innerRightEquiExprIds = + equiPairs.map(_._2).collect { case a: Attribute => a.exprId }.toSet + + val projectListResolved: Option[Seq[NamedExpression]] = projectListOpt match { + case None => + Some(innerLeftEquiAttrs.map(_.asInstanceOf[NamedExpression])) + case Some(pl) => + val remapped = pl.map { + case a: Attribute + if nameToInnerLeft.contains(a.name) && + (innerLeftEquiExprIds.contains(a.exprId) || + innerRightEquiExprIds.contains(a.exprId)) => + Some(Alias(nameToInnerLeft(a.name), a.name)(a.exprId).asInstanceOf[NamedExpression]) + case al @ Alias(a: Attribute, _) + if nameToInnerLeft.contains(a.name) && + (innerLeftEquiExprIds.contains(a.exprId) || + innerRightEquiExprIds.contains(a.exprId)) => + Some(Alias(nameToInnerLeft(a.name), al.name)(al.exprId).asInstanceOf[NamedExpression]) + case _ => + None + } + if (remapped.exists(_.isEmpty)) None + else Some(remapped.flatten) + } + + projectListResolved.map { + pl => + logInfo( + s"RewriteSelfJoinInequalityToAggregate: Pattern A' - rewrote subquery Project(Inner " + + s"self-join) to Project + Filter(count_distinct>1) + Aggregate. " + + s"equiKeys=[${innerLeftEquiAttrs.map(_.name).mkString(",")}], " + + s"neqCol=${innerLeftNeqAttr.name}, outCols=[${pl.map(_.name).mkString(",")}]") + Project(pl, filtered) + } + } + + /** + * Pattern A2: the top-level InnerJoin is NOT a self-join, but one of its children IS a self-join + * (possibly wrapped in Project). Example: `Join(Inner, web_returns, Project(self-join))`. + * + * We replace the self-join child with Aggregate + Filter, preserving the outer join. This is safe + * because: + * 1. The outer join condition connects the other table to the self-join's equi-key. + * 2. Replacing the self-join with GROUP BY + HAVING preserves the set of distinct keys (only + * row-count multiplicity changes), and the other table joins on that key. + * 3. The entire subquery is consumed by InSubquery/Exists (existence semantics), so the final + * output is still just a set of distinct keys. + */ + private def rewriteNestedSelfJoin( + plan: LogicalPlan, + projectListOpt: Option[Seq[NamedExpression]], + outerJoin: Join): Option[LogicalPlan] = { + + // Use field accessors instead of case-class unapply (portable across Spark + // versions). Caller has already pattern-matched outerJoin as + // Join(_, _, Inner, Some(_), _). + val outerLeft = outerJoin.left + val outerRight = outerJoin.right + val outerCond = outerJoin.condition.get + val outerHint = outerJoin.hint + + val (selfJoinSide, otherSide, selfJoinOnRight) = + tryExtractSelfJoin(outerRight) match { + case Some(_) => (outerRight, outerLeft, true) + case None => + tryExtractSelfJoin(outerLeft) match { + case Some(_) => (outerLeft, outerRight, false) + case None => return None + } + } + + val (selfJoinProjectOpt, selfJoin) = selfJoinSide 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 + } + + 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 + val (equiPairs, neqPairs) = parsed.get + + val sjLeftEquiAttrs = equiPairs.map(_._1).collect { case a: Attribute => a } + if (sjLeftEquiAttrs.size != equiPairs.size) return None + + val sjLeftNeqAttr = neqPairs.head._1 match { + case a: Attribute => a + case _ => return None + } + + val selfJoinOutputSet = selfJoinSide.outputSet + val outerCondRefs = outerCond.references.filter(selfJoinOutputSet.contains) + val sjEquiExprIds = equiPairs.flatMap { + case (l: Attribute, r: Attribute) => Seq(l.exprId, r.exprId) + case _ => Seq.empty + }.toSet + val wrapperEquiExprIds: Set[ExprId] = selfJoinProjectOpt match { + case Some(pl) => + pl.flatMap { + case al @ Alias(a: Attribute, _) if sjEquiExprIds.contains(a.exprId) => + Some(al.exprId) + case a: Attribute if sjEquiExprIds.contains(a.exprId) => + Some(a.exprId) + case _ => None + }.toSet + case None => Set.empty + } + val allEquiExprIds = sjEquiExprIds ++ wrapperEquiExprIds + if (!outerCondRefs.forall(a => allEquiExprIds.contains(a.exprId))) return None + + // The rewrite replaces the self-join subtree with Project(equiKeys, Filter(Aggregate)), + // so any reference to non-equi self-join attributes in the subquery's top-level + // projectList would become unresolved. Bail out unless the top-level projectList + // depends only on equi-key attributes from the self-join side. + val projectOk = projectListOpt.forall { + pl => + val projRefsFromSelfJoin = pl.flatMap(_.references).filter(selfJoinOutputSet.contains) + projRefsFromSelfJoin.forall(a => allEquiExprIds.contains(a.exprId)) + } + if (!projectOk) return None + + val countDistinctExpr = AggregateExpression( + Count(Seq(sjLeftNeqAttr)), + mode = Complete, + isDistinct = true, + filter = None, + NamedExpression.newExprId) + val countAlias = Alias(countDistinctExpr, "_gluten_rw_selfjoin_cnt_distinct")() + + val groupingExprs: Seq[Expression] = sjLeftEquiAttrs + val aggExprs: Seq[NamedExpression] = + sjLeftEquiAttrs.map(_.asInstanceOf[NamedExpression]) :+ countAlias + val aggregate = Aggregate(groupingExprs, aggExprs, sjLeft) + + val filterExpr = GreaterThan(countAlias.toAttribute, Literal(1L, LongType)) + val filtered = Filter(filterExpr, aggregate) + + val sjLeftEquiByName: Map[String, Attribute] = sjLeftEquiAttrs.map(a => a.name -> a).toMap + val sjLeftEquiExprIds = sjLeftEquiAttrs.map(_.exprId).toSet + val sjRightEquiExprIds = equiPairs.map(_._2).collect { case a: Attribute => a.exprId }.toSet + + val newSelfJoinSide: LogicalPlan = selfJoinProjectOpt match { + case Some(pl) => + // Use flatMap + size check instead of `return None` inside the map lambda, + // which triggers Scala's nonlocal return (unsafe / lint-flagged). + val remapped: Seq[NamedExpression] = pl.flatMap { + case a: Attribute + if sjLeftEquiByName.contains(a.name) && + (sjLeftEquiExprIds.contains(a.exprId) || + sjRightEquiExprIds.contains(a.exprId)) => + Some(Alias(sjLeftEquiByName(a.name), a.name)(a.exprId): NamedExpression) + case al @ Alias(a: Attribute, _) + if sjLeftEquiByName.contains(a.name) && + (sjLeftEquiExprIds.contains(a.exprId) || + sjRightEquiExprIds.contains(a.exprId)) => + Some(Alias(sjLeftEquiByName(a.name), al.name)(al.exprId): NamedExpression) + case _ => None + } + if (remapped.size != pl.size) return None + Project(remapped, filtered) + case None => + Project(sjLeftEquiAttrs.map(_.asInstanceOf[NamedExpression]), filtered) + } + + val newOuterJoin = if (selfJoinOnRight) { + outerJoin.copy(right = newSelfJoinSide) + } else { + outerJoin.copy(left = newSelfJoinSide) + } + + val result = projectListOpt match { + case Some(pl) => Project(pl, newOuterJoin) + case None => newOuterJoin + } + + logInfo( + s"RewriteSelfJoinInequalityToAggregate: Pattern A2 - rewrote nested self-join inside " + + s"subquery InnerJoin. Self-join replaced with GROUP BY HAVING COUNT(DISTINCT) > 1. " + + s"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 + } + val l = join.left + val r = join.right + val cond = join.condition.get + if (!isSameBaseRelation(l, r)) return None + val parsed = parseSelfJoinCondition(cond, l, r) + if (parsed.isEmpty) return None + Some(join) + } + + /** + * Pattern A: The LeftSemi/LeftAnti right child is itself an Inner self-join (with or without a + * wrapping Project). The semi-join condition is pure equi. We replace the right child with GROUP + * BY + HAVING COUNT(DISTINCT) > 1. + * + * Matches q95 structure: Join(LeftSemi, left = outer query (filtered web_sales + dims), right = + * Join(Inner, ws_L, ws_R, equi(order_number) AND neq(warehouse_sk)), condition = + * outer.order_number = inner.order_number) + */ + private def tryRewriteSemiWithSelfJoinChild( + original: Join, + left: LogicalPlan, + right: LogicalPlan, + joinType: JoinType, + semiCondition: Expression, + hint: JoinHint): Option[LogicalPlan] = { + + // Unwrap the right child: may be Project(Join(Inner,...)) or bare Join(Inner,...) + val (innerJoin, wrapper) = 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 + } + + // Use field accessors (portable across Spark versions) instead of + // case-class unapply with a fixed parameter list. The caller has already + // pattern-matched innerJoin as Join(_, _, Inner, Some(_), _). + val innerLeft = innerJoin.left + val innerRight = innerJoin.right + val innerCond = innerJoin.condition.get + + // Both sides of inner join must scan the same base relation + if (!isSameBaseRelation(innerLeft, innerRight)) return None + + // Parse inner join condition: must be equi + neq only + val parsed = parseSelfJoinCondition(innerCond, innerLeft, innerRight) + if (parsed.isEmpty) return None + val (innerEquiPairs, innerNeqPairs) = parsed.get + + // The semi-join condition must reference an output attribute from the inner join + // that corresponds to an equi-key. This is how the outer query connects to the CTE. + val semiPreds = splitConjunctivePredicates(semiCondition) + val rightOutputSet = right.outputSet + + // All semi predicates must be pure equi-join (no inequality in semi condition) + val semiEquiPairs = semiPreds.collect { + case EqualTo(l: Attribute, r: Attribute) + if left.outputSet.contains(l) && rightOutputSet.contains(r) => (l, r) + case EqualTo(r: Attribute, l: Attribute) + if left.outputSet.contains(l) && rightOutputSet.contains(r) => (l, r) + } + if (semiEquiPairs.size != semiPreds.size) return None + if (semiEquiPairs.isEmpty) return None + + // The semi-join right-side keys must be derivable from the inner join's equi-keys + val innerEquiLeftAttrIds = + innerEquiPairs.map(_._1).collect { case a: Attribute => a.exprId }.toSet + val innerEquiRightAttrIds = + innerEquiPairs.map(_._2).collect { case a: Attribute => a.exprId }.toSet + val innerEquiAllIds = innerEquiLeftAttrIds ++ innerEquiRightAttrIds + + // Check that all semi-join right keys reference inner equi-key attributes + val semiRightKeys = semiEquiPairs.map(_._2) + // After ColumnPruning, the right side might output only the equi-key columns + // The semi right keys must be from the inner join's equi-key set + val rightKeyIds = semiRightKeys.map(_.exprId).toSet + val validSemiKeys = rightKeyIds.forall { + id => + innerEquiAllIds.contains(id) || { + // Semi-key may be exposed via a wrapper Project. For an Alias, + // the alias's *output* ExprId must match the caller's `id`, AND + // the underlying attribute must be an inner equi-key. Using the + // alias's input ExprId here compared with `id` compares two + // unrelated ids (alias always mints a new output id). + wrapper.exists { + case Project(pl, _) => + pl.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 + } + case _ => false + } + } + } + if (!validSemiKeys) return None + + // Build replacement: GROUP BY equi_keys HAVING COUNT(DISTINCT neq_col) > 1 + val innerLeftEquiAttrs = innerEquiPairs.map(_._1).collect { case a: Attribute => a } + val innerNeqAttr = innerNeqPairs.head._1 match { + case a: Attribute => a + case _ => return None + } + + val countDistinctExpr = AggregateExpression( + Count(Seq(innerNeqAttr)), + mode = Complete, + isDistinct = true, + filter = None, + NamedExpression.newExprId) + val countAlias = Alias(countDistinctExpr, "_gluten_rw_selfjoin_cnt_distinct")() + + val groupingExprs: Seq[Expression] = innerLeftEquiAttrs + val aggExprs: Seq[NamedExpression] = + innerLeftEquiAttrs.map(_.asInstanceOf[NamedExpression]) :+ countAlias + val aggregate = Aggregate(groupingExprs, aggExprs, innerLeft) + + val filterExpr = GreaterThan(countAlias.toAttribute, Literal(1L, LongType)) + val filtered = Filter(filterExpr, aggregate) + + // Project only the equi-key columns (to match the original right-side output schema) + val projectedKeys = Project( + innerLeftEquiAttrs.map(_.asInstanceOf[NamedExpression]), + filtered) + + // Rebuild semi condition: map any exprId reachable from the old right side to + // the corresponding innerLeft equi attr. Cover three cases: + // (1) direct innerRight equi attr + // (2) innerLeft equi attr (identity) + // (3) wrapper Project's output attr backed by Alias(innerEquiAttr, _) - new exprId + val innerRightEquiAttrs = innerEquiPairs.map(_._2).collect { case a: Attribute => a } + val nameToInnerLeft: Map[String, Attribute] = + innerLeftEquiAttrs.map(a => a.name -> a).toMap + val wrapperRemap: Map[ExprId, Attribute] = wrapper match { + case Some(Project(pl, _)) => + pl.flatMap { + case al @ Alias(a: Attribute, _) + if (innerEquiLeftAttrIds.contains(a.exprId) || + innerEquiRightAttrIds.contains(a.exprId)) && + nameToInnerLeft.contains(a.name) => + Some(al.exprId -> nameToInnerLeft(a.name)) + case _ => None + }.toMap + case _ => Map.empty + } + val oldToNewMap: Map[ExprId, Attribute] = + innerRightEquiAttrs.map(_.exprId).zip(innerLeftEquiAttrs).toMap ++ + innerLeftEquiAttrs.map(a => a.exprId -> a).toMap ++ + wrapperRemap + + // Safety: every attr in semiCondition that came from the old right output must + // resolve to something in oldToNewMap. If not, refuse to rewrite. + 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) + } + + val newJoin = original.copy( + right = projectedKeys, + condition = Some(newSemiCondition)) + + logInfo( + s"RewriteSelfJoinInequalityToAggregate: Pattern A - rewrote $joinType with Inner " + + s"self-join child to $joinType + GROUP BY HAVING COUNT(DISTINCT) > 1. " + + s"equiKeys=[${innerLeftEquiAttrs.map(_.name).mkString(",")}], " + + s"neqCol=${innerNeqAttr.name}") + + Some(newJoin) + } + + /** + * Parse a join condition into equi-pairs and inequality-pairs. Accepts only: EqualTo(attr, attr) + * and Not(EqualTo(attr, attr)). Ignores IsNotNull predicates (added by + * InferFiltersFromConstraints). Returns None if there are unrecognized predicates beyond equi + + * neq + IsNotNull. + */ + private def parseSelfJoinCondition( + condition: Expression, + left: LogicalPlan, + right: LogicalPlan): Option[(Seq[(Attribute, Attribute)], Seq[(Attribute, Attribute)])] = { + + val leftOutput = left.outputSet + val rightOutput = right.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 equi/neq attributes are safe to drop + // (they're redundant with the join semantics or auto-added by + // InferFiltersFromConstraints). IsNotNull on other columns would be + // silently dropped by the rewrite and change query semantics -- bail out. + val joinAttrIds: Set[org.apache.spark.sql.catalyst.expressions.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 + } + + // Accept only equi + neq + IsNotNull-on-join-cols; reject anything else. + 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: COUNT(DISTINCT single_col) > 1 semantics + // is strictly wider than "exists row with ALL of multiple != predicates". + if (neqPairs.size != 1) return None + + // Self-join invariant: equi and neq must be on same-named columns from both sides. + val equiValid = equiPairs.forall { case (l, r) => l.name == r.name } + val neqValid = neqPairs.forall { case (l, r) => l.name == r.name } + if (!equiValid || !neqValid) return None + + Some((equiPairs, neqPairs)) Review Comment: Thanks for flagging this. The claimed failure mode is actually not a correctness issue. If the inequality column is also a grouping/equi-key, e.g. t1.k = t2.k AND t1.k <> t2.k, the original join is unsatisfiable. The rewritten form is also empty: within each GROUP BY k group, COUNT(DISTINCT k) is at most 1, so HAVING COUNT(DISTINCT k) > 1 can never pass. That said, the rule already rejects this shape explicitly. The guard is intentionally fail-closed: it keeps the supported rewrite shape limited to cases where the inequality column is distinct from the equi-key columns, rather than relying on the degenerate-case equivalence above. -- 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]
