wangyum commented on code in PR #59027: URL: https://github.com/apache/spark/pull/59027#discussion_r4369773400
########## sql/core/src/test/scala/org/apache/spark/sql/execution/joins/RangeJoinSQLSuite.scala: ########## @@ -0,0 +1,793 @@ +/* + * 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 +import org.apache.spark.sql.types.DoubleType + +/** + * 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( + s"SELECT s.id, l.zip_id FROM signin_ip s JOIN ip_lookup l ON $pointCond", + s"SELECT s.id, l.zip_id FROM signin_ip s LEFT JOIN ip_lookup l ON $pointCond", + s"SELECT s.id, l.zip_id FROM signin_ip s RIGHT JOIN ip_lookup l ON $pointCond", + s"SELECT s.id FROM signin_ip s LEFT SEMI JOIN ip_lookup l ON $pointCond", + s"SELECT s.id FROM signin_ip s LEFT ANTI JOIN ip_lookup l ON $pointCond", + "SELECT p.p, r.lo, r.hi FROM dec_points p JOIN dec_ranges r ON p.p >= r.lo AND p.p <= r.hi", + "SELECT /*+ BROADCAST(r) */ p.s, r.lo, r.hi FROM str_points p JOIN str_ranges r " + + "ON p.s >= r.lo AND p.s <= r.hi", + "SELECT a.lo, a.hi, b.lo, b.hi FROM intervals_a a JOIN intervals_b b " + + "ON a.lo < b.hi AND b.lo < a.hi", + "SELECT a.lo, a.hi, b.lo, b.hi FROM intervals_a a LEFT JOIN intervals_b b ON a.hi < b.lo", + "SELECT /*+ BROADCAST(a) */ a.lo, b.hi FROM intervals_a a JOIN intervals_b b ON a.lo < b.hi", + // DOUBLE keys stay on the boxed interval index (see setupDoubleIntervalViews), unlike + // the INT two-key overlap above, which is long-specialized. + "SELECT a.lo, a.hi, b.lo, b.hi FROM dbl_ranges_a a JOIN dbl_ranges_b b " + + "ON a.lo < b.hi AND b.lo < a.hi", + // Double and string keys stay on the boxed point index, so these queries reach + // its generated probe code (or, interpreted, the per-index dispatch in + // BroadcastRangeJoinExec.makeCandidateFn), which int partial ranges never do. + "SELECT p.d, b.lo FROM dbl_partial_points p JOIN dbl_partial_bounds b ON p.d < b.lo", Review Comment: Added `<=` and `>=` queries for DOUBLE and STRING in both directions to `probeCoverageQueries`, with boundary values equal to build keys added in `setupPartialRangeViews`. ########## sql/core/src/test/scala/org/apache/spark/sql/execution/joins/RangeIndexSuite.scala: ########## @@ -0,0 +1,1086 @@ +/* + * 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 java.io.{ByteArrayInputStream, ByteArrayOutputStream, ObjectInputStream, ObjectOutputStream} + +import org.apache.spark.SparkFunSuite +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.BoundReference +import org.apache.spark.sql.catalyst.types.PhysicalDataType +import org.apache.spark.sql.types._ +import org.apache.spark.unsafe.Platform +import org.apache.spark.unsafe.types.UTF8String + +/** + * Unit tests for the range-join indexes: the boxed pair ([[IntervalIndex]], + * [[PointIndex]]), the raw-long pair ([[LongIntervalIndex]], [[LongPointIndex]]), + * and the broadcast-mode routing that picks between the two families. + * Operator-level behavior is covered by `RangeJoinSuite`. + */ +class RangeIndexSuite extends SparkFunSuite { + + private def row(id: Int): InternalRow = InternalRow(id) + + private def span(lo: Any, hi: Any, id: Int): (Any, Any, InternalRow) = (lo, hi, row(id)) + + private def buildIntervals(spans: Seq[(Any, Any, InternalRow)]): IntervalIndex = + buildIntervals(IntegerType, spans) + + /** + * Builds the boxed interval index through the production build. The row + * layout puts the asserted id at ordinal 0, exactly where a real build row + * carries its columns; the keys bind to ordinals 1 and 2. + */ + private def buildIntervals( + keyType: DataType, + spans: Seq[(Any, Any, InternalRow)]): IntervalIndex = { + val keys = Seq( + BoundReference(1, keyType, nullable = true), + BoundReference(2, keyType, nullable = true)) + IntervalIndex.build(keyType, keys, + spans.iterator.map(t => InternalRow(t._3.getInt(0), t._1, t._2)), + Some(spans.size.toLong)) + } + + /** The point-index build with the same row layout: id at 0, key at 1. */ + private def buildPoints(keyType: DataType, keyed: Seq[(Any, InternalRow)]): PointIndex = { + val key = Seq(BoundReference(1, keyType, nullable = true)) + PointIndex.build(keyType, key, + keyed.iterator.map(t => InternalRow(t._2.getInt(0), t._1)), + Some(keyed.size.toLong)) + } + + private def ids(rows: Iterator[InternalRow]): List[Int] = rows.map(_.getInt(0)).toList + + private def probeIds(index: IntervalIndex, low: Any, high: Any): List[Int] = + ids(index.overlapping(low, high)) + + /** + * An iterator that hands back the same mutable row on every `next()` call, each field + * set to the current index. Mirrors how Spark's row iterators commonly reuse a single + * buffer (e.g. `UnsafeRow` deserialization) instead of allocating a fresh row per record. + */ + private def reusingRowIterator( + fields: Int, + n: Int): (org.apache.spark.sql.catalyst.expressions.GenericInternalRow, + Iterator[InternalRow]) = { + val shared = new org.apache.spark.sql.catalyst.expressions.GenericInternalRow(fields) + val iter = new Iterator[InternalRow] { + private var i = 0 + override def hasNext: Boolean = i < n + override def next(): InternalRow = { + var c = 0 + while (c < fields) { shared.update(c, i); c += 1 } + i += 1 + shared + } + } + (shared, iter) + } + + private def reusingByteBufferStringIterator( + n: Int): (Array[Byte], Iterator[InternalRow]) = { + val sharedBytes = new Array[Byte](4) + val sharedRow = new org.apache.spark.sql.catalyst.expressions.GenericInternalRow(2) + val iter = new Iterator[InternalRow] { + private var i = 0 + override def hasNext: Boolean = i < n + override def next(): InternalRow = { + val s = f"$i%04d" + java.lang.System.arraycopy(s.getBytes("UTF-8"), 0, sharedBytes, 0, 4) + val utf8 = UTF8String.fromAddress(sharedBytes, Platform.BYTE_ARRAY_OFFSET, 4) + sharedRow.update(0, utf8) + sharedRow.update(1, i) + i += 1 + sharedRow + } + } + (sharedBytes, iter) + } + + private def javaRoundTrip(index: IntervalIndex): (Array[Byte], IntervalIndex) = { + val bytes = new ByteArrayOutputStream() + val out = new ObjectOutputStream(bytes) + out.writeObject(index) + out.close() + val raw = bytes.toByteArray + val in = new ObjectInputStream(new ByteArrayInputStream(raw)) + try { + (raw, in.readObject().asInstanceOf[IntervalIndex]) + } finally { + in.close() + } + } + + test("build keeps a point, orders an inverted interval, and drops a null bound") { + val index = buildIntervals(Seq( + span(5, 5, 5), + span(1, 3, 1), + span(3, 1, 3), + (null, 1, row(1)))) + // (3, 1) is stored as [1, 3]. The null bound is absent. + assert(ids(index.overlapping(2, 2)).sorted == List(1, 3)) + assert(ids(index.overlapping(5, 5)) == List(5)) + assert(ids(index.overlapping(4, 4)).isEmpty) + } + + test("IntervalIndex.build defends against a build-side iterator reusing its row buffer") { + // If build() captured keys or rows by reference instead of copying them, every stored + // entry would collapse to the iterator's last value once the shared buffer is mutated + // again after the build-side iterator is drained -- the hazard guarded against is any + // producer whose rows or projected keys alias a reused buffer (e.g. a typed UDF result), + // not specifically BroadcastExchangeExec: its own row iterator decodes a fresh UnsafeRow + // per record. + val n = 50 + val buildKeys = Seq( + BoundReference(0, IntegerType, nullable = true), + BoundReference(1, IntegerType, nullable = true)) + val (shared, rows) = reusingRowIterator(fields = 3, n) + val index = IntervalIndex.build(IntegerType, buildKeys, rows, Some(n)) + + // Simulate the caller reusing the buffer after the build-side iterator is exhausted. + shared.update(0, -1) + shared.update(1, -1) + shared.update(2, -1) + + (0 until n).foreach { i => + val payloads = index.overlapping(i, i).map(_.getInt(2)).toList + assert(payloads == List(i), s"interval $i was corrupted by later reuse of the row buffer") + } + } + + test("IntervalIndex.build defends against string keys sharing a reused byte buffer") { + val n = 50 + val buildKeys = Seq( + BoundReference(0, StringType, nullable = true), + BoundReference(0, StringType, nullable = true)) + val (sharedBytes, rows) = reusingByteBufferStringIterator(n) + val index = IntervalIndex.build(StringType, buildKeys, rows, Some(n)) + + java.util.Arrays.fill(sharedBytes, 'X'.toByte) + + (0 until n).foreach { i => + val key = UTF8String.fromString(f"$i%04d") + val matching = index.overlapping(key, key).toList + assert(matching.size == 1, s"interval $i was corrupted by shared buffer reuse") + assert(matching.head.getInt(1) == i) + } + } + + test("a one-key interval build owns a single copy of the point key") { + // One build key is one key used as both ends. The build must copy it once + // and serve both bounds from that copy, instead of copying twice. + val n = 50 + val oneKey = Seq(BoundReference(0, StringType, nullable = true)) + val (sharedBytes, rows) = reusingByteBufferStringIterator(n) + val index = IntervalIndex.build(StringType, oneKey, rows, Some(n)) + + java.util.Arrays.fill(sharedBytes, 'X'.toByte) + + (0 until n).foreach { i => + val key = UTF8String.fromString(f"$i%04d") + val matching = index.overlapping(key, key).toList + assert(matching.size == 1, s"interval $i was corrupted by shared buffer reuse") + assert(matching.head.getInt(1) == i) + } + + // SizeEstimator's identity tracking charges the shared bounds once, so the + // estimate stays below the same build made from two independently copied keys. + val twoKeys = Seq( + BoundReference(0, StringType, nullable = true), + BoundReference(0, StringType, nullable = true)) + val (_, moreRows) = reusingByteBufferStringIterator(n) + val twoCopy = IntervalIndex.build(StringType, twoKeys, moreRows, Some(n)) + assert(index.estimatedSize < twoCopy.estimatedSize) + } + + test("an inverted build interval reaches the windows the condition accepts") { + // Int keys route to the long index, so this also pins swap normalization there. + val index = RangeRelationBroadcastMode( + Seq(BoundReference(0, IntegerType, nullable = true), + BoundReference(1, IntegerType, nullable = true)), + IntervalIndexKind) + .transform(Array(InternalRow(3, 1), InternalRow(0, 5))) + .asInstanceOf[LongIntervalIndex] + assert(index.overlapping(0L, 4L).map(_.getInt(0)).toSet == Set(3, 0)) + assert(index.overlapping(4L, 9L).map(_.getInt(0)).toList == List(0)) + } + + test("interval tree returns each overlapping row once") { + // Spans arrive out of low order: [10, 20] then [0, 5]. + val unsorted = buildIntervals(Seq(span(10, 20, 0), span(0, 5, 1))) + assert(probeIds(unsorted, 3, 3) == List(1)) + assert(probeIds(unsorted, 15, 15) == List(0)) + + // Probe 50 is a build key. Row 0 spans it; row 1 starts there. + val spanning = buildIntervals(Seq(span(0, 100, 0), span(50, 200, 1))) + assert(probeIds(spanning, 50, 50).toSet == Set(0, 1)) + assert(probeIds(spanning, 25, 25) == List(0)) + assert(probeIds(spanning, 150, 150) == List(1)) + + val duplicated = buildIntervals(Seq(span(0, 100, 0), span(0, 100, 1))) + assert(probeIds(duplicated, 50, 50).sorted == List(0, 1)) + + val nested = buildIntervals(Seq( + span(0, 100, 0), span(10, 20, 1), span(10, 50, 2), span(90, 110, 3))) + def assertOnce(low: Int, expected: Set[Int]): Unit = { + val actual = probeIds(nested, low, low) + assert(actual.toSet == expected, s"probe $low -> $actual") + assert(actual.size == expected.size, s"probe $low duplicated $actual") + } + assertOnce(0, Set(0)) + assertOnce(10, Set(0, 1, 2)) + assertOnce(15, Set(0, 1, 2)) + assertOnce(30, Set(0, 2)) + assertOnce(95, Set(0, 3)) + assertOnce(105, Set(3)) + assertOnce(200, Set.empty) + + assert(probeIds(buildIntervals(Seq.empty), 1, 1).isEmpty) + + val withNull = buildIntervals(Seq(span(0, 10, 0), span(5, 5, 1), (null, 4, row(9)))) + assert(probeIds(withNull, null, 5).isEmpty) + assert(probeIds(withNull, 5, null).isEmpty) + // The null bound is dropped, so two intervals remain and the estimate covers both. + val one = buildIntervals(Seq(span(0, 10, 0))) + assert(withNull.estimatedSize > one.estimatedSize) + + // At an interior endpoint, the range that ends and the one that starts + // are each a candidate once. + val chained = buildIntervals(Seq(span(0, 1, 0), span(1, 2, 1), span(2, 3, 2))) + assert(probeIds(chained, 0, 0) == List(0)) + assert(probeIds(chained, 1, 1).sorted == List(0, 1)) + assert(probeIds(chained, 2, 2).sorted == List(1, 2)) + assert(probeIds(chained, 3, 3) == List(2)) + + // A probe whose endpoints are reversed is the closed window between them. + // [10, 25] meets [0, 30], [15, 40], and the point at 15. + val inverted = buildIntervals(Seq(span(0, 30, 0), span(15, 40, 1), span(15, 15, 2))) + assert(probeIds(inverted, 25, 10).sorted == List(0, 1, 2)) + assert(probeIds(inverted, 10, 20).sorted == List(0, 1, 2)) + } + + test("interval tree is correct under heavy build-side overlap") { + // Every interval nests inside the one before it, so a point in the middle + // hits about half of the rows. The tree still stores one node per row. Review Comment: Corrected the comment in `RangeIndexSuite` to clarify that the probe at $n$ (300) hits all 300 rows and the probe at $n/2$ (150) hits 151 rows. ########## sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/planning/ExtractRangeJoinKeysSuite.scala: ########## @@ -0,0 +1,214 @@ +/* + * 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 `RangeJoinKind` 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: RangeJoinKind, + leftKeys: Seq[Expression], + rightKeys: Seq[Expression]): Unit = { + join(left, right, cond) match { + case ExtractRangeJoinKeys(_, actualLeft, actualRight, rangeJoinKind) => + assert(rangeJoinKind === 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 Review Comment: Updated the comment in `ExtractRangeJoinKeysSuite` explaining that `(P1, P2)` shares an endpoint `p` but fails point-in-range chaining and fails `place` during interval-overlap regrouping. -- 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]
