wangyum commented on code in PR #59027: URL: https://github.com/apache/spark/pull/59027#discussion_r4355925219
########## sql/core/src/main/scala/org/apache/spark/sql/execution/joins/BroadcastRangeJoinExec.scala: ########## @@ -0,0 +1,597 @@ +/* + * 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.joins + +import org.apache.spark.SparkException +import org.apache.spark.broadcast.Broadcast +import org.apache.spark.rdd.RDD +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions._ +import org.apache.spark.sql.catalyst.expressions.codegen._ +import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide} +import org.apache.spark.sql.catalyst.planning.RangeJoin +import org.apache.spark.sql.catalyst.plans._ +import org.apache.spark.sql.catalyst.plans.physical.{BroadcastDistribution, Distribution, Partitioning, UnspecifiedDistribution} +import org.apache.spark.sql.execution.{CodegenSupport, ExplainUtils, SparkPlan} +import org.apache.spark.sql.execution.metric.SQLMetrics + +/** + * Performs a range join on two tables. Supported types are inner, left outer, right outer, + * left semi, and left anti. The preserved side is streamed: unmatched build rows are not + * emitted, because the broadcast is copied to every partition. See `ExtractRangeJoinKeys` + * for the exact set of join-condition shapes that are recognized as a range join. + * + * The build side is an index. Point-in-range and interval overlap both build an + * [[IntervalIndex]] and probe it with `overlapping`. A single inequality builds a + * [[PointIndex]] and scans the side of the bound the build side holds. Keys that + * map to one order-preserving long instead build [[LongIntervalIndex]] or + * [[LongPointIndex]], so their probes compare raw longs without boxing. Either + * probe returns a superset. A candidate is kept only when the original join + * condition is true, so inclusivity stays on the predicate. Heavy overlap can + * make one probe return almost every build row. + */ +case class BroadcastRangeJoinExec( + left: SparkPlan, + right: SparkPlan, + buildSide: BuildSide, + joinType: JoinType, + leftKeys: Seq[Expression], + rightKeys: Seq[Expression], + condition: Option[Expression], + rangeJoin: RangeJoin) + extends JoinCodegenSupport { + + override lazy val metrics = Map( + "numOutputRows" -> SQLMetrics.createMetric(sparkContext, "number of output rows")) + + override lazy val requiredChildDistribution: Seq[Distribution] = { + val mode = RangeRelationBroadcastMode( + BindReferences.bindReferences(buildKeys, buildOutput), + RangeIndexKind(rangeJoin)) + buildSide match { + case BuildLeft => + BroadcastDistribution(mode) :: UnspecifiedDistribution :: Nil + case BuildRight => + UnspecifiedDistribution :: BroadcastDistribution(mode) :: Nil + } + } + + override def output: Seq[Attribute] = { + joinType match { + case _: InnerLike => + left.output ++ right.output + case LeftOuter => + left.output ++ right.output.map(_.withNullability(true)) + case RightOuter => + left.output.map(_.withNullability(true)) ++ right.output + case LeftSemi | LeftAnti => + left.output + case t => + throw SparkException.internalError(s"BroadcastRangeJoin should not take $t as the JoinType") + } + } + + // The preserved side is streamed, so its partitioning and ordering pass through. + override def outputPartitioning: Partitioning = streamedPlan.outputPartitioning + + override def outputOrdering: Seq[SortOrder] = streamedPlan.outputOrdering + + override def simpleStringWithNodeId(): String = { + val opId = ExplainUtils.getOpId(this) + s"$nodeName $joinType $buildSide ($opId)".trim + } + + private[this] lazy val buildPlan = if (buildSide == BuildLeft) left else right + private[this] lazy val streamedPlan = if (buildSide == BuildLeft) right else left + + private[this] lazy val buildKeys = if (buildSide == BuildLeft) leftKeys else rightKeys + private[this] lazy val streamedKeys = if (buildSide == BuildLeft) rightKeys else leftKeys + + @transient private[this] lazy val buildOutput = buildPlan.output + @transient private[this] lazy val streamedOutput = streamedPlan.output + + // The point-index scan this join's build side probes with. See `PointScan.forJoin`. + private lazy val pointScan: PointScan = PointScan.forJoin(rangeJoin, buildSide) + + // Alias so `innerJoin`/`outerJoin`/`leftExistenceJoin` need no change below; the class + // itself lives on the companion object (see there for why). + private type ProbeContext = BroadcastRangeJoinExec.ProbeContext + + private def perPartition( + relation: Broadcast[RangeRelation])( + f: ProbeContext => Iterator[InternalRow]): RDD[InternalRow] = { + val boundStreamKeys = BindReferences.bindReferences(streamedKeys, streamedOutput) + val cond = condition.getOrElse( + throw SparkException.internalError("BroadcastRangeJoin requires a join condition.")) + val schema = left.output ++ right.output + val numBuildFields = buildPlan.output.length + val join = rangeJoin + val side = buildSide + streamedPlan.execute().mapPartitionsInternal { stream => + // relation.value is read here, inside the task closure, so each executor fetches + // the index through the broadcast machinery (and shares it across the partitions + // that closure runs over) instead of the driver baking the materialized index + // into the closure that gets shipped out. PointScan.forJoin is passed by-name into + // makeCandidateFn, so it is only evaluated - and only throws for a non-partial-range + // join - inside the point-index branches that actually need it, same as before. It Review Comment: Fixed: removed the stale "same as before"/"need no change" comments. ########## sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/planning/patterns.scala: ########## @@ -421,6 +422,197 @@ object ExtractSingleColumnNullAwareAntiJoin extends JoinSelectionHelper with Pre } } +/** + * Finds joins whose condition is a range predicate usable by a broadcast range join. Review Comment: Fixed: doc moved above the extractor and now mentions the strict flag's use. ########## sql/core/src/test/scala/org/apache/spark/sql/execution/joins/RangeJoinSQLSuite.scala: ########## @@ -0,0 +1,745 @@ +/* + * 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.joins + +import org.apache.spark.SparkException +import org.apache.spark.sql.QueryTest +import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight} +import org.apache.spark.sql.catalyst.planning.{IntervalOverlapJoin, LessPartialRangeJoin} +import org.apache.spark.sql.catalyst.plans.{Inner, LeftAnti, LeftOuter, LeftSemi} +import org.apache.spark.sql.execution._ +import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.test.SharedSparkSession + +/** + * Planner and end-to-end coverage for broadcast range join. + * + * Which condition is recognized: `ExtractRangeJoinKeysSuite`. + * What the index returns: `RangeIndexSuite`. + * What the operator returns: `RangeJoinSuite`. + * This suite checks that the planner chooses the operator, and that the SQL + * result matches broadcast nested loop join where that comparison is the point. + */ +class RangeJoinSQLSuite extends QueryTest with SharedSparkSession with AdaptiveSparkPlanHelper { + + private val pointCond = + "s.ip_add_int >= l.BEGIN_IP2LONG AND s.ip_add_int <= l.END_IP2LONG" + + private def findRangeJoin(plan: SparkPlan): Seq[BroadcastRangeJoinExec] = { + collect(plan) { case j: BroadcastRangeJoinExec => j } + } + + /** `n` range joins sit in one whole-stage, not in nested stages or input adapters. */ + private def oneCodegenStageOwns(plan: SparkPlan, n: Int): Boolean = { + exists(plan) { + case w: WholeStageCodegenExec => + def count(p: SparkPlan): Int = p match { + case _: InputAdapter | _: WholeStageCodegenExec => 0 + case j: BroadcastRangeJoinExec => 1 + j.children.map(count).sum + case other => other.children.map(count).sum + } + count(w.child) == n + case _ => false + } + } + + /** The range join is compiled into a whole-stage, not hidden under an input adapter. */ + private def rangeJoinIsCodegen(plan: SparkPlan): Boolean = { + exists(plan) { + case w: WholeStageCodegenExec => + def contains(p: SparkPlan): Boolean = p match { + case _: InputAdapter => false + case _: BroadcastRangeJoinExec => true + case other => other.children.exists(contains) + } + contains(w.child) + case _ => false + } + } + + private def findBroadcastNestedLoopJoin(plan: SparkPlan): Seq[BroadcastNestedLoopJoinExec] = { + collect(plan) { case j: BroadcastNestedLoopJoinExec => j } + } + + private def findCartesianProduct(plan: SparkPlan): Seq[CartesianProductExec] = { + collect(plan) { case c: CartesianProductExec => c } + } + + private def withRangeJoin(extra: (String, String)*)(f: => Unit): Unit = { + withSQLConf((Seq(SQLConf.BROADCAST_RANGE_JOIN_ENABLED.key -> "true") ++ extra): _*)(f) + } + + /** + * One query per key family and join type that a range join can plan, shared by + * "whole-stage codegen covers each range-join probe" and "interpreted execution + * covers each range-join probe" so the two coverage lists cannot drift apart. + */ + private def probeCoverageQueries: Seq[String] = Seq( Review Comment: Added a boxed two-key interval query exercising the IntervalIndex branch. ########## sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/planning/ExtractRangeJoinKeysSuite.scala: ########## @@ -0,0 +1,208 @@ +/* + * 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.catalyst.planning + +import org.apache.spark.sql.catalyst.dsl.expressions._ +import org.apache.spark.sql.catalyst.expressions._ +import org.apache.spark.sql.catalyst.plans.{Inner, PlanTest} +import org.apache.spark.sql.catalyst.plans.logical.{Join, JoinHint, LocalRelation} +import org.apache.spark.sql.types.{DayTimeIntervalType, DoubleType, IntegerType} + +/** + * Plan-level coverage for [[ExtractRangeJoinKeys]]: which condition shapes are + * recognized, which `RangeJoin` tag is produced, and what residual remains. + */ +class ExtractRangeJoinKeysSuite extends PlanTest { + + private val points = LocalRelation($"p".int) + private val ranges = LocalRelation($"lo".int, $"hi".int) + private val intervalsA = LocalRelation($"lo".int, $"hi".int) + private val intervalsB = LocalRelation($"lo".int, $"hi".int) + + private val p = points.output.head + private val rLo = ranges.output(0) + private val rHi = ranges.output(1) + private val aLo = intervalsA.output(0) + private val aHi = intervalsA.output(1) + private val bLo = intervalsB.output(0) + private val bHi = intervalsB.output(1) + + private def join(left: LocalRelation, right: LocalRelation, cond: Expression): Join = + Join(left, right, Inner, Some(cond), JoinHint.NONE) + + private def assertKeys( + left: LocalRelation, + right: LocalRelation, + cond: Expression, + expected: RangeJoin, + leftKeys: Seq[Expression], + rightKeys: Seq[Expression]): Unit = { + join(left, right, cond) match { + case ExtractRangeJoinKeys(_, actualLeft, actualRight, rangeJoin) => + assert(rangeJoin === expected) + assert(actualLeft === leftKeys) + assert(actualRight === rightKeys) + case other => fail(s"expected $expected, got $other") + } + } + + private def assertUnrecognized(plan: Join): Unit = { + assert(ExtractRangeJoinKeys.unapply(plan).isEmpty) + } + + test("point-in-range keys are the point and the two bounds") { + // Inclusivity and conjunct order stay on the condition. BETWEEN lowers to a With, + // but RewriteWithExpression resolves it to a plain And of comparisons before physical + // planning runs, so by the time this extractor sees the condition it is already one + // of these shapes. + val inclusive = And(GreaterThanOrEqual(p, rLo), LessThanOrEqual(p, rHi)) + val exclusive = And(GreaterThan(p, rLo), LessThan(p, rHi)) + val highFirst = And(LessThanOrEqual(p, rHi), GreaterThan(p, rLo)) + val lowFirst = And(GreaterThanOrEqual(p, rLo), LessThan(p, rHi)) + Seq(inclusive, exclusive, highFirst, lowFirst).foreach { cond => + assertKeys(points, ranges, cond, PointInRangeJoin, Seq(p), Seq(rLo, rHi)) + } + assertKeys( + points, ranges, And(inclusive, Not(EqualTo(p, rLo))), + PointInRangeJoin, Seq(p), Seq(rLo, rHi)) + } + + test("point-in-range tie-break: the earliest candidate pair wins") { + // Three predicates give three candidate pairs under predicates.combinations(2): + // (P0, P1), (P0, P2), (P1, P2). (P0, P1) and (P0, P2) are each independently a + // valid point-in-range join -- p bounded by (rLo3, rHi3), and separately by + // (rLo3, rHi2) -- while (P1, P2) shares no endpoint. combinations(2) enumerates + // (P0, P1) first, so that is the pair that must win, not the otherwise-equally + // valid (P0, P2). + val ranges3 = LocalRelation($"lo".int, $"hi".int, $"hi2".int) + val rLo3 = ranges3.output(0) + val rHi3 = ranges3.output(1) + val rHi2 = ranges3.output(2) + val cond = And( + And(GreaterThanOrEqual(p, rLo3), LessThanOrEqual(p, rHi3)), + LessThanOrEqual(p, rHi2)) + assertKeys(points, ranges3, cond, PointInRangeJoin, Seq(p), Seq(rLo3, rHi3)) + } + + test("non-deterministic expressions are not range keys") { Review Comment: Added a referencing-and-nondeterministic key case to the test. ########## sql/core/src/test/scala/org/apache/spark/sql/execution/joins/RangeJoinSuite.scala: ########## @@ -0,0 +1,435 @@ +/* + * 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.joins + +import org.apache.spark.sql.{DataFrame, QueryTest, Row} +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{And, Expression, GreaterThan, GreaterThanOrEqual, LessThan, LessThanOrEqual, UnaryExpression} +import org.apache.spark.sql.catalyst.expressions.codegen._ +import org.apache.spark.sql.catalyst.expressions.codegen.Block._ +import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide} +import org.apache.spark.sql.catalyst.planning.{GreaterPartialRangeJoin, IntervalOverlapJoin, LessPartialRangeJoin, PointInRangeJoin, RangeJoin} +import org.apache.spark.sql.catalyst.plans.{Inner, JoinType, LeftAnti, LeftOuter, LeftSemi, RightOuter} +import org.apache.spark.sql.execution.SparkPlan +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.test.SharedSparkSession +import org.apache.spark.sql.types.{DataType, IntegerType, StructType} +import org.apache.spark.util.LongAccumulator + +/** + * Wraps `child` and adds to `counter` each time it is evaluated, interpreted or + * generated. Lets a test count how many build-side candidates a join actually probes, + * without needing a metric on the operator under test. + */ +private[joins] case class CountingPredicate(counter: LongAccumulator, child: Expression) + extends UnaryExpression { + + override def dataType: DataType = child.dataType + override def nullable: Boolean = child.nullable + + override def eval(input: InternalRow): Any = { + counter.add(1) + child.eval(input) + } + + override protected def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = { + val childGen = child.genCode(ctx) + val counterTerm = ctx.addReferenceObj("counter", counter) + val code = code""" + |$counterTerm.add(1); + |${childGen.code} + |""" + ev.copy(code = code, isNull = childGen.isNull, value = childGen.value) + } + + override protected def withNewChildInternal(newChild: Expression): CountingPredicate = + copy(child = newChild) +} + +/** + * Operator results for [[BroadcastRangeJoinExec]]. The planner is covered by + * `RangeJoinSQLSuite`; recognition and the index are covered below that. + */ +class RangeJoinSuite extends QueryTest with SharedSparkSession { + private lazy val intervals1: DataFrame = spark.createDataFrame( + sparkContext.parallelize(Seq( + Row(-1, 0), + Row(0, 1), + Row(0, 2), + Row(1, 5) + )), new StructType().add("low1", IntegerType).add("high1", IntegerType)) + + private lazy val intervals2: DataFrame = spark.createDataFrame( + sparkContext.parallelize(Seq( + Row(-2, -1), + Row(1, 3), + Row(5, 7) + )), new StructType().add("low2", IntegerType).add("high2", IntegerType)) + + private lazy val points: DataFrame = spark.createDataFrame( + sparkContext.parallelize(Seq( + Row(-3), + Row(1), + Row(3), + Row(6) + )), new StructType().add("point", IntegerType)) + + /** + * Operator-level helper. Keys are the original expressions from each child. + * `condition` is the original join predicate and the only accept/reject check. + */ + private def rangeJoinExec( + left: SparkPlan, + right: SparkPlan, + buildSide: BuildSide, + joinType: JoinType, + buildKeys: Seq[Expression], + streamedKeys: Seq[Expression], + condition: Expression, + rangeJoin: RangeJoin = PointInRangeJoin): BroadcastRangeJoinExec = { + val (leftKeys, rightKeys) = buildSide match { + case BuildLeft => (buildKeys, streamedKeys) + case BuildRight => (streamedKeys, buildKeys) + } + BroadcastRangeJoinExec( + left, right, buildSide, joinType, leftKeys, rightKeys, Some(condition), rangeJoin) + } + + private def pointCondition( + point: Expression, + low: Expression, + high: Expression, + lowInclusive: Boolean, + highInclusive: Boolean, + highFirst: Boolean = false): Expression = { + val lowPred = if (lowInclusive) GreaterThanOrEqual(point, low) else GreaterThan(point, low) + val highPred = if (highInclusive) LessThanOrEqual(point, high) else LessThan(point, high) + if (highFirst) And(highPred, lowPred) else And(lowPred, highPred) + } + + private def checkAnswerBothModes( + left: DataFrame, + right: DataFrame, + plan: (SparkPlan, SparkPlan) => SparkPlan, + expected: Seq[Row]): Unit = { + Seq(true, false).foreach { codegen => + withSQLConf( + org.apache.spark.sql.internal.SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> codegen.toString, + org.apache.spark.sql.internal.SQLConf.CODEGEN_FALLBACK.key -> "false") { + checkAnswer2(left, right, plan, expected) + } + } + } + + private def checkPointJoin( + left: DataFrame, + right: DataFrame, + pointOnLeft: Boolean, + buildSide: BuildSide, + joinType: JoinType, + lowInclusive: Boolean, + highInclusive: Boolean, + expected: Seq[Row], + highFirst: Boolean = false): Unit = { + checkAnswerBothModes(left, right, (leftPlan: SparkPlan, rightPlan: SparkPlan) => { + val pointPlan = if (pointOnLeft) leftPlan else rightPlan + val boundPlan = if (pointOnLeft) rightPlan else leftPlan + val point = pointPlan.output.head + val bounds = boundPlan.output + val pointKeys = point :: Nil + val buildIsPoint = buildSide match { + case BuildRight => !pointOnLeft + case BuildLeft => pointOnLeft + } + val (buildKeys, streamedKeys): (Seq[Expression], Seq[Expression]) = + if (buildIsPoint) (pointKeys, bounds) else (bounds, pointKeys) + rangeJoinExec( + leftPlan, rightPlan, buildSide, joinType, buildKeys, streamedKeys, + pointCondition(point, bounds(0), bounds(1), lowInclusive, highInclusive, highFirst)) + }, expected) + } + + test("point-in-range keeps boundary inclusivity on the condition") { + val lowInclHighExcl = Seq((0, 2, 1), (1, 5, 1), (1, 5, 3)).map(Row.fromTuple) + val bothExcl = Seq((0, 2, 1), (1, 5, 3)).map(Row.fromTuple) + val bothIncl = Seq((1, 0, 1), (1, 0, 2), (1, 1, 5), (3, 1, 5)).map(Row.fromTuple) + + Seq(false, true).foreach { highFirst => + checkPointJoin( + intervals1, points, pointOnLeft = false, BuildRight, Inner, + lowInclusive = true, highInclusive = false, lowInclHighExcl, highFirst) + } + checkPointJoin( + intervals1, points, pointOnLeft = false, BuildRight, Inner, + lowInclusive = false, highInclusive = false, bothExcl) + checkPointJoin( + points, intervals1, pointOnLeft = true, BuildRight, Inner, + lowInclusive = true, highInclusive = true, bothIncl) + checkPointJoin( + intervals1, points, pointOnLeft = false, BuildLeft, Inner, + lowInclusive = false, highInclusive = false, bothExcl) + checkPointJoin( + intervals1, points, pointOnLeft = false, BuildLeft, Inner, + lowInclusive = true, highInclusive = false, lowInclHighExcl) + } + + test("outer, semi, and anti preserve unmatched streamed rows") { + // low <= point && point < high. (-1, 0) and (0, 1) contain no point. + checkPointJoin( + intervals1, points, pointOnLeft = false, BuildRight, LeftOuter, + lowInclusive = true, highInclusive = false, + Seq((-1, 0, null), (0, 1, null), (0, 2, 1), (1, 5, 1), (1, 5, 3)).map(Row.fromTuple)) + checkPointJoin( + intervals1, points, pointOnLeft = false, BuildLeft, RightOuter, + lowInclusive = true, highInclusive = false, + Seq((null, null, -3), (0, 2, 1), (1, 5, 1), (1, 5, 3), (null, null, 6)).map(Row.fromTuple)) + checkPointJoin( + intervals1, points, pointOnLeft = false, BuildRight, LeftSemi, + lowInclusive = true, highInclusive = false, + Seq((0, 2), (1, 5)).map(Row.fromTuple)) + checkPointJoin( + intervals1, points, pointOnLeft = false, BuildRight, LeftAnti, + lowInclusive = true, highInclusive = false, + Seq((-1, 0), (0, 1)).map(Row.fromTuple)) + } + + test("a null streamed bound is not a match") { + val nullInterval = spark.createDataFrame( + sparkContext.parallelize(Seq(Row(null, 5))), + new StructType().add("low1", IntegerType).add("high1", IntegerType)) + checkPointJoin( + nullInterval, points, pointOnLeft = false, BuildRight, LeftOuter, + lowInclusive = true, highInclusive = false, Seq(Row(null, 5, null))) + checkPointJoin( + nullInterval, points, pointOnLeft = false, BuildRight, LeftSemi, + lowInclusive = true, highInclusive = false, Seq.empty) + checkPointJoin( + nullInterval, points, pointOnLeft = false, BuildRight, LeftAnti, + lowInclusive = true, highInclusive = false, Seq(Row(null, 5))) + } + + test("partial range scans upTo or from according to the lower bound") { + // a.lo < b.hi builds the high side and scans from the stream key. + // a.lo > b.hi builds the low side and scans up to the stream key. + def check( + rangeJoin: RangeJoin, + condition: (SparkPlan, SparkPlan) => Expression, + expected: Seq[Row], + buildSide: BuildSide = BuildRight): Unit = { + checkAnswerBothModes(intervals1, intervals2, (left: SparkPlan, right: SparkPlan) => { + val (buildKeys, streamedKeys) = buildSide match { + case BuildRight => (right.output(1) :: Nil, left.output.head :: Nil) + case BuildLeft => (left.output.head :: Nil, right.output(1) :: Nil) + } + rangeJoinExec( + left, right, buildSide, Inner, + buildKeys = buildKeys, + streamedKeys = streamedKeys, + condition = condition(left, right), + rangeJoin = rangeJoin) + }, expected) + } + check( + LessPartialRangeJoin(strict = true), + (left, right) => LessThan(left.output.head, right.output(1)), + Seq( + (-1, 0, 1, 3), (-1, 0, 5, 7), + (0, 1, 1, 3), (0, 1, 5, 7), + (0, 2, 1, 3), (0, 2, 5, 7), + (1, 5, 1, 3), (1, 5, 5, 7)).map(Row.fromTuple)) + check( + GreaterPartialRangeJoin(strict = true), + (left, right) => GreaterThan(left.output.head, right.output(1)), + Seq((0, 1, -2, -1), (0, 2, -2, -1), (1, 5, -2, -1)).map(Row.fromTuple)) + // Same join, built from the other side: GreaterPartialRangeJoin + BuildLeft takes the + // opposite PointIndex scan arm (from, not upTo) from the BuildRight case above. + check( + GreaterPartialRangeJoin(strict = true), + (left, right) => GreaterThan(left.output.head, right.output(1)), + Seq((0, 1, -2, -1), (0, 2, -2, -1), (1, 5, -2, -1)).map(Row.fromTuple), + buildSide = BuildLeft) + } + + test("interval overlap probes both bounds") { + val condition = (left: SparkPlan, right: SparkPlan) => And( + LessThan(left.output(0), right.output(1)), + LessThan(right.output(0), left.output(1))) + checkAnswerBothModes(intervals1, intervals2, (left: SparkPlan, right: SparkPlan) => { + rangeJoinExec( + left, right, BuildRight, Inner, + buildKeys = right.output, + streamedKeys = left.output, + condition = condition(left, right), + rangeJoin = IntervalOverlapJoin) + }, + Seq( + (0, 2, 1, 3), + (1, 5, 1, 3) + ).map(Row.fromTuple)) + } + + test("strict partial scans skip duplicated build keys instead of scanning them") { + // A strict inequality rejects build keys equal to the probe value, and the + // scan arm must skip them rather than surface them for the condition to + // reject one by one. With 10 duplicated rows per date and the probe at + // date 50, `<` must evaluate the condition on exactly the 49 later + // dates' rows and `>` on exactly the 50 earlier ones: an inclusive scan + // wired by mistake surfaces the probe's own 10 duplicates too, and an + // over-exclusive one silently drops passing rows. Both strict invariant + // answers are checked elsewhere; this pins the surfaced count itself. + val dup = 10 + val n = 100 + val probe = 50 + val bounds = spark.createDataFrame( + sparkContext.parallelize((0 until n).flatMap(d => (0 until dup).map(_ => Row(d))), 1), + new StructType().add("b", IntegerType)) + val points = spark.createDataFrame( + sparkContext.parallelize(Seq(Row(probe))), + new StructType().add("a", IntegerType)) + + val leftPlan = points.queryExecution.sparkPlan + val rightPlan = bounds.queryExecution.sparkPlan + val a = leftPlan.output.head + val b = rightPlan.output.head + + // Strict `<` with the bound on the build side scans from the exclusive + // FROM bound; strict `>` scans up to the exclusive UP TO bound. The inner + // join evaluates every surfaced candidate, so each counter is exact. + val expectations = Seq( + (LessPartialRangeJoin(strict = true), LessThan(a, b), (n - 1 - probe) * dup), + (GreaterPartialRangeJoin(strict = true), GreaterThan(a, b), probe * dup)) + Seq(true, false).foreach { codegenEnabled => + val factoryMode = if (codegenEnabled) "CODEGEN_ONLY" else "NO_CODEGEN" + withSQLConf( + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> codegenEnabled.toString, + SQLConf.CODEGEN_FALLBACK.key -> "false", + SQLConf.CODEGEN_FACTORY_MODE.key -> factoryMode) { + expectations.foreach { case (tag, cond, expected) => + val counter = spark.sparkContext.longAccumulator("numCandidates") + val exec = rangeJoinExec( + leftPlan, rightPlan, BuildRight, Inner, + buildKeys = b :: Nil, + streamedKeys = a :: Nil, + condition = CountingPredicate(counter, cond), + rangeJoin = tag) + QueryTest.executePlan(exec, spark.sqlContext) + assert(counter.value == expected, + s"$tag (codegen=$codegenEnabled): examined ${counter.value} candidates, " + + s"expected exactly $expected, so the probe's own duplicated key was " + + "scanned or passing rows were dropped") + } + } + } + } + + test("left semi and left anti stop probing once the first candidate matches") { + // Every build row is the same fully-overlapping interval, so every one of them + // satisfies the point-in-range condition. Iterator.exists is then guaranteed, + // deterministically, to stop at the first candidate it pulls from the index - + // this is the mechanism leftExistenceJoin relies on for its early exit. + val n = 2000 + val overlapping = spark.createDataFrame( + sparkContext.parallelize((0 until n).map(_ => Row(0, 100))), + new StructType().add("low2", IntegerType).add("high2", IntegerType)) + val onePoint = spark.createDataFrame( + sparkContext.parallelize(Seq(Row(50))), new StructType().add("point", IntegerType)) + + val leftPlan = onePoint.queryExecution.sparkPlan + val rightPlan = overlapping.queryExecution.sparkPlan + val point = leftPlan.output.head + val bounds = rightPlan.output + + Seq(true, false).foreach { codegenEnabled => + val factoryMode = if (codegenEnabled) "CODEGEN_ONLY" else "NO_CODEGEN" + withSQLConf( + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> codegenEnabled.toString, + SQLConf.CODEGEN_FALLBACK.key -> "false", + SQLConf.CODEGEN_FACTORY_MODE.key -> factoryMode) { + Seq(LeftSemi, LeftAnti).foreach { joinType => + val counter = spark.sparkContext.longAccumulator("numCandidates") + val baseCond = + pointCondition(point, bounds(0), bounds(1), lowInclusive = true, highInclusive = true) + val exec = rangeJoinExec( + leftPlan, rightPlan, BuildRight, joinType, + buildKeys = bounds, streamedKeys = point :: Nil, + condition = CountingPredicate(counter, baseCond)) + QueryTest.executePlan(exec, spark.sqlContext) + assert(counter.value < n, + s"$joinType (codegen=$codegenEnabled) examined ${counter.value} of $n " + + "fully overlapping candidates; the first candidate should have ended the probe") + } + } + } + } + + test("left semi and left anti lazy probes over a build with no common point") { + // Staggered build intervals share no common point, so the probe cannot take + // the all-rows shortcut and must walk the interval tree instead. With codegen + // both on and off, the walk is lazy: the first matching candidate ends the probe. + val n = 2000 + val staggered = spark.createDataFrame( + sparkContext.parallelize((0 until n).map(i => Row(i, i + 100))), + new StructType().add("low2", IntegerType).add("high2", IntegerType)) + val onePoint = spark.createDataFrame( + sparkContext.parallelize(Seq(Row(150))), new StructType().add("point", IntegerType)) + + val leftPlan = onePoint.queryExecution.sparkPlan + val rightPlan = staggered.queryExecution.sparkPlan + val point = leftPlan.output.head + val bounds = rightPlan.output + + Seq(true, false).foreach { codegenEnabled => + val factoryMode = if (codegenEnabled) "CODEGEN_ONLY" else "NO_CODEGEN" + withSQLConf( + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> codegenEnabled.toString, + SQLConf.CODEGEN_FALLBACK.key -> "false", + SQLConf.CODEGEN_FACTORY_MODE.key -> factoryMode) { + Seq(LeftSemi, LeftAnti).foreach { joinType => + val counter = spark.sparkContext.longAccumulator("numCandidates") + val baseCond = + pointCondition(point, bounds(0), bounds(1), lowInclusive = true, highInclusive = true) + val exec = rangeJoinExec( + leftPlan, rightPlan, BuildRight, joinType, + buildKeys = bounds, streamedKeys = point :: Nil, + condition = CountingPredicate(counter, baseCond)) + QueryTest.executePlan(exec, spark.sqlContext) + assert(counter.value < n, Review Comment: Added pending-work-based laziness tests for both iterators in RangeIndexSuite. ########## sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/planning/patterns.scala: ########## @@ -421,6 +422,197 @@ object ExtractSingleColumnNullAwareAntiJoin extends JoinSelectionHelper with Pre } } +/** + * Tags the shape of a range-predicate join recognized by [[ExtractRangeJoinKeys]]. + * Point-in-range and interval overlap both probe an interval index. A partial range + * is one inequality; Less vs Greater records which side holds the lower bound. + */ +sealed abstract class RangeJoin + +/** A point-in-range join: a value from one side falls within a `[low, high]` range on the other. */ +case object PointInRangeJoin extends RangeJoin + +/** + * An interval-overlap join: `(low, high)` on each side. The build side is an + * interval index and the stream side probes it with `overlapping(low, high)`. + */ +case object IntervalOverlapJoin extends RangeJoin + +/** + * A partial-range join: one inequality, one column per side. + * `leftIsLower` says which child holds the lower bound. + */ +sealed abstract class PartialRangeJoin extends RangeJoin { + def leftIsLower: Boolean + + /** + * The inequality is `<`/`>` and not `<=`/`>=`. Inclusivity stays on the + * original condition, but the strict form also selects the exclusive + * point-index bound, so duplicated build keys equal to the probe value + * are skipped instead of failing the condition one by one. + */ + def strict: Boolean +} + +/** A `low <[=] high` partial-range join where the left side is the low bound. */ +final case class LessPartialRangeJoin(override val strict: Boolean) extends PartialRangeJoin { + override def leftIsLower: Boolean = true +} + +/** A `low <[=] high` partial-range join where the left side is the high bound. */ +final case class GreaterPartialRangeJoin(override val strict: Boolean) extends PartialRangeJoin { + override def leftIsLower: Boolean = false +} + +/** + * Finds joins whose condition is a range predicate usable by a broadcast range join. + * + * Tries the two-predicate shapes first (point-in-range, then interval overlap), + * which capture more of the condition, and falls back to the single + * cross-side inequality: keeps the first that matches. + * - point-in-range: one key probed as both ends of a `(low, high)` window on the other side. + * - interval overlap: a `(low, high)` key on each side. + * - partial range: one inequality, one column per side. + * Leftover conjuncts, including inclusivity ([[RangePredicate]] only extracts the Review Comment: Fixed: doc now mentions the strict flag and its use for exclusive point-index scans. -- 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]
