Copilot commented on code in PR #12756: URL: https://github.com/apache/gluten/pull/12756#discussion_r3765744620
########## backends-velox/src/main/scala/org/apache/gluten/extension/RewriteSelfJoinInequalityToAggregate.scala: ########## @@ -0,0 +1,646 @@ +/* + * 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 Review Comment: `nameToInnerLeft` is built with `.toMap` keyed by attribute name. If equi-key attributes contain duplicate names (possible with projections/aliases), `.toMap` will silently drop earlier entries and the remapping below can bind the wrong attribute. Since this rule is intended to be fail-closed, add a guard that equi-key names are distinct before constructing the map (or switch to ExprId-based mapping). This issue also appears in the following locations of the same file: - line 322 - line 509 -- 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]
