cloud-fan commented on code in PR #59027: URL: https://github.com/apache/spark/pull/59027#discussion_r4340859486
########## 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: **Nit (P3):** `same as before` here, and `need no change below` on the `ProbeContext` alias (line 111), describe an earlier revision of this file that nobody reading the merged code can see. Could both comments state the invariant instead? `PointScan.forJoin`, which throws for non-partial-range joins, is forced only inside the point-index branches of `makeCandidateFn`, and the alias lets the join methods use the companion-object `ProbeContext` unqualified. **Verification:** - **Inspection:** The perPartition and ProbeContext comments describe only the current code's invariants, with no reference to a previous version. ########## sql/core/src/test/scala/org/apache/spark/sql/execution/benchmark/RangeJoinBenchmark.scala: ########## @@ -0,0 +1,492 @@ +/* + * 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.benchmark + +import java.sql.Date + +import org.apache.spark.benchmark.Benchmark +import org.apache.spark.sql.{Column, DataFrame} +import org.apache.spark.sql.catalyst.planning.{GreaterPartialRangeJoin, IntervalOverlapJoin, LessPartialRangeJoin, PointInRangeJoin, RangeJoin} +import org.apache.spark.sql.execution.joins.{BroadcastNestedLoopJoinExec, BroadcastRangeJoinExec} +import org.apache.spark.sql.functions.{broadcast, col, date_add, lit, lpad, when} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.StringType + +/** + * Benchmark broadcast range join against broadcast nested loop join. + * + * Long keys cross every supported shape with every supported join type. + * Inner joins broadcast each side once. Left anti probes values that match + * nothing. A partial range whose bound sits at the median is not selective. + * Wide interval overlap is mostly selective: a probe overlaps a build + * interval only when their starts sit within `wideSpan`, so about 2% of + * probes hit roughly `wideSpan / width` intervals each and the rest hit + * nothing. Point-in-range inner is repeated for date, timestamp, decimal, + * and UNICODE_CI strings. + * + * More cases probe the cost model directly: a streamed side of 1, 10, 100, + * 1000, 5000, or 10000 rows, where range join pays to sort the whole build + * side instead of one nested-loop scan, and a fully overlapping inner join, + * where every pair matches and the index prunes nothing. + * + * Staggered high-overlap builds cover the regime between those: most, but not + * all, build intervals overlap each probe window, miss intervals are + * interleaved so no subtree's max high falls under the probe, and the build + * has no common core, so the tree is walked without pruning and evaluated per + * candidate. Long and double (boxed) keys both run at 50% and 90% overlap. + * + * To run this benchmark: + * {{{ + * 1. build/sbt "sql/Test/runMain org.apache.spark.sql.execution.benchmark.RangeJoinBenchmark" + * 2. generate result: + * SPARK_GENERATE_BENCHMARK_FILES=1 build/sbt \ + * "sql/Test/runMain org.apache.spark.sql.execution.benchmark.RangeJoinBenchmark" + * Results will be written to "benchmarks/RangeJoinBenchmark-results.txt". + * }}} + */ +object RangeJoinBenchmark extends SqlBasedBenchmark { + + // An optional regex in the environment variable RANGE_BENCH_FILTER runs only the + // sections whose name matches, e.g. RANGE_BENCH_FILTER=".*tiny probe.*|.*full overlap.*". + // A workaround for sbt splitting runMain arguments on spaces. + private var nameFilter: Option[String] = Option(System.getenv("RANGE_BENCH_FILTER")) + + private def shouldRun(name: String): Boolean = nameFilter.forall(name.matches) + + private val buildRows = 4000 + private val probeRows = 400000 + private val width = 10 + // Wide overlap and the median bound concentrate the output, so they stay smaller. + private val heavyBuildRows = 2000 + private val heavyProbeRows = 100000 + private val halfBuildRows = 1000 + private val halfProbeRows = 10000 + private val wideSpan = 20 * width + + private val epoch = lit(Date.valueOf("1970-01-01")) + + // [0, buildRows * width). Each range covers `width` values and none overlap. + private lazy val ranges: DataFrame = cached(spark.range(buildRows).select( + (col("id") * width).as("lo"), + (col("id") * width + (width - 1)).as("hi"))) + + // One hit per point: the value falls in exactly one range above. + private lazy val points: DataFrame = cached(spark.range(probeRows).select( + (col("id") % (buildRows.toLong * width)).as("p"))) + + // Above every range, so a point-in-range probe matches nothing. + private lazy val misses: DataFrame = cached(spark.range(probeRows).select( + (lit(buildRows.toLong * width) + (col("id") % width)).as("p"))) + + // Short intervals. A probe overlaps the build intervals whose starts are near it. + private lazy val probeIntervals: DataFrame = + cached(intervals(probeRows, "a", width + width / 2)) + private lazy val buildIntervals: DataFrame = + cached(intervals(buildRows, "b", width + width / 2)) + // Starts after every build interval. + private lazy val missIntervals: DataFrame = cached(spark.range(probeRows).select( + (lit(buildRows.toLong * width + width) + col("id") * width).as("alo"), + (lit(buildRows.toLong * width + width * 2) + col("id") * width).as("ahi"))) + + // About 40 build intervals overlap each probe. Review Comment: **Nit (P3):** The tiny-probe sizes are fixed, but the wide-overlap text still doesn't match the generators. Both sides are `[10*id, 10*id + 200]`, so a probe overlaps a build interval only when `|i - j| <= 20`: about 2% of the 100,000 probes hit up to 41 of the 2,000 build intervals each, and the rest hit none. This line still says about 40 intervals overlap each probe. The class doc (lines 37-38) now says a hitting probe meets roughly `wideSpan / width`, which is 20, and line 74 still says wide overlap concentrates the output. The strict-scan change also left the partial-range comments at lines 382, 386, 417, and 421 saying the point index scans with `from()` or `upTo()`. Every one of those cases is a strict `<` or `>`, which `PointScan.forJoin` sends to `fromExclusive` or `upToExclusive`. Could these comments describe what the generators and scans actually do? This extends an [existing discussion](https://github.com/apache/spark/pull/59027#discussion_r4333126475). **Verification:** - **Inspection:** Each wide-overlap comment matches the generators: about 2% of the 100,000 probes overlap up to 41 of the 2,000 build intervals, and the rest overlap none. - **Inspection:** Each partial-range case comment names the scan selected for its strict inequality and build side. ########## 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: **Nit (P3):** Since the `RangeJoin` tags moved here from `optimizer/joins.scala`, this block is followed directly by `RangeJoin`'s own `/** ... */`, with no declaration between them. Scaladoc attaches only the second block, so this one documents nothing, and `object ExtractRangeJoinKeys` (line 480) has no doc, even though `BroadcastRangeJoinExec` and the planner docs point readers to it. The shape-order wording fixed last round is intact, just detached. The parenthetical at lines 434-435 is also stale. `RangePredicate` now returns `(dataType, low, high, strict)`, and the strict flag picks `fromExclusive`/`upToExclusive`. The original condition still re-checks inclusivity, but strictness is no longer left only to it. Could you move this block directly above `object ExtractRangeJoinKeys` and update that clause? This extends an [existing discussion](https://github.com/apache/spark/pull/59027#discussion_r4314225433). **Verification:** - **Inspection:** Each of ExtractRangeJoinKeys and RangeJoin is immediately preceded by exactly one doc comment describing it. - **Inspection:** The extractor doc's description of RangePredicate matches its four-field result and the strict flag's use in point-index scans. ########## 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: **Nit (P3):** The interpreted twin of this list says it covers every key family. But the only query here that builds a boxed `IntervalIndex` is `BROADCAST(r)` on collated strings, which streams one key, and the two-key interval probes use bigint, int, or `decimal(10,2)`, all of which build `LongIntervalIndex`. So with codegen off, the `case intervals: IntervalIndex` branch of `makeCandidateFn` never evaluates `highExpr`. The collated `BROADCAST(p)` case runs only with codegen on, and `RangeJoinSuite` uses ints. If that branch probed `[low, low]`, double, string, binary, or wide-decimal overlap joins would drop rows in interpreted execution and after codegen fallback, and every test would still pass. Could you add a boxed query whose streamed side supplies both bounds to this list? For example: `SELECT /*+ BROADCAST(p) */ p.s, r.lo, r.hi FROM str_points p JOIN str_ranges r ON p.s >= r.lo AND p.s <= r.hi`. **Verification:** - **Regression:** With whole-stage codegen disabled, a boxed-key join whose streamed side supplies both interval bounds returns the nested-loop result, and the check fails if the interpreted probe uses the low bound for both ends. - **Behavior:** The same boxed two-key query also returns the nested-loop result with whole-stage codegen enabled. ########## 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: **Nit (P3):** Every nondeterministic key here is `Rand(Literal(0L))` or `Cast(Rand(...), IntegerType)`, neither of which has references, and `supportedKeys` checks `key.references.nonEmpty` before `key.deterministic`. The references check rejects them all, so deleting `key.deterministic &&` leaves all three assertions green. CheckAnalysis already rejects nondeterministic join conditions, so the guard is defense in depth, but this is the test named for it. Could you add a key that references a column and is also nondeterministic, for example `LessThan(Cast(aLo, DoubleType), Add(Cast(bHi, DoubleType), Rand(Literal(0L))))`, and assert that it is not extracted? **Verification:** - **Regression:** A nondeterministic key expression that references a join side is not extracted as a range key, and the assertion fails if the deterministic guard is removed. ########## 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: **Nit (P3):** This revision removed the `steps` counters from `OverlapIterator` and `LongOverlapIterator`, and with them the two `RangeIndexSuite` tests that bounded steps after taking a few matches, one per index family. That leaves this test as the only laziness coverage, but its counter wraps the join condition, so it counts condition evaluations. Those are controlled by `exists` and the `!found` guard, and the fully overlapping test above already pins that early exit (it would read 2000). The counter can't see how far the walk advanced. If `LongOverlapIterator` drained its stack into a buffer up front, this fixture would still evaluate the condition once, so the counter reads 1 and the test passes, even with a tighter bound. The bound is also loose on its own terms: the walk surfaces only the 101 intervals with `50 <= i <= 150`, all of them match, and `101 < 2000` passes with no early exit at all. Could you restore a laziness check for both iterators in `RangeIndexSuite`? A read-only view of pending walk work, such as the number of spans left on the stack, would show that taking the first match leaves work pending, without adding writes to the probe loop. Then this assertion can be tightened to one evaluation per probe, with the comment limited to what the counter observes. This extends existing discussions [1](https://github.com/apache/spark/pull/59027#discussion_r4237663579), [2](https://github.com/apache/spark/pull/59027#discussion_r4314225464). **Verification:** - **Regression:** Over a staggered build with no common point, taking the first match from a boxed and from a long interval index leaves walk work pending, and the assertion fails if either overlap iterator is changed to collect every match before returning the first. - **Regression:** With codegen on and off, left semi and left anti joins over a build with no common point evaluate the condition once per probe when the first candidate matches, and the assertion fails if early exit is removed. - **Behavior:** Fully consumed, both overlap iterators still return the brute-force overlap set, including when the pending-work view is read partway through the walk. -- 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]
