cloud-fan commented on code in PR #58424: URL: https://github.com/apache/spark/pull/58424#discussion_r4015377438
########## sql/core/src/main/scala/org/apache/spark/sql/execution/RewriteSelfJoinInequalityToAggregate.scala: ########## @@ -0,0 +1,579 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.execution + +import org.apache.spark.sql.catalyst.expressions._ +import org.apache.spark.sql.catalyst.expressions.aggregate._ +import org.apache.spark.sql.catalyst.plans._ +import org.apache.spark.sql.catalyst.plans.logical._ +import org.apache.spark.sql.catalyst.rules.Rule +import org.apache.spark.sql.catalyst.trees.TreePattern.IN_SUBQUERY +import org.apache.spark.sql.catalyst.util.CharVarcharUtils +import org.apache.spark.sql.execution.datasources.{HadoopFsRelation, LogicalRelation} +import org.apache.spark.sql.execution.datasources.parquet.ParquetFileFormat +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types._ + +/** + * Rewrites supported uncorrelated IN-subquery inequality self-joins into + * `GROUP BY + HAVING MIN(neq) <> MAX(neq)`, avoiding the self-join cross-product. + * + * Supports a direct self-join (Pattern A') and a self-join nested under an outer inner join + * (Pattern A2, where only the self-join child becomes an Aggregate). Unsupported and correlated + * shapes fail closed. + * + * Runs in `extendedOperatorOptimizationRules`, before `RewritePredicateSubquery` turns the + * predicate subquery into a semi/anti/existence join, so it only sees the uncorrelated + * `InSubquery` shape. + * + * Controlled by `spark.sql.optimizer.rewriteSelfJoinInequalityToAggregate.enabled` + * (default false, opt-in). + */ +object RewriteSelfJoinInequalityToAggregate extends Rule[LogicalPlan] with PredicateHelper { + + private val MinNeqAliasName = "_rewrite_selfjoin_inequality_min" + private val MaxNeqAliasName = "_rewrite_selfjoin_inequality_max" + + override def apply(plan: LogicalPlan): LogicalPlan = { + if (!conf.getConf(SQLConf.REWRITE_SELF_JOIN_INEQUALITY_TO_AGGREGATE_ENABLED)) { + return plan + } + + // Fail closed on correlated subqueries: `lq.children` holds the outer references this rule + // does not remap. + plan.transformAllExpressionsWithPruning(_.containsPattern(IN_SUBQUERY)) { + case in @ InSubquery(_, lq: ListQuery) if lq.children.isEmpty => + rewriteSubqueryPlan(lq.plan) match { + case Some(newSub) => in.copy(query = lq.copy(plan = newSub)) + case None => in + } + } + } + + // ============================================================================ + // Shared helpers + // ============================================================================ + + /** + * Build `Filter(min <> max, Aggregate(equiKeys, child))`, taking MIN and MAX over the neq column. + * `MIN(neqCol) <> MAX(neqCol)` is true exactly when the group holds two or more distinct non-null + * values. + * + * The `IsNotNull(equiKeys)` filter preserves the equi-join's NULL semantics: `=` never matches a + * NULL key, but GROUP BY would fold all NULL keys into one group that can leak NULL into a + * `NOT IN`. The neq column needs no filter -- MIN/MAX ignore NULL, and a group with fewer than + * two non-null values has `min = max` (or both NULL, which makes `<>` NULL), so `<>` is never + * true for it and the group is dropped. + * + * Sound only because the rule fires under `InSubquery`: IN / NOT IN membership is insensitive to + * duplicate rows in the subquery result, and it binds columns by position (`equalsStructurally`), + * so the A2 branch's column rename is harmless. + */ + private def buildAggregateHavingMultipleDistinct( + equiKeys: Seq[Attribute], + neqCol: Attribute, + child: LogicalPlan): LogicalPlan = { + val minAlias = Alias(Min(neqCol).toAggregateExpression(), MinNeqAliasName)() + val maxAlias = Alias(Max(neqCol).toAggregateExpression(), MaxNeqAliasName)() + val aggExprs: Seq[NamedExpression] = equiKeys :+ minAlias :+ maxAlias + val nonNullChild = equiKeys + .map(a => IsNotNull(a): Expression) + .reduceOption(And) + .map(Filter(_, child)) + .getOrElse(child) + val agg = Aggregate(equiKeys, aggExprs, nonNullChild) + Filter(Not(EqualTo(minAlias.toAttribute, maxAlias.toAttribute)), agg) + } + + /** + * Rebuild the wrapper Project so every equi-key reference points at the sjLeft attribute with a + * fresh output ExprId, returning `oldOutputExprId -> newOutputAttr` for downstream references + * (outer join condition, top-level Project). Lookup is by ExprId (Catalyst attribute identity), + * not name. Fails closed when an entry is neither an equi-key Attribute nor `Alias(equi-key, _)`. + */ + private def canonicalizeWrapper( + projectList: Seq[NamedExpression], + equiPairs: Seq[(Attribute, Attribute)], + newChild: LogicalPlan): Option[(Project, Map[ExprId, Attribute])] = { + val exprIdToLeft: Map[ExprId, Attribute] = + equiPairs.flatMap { case (l, r) => Seq(l.exprId -> l, r.exprId -> l) }.toMap + val oldOutput: Seq[Attribute] = projectList.map(_.toAttribute) + val mapped: Seq[Option[NamedExpression]] = projectList.map { + case a: Attribute if exprIdToLeft.contains(a.exprId) => + // Fresh exprId, but carry over qualifier / metadata so this branch stays consistent with + // the Alias branch and a column keeps its metadata. + Some( + Alias(exprIdToLeft(a.exprId), a.name)( + qualifier = a.qualifier, + explicitMetadata = Some(a.metadata)): NamedExpression) + case al @ Alias(a: Attribute, _) if exprIdToLeft.contains(a.exprId) => + // withNewChild preserves name/qualifier/metadata and exprId; newInstance then re-stamps a + // fresh exprId, so Alias keeps ownership of its own metadata contract instead of us + // re-listing its fields (which drift when Alias gains one). + Some(al.withNewChild(exprIdToLeft(a.exprId)).newInstance()) + case _ => None + } + if (mapped.exists(_.isEmpty)) { + None + } else { + val newProjectList = mapped.flatten + val newWrapper = Project(newProjectList, newChild) + val remap: Map[ExprId, Attribute] = + oldOutput.zip(newWrapper.output).map { case (o, n) => o.exprId -> n }.toMap + Some((newWrapper, remap)) + } + } + + /** + * Replace equi-key references inside a NamedExpression per `remap`, preserving Attribute/Alias + * shape. Any other expression still referencing a replaced output returns None (fail-closed) to + * avoid a dangling ExprId. + */ + private def remapNamedExpressionAttributes( + ne: NamedExpression, + remap: Map[ExprId, Attribute]): Option[NamedExpression] = ne match { + case a: Attribute if remap.contains(a.exprId) => Some(remap(a.exprId)) + case a: Attribute => Some(a) + case al: Alias => + val newChild = al.child.transformUp { + case a: Attribute if remap.contains(a.exprId) => remap(a.exprId) + } + // withNewChild preserves the same exprId/qualifier/metadata the manual copy did. + Some(if (newChild eq al.child) al else al.withNewChild(newChild)) + case other if other.references.exists(a => remap.contains(a.exprId)) => + None + case other => Some(other) + } + + // ============================================================================ + // Pattern A' / A2 dispatch (subquery plans of InSubquery) + // ============================================================================ + + private def rewriteSubqueryPlan(plan: LogicalPlan): Option[LogicalPlan] = { + // Match the candidate shape first -- a top-level Inner Join, optionally under one wrapper + // Project. This structural match is cheap, so run it before the whole-subquery + // `isRepeatablePlan` walk and skip that walk entirely for the many subqueries that are not + // even shaped like a self-join. + val (projectListOpt, innerJoin): (Option[Seq[NamedExpression]], Join) = plan match { + case Project(pl, j: Join) if j.joinType == Inner && j.condition.isDefined => + (Some(pl), j) + case j: Join if j.joinType == Inner && j.condition.isDefined => + (None, j) + case _ => return None + } + + // Candidate-level guard: reject if any node in the whole subquery is non-repeatable, catching + // nondeterminism hoisted above the self-join that the per-side `isSameBaseRelation` misses. + if (!isRepeatablePlan(plan)) return None + + if (isSameBaseRelation(innerJoin.left, innerJoin.right)) { + rewriteDirectSelfJoin(projectListOpt, innerJoin) + } else { + rewriteNestedSelfJoin(projectListOpt, innerJoin) + } + } + + // ============================================================================ + // Pattern A' : direct self-join at subquery top level + // ============================================================================ + + private def rewriteDirectSelfJoin( + projectListOpt: Option[Seq[NamedExpression]], + innerJoin: Join): Option[LogicalPlan] = { + // Fail closed on an explicit join hint: it is a directive about the join this rule deletes. + if (!innerJoin.hint.isEmpty) return None + + val innerLeft = innerJoin.left + val innerCond = innerJoin.condition.get + + val parsed = parseSelfJoinCondition(innerCond, innerLeft, innerJoin.right) + if (parsed.isEmpty) return None + val (equiPairs, neqPairs) = parsed.get + + val innerLeftEquiAttrs: Seq[Attribute] = equiPairs.map(_._1) + val innerLeftNeqAttr: Attribute = neqPairs.head._1 + val filtered = + buildAggregateHavingMultipleDistinct(innerLeftEquiAttrs, innerLeftNeqAttr, innerLeft) + + // Fail closed on a bare-Join subquery: with no wrapper Project, replacing the self-join output + // with `Project(equiKeys, filtered)` shrinks arity and RewritePredicateSubquery's positional + // `values.zip(sub.output)` would misbind semi predicates. Q95 subqueries always have a Project. + projectListOpt match { + case None => None + case Some(pl) => + canonicalizeWrapper(pl, equiPairs, filtered).map { case (newWrapper, _) => newWrapper } + } + } + + // ============================================================================ + // Pattern A2 : self-join nested inside another InnerJoin in the subquery + // ============================================================================ + + private def rewriteNestedSelfJoin( + projectListOpt: Option[Seq[NamedExpression]], + outerJoin: Join): Option[LogicalPlan] = { + val outerCond = outerJoin.condition.get + + val (selfJoinSide, selfJoinOnRight) = + tryExtractSelfJoin(outerJoin.right) match { + case Some(_) => (outerJoin.right, true) + case None => + tryExtractSelfJoin(outerJoin.left) match { + case Some(_) => (outerJoin.left, false) + case None => return None + } + } + + val (selfJoinProjectOpt, selfJoin) = selfJoinSide match { + case p @ Project(_, j: Join) if j.joinType == Inner && j.condition.isDefined => + (Some(p), j) + case j: Join if j.joinType == Inner && j.condition.isDefined => + (None, j) + case _ => return None + } + + val sjLeft = selfJoin.left + val sjCond = selfJoin.condition.get + if (!isSameBaseRelation(sjLeft, selfJoin.right)) return None + + val parsed = parseSelfJoinCondition(sjCond, sjLeft, selfJoin.right) + if (parsed.isEmpty) return None + val (equiPairs, neqPairs) = parsed.get + + val sjLeftEquiAttrs: Seq[Attribute] = equiPairs.map(_._1) + val sjLeftNeqAttr: Attribute = neqPairs.head._1 + + val selfJoinOutputSet = selfJoinSide.outputSet + val sjEquiExprIds: Set[ExprId] = + equiPairs.flatMap { case (l, r) => Seq(l.exprId, r.exprId) }.toSet + // A wrapper Project may reproject equi-keys under fresh alias exprIds; include those. + val wrapperEquiExprIds: Set[ExprId] = selfJoinProjectOpt.toSeq.flatMap { p => + p.projectList.flatMap { + case a: Attribute if sjEquiExprIds.contains(a.exprId) => Some(a.exprId) + case al @ Alias(a: Attribute, _) if sjEquiExprIds.contains(a.exprId) => Some(al.exprId) + case _ => None + } + }.toSet + val allEquiExprIds = sjEquiExprIds ++ wrapperEquiExprIds + + // The outer join condition and any top-level Project may reference only equi-key attrs from the + // self-join side (the neq column does not survive the rewrite). + val outerCondRefs = outerCond.references.filter(selfJoinOutputSet.contains) + if (!outerCondRefs.forall(a => allEquiExprIds.contains(a.exprId))) return None + val projectOk = projectListOpt.forall { pl => + val refs = pl.flatMap(_.references).filter(selfJoinOutputSet.contains) + refs.forall(a => allEquiExprIds.contains(a.exprId)) + } + if (!projectOk) return None + + val filtered = buildAggregateHavingMultipleDistinct(sjLeftEquiAttrs, sjLeftNeqAttr, sjLeft) + + val (newSelfJoinSide, outputRemap): (LogicalPlan, Map[ExprId, Attribute]) = + selfJoinProjectOpt match { + case Some(wp) => + canonicalizeWrapper(wp.projectList, equiPairs, filtered) match { + case Some((newWrapper, remap)) => (newWrapper, remap) + case None => return None + } + case None => + // A bare self-join side (no wrapper Project) is not produced for a fireable A2 by the + // normal optimizer pipeline: only equi keys are referenced above the self-join, so + // ColumnPruning inserts a wrapper Project to drop the unused neq column, leaving + // selfJoinProjectOpt = Some. Fail closed on the non-standard bare shape. + return None + } + + val newOuterCond = outerCond.transformUp { + case a: Attribute if outputRemap.contains(a.exprId) => outputRemap(a.exprId) + } + + val newOuterJoin = if (selfJoinOnRight) { + outerJoin.copy(right = newSelfJoinSide, condition = Some(newOuterCond)) + } else { + outerJoin.copy(left = newSelfJoinSide, condition = Some(newOuterCond)) + } + + projectListOpt match { + case Some(pl) => + val remapped = pl.map(ne => remapNamedExpressionAttributes(ne, outputRemap)) + if (remapped.exists(_.isEmpty)) return None + Some(Project(remapped.flatten, newOuterJoin)) + case None => Some(newOuterJoin) + } + } + + private def tryExtractSelfJoin(plan: LogicalPlan): Option[Join] = { + val join = plan match { + case Project(_, j: Join) if j.joinType == Inner && j.condition.isDefined => j + case j: Join if j.joinType == Inner && j.condition.isDefined => j + case _ => return None + } + // A hinted self-join is not an extraction candidate; see `rewriteDirectSelfJoin`. + if (!join.hint.isEmpty) return None + if (!isSameBaseRelation(join.left, join.right)) return None + if (parseSelfJoinCondition(join.condition.get, join.left, join.right).isEmpty) return None + Some(join) + } + + // ============================================================================ + // parseSelfJoinCondition + isSameBaseRelation + // ============================================================================ + + private def outputOrdinal(plan: LogicalPlan, attr: Attribute): Int = + plan.output.indexWhere(_.exprId == attr.exprId) + + private def sameOutputPosition( + leftPlan: LogicalPlan, + rightPlan: LogicalPlan, + leftAttr: Attribute, + rightAttr: Attribute): Boolean = { + val leftPos = outputOrdinal(leftPlan, leftAttr) + val rightPos = outputOrdinal(rightPlan, rightAttr) + leftPos >= 0 && rightPos >= 0 && leftPos == rightPos + } + + /** + * Parse a join condition into equi-pairs and inequality-pairs. Accepts only `EqualTo(attr, attr)` + * and `Not(EqualTo(attr, attr))` across opposite sides, and `IsNotNull(attr)` on a join column; + * anything else fails the whole rewrite closed. + */ + private def parseSelfJoinCondition( + condition: Expression, + leftPlan: LogicalPlan, + rightPlan: LogicalPlan) + : Option[(Seq[(Attribute, Attribute)], Seq[(Attribute, Attribute)])] = { + + val leftOutput = leftPlan.outputSet + val rightOutput = rightPlan.outputSet + val predicates = splitConjunctivePredicates(condition) + + val equiPairs = predicates.collect { + case EqualTo(l: Attribute, r: Attribute) + if leftOutput.contains(l) && rightOutput.contains(r) => + (l, r) + case EqualTo(r: Attribute, l: Attribute) + if leftOutput.contains(l) && rightOutput.contains(r) => + (l, r) + } + + val neqPairs = predicates.collect { + case Not(EqualTo(l: Attribute, r: Attribute)) + if leftOutput.contains(l) && rightOutput.contains(r) => + (l, r) + case Not(EqualTo(r: Attribute, l: Attribute)) + if leftOutput.contains(l) && rightOutput.contains(r) => + (l, r) + } + + // Only IsNotNull on a join column is safe to drop -- redundant with the join or auto-added by + // InferFiltersFromConstraints. IsNotNull on any other column changes semantics; bail out. + val joinAttrIds: Set[ExprId] = + (equiPairs ++ neqPairs).flatMap { case (l, r) => Seq(l.exprId, r.exprId) }.toSet + val isNotNullOnJoinCols = predicates.count { + case IsNotNull(a: Attribute) if joinAttrIds.contains(a.exprId) => true + case _ => false + } + + val totalMatched = equiPairs.size + neqPairs.size + isNotNullOnJoinCols + if (totalMatched != predicates.size) return None + if (equiPairs.isEmpty || neqPairs.isEmpty) return None + + // A single inequality only: MIN/MAX over one column cannot represent multiple neqs. + if (neqPairs.size != 1) return None + + // Equi-keys move into grouping equality, the neq column into MIN/MAX ordering equality -- two + // different gates (below). Check both ends of each pair, since canonicalization can drop the + // metadata the gate reads. Fail closed on anything not proven. + val equiAttrs = equiPairs.flatMap { case (l, r) => Seq(l, r) } + val neqAttrs = neqPairs.flatMap { case (l, r) => Seq(l, r) } + if (!equiAttrs.forall(isSafeEquiKeyAttribute)) return None + if (!neqAttrs.forall(isSafeNeqColumnAttribute)) return None + + // The rewrite expresses "two or more distinct values" as MIN(neq) <> MAX(neq), so the neq + // column must be orderable. The allowlist above already implies this, but assert Spark's own + // MIN/MAX input contract (RowOrdering.isOrderable, the same check Min/Max run) explicitly, so + // the requirement is visible at the rewrite site. + if (!RowOrdering.isOrderable(neqPairs.head._1.dataType)) return None + + // Resolve each predicate end by ExprId and require matching output ordinals, not name equality + // (canonicalization erases cosmetic Alias names). + val equiValid = + equiPairs.forall { case (l, r) => sameOutputPosition(leftPlan, rightPlan, l, r) } + val neqValid = neqPairs.forall { case (l, r) => sameOutputPosition(leftPlan, rightPlan, l, r) } + if (!equiValid || !neqValid) return None + + // Equi-key output positions must be distinct, so swapped/duplicate aliases cannot collide. + val leftEquiOrdinals = equiPairs.map { case (l, _) => outputOrdinal(leftPlan, l) } + if (leftEquiOrdinals.exists(_ < 0)) return None + if (leftEquiOrdinals.distinct.size != leftEquiOrdinals.size) return None + + // The neq column must map to an output position of the left self-join input. + val neqLeftOrdinal = outputOrdinal(leftPlan, neqPairs.head._1) + if (neqLeftOrdinal < 0) return None + Some((equiPairs, neqPairs)) + } + + // CHAR/VARCHAR reach the optimizer as StringType with the declared type in metadata; read it back + // (falling back to `dataType`) so they don't slip through the StringType branch of the gates. + private def rawType(attr: Attribute): DataType = + CharVarcharUtils.getRawType(attr.metadata).getOrElse(attr.dataType) + + private def isSafeEquiKeyAttribute(attr: Attribute): Boolean = + isSafeEquiKeyType(rawType(attr)) + + private def isSafeNeqColumnAttribute(attr: Attribute): Boolean = + isSafeNeqColumnType(rawType(attr)) + + // Allowlist for equi-keys, which move into GROUP BY: `=` must coincide with grouping equality; + // fail closed otherwise. Float/Double are excluded conservatively (NormalizeFloatingNumbers + // already reconciles NaN/signed zero), not out of necessity. + private def isSafeEquiKeyType(dt: DataType): Boolean = dt match { + case ByteType | ShortType | IntegerType | LongType => true + case _: DecimalType => true + case BooleanType => true + case DateType => true + case TimestampType | TimestampNTZType => true + case BinaryType => true + case _: CharType | _: VarcharType => false + case st: StringType if st.supportsBinaryEquality => true + case _ => false + } + + // Neq column moves into MIN/MAX, which compare by ORDERING, so a collated string needs + // supportsBinaryOrdering (not the weaker supportsBinaryEquality). + private def isSafeNeqColumnType(dt: DataType): Boolean = dt match { + case ByteType | ShortType | IntegerType | LongType => true + case _: DecimalType => true + case BooleanType => true + case DateType => true + case TimestampType | TimestampNTZType => true + case BinaryType => true + case _: CharType | _: VarcharType => false + case st: StringType if st.supportsBinaryOrdering => true + case _ => false + } + + /** + * Primary safety guard: the rewrite folds two occurrences of one subtree into a single aggregate, + * so a plan qualifies only when its operators, leaves and expressions are all allowlisted as + * repeatable. `plan.deterministic` alone is insufficient -- Aggregate(First), Window row_number + * over a non-total order and Limit/Sample are row-bag nondeterministic yet report deterministic. + * Embedded expression subqueries also fail closed. + */ + private def isRepeatablePlan(plan: LogicalPlan): Boolean = { + plan.deterministic && + !plan.isStreaming && + plan.subqueriesAll.isEmpty && + isRowBagRepeatable(plan) && + hasRepeatableExpressions(plan) + } + + /** + * Operator/leaf allowlist for repeatable row bags; everything unknown fails closed. Kept narrow: + * the target shape needs only a Parquet scan optionally wrapped in Project / Filter / + * SubqueryAlias plus the self-join. The exact-`ParquetFileFormat` leaf check is deliberate: + * other formats (ORC/JSON/CSV) reach the same scan but stay rejected until separately validated. + */ + private def isRowBagRepeatable(plan: LogicalPlan): Boolean = !plan.exists { + // Whitelisted operator => false ("does not break repeatability"); negating `exists` then means + // "every operator is whitelisted". + case _: Project => false + case _: Filter => false + case _: SubqueryAlias => false + case _: Join => false + case _: Range => false + case _: LocalRelation => false + case relation: LogicalRelation => + // Trust a Parquet scan only: exact `ParquetFileFormat` (getClass, not isInstanceOf, since it + // is non-final); any other FileFormat is not provably repeatable. + relation.relation match { + case h: HadoopFsRelation if h.fileFormat.getClass == classOf[ParquetFileFormat] => false + case _ => true + } + case _ => true + } + + private def hasRepeatableExpressions(plan: LogicalPlan): Boolean = { + !plan.exists(node => node.expressions.exists(expr => !isRepeatableExpression(expr))) + } + + /** + * Expression allowlist: repeatable only when the root type is allowlisted and all children are, + * so `Add(v, Abs(w))` is rejected. Unknown types fail closed (a missed optimization, not a bug). + * Decimal wrappers such as `PromotePrecision` / `CheckOverflow` are absent and may fail closed. + */ + private def isRepeatableExpression(expr: Expression): Boolean = expr match { + case _: Attribute | _: Literal => + true + // A clock-dependent STRING/VARIANT -> TIMESTAMP_LTZ cast is not repeatable at any + // nesting level; see castHasClockDependentStringToTimestamp. Fail closed. + case c: Cast if castHasClockDependentStringToTimestamp(c.child.dataType, c.dataType) => + false + case _: Alias | _: Cast | _: Add | _: Subtract | _: Multiply | _: Divide | _: Remainder | + _: And | _: Or | _: Not | _: EqualTo | _: EqualNullSafe | _: LessThan | + _: LessThanOrEqual | _: GreaterThan | _: GreaterThanOrEqual | _: IsNull | _: IsNotNull => + expr.children.forall(isRepeatableExpression) + case _ => + false + } + + /** + * True when casting `from` to `to` does a clock-dependent STRING -> TIMESTAMP_LTZ conversion at + * any nesting level (directly or inside array/map/struct casts). A time-only string's missing + * date comes from the runtime clock, so two folded scans can disagree; TIMESTAMP_NTZ is clock- + * free. The scalar decision is delegated to `Cast.needsTimeZone`. A VariantType source is handled + * separately: `Cast.needsTimeZone(VariantType, _)` is unconditionally true, but a Variant cast + * may parse a runtime String. + */ + private def castHasClockDependentStringToTimestamp(from: DataType, to: DataType): Boolean = + (from, to) match { + case (s: StringType, t) => Cast.needsTimeZone(s, t) + case (_: VariantType, t) => variantTargetHasClockDependentStringToTimestamp(t) + case (ArrayType(fromEl, _), ArrayType(toEl, _)) => + castHasClockDependentStringToTimestamp(fromEl, toEl) + case (MapType(fromKey, fromVal, _), MapType(toKey, toVal, _)) => + castHasClockDependentStringToTimestamp(fromKey, toKey) || + castHasClockDependentStringToTimestamp(fromVal, toVal) + case (StructType(fromFields), StructType(toFields)) + if fromFields.length == toFields.length => + fromFields.zip(toFields).exists { case (f, t) => + castHasClockDependentStringToTimestamp(f.dataType, t.dataType) + } + case _ => false Review Comment: **Non-blocking (P2):** This recursive check misses `TimeType` leaves. `Cast.needsTimeZone` treats TIME-to-TIMESTAMP_NTZ/LTZ as current-date dependent, while `ComputeCurrentTime` only stabilizes a direct timestamp-target cast; for an outer array/map/struct cast the two original scans can therefore evaluate on opposite sides of midnight, but the rewrite keeps one evaluation. Please reject these residual nested TIME-to-timestamp paths and add paired firing controls plus scalar/nested coverage. ########## sql/core/src/test/scala/org/apache/spark/sql/execution/RewriteSelfJoinInequalityToAggregateSuite.scala: ########## @@ -0,0 +1,1400 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.spark.sql.execution + +import org.apache.spark.SparkThrowable +import org.apache.spark.sql.{QueryTest, Row} +import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, EqualTo, InSubquery, ListQuery, Not} +import org.apache.spark.sql.catalyst.optimizer.ReorderJoin +import org.apache.spark.sql.catalyst.plans.{Inner, LeftOuter} +import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Filter, GlobalLimit, Join, LocalLimit, LogicalPlan} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.test.SharedSparkSession +import org.apache.spark.sql.types.{IntegerType, StructField, StructType} + +/** + * Correctness tests for [[RewriteSelfJoinInequalityToAggregate]]. + * + * Positive A' / A2 cases assert both result equivalence and that the rewrite actually fired. + * + * `assert(!ruleFired(plan))` on its own only proves the rewrite did not happen -- not that it was + * the guard under test that stopped it. A fixture whose two self-join sides are not structurally + * identical is rejected by `isSameBaseRelation` before any predicate is parsed, so such a test + * passes while covering nothing. Rejection paths are therefore tested as single-variable pairs: + * the same fixture and query shape, one control that must fire and one variant changing only the + * tested feature that must not. A firing control does not pin the rejection to a line, but rules + * out an unrelated fixture mismatch as why its partner was rejected. (The `LIMIT` case is such a + * pair: a Limit's expressions are all allowlisted, so the operator whitelist alone can reject it.) + * + * Self-joined fixtures are real tables, not temp views over VALUES. Spark deduplicates a self-join + * over a [[org.apache.spark.sql.catalyst.analysis.MultiInstanceRelation]] via `newInstance()`, + * which refreshes one side's ExprIds without inserting a rename-only Project, so both sides stay + * structurally identical. A temp view over VALUES cannot, and Spark renames one side with a Project Review Comment: **Nit (P3):** This explanation does not match Catalyst behavior. `LocalRelation` implements `MultiInstanceRelation.newInstance`; although a temp view over VALUES adds View/Project wrappers, the two sides canonicalize to identical LocalRelations and `sameResult` remains true. Please remove or correct this rationale; the existing fixtures do not need to be rewritten solely to use VALUES. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/RewriteSelfJoinInequalityToAggregate.scala: ########## @@ -0,0 +1,579 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.execution + +import org.apache.spark.sql.catalyst.expressions._ +import org.apache.spark.sql.catalyst.expressions.aggregate._ +import org.apache.spark.sql.catalyst.plans._ +import org.apache.spark.sql.catalyst.plans.logical._ +import org.apache.spark.sql.catalyst.rules.Rule +import org.apache.spark.sql.catalyst.trees.TreePattern.IN_SUBQUERY +import org.apache.spark.sql.catalyst.util.CharVarcharUtils +import org.apache.spark.sql.execution.datasources.{HadoopFsRelation, LogicalRelation} +import org.apache.spark.sql.execution.datasources.parquet.ParquetFileFormat +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types._ + +/** + * Rewrites supported uncorrelated IN-subquery inequality self-joins into + * `GROUP BY + HAVING MIN(neq) <> MAX(neq)`, avoiding the self-join cross-product. + * + * Supports a direct self-join (Pattern A') and a self-join nested under an outer inner join + * (Pattern A2, where only the self-join child becomes an Aggregate). Unsupported and correlated + * shapes fail closed. + * + * Runs in `extendedOperatorOptimizationRules`, before `RewritePredicateSubquery` turns the + * predicate subquery into a semi/anti/existence join, so it only sees the uncorrelated + * `InSubquery` shape. + * + * Controlled by `spark.sql.optimizer.rewriteSelfJoinInequalityToAggregate.enabled` + * (default false, opt-in). + */ +object RewriteSelfJoinInequalityToAggregate extends Rule[LogicalPlan] with PredicateHelper { + + private val MinNeqAliasName = "_rewrite_selfjoin_inequality_min" + private val MaxNeqAliasName = "_rewrite_selfjoin_inequality_max" + + override def apply(plan: LogicalPlan): LogicalPlan = { + if (!conf.getConf(SQLConf.REWRITE_SELF_JOIN_INEQUALITY_TO_AGGREGATE_ENABLED)) { + return plan + } + + // Fail closed on correlated subqueries: `lq.children` holds the outer references this rule + // does not remap. + plan.transformAllExpressionsWithPruning(_.containsPattern(IN_SUBQUERY)) { + case in @ InSubquery(_, lq: ListQuery) if lq.children.isEmpty => + rewriteSubqueryPlan(lq.plan) match { + case Some(newSub) => in.copy(query = lq.copy(plan = newSub)) + case None => in + } + } + } + + // ============================================================================ + // Shared helpers + // ============================================================================ + + /** + * Build `Filter(min <> max, Aggregate(equiKeys, child))`, taking MIN and MAX over the neq column. + * `MIN(neqCol) <> MAX(neqCol)` is true exactly when the group holds two or more distinct non-null + * values. + * + * The `IsNotNull(equiKeys)` filter preserves the equi-join's NULL semantics: `=` never matches a + * NULL key, but GROUP BY would fold all NULL keys into one group that can leak NULL into a + * `NOT IN`. The neq column needs no filter -- MIN/MAX ignore NULL, and a group with fewer than + * two non-null values has `min = max` (or both NULL, which makes `<>` NULL), so `<>` is never + * true for it and the group is dropped. + * + * Sound only because the rule fires under `InSubquery`: IN / NOT IN membership is insensitive to + * duplicate rows in the subquery result, and it binds columns by position (`equalsStructurally`), + * so the A2 branch's column rename is harmless. + */ + private def buildAggregateHavingMultipleDistinct( + equiKeys: Seq[Attribute], + neqCol: Attribute, + child: LogicalPlan): LogicalPlan = { + val minAlias = Alias(Min(neqCol).toAggregateExpression(), MinNeqAliasName)() + val maxAlias = Alias(Max(neqCol).toAggregateExpression(), MaxNeqAliasName)() + val aggExprs: Seq[NamedExpression] = equiKeys :+ minAlias :+ maxAlias + val nonNullChild = equiKeys + .map(a => IsNotNull(a): Expression) + .reduceOption(And) + .map(Filter(_, child)) + .getOrElse(child) + val agg = Aggregate(equiKeys, aggExprs, nonNullChild) + Filter(Not(EqualTo(minAlias.toAttribute, maxAlias.toAttribute)), agg) + } + + /** + * Rebuild the wrapper Project so every equi-key reference points at the sjLeft attribute with a + * fresh output ExprId, returning `oldOutputExprId -> newOutputAttr` for downstream references + * (outer join condition, top-level Project). Lookup is by ExprId (Catalyst attribute identity), + * not name. Fails closed when an entry is neither an equi-key Attribute nor `Alias(equi-key, _)`. + */ + private def canonicalizeWrapper( + projectList: Seq[NamedExpression], + equiPairs: Seq[(Attribute, Attribute)], + newChild: LogicalPlan): Option[(Project, Map[ExprId, Attribute])] = { + val exprIdToLeft: Map[ExprId, Attribute] = + equiPairs.flatMap { case (l, r) => Seq(l.exprId -> l, r.exprId -> l) }.toMap + val oldOutput: Seq[Attribute] = projectList.map(_.toAttribute) + val mapped: Seq[Option[NamedExpression]] = projectList.map { + case a: Attribute if exprIdToLeft.contains(a.exprId) => + // Fresh exprId, but carry over qualifier / metadata so this branch stays consistent with + // the Alias branch and a column keeps its metadata. + Some( + Alias(exprIdToLeft(a.exprId), a.name)( + qualifier = a.qualifier, + explicitMetadata = Some(a.metadata)): NamedExpression) + case al @ Alias(a: Attribute, _) if exprIdToLeft.contains(a.exprId) => + // withNewChild preserves name/qualifier/metadata and exprId; newInstance then re-stamps a + // fresh exprId, so Alias keeps ownership of its own metadata contract instead of us + // re-listing its fields (which drift when Alias gains one). + Some(al.withNewChild(exprIdToLeft(a.exprId)).newInstance()) + case _ => None + } + if (mapped.exists(_.isEmpty)) { + None + } else { + val newProjectList = mapped.flatten + val newWrapper = Project(newProjectList, newChild) + val remap: Map[ExprId, Attribute] = + oldOutput.zip(newWrapper.output).map { case (o, n) => o.exprId -> n }.toMap + Some((newWrapper, remap)) + } + } + + /** + * Replace equi-key references inside a NamedExpression per `remap`, preserving Attribute/Alias + * shape. Any other expression still referencing a replaced output returns None (fail-closed) to + * avoid a dangling ExprId. + */ + private def remapNamedExpressionAttributes( + ne: NamedExpression, + remap: Map[ExprId, Attribute]): Option[NamedExpression] = ne match { + case a: Attribute if remap.contains(a.exprId) => Some(remap(a.exprId)) + case a: Attribute => Some(a) + case al: Alias => + val newChild = al.child.transformUp { + case a: Attribute if remap.contains(a.exprId) => remap(a.exprId) + } + // withNewChild preserves the same exprId/qualifier/metadata the manual copy did. + Some(if (newChild eq al.child) al else al.withNewChild(newChild)) + case other if other.references.exists(a => remap.contains(a.exprId)) => + None + case other => Some(other) + } + + // ============================================================================ + // Pattern A' / A2 dispatch (subquery plans of InSubquery) + // ============================================================================ + + private def rewriteSubqueryPlan(plan: LogicalPlan): Option[LogicalPlan] = { + // Match the candidate shape first -- a top-level Inner Join, optionally under one wrapper + // Project. This structural match is cheap, so run it before the whole-subquery + // `isRepeatablePlan` walk and skip that walk entirely for the many subqueries that are not + // even shaped like a self-join. + val (projectListOpt, innerJoin): (Option[Seq[NamedExpression]], Join) = plan match { + case Project(pl, j: Join) if j.joinType == Inner && j.condition.isDefined => + (Some(pl), j) + case j: Join if j.joinType == Inner && j.condition.isDefined => + (None, j) + case _ => return None + } + + // Candidate-level guard: reject if any node in the whole subquery is non-repeatable, catching + // nondeterminism hoisted above the self-join that the per-side `isSameBaseRelation` misses. + if (!isRepeatablePlan(plan)) return None + + if (isSameBaseRelation(innerJoin.left, innerJoin.right)) { + rewriteDirectSelfJoin(projectListOpt, innerJoin) + } else { + rewriteNestedSelfJoin(projectListOpt, innerJoin) + } + } + + // ============================================================================ + // Pattern A' : direct self-join at subquery top level + // ============================================================================ + + private def rewriteDirectSelfJoin( + projectListOpt: Option[Seq[NamedExpression]], + innerJoin: Join): Option[LogicalPlan] = { + // Fail closed on an explicit join hint: it is a directive about the join this rule deletes. + if (!innerJoin.hint.isEmpty) return None + + val innerLeft = innerJoin.left + val innerCond = innerJoin.condition.get + + val parsed = parseSelfJoinCondition(innerCond, innerLeft, innerJoin.right) + if (parsed.isEmpty) return None + val (equiPairs, neqPairs) = parsed.get + + val innerLeftEquiAttrs: Seq[Attribute] = equiPairs.map(_._1) + val innerLeftNeqAttr: Attribute = neqPairs.head._1 + val filtered = + buildAggregateHavingMultipleDistinct(innerLeftEquiAttrs, innerLeftNeqAttr, innerLeft) + + // Fail closed on a bare-Join subquery: with no wrapper Project, replacing the self-join output + // with `Project(equiKeys, filtered)` shrinks arity and RewritePredicateSubquery's positional + // `values.zip(sub.output)` would misbind semi predicates. Q95 subqueries always have a Project. + projectListOpt match { + case None => None + case Some(pl) => + canonicalizeWrapper(pl, equiPairs, filtered).map { case (newWrapper, _) => newWrapper } + } + } + + // ============================================================================ + // Pattern A2 : self-join nested inside another InnerJoin in the subquery + // ============================================================================ + + private def rewriteNestedSelfJoin( + projectListOpt: Option[Seq[NamedExpression]], + outerJoin: Join): Option[LogicalPlan] = { + val outerCond = outerJoin.condition.get + + val (selfJoinSide, selfJoinOnRight) = + tryExtractSelfJoin(outerJoin.right) match { + case Some(_) => (outerJoin.right, true) + case None => + tryExtractSelfJoin(outerJoin.left) match { + case Some(_) => (outerJoin.left, false) + case None => return None + } + } + + val (selfJoinProjectOpt, selfJoin) = selfJoinSide match { + case p @ Project(_, j: Join) if j.joinType == Inner && j.condition.isDefined => + (Some(p), j) + case j: Join if j.joinType == Inner && j.condition.isDefined => + (None, j) + case _ => return None + } + + val sjLeft = selfJoin.left + val sjCond = selfJoin.condition.get + if (!isSameBaseRelation(sjLeft, selfJoin.right)) return None + + val parsed = parseSelfJoinCondition(sjCond, sjLeft, selfJoin.right) + if (parsed.isEmpty) return None + val (equiPairs, neqPairs) = parsed.get + + val sjLeftEquiAttrs: Seq[Attribute] = equiPairs.map(_._1) + val sjLeftNeqAttr: Attribute = neqPairs.head._1 + + val selfJoinOutputSet = selfJoinSide.outputSet + val sjEquiExprIds: Set[ExprId] = + equiPairs.flatMap { case (l, r) => Seq(l.exprId, r.exprId) }.toSet + // A wrapper Project may reproject equi-keys under fresh alias exprIds; include those. + val wrapperEquiExprIds: Set[ExprId] = selfJoinProjectOpt.toSeq.flatMap { p => + p.projectList.flatMap { + case a: Attribute if sjEquiExprIds.contains(a.exprId) => Some(a.exprId) + case al @ Alias(a: Attribute, _) if sjEquiExprIds.contains(a.exprId) => Some(al.exprId) + case _ => None + } + }.toSet + val allEquiExprIds = sjEquiExprIds ++ wrapperEquiExprIds + + // The outer join condition and any top-level Project may reference only equi-key attrs from the + // self-join side (the neq column does not survive the rewrite). + val outerCondRefs = outerCond.references.filter(selfJoinOutputSet.contains) + if (!outerCondRefs.forall(a => allEquiExprIds.contains(a.exprId))) return None + val projectOk = projectListOpt.forall { pl => + val refs = pl.flatMap(_.references).filter(selfJoinOutputSet.contains) + refs.forall(a => allEquiExprIds.contains(a.exprId)) + } + if (!projectOk) return None + + val filtered = buildAggregateHavingMultipleDistinct(sjLeftEquiAttrs, sjLeftNeqAttr, sjLeft) + + val (newSelfJoinSide, outputRemap): (LogicalPlan, Map[ExprId, Attribute]) = + selfJoinProjectOpt match { + case Some(wp) => + canonicalizeWrapper(wp.projectList, equiPairs, filtered) match { + case Some((newWrapper, remap)) => (newWrapper, remap) + case None => return None + } + case None => + // A bare self-join side (no wrapper Project) is not produced for a fireable A2 by the + // normal optimizer pipeline: only equi keys are referenced above the self-join, so + // ColumnPruning inserts a wrapper Project to drop the unused neq column, leaving + // selfJoinProjectOpt = Some. Fail closed on the non-standard bare shape. + return None + } + + val newOuterCond = outerCond.transformUp { + case a: Attribute if outputRemap.contains(a.exprId) => outputRemap(a.exprId) + } + + val newOuterJoin = if (selfJoinOnRight) { + outerJoin.copy(right = newSelfJoinSide, condition = Some(newOuterCond)) + } else { + outerJoin.copy(left = newSelfJoinSide, condition = Some(newOuterCond)) + } + + projectListOpt match { + case Some(pl) => + val remapped = pl.map(ne => remapNamedExpressionAttributes(ne, outputRemap)) + if (remapped.exists(_.isEmpty)) return None + Some(Project(remapped.flatten, newOuterJoin)) + case None => Some(newOuterJoin) + } + } + + private def tryExtractSelfJoin(plan: LogicalPlan): Option[Join] = { + val join = plan match { + case Project(_, j: Join) if j.joinType == Inner && j.condition.isDefined => j + case j: Join if j.joinType == Inner && j.condition.isDefined => j + case _ => return None + } + // A hinted self-join is not an extraction candidate; see `rewriteDirectSelfJoin`. + if (!join.hint.isEmpty) return None + if (!isSameBaseRelation(join.left, join.right)) return None + if (parseSelfJoinCondition(join.condition.get, join.left, join.right).isEmpty) return None + Some(join) + } + + // ============================================================================ + // parseSelfJoinCondition + isSameBaseRelation + // ============================================================================ + + private def outputOrdinal(plan: LogicalPlan, attr: Attribute): Int = + plan.output.indexWhere(_.exprId == attr.exprId) + + private def sameOutputPosition( + leftPlan: LogicalPlan, + rightPlan: LogicalPlan, + leftAttr: Attribute, + rightAttr: Attribute): Boolean = { + val leftPos = outputOrdinal(leftPlan, leftAttr) + val rightPos = outputOrdinal(rightPlan, rightAttr) + leftPos >= 0 && rightPos >= 0 && leftPos == rightPos + } + + /** + * Parse a join condition into equi-pairs and inequality-pairs. Accepts only `EqualTo(attr, attr)` + * and `Not(EqualTo(attr, attr))` across opposite sides, and `IsNotNull(attr)` on a join column; + * anything else fails the whole rewrite closed. + */ + private def parseSelfJoinCondition( + condition: Expression, + leftPlan: LogicalPlan, + rightPlan: LogicalPlan) + : Option[(Seq[(Attribute, Attribute)], Seq[(Attribute, Attribute)])] = { + + val leftOutput = leftPlan.outputSet + val rightOutput = rightPlan.outputSet + val predicates = splitConjunctivePredicates(condition) + + val equiPairs = predicates.collect { + case EqualTo(l: Attribute, r: Attribute) + if leftOutput.contains(l) && rightOutput.contains(r) => + (l, r) + case EqualTo(r: Attribute, l: Attribute) + if leftOutput.contains(l) && rightOutput.contains(r) => + (l, r) + } + + val neqPairs = predicates.collect { + case Not(EqualTo(l: Attribute, r: Attribute)) + if leftOutput.contains(l) && rightOutput.contains(r) => + (l, r) + case Not(EqualTo(r: Attribute, l: Attribute)) + if leftOutput.contains(l) && rightOutput.contains(r) => + (l, r) + } + + // Only IsNotNull on a join column is safe to drop -- redundant with the join or auto-added by + // InferFiltersFromConstraints. IsNotNull on any other column changes semantics; bail out. + val joinAttrIds: Set[ExprId] = + (equiPairs ++ neqPairs).flatMap { case (l, r) => Seq(l.exprId, r.exprId) }.toSet + val isNotNullOnJoinCols = predicates.count { + case IsNotNull(a: Attribute) if joinAttrIds.contains(a.exprId) => true + case _ => false + } + + val totalMatched = equiPairs.size + neqPairs.size + isNotNullOnJoinCols + if (totalMatched != predicates.size) return None + if (equiPairs.isEmpty || neqPairs.isEmpty) return None + + // A single inequality only: MIN/MAX over one column cannot represent multiple neqs. + if (neqPairs.size != 1) return None + + // Equi-keys move into grouping equality, the neq column into MIN/MAX ordering equality -- two + // different gates (below). Check both ends of each pair, since canonicalization can drop the + // metadata the gate reads. Fail closed on anything not proven. + val equiAttrs = equiPairs.flatMap { case (l, r) => Seq(l, r) } + val neqAttrs = neqPairs.flatMap { case (l, r) => Seq(l, r) } + if (!equiAttrs.forall(isSafeEquiKeyAttribute)) return None + if (!neqAttrs.forall(isSafeNeqColumnAttribute)) return None + + // The rewrite expresses "two or more distinct values" as MIN(neq) <> MAX(neq), so the neq + // column must be orderable. The allowlist above already implies this, but assert Spark's own + // MIN/MAX input contract (RowOrdering.isOrderable, the same check Min/Max run) explicitly, so + // the requirement is visible at the rewrite site. + if (!RowOrdering.isOrderable(neqPairs.head._1.dataType)) return None + + // Resolve each predicate end by ExprId and require matching output ordinals, not name equality + // (canonicalization erases cosmetic Alias names). + val equiValid = + equiPairs.forall { case (l, r) => sameOutputPosition(leftPlan, rightPlan, l, r) } + val neqValid = neqPairs.forall { case (l, r) => sameOutputPosition(leftPlan, rightPlan, l, r) } + if (!equiValid || !neqValid) return None + + // Equi-key output positions must be distinct, so swapped/duplicate aliases cannot collide. + val leftEquiOrdinals = equiPairs.map { case (l, _) => outputOrdinal(leftPlan, l) } + if (leftEquiOrdinals.exists(_ < 0)) return None + if (leftEquiOrdinals.distinct.size != leftEquiOrdinals.size) return None + + // The neq column must map to an output position of the left self-join input. + val neqLeftOrdinal = outputOrdinal(leftPlan, neqPairs.head._1) + if (neqLeftOrdinal < 0) return None + Some((equiPairs, neqPairs)) + } + + // CHAR/VARCHAR reach the optimizer as StringType with the declared type in metadata; read it back + // (falling back to `dataType`) so they don't slip through the StringType branch of the gates. + private def rawType(attr: Attribute): DataType = + CharVarcharUtils.getRawType(attr.metadata).getOrElse(attr.dataType) + + private def isSafeEquiKeyAttribute(attr: Attribute): Boolean = + isSafeEquiKeyType(rawType(attr)) + + private def isSafeNeqColumnAttribute(attr: Attribute): Boolean = + isSafeNeqColumnType(rawType(attr)) + + // Allowlist for equi-keys, which move into GROUP BY: `=` must coincide with grouping equality; + // fail closed otherwise. Float/Double are excluded conservatively (NormalizeFloatingNumbers + // already reconciles NaN/signed zero), not out of necessity. + private def isSafeEquiKeyType(dt: DataType): Boolean = dt match { + case ByteType | ShortType | IntegerType | LongType => true + case _: DecimalType => true + case BooleanType => true + case DateType => true + case TimestampType | TimestampNTZType => true + case BinaryType => true + case _: CharType | _: VarcharType => false + case st: StringType if st.supportsBinaryEquality => true + case _ => false + } + + // Neq column moves into MIN/MAX, which compare by ORDERING, so a collated string needs + // supportsBinaryOrdering (not the weaker supportsBinaryEquality). + private def isSafeNeqColumnType(dt: DataType): Boolean = dt match { + case ByteType | ShortType | IntegerType | LongType => true + case _: DecimalType => true + case BooleanType => true + case DateType => true + case TimestampType | TimestampNTZType => true + case BinaryType => true + case _: CharType | _: VarcharType => false + case st: StringType if st.supportsBinaryOrdering => true + case _ => false + } + + /** + * Primary safety guard: the rewrite folds two occurrences of one subtree into a single aggregate, + * so a plan qualifies only when its operators, leaves and expressions are all allowlisted as + * repeatable. `plan.deterministic` alone is insufficient -- Aggregate(First), Window row_number + * over a non-total order and Limit/Sample are row-bag nondeterministic yet report deterministic. + * Embedded expression subqueries also fail closed. + */ + private def isRepeatablePlan(plan: LogicalPlan): Boolean = { + plan.deterministic && + !plan.isStreaming && + plan.subqueriesAll.isEmpty && + isRowBagRepeatable(plan) && + hasRepeatableExpressions(plan) + } + + /** + * Operator/leaf allowlist for repeatable row bags; everything unknown fails closed. Kept narrow: + * the target shape needs only a Parquet scan optionally wrapped in Project / Filter / + * SubqueryAlias plus the self-join. The exact-`ParquetFileFormat` leaf check is deliberate: + * other formats (ORC/JSON/CSV) reach the same scan but stay rejected until separately validated. + */ + private def isRowBagRepeatable(plan: LogicalPlan): Boolean = !plan.exists { + // Whitelisted operator => false ("does not break repeatability"); negating `exists` then means + // "every operator is whitelisted". + case _: Project => false + case _: Filter => false + case _: SubqueryAlias => false + case _: Join => false + case _: Range => false + case _: LocalRelation => false + case relation: LogicalRelation => + // Trust a Parquet scan only: exact `ParquetFileFormat` (getClass, not isInstanceOf, since it + // is non-final); any other FileFormat is not provably repeatable. + relation.relation match { + case h: HadoopFsRelation if h.fileFormat.getClass == classOf[ParquetFileFormat] => false Review Comment: **Non-blocking (P2):** `sameResult` is not sufficient to prove the same Parquet snapshot here: `InMemoryFileIndex.equals` compares root paths but not the cached file list. A view captured before an append and another captured afterward can therefore pass this check even though the original join reads different row bags; aggregating only the first side can remove an IN match. Please additionally require the two Parquet leaves to share the same analyzed scan/FileIndex snapshot identity, with a same-root different-snapshot regression and an ordinary self-join positive control. ########## sql/core/src/test/scala/org/apache/spark/sql/execution/RewriteSelfJoinInequalityToAggregateSuite.scala: ########## @@ -0,0 +1,1400 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.spark.sql.execution + +import org.apache.spark.SparkThrowable +import org.apache.spark.sql.{QueryTest, Row} +import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, EqualTo, InSubquery, ListQuery, Not} +import org.apache.spark.sql.catalyst.optimizer.ReorderJoin +import org.apache.spark.sql.catalyst.plans.{Inner, LeftOuter} +import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Filter, GlobalLimit, Join, LocalLimit, LogicalPlan} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.test.SharedSparkSession +import org.apache.spark.sql.types.{IntegerType, StructField, StructType} + +/** + * Correctness tests for [[RewriteSelfJoinInequalityToAggregate]]. + * + * Positive A' / A2 cases assert both result equivalence and that the rewrite actually fired. + * + * `assert(!ruleFired(plan))` on its own only proves the rewrite did not happen -- not that it was + * the guard under test that stopped it. A fixture whose two self-join sides are not structurally + * identical is rejected by `isSameBaseRelation` before any predicate is parsed, so such a test + * passes while covering nothing. Rejection paths are therefore tested as single-variable pairs: + * the same fixture and query shape, one control that must fire and one variant changing only the + * tested feature that must not. A firing control does not pin the rejection to a line, but rules + * out an unrelated fixture mismatch as why its partner was rejected. (The `LIMIT` case is such a + * pair: a Limit's expressions are all allowlisted, so the operator whitelist alone can reject it.) + * + * Self-joined fixtures are real tables, not temp views over VALUES. Spark deduplicates a self-join + * over a [[org.apache.spark.sql.catalyst.analysis.MultiInstanceRelation]] via `newInstance()`, + * which refreshes one side's ExprIds without inserting a rename-only Project, so both sides stay + * structurally identical. A temp view over VALUES cannot, and Spark renames one side with a Project + * instead, which would make `isSameBaseRelation` false for every self-join below. `range()` needs + * no such treatment -- Range is a MultiInstanceRelation already. + */ +class RewriteSelfJoinInequalityToAggregateSuite extends QueryTest with SharedSparkSession { + + private val rewriteConf = SQLConf.REWRITE_SELF_JOIN_INEQUALITY_TO_AGGREGATE_ENABLED.key + + /** Signature aliases produced by the rewrite; presence of both => rule definitely fired. */ + private val MinNeqAlias = "_rewrite_selfjoin_inequality_min" + private val MaxNeqAlias = "_rewrite_selfjoin_inequality_max" + + // Descends into subqueries: `QueryPlan.exists` does not, and the rewrite's signature alias lives + // inside the IN-subquery when the rule runs on an analyzed (not-yet-rewritten) plan. + private def hasAlias(plan: LogicalPlan, name: String): Boolean = + plan.collectFirstWithSubqueries { + case p if p.expressions.exists(_.exists { + case a: Alias if a.name == name => true + case _ => false + }) => () + }.isDefined + + /** Require BOTH aliases: a rewrite that emitted MIN but dropped MAX is still a bug. */ + private def ruleFired(plan: LogicalPlan): Boolean = + hasAlias(plan, MinNeqAlias) && hasAlias(plan, MaxNeqAlias) + + /** Total count of a signature alias across the plan and its subqueries: one per rewrite. */ + private def countAlias(plan: LogicalPlan, name: String): Int = + plan.collectWithSubqueries { + case p => p.expressions.map(_.collect { case a: Alias if a.name == name => a }.size).sum + }.sum + + private def assertRuleFired(sql: String): Unit = { + withSQLConf(rewriteConf -> "true") { + val plan = spark.sql(sql).queryExecution.optimizedPlan + assert(ruleFired(plan), s"self-join inequality rewrite should fire:\n$plan") + } + } + + private def assertRuleNotFired(sql: String): Unit = { + withSQLConf(rewriteConf -> "true") { + val plan = spark.sql(sql).queryExecution.optimizedPlan + assert(!ruleFired(plan), s"self-join inequality rewrite must not fire:\n$plan") + } + } + + /** + * Optimize just the IN-subquery plan (rewrite left at its default-off) and return the result, so + * a test can prove what shape the subquery reaches the rule as -- e.g. a bare Join with no + * wrapper Project -- before asserting the rule declines it. Without this, `assertRuleNotFired` + * alone can pass merely because the fixture never produced the shape the guard means to reject. + */ + private def optimizedInSubqueryPlan(sql: String): LogicalPlan = { + val analyzed = spark.sql(sql).queryExecution.analyzed + val subqueries = analyzed.subqueriesAll + assert( + subqueries.length == 1, + s"expected exactly one subquery in analyzed plan, got ${subqueries.length}:\n$analyzed") + spark.sessionState.optimizer.execute(subqueries.head) + } + + /** + * A real table, so that a self-join of it dedups into two structurally identical sides. See the + * class comment for why a temp view over VALUES cannot be used for a self-joined fixture. + */ + private def createTable(name: String, schema: String, values: String): Unit = { + spark.sql(s"DROP TABLE IF EXISTS $name") + spark.sql(s"CREATE TABLE $name($schema) USING parquet") + spark.sql(s"INSERT INTO $name SELECT * FROM VALUES $values") + } + + /** + * Run `sql` twice, first with rewrite ON then OFF. The rewrite must not change the outer query's + * row multiplicity, so assert ON and OFF agree as MULTISETS (a `.toSet` here would hide a + * duplicated or dropped row) before handing callers the row sets their fixed-value assertions + * compare against. `QueryTest.sameRows` is Spark's own multiset comparison (order-insensitive, + * duplicate-sensitive) and formats the offending rows on mismatch. + */ + private def runBoth(sql: String): (Set[Row], Set[Row]) = { + val on = withSQLConf(rewriteConf -> "true") { + spark.sql(sql).collect().toSeq + } + val off = withSQLConf(rewriteConf -> "false") { + spark.sql(sql).collect().toSeq + } + QueryTest.sameRows(on, off).foreach { error => + fail(s"rewrite changed row multiplicity between ON and OFF:\n$error") + } + (on.toSet, off.toSet) + } + + private def setupTable(): Unit = { + // k=1: distinct v={10,20} -> matches (has 2 non-null distinct) + // k=2: distinct v={30} -> no match (only 1) + // k=3: distinct v={40,50,60} -> matches + // k=4: v={70, NULL} -> no match (only 1 non-null) + // k=5: v={NULL, NULL} -> no match (0 non-null) + // k=6: v={80, 90, NULL} -> matches + // k=7: v={100,100} -> no match: duplicate-only, min(v)==max(v)==100 so min<>max is + // false. A plain COUNT(v) > 1 would wrongly match this group. + createTable( + "T", + "k INT, v INT", + """ (1, 10), (1, 10), (1, 20), + | (2, 30), + | (3, 40), (3, 50), (3, 60), + | (4, 70), (4, CAST(NULL AS INT)), + | (5, CAST(NULL AS INT)), (5, CAST(NULL AS INT)), + | (6, 80), (6, 90), (6, CAST(NULL AS INT)), + | (7, 100), (7, 100)""".stripMargin + ) + } + + // ==================== Positive: rewrite fires and is semantically equivalent =============== + + test("Pattern A': direct InSubquery self-join is rewritten") { + withTable("T") { + setupTable() + val sql = + """SELECT k FROM T outer_t WHERE k IN ( + | SELECT s1.k FROM T s1 JOIN T s2 + | ON s1.k = s2.k AND s1.v <> s2.v)""".stripMargin + + assertRuleFired(sql) + assertMinMaxRewriteShape(optimizedPlanWith(sql, rewrite = true)) + val (on, off) = runBoth(sql) + assert(on == off, s"rewrite ON $on != OFF $off") + // `setupTable()` seeds duplicate-only groups and NULL neq values, so this also pins the neq + // 3VL: k=4 (v={70,NULL}) and k=5 (v={NULL,NULL}) do not satisfy SQL `<>`, leaving {1,3,6}. + assert(on == Set(Row(1), Row(3), Row(6)), s"expected {1,3,6}, got $on") + } + } + + test("InSubquery in a SELECT-list CASE WHEN is rewritten, not only in a WHERE predicate") { + withTable("T") { + // `apply` rewrites via transformAllExpressionsWithPruning, so an InSubquery anywhere in the + // plan is a candidate -- not just a WHERE filter. A Project is a valid host (see + // ValidateSubqueryExpression), so pin that the rewrite fires with the same self-join subquery + // inside a projected CASE WHEN. Moving it out of WHERE is the only change from Pattern A'. + setupTable() + val subquery = + """SELECT s1.k FROM T s1 JOIN T s2 + | ON s1.k = s2.k AND s1.v <> s2.v""".stripMargin + val sql = + s"""SELECT k, CASE WHEN k IN ($subquery) THEN 1 ELSE 0 END AS flag + |FROM T outer_t""".stripMargin + + assertRuleFired(sql) + val (on, off) = runBoth(sql) + assert(on == off, s"SELECT-list InSubquery rewrite ON $on != OFF $off") + } + } + + test("Rewrite is idempotent: a second application on the rewritten plan is a no-op") { + withTable("T") { + // After the rewrite the outer `k IN (...)` is still an InSubquery, now over the aggregate, so + // the rule revisits it on later passes. Its output must be a no-op: the aggregate no longer + // matches the self-join shape, so the second pass returns the plan unchanged. + setupTable() + val sql = + """SELECT k FROM T outer_t WHERE k IN ( + | SELECT s1.k FROM T s1 JOIN T s2 + | ON s1.k = s2.k AND s1.v <> s2.v)""".stripMargin + withSQLConf(rewriteConf -> "true") { + val analyzed = spark.sql(sql).queryExecution.analyzed + val once = RewriteSelfJoinInequalityToAggregate(analyzed) + assert(ruleFired(once), s"precondition: first pass should fire:\n$once") + val twice = RewriteSelfJoinInequalityToAggregate(once) + assert(twice == once, s"rule is not idempotent:\n$once\n-- second pass -->\n$twice") + } + } + } + + test("Pattern A2: nested self-join is rewritten") { + withTable("T") { + withTempView("D") { + setupTable() + spark.sql( + """CREATE OR REPLACE TEMP VIEW D AS SELECT * FROM VALUES + | (1), (3), (6) AS D(k)""".stripMargin) + val sql = + """SELECT k FROM T outer_t WHERE k IN ( + | SELECT d.k + | FROM D d, (SELECT s1.k FROM T s1 JOIN T s2 + | ON s1.k = s2.k AND s1.v <> s2.v) sj + | WHERE d.k = sj.k)""".stripMargin + + assertRuleFired(sql) + assertMinMaxRewriteShape(optimizedPlanWith(sql, rewrite = true)) + val (on, off) = runBoth(sql) + assert(on == off, s"Pattern A2 rewrite ON $on != OFF $off") + assert(on == Set(Row(1), Row(3), Row(6))) + } + } + } + + test("Pattern A2: self-join on the LEFT of the outer join is rewritten") { + withTable("T") { + withTempView("D") { + // Mirror of the Pattern A2 test above. There the self-join is the RIGHT child of the outer + // join (`selfJoinOnRight = true`); here it is the LEFT child (`selfJoinOnRight = false`). + // The rule has an explicit branch for each side, so both are covered. + setupTable() + spark.sql( + """CREATE OR REPLACE TEMP VIEW D AS SELECT * FROM VALUES + | (1), (3), (6) AS D(k)""".stripMargin) + val sql = + """SELECT k FROM T outer_t WHERE k IN ( + | SELECT d.k + | FROM (SELECT s1.k FROM T s1 JOIN T s2 + | ON s1.k = s2.k AND s1.v <> s2.v) sj, D d + | WHERE sj.k = d.k)""".stripMargin + + assertRuleFired(sql) + val (on, off) = runBoth(sql) + assert(on == off, s"Pattern A2 (self-join on left) rewrite ON $on != OFF $off") + assert(on == Set(Row(1), Row(3), Row(6))) + } + } + } + + test("Q95-shaped query: both IN subqueries (Pattern A' and A2) are rewritten") { + withTable("T") { + withTempView("D") { + // TPC-DS Q95 feeds a self-join-inequality CTE into two IN subqueries: one selects the CTE + // directly (Pattern A') and one joins it with another relation (Pattern A2). Like real Q95, + // the CTE also projects the two neq-side columns (wh1/wh2) that both INs never consume, so + // the rewrite fires only if OptimizeSubqueries first drops them (ColumnPruning) and merges + // the projections (CollapseProject). Asserting two MIN and two MAX aliases pins both that + // upstream chain and that both subqueries rewrite. + setupTable() + spark.sql( + """CREATE OR REPLACE TEMP VIEW D AS SELECT * FROM VALUES + | (1), (3), (6) AS D(k)""".stripMargin) + val sql = + """WITH ws_wh AS ( + | SELECT s1.k AS ordn, s1.v AS wh1, s2.v AS wh2 FROM T s1 JOIN T s2 + | ON s1.k = s2.k AND s1.v <> s2.v) + |SELECT o.k FROM T o + |WHERE o.k IN (SELECT ordn FROM ws_wh) + | AND o.k IN (SELECT d.k FROM D d, ws_wh WHERE d.k = ws_wh.ordn)""".stripMargin + + val plan = optimizedPlanWith(sql, rewrite = true) + assert( + countAlias(plan, MinNeqAlias) == 2 && countAlias(plan, MaxNeqAlias) == 2, + s"expected both IN subqueries rewritten (2 MIN + 2 MAX signature aliases):\n$plan") + val (on, off) = runBoth(sql) + assert(on == off, s"Q95-shaped rewrite ON $on != OFF $off") + assert(on == Set(Row(1), Row(3), Row(6)), s"expected {1,3,6}, got $on") + } + } + } + + test("Pattern A2: nondeterminism in the outer join condition must not be rewritten") { + withTable("T") { + withTempView("D") { + // Both self-join sides are still repeatable here, so the per-side `isSameBaseRelation` + // check would pass; the `rand()` conjunct lives on the outer join ABOVE the self-join. Only + // the candidate-level `isRepeatablePlan` walk over the whole subquery catches it, so the + // rule must fail closed. This is the case the candidate-level guard exists for. + setupTable() + spark.sql( + """CREATE OR REPLACE TEMP VIEW D AS SELECT * FROM VALUES + | (1), (3), (6) AS D(k)""".stripMargin) + val sql = + """SELECT k FROM T outer_t WHERE k IN ( + | SELECT d.k + | FROM D d, (SELECT s1.k FROM T s1 JOIN T s2 + | ON s1.k = s2.k AND s1.v <> s2.v) sj + | WHERE d.k = sj.k AND rand() < 0.5)""".stripMargin + + assertRuleNotFired(sql) + } + } + } + + test("Clock-dependent STRING -> TIMESTAMP cast in the subquery is fail-closed") { + withTable("TTs") { + // A time-only STRING cast to TIMESTAMP takes its missing date from the runtime clock + // (LocalDate.now), so two self-join scans that straddle midnight can disagree while the + // rewrite folds them into one evaluation. The neq column here is derived by such a cast, so + // the rewrite must not fire. Isolate the cast as the sole cause with a same-shape control + // whose derived neq column uses a benign INT -> BIGINT cast (which does fire). + createTable( + "TTs", + "k INT, v INT, t STRING", + """ (1, 10, '01:00:00'), (1, 20, '02:00:00'), + | (2, 30, '03:00:00')""".stripMargin) + + def wrapped(neqExpr: String): String = + s"""SELECT k FROM TTs outer_t WHERE k IN ( + | SELECT s1.k FROM + | (SELECT k, $neqExpr AS nc FROM TTs) s1 + | JOIN (SELECT k, $neqExpr AS nc FROM TTs) s2 + | ON s1.k = s2.k AND s1.nc <> s2.nc)""".stripMargin + + // Same wrapper shape with a repeatable cast fires, proving the shape itself is supported. + assertRuleFired(wrapped("CAST(v AS BIGINT)")) + + // A direct STRING -> TIMESTAMP cast is fail-closed. + assertRuleNotFired(wrapped("CAST(t AS TIMESTAMP)")) + + // STRING -> TIMESTAMP_NTZ stays supported: it does not consult the session time zone and a + // time-only string parses to NULL deterministically rather than borrowing the runtime date, + // so the guard only rejects the clock-dependent LTZ conversion, not all string-to-timestamp. + assertRuleFired(wrapped("CAST(t AS TIMESTAMP_NTZ)")) + } + } + + test("Clock-dependent nested STRING -> TIMESTAMP casts (array/map/struct) are fail-closed") { + withTable("TNest") { + // castHasClockDependentStringToTimestamp recurses into ARRAY/MAP/STRUCT casts, so a + // clock-dependent STRING -> TIMESTAMP_LTZ hidden at any nesting level must fail closed. The + // cast rides in a WHERE filter as `CAST(nested AS <ts>) IS NOT NULL` (tree IsNotNull -> Cast + // -> Attribute, all allowlisted) rather than the neq column, which the neq-column type gate + // would reject wholesale and mask the recursion. So only the clock-dependent-cast guard can + // decline it. For each shape the TIMESTAMP_NTZ control must fire before the LTZ variant is + // required not to, so deleting a recursion arm turns the matching negative red. + createTable( + "TNest", + "k INT, v INT, arr ARRAY<STRING>, mp MAP<STRING, STRING>, st STRUCT<x: STRING>", + """ (1, 10, ARRAY('01:00:00'), MAP('a', '01:00:00'), NAMED_STRUCT('x', '01:00:00')), + | (1, 20, ARRAY('02:00:00'), MAP('a', '02:00:00'), NAMED_STRUCT('x', '02:00:00')), + | (2, 30, ARRAY('03:00:00'), MAP('a', '03:00:00'), NAMED_STRUCT('x', '03:00:00'))""" + .stripMargin) + + def wrapped(filterExpr: String): String = + s"""SELECT k FROM TNest outer_t WHERE k IN ( + | SELECT s1.k FROM + | (SELECT k, v FROM TNest WHERE $filterExpr) s1 + | JOIN (SELECT k, v FROM TNest WHERE $filterExpr) s2 + | ON s1.k = s2.k AND s1.v <> s2.v)""".stripMargin + + Seq( + "CAST(arr AS ARRAY<TIMESTAMP>)" -> "CAST(arr AS ARRAY<TIMESTAMP_NTZ>)", + "CAST(mp AS MAP<STRING, TIMESTAMP>)" -> "CAST(mp AS MAP<STRING, TIMESTAMP_NTZ>)", + "CAST(st AS STRUCT<x: TIMESTAMP>)" -> "CAST(st AS STRUCT<x: TIMESTAMP_NTZ>)" + ).foreach { case (ltz, ntz) => + assertRuleFired(wrapped(s"$ntz IS NOT NULL")) + assertRuleNotFired(wrapped(s"$ltz IS NOT NULL")) + } + } + } + + test("VARIANT casts to LTZ-containing targets are fail-closed") { + // Same WHERE-filter shape and NTZ-control-fires-first setup as the nested STRING-cast test + // above (guard coverage at plan level, not a runtime nested-Variant conversion). Variant + // specifics: a Variant cast parses a runtime string, so a target with a TIMESTAMP_LTZ leaf at + // any nesting level is clock-dependent and must fail closed, while TIMESTAMP_NTZ stays allowed. + // `Cast.needsTimeZone(VariantType, _)` is unconditionally true, so the guard uses + // variantTargetHasClockDependentStringToTimestamp, which recurses the TARGET type only. + // pushVariantIntoScan is disabled so the Cast reaches the rule rather than being folded into a + // scan-level struct-field extraction. + withSQLConf("spark.sql.variant.pushVariantIntoScan" -> "false") { + withTable("TVar") { + createTable( + "TVar", + "k INT, v INT, vt VARIANT", + """ (1, 10, parse_json('"01:00:00"')), (1, 20, parse_json('"02:00:00"')), + | (2, 30, parse_json('"03:00:00"'))""".stripMargin) + + def wrapped(filterExpr: String): String = + s"""SELECT k FROM TVar outer_t WHERE k IN ( + | SELECT s1.k FROM + | (SELECT k, v FROM TVar WHERE $filterExpr) s1 + | JOIN (SELECT k, v FROM TVar WHERE $filterExpr) s2 + | ON s1.k = s2.k AND s1.v <> s2.v)""".stripMargin + + Seq( + "CAST(vt AS TIMESTAMP)" -> "CAST(vt AS TIMESTAMP_NTZ)", + "CAST(vt AS ARRAY<TIMESTAMP>)" -> "CAST(vt AS ARRAY<TIMESTAMP_NTZ>)", + "CAST(vt AS MAP<STRING, TIMESTAMP>)" -> "CAST(vt AS MAP<STRING, TIMESTAMP_NTZ>)", + "CAST(vt AS STRUCT<x: TIMESTAMP>)" -> "CAST(vt AS STRUCT<x: TIMESTAMP_NTZ>)" + ).foreach { case (ltz, ntz) => + assertRuleFired(wrapped(s"$ntz IS NOT NULL")) + assertRuleNotFired(wrapped(s"$ltz IS NOT NULL")) + } + } + } + } + + test("Pattern A': multi-equi tuple IN with sjRight key remap is rewritten") { + withTable("TM") { + // Exercises the multi-equi-key path: two equi keys (k1, k2) drive the GROUP BY, and the tuple + // IN projects `s1.k1, s2.k2` -- so the second output column comes from the RIGHT self-join + // side and must be remapped to its sjLeft counterpart by `canonicalizeWrapper`. Covers + // multiple equi keys, tuple IN arity, the two injected IsNotNull(equiKey) filters, and the + // sjRight remap at once. + createTable( + "TM", + "k1 INT, k2 INT, v INT", + """ (1, 1, 10), (1, 1, 20), + | (1, 2, 30), (1, 2, 30), + | (2, 1, 40), (2, 1, 50), + | (CAST(NULL AS INT), 1, 60), (CAST(NULL AS INT), 1, 70), + | (3, CAST(NULL AS INT), 80), (3, CAST(NULL AS INT), 90)""".stripMargin) + val sql = + """SELECT k1, k2 FROM TM outer_t WHERE (k1, k2) IN ( + | SELECT s1.k1, s2.k2 FROM TM s1 JOIN TM s2 + | ON s1.k1 = s2.k1 AND s1.k2 = s2.k2 AND s1.v <> s2.v)""".stripMargin + + assertRuleFired(sql) + val (on, off) = runBoth(sql) + assert(on == off, s"multi-equi tuple IN rewrite ON $on != OFF $off") + // (1,1): distinct v={10,20} -> matches; (1,2): v={30} -> no; (2,1): v={40,50} -> matches; + // (NULL,1) and (3,NULL): NULL equi key filtered out by the injected IsNotNull. -> + // {(1,1),(2,1)} + assert(on == Set(Row(1, 1), Row(2, 1)), s"expected {(1,1),(2,1)}, got $on") + } + } + + test("NULL equi-key is filtered before aggregation for NOT IN") { + withTable("TN") { + withTempView("OuterKeys") { + createTable( + "TN", + "k INT, v INT", + """ (CAST(NULL AS INT), 10), + | (CAST(NULL AS INT), 20), + | (1, 10), (1, 20), + | (2, 30)""".stripMargin + ) + spark.sql( + """CREATE OR REPLACE TEMP VIEW OuterKeys AS SELECT * FROM VALUES + | (1), (2), (3) AS OuterKeys(k)""".stripMargin) + val sql = + """SELECT k FROM OuterKeys o WHERE k NOT IN ( + | SELECT s1.k FROM TN s1 JOIN TN s2 + | ON s1.k = s2.k AND s1.v <> s2.v)""".stripMargin + + assertRuleFired(sql) + val (on, off) = runBoth(sql) + assert(on == off, s"NULL equi-key NOT IN semantics diverge: ON=$on OFF=$off") + assert(on == Set(Row(2), Row(3)), s"expected {2,3}, got $on") + } + } + } + + test("Swapped aliases must not be treated as the same self-join columns") { + withTable("AliasBase") { + // The guard under test is `sameOutputPosition`. Both queries alias the same two base columns + // to the names `k` and `v` on both sides, so a rule that compares attribute names would fire + // on both; only the output ordinal tells them apart. + createTable("AliasBase", "a INT, b INT", " (1, 10), (1, 20), (2, 30)") + + val alignedSql = + """SELECT a FROM AliasBase outer_t WHERE a IN ( + | SELECT s1.k + | FROM (SELECT a AS k, b AS v FROM AliasBase) s1 + | JOIN (SELECT a AS k, b AS v FROM AliasBase) s2 + | ON s1.k = s2.k AND s1.v <> s2.v)""".stripMargin + assertRuleFired(alignedSql) + val (alignedOn, alignedOff) = runBoth(alignedSql) + assert( + alignedOn == alignedOff, + s"aligned-alias control diverges: ON=$alignedOn OFF=$alignedOff") + assert(alignedOn == Set(Row(1)), s"aligned-alias control expected {1}, got $alignedOn") + + // s1.k is `a` (output position 0) but s2.k is `b` (output position 1): same name, different + // column. Rewriting this would count distinct `b` per `a`, which is a different query. + val swappedSql = + """SELECT a FROM AliasBase outer_t WHERE a IN ( + | SELECT s1.k + | FROM (SELECT a AS k, b AS v FROM AliasBase) s1 + | JOIN (SELECT a AS v, b AS k FROM AliasBase) s2 + | ON s1.k = s2.k AND s1.v <> s2.v)""".stripMargin + assertRuleNotFired(swappedSql) + val (on, off) = runBoth(swappedSql) + assert(on == off, s"swapped-alias semantics diverge: ON=$on OFF=$off") + assert(on.isEmpty, s"swapped-alias baseline should be empty, got $on") + } + } + + test("Two different relations with the same schema must not be treated as a self-join") { + withTable("TLeft", "TRight") { + // The guard under test is `isSameBaseRelation`: it must reject a join between two DIFFERENT + // base tables even when they share a schema and column names. Distinct Parquet tables + // canonicalize to distinct `rootPaths`, so `left.canonicalized == right.canonicalized` is + // false and the rewrite must not fire. This is a correctness boundary, not a missed + // optimization: rewriting `TLeft JOIN TRight` as MIN(v) <> MAX(v) over TLeft alone would drop + // TRight's rows and change the answer. + createTable("TLeft", "k INT, v INT", " (1, 10), (1, 10), (2, 30)") + createTable("TRight", "k INT, v INT", " (1, 20), (1, 20), (2, 30)") + val sql = + """SELECT k FROM TLeft outer_t WHERE k IN ( + | SELECT s1.k FROM TLeft s1 JOIN TRight s2 + | ON s1.k = s2.k AND s1.v <> s2.v)""".stripMargin + assertRuleNotFired(sql) + val (on, off) = runBoth(sql) + assert(on == off, s"different-relation join semantics diverge: ON=$on OFF=$off") + // k=1: TLeft v={10} vs TRight v={20} -> 10<>20 true -> qualifies; k=2: 30<>30 false -> no. + assert(on == Set(Row(1)), s"expected {1}, got $on") + } + } + + // ==================== Positive: rewritten plan is structurally the aggregate ================ + // + // Result parity (ON == OFF) does not prove the rewrite produced the GROUP BY + HAVING + // MIN(v) <> MAX(v) shape rather than leaving the self-join and happening to agree, so these + // controls assert the shape directly. A full plan comparison against hand-written aggregate SQL + // would be wrong: InferFiltersFromConstraints adds isnotnull(min)/isnotnull(max) to a + // hand-written HAVING but not to the rule's filter (a later batch), so it would fail on redundant + // null predicates. + + private def optimizedPlanWith(sql: String, rewrite: Boolean): LogicalPlan = + withSQLConf(rewriteConf -> rewrite.toString) { + spark.sql(sql).queryExecution.optimizedPlan + } + + // The rewrite shape: an Aggregate emitting both signature aliases, with a Filter on top whose + // condition includes Not(EqualTo(min, max)). Match the two operands by ExprId, not by name, so a + // same-named attribute from elsewhere cannot satisfy it. `exists` on the condition (not exact + // match) because InferFiltersFromConstraints may fold redundant isnotnull(min/max) into it. + private def assertMinMaxRewriteShape(plan: LogicalPlan): Unit = { + val found = plan.collectFirstWithSubqueries { + case Filter(cond, agg: Aggregate) + if { + // Key by alias name so the shape requires exactly one MIN alias AND one MAX alias: two + // same-named aliases collapse to a single map key and fail the keySet check, which a + // bare `size == 2` on exprIds would not catch. + val signatureAttrs = agg.aggregateExpressions.collect { + case a: Alias if a.name == MinNeqAlias || a.name == MaxNeqAlias => + a.name -> a.toAttribute + }.toMap + signatureAttrs.keySet == Set(MinNeqAlias, MaxNeqAlias) && cond.exists { + case Not(EqualTo(l: Attribute, r: Attribute)) => + Set(l.exprId, r.exprId) == signatureAttrs.values.map(_.exprId).toSet + case _ => false + } + } => () + } + assert(found.isDefined, s"expected a MIN(v) <> MAX(v) aggregate rewrite shape:\n$plan") + } + + // ==================== Negative: rewrite must produce equivalent results (or bail) ========== + + test("Bare-Join subquery (no wrapper Project) fails closed: Pattern A' arity guard") { + withTable("T") { + // The guard under test is the `projectListOpt match { case None => None }` bail in + // `rewriteDirectSelfJoin`. With no wrapper Project the self-join output is `left ++ right`; + // replacing it with `Project(equiKeys, aggregate)` would shrink the arity that + // RewritePredicateSubquery later positionally zips against, misbinding the semi predicates. A + // `SELECT *` over the self-join whose tuple IN references every output column lets + // RemoveNoopOperators strip the identity Project, so the subquery reaches the rule as a bare + // Join -- proven below before asserting the rule declines it. + setupTable() + val sql = + """SELECT k FROM T outer_t WHERE (k, v, k, v) IN ( + | SELECT * FROM T s1 JOIN T s2 + | ON s1.k = s2.k AND s1.v <> s2.v)""".stripMargin + optimizedInSubqueryPlan(sql) match { + case _: Join => + case other => fail(s"expected a bare Join subquery, got:\n$other") + } + assertRuleNotFired(sql) + val (on, off) = runBoth(sql) + assert(on == off, s"bare-Join A' arity guard semantics diverge: ON=$on OFF=$off") + } + } + + test("Bare-Join nested self-join (no Project anywhere) fails closed: Pattern A2 arity guard") { + withTable("T") { + withTempView("D") { + // The guard under test is the `case None => return None` bail in `rewriteNestedSelfJoin`: + // with no wrapper Project over the self-join, changing its output arity would misbind the + // positional semi-predicate zip, so the rule fails closed. `SELECT *` plus a + // tuple IN over every column lets RemoveNoopOperators expose the bare nested self-join. + // ReorderJoin is excluded (a valid config) so the bare nested self-join deterministically + // reaches this rule rather than being reshaped away; the shape is then asserted so the test + // fails loudly if it stops reaching the guard. The two sides use renaming subqueries (ka/va + // vs kb/vb) so `SELECT *` yields distinct names -- a plain `T s1 JOIN T s2` would emit two + // columns named `k` and fail analysis; the renames are children of the self-join, so the + // identity `SELECT *` is still stripped. + setupTable() + spark.sql( + """CREATE OR REPLACE TEMP VIEW D AS SELECT * FROM VALUES + | (1), (3), (6) AS D(k)""".stripMargin) + val sql = + """SELECT k FROM T outer_t WHERE (k, v, k, k, v) IN ( + | SELECT * FROM D d JOIN ( + | SELECT * FROM (SELECT k AS ka, v AS va FROM T) s1 + | JOIN (SELECT k AS kb, v AS vb FROM T) s2 + | ON s1.ka = s2.kb AND s1.va <> s2.vb) sj + | ON d.k = sj.ka)""".stripMargin + withSQLConf(SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> ReorderJoin.ruleName) { + // A child that is a `sameResult` Inner Join carrying an inequality is exactly the + // self-join rewriteNestedSelfJoin extracts; asserting it proves execution reaches the + // arity guard. + def isTargetSelfJoin(p: LogicalPlan): Boolean = p match { + case j: Join if j.joinType == Inner => + j.left.sameResult(j.right) && + j.condition.exists(_.exists { case _: Not => true; case _ => false }) + case _ => false + } + optimizedInSubqueryPlan(sql) match { + case j: Join if isTargetSelfJoin(j.left) || isTargetSelfJoin(j.right) => + case other => fail(s"expected a bare nested self-join child, got:\n$other") + } + assertRuleNotFired(sql) + val (on, off) = runBoth(sql) + assert(on == off, s"bare-Join A2 arity guard semantics diverge: ON=$on OFF=$off") + } + } + } + } + + test("IS DISTINCT FROM is rejected by the self-join condition parser") { + withTable("T") { + setupTable() + val sql = + """SELECT k FROM T outer_t WHERE k IN ( + | SELECT s1.k FROM T s1 JOIN T s2 + | ON s1.k = s2.k AND s1.v IS DISTINCT FROM s2.v)""".stripMargin + val (on, off) = runBoth(sql) + assert(on == off, s"IS DISTINCT FROM semantics diverge: ON=$on OFF=$off") + // Assert the full result, not just contains(4): unlike `<>`, `IS DISTINCT FROM` treats NULL + // as a value, so k=4 (v={70,NULL}) qualifies alongside k=1, k=3 and k=6. + assert(on == Set(Row(1), Row(3), Row(4), Row(6)), s"expected {1,3,4,6}, got $on") + assertRuleNotFired(sql) + } + } + + test("IsNotNull on a non-join column is rejected") { + withTable("T3") { + // The guard under test is the predicate parser: it accepts IsNotNull only on a column the + // join condition already references, because such a predicate is implied by the equi-key or + // the inequality and can be dropped, while IsNotNull(w) filters rows the aggregate would + // otherwise count. The control is the same query without that one conjunct. + createTable( + "T3", + "k INT, v INT, w INT", + """ (1, 10, 100), (1, 20, 200), + | (2, 30, CAST(NULL AS INT)), (2, 40, CAST(NULL AS INT))""".stripMargin) + + val controlSql = + """SELECT k FROM T3 outer_t WHERE k IN ( + | SELECT s1.k FROM T3 s1 JOIN T3 s2 + | ON s1.k = s2.k AND s1.v <> s2.v)""".stripMargin + assertRuleFired(controlSql) + val (controlOn, controlOff) = runBoth(controlSql) + assert(controlOn == controlOff, s"T3 control diverges: ON=$controlOn OFF=$controlOff") + assert(controlOn == Set(Row(1), Row(2)), s"T3 control expected {1,2}, got $controlOn") + + val sql = + """SELECT k FROM T3 outer_t WHERE k IN ( + | SELECT s1.k FROM T3 s1 JOIN T3 s2 + | ON s1.k = s2.k AND s1.v <> s2.v AND s1.w IS NOT NULL)""".stripMargin + assertRuleNotFired(sql) + val (on, off) = runBoth(sql) + assert(on == off, s"IsNotNull(non-join-col) semantics diverge: ON=$on OFF=$off") + assert(on == Set(Row(1)), s"expected {1}, got $on") + } + } + + test("Multiple inequality columns are rejected") { + withTable("T2") { + // The guard under test is `neqPairs.size != 1`. Two inequalities need "at least two rows + // differing in v AND in w", which no MIN/MAX over a single column can express. The control is + // the same query with only the first inequality. + createTable( + "T2", + "k INT, v INT, w INT", + """ (1, 10, 100), (1, 20, 200), + | (2, 30, 300), + | (3, 40, 100), (3, 50, 100)""".stripMargin) + + val controlSql = + """SELECT k FROM T2 outer_t WHERE k IN ( + | SELECT s1.k FROM T2 s1 JOIN T2 s2 + | ON s1.k = s2.k AND s1.v <> s2.v)""".stripMargin + assertRuleFired(controlSql) + val (controlOn, controlOff) = runBoth(controlSql) + assert(controlOn == controlOff, s"T2 control diverges: ON=$controlOn OFF=$controlOff") + assert(controlOn == Set(Row(1), Row(3)), s"T2 control expected {1,3}, got $controlOn") + + val sql = + """SELECT k FROM T2 outer_t WHERE k IN ( + | SELECT s1.k FROM T2 s1 JOIN T2 s2 + | ON s1.k = s2.k AND s1.v <> s2.v AND s1.w <> s2.w)""".stripMargin + assertRuleNotFired(sql) + val (on, off) = runBoth(sql) + assert(on == off, s"multi-column neq semantics diverge: ON=$on OFF=$off") + assert(on == Set(Row(1)), s"expected {1}, got $on") + } + } + + test("LeftOuter self-join inside the IN subquery is rejected by the join-type guard") { + withTable("T") { + // The guard under test is `joinType == Inner`. The MIN(v) <> MAX(v) collapse is valid only + // for the INNER shape: a LEFT OUTER self-join preserves unmatched left rows, so its subquery + // returns every left key; rewriting it to a per-key aggregate drops keys and changes IN + // membership. The INNER control and LEFT OUTER variant differ only in join type. + setupTable() + val controlSql = + """SELECT k FROM T outer_t WHERE k IN ( + | SELECT s1.k FROM T s1 JOIN T s2 + | ON s1.k = s2.k AND s1.v <> s2.v)""".stripMargin + assertRuleFired(controlSql) + + val leftOuterSql = + """SELECT k FROM T outer_t WHERE k IN ( + | SELECT s1.k FROM T s1 LEFT OUTER JOIN T s2 + | ON s1.k = s2.k AND s1.v <> s2.v)""".stripMargin + // Confirm the optimizer keeps it LEFT OUTER, so the join-type guard is what declines it and + // not some rule having reshaped the join away. + assert( + optimizedInSubqueryPlan(leftOuterSql).exists { + case j: Join if j.joinType == LeftOuter => true + case _ => false + }, + "expected the subquery to still contain a LEFT OUTER join") + assertRuleNotFired(leftOuterSql) + val (on, off) = runBoth(leftOuterSql) + assert(on == off, s"LeftOuter-in-IN semantics diverge: ON=$on OFF=$off") + } + } + + test("Join hint on the self-join is fail-closed (direct and nested)") { + withTable("T") { + withTempView("D") { + // The guard is `hint.isEmpty`, checked in two places: rewriteDirectSelfJoin for a top-level + // self-join (Pattern A') and tryExtractSelfJoin for a self-join nested under an outer join + // (Pattern A2). The rewrite deletes the self-join, so a hint on it is a directive about a + // join that would vanish -- fail closed in both. A hint never changes rows, so results are + // identical. + setupTable() + spark.sql( + """CREATE OR REPLACE TEMP VIEW D AS SELECT * FROM VALUES + | (1), (3), (6) AS D(k)""".stripMargin) + + // Pattern A': hint on the top-level self-join (rewriteDirectSelfJoin). + val directControl = + """SELECT k FROM T outer_t WHERE k IN ( + | SELECT s1.k FROM T s1 JOIN T s2 + | ON s1.k = s2.k AND s1.v <> s2.v)""".stripMargin + assertRuleFired(directControl) + val directHinted = + """SELECT k FROM T outer_t WHERE k IN ( + | SELECT /*+ BROADCAST(s2) */ s1.k FROM T s1 JOIN T s2 + | ON s1.k = s2.k AND s1.v <> s2.v)""".stripMargin + assertRuleNotFired(directHinted) + val (dOn, dOff) = runBoth(directHinted) + assert(dOn == dOff, s"direct-hint semantics diverge: ON=$dOn OFF=$dOff") + assert(dOn == Set(Row(1), Row(3), Row(6)), s"expected {1,3,6}, got $dOn") + + // Pattern A2: hint on the nested self-join (tryExtractSelfJoin). + val nestedControl = + """SELECT k FROM T outer_t WHERE k IN ( + | SELECT d.k + | FROM D d, (SELECT s1.k FROM T s1 JOIN T s2 + | ON s1.k = s2.k AND s1.v <> s2.v) sj + | WHERE d.k = sj.k)""".stripMargin + assertRuleFired(nestedControl) + val nestedHinted = + """SELECT k FROM T outer_t WHERE k IN ( + | SELECT d.k + | FROM D d, (SELECT /*+ BROADCAST(s2) */ s1.k FROM T s1 JOIN T s2 + | ON s1.k = s2.k AND s1.v <> s2.v) sj + | WHERE d.k = sj.k)""".stripMargin + // Confirm a hinted self-join actually survives into the subquery the rule inspects, so the + // decline is attributable to the `hint.isEmpty` guard and not to the hint reshaping the + // plan so that tryExtractSelfJoin never sees a candidate. + assert( + optimizedInSubqueryPlan(nestedHinted).exists { + case j: Join if j.left.sameResult(j.right) && !j.hint.isEmpty => true + case _ => false + }, + "expected a hinted nested self-join to survive into the optimized subquery") + assertRuleNotFired(nestedHinted) + val (nOn, nOff) = runBoth(nestedHinted) + assert(nOn == nOff, s"nested-hint semantics diverge: ON=$nOn OFF=$nOff") + assert(nOn == Set(Row(1), Row(3), Row(6)), s"expected {1,3,6}, got $nOn") + } + } + } + + test("Config gate: rewrite disabled leaves a valid A' candidate untouched") { + withTable("T") { + setupTable() + val sql = + """SELECT k FROM T outer_t WHERE k IN ( + | SELECT s1.k FROM T s1 JOIN T s2 + | ON s1.k = s2.k AND s1.v <> s2.v)""".stripMargin + withSQLConf(rewriteConf -> "false") { + val plan = spark.sql(sql).queryExecution.optimizedPlan + assert(!ruleFired(plan), s"config off must not fire rewrite:\n$plan") + val res = spark.sql(sql).collect().toSet + assert(res == Set(Row(1), Row(3), Row(6)), s"config off correctness broken: $res") + } + } + } + + // ==================== Correlated subquery: rule must fail-closed ==================== + + private def setupOuterT(): Unit = { + spark.sql( + """CREATE OR REPLACE TEMP VIEW OuterT AS SELECT * FROM VALUES + | (1), (3), (6) AS OuterT(k)""".stripMargin) + } + + test("Correlated InSubquery is fail-closed") { + withTable("T") { + withTempView("OuterT") { + setupTable() + setupOuterT() + val sql = + """SELECT o.k FROM OuterT o WHERE o.k IN ( + | SELECT s1.k FROM T s1 JOIN T s2 + | ON s1.k = s2.k AND s1.v <> s2.v + | WHERE s2.k = o.k)""".stripMargin + // Precondition: the InSubquery is genuinely correlated (its ListQuery carries outer + // references), so the not-fired result exercises the `lq.children.isEmpty` guard in `apply` + // rather than an unrelated shape mismatch. + val analyzed = spark.sql(sql).queryExecution.analyzed + var sawCorrelated = false + analyzed.foreach { node => + node.expressions.foreach(_.foreach { + case InSubquery(_, lq: ListQuery) => sawCorrelated ||= lq.children.nonEmpty + case _ => + }) + } + assert(sawCorrelated, s"expected a correlated InSubquery in analyzed plan:\n$analyzed") + val (on, off) = runBoth(sql) + assert(on == off, s"correlated IN parity: ON=$on OFF=$off") + assert(on == Set(Row(1), Row(3), Row(6)), s"expected {1,3,6}, got $on") + assertRuleNotFired(sql) + } + } + } + + // ==================== Repeatability whitelist: unknown operators fail-closed ============== + + test("LIMIT inside the self-join subtrees breaks row-bag repeatability: rule bails out") { + withTable("T") { + // The guard under test is the operator whitelist in `isRowBagRepeatable`: a `LIMIT` without a + // total order returns an arbitrary row subset that two scans need not agree on, so it is not + // row-bag repeatable. A `LIMIT` node carries only a `Literal`, so every expression is + // allowlisted and `hasRepeatableExpressions` passes; with the precondition below proving the + // two self-join inputs are still `sameResult`, the operator whitelist (which does not list + // Limit) is the relevant remaining rejection point. A control without `LIMIT` fires on the + // same shape. + setupTable() + val controlSql = + """SELECT k FROM T outer_t WHERE k IN ( + | SELECT s1.k FROM (SELECT k, v FROM T) s1 + | JOIN (SELECT k, v FROM T) s2 + | ON s1.k = s2.k AND s1.v <> s2.v)""".stripMargin + assertRuleFired(controlSql) + val (controlOn, controlOff) = runBoth(controlSql) + assert(controlOn == controlOff, s"LIMIT control diverges: ON=$controlOn OFF=$controlOff") + assert(controlOn == Set(Row(1), Row(3), Row(6)), s"expected {1,3,6}, got $controlOn") + + val sql = + """SELECT k FROM T outer_t WHERE k IN ( + | SELECT s1.k FROM (SELECT k, v FROM T LIMIT 2) s1 + | JOIN (SELECT k, v FROM T LIMIT 2) s2 + | ON s1.k = s2.k AND s1.v <> s2.v)""".stripMargin + // Precondition: both self-join inputs stay `sameResult`, ruling out a structural mismatch, + // and both still contain a Limit. Since Limit carries only an allowlisted Literal expression, + // the non-firing result pins the row-bag operator whitelist. + def containsLimit(plan: LogicalPlan): Boolean = + plan.exists { + case _: GlobalLimit | _: LocalLimit => true + case _ => false + } + val before = optimizedInSubqueryPlan(sql) + assert( + before.exists { + case j: Join if j.joinType == Inner && j.left.sameResult(j.right) && + containsLimit(j.left) && containsLimit(j.right) => true + case _ => false + }, + s"expected a sameResult self-join whose inputs both contain a Limit:\n$before") + assertRuleNotFired(sql) + } + } + + test("Nondeterministic self-join input is rejected") { + // The guard under test is `plan.deterministic` inside `isRepeatablePlan`. Both sides use the + // same explicit seed, so the subplans share a canonical shape and the rejection is not from + // `isSameBaseRelation`. The control replaces `rand(41) < 0.5` with a deterministic filter, + // proving this Range/Filter/Project shape reaches the rewrite. + val controlSql = + """SELECT k FROM (SELECT CAST(id AS INT) AS k, CAST(id AS INT) AS v FROM range(100)) t + |WHERE k IN ( + | SELECT s1.k FROM ( + | SELECT CAST(id % 10 AS INT) AS k, CAST(id AS INT) AS v + | FROM range(1000) WHERE id % 2 = 0 + | ) s1 + | JOIN ( + | SELECT CAST(id % 10 AS INT) AS k, CAST(id AS INT) AS v + | FROM range(1000) WHERE id % 2 = 0 + | ) s2 + | ON s1.k = s2.k AND s1.v <> s2.v)""".stripMargin + assertRuleFired(controlSql) + val (controlOn, controlOff) = runBoth(controlSql) + assert(controlOn == controlOff, s"range control diverges: ON=$controlOn OFF=$controlOff") + assert( + controlOn == Set(Row(0), Row(2), Row(4), Row(6), Row(8)), + s"range control expected the even keys, got $controlOn") + + val sql = + """SELECT k FROM (SELECT CAST(id AS INT) AS k, CAST(id AS INT) AS v FROM range(100)) t + |WHERE k IN ( + | SELECT s1.k FROM ( + | SELECT CAST(id % 10 AS INT) AS k, CAST(id AS INT) AS v + | FROM range(1000) WHERE rand(41) < 0.5 + | ) s1 + | JOIN ( + | SELECT CAST(id % 10 AS INT) AS k, CAST(id AS INT) AS v + | FROM range(1000) WHERE rand(41) < 0.5 + | ) s2 + | ON s1.k = s2.k AND s1.v <> s2.v)""".stripMargin + assertRuleNotFired(sql) + } + + test("LogicalRDD leaf is not a trusted repeatable source: rule bails out") { + withTable("RddCtl") { + withTempView("RddT") { + // The guard under test is the leaf allowlist in `isRowBagRepeatable`: a Parquet + // `LogicalRelation` is trusted, but a `LogicalRDD` (createDataFrame over an RDD) wraps an + // arbitrary RDD lineage whose runtime row bag Catalyst cannot prove repeatable, so it must + // fail closed even though `plan.deterministic` is true. The Parquet control uses the same + // schema, data and query shape and fires, pinning the rejection to the leaf allowlist. + createTable("RddCtl", "k INT, v INT", " (1, 10), (1, 20), (2, 30)") + val controlSql = + """SELECT k FROM RddCtl outer_t WHERE k IN ( + | SELECT s1.k FROM RddCtl s1 JOIN RddCtl s2 + | ON s1.k = s2.k AND s1.v <> s2.v)""".stripMargin + assertRuleFired(controlSql) + val (controlOn, controlOff) = runBoth(controlSql) + assert(controlOn == controlOff, s"Parquet control diverges: ON=$controlOn OFF=$controlOff") + assert(controlOn == Set(Row(1)), s"Parquet control expected {1}, got $controlOn") + + val schema = StructType(Seq(StructField("k", IntegerType), StructField("v", IntegerType))) + val rows = spark.sparkContext.parallelize(Seq(Row(1, 10), Row(1, 20), Row(2, 30))) + spark.createDataFrame(rows, schema).createOrReplaceTempView("RddT") + val sql = + """SELECT k FROM RddT outer_t WHERE k IN ( + | SELECT s1.k FROM RddT s1 JOIN RddT s2 + | ON s1.k = s2.k AND s1.v <> s2.v)""".stripMargin + assertRuleNotFired(sql) + val (on, off) = runBoth(sql) + assert(on == off, s"LogicalRDD semantics diverge: ON=$on OFF=$off") + assert(on == Set(Row(1)), s"expected {1}, got $on") + } + } + } + + test("Non-allowlisted deterministic expression (Abs) fails closed") { + withTable("T") { + // The guard under test is the expression allowlist in `isRepeatableExpression`: it trusts + // expression TYPES, not merely `deterministic`. Abs is deterministic but not (yet) + // allowlisted, so a self-join side projecting abs(v) fails closed -- a missed optimization, + // not a bug. The control uses `v + 1` (not `+ 0`, so the Add survives arithmetic + // simplification) and fires; wrapping the same column in abs() is the sole change and makes + // it not fire. v is INT, so the Add is a plain `Add(v, 1)` with no decimal PromotePrecision / + // CheckOverflow wrappers. + setupTable() + + val controlSql = + """SELECT k FROM T outer_t WHERE k IN ( + | SELECT s1.k + | FROM (SELECT k, v + 1 AS x FROM T) s1 + | JOIN (SELECT k, v + 1 AS x FROM T) s2 + | ON s1.k = s2.k AND s1.x <> s2.x)""".stripMargin + assertRuleFired(controlSql) + val (controlOn, controlOff) = runBoth(controlSql) + assert(controlOn == controlOff, s"Add control diverges: ON=$controlOn OFF=$controlOff") + // v+1 is injective over the (non-null) v values, so distinctness per k is unchanged: {1,3,6}. + assert( + controlOn == Set(Row(1), Row(3), Row(6)), + s"Add control expected {1,3,6}, got $controlOn") + + val sql = + """SELECT k FROM T outer_t WHERE k IN ( + | SELECT s1.k + | FROM (SELECT k, abs(v) AS x FROM T) s1 + | JOIN (SELECT k, abs(v) AS x FROM T) s2 + | ON s1.k = s2.k AND s1.x <> s2.x)""".stripMargin + assertRuleNotFired(sql) + val (on, off) = runBoth(sql) + assert(on == off, s"Abs-projected self-join semantics diverge: ON=$on OFF=$off") + } + } + + // ==================== Data-type safety: comparison vs grouping/MIN-MAX equality ============= + // + // The rewrite turns `<>` into MIN(v) <> MAX(v) and `=` into GROUP BY, so it is only sound where + // comparison equality coincides with grouping/MIN-MAX ordering equality. The two roles differ: an + // equi key needs grouping equality (binary equality for strings), the neq column needs MIN/MAX + // ordering equality (binary ordering for strings). Both are positive allowlists, not + // `RowOrdering.isOrderable`. These tests pin the boundary for the risky types. + + test("Float/Double neq column is rejected (comparison-vs-MIN-MAX contract, defensive)") { + withTable("TFloat") { + // This pins the neq-column type gate: `vi` (Int, allowlisted) fires, `vd` (Double) does not. + // Double fails closed defensively -- comparison vs grouping/MIN-MAX agreement on signed zero + // and NaN rests on normalization details (NormalizeFloatingNumbers) that need not match + // across Spark versions or native backends. Current Spark aligns them, so OFF is {1} and the + // non-firing ON matches it: the test pins fail-closed behavior, not a divergence. + createTable( + "TFloat", + "k INT, vi INT, vd DOUBLE", + """ (1, 10, 1.0), (1, 20, 2.0), + | (2, 30, 0.0), (2, 30, -0.0), + | (3, 40, CAST('NaN' AS DOUBLE)), (3, 50, CAST('NaN' AS DOUBLE))""".stripMargin) + + val controlSql = + """SELECT k FROM TFloat outer_t WHERE k IN ( + | SELECT s1.k FROM TFloat s1 JOIN TFloat s2 + | ON s1.k = s2.k AND s1.vi <> s2.vi)""".stripMargin + assertRuleFired(controlSql) + val (controlOn, controlOff) = runBoth(controlSql) + assert(controlOn == controlOff, s"Int-neq control diverges: ON=$controlOn OFF=$controlOff") + assert(controlOn == Set(Row(1), Row(3)), s"Int-neq control expected {1,3}, got $controlOn") + + val sql = + """SELECT k FROM TFloat outer_t WHERE k IN ( + | SELECT s1.k FROM TFloat s1 JOIN TFloat s2 + | ON s1.k = s2.k AND s1.vd <> s2.vd)""".stripMargin + assertRuleNotFired(sql) + val (on, off) = runBoth(sql) + assert(on == off, s"Double-neq semantics diverge: ON=$on OFF=$off") + // k=1 matches (1.0 <> 2.0); k=2 does not (0.0 = -0.0); k=3 does not (Spark NaN = NaN). + assert(on == Set(Row(1)), s"Double-neq baseline expected {1}, got $on") + } + } + + test("Float/Double equi-key is rejected (defensive fail-closed)") { + withTable("TFloatKey") { + // This pins the equi-key type gate: `ki` (Int) fires, `kd` (Double) does not. Current Spark + // aligns floating-point comparison with grouping normalization here, but the rule does not + // depend on that implementation contract, so it fails closed. + createTable( + "TFloatKey", + "kd DOUBLE, ki INT, v INT", + """ (1.0, 1, 10), (1.0, 1, 20), + | (2.0, 2, 30), + | (0.0, 3, 40), (-0.0, 3, 50)""".stripMargin) + + val controlSql = + """SELECT ki FROM TFloatKey outer_t WHERE ki IN ( + | SELECT s1.ki FROM TFloatKey s1 JOIN TFloatKey s2 + | ON s1.ki = s2.ki AND s1.v <> s2.v)""".stripMargin + assertRuleFired(controlSql) + val (controlOn, controlOff) = runBoth(controlSql) + assert(controlOn == controlOff, s"Int-key control diverges: ON=$controlOn OFF=$controlOff") + assert(controlOn == Set(Row(1), Row(3)), s"Int-key control expected {1,3}, got $controlOn") + + val sql = + """SELECT kd FROM TFloatKey outer_t WHERE kd IN ( + | SELECT s1.kd FROM TFloatKey s1 JOIN TFloatKey s2 + | ON s1.kd = s2.kd AND s1.v <> s2.v)""".stripMargin + assertRuleNotFired(sql) + val (on, off) = runBoth(sql) + assert(on == off, s"Double-key semantics diverge: ON=$on OFF=$off") + } + } + + test("Complex-type neq column (array/struct) is rejected wholesale") { + withTable("TCplx") { + // The neq-column type gate rejects complex types wholesale, which also covers any + // Float/Double nested inside them. + // Control (`v` Int) fires; the ARRAY<DOUBLE> and STRUCT<..DOUBLE> variants -- identical query + // shape, only the neq column changed -- do not. + createTable( + "TCplx", + "k INT, v INT, a ARRAY<DOUBLE>, s STRUCT<x: INT, y: DOUBLE>", + """ (1, 10, ARRAY(1.0), NAMED_STRUCT('x', 1, 'y', 1.0)), + | (1, 20, ARRAY(2.0), NAMED_STRUCT('x', 2, 'y', 2.0)), + | (2, 30, ARRAY(1.0), NAMED_STRUCT('x', 1, 'y', 1.0))""".stripMargin) + + val controlSql = + """SELECT k FROM TCplx outer_t WHERE k IN ( + | SELECT s1.k FROM TCplx s1 JOIN TCplx s2 + | ON s1.k = s2.k AND s1.v <> s2.v)""".stripMargin + assertRuleFired(controlSql) + val (controlOn, controlOff) = runBoth(controlSql) + assert(controlOn == controlOff, s"Int-neq control diverges: ON=$controlOn OFF=$controlOff") + assert(controlOn == Set(Row(1)), s"Int-neq control expected {1}, got $controlOn") + + val arraySql = + """SELECT k FROM TCplx outer_t WHERE k IN ( + | SELECT s1.k FROM TCplx s1 JOIN TCplx s2 + | ON s1.k = s2.k AND s1.a <> s2.a)""".stripMargin + assertRuleNotFired(arraySql) + val (arrayOn, arrayOff) = runBoth(arraySql) + assert(arrayOn == arrayOff, s"array-neq semantics diverge: ON=$arrayOn OFF=$arrayOff") + + val structSql = + """SELECT k FROM TCplx outer_t WHERE k IN ( + | SELECT s1.k FROM TCplx s1 JOIN TCplx s2 + | ON s1.k = s2.k AND s1.s <> s2.s)""".stripMargin + assertRuleNotFired(structSql) + val (structOn, structOff) = runBoth(structSql) + assert(structOn == structOff, s"struct-neq semantics diverge: ON=$structOn OFF=$structOff") + } + } + + test("String neq/equi key: default (UTF8_BINARY) fires, non-binary collation fails closed") { + withTable("TStrBin", "TStrCiNeq", "TStrCiEqui") { + // On strings the two type gates diverge: an equi key is admitted only under binary equality + // and a neq column only under binary ordering (both byte-wise), because a non-binary + // collation (Spark 4.0+) routes comparison and grouping through different code paths. + // UTF8_LCASE satisfies neither, so it is rejected in either role; the default-collation + // control fires. The two negatives collate exactly ONE column each so each pins its gate: a + // UTF8_LCASE NEQ column the neq-side check, a UTF8_LCASE EQUI key the equi-side check. + // Collating both at once would leave it ambiguous which gate fired. + createTable( + "TStrBin", + "k STRING, v STRING", + """ ('a', 'x'), ('a', 'y'), + | ('b', 'z')""".stripMargin) + val binSql = + """SELECT k FROM TStrBin outer_t WHERE k IN ( + | SELECT s1.k FROM TStrBin s1 JOIN TStrBin s2 + | ON s1.k = s2.k AND s1.v <> s2.v)""".stripMargin + assertRuleFired(binSql) + val (binOn, binOff) = runBoth(binSql) + assert(binOn == binOff, s"binary-string control diverges: ON=$binOn OFF=$binOff") + assert(binOn == Set(Row("a")), s"binary-string control expected {a}, got $binOn") + + // Negative 1: only the NEQ column is non-binary collated -> neq-side type gate rejects. + createTable( + "TStrCiNeq", + "k STRING, v STRING COLLATE UTF8_LCASE", + """ ('a', 'x'), ('a', 'y'), + | ('b', 'z')""".stripMargin) + val ciNeqSql = + """SELECT k FROM TStrCiNeq outer_t WHERE k IN ( + | SELECT s1.k FROM TStrCiNeq s1 JOIN TStrCiNeq s2 + | ON s1.k = s2.k AND s1.v <> s2.v)""".stripMargin + assertRuleNotFired(ciNeqSql) + val (neqOn, neqOff) = runBoth(ciNeqSql) + assert(neqOn == neqOff, s"collated-neq semantics diverge: ON=$neqOn OFF=$neqOff") + + // Negative 2: only the EQUI key is non-binary collated -> equi-side type gate rejects. + createTable( + "TStrCiEqui", + "k STRING COLLATE UTF8_LCASE, v STRING", + """ ('a', 'x'), ('a', 'y'), + | ('b', 'z')""".stripMargin) + val ciEquiSql = + """SELECT k FROM TStrCiEqui outer_t WHERE k IN ( + | SELECT s1.k FROM TStrCiEqui s1 JOIN TStrCiEqui s2 + | ON s1.k = s2.k AND s1.v <> s2.v)""".stripMargin + assertRuleNotFired(ciEquiSql) + val (equiOn, equiOff) = runBoth(ciEquiSql) + assert(equiOn == equiOff, s"collated-equi semantics diverge: ON=$equiOn OFF=$equiOff") + } + } + + test("CHAR/VARCHAR join keys are rejected via declared-type metadata") { + withTable("TStr", "TCharVarchar") { + // CHAR/VARCHAR columns reach the optimizer as annotated StringType (CharVarcharUtils records + // the declared type in the metadata). Recovering the declared type from the metadata keeps + // them out of the StringType allowlist, where a dataType-only check would admit them. The + // plain-STRING control fires while the CHAR(5) and VARCHAR(5) variants fail closed. (Both `k` + // and `v` are CHAR/VARCHAR, so this exercises the metadata read, not which gate rejects.) + createTable( + "TStr", + "k STRING, v STRING", + """ ('a', 'x'), ('a', 'y'), + | ('b', 'z')""".stripMargin) + val stringSql = + """SELECT k FROM TStr outer_t WHERE k IN ( + | SELECT s1.k FROM TStr s1 JOIN TStr s2 + | ON s1.k = s2.k AND s1.v <> s2.v)""".stripMargin + assertRuleFired(stringSql) + val (strOn, strOff) = runBoth(stringSql) + assert(strOn == strOff, s"string control diverges: ON=$strOn OFF=$strOff") + assert(strOn == Set(Row("a")), s"string control expected {a}, got $strOn") + + Seq("CHAR(5)", "VARCHAR(5)").foreach { keyType => + createTable( + "TCharVarchar", + s"k $keyType, v $keyType", + """ ('a', 'x'), ('a', 'y'), + | ('b', 'z')""".stripMargin) + val sql = + """SELECT k FROM TCharVarchar outer_t WHERE k IN ( + | SELECT s1.k FROM TCharVarchar s1 JOIN TCharVarchar s2 + | ON s1.k = s2.k AND s1.v <> s2.v)""".stripMargin + assertRuleNotFired(sql) + val (on, off) = runBoth(sql) + assert(on == off, s"$keyType not-fired query diverges: ON=$on OFF=$off") + } + } + } + + test("ANSI: rewrite preserves observable error behavior (throw-or-succeed parity)") { + withTable("TAnsi") { + // Not a rejection but a parity property: the allowlist admits Cast, which can throw under + // ANSI. It evaluates `CAST(s AS INT)` once per row inside the Aggregate, the baseline + // self-join per row on each side -- the same rows -- so a malformed value must make + // BOTH forms behave the same (both succeed with equal rows, or both throw the same error + // class; a one-sided throw is a blocker). k=2 holds a non-numeric 's'. + createTable( + "TAnsi", + "k INT, s STRING", + """ (1, '10'), (1, '20'), + | (2, '30'), (2, 'xyz')""".stripMargin) + val sql = + """SELECT k FROM TAnsi outer_t WHERE k IN ( + | SELECT s1.k + | FROM (SELECT k, CAST(s AS INT) AS x FROM TAnsi) s1 + | JOIN (SELECT k, CAST(s AS INT) AS x FROM TAnsi) s2 + | ON s1.k = s2.k AND s1.x <> s2.x)""".stripMargin + + Seq("false", "true").foreach { ansi => + withSQLConf(SQLConf.ANSI_ENABLED.key -> ansi) { + // The rule fires at plan level regardless of ANSI (the cast throws only at runtime). + assertRuleFired(sql) + val on = runOutcome(sql, rewrite = true) + val off = runOutcome(sql, rewrite = false) + (on, off) match { + case (Right(onRows), Right(offRows)) => + assert(onRows == offRows, + s"ANSI=$ansi both succeeded but diverged: ON=$onRows OFF=$offRows") + // Positive signal under ANSI off: the cast is defined (yields NULL) for every + // surviving row, so ON must return the real membership {1}, not just match OFF. + if (ansi == "false") { + assert(onRows == Set(Row(1)), s"ANSI=false expected {1}, got $onRows") + } + case (Left(onErr), Left(offErr)) => Review Comment: **Non-blocking (P2):** For `ansi=false`, this match still accepts `Left/Left` as long as the error classes agree, so both executions can throw and the expected successful rows are never checked. Please make the non-ANSI Cast and Remainder cases require `Right(expectedRows)`, while the ANSI cases require `Left(expectedErrorClass)`, retaining ON/OFF parity within those required outcomes. -- 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]
