spark-test commented on code in PR #59027:
URL: https://github.com/apache/spark/pull/59027#discussion_r4359606313
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala:
##########
@@ -1220,6 +1220,21 @@ object SQLConf {
.bytesConf(ByteUnit.BYTE)
.createWithDefaultString("10MB")
+ val BROADCAST_RANGE_JOIN_ENABLED =
buildConf("spark.sql.join.broadcastRangeJoin.enabled")
+ .internal()
+ .doc("When true, joins with range predicates (such as point-in-range,
interval-overlap, " +
+ "or single inequalities) are planned as broadcast range joins instead of
broadcast " +
+ "nested loop joins if the build side is eligible for broadcast.
Supported join types " +
+ "include inner, left outer, right outer, left semi, and left anti.
Broadcast range joins " +
+ "build an index on the broadcast side to significantly accelerate
selective range probes, " +
+ "but perform worse than nested loop joins when half or more of the build
intervals " +
Review Comment:
**Nit (P3):** Thanks for restoring the ratios. "Half or more" still takes in
the fully overlapping case, though, and that case runs at 1.0X-1.1X in all
three results files. A window that overlaps every stored interval satisfies
`probeLow <= coreHigh && probeHigh >= coreLow`, so `overlapping` returns the
row array without walking. The 0.5X-0.8X rows all come from the staggered
builds, where most but not all intervals overlap each probe. This doc is the
only guidance on when to leave the flag off. Could it say "most, but not all"
as the earlier wording did, and note that fully overlapping builds run at about
parity?
This extends an [existing
discussion](https://github.com/apache/spark/pull/59027#discussion_r4346073069).
**Verification:**
- **Inspection:** The config doc's description of when range joins are
slower names the case where most but not all build intervals overlap each
probe, keeps the measured ratios, and does not imply that fully overlapping
builds regress, consistent with every committed benchmark result file.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/joins/RangeJoinSQLSuite.scala:
##########
@@ -0,0 +1,771 @@
+/*
+ * 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(
+ 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",
+ "SELECT p.d, b.lo FROM dbl_partial_bounds b JOIN dbl_partial_points p ON
p.d > b.lo",
+ "SELECT p.s, b.lo FROM str_partial_points p JOIN str_partial_bounds b ON
p.s < b.lo",
+ "SELECT p.s, b.lo FROM str_partial_bounds b JOIN str_partial_points p ON
p.s > b.lo")
+
+ private def setupIpLookupViews(): Unit = {
+ sql(
+ """
+ |CREATE OR REPLACE TEMP VIEW signin_ip(id, ip_add_int) AS VALUES
+ | ('1', 1467630666L),
+ | ('2', 1186929717L),
+ | ('3', 3501544352L),
+ | ('31', 3501544320L),
+ | ('32', 3501544383L),
+ | ('33', 3501544330L),
+ | ('34', 3501544390L),
+ | ('4', 1626243295L),
+ | ('41', 1626243072L),
+ | ('42', 1626243327L),
+ | ('5', 1206355995L),
+ | ('5', 1206355995L),
+ | ('6', CAST(NULL AS BIGINT)),
+ | ('7', 0L)
+ |""".stripMargin)
+ sql(
+ """
+ |CREATE OR REPLACE TEMP VIEW ip_lookup(
+ | BEGIN_IP2LONG, END_IP2LONG, TWO_CHAR_CNTRY_CD, zip_id) AS VALUES
+ | (1467630592L, 1467630847L, 'de', 2380520L),
+ | (1186929664L, 1186930175L, 'us', 6723936L),
+ | (1186929664L, 1186930175L, 'us', 6723936L),
+ | (3501544320L, 3501544383L, 'ca', 12804631L),
+ | (3501544330L, 3501544390L, 'ca', 12804631001L),
+ | (1626243072L, 1626243327L, 'us', 8849753L),
+ | (1206355968L, 1206356223L, 'us', 10138801L),
+ | (CAST(NULL AS BIGINT), 1206356223L, 'us', 10138801L),
+ | (1206355968L, CAST(NULL AS BIGINT), 'us', 10138801L),
+ | (CAST(NULL AS BIGINT), CAST(NULL AS BIGINT), 'us', 10138801L)
+ |""".stripMargin)
+ }
+
+ private def setupIntervalViews(): Unit = {
+ sql(
+ """
+ |CREATE OR REPLACE TEMP VIEW intervals_a(lo, hi) AS VALUES
+ | (-1, 0), (0, 1), (0, 2), (1, 5)
+ |""".stripMargin)
+ sql(
+ """
+ |CREATE OR REPLACE TEMP VIEW intervals_b(lo, hi) AS VALUES
+ | (-2, -1), (-4, -2), (1, 3), (5, 7)
+ |""".stripMargin)
+ }
+
+ /**
+ * Two-column DOUBLE ranges, so a two-key interval-overlap query here builds
a boxed
+ * `IntervalIndex` (DoubleType has no bit-preserving long encoding) with
both bounds
+ * supplied by the streamed side - the `case intervals: IntervalIndex`
branch of
+ * `BroadcastRangeJoinExec.makeCandidateFn` that `intervals_a`/`intervals_b`
(INT,
+ * long-specialized) never reaches.
+ */
+ private def setupDoubleIntervalViews(): Unit = {
+ sql(
+ """
+ |CREATE OR REPLACE TEMP VIEW dbl_ranges_a(lo, hi) AS VALUES
+ | (-1.0, 0.0), (0.0, 1.0), (0.0, 2.0), (1.0, 5.0)
Review Comment:
**Nit (P3):** These literals are decimals, not doubles. Spark SQL types an
unsuffixed `1.0` as `DECIMAL(2,1)` (`misc-functions.sql.out` records
`typeof(1.2)` as `decimal(2,1)`); only `1.0D` or a cast gives a double.
`RangeIndex.longKey` maps `DECIMAL(2,1)` to an unscaled long, so the
`dbl_ranges` query builds a `LongIntervalIndex`, not the boxed `IntervalIndex`
that the comments here and at line 106 describe.
With codegen off, the `IntervalIndex` case of `makeCandidateFn`, which
evaluates the second streamed bound, still runs in no test. That is the
coverage the earlier round asked for. The collated `BROADCAST(p)` query is the
only one that streams two boxed bounds, and it runs with codegen on, while
`RangeJoinSuite` uses int keys. If that branch probed `[low, low]`, interpreted
double, string, binary, and wide-decimal overlap joins would drop rows and
every test would still pass.
Could these views use `-1.0D`-style literals or a cast, with a setup check
that the columns are DOUBLE? Adding the collated `BROADCAST(p)` query to the
shared list would also cover the branch.
**Verification:**
- **Regression:** With whole-stage codegen disabled, a two-key
interval-overlap join whose build side is a boxed interval index returns the
nested-loop answer, and the check fails if the interpreted probe uses the low
bound for both ends.
- **Inspection:** The setupDoubleIntervalViews doc and the probe-coverage
comment describe a key type that matches the views' declared column types.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/joins/RangeIndexSuite.scala:
##########
@@ -0,0 +1,1066 @@
+/*
+ * 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 -- exactly what a
broadcast's row
+ // iterator does once BroadcastExchangeExec finishes feeding it to
RangeRelationBroadcastMode.
+ 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.
+ val n = 300
+ val nested = buildIntervals((0 until n).map(i => span(i, 2 * n - i, i)))
+
+ def expected(probe: Int): Set[Int] =
+ (0 until n).filter(i => i <= probe && probe <= 2 * n - i).toSet
+
+ Seq(0, 1, n / 2, n, 2 * n - 1, 2 * n).foreach { probe =>
+ assert(probeIds(nested, probe, probe).toSet == expected(probe), s"probe
$probe")
+ }
+ }
+
+ test("broadcast of heavily overlapping intervals stays linear") {
+ // [i, i + n] for i in 0 until n. Every interval contains n. The broadcast
+ // is the flat arrays, one slot per interval, so serialization stays
linear.
+ val n = 200
+ val index = buildIntervals((0 until n).map(i => span(i, i + n, i)))
+ val (raw, restored) = javaRoundTrip(index)
+ assert(raw.length < n * 1024, s"serialized ${raw.length} bytes")
+ val probe = n
+ assert(probeIds(restored, probe, probe).toSet == probeIds(index, probe,
probe).toSet)
+ assert(probeIds(restored, probe, probe).size == n)
+ }
+
+ test("estimatedSize depends on the interval count, not the overlap depth") {
+ val n = 200
+ val one = buildIntervals(Seq(span(0, 1, 0)))
+ val disjoint = buildIntervals((0 until n).map(i => span(2 * i, 2 * i + 1,
i)))
+ val overlap = buildIntervals((0 until n).map(i => span(i, i + n, i)))
+ val again = buildIntervals((0 until n).map(i => span(2 * i, 2 * i + 1, i)))
+ assert(disjoint.estimatedSize > one.estimatedSize)
+ assert(overlap.estimatedSize > one.estimatedSize)
+ assert(disjoint.estimatedSize == again.estimatedSize)
+ // Nesting stores no extra copy of the rows. Key objects differ, so the two
+ // indexes need not be equal, but overlap stays the same order of
magnitude.
+ assert(overlap.estimatedSize < disjoint.estimatedSize * 3)
+ }
+
+ test("estimatedSize includes distinct key objects") {
+ val ints = buildIntervals(Seq(span(0, 1, 0)))
+ val wide = UTF8String.fromString("x" * 64)
+ val strings = buildIntervals(StringType, Seq((wide, wide, row(0))))
+ assert(strings.estimatedSize > ints.estimatedSize)
+ }
+
+ test("interval tree treats NaN as its own point and Infinity as a bound") {
+ // PhysicalDoubleType's ordering (SQLOrderingUtil.compareDoubles) totally
orders
+ // -Inf < ... < +Inf < NaN, with NaN comparing equal only to itself. So
[1.0, +Inf]
+ // includes +Inf under inclusive probe, and a NaN key overlaps only a
probe that
+ // also lands exactly on NaN.
+ val index = buildIntervals(DoubleType, Seq(
+ (1.0, Double.PositiveInfinity, row(1)),
+ (Double.NegativeInfinity, -1.0, row(2)),
+ (Double.NaN, Double.NaN, row(3))))
+
+ assert(ids(index.overlapping(5.0, 5.0)) == List(1))
+ assert(ids(index.overlapping(Double.PositiveInfinity,
Double.PositiveInfinity)) == List(1))
+ assert(ids(index.overlapping(-5.0, -5.0)) == List(2))
+ assert(ids(index.overlapping(Double.NegativeInfinity,
Double.NegativeInfinity)) == List(2))
+ assert(ids(index.overlapping(Double.NaN, Double.NaN)) == List(3))
+ assert(ids(index.overlapping(0.0, 0.0)).isEmpty)
+ }
+
+ test("interval index merges keys that compare equal but are not equal") {
+ // Array equals is identity, and UTF8String equals is binary, so these keys
+ // compare equal under the index ordering without being equal. The tree
+ // orders them with that comparator, so a probe on either key returns both.
+ // (Signed zero is the same kind of merge, but Double's == already treats
-0.0
+ // and 0.0 as equal, which fails this helper's != precondition - see the
+ // dedicated "-0.0 and 0.0 are treated as one key" test below instead.)
+ def assertMerged(
+ dataType: DataType,
+ low: Any,
+ leftKey: Any,
+ rightKey: Any,
+ high: Any,
+ earlier: Seq[Any]): Unit = {
+ val ordering = PhysicalDataType.ordering(dataType)
+ assert(leftKey != rightKey)
+ assert(ordering.compare(leftKey, rightKey) == 0)
+ val points = Array(
+ (leftKey, leftKey, row(0)),
+ (rightKey, rightKey, row(1)))
+ val pointIdx = buildIntervals(dataType, points.toSeq)
+ assert(ids(pointIdx.overlapping(leftKey, leftKey)).toSet == Set(0, 1))
+ assert(ids(pointIdx.overlapping(rightKey, rightKey)).toSet == Set(0, 1))
+
+ val intervals = earlier.zipWithIndex.map { case (key, i) =>
+ (key, key, row(10 + i))
+ }.toArray ++ Array(
+ (low, leftKey, row(0)),
+ (rightKey, high, row(1)))
+ val index = buildIntervals(dataType, intervals.toSeq)
+ assert(ids(index.overlapping(leftKey, leftKey)).sorted == List(0, 1))
+ assert(ids(index.overlapping(rightKey, rightKey)).sorted == List(0, 1))
+ }
+
+ val bytes = Array[Byte](1, 2)
+ val sameBytes = Array[Byte](1, 2)
+ assertMerged(
+ BinaryType,
+ Array[Byte](1, 0),
+ bytes,
+ sameBytes,
+ Array[Byte](2),
+ Seq(Array[Byte](0), Array[Byte](0, 1)))
+
+ val upper = UTF8String.fromString("M")
+ val lower = UTF8String.fromString("m")
+ assertMerged(
+ StringType("UNICODE_CI"),
+ UTF8String.fromString("c"),
+ upper,
+ lower,
+ UTF8String.fromString("z"),
+ Seq(UTF8String.fromString("a"), UTF8String.fromString("b")))
+ }
+
+ test("-0.0 and 0.0 are treated as one key") {
+ // Double's own `==` already treats -0.0 and 0.0 as equal, so
`PhysicalDataType.ordering`
+ // (which SQL requires to distinguish them from no other value, but not
from each other)
+ // merges them exactly like the UTF8String/Array[Byte] keys above: a probe
on either
+ // sign finds rows stored under both.
+ val negZero: Any = -0.0
+ val posZero: Any = 0.0
+ val ordering = PhysicalDataType.ordering(DoubleType)
+ assert(ordering.compare(negZero, posZero) == 0)
+
+ val points = Array((negZero, negZero, row(0)), (posZero, posZero, row(1)))
+ val pointIdx = buildIntervals(DoubleType, points.toSeq)
+ assert(ids(pointIdx.overlapping(negZero, negZero)).toSet == Set(0, 1))
+ assert(ids(pointIdx.overlapping(posZero, posZero)).toSet == Set(0, 1))
+
+ val intervals = Array((-5.0: Any, negZero, row(0)), (posZero, 5.0: Any,
row(1)))
+ val index = buildIntervals(DoubleType, intervals.toSeq)
+ assert(ids(index.overlapping(negZero, negZero)).sorted == List(0, 1))
+ assert(ids(index.overlapping(posZero, posZero)).sorted == List(0, 1))
+ }
+
+ test("a nested probe does not drop rows from the outer iterator") {
+ val index = buildIntervals((0 until 32).map(i => span(i, i + 10, i)))
+ val outer = index.overlapping(0, 100)
Review Comment:
**Nit (P3):** For `[i, i + 10]` with `i` in `0 until 32`, `coreLow` is 31
and `coreHigh` is 10. So `overlapping(0, 100)` passes the all-covering check
and returns the row slice, not an `OverlapIterator`. Only the inner
`overlapping(0, 5)` walks, so two walks never run on the index at the same
time. A change that shared walk state across probes on one index (the stacks,
`stackTop`, or `nextRow`), for example to save the per-probe allocation, would
still pass. Fused interval joins on one broadcast depend on exactly that
isolation. Could the outer window be one that has to walk, such as
`overlapping(0, 15)`, with an `isInstanceOf[OverlapIterator]` check? The same
case for `LongIntervalIndex` would cover the other family.
This extends an [existing
discussion](https://github.com/apache/spark/pull/59027#discussion_r4249393630).
**Verification:**
- **Regression:** For the boxed and long interval index families, an outer
overlap probe whose window does not cover every interval is confirmed to walk
the tree, runs interleaved with inner probes on the same index, and still
yields exactly the brute-force overlap set, failing if probes share walk state.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/joins/RangeIndexSuite.scala:
##########
@@ -0,0 +1,1066 @@
+/*
+ * 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 -- exactly what a
broadcast's row
+ // iterator does once BroadcastExchangeExec finishes feeding it to
RangeRelationBroadcastMode.
+ 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.
+ val n = 300
+ val nested = buildIntervals((0 until n).map(i => span(i, 2 * n - i, i)))
+
+ def expected(probe: Int): Set[Int] =
+ (0 until n).filter(i => i <= probe && probe <= 2 * n - i).toSet
+
+ Seq(0, 1, n / 2, n, 2 * n - 1, 2 * n).foreach { probe =>
+ assert(probeIds(nested, probe, probe).toSet == expected(probe), s"probe
$probe")
+ }
+ }
+
+ test("broadcast of heavily overlapping intervals stays linear") {
+ // [i, i + n] for i in 0 until n. Every interval contains n. The broadcast
+ // is the flat arrays, one slot per interval, so serialization stays
linear.
+ val n = 200
+ val index = buildIntervals((0 until n).map(i => span(i, i + n, i)))
+ val (raw, restored) = javaRoundTrip(index)
+ assert(raw.length < n * 1024, s"serialized ${raw.length} bytes")
+ val probe = n
+ assert(probeIds(restored, probe, probe).toSet == probeIds(index, probe,
probe).toSet)
+ assert(probeIds(restored, probe, probe).size == n)
+ }
+
+ test("estimatedSize depends on the interval count, not the overlap depth") {
+ val n = 200
+ val one = buildIntervals(Seq(span(0, 1, 0)))
+ val disjoint = buildIntervals((0 until n).map(i => span(2 * i, 2 * i + 1,
i)))
+ val overlap = buildIntervals((0 until n).map(i => span(i, i + n, i)))
+ val again = buildIntervals((0 until n).map(i => span(2 * i, 2 * i + 1, i)))
+ assert(disjoint.estimatedSize > one.estimatedSize)
+ assert(overlap.estimatedSize > one.estimatedSize)
+ assert(disjoint.estimatedSize == again.estimatedSize)
+ // Nesting stores no extra copy of the rows. Key objects differ, so the two
+ // indexes need not be equal, but overlap stays the same order of
magnitude.
+ assert(overlap.estimatedSize < disjoint.estimatedSize * 3)
+ }
+
+ test("estimatedSize includes distinct key objects") {
+ val ints = buildIntervals(Seq(span(0, 1, 0)))
+ val wide = UTF8String.fromString("x" * 64)
+ val strings = buildIntervals(StringType, Seq((wide, wide, row(0))))
+ assert(strings.estimatedSize > ints.estimatedSize)
+ }
+
+ test("interval tree treats NaN as its own point and Infinity as a bound") {
+ // PhysicalDoubleType's ordering (SQLOrderingUtil.compareDoubles) totally
orders
+ // -Inf < ... < +Inf < NaN, with NaN comparing equal only to itself. So
[1.0, +Inf]
+ // includes +Inf under inclusive probe, and a NaN key overlaps only a
probe that
+ // also lands exactly on NaN.
+ val index = buildIntervals(DoubleType, Seq(
+ (1.0, Double.PositiveInfinity, row(1)),
+ (Double.NegativeInfinity, -1.0, row(2)),
+ (Double.NaN, Double.NaN, row(3))))
+
+ assert(ids(index.overlapping(5.0, 5.0)) == List(1))
+ assert(ids(index.overlapping(Double.PositiveInfinity,
Double.PositiveInfinity)) == List(1))
+ assert(ids(index.overlapping(-5.0, -5.0)) == List(2))
+ assert(ids(index.overlapping(Double.NegativeInfinity,
Double.NegativeInfinity)) == List(2))
+ assert(ids(index.overlapping(Double.NaN, Double.NaN)) == List(3))
+ assert(ids(index.overlapping(0.0, 0.0)).isEmpty)
+ }
+
+ test("interval index merges keys that compare equal but are not equal") {
+ // Array equals is identity, and UTF8String equals is binary, so these keys
+ // compare equal under the index ordering without being equal. The tree
+ // orders them with that comparator, so a probe on either key returns both.
+ // (Signed zero is the same kind of merge, but Double's == already treats
-0.0
+ // and 0.0 as equal, which fails this helper's != precondition - see the
+ // dedicated "-0.0 and 0.0 are treated as one key" test below instead.)
+ def assertMerged(
+ dataType: DataType,
+ low: Any,
+ leftKey: Any,
+ rightKey: Any,
+ high: Any,
+ earlier: Seq[Any]): Unit = {
+ val ordering = PhysicalDataType.ordering(dataType)
+ assert(leftKey != rightKey)
+ assert(ordering.compare(leftKey, rightKey) == 0)
+ val points = Array(
+ (leftKey, leftKey, row(0)),
+ (rightKey, rightKey, row(1)))
+ val pointIdx = buildIntervals(dataType, points.toSeq)
+ assert(ids(pointIdx.overlapping(leftKey, leftKey)).toSet == Set(0, 1))
+ assert(ids(pointIdx.overlapping(rightKey, rightKey)).toSet == Set(0, 1))
+
+ val intervals = earlier.zipWithIndex.map { case (key, i) =>
+ (key, key, row(10 + i))
+ }.toArray ++ Array(
+ (low, leftKey, row(0)),
+ (rightKey, high, row(1)))
+ val index = buildIntervals(dataType, intervals.toSeq)
+ assert(ids(index.overlapping(leftKey, leftKey)).sorted == List(0, 1))
+ assert(ids(index.overlapping(rightKey, rightKey)).sorted == List(0, 1))
+ }
+
+ val bytes = Array[Byte](1, 2)
+ val sameBytes = Array[Byte](1, 2)
+ assertMerged(
+ BinaryType,
+ Array[Byte](1, 0),
+ bytes,
+ sameBytes,
+ Array[Byte](2),
+ Seq(Array[Byte](0), Array[Byte](0, 1)))
+
+ val upper = UTF8String.fromString("M")
+ val lower = UTF8String.fromString("m")
+ assertMerged(
+ StringType("UNICODE_CI"),
+ UTF8String.fromString("c"),
+ upper,
+ lower,
+ UTF8String.fromString("z"),
+ Seq(UTF8String.fromString("a"), UTF8String.fromString("b")))
+ }
+
+ test("-0.0 and 0.0 are treated as one key") {
+ // Double's own `==` already treats -0.0 and 0.0 as equal, so
`PhysicalDataType.ordering`
+ // (which SQL requires to distinguish them from no other value, but not
from each other)
+ // merges them exactly like the UTF8String/Array[Byte] keys above: a probe
on either
+ // sign finds rows stored under both.
+ val negZero: Any = -0.0
+ val posZero: Any = 0.0
+ val ordering = PhysicalDataType.ordering(DoubleType)
+ assert(ordering.compare(negZero, posZero) == 0)
+
+ val points = Array((negZero, negZero, row(0)), (posZero, posZero, row(1)))
+ val pointIdx = buildIntervals(DoubleType, points.toSeq)
+ assert(ids(pointIdx.overlapping(negZero, negZero)).toSet == Set(0, 1))
+ assert(ids(pointIdx.overlapping(posZero, posZero)).toSet == Set(0, 1))
+
+ val intervals = Array((-5.0: Any, negZero, row(0)), (posZero, 5.0: Any,
row(1)))
+ val index = buildIntervals(DoubleType, intervals.toSeq)
+ assert(ids(index.overlapping(negZero, negZero)).sorted == List(0, 1))
+ assert(ids(index.overlapping(posZero, posZero)).sorted == List(0, 1))
+ }
+
+ test("a nested probe does not drop rows from the outer iterator") {
+ val index = buildIntervals((0 until 32).map(i => span(i, i + 10, i)))
+ val outer = index.overlapping(0, 100)
+ val seen = scala.collection.mutable.ArrayBuffer.empty[Int]
+ // The inner probe runs while the outer iterator still has rows, the same
+ // shape as two interval joins fused in one whole-stage.
+ while (outer.hasNext) {
+ seen += outer.next().getInt(0)
+ assert(ids(index.overlapping(0, 5)).sorted == (0 to 5).toList)
+ }
+ assert(seen.sorted == (0 until 32).toList)
+ }
+
+ test("a point probe of disjoint intervals returns only the containing row") {
+ val n = 1000
+ val index = buildIntervals((0 until n).map(i => span(i * 2, i * 2 + 1, i)))
+ assert(probeIds(index, 2, 2) == List(1))
+ assert(probeIds(index, 3, 3) == List(1))
+ assert(probeIds(index, 4, 4) == List(2))
+ assert(probeIds(index, -1, -1).isEmpty)
+ assert(probeIds(index, 100000, 100000).isEmpty)
+ }
+
+ test("point index answers both sides of a bound, including equals") {
+ // Unsorted, with two rows on key 5. Equals stay, in input order.
+ val keyed = Array[(Any, InternalRow)](
+ (8, row(8)),
+ (1, row(1)),
+ (5, row(5)),
+ (5, row(50)),
+ (3, row(3)),
+ (null, row(-1)))
+ val index = buildPoints(IntegerType, keyed.toSeq)
+ assert(ids(index.upTo(5)) == List(1, 3, 5, 50))
+ assert(ids(index.upTo(4)) == List(1, 3))
+ assert(ids(index.upTo(0)).isEmpty)
+ assert(ids(index.upTo(8)) == List(1, 3, 5, 50, 8))
+ assert(ids(index.from(5)) == List(5, 50, 8))
+ assert(ids(index.from(6)) == List(8))
+ assert(ids(index.from(9)).isEmpty)
+ assert(ids(index.from(1)) == List(1, 3, 5, 50, 8))
+ assert(index.upTo(null).isEmpty)
+ assert(index.from(null).isEmpty)
+
+ val empty = buildPoints(IntegerType, Seq.empty)
+ assert(empty.upTo(1).isEmpty)
+ assert(empty.from(1).isEmpty)
+ // SizeEstimator charges the two empty arrays' shells, and the five kept
+ // points on top of them.
+ assert(empty.estimatedSize > 0)
+ assert(index.estimatedSize > empty.estimatedSize)
+ }
+
+ test("exclusive point bounds skip duplicated build keys") {
+ // 100 duplicate rows per key on both families. A strict partial predicate
+ // rejects the probe's own key, so the exclusive bound must skip it
+ // instead of surfacing 100 rows that all fail the condition.
+ val keyed = (1 to 5).flatMap(k => Seq.fill(100)((k: Any, row(k))))
+ val boxed = buildPoints(IntegerType, keyed)
+ val longPoints = RangeRelationBroadcastMode(
+ Seq(BoundReference(1, IntegerType, nullable = true)), PointIndexKind)
+ .transform(keyed.map(t => InternalRow(t._2.getInt(0), t._1)).toArray)
+ .asInstanceOf[LongPointIndex]
+
+ assert(ids(boxed.upTo(3)).size == 300)
+ assert(ids(boxed.upToExclusive(3)).size == 200)
+ assert(ids(boxed.from(3)).size == 300)
+ assert(ids(boxed.fromExclusive(3)).size == 200)
+ assert(ids(boxed.upToExclusive(1)).isEmpty)
+ assert(ids(boxed.fromExclusive(5)).isEmpty)
+
+ assert(ids(longPoints.upTo(3L)).size == 300)
+ assert(ids(longPoints.upToExclusive(3L)).size == 200)
+ assert(ids(longPoints.from(3L)).size == 300)
+ assert(ids(longPoints.fromExclusive(3L)).size == 200)
+ }
+
+ test("PointIndex.build defends against a build-side iterator reusing its row
buffer") {
+ // Same hazard as the IntervalIndex case above: a point index built
straight off a
+ // reused row buffer must copy both the key and the row, not just
reference them.
+ val n = 50
+ val buildKeys = Seq(BoundReference(0, IntegerType, nullable = true))
+ val (shared, rows) = reusingRowIterator(fields = 2, n)
+ val index = PointIndex.build(IntegerType, buildKeys, rows, Some(n))
+
+ shared.update(0, -1)
+ shared.update(1, -1)
+
+ (0 until n).foreach { i =>
+ val payload = index.from(i).next().getInt(1)
+ assert(payload == i, s"point $i was corrupted by later reuse of the row
buffer")
+ }
+ }
+
+ test("PointIndex.build defends against string keys sharing a reused byte
buffer") {
+ val n = 50
+ val buildKeys = Seq(BoundReference(0, StringType, nullable = true))
+ val (sharedBytes, rows) = reusingByteBufferStringIterator(n)
+ val index = PointIndex.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.from(key).take(1).toList
+ assert(matching.nonEmpty && matching.head.getInt(1) == i,
+ s"point $i was corrupted by shared buffer reuse")
+ }
+ }
+
+ test("point index orders NaN above +Infinity and -Infinity below every
finite value") {
+ // Total order: -Inf < -1.0 < 0.0 < 1.0 < +Inf < NaN, with NaN the unique
maximum.
+ val keyed = Array[(Any, InternalRow)](
+ (Double.NaN, row(90)),
+ (Double.PositiveInfinity, row(80)),
+ (1.0, row(1)),
+ (Double.NegativeInfinity, row(-80)),
+ (-1.0, row(-1)),
+ (0.0, row(0)))
+ val index = buildPoints(DoubleType, keyed.toSeq)
+ assert(ids(index.upTo(0.0)) == List(-80, -1, 0))
+ assert(ids(index.from(1.0)) == List(1, 80, 90))
+ // NaN is the unique maximum: only a NaN key is at or above it.
+ assert(ids(index.from(Double.NaN)) == List(90))
+ assert(ids(index.upTo(Double.NaN)) == List(-80, -1, 0, 1, 80, 90))
+ // +Infinity is below NaN but above every finite value.
+ assert(ids(index.upTo(Double.PositiveInfinity)) == List(-80, -1, 0, 1, 80))
+ assert(ids(index.from(Double.NegativeInfinity)) == List(-80, -1, 0, 1, 80,
90))
+ }
+
+ test("upTo/from matches a brute-force reference across randomized point
sets") {
+ // PointIndex's counterpart to "overlapping matches a brute-force reference
+ // across randomized interval sets": cross-checks the binary-search bound
+ // scan against an independent O(n) scan, across many random key sets and
+ // probe values, including duplicate keys.
+ val rand = new scala.util.Random(42)
+ val bound = 50
+ for (_ <- 0 until 30) {
+ val n = 1 + rand.nextInt(300)
+ val raw = (0 until n).map(i => (rand.nextInt(bound), i))
+ val index = buildPoints(IntegerType, raw.map { case (k, i) => (k,
row(i)) })
+
+ for (_ <- 0 until 20) {
+ val probe = rand.nextInt(bound)
+ val expectedUpTo = raw.collect { case (k, i) if k <= probe => i }.toSet
+ val expectedFrom = raw.collect { case (k, i) if k >= probe => i }.toSet
+ assert(ids(index.upTo(probe)).toSet == expectedUpTo,
+ s"n=$n probe=$probe expected=$expectedUpTo")
+ assert(ids(index.from(probe)).toSet == expectedFrom,
+ s"n=$n probe=$probe expected=$expectedFrom")
+ }
+ }
+ }
+
+ test("indexes round-trip through Java and strict Kryo serialization") {
+ val serializers = Seq(
+ "Java" -> new org.apache.spark.serializer.JavaSerializer(
+ new org.apache.spark.SparkConf()).newInstance(),
+ "Kryo" -> {
+ val conf = new org.apache.spark.SparkConf()
+
.set(org.apache.spark.internal.config.Kryo.KRYO_REGISTRATION_REQUIRED, true)
+ new org.apache.spark.serializer.KryoSerializer(conf).newInstance()
+ })
+
+ // (key type, lo, hi, probe) covering each supported key family: fixed
+ // width, floating point, decimal, binary, and collated strings. The
+ // collated probe mixes cases, so a restored index only answers it when
+ // its collation ordering survived the round trip. The key type itself
+ // never reaches a serializer: writes cross it as JSON, so this exercises
+ // the index registrations, not any DataType registration.
+ val cases: Seq[(DataType, Any, Any, Any)] = Seq(
+ (ByteType, 0.toByte, 10.toByte, 5.toByte),
+ (ShortType, 0.toShort, 10.toShort, 5.toShort),
+ (IntegerType, 0, 10, 5),
+ (LongType, 0L, 10L, 5L),
+ (FloatType, 0.0f, 10.0f, 5.0f),
+ (DoubleType, 0.0, 10.0, 5.0),
+ (DecimalType(10, 2), Decimal(0.0), Decimal(10.0), Decimal(5.0)),
+ (StringType("UTF8_LCASE"),
+ UTF8String.fromString("abc"), UTF8String.fromString("zzz"),
UTF8String.fromString("MMM")),
+ (CharType(3),
+ UTF8String.fromString("abc"), UTF8String.fromString("zzz"),
UTF8String.fromString("mmm")),
+ (VarcharType(10),
+ UTF8String.fromString("abc"), UTF8String.fromString("zzz"),
UTF8String.fromString("mmm")),
+ (BinaryType, Array[Byte](0), Array[Byte](10), Array[Byte](5)),
+ (DateType, 100, 200, 150),
+ (TimestampType, 1000L, 2000L, 1500L),
+ (TimestampNTZType, 1000L, 2000L, 1500L),
+ (BooleanType, false, true, true))
+
+ def roundTrip[A <: AnyRef: scala.reflect.ClassTag](
+ ser: org.apache.spark.serializer.SerializerInstance, index: A): A =
+ ser.deserialize[A](ser.serialize(index))
+
+ // One reader for the sealed family, mirroring
BroadcastRangeJoinExec.makeCandidateFn dispatch:
+ // the interval kinds answer a window, the point kinds answer a
+ // lower-bound scan, and every relation this test builds answers at
+ // row ordinal 0.
+ def answer(rel: RangeRelation, low: Any, high: Any): List[Int] = rel match
{
+ case r: LongIntervalIndex => ids(r.overlapping(r.toLong(low),
r.toLong(high)))
+ case r: IntervalIndex => ids(r.overlapping(low, high))
+ case r: LongPointIndex => ids(r.from(r.toLong(low)))
+ case r: PointIndex => ids(r.from(low))
+ }
+
+ cases.foreach { case (dt, lo, hi, probe) =>
+ // One straight span and one inverted span that normalizes to the same
+ // shape, plus the family the broadcast build routes the type to. Rows
+ // carry the asserted id at ordinal 0 and the keys at 1 and 2, so one
+ // reader answers every family.
+ val interval = buildIntervals(dt, Seq((lo, hi, row(5)), (hi, lo,
row(6))))
+ val point = buildPoints(dt, Seq((probe, row(7))))
+ val routedInterval = RangeRelationBroadcastMode(
+ Seq(BoundReference(1, dt, nullable = true), BoundReference(2, dt,
nullable = true)),
+ IntervalIndexKind)
+ .transform(Array(InternalRow(5, lo, hi), InternalRow(6, hi, lo)))
+ val routedPoint = RangeRelationBroadcastMode(
+ Seq(BoundReference(1, dt, nullable = true)), PointIndexKind)
+ .transform(Array(InternalRow(7, probe)))
+
+ val intervalWindow = answer(interval, probe, probe)
+ val pointScan = answer(point, probe, probe)
+ assert(intervalWindow.sorted == List(5, 6), s"$dt window")
+ assert(pointScan.sorted == List(7), s"$dt bound")
+ // The routed family answers the same query on the same rows.
+ assert(answer(routedInterval, probe, probe).sorted ==
intervalWindow.sorted, s"$dt routed")
+ assert(answer(routedPoint, probe, probe).sorted == pointScan.sorted,
s"$dt routed point")
+
+ serializers.foreach { case (name, ser) =>
+ Seq(interval, routedInterval).foreach { rel =>
+ assert(answer(roundTrip(ser, rel), probe, probe).sorted ==
intervalWindow.sorted,
+ s"$dt interval via $name")
+ }
+ Seq(point, routedPoint).foreach { rel =>
+ assert(answer(roundTrip(ser, rel), probe, probe).sorted ==
pointScan.sorted,
+ s"$dt point via $name")
+ }
+ }
+ }
+ }
+
+ test("restored interval indexes answer windows that need the walk") {
+ // The round-trip fixtures above place every probe inside a common core, so
+ // the restored index answers from the row slice without reading the tree
+ // state. These intervals are disjoint: each probe must walk maxima, lows,
+ // and highs on the restored index, under every serializer.
+ val serializers = Seq(
+ "Java" -> new org.apache.spark.serializer.JavaSerializer(
+ new org.apache.spark.SparkConf()).newInstance(),
+ "Kryo" -> {
+ val conf = new org.apache.spark.SparkConf()
+
.set(org.apache.spark.internal.config.Kryo.KRYO_REGISTRATION_REQUIRED, true)
+ new org.apache.spark.serializer.KryoSerializer(conf).newInstance()
+ },
+ // Reference tracking off duplicates any aliased object on the wire, so
it
+ // also probes a restored graph that re-materialized every array entry.
+ "Kryo without references" -> {
+ val conf = new org.apache.spark.SparkConf()
+
.set(org.apache.spark.internal.config.Kryo.KRYO_REGISTRATION_REQUIRED, true)
+ .set(org.apache.spark.internal.config.Kryo.KRYO_REFERENCE_TRACKING,
false)
+ new org.apache.spark.serializer.KryoSerializer(conf).newInstance()
+ })
+
+ def roundTrip[A <: AnyRef: scala.reflect.ClassTag](
+ ser: org.apache.spark.serializer.SerializerInstance, index: A): A =
+ ser.deserialize[A](ser.serialize(index))
+
+ val windows = Seq((5.0, 5.0, List(0)), (25.0, 25.0, List(1)), (45.0, 45.0,
List(2)),
+ (5.0, 25.0, List(0, 1)), (25.0, 45.0, List(1, 2)))
+ val spans = Seq(span(0.0, 10.0, 0), span(20.0, 30.0, 1), span(40.0, 50.0,
2))
+ val boxed = buildIntervals(DoubleType, spans)
+ val longIdx = RangeRelationBroadcastMode(
+ Seq(BoundReference(1, IntegerType, nullable = true),
+ BoundReference(2, IntegerType, nullable = true)),
+ IntervalIndexKind)
+ .transform(Array(InternalRow(0, 0, 10), InternalRow(1, 20, 30),
InternalRow(2, 40, 50)))
+ .asInstanceOf[LongIntervalIndex]
+
+ windows.foreach { case (low, high, expected) =>
+ assert(probeIds(boxed, low, high).sorted == expected, s"fresh boxed
$low-$high")
+ assert(longIdx.overlapping(low.toLong,
high.toLong).map(_.getInt(0)).toList.sorted ==
+ expected, s"fresh long $low-$high")
+ }
+
+ serializers.foreach { case (name, ser) =>
+ val restoredBoxed = roundTrip(ser, boxed)
+ val restoredLong = roundTrip(ser, longIdx)
+ windows.foreach { case (low, high, expected) =>
+ // A disjoint build has no all-covering shortcut for these windows, so
+ // the restored index must answer through a fresh walk iterator.
+ assert(restoredBoxed.overlapping(low,
high).isInstanceOf[OverlapIterator],
+ s"boxed $low-$high did not walk after $name")
+ assert(restoredLong.overlapping(low.toLong, high.toLong)
+ .isInstanceOf[LongOverlapIterator], s"long $low-$high did not walk
after $name")
+ assert(probeIds(restoredBoxed, low, high).sorted == expected,
+ s"boxed $low-$high via $name")
+ assert(restoredLong.overlapping(low.toLong, high.toLong)
+ .map(_.getInt(0)).toList.sorted == expected, s"long $low-$high via
$name")
+ }
+ }
+ }
+
+ test("overlapping uses the common-core short cut without changing the
answer") {
+ // [0, 10], [5, 15], and [8, 20] share the core [8, 10] (1-D Helly). Any
+ // window that reaches it overlaps every interval, so the index returns
+ // all rows directly; windows outside keep the walk's selective answer.
+ val index = buildIntervals(Seq(span(0, 10, 0), span(5, 15, 1), span(8, 20,
2)))
+ assert(probeIds(index, 9, 9).sorted == List(0, 1, 2))
+ assert(probeIds(index, 8, 10).sorted == List(0, 1, 2))
+ assert(probeIds(index, 6, 12).sorted == List(0, 1, 2))
+ // Endpoints are normalized first, so a reversed window reaches the core
too.
+ assert(probeIds(index, 12, 6).sorted == List(0, 1, 2))
+ // The short cut skips the walk's iterator for a core window, and only
there.
+ assert(!index.overlapping(9, 9).isInstanceOf[OverlapIterator])
+ assert(probeIds(index, 0, 2) == List(0))
+ assert(probeIds(index, 6, 7).sorted == List(0, 1))
+ assert(probeIds(index, 11, 12).sorted == List(1, 2))
+ assert(probeIds(index, 16, 30) == List(2))
+ }
+
+ test("overlapping short circuits all-covering windows over builds without a
common core") {
+ // Disjoint intervals and distinct points share no core, but a window
+ // spanning max-low and min-high still covers every interval, so the
+ // direct all-rows scan answers it and no walk iterator is built. A
+ // window covering some but not all intervals keeps the walk.
+ val disjoint = buildIntervals(Seq(span(0, 10, 0), span(20, 30, 1)))
+ assert(probeIds(disjoint, -5, 35).sorted == List(0, 1))
+ assert(!disjoint.overlapping(-5, 35).isInstanceOf[OverlapIterator])
+ assert(probeIds(disjoint, 5, 25).sorted == List(0, 1))
+ assert(!disjoint.overlapping(5, 25).isInstanceOf[OverlapIterator])
+ assert(probeIds(disjoint, 5, 15) == List(0))
+ assert(disjoint.overlapping(5, 15).isInstanceOf[OverlapIterator])
+
+ val points = buildIntervals(Seq(span(1, 1, 1), span(5, 5, 5)))
+ assert(probeIds(points, 0, 6).sorted == List(1, 5))
+ assert(!points.overlapping(0, 6).isInstanceOf[OverlapIterator])
+ assert(probeIds(points, 4, 5) == List(5))
+ assert(points.overlapping(4, 5).isInstanceOf[OverlapIterator])
+
+ val longIdx = RangeRelationBroadcastMode(
+ Seq(BoundReference(0, IntegerType, nullable = true),
+ BoundReference(1, IntegerType, nullable = true)),
+ IntervalIndexKind)
+ .transform(Array(InternalRow(0, 10), InternalRow(20, 30)))
+ .asInstanceOf[LongIntervalIndex]
+ assert(longIdx.overlapping(-5L, 35L).map(_.getInt(0)).toSet == Set(0, 20))
+ assert(!longIdx.overlapping(-5L, 35L).isInstanceOf[LongOverlapIterator])
+ assert(longIdx.overlapping(5L, 15L).map(_.getInt(0)).toList == List(0))
+ assert(longIdx.overlapping(5L, 15L).isInstanceOf[LongOverlapIterator])
+ }
+
+ test("overlap iterator sizes its scratch stack to tree depth, not row
count") {
+ // The scratch stack is pre-sized off `lows.length` alone (tree depth),
independent of
+ // data distribution. Bounding against that same depth formula (rather
than a fixed
+ // magic number) still catches a regression to O(row count) sizing,
without breaking
+ // if the depth-to-capacity multiplier or the floor constant is tuned to a
different
+ // (but still O(log n)) value.
+ def depthOf(n: Int): Int = 32 - java.lang.Integer.numberOfLeadingZeros(n)
+
+ val oneRowIndex = buildIntervals(Seq(span(0, 10, 0)))
+ val oneRowIter = oneRowIndex.overlapping(15,
15).asInstanceOf[OverlapIterator]
+ assert(oneRowIter.stackCapacity <= math.max(8, depthOf(1) * 4),
+ s"a one-row index should need only a handful of scratch slots, " +
+ s"got ${oneRowIter.stackCapacity}")
+
+ val n = 10000
+ val largeIndex = buildIntervals((0 until n).map(i => span(i, i + 100, i)))
+ val largeIter = largeIndex.overlapping(50,
50).asInstanceOf[OverlapIterator]
+ // log2(10000) ~= 13.3: a handful of tree levels, nowhere near n.
+ assert(largeIter.stackCapacity <= depthOf(n) * 4,
+ s"a $n-row index should size its scratch stack to tree depth
(~${depthOf(n)}), " +
+ s"not row count; got ${largeIter.stackCapacity}")
+ }
+
+ test("overlapping traversal is lazy and supports early exit under high
overlap") {
+ // The pending-work stack tracks tree depth, not match count: draining a
handful
+ // of candidates out of ~100 matches among n=10,000 intervals must not
grow it
+ // toward n. Staggered spans also avoid the common-core short cut, which
would
+ // skip the walk entirely. Reading `pendingSpans` confirms work remains on
the
+ // stack after taking initial matches, proving the walk did not drain
eagerly.
+ val n = 10000
+ val overlappingSpans = (0 until n).map(i => span(i, i + 100, i))
+ val index = buildIntervals(overlappingSpans)
+ val iter = index.overlapping(5000, 5000).asInstanceOf[OverlapIterator]
+ val initialCapacity = iter.stackCapacity
+
+ assert(iter.take(3).size == 3)
+ assert(iter.pendingSpans > 0,
+ s"lazy traversal must leave work pending on the stack, but had
${iter.pendingSpans} spans")
+ assert(iter.stackCapacity == initialCapacity,
+ s"stack grew from $initialCapacity to ${iter.stackCapacity} after taking
3 of $n matches")
+ }
+
+ test("overlapping matches a brute-force reference across randomized interval
sets") {
+ // Cross-checks the tree-pruning search in overlapping() against an
independent
+ // O(n) scan, across many random interval sets and probe windows.
Hand-written
+ // cases above cover specific shapes; this covers combinations they would
miss.
+ val rand = new scala.util.Random(42)
+ val bound = 50
+ for (_ <- 0 until 30) {
+ val n = 1 + rand.nextInt(300)
+ // Endpoints are not pre-sorted, exercising build()'s lo/hi-swap
normalization.
+ val raw = (0 until n).map(i => (rand.nextInt(bound),
rand.nextInt(bound), i))
+ val index = buildIntervals(raw.map { case (a, b, i) => span(a, b, i) })
+ val normalized = raw.map { case (a, b, i) => (math.min(a, b),
math.max(a, b), i) }
+
+ for (_ <- 0 until 20) {
+ val a = rand.nextInt(bound)
+ val b = rand.nextInt(bound)
+ val (low, high) = (math.min(a, b), math.max(a, b))
+ val expected = normalized.collect {
+ case (lo, hi, i) if lo <= high && low <= hi => i
+ }.toSet
+ assert(probeIds(index, low, high).toSet == expected,
+ s"n=$n low=$low high=$high expected=$expected")
+ }
+ }
+ }
+
+ test("broadcast modes route keys that map to one long to the long indexes") {
+ // Both build keys are long-convertible: the long interval index.
+ assert(RangeRelationBroadcastMode(
+ Seq(BoundReference(0, IntegerType, nullable = true),
+ BoundReference(1, IntegerType, nullable = true)),
+ IntervalIndexKind)
+ .transform(Array(InternalRow(0, 10, 1))).isInstanceOf[LongIntervalIndex])
+ assert(RangeRelationBroadcastMode(
+ Seq(BoundReference(0, IntegerType, nullable = true)), PointIndexKind)
+ .transform(Array(InternalRow(5, 1))).isInstanceOf[LongPointIndex])
+ // Floating point has no bit-preserving SQL order, and strings are not one
long.
+ // Pinned for both kinds and both widths: signed zero and NaN need the
boxed
+ // ordering's special-cased comparator, not a bit-preserving long encoding.
+ assert(RangeRelationBroadcastMode(
+ Seq(BoundReference(0, FloatType, nullable = true)), PointIndexKind)
+ .transform(Array(InternalRow(1.0f, 1))).isInstanceOf[PointIndex])
+ assert(RangeRelationBroadcastMode(
+ Seq(BoundReference(0, DoubleType, nullable = true)), PointIndexKind)
+ .transform(Array(InternalRow(1.0, 1))).isInstanceOf[PointIndex])
+ assert(RangeRelationBroadcastMode(
+ Seq(BoundReference(0, FloatType, nullable = true),
+ BoundReference(1, FloatType, nullable = true)),
+ IntervalIndexKind)
+ .transform(Array(InternalRow(1.0f, 2.0f,
1))).isInstanceOf[IntervalIndex])
+ assert(RangeRelationBroadcastMode(
+ Seq(BoundReference(0, DoubleType, nullable = true),
+ BoundReference(1, DoubleType, nullable = true)),
+ IntervalIndexKind)
+ .transform(Array(InternalRow(1.0, 2.0, 1))).isInstanceOf[IntervalIndex])
+ assert(RangeRelationBroadcastMode(
+ Seq(BoundReference(0, StringType, nullable = true),
+ BoundReference(1, StringType, nullable = true)),
+ IntervalIndexKind)
+ .transform(Array(InternalRow(
+ UTF8String.fromString("a"), UTF8String.fromString("b"),
1))).isInstanceOf[IntervalIndex])
+ // A compact decimal (precision <= 18) is representable as a scaled long.
+ assert(RangeRelationBroadcastMode(
+ Seq(BoundReference(0, DecimalType(10, 2), nullable = true)),
PointIndexKind)
+ .transform(Array(InternalRow(Decimal(1.0),
1))).isInstanceOf[LongPointIndex])
+ assert(RangeRelationBroadcastMode(
+ Seq(BoundReference(0, DecimalType(18, 2), nullable = true)),
PointIndexKind)
+ .transform(Array(InternalRow(Decimal(1.0),
1))).isInstanceOf[LongPointIndex])
+ // Precision above 18 cannot be read back as an unscaled long: stays boxed.
+ assert(RangeRelationBroadcastMode(
+ Seq(BoundReference(0, DecimalType(19, 2), nullable = true)),
PointIndexKind)
+ .transform(Array(InternalRow(Decimal(1.0), 1))).isInstanceOf[PointIndex])
+ }
+
+ test("the long interval index answers the same windows as the generic one") {
+ def longIndex(spans: Seq[(Any, Any, InternalRow)]): LongIntervalIndex =
+ RangeRelationBroadcastMode(
+ Seq(BoundReference(0, IntegerType, nullable = true),
+ BoundReference(1, IntegerType, nullable = true)),
+ IntervalIndexKind)
+ .transform(spans.map(t => InternalRow(t._1, t._2,
t._3.getInt(0))).toArray)
+ .asInstanceOf[LongIntervalIndex]
+ def assertSameWindow(
+ generic: IntervalIndex,
+ longIdx: LongIntervalIndex,
+ lo: Int,
+ hi: Int): Unit = {
+ val expected = probeIds(generic, lo, hi).sorted
+ val actual = longIdx.overlapping(lo.toLong,
hi.toLong).map(_.getInt(2)).toList.sorted
+ assert(actual == expected, s"window ($lo, $hi): expected $expected, got
$actual")
+ }
+
+ // Includes an inverted span and a null bound, and shares the core [8, 10].
+ val spans = Seq(span(0, 10, 0), span(5, 15, 1), (null, 4, row(9)),
+ span(8, 20, 2), span(15, 5, 3))
+ val generic = buildIntervals(spans)
+ val longIdx = longIndex(spans)
+ assert(longIdx.estimatedSize > 0)
+ // A window that meets the core skips the walk and still answers
everything.
+ assert(!longIdx.overlapping(9L, 9L).isInstanceOf[LongOverlapIterator])
+ assert(longIdx.overlapping(9L, 9L).map(_.getInt(2)).toList.sorted ==
List(0, 1, 2, 3))
+ Seq((2, 2), (5, 7), (9, 9), (13, 12), (11, 11), (15, 15), (20, 30))
+ .foreach { case (lo, hi) => assertSameWindow(generic, longIdx, lo, hi) }
+
+ // Randomized sets cross-check the families independently: the generic
index
+ // is itself brute-force verified, so this pins the walk, the lo/hi swap,
+ // and null drops without recomputing expectations by hand.
+ val rand = new scala.util.Random(42)
+ val bound = 50
+ for (_ <- 0 until 20) {
+ val n = 1 + rand.nextInt(300)
+ val raw = (0 until n).map(i => (rand.nextInt(bound),
rand.nextInt(bound), i))
+ val randomSpans = raw.map { case (a, b, i) => span(a, b, i) }
+ val randomGeneric = buildIntervals(randomSpans)
+ val randomLong = longIndex(randomSpans)
+ for (_ <- 0 until 15) {
+ val a = rand.nextInt(bound)
+ val b = rand.nextInt(bound)
+ assertSameWindow(randomGeneric, randomLong, math.min(a, b),
math.max(a, b))
+ }
+ }
+ }
+
+ test("LongIntervalIndex.coreLow/coreHigh use the empty-set identities, like
the boxed family") {
Review Comment:
**Nit (P3):** The name says the long family's empty-set identities are "like
the boxed family", and the comment says the two families "should agree on that
convention". The same comment notes that `IntervalIndex.coreLow`/`coreHigh`
return null for an empty index, while this family returns
`Long.MinValue`/`Long.MaxValue`, so there is no shared convention. Both values
are also unobservable, because `overlapping` returns early on an empty index.
Could the name and comment describe only the long family's identities?
This extends an [existing
discussion](https://github.com/apache/spark/pull/59027#discussion_r4346073108).
**Verification:**
- **Inspection:** The long-family empty-identity test no longer claims
parity with the boxed family, and its description matches the null and
Long.MinValue/Long.MaxValue values the two families return for an empty index.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/joins/RangeIndexSuite.scala:
##########
@@ -0,0 +1,1066 @@
+/*
+ * 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 -- exactly what a
broadcast's row
Review Comment:
**Nit (P3):** `BroadcastExchangeExec` builds the relation from
`child.executeCollectIterator()`, and `SparkPlan.executeCollectIterator`
decodes every record through `decodeUnsafeRows` into a new `byte[]` and a new
`UnsafeRow`. So the broadcast's row iterator never reuses or later mutates a
buffer. The hazard the builders actually defend against is a projected key
value that points into a reused buffer, such as a typed ScalaUDF result; the
row copy is defensive. As written, the comment could lead someone to rely on
broadcast-iterator behavior that doesn't exist, or to drop the key copy
thinking only rows are at risk. Could it describe the copies as guarding
against any producer that reuses its row or key buffers?
**Verification:**
- **Inspection:** The reused-row-buffer test comment describes buffer reuse
as a defended producer behavior and makes no claim about the broadcast
exchange's row iterator that the collect path contradicts.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/joins/BroadcastRangeJoinExec.scala:
##########
@@ -0,0 +1,592 @@
+/*
+ * 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.RangeJoinKind
+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],
+ rangeJoinKind: RangeJoinKind)
+ extends JoinCodegenSupport {
+
+ if (condition.isEmpty) {
+ throw SparkException.internalError("BroadcastRangeJoinExec requires a join
condition.")
+ }
+
+ override lazy val metrics = Map(
+ "numOutputRows" -> SQLMetrics.createMetric(sparkContext, "number of output
rows"))
Review Comment:
**Nit (P3):** `numOutputRows` is incremented in four places: the interpreted
result map in `doExecute`, `codegenInner`, `codegenOuter` inside the
`shouldOutputRow` guard, and `codegenLeftExistence`. `SQLMetricsSuite` pins
"number of output rows" for `SortMergeJoin`, `SortMergeAsOfJoin`,
`BroadcastHashJoin`, `ShuffledHashJoin`, and `BroadcastNestedLoopJoin`, but
nothing reads it for the range join, and the range join suites assert only
answers and plan shape. If `codegenOuter` counted both a matched row and its
padded row, or the interpreted path lost its increment, the SQL UI would show
the wrong count and every test would still pass. Could you add a range join
case to `SQLMetricsSuite` next to the BNLJ one?
**Verification:**
- **Regression:** With whole-stage codegen enabled and disabled, inner,
outer (including null-padded rows), left semi, and left anti broadcast range
joins report a 'number of output rows' metric equal to the number of rows they
return, and the assertion fails if any path skips or double-counts an output
row.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/joins/RangeJoinSQLSuite.scala:
##########
@@ -0,0 +1,771 @@
+/*
+ * 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(
+ 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",
+ "SELECT p.d, b.lo FROM dbl_partial_bounds b JOIN dbl_partial_points p ON
p.d > b.lo",
+ "SELECT p.s, b.lo FROM str_partial_points p JOIN str_partial_bounds b ON
p.s < b.lo",
+ "SELECT p.s, b.lo FROM str_partial_bounds b JOIN str_partial_points p ON
p.s > b.lo")
+
+ private def setupIpLookupViews(): Unit = {
+ sql(
+ """
+ |CREATE OR REPLACE TEMP VIEW signin_ip(id, ip_add_int) AS VALUES
+ | ('1', 1467630666L),
+ | ('2', 1186929717L),
+ | ('3', 3501544352L),
+ | ('31', 3501544320L),
+ | ('32', 3501544383L),
+ | ('33', 3501544330L),
+ | ('34', 3501544390L),
+ | ('4', 1626243295L),
+ | ('41', 1626243072L),
+ | ('42', 1626243327L),
+ | ('5', 1206355995L),
+ | ('5', 1206355995L),
+ | ('6', CAST(NULL AS BIGINT)),
+ | ('7', 0L)
+ |""".stripMargin)
+ sql(
+ """
+ |CREATE OR REPLACE TEMP VIEW ip_lookup(
+ | BEGIN_IP2LONG, END_IP2LONG, TWO_CHAR_CNTRY_CD, zip_id) AS VALUES
+ | (1467630592L, 1467630847L, 'de', 2380520L),
+ | (1186929664L, 1186930175L, 'us', 6723936L),
+ | (1186929664L, 1186930175L, 'us', 6723936L),
+ | (3501544320L, 3501544383L, 'ca', 12804631L),
+ | (3501544330L, 3501544390L, 'ca', 12804631001L),
+ | (1626243072L, 1626243327L, 'us', 8849753L),
+ | (1206355968L, 1206356223L, 'us', 10138801L),
+ | (CAST(NULL AS BIGINT), 1206356223L, 'us', 10138801L),
+ | (1206355968L, CAST(NULL AS BIGINT), 'us', 10138801L),
+ | (CAST(NULL AS BIGINT), CAST(NULL AS BIGINT), 'us', 10138801L)
+ |""".stripMargin)
+ }
+
+ private def setupIntervalViews(): Unit = {
+ sql(
+ """
+ |CREATE OR REPLACE TEMP VIEW intervals_a(lo, hi) AS VALUES
+ | (-1, 0), (0, 1), (0, 2), (1, 5)
+ |""".stripMargin)
+ sql(
+ """
+ |CREATE OR REPLACE TEMP VIEW intervals_b(lo, hi) AS VALUES
+ | (-2, -1), (-4, -2), (1, 3), (5, 7)
+ |""".stripMargin)
+ }
+
+ /**
+ * Two-column DOUBLE ranges, so a two-key interval-overlap query here builds
a boxed
+ * `IntervalIndex` (DoubleType has no bit-preserving long encoding) with
both bounds
+ * supplied by the streamed side - the `case intervals: IntervalIndex`
branch of
+ * `BroadcastRangeJoinExec.makeCandidateFn` that `intervals_a`/`intervals_b`
(INT,
+ * long-specialized) never reaches.
+ */
+ private def setupDoubleIntervalViews(): Unit = {
+ sql(
+ """
+ |CREATE OR REPLACE TEMP VIEW dbl_ranges_a(lo, hi) AS VALUES
+ | (-1.0, 0.0), (0.0, 1.0), (0.0, 2.0), (1.0, 5.0)
+ |""".stripMargin)
+ sql(
+ """
+ |CREATE OR REPLACE TEMP VIEW dbl_ranges_b(lo, hi) AS VALUES
+ | (-2.0, -1.0), (-4.0, -2.0), (1.0, 3.0), (5.0, 7.0)
+ |""".stripMargin)
+ }
+
+ private def setupDecimalViews(): Unit = {
+ sql(
+ """
+ |CREATE OR REPLACE TEMP VIEW dec_ranges(lo, hi) AS VALUES
+ | (CAST(0.00 AS DECIMAL(10, 2)), CAST(2.50 AS DECIMAL(10, 2))),
+ | (CAST(3.00 AS DECIMAL(10, 2)), CAST(5.00 AS DECIMAL(10, 2)))
+ |""".stripMargin)
+ sql(
+ """
+ |CREATE OR REPLACE TEMP VIEW dec_points(p) AS VALUES
+ | (CAST(1.25 AS DECIMAL(10, 2))),
+ | (CAST(3.00 AS DECIMAL(10, 2))),
+ | (CAST(9.00 AS DECIMAL(10, 2)))
+ |""".stripMargin)
+ }
+
+ private def setupCollatedStrings(): Unit = {
+ // 'M' and 'm' compare equal under UNICODE_CI and differ as bytes.
+ sql(
+ """
+ |CREATE OR REPLACE TEMP VIEW str_ranges(lo, hi) AS VALUES
+ | (CAST('a' AS STRING COLLATE UNICODE_CI), CAST('a' AS STRING COLLATE
UNICODE_CI)),
+ | (CAST('b' AS STRING COLLATE UNICODE_CI), CAST('b' AS STRING COLLATE
UNICODE_CI)),
+ | (CAST('c' AS STRING COLLATE UNICODE_CI), CAST('M' AS STRING COLLATE
UNICODE_CI)),
+ | (CAST('m' AS STRING COLLATE UNICODE_CI), CAST('z' AS STRING COLLATE
UNICODE_CI))
+ |""".stripMargin)
+ sql(
+ """
+ |CREATE OR REPLACE TEMP VIEW str_points(s) AS VALUES
+ | (CAST('M' AS STRING COLLATE UNICODE_CI)),
+ | (CAST('m' AS STRING COLLATE UNICODE_CI))
+ |""".stripMargin)
+ }
+
+ /** Points and bounds for single-inequality joins on boxed key families. */
+ private def setupPartialRangeViews(): Unit = {
+ sql(
+ """
+ |CREATE OR REPLACE TEMP VIEW dbl_partial_points(d) AS VALUES
+ | (0.5), (3.0), (6.0),
+ | (CAST(NULL AS DOUBLE))
+ |""".stripMargin)
+ sql(
+ """
+ |CREATE OR REPLACE TEMP VIEW dbl_partial_bounds(lo) AS VALUES
+ | (1.0), (4.0), (7.5)
+ |""".stripMargin)
+ sql(
+ """
+ |CREATE OR REPLACE TEMP VIEW str_partial_points(s) AS VALUES
+ | ('a'), ('m'), ('z')
+ |""".stripMargin)
+ sql(
+ """
+ |CREATE OR REPLACE TEMP VIEW str_partial_bounds(lo) AS VALUES
+ | ('b'), ('n')
+ |""".stripMargin)
+ }
+
+ private def setupDateAndTimestampViews(): Unit = {
+ sql(
+ """
+ |CREATE OR REPLACE TEMP VIEW dt_ranges(lo, hi) AS VALUES
+ | (DATE'2024-01-01', DATE'2024-01-31'),
+ | (DATE'2024-06-01', DATE'2024-06-30')
+ |""".stripMargin)
+ sql(
+ """
+ |CREATE OR REPLACE TEMP VIEW dt_points(d) AS VALUES
+ | (DATE'2024-01-15'),
+ | (DATE'2024-05-01'),
+ | (DATE'2024-06-15')
+ |""".stripMargin)
+ sql(
+ """
+ |CREATE OR REPLACE TEMP VIEW ts_ranges(lo, hi) AS VALUES
+ | (TIMESTAMP'2024-01-01 00:00:00', TIMESTAMP'2024-01-01 12:00:00')
+ |""".stripMargin)
+ sql(
+ """
+ |CREATE OR REPLACE TEMP VIEW ts_points(t) AS VALUES
+ | (TIMESTAMP'2024-01-01 06:00:00'),
+ | (TIMESTAMP'2024-01-01 18:00:00')
+ |""".stripMargin)
+ }
+
+ test("range join is disabled by default and falls back to nested loop join")
{
+ setupIpLookupViews()
+ val df = sql(s"SELECT s.id, l.zip_id FROM signin_ip s JOIN ip_lookup l ON
$pointCond")
+ assert(findRangeJoin(df.queryExecution.executedPlan).isEmpty)
+ assert(findBroadcastNestedLoopJoin(df.queryExecution.executedPlan).size ==
1)
+ }
+
+ test("point-in-range matches the nested loop result") {
+ setupIpLookupViews()
+ val query = s"SELECT s.id, l.zip_id FROM signin_ip s JOIN ip_lookup l ON
$pointCond"
+ val expected = sql(query).collect()
+ withRangeJoin() {
+ val df = sql(query)
+ assert(findRangeJoin(df.queryExecution.executedPlan).size == 1)
+ checkAnswer(df, expected)
+ }
+ }
+
+ test("BETWEEN is planned as a point-in-range join") {
+ setupIpLookupViews()
+ val query =
+ "SELECT s.id, l.zip_id FROM signin_ip s JOIN ip_lookup l " +
+ "ON s.ip_add_int BETWEEN l.BEGIN_IP2LONG AND l.END_IP2LONG"
+ val expected = sql(query).collect()
+ withRangeJoin(SQLConf.ALWAYS_INLINE_COMMON_EXPR.key -> "false") {
+ val df = sql(query)
+ val planned = findRangeJoin(df.queryExecution.executedPlan)
+ assert(planned.size == 1)
+ // BETWEEN lowers to one point expression. The other side keeps (low,
high).
+ val keyCounts = Set(planned.head.leftKeys.length,
planned.head.rightKeys.length)
+ assert(keyCounts == Set(1, 2))
+ checkAnswer(df, expected)
+ }
+ }
+
+ test("hints select range join, nested loop, or cartesian") {
+ setupIpLookupViews()
+ setupIntervalViews()
+ val expected = sql(
+ s"SELECT s.id, l.zip_id FROM signin_ip s JOIN ip_lookup l ON
$pointCond").collect()
+ withRangeJoin() {
+ val broadcast = sql(
+ "SELECT /*+ BROADCAST(l) */ s.id, l.zip_id FROM signin_ip s JOIN
ip_lookup l " +
+ s"ON $pointCond")
+ val joins = findRangeJoin(broadcast.queryExecution.executedPlan)
+ assert(joins.size == 1 && joins.head.buildSide == BuildRight)
+ checkAnswer(broadcast, expected)
+
+ val dualHintInner = sql(
+ "SELECT /*+ BROADCAST(s), BROADCAST(l) */ s.id, l.zip_id " +
+ s"FROM signin_ip s JOIN ip_lookup l ON $pointCond")
+ val dualJoins = findRangeJoin(dualHintInner.queryExecution.executedPlan)
+ assert(dualJoins.size == 1 && dualJoins.head.buildSide == BuildRight)
+ checkAnswer(dualHintInner, expected)
+
+ val dualHintLeftOuter = sql(
+ "SELECT /*+ BROADCAST(s), BROADCAST(l) */ s.id, l.zip_id " +
+ s"FROM signin_ip s LEFT JOIN ip_lookup l ON $pointCond")
+ val dualOuterJoins =
findRangeJoin(dualHintLeftOuter.queryExecution.executedPlan)
+ assert(dualOuterJoins.size == 1 && dualOuterJoins.head.buildSide ==
BuildRight)
+
+ val unsupported = sql(
+ "SELECT /*+ BROADCAST(a) */ a.lo, b.hi FROM intervals_a a LEFT JOIN
intervals_b b " +
+ "ON a.lo < b.hi")
+ val loops =
findBroadcastNestedLoopJoin(unsupported.queryExecution.executedPlan)
+ assert(findRangeJoin(unsupported.queryExecution.executedPlan).isEmpty)
+ assert(loops.size == 1 && loops.head.buildSide == BuildLeft)
+
+ val cartesian = sql(
+ "SELECT /*+ SHUFFLE_REPLICATE_NL(s) */ s.id, l.zip_id " +
+ s"FROM signin_ip s JOIN ip_lookup l ON $pointCond")
+ val plan = cartesian.queryExecution.executedPlan
+ assert(findRangeJoin(plan).isEmpty)
+ assert(findCartesianProduct(plan).size == 1)
+ }
+ }
+
+ test("supported join types match nested loop; full outer does not") {
+ setupIpLookupViews()
+ Seq(
+ 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"
+ ).foreach { query =>
+ val expected = sql(query).collect()
+ withRangeJoin() {
+ val df = sql(query)
+ assert(findRangeJoin(df.queryExecution.executedPlan).size == 1, query)
+ checkAnswer(df, expected)
+ }
+ }
+ val fullOuter =
+ s"SELECT s.id, l.zip_id FROM signin_ip s FULL OUTER JOIN ip_lookup l ON
$pointCond"
+ withRangeJoin() {
+ assert(findRangeJoin(sql(fullOuter).queryExecution.executedPlan).isEmpty)
+ }
+ }
+
+ test("partial range left outer broadcasts the right side") {
+ setupIntervalViews()
+ // No b.lo is strictly greater than a.hi (5), so the left row (1, 5) is
kept with nulls.
+ val query =
+ "SELECT a.lo, a.hi, b.lo, b.hi FROM intervals_a a LEFT JOIN intervals_b
b " +
+ "ON a.hi < b.lo"
+ val expected = sql(query).collect()
+ assert(expected.exists(_.anyNull))
+ withRangeJoin() {
+ val df = sql(query)
+ val planned = findRangeJoin(df.queryExecution.executedPlan)
+ assert(planned.size == 1)
+ assert(planned.head.joinType == LeftOuter)
+ assert(planned.head.buildSide == BuildRight)
+ assert(planned.head.rangeJoinKind == LessPartialRangeJoin(strict = true))
+ checkAnswer(df, expected)
+ }
+ }
+
+ test("strict inequalities over duplicated build keys match the nested loop
join") {
Review Comment:
**Nit (P3):** The comment says a point index that kept the probe's own key
would surface 100 duplicates, but that only holds when `dup_dates` is the build
side. The four inner joins have no hint, so `getSmallerSide` builds
`probe_dates` (100 distinct keys) and the duplicates end up on the streamed
side. Every case here also asserts only answers. Answers can't tell an
inclusive scan from an exclusive one, because the condition rechecks each
candidate. The candidate-count test in `RangeJoinSuite` builds the right side
only.
So if `PointScan.forJoin` returned `UpTo` or `From` for a strict BuildLeft
join, every test would still pass, and duplicate-heavy BuildLeft joins would go
back to scanning every equal key. Could the count test cover BuildLeft in both
strict directions? And could this test either hint the duplicated side or
describe what it actually runs?
This extends an [existing
discussion](https://github.com/apache/spark/pull/59027#discussion_r4333126458).
**Verification:**
- **Regression:** With whole-stage codegen enabled and disabled, strict
less-than and greater-than partial-range joins that build the left side over
duplicated keys evaluate the condition only on keys that can satisfy the strict
predicate, and the check fails if either build-left scan becomes inclusive.
- **Behavior:** Strict and inclusive partial-range joins over duplicated
build keys still return the nested-loop answer on both build sides.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/joins/RangeIndexSuite.scala:
##########
@@ -0,0 +1,1066 @@
+/*
+ * 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 -- exactly what a
broadcast's row
+ // iterator does once BroadcastExchangeExec finishes feeding it to
RangeRelationBroadcastMode.
+ 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.
+ val n = 300
+ val nested = buildIntervals((0 until n).map(i => span(i, 2 * n - i, i)))
+
+ def expected(probe: Int): Set[Int] =
+ (0 until n).filter(i => i <= probe && probe <= 2 * n - i).toSet
+
+ Seq(0, 1, n / 2, n, 2 * n - 1, 2 * n).foreach { probe =>
+ assert(probeIds(nested, probe, probe).toSet == expected(probe), s"probe
$probe")
+ }
+ }
+
+ test("broadcast of heavily overlapping intervals stays linear") {
+ // [i, i + n] for i in 0 until n. Every interval contains n. The broadcast
+ // is the flat arrays, one slot per interval, so serialization stays
linear.
+ val n = 200
+ val index = buildIntervals((0 until n).map(i => span(i, i + n, i)))
+ val (raw, restored) = javaRoundTrip(index)
+ assert(raw.length < n * 1024, s"serialized ${raw.length} bytes")
+ val probe = n
+ assert(probeIds(restored, probe, probe).toSet == probeIds(index, probe,
probe).toSet)
+ assert(probeIds(restored, probe, probe).size == n)
+ }
+
+ test("estimatedSize depends on the interval count, not the overlap depth") {
+ val n = 200
+ val one = buildIntervals(Seq(span(0, 1, 0)))
+ val disjoint = buildIntervals((0 until n).map(i => span(2 * i, 2 * i + 1,
i)))
+ val overlap = buildIntervals((0 until n).map(i => span(i, i + n, i)))
+ val again = buildIntervals((0 until n).map(i => span(2 * i, 2 * i + 1, i)))
+ assert(disjoint.estimatedSize > one.estimatedSize)
+ assert(overlap.estimatedSize > one.estimatedSize)
+ assert(disjoint.estimatedSize == again.estimatedSize)
+ // Nesting stores no extra copy of the rows. Key objects differ, so the two
+ // indexes need not be equal, but overlap stays the same order of
magnitude.
+ assert(overlap.estimatedSize < disjoint.estimatedSize * 3)
+ }
+
+ test("estimatedSize includes distinct key objects") {
+ val ints = buildIntervals(Seq(span(0, 1, 0)))
+ val wide = UTF8String.fromString("x" * 64)
+ val strings = buildIntervals(StringType, Seq((wide, wide, row(0))))
+ assert(strings.estimatedSize > ints.estimatedSize)
+ }
+
+ test("interval tree treats NaN as its own point and Infinity as a bound") {
+ // PhysicalDoubleType's ordering (SQLOrderingUtil.compareDoubles) totally
orders
+ // -Inf < ... < +Inf < NaN, with NaN comparing equal only to itself. So
[1.0, +Inf]
+ // includes +Inf under inclusive probe, and a NaN key overlaps only a
probe that
+ // also lands exactly on NaN.
+ val index = buildIntervals(DoubleType, Seq(
+ (1.0, Double.PositiveInfinity, row(1)),
+ (Double.NegativeInfinity, -1.0, row(2)),
+ (Double.NaN, Double.NaN, row(3))))
+
+ assert(ids(index.overlapping(5.0, 5.0)) == List(1))
+ assert(ids(index.overlapping(Double.PositiveInfinity,
Double.PositiveInfinity)) == List(1))
+ assert(ids(index.overlapping(-5.0, -5.0)) == List(2))
+ assert(ids(index.overlapping(Double.NegativeInfinity,
Double.NegativeInfinity)) == List(2))
+ assert(ids(index.overlapping(Double.NaN, Double.NaN)) == List(3))
+ assert(ids(index.overlapping(0.0, 0.0)).isEmpty)
+ }
+
+ test("interval index merges keys that compare equal but are not equal") {
+ // Array equals is identity, and UTF8String equals is binary, so these keys
+ // compare equal under the index ordering without being equal. The tree
+ // orders them with that comparator, so a probe on either key returns both.
+ // (Signed zero is the same kind of merge, but Double's == already treats
-0.0
+ // and 0.0 as equal, which fails this helper's != precondition - see the
+ // dedicated "-0.0 and 0.0 are treated as one key" test below instead.)
+ def assertMerged(
+ dataType: DataType,
+ low: Any,
+ leftKey: Any,
+ rightKey: Any,
+ high: Any,
+ earlier: Seq[Any]): Unit = {
+ val ordering = PhysicalDataType.ordering(dataType)
+ assert(leftKey != rightKey)
+ assert(ordering.compare(leftKey, rightKey) == 0)
+ val points = Array(
+ (leftKey, leftKey, row(0)),
+ (rightKey, rightKey, row(1)))
+ val pointIdx = buildIntervals(dataType, points.toSeq)
+ assert(ids(pointIdx.overlapping(leftKey, leftKey)).toSet == Set(0, 1))
+ assert(ids(pointIdx.overlapping(rightKey, rightKey)).toSet == Set(0, 1))
+
+ val intervals = earlier.zipWithIndex.map { case (key, i) =>
+ (key, key, row(10 + i))
+ }.toArray ++ Array(
+ (low, leftKey, row(0)),
+ (rightKey, high, row(1)))
+ val index = buildIntervals(dataType, intervals.toSeq)
+ assert(ids(index.overlapping(leftKey, leftKey)).sorted == List(0, 1))
+ assert(ids(index.overlapping(rightKey, rightKey)).sorted == List(0, 1))
+ }
+
+ val bytes = Array[Byte](1, 2)
+ val sameBytes = Array[Byte](1, 2)
+ assertMerged(
+ BinaryType,
+ Array[Byte](1, 0),
+ bytes,
+ sameBytes,
+ Array[Byte](2),
+ Seq(Array[Byte](0), Array[Byte](0, 1)))
+
+ val upper = UTF8String.fromString("M")
+ val lower = UTF8String.fromString("m")
+ assertMerged(
+ StringType("UNICODE_CI"),
+ UTF8String.fromString("c"),
+ upper,
+ lower,
+ UTF8String.fromString("z"),
+ Seq(UTF8String.fromString("a"), UTF8String.fromString("b")))
+ }
+
+ test("-0.0 and 0.0 are treated as one key") {
+ // Double's own `==` already treats -0.0 and 0.0 as equal, so
`PhysicalDataType.ordering`
+ // (which SQL requires to distinguish them from no other value, but not
from each other)
Review Comment:
**Nit (P3):** "Distinguish them from no other value" reads as if signed zero
compares equal to everything. `SQLOrderingUtil.compareDoubles` keeps -0.0 and
0.0 distinct from every other value and merges only the two of them, which is
what this test asserts. I think the intended wording is "from every other
value, but not from each other".
This extends an [existing
discussion](https://github.com/apache/spark/pull/59027#discussion_r4346073093).
**Verification:**
- **Inspection:** The signed-zero test comment says signed zero is merged
only with itself and kept distinct from every other value, matching SQL
floating-point ordering and the test's assertions.
--
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]
