wangyum commented on code in PR #59027:
URL: https://github.com/apache/spark/pull/59027#discussion_r4436520771
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/planning/patterns.scala:
##########
@@ -421,6 +422,198 @@ object ExtractSingleColumnNullAwareAntiJoin extends
JoinSelectionHelper with Pre
}
}
+/**
+ * Tags the shape of a range-predicate join recognized by
[[ExtractRangeJoinKeys]].
+ * Point-in-range and interval overlap both probe an interval index. A partial
range
+ * is one inequality; Less vs Greater records which side holds the lower bound.
+ */
+sealed abstract class RangeJoinKind
+
+/** A point-in-range join: a value from one side falls within a `[low, high]`
range on the other. */
+case object PointInRangeJoin extends RangeJoinKind
+
+/**
+ * An interval-overlap join: `(low, high)` on each side. The build side is an
+ * interval index and the stream side probes it with `overlapping(low, high)`.
+ */
+case object IntervalOverlapJoin extends RangeJoinKind
+
+/**
+ * A partial-range join: one inequality, one expression per side.
+ * `leftIsLower` says which child holds the lower bound.
+ */
+sealed abstract class PartialRangeJoin extends RangeJoinKind {
+ def leftIsLower: Boolean
+
+ /**
+ * The inequality is `<`/`>` and not `<=`/`>=`. Inclusivity stays on the
+ * original condition, but the strict form also selects the exclusive
+ * point-index bound, so duplicated build keys equal to the probe value
+ * are skipped instead of failing the condition one by one.
+ */
+ def strict: Boolean
+}
+
+/** A `low <[=] high` partial-range join where the left side is the low bound.
*/
+final case class LessPartialRangeJoin(override val strict: Boolean) extends
PartialRangeJoin {
+ override def leftIsLower: Boolean = true
+}
+
+/** A `low <[=] high` partial-range join where the left side is the high
bound. */
+final case class GreaterPartialRangeJoin(override val strict: Boolean) extends
PartialRangeJoin {
+ override def leftIsLower: Boolean = false
+}
+
+/**
+ * Finds joins whose condition is a range predicate usable by a broadcast
range join.
+ *
+ * Tries the two-predicate shapes first (point-in-range, then interval
overlap),
+ * which capture more of the condition, and falls back to the single
+ * cross-side inequality: keeps the first that matches.
+ * - point-in-range: one key probed as both ends of a `(low, high)` window on
the other side.
+ * - interval overlap: a `(low, high)` key on each side.
+ * - partial range: one inequality, one expression per side.
+ * [[RangePredicate]] normalizes `<`/`<=`/`>`/`>=` to `(dataType, low, high,
strict)`;
+ * `strict` selects the exclusive point-index scan for `partialRange` below,
but
Review Comment:
Fixed with your suggested wording: `strict` only selects the exclusive
point-index scan in `partialRange`.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/joins/RangeIndexSuite.scala:
##########
@@ -0,0 +1,1087 @@
+/*
+ * 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 `IntervalIndex.apply`. 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(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(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.apply defends against a build-side iterator reusing its
row buffer") {
+ // If apply() captured keys or rows by reference instead of copying them,
every stored
+ // entry would collapse to the iterator's last value once the shared
buffer is mutated
+ // again after the build-side iterator is drained -- the hazard guarded
against is any
+ // producer whose rows or projected keys alias a reused buffer (e.g. a
typed UDF result),
+ // not specifically BroadcastExchangeExec: its own row iterator decodes a
fresh UnsafeRow
+ // per record.
+ val n = 50
+ val buildKeys = Seq(
+ BoundReference(0, IntegerType, nullable = true),
+ BoundReference(1, IntegerType, nullable = true))
+ val (shared, rows) = reusingRowIterator(fields = 3, n)
+ val index = IntervalIndex(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.apply 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(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(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(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 and contains n, so the
probe
+ // at n hits all 300 rows; the probe at n/2 hits about half (151), since
only
+ // intervals with i <= n/2 reach that far. 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 every 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 spans = (0 until 32).map(i => span(i, i + 10, i))
+ val index = buildIntervals(spans)
+ val longIdx = 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]
+
+ // [0, 15] does not cover every span (the core is [31, 10]), so this walks
the tree
+ // instead of taking the all-covering shortcut: proof that two concurrent
walks can
+ // run on one index, the shape two interval joins fused in one whole-stage
rely on.
+ val expectedOuter = (0 to 15).toList
+ val outer = index.overlapping(0, 15)
+ assert(outer.isInstanceOf[OverlapIterator])
+ val seen = scala.collection.mutable.ArrayBuffer.empty[Int]
+ while (outer.hasNext) {
+ seen += outer.next().getInt(0)
+ assert(ids(index.overlapping(0, 5)).sorted == (0 to 5).toList)
+ }
+ assert(seen.sorted == expectedOuter)
+
+ val longOuter = longIdx.overlapping(0L, 15L)
+ assert(longOuter.isInstanceOf[LongOverlapIterator])
+ val longSeen = scala.collection.mutable.ArrayBuffer.empty[Int]
+ while (longOuter.hasNext) {
+ longSeen += longOuter.next().getInt(2)
+ assert(longIdx.overlapping(0L, 5L).map(_.getInt(2)).toList.sorted == (0
to 5).toList)
+ }
+ assert(longSeen.sorted == expectedOuter)
+ }
+
+ 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.apply 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(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.apply 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(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))))
Review Comment:
Fixed: added `restored point indexes answer multi-entry bound scans`,
covering duplicate and collation-equal keys through all four bound scans, both
point families, and all three serializers.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/joins/RangeRelation.scala:
##########
@@ -0,0 +1,1225 @@
+/*
+ * 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.{Externalizable, ObjectInput, ObjectOutput}
+
+import scala.collection.mutable.ArrayBuffer
+
+import com.esotericsoftware.kryo.{Kryo, KryoSerializable}
+import com.esotericsoftware.kryo.io.{Input, Output}
+
+import org.apache.spark.SparkException
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{BoundReference, Expression}
+import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildSide}
+import org.apache.spark.sql.catalyst.planning.{IntervalOverlapJoin,
PartialRangeJoin, PointInRangeJoin, RangeJoinKind}
+import org.apache.spark.sql.catalyst.plans.physical.BroadcastMode
+import org.apache.spark.sql.catalyst.types.PhysicalDataType
+import org.apache.spark.sql.types.{BooleanType, ByteType, DataType, DateType,
Decimal, DecimalType, IntegerType, LongType, ShortType, TimestampNTZType,
TimestampType}
+import org.apache.spark.util.{KnownSizeEstimation, SizeEstimator, Utils}
+
+/**
+ * Broadcast payload of a range join. An index only names candidate rows.
+ * [[BroadcastRangeJoinExec]] evaluates the original join condition to filter
+ * candidates. Keys that map to one order-preserving long are stored raw and
+ * compared unboxed ([[LongIntervalIndex]], [[LongPointIndex]]); every other
+ * supported type is compared as boxed values ([[IntervalIndex]],
+ * [[PointIndex]]). [[RangeRelationBroadcastMode]] picks the family from the
+ * build keys' data types.
+ *
+ * - The interval indexes answer "which ranges overlap this window?"
+ * - The point indexes answer "which points lie on this side of a bound?"
+ */
+private[execution] sealed trait RangeRelation extends KnownSizeEstimation with
Serializable {
+ def keyType: DataType
+}
+
+/**
+ * Which [[RangeRelation]] a broadcast materializes.
+ *
+ * A point index has no direction, so one broadcast of a column serves both
+ * `<` and `>`. The join records which side holds the lower bound and scans
+ * accordingly.
+ */
+private[execution] sealed abstract class RangeIndexKind {
+ def matches(rangeJoinKind: RangeJoinKind): Boolean =
RangeIndexKind(rangeJoinKind) == this
+}
+
+private[execution] case object IntervalIndexKind extends RangeIndexKind
+
+private[execution] case object PointIndexKind extends RangeIndexKind
+
+private[execution] object RangeIndexKind {
+ def apply(rangeJoinKind: RangeJoinKind): RangeIndexKind = rangeJoinKind
match {
+ case PointInRangeJoin | IntervalOverlapJoin => IntervalIndexKind
+ case _: PartialRangeJoin => PointIndexKind
+ }
+}
+
+/**
+ * The boxed-key index family's shared ordering cache, derived from `keyType`
+ * on first comparison and kept out of the serialized form, so each executor
+ * re-derives its own instance instead of receiving one per broadcast. The
+ * long index family compares raw keys and has no ordering to cache.
+ */
+private[joins] sealed trait BoxedKeyOrdering extends RangeRelation {
+ @transient
+ @volatile private[this] var _ordering: Ordering[Any] = _
+
+ private[joins] def ordering: Ordering[Any] = {
+ var ord = _ordering
+ if (ord == null) {
+ ord = PhysicalDataType.ordering(keyType)
+ _ordering = ord
+ }
+ ord
+ }
+}
+
+/**
+ * The long index family's shared key conversion, derived from `keyType` on
+ * first probe and kept out of the serialized form, so each executor derives
+ * its own converter after deserialization. The interpreted probe path calls
+ * [[toLong]] once per streamed key; generated probe code keeps its own
+ * conversion, emitted at plan time.
+ */
+private[joins] sealed trait LongKeyConversion extends RangeRelation {
+ @transient
+ @volatile private[this] var _toLong: Any => Long = _
+
+ /**
+ * One evaluated streamed key to its order-preserving long. Within a
predicate, the
+ * analyzer's comparison type coercion gives both sides one type; across the
two
+ * predicates an interval-overlap join pairs, `ExtractRangeJoinKeys`' own
same-type
+ * guards keep the streamed key on the build side's type, so the streamed
key type
+ * is this index's `keyType`.
+ */
+ private[joins] def toLong(key: Any): Long = {
+ var f = _toLong
+ if (f == null) {
+ f = RangeIndex.requiredLongKey(keyType).value
+ _toLong = f
+ }
+ f(key)
+ }
+}
+
+private[joins] object RangeIndex {
+
+ /**
+ * Capacity for a build buffer from the driver's row-count hint, capped at
2^20 for any
+ * larger hint (well under a hint that large) and equal to the hint itself
below that.
+ * `BroadcastExchangeExec` only checks the row count against its row-count
+ * limit (up to half a billion for this mode) before calling `transform`,
and only
+ * checks the built relation's byte size afterward, so sizing a buffer
directly off an
+ * inflated hint can allocate hundreds of megabytes before a single row is
read - ahead
+ * of the size check that is supposed to catch an over-large broadcast.
Starting small
+ * and letting the buffer grow from there bounds only the initial
allocation; the final
+ * size still tracks the actual row count, which is what the post-build size
check sees.
+ */
+ def initialCapacity(sizeHint: Option[Long]): Int =
+ sizeHint.map(n => math.min(n, 1L << 20).toInt).getOrElse(64)
+
+ /**
+ * One key type's bridge to an order-preserving long: converting an evaluated
+ * key value (from either the build or the streamed side), and turning
+ * generated Java for a typed key into generated Java for a long. The two
+ * agree, so a build key and a probe key of the same type become the same
+ * long.
+ */
+ final case class LongKey(value: Any => Long, code: String => String)
+
+ /**
+ * An order-preserving long for a key type that has one, or None when the
type
+ * must stay on the generic boxed index: variable-length keys, floating
point,
+ * whose SQL equality for -0.0 and NaN has no bit-preserving order, and
+ * decimals wider than `Decimal.MAX_LONG_DIGITS`, whose unscaled value does
+ * not fit a long. This is the only place a key type is specialized; the long
+ * indexes and the generated probe code both consult it.
+ */
+ def longKey(dt: DataType): Option[LongKey] = {
+ def key(value: Any => Long): LongKey = LongKey(value, v => s"((long) $v)")
+ dt match {
+ case _: BooleanType =>
+ Some(LongKey(v => if (v.asInstanceOf[Boolean]) 1L else 0L, v => s"($v
? 1L : 0L)"))
+ case _: ByteType =>
+ Some(key(v => v.asInstanceOf[Byte].toLong))
+ case _: ShortType =>
+ Some(key(v => v.asInstanceOf[Short].toLong))
+ case _: IntegerType | _: DateType =>
+ Some(key(v => v.asInstanceOf[Int].toLong))
+ case _: LongType | _: TimestampType | _: TimestampNTZType =>
+ Some(key(v => v.asInstanceOf[Long]))
+ case d: DecimalType if d.precision <= Decimal.MAX_LONG_DIGITS =>
+ Some(LongKey(v => v.asInstanceOf[Decimal].toUnscaledLong, v =>
s"$v.toUnscaledLong()"))
+ case _ => None
+ }
+ }
+
+ /**
+ * Every key maps to one order-preserving long.
[[RangeRelationBroadcastMode]] picks
+ * the index family by this test and [[BroadcastRangeJoinExec]] picks its
probe
+ * code by the same one, so the two cannot drift apart. Within a predicate,
the
+ * analyzer's comparison type coercion gives both sides one type; across
predicates,
+ * `ExtractRangeJoinKeys`' own same-type guards keep build and streamed keys
aligned,
+ * so testing the build keys also settles the streamed keys.
+ */
+ def canMapKeysToLong(keys: Seq[Expression]): Boolean =
+ keys.forall(k => longKey(k.dataType).isDefined)
+
+ /** longKey for a type the caller has routed to the long family. */
+ def requiredLongKey(dt: DataType): LongKey = {
+ longKey(dt).getOrElse(
+ throw new IllegalStateException(s"$dt has no order-preserving long
key."))
+ }
+}
+
+/**
+ * Array-backed candidate scan. The index hands rows to the join in array
order,
+ * without a closure or a tuple hop: the per-pair overhead then matches a bare
+ * loop over a broadcast row array. One instance per probe; never serialized.
+ */
+private[joins] final class RowSlice(
+ rows: Array[InternalRow],
+ from: Int,
+ until: Int) extends Iterator[InternalRow] {
+ private var i = from
+ def hasNext: Boolean = i < until
+ def next(): InternalRow = {
+ val r = rows(i)
+ i += 1
+ r
+ }
+}
+
+private[joins] object RowSlice {
+ /** A shared empty scan: a null probe bound matches nothing. */
+ val empty: RowSlice = new RowSlice(new Array[InternalRow](0), 0, 0)
+}
+
+// ===== Boxed-key indexes ====================================================
+//
+// Keys are the boxed values each build key expression evaluates to, compared
+// through `BoxedKeyOrdering`. Floating point, strings, binary, and decimals
wider
+// than 18 digits stay here: these types have no order-preserving
+// single-long form, which is the property that routes a key type to the
+// long indexes below (`RangeIndex.longKey`).
+
+/**
+ * Augmented interval tree of build-side intervals. Lows and highs sit in
+ * parallel arrays sorted by low bound, with a row array in the same order
+ * feeding the scan path and `maxHighIdx(mid)` indexing the maximum high
+ * bound in the span rooted at mid. Each key and row object is owned once:
+ * the tree holds one reference per value plus the single row array. Java
+ * serialization keeps these aliases; Kryo's reference resolver does not
+ * track primitive wrappers, so a key shared by `lows` and `highs`, such as
+ * the one key of a degenerate build, can deserialize as two objects.
+ * Subtrees whose max high is below the probe window, or whose low starts
+ * after it, are pruned during search, and a window covering every interval
+ * skips the walk entirely.
+ */
+private[joins] class IntervalIndex private (
+ private[this] var _keyType: DataType,
+ private[this] var lows: Array[Any],
+ private[this] var highs: Array[Any],
+ private[this] var maxHighIdx: Array[Int],
+ private[this] var rows: Array[InternalRow])
+ extends BoxedKeyOrdering with Externalizable with KryoSerializable {
+
+ // Needed for Java/Kryo deserialization; see readExternal/read below.
+ def this() = this(null, null, null, null, null)
+
+ def keyType: DataType = _keyType
+
+ // All-covering window bounds. `lows` is sorted, so the last entry is the
+ // largest low: it bounds every low from above, and coreHigh is the minimum
+ // high, which bounds every high from below. A window overlaps every stored
+ // interval exactly when it holds probeLow <= coreHigh and probeHigh >=
coreLow,
+ // whether or not the intervals share a common core.
+ // Recomputed in `read` too, since `lows`/`highs` are reassigned there.
+ private[this] var coreLow: Any = IntervalIndex.coreLow(lows)
+ private[this] var coreHigh: Any = IntervalIndex.coreHigh(highs, ordering)
+
+ /**
+ * Measured over the retained arrays, not the index object: SizeEstimator
+ * would consult `estimatedSize` for a KnownSizeEstimation and recurse.
+ * Its identity tracking charges a key shared by `lows` and `highs`, or
+ * aliased across indexes, exactly once.
+ */
+ override def estimatedSize: Long =
+ SizeEstimator.estimate(Array[AnyRef](lows, highs, maxHighIdx, rows))
+
+ def overlapping(low: Any, high: Any): Iterator[InternalRow] = {
+ if (lows.isEmpty || low == null || high == null) return Iterator.empty
+ val ord = ordering
+ val (probeLow, probeHigh) = if (ord.lteq(low, high)) (low, high) else
(high, low)
+ // A window that covers every interval gets every row the walk would
+ // return, without visiting the tree. Returns all rows iff every interval
+ // overlaps the probe; the condition still re-checks each candidate.
+ if (ord.lteq(probeLow, coreHigh) && ord.gteq(probeHigh, coreLow)) {
+ return new RowSlice(rows, 0, rows.length)
+ }
+ new OverlapIterator(lows, highs, maxHighIdx, rows, ord, probeLow,
probeHigh)
+ }
+
+ override def writeExternal(out: ObjectOutput): Unit = Utils.tryOrIOException
{
+ write(out.writeUTF, out.writeInt, (o: Any) =>
out.writeObject(o.asInstanceOf[AnyRef]))
+ }
+
+ override def write(kryo: Kryo, out: Output): Unit = Utils.tryOrIOException {
+ write(out.writeString, out.writeInt,
+ (o: Any) => kryo.writeClassAndObject(out, o.asInstanceOf[AnyRef]))
+ }
+
+ private def write(
+ writeString: String => Unit,
+ writeInt: Int => Unit,
+ writeObj: Any => Unit): Unit = {
+ writeString(keyType.json)
+ val n = lows.length
+ writeInt(n)
+ var i = 0
+ while (i < n) {
+ writeObj(lows(i))
+ writeObj(highs(i))
+ i += 1
+ }
+ i = 0
+ while (i < n) {
+ writeInt(maxHighIdx(i))
+ i += 1
+ }
+ i = 0
+ while (i < n) {
+ writeObj(rows(i))
+ i += 1
+ }
+ }
+
+ override def readExternal(in: ObjectInput): Unit = Utils.tryOrIOException {
+ read(() => in.readUTF(), () => in.readInt(), () => in.readObject())
+ }
+
+ override def read(kryo: Kryo, in: Input): Unit = Utils.tryOrIOException {
+ read(() => in.readString(), () => in.readInt(), () =>
kryo.readClassAndObject(in))
+ }
+
+ private def read(
+ readString: () => String,
+ readInt: () => Int,
+ readObj: () => AnyRef): Unit = {
+ _keyType = DataType.fromJson(readString())
+ val n = readInt()
+ lows = new Array[Any](n)
+ highs = new Array[Any](n)
+ maxHighIdx = new Array[Int](n)
+ rows = new Array[InternalRow](n)
+ var i = 0
+ while (i < n) {
+ lows(i) = readObj()
+ highs(i) = readObj()
+ i += 1
+ }
+ i = 0
+ while (i < n) {
+ maxHighIdx(i) = readInt()
+ i += 1
+ }
+ i = 0
+ while (i < n) {
+ rows(i) = readObj().asInstanceOf[InternalRow]
+ i += 1
+ }
+ coreLow = IntervalIndex.coreLow(lows)
+ coreHigh = IntervalIndex.coreHigh(highs, ordering)
+ }
+}
+
+private[joins] object IntervalIndex {
+
+ // Shared by the primary constructor and `read`, which recomputes these
+ // caches after reassigning `lows`/`highs` on deserialization. `lows`/
+ // `highs` are null right after the zero-arg constructor runs, before
+ // `readExternal`/`read` populate them.
+ private[joins] def coreLow(lows: Array[Any]): Any =
+ if (lows == null || lows.isEmpty) null else lows(lows.length - 1)
+
+ // `ordering` is by-name: the zero-arg constructor calls this with `highs ==
+ // null`, before `keyType` is set, and forcing `ordering` eagerly would
+ // throw trying to derive an ordering for a not-yet-known key type.
+ private[joins] def coreHigh(highs: Array[Any], ordering: => Ordering[Any]):
Any = {
+ if (highs == null || highs.isEmpty) {
+ null
+ } else {
+ val ord = ordering
+ var min = highs(0)
+ var i = 1
+ while (i < highs.length) {
+ if (ord.lt(highs(i), min)) min = highs(i)
+ i += 1
+ }
+ min
+ }
+ }
+
+ def apply(
+ keyType: DataType,
+ buildKeys: Seq[Expression],
+ rows: Iterator[InternalRow],
+ sizeHint: Option[Long]): IntervalIndex = {
+ val buffer = new ArrayBuffer[(Any, Any,
InternalRow)](RangeIndex.initialCapacity(sizeHint))
+ val ordering = PhysicalDataType.ordering(keyType)
+
+ if (buildKeys.length == 1) {
+ // A point build key is one key used as both ends, so it is copied once
+ // and serves as both bounds, matching `LongIntervalIndex.apply`.
+ val lowExpr = buildKeys.head
+ while (rows.hasNext) {
+ val r = rows.next()
+ val rawLo = lowExpr.eval(r)
+ if (rawLo != null) {
+ val key = InternalRow.copyValue(rawLo)
+ buffer += ((key, key, r.copy()))
+ }
+ }
+ } else {
+ val lowExpr = buildKeys.head
+ val highExpr = buildKeys(1)
+ while (rows.hasNext) {
+ val r = rows.next()
+ val rawLo = lowExpr.eval(r)
+ val rawHi = highExpr.eval(r)
+ if (rawLo != null && rawHi != null) {
+ val lo = InternalRow.copyValue(rawLo)
+ val hi = InternalRow.copyValue(rawHi)
+ val (normLo, normHi) = if (ordering.lteq(lo, hi)) (lo, hi) else (hi,
lo)
+ buffer += ((normLo, normHi, r.copy()))
+ }
+ }
+ }
+
+ sortAndBuild(keyType, ordering, buffer.toArray)
+ }
+
+ private def sortAndBuild(
+ keyType: DataType,
+ ordering: Ordering[Any],
+ valid: Array[(Any, Any, InternalRow)]): IntervalIndex = {
+ java.util.Arrays.sort(valid, Ordering.by[(Any, Any, InternalRow),
Any](_._1)(ordering))
+ val length = valid.length
+ val lows = new Array[Any](length)
+ val highs = new Array[Any](length)
+ val rows = new Array[InternalRow](length)
+ var i = 0
+ while (i < length) {
+ lows(i) = valid(i)._1
+ highs(i) = valid(i)._2
+ rows(i) = valid(i)._3
+ i += 1
+ }
+
+ // `maxHighIdx(mid)` indexes the maximum high bound in the span rooted at
+ // mid. Storing an index instead of aliasing the key keeps each key owned,
+ // and serialized, exactly once.
+ val maxHighIdx = new Array[Int](length)
+
+ def fillMaxHigh(start: Int, end: Int): Int = {
+ if (start >= end) return -1
+ val mid = (start + end) >>> 1
+ var best = mid
+ val left = fillMaxHigh(start, mid)
+ if (left >= 0 && ordering.gt(highs(left), highs(best))) best = left
+ val right = fillMaxHigh(mid + 1, end)
+ if (right >= 0 && ordering.gt(highs(right), highs(best))) best = right
+ maxHighIdx(mid) = best
+ best
+ }
+
+ fillMaxHigh(0, length)
+ new IntervalIndex(keyType, lows, highs, maxHighIdx, rows)
+ }
+}
+
+private[joins] class OverlapIterator(
+ lows: Array[Any],
+ highs: Array[Any],
+ maxHighIdx: Array[Int],
+ rows: Array[InternalRow],
+ ord: Ordering[Any],
+ probeLow: Any,
+ probeHigh: Any) extends Iterator[InternalRow] {
+
+ private[this] val depth = 32 -
java.lang.Integer.numberOfLeadingZeros(lows.length)
+ private[this] var stackStart = new Array[Int](math.max(8, depth * 3))
+ private[this] var stackEnd = new Array[Int](math.max(8, depth * 3))
+ private[this] var stackTop = 0
+ private[this] var nextRow: InternalRow = null
+
+ push(0, lows.length)
+
+ private[joins] def stackCapacity: Int = stackStart.length
+ private[joins] def pendingSpans: Int = stackTop
+
+ private def push(s: Int, e: Int): Unit = {
+ if (stackTop >= stackStart.length) {
+ stackStart = java.util.Arrays.copyOf(stackStart, stackStart.length * 2)
+ stackEnd = java.util.Arrays.copyOf(stackEnd, stackEnd.length * 2)
+ }
+ stackStart(stackTop) = s
+ stackEnd(stackTop) = e
+ stackTop += 1
+ }
+
+ private def advance(): InternalRow = {
+ while (stackTop > 0) {
+ stackTop -= 1
+ val start = stackStart(stackTop)
+ val end = stackEnd(stackTop)
+ if (start < end) {
+ val mid = (start + end) >>> 1
+ if (ord.gteq(highs(maxHighIdx(mid)), probeLow)) {
+ if (ord.lteq(lows(mid), probeHigh)) {
+ push(mid + 1, end)
+ }
+ push(mid, -1)
+ push(start, mid)
+ }
+ } else if (end == -1) {
+ if (ord.gteq(highs(start), probeLow) && ord.lteq(lows(start),
probeHigh)) {
+ return rows(start)
+ }
+ }
+ }
+ null
+ }
+
+ override def hasNext: Boolean = {
+ if (nextRow != null) {
+ true
+ } else {
+ nextRow = advance()
+ nextRow != null
+ }
+ }
+
+ override def next(): InternalRow = {
+ if (!hasNext) throw new NoSuchElementException("empty iterator")
+ val r = nextRow
+ nextRow = null
+ r
+ }
+}
+
+/**
+ * Sorted build-side points for single cross-side inequalities (`<`, `<=`,
`>`, `>=`).
+ * Keys sit beside the rows they name, one owned copy each, in key order.
+ */
+private[joins] class PointIndex private (
+ private[this] var _keyType: DataType,
+ private[this] var keys: Array[Any],
+ private[this] var rows: Array[InternalRow])
+ extends BoxedKeyOrdering with Externalizable with KryoSerializable {
+
+ // Needed for Java/Kryo deserialization; see readExternal/read below.
+ def this() = this(null, null, null)
+
+ def keyType: DataType = _keyType
+
+ private[this] def length: Int = keys.length
+
+ override def estimatedSize: Long =
SizeEstimator.estimate(Array[AnyRef](keys, rows))
+
+ /** Points whose key is less than or equal to `value`. */
+ def upTo(value: Any): Iterator[InternalRow] =
+ if (value == null) Iterator.empty else slice(0, searchBound(value,
skipEquals = true))
+
+ /** Points whose key is greater than or equal to `value`. */
+ def from(value: Any): Iterator[InternalRow] =
+ if (value == null) Iterator.empty else slice(searchBound(value, skipEquals
= false), length)
+
+ /**
+ * Points whose key is strictly less than `value`. For a strict partial-range
+ * predicate the original condition rejects keys equal to the probe value,
+ * so the scan can skip them: a build heavy in duplicated keys would only
+ * hand the join rows that all fail.
+ */
+ def upToExclusive(value: Any): Iterator[InternalRow] =
+ if (value == null) Iterator.empty else slice(0, searchBound(value,
skipEquals = false))
+
+ /** Points whose key is strictly greater than `value`. */
+ def fromExclusive(value: Any): Iterator[InternalRow] =
+ if (value == null) Iterator.empty else slice(searchBound(value, skipEquals
= true), length)
+
+ private def searchBound(value: Any, skipEquals: Boolean): Int = {
+ val ord = ordering
+ var lo = 0
+ var hi = length
+ while (lo < hi) {
+ val mid = (lo + hi) >>> 1
+ val c = ord.compare(keys(mid), value)
+ if (c < 0 || (skipEquals && c == 0)) lo = mid + 1 else hi = mid
+ }
+ lo
+ }
+
+ private def slice(from: Int, until: Int): Iterator[InternalRow] = new
RowSlice(rows, from, until)
+
+ override def writeExternal(out: ObjectOutput): Unit = Utils.tryOrIOException
{
+ write(out.writeUTF, out.writeInt, (o: Any) =>
out.writeObject(o.asInstanceOf[AnyRef]))
+ }
+
+ override def write(kryo: Kryo, out: Output): Unit = Utils.tryOrIOException {
+ write(out.writeString, out.writeInt,
+ (o: Any) => kryo.writeClassAndObject(out, o.asInstanceOf[AnyRef]))
+ }
+
+ private def write(
+ writeString: String => Unit,
+ writeInt: Int => Unit,
+ writeObj: Any => Unit): Unit = {
+ writeString(keyType.json)
+ val n = keys.length
+ writeInt(n)
+ var i = 0
+ while (i < n) {
+ writeObj(keys(i))
+ i += 1
+ }
+ i = 0
+ while (i < n) {
+ writeObj(rows(i))
+ i += 1
+ }
+ }
+
+ override def readExternal(in: ObjectInput): Unit = Utils.tryOrIOException {
+ read(() => in.readUTF(), () => in.readInt(), () => in.readObject())
+ }
+
+ override def read(kryo: Kryo, in: Input): Unit = Utils.tryOrIOException {
+ read(() => in.readString(), () => in.readInt(), () =>
kryo.readClassAndObject(in))
+ }
+
+ private def read(
+ readString: () => String,
+ readInt: () => Int,
+ readObj: () => AnyRef): Unit = {
+ _keyType = DataType.fromJson(readString())
+ val n = readInt()
+ keys = new Array[Any](n)
+ rows = new Array[InternalRow](n)
+ var i = 0
+ while (i < n) {
+ keys(i) = readObj()
+ i += 1
+ }
+ i = 0
+ while (i < n) {
+ rows(i) = readObj().asInstanceOf[InternalRow]
+ i += 1
+ }
+ }
+}
+
+private[joins] object PointIndex {
+
+ def apply(
+ keyType: DataType,
+ buildKeys: Seq[Expression],
+ rows: Iterator[InternalRow],
+ sizeHint: Option[Long]): PointIndex = {
+ val buffer = new ArrayBuffer[(Any,
InternalRow)](RangeIndex.initialCapacity(sizeHint))
+ val keyExpr = buildKeys.head
+ val ordering = PhysicalDataType.ordering(keyType)
+
+ while (rows.hasNext) {
+ val r = rows.next()
+ val rawKey = keyExpr.eval(r)
+ if (rawKey != null) {
+ val key = InternalRow.copyValue(rawKey)
+ buffer += ((key, r.copy()))
+ }
+ }
+
+ sortAndBuild(keyType, ordering, buffer.toArray)
+ }
+
+ private def sortAndBuild(
+ keyType: DataType,
+ ordering: Ordering[Any],
+ valid: Array[(Any, InternalRow)]): PointIndex = {
+ java.util.Arrays.sort(valid, Ordering.by[(Any, InternalRow),
Any](_._1)(ordering))
+ val length = valid.length
+ val keys = new Array[Any](length)
+ val rows = new Array[InternalRow](length)
+ var i = 0
+ while (i < length) {
+ keys(i) = valid(i)._1
+ rows(i) = valid(i)._2
+ i += 1
+ }
+ new PointIndex(keyType, keys, rows)
+ }
+}
+
+// ===== Long-specialized indexes for keys that map to one long ==============
+//
+// Indexes for the key types `RangeIndex.longKey` maps to one order-preserving
+// long, stored as raw longs converted at the build boundary. The key
comparison
+// itself touches no Spark types once built: the tree walk and the binary
searches
+// compare raw longs, and generated probe code passes long keys without boxing.
+// The row payload still does - `rows: Array[InternalRow]` and the serialized
+// `_keyType: DataType` are retained alongside the long arrays. Fixed width
alone
+// is not the routed property: floating point is fixed-width but stays boxed,
+// while compact decimals route here as unscaled longs.
+// `RangeRelationBroadcastMode` picks between the two families once, from the
+// build keys' data types.
+
+/**
+ * Augmented interval tree of build-side intervals, on raw keys. Lows and
highs sit in
+ * parallel arrays sorted by low bound, with a row array in the same order
feeding the
+ * scan path and `maxHighs(mid)` holding the maximum high bound in the span
rooted at
+ * mid - the value itself, unlike [[IntervalIndex.maxHighIdx]], which stores
an index
+ * into `highs` because boxed values there are shared references worth not
duplicating.
+ * `highs` here is already an unboxed `Array[Long]`, so storing the max value
directly
+ * costs 4 bytes/row more than an Int index would (a Long vs. an Int), in
exchange for
+ * comparing it without the extra `highs(idx)` indirection on every node a
probe visits.
+ */
+private[joins] class LongIntervalIndex private (
+ private[this] var _keyType: DataType,
+ private[this] var lows: Array[Long],
+ private[this] var highs: Array[Long],
+ private[this] var maxHighs: Array[Long],
+ private[this] var rows: Array[InternalRow])
+ extends LongKeyConversion with Externalizable with KryoSerializable {
+
+ // Needed for Java/Kryo deserialization; see readExternal/read below.
+ def this() = this(null, null, null, null, null)
+
+ def keyType: DataType = _keyType
+
+ // All-covering window bounds on raw keys: `lows` is sorted, so the last
entry
+ // is the largest low, bounding every low from above, and coreHigh is the
+ // minimum high, bounding every high from below. A window overlaps every
+ // stored interval exactly when it holds probeLow <= coreHigh and
+ // probeHigh >= coreLow, whether or not a common core exists.
+ // Recomputed in `read` too, since `lows`/`highs` are reassigned there.
+ private[this] var coreLow: Long = LongIntervalIndex.coreLow(lows)
+ private[this] var coreHigh: Long = LongIntervalIndex.coreHigh(highs)
+
+ override def estimatedSize: Long =
+ SizeEstimator.estimate(Array[AnyRef](lows, highs, maxHighs, rows))
+
+ def overlapping(low: Long, high: Long): Iterator[InternalRow] = {
+ if (lows.isEmpty) return Iterator.empty
+ var probeLow = low
+ var probeHigh = high
+ if (probeLow > probeHigh) {
+ val t = probeLow
+ probeLow = probeHigh
+ probeHigh = t
+ }
+ // A window that covers every interval gets every row the walk would
+ // return, without visiting the tree. Returns all rows iff every interval
+ // overlaps the probe; the condition still re-checks each candidate.
+ if (probeLow <= coreHigh && probeHigh >= coreLow) {
+ return new RowSlice(rows, 0, rows.length)
+ }
+ new LongOverlapIterator(lows, highs, rows, maxHighs, probeLow, probeHigh)
+ }
+
+ override def writeExternal(out: ObjectOutput): Unit = Utils.tryOrIOException
{
+ write(out.writeUTF, out.writeInt, out.writeLong,
+ (o: Any) => out.writeObject(o.asInstanceOf[AnyRef]))
+ }
+
+ override def write(kryo: Kryo, out: Output): Unit = Utils.tryOrIOException {
+ write(out.writeString, out.writeInt, out.writeLong,
+ (o: Any) => kryo.writeClassAndObject(out, o.asInstanceOf[AnyRef]))
+ }
+
+ private def write(
+ writeString: String => Unit,
+ writeInt: Int => Unit,
+ writeLong: Long => Unit,
+ writeObj: Any => Unit): Unit = {
+ writeString(keyType.json)
+ val n = lows.length
+ writeInt(n)
+ var i = 0
+ while (i < n) {
+ writeLong(lows(i))
+ writeLong(highs(i))
+ writeLong(maxHighs(i))
+ i += 1
+ }
+ i = 0
+ while (i < n) {
+ writeObj(rows(i))
+ i += 1
+ }
+ }
+
+ override def readExternal(in: ObjectInput): Unit = Utils.tryOrIOException {
+ read(() => in.readUTF(), () => in.readInt(), () => in.readLong(), () =>
in.readObject())
+ }
+
+ override def read(kryo: Kryo, in: Input): Unit = Utils.tryOrIOException {
+ read(() => in.readString(), () => in.readInt(), () => in.readLong(),
+ () => kryo.readClassAndObject(in))
+ }
+
+ private def read(
+ readString: () => String,
+ readInt: () => Int,
+ readLong: () => Long,
+ readObj: () => AnyRef): Unit = {
+ _keyType = DataType.fromJson(readString())
+ val n = readInt()
+ lows = new Array[Long](n)
+ highs = new Array[Long](n)
+ maxHighs = new Array[Long](n)
+ rows = new Array[InternalRow](n)
+ var i = 0
+ while (i < n) {
+ lows(i) = readLong()
+ highs(i) = readLong()
+ maxHighs(i) = readLong()
+ i += 1
+ }
+ i = 0
+ while (i < n) {
+ rows(i) = readObj().asInstanceOf[InternalRow]
+ i += 1
+ }
+ coreLow = LongIntervalIndex.coreLow(lows)
+ coreHigh = LongIntervalIndex.coreHigh(highs)
+ }
+}
+
+private[joins] object LongIntervalIndex {
+
+ // Shared by the primary constructor and `read`, which recomputes these
+ // caches after reassigning `lows`/`highs` on deserialization. `lows`/
+ // `highs` are null right after the zero-arg constructor runs, before
+ // `readExternal`/`read` populate them.
+ private[joins] def coreLow(lows: Array[Long]): Long =
+ if (lows == null || lows.isEmpty) {
+ Long.MinValue
+ } else {
+ lows(lows.length - 1)
+ }
+
+ private[joins] def coreHigh(highs: Array[Long]): Long = {
+ if (highs == null || highs.isEmpty) {
+ Long.MaxValue
+ } else {
+ var min = Long.MaxValue
+ var i = 0
+ while (i < highs.length) {
+ if (highs(i) < min) min = highs(i)
+ i += 1
+ }
+ min
+ }
+ }
+
+ def apply(
+ keyType: DataType,
+ buildKeys: Seq[Expression],
+ rows: Iterator[InternalRow],
+ sizeHint: Option[Long]): LongIntervalIndex = {
+ val toLongs: Seq[Any => Long] = buildKeys.map(k =>
RangeIndex.requiredLongKey(k.dataType).value)
+ val valid = new ArrayBuffer[(Long, Long,
InternalRow)](RangeIndex.initialCapacity(sizeHint))
+ val lowExpr = buildKeys.head
+ val highExpr = if (buildKeys.length == 1) null else buildKeys(1)
+
+ while (rows.hasNext) {
+ val r = rows.next()
+ val rawLo = lowExpr.eval(r)
+ val rawHi = if (buildKeys.length == 1) rawLo else highExpr.eval(r)
+ if (rawLo != null && rawHi != null) {
+ // A point build key is one key used as both ends, so it is read once.
+ var lo = toLongs.head(rawLo)
+ var hi = if (buildKeys.length == 1) lo else toLongs(1)(rawHi)
+ if (lo > hi) {
+ val t = lo
+ lo = hi
+ hi = t
+ }
+ valid += ((lo, hi, r.copy()))
+ }
+ }
+
+ val staged = valid.toArray
+ java.util.Arrays.sort(staged, Ordering.by[(Long, Long, InternalRow),
Long](_._1))
+ val length = staged.length
+ val lows = new Array[Long](length)
+ val highs = new Array[Long](length)
+ val rowArr = new Array[InternalRow](length)
+ var i = 0
+ while (i < length) {
+ lows(i) = staged(i)._1
+ highs(i) = staged(i)._2
+ rowArr(i) = staged(i)._3
+ i += 1
+ }
+
+ val maxHighs = new Array[Long](length)
+
+ def fillMaxHigh(start: Int, end: Int): Long = {
+ if (start >= end) return Long.MinValue
+ val mid = (start + end) >>> 1
+ var max = highs(mid)
+ val left = fillMaxHigh(start, mid)
+ val right = fillMaxHigh(mid + 1, end)
+ if (left > max) max = left
+ if (right > max) max = right
+ maxHighs(mid) = max
+ max
+ }
+
+ fillMaxHigh(0, length)
+ new LongIntervalIndex(keyType, lows, highs, maxHighs, rowArr)
+ }
+}
+
+private[joins] class LongOverlapIterator(
+ lows: Array[Long],
+ highs: Array[Long],
+ rows: Array[InternalRow],
+ maxHighs: Array[Long],
+ probeLow: Long,
+ probeHigh: Long) extends Iterator[InternalRow] {
+
+ private[this] val depth = 32 -
java.lang.Integer.numberOfLeadingZeros(lows.length)
+ private[this] var stackStart = new Array[Int](math.max(8, depth * 3))
+ private[this] var stackEnd = new Array[Int](math.max(8, depth * 3))
+ private[this] var stackTop = 0
+ private[this] var nextRow: InternalRow = null
+
+ push(0, lows.length)
+
+ private[joins] def stackCapacity: Int = stackStart.length
+ private[joins] def pendingSpans: Int = stackTop
+
+ private def push(s: Int, e: Int): Unit = {
+ if (stackTop >= stackStart.length) {
+ stackStart = java.util.Arrays.copyOf(stackStart, stackStart.length * 2)
+ stackEnd = java.util.Arrays.copyOf(stackEnd, stackEnd.length * 2)
+ }
+ stackStart(stackTop) = s
+ stackEnd(stackTop) = e
+ stackTop += 1
+ }
+
+ private def advance(): InternalRow = {
+ while (stackTop > 0) {
+ stackTop -= 1
+ val start = stackStart(stackTop)
+ val end = stackEnd(stackTop)
+ if (start < end) {
+ val mid = (start + end) >>> 1
+ if (maxHighs(mid) >= probeLow) {
+ if (lows(mid) <= probeHigh) {
+ push(mid + 1, end)
+ }
+ push(mid, -1)
+ push(start, mid)
+ }
+ } else if (end == -1) {
+ if (highs(start) >= probeLow && lows(start) <= probeHigh) {
+ return rows(start)
+ }
+ }
+ }
+ null
+ }
+
+ override def hasNext: Boolean = {
+ if (nextRow != null) {
+ true
+ } else {
+ nextRow = advance()
+ nextRow != null
+ }
+ }
+
+ override def next(): InternalRow = {
+ if (!hasNext) throw new NoSuchElementException("empty iterator")
+ val r = nextRow
+ nextRow = null
+ r
+ }
+}
+
+private[joins] class LongPointIndex private (
+ private[this] var _keyType: DataType,
+ private[this] var keys: Array[Long],
+ private[this] var rows: Array[InternalRow])
+ extends LongKeyConversion with Externalizable with KryoSerializable {
+
+ // Needed for Java/Kryo deserialization; see readExternal/read below.
+ def this() = this(null, null, null)
+
+ def keyType: DataType = _keyType
+
+ private[this] def length: Int = keys.length
+
+ override def estimatedSize: Long =
SizeEstimator.estimate(Array[AnyRef](keys, rows))
+
+ /** Points whose key is less than or equal to `value`. */
+ def upTo(value: Long): Iterator[InternalRow] =
+ new RowSlice(rows, 0, searchBound(value, skipEquals = true))
+
+ /** Points whose key is greater than or equal to `value`. */
+ def from(value: Long): Iterator[InternalRow] =
+ new RowSlice(rows, searchBound(value, skipEquals = false), length)
+
+ /**
+ * Points whose key is strictly less than `value`. For a strict partial-range
+ * predicate the original condition rejects keys equal to the probe value,
+ * so the scan can skip them; see [[PointIndex.upToExclusive]].
+ */
+ def upToExclusive(value: Long): Iterator[InternalRow] =
+ new RowSlice(rows, 0, searchBound(value, skipEquals = false))
+
+ /** Points whose key is strictly greater than `value`. */
+ def fromExclusive(value: Long): Iterator[InternalRow] =
+ new RowSlice(rows, searchBound(value, skipEquals = true), length)
+
+ // java.util.Arrays.binarySearch returns an arbitrary hit; a bound scan needs
+ // the first or last position, so lower and upper bounds are hand-rolled.
+ private def searchBound(value: Long, skipEquals: Boolean): Int = {
+ var lo = 0
+ var hi = length
+ while (lo < hi) {
+ val mid = (lo + hi) >>> 1
+ val c = java.lang.Long.compare(keys(mid), value)
+ if (c < 0 || (skipEquals && c == 0)) lo = mid + 1 else hi = mid
+ }
+ lo
+ }
+
+ override def writeExternal(out: ObjectOutput): Unit = Utils.tryOrIOException
{
+ write(out.writeUTF, out.writeInt, out.writeLong,
+ (o: Any) => out.writeObject(o.asInstanceOf[AnyRef]))
+ }
+
+ override def write(kryo: Kryo, out: Output): Unit = Utils.tryOrIOException {
+ write(out.writeString, out.writeInt, out.writeLong,
+ (o: Any) => kryo.writeClassAndObject(out, o.asInstanceOf[AnyRef]))
+ }
+
+ private def write(
+ writeString: String => Unit,
+ writeInt: Int => Unit,
+ writeLong: Long => Unit,
+ writeObj: Any => Unit): Unit = {
+ writeString(keyType.json)
+ val n = keys.length
+ writeInt(n)
+ var i = 0
+ while (i < n) {
+ writeLong(keys(i))
+ i += 1
+ }
+ i = 0
+ while (i < n) {
+ writeObj(rows(i))
+ i += 1
+ }
+ }
+
+ override def readExternal(in: ObjectInput): Unit = Utils.tryOrIOException {
+ read(() => in.readUTF(), () => in.readInt(), () => in.readLong(), () =>
in.readObject())
+ }
+
+ override def read(kryo: Kryo, in: Input): Unit = Utils.tryOrIOException {
+ read(() => in.readString(), () => in.readInt(), () => in.readLong(),
+ () => kryo.readClassAndObject(in))
+ }
+
+ private def read(
+ readString: () => String,
+ readInt: () => Int,
+ readLong: () => Long,
+ readObj: () => AnyRef): Unit = {
+ _keyType = DataType.fromJson(readString())
+ val n = readInt()
+ keys = new Array[Long](n)
+ rows = new Array[InternalRow](n)
+ var i = 0
+ while (i < n) {
+ keys(i) = readLong()
+ i += 1
+ }
+ i = 0
+ while (i < n) {
+ rows(i) = readObj().asInstanceOf[InternalRow]
+ i += 1
+ }
+ }
+}
+
+private[joins] object LongPointIndex {
+
+ def apply(
+ keyType: DataType,
+ buildKeys: Seq[Expression],
+ rows: Iterator[InternalRow],
+ sizeHint: Option[Long]): LongPointIndex = {
+ require(buildKeys.length == 1, "A long point index expects one build key.")
+ val toLong = RangeIndex.requiredLongKey(buildKeys.head.dataType).value
+ val valid = new ArrayBuffer[(Long,
InternalRow)](RangeIndex.initialCapacity(sizeHint))
+ val keyExpr = buildKeys.head
+
+ while (rows.hasNext) {
+ val r = rows.next()
+ val rawKey = keyExpr.eval(r)
+ if (rawKey != null) {
+ valid += ((toLong(rawKey), r.copy()))
+ }
+ }
+
+ val staged = valid.toArray
+ java.util.Arrays.sort(staged, Ordering.by[(Long, InternalRow), Long](_._1))
+ val length = staged.length
+ val keys = new Array[Long](length)
+ val rowArr = new Array[InternalRow](length)
+ var i = 0
+ while (i < length) {
+ keys(i) = staged(i)._1
+ rowArr(i) = staged(i)._2
+ i += 1
+ }
+ new LongPointIndex(keyType, keys, rowArr)
+ }
+}
+
+/**
+ * Which side of a single-inequality range join's bound the build side holds,
and
+ * whether the comparison is strict - the four spellings a [[PointIndex]] /
+ * [[LongPointIndex]] expose as scan methods. A `match` on this sealed
hierarchy is
+ * exhaustiveness-checked by the compiler, so both the interpreted probe
+ * ([[BroadcastRangeJoinExec]]'s `makeCandidateFn`) and generated probe code
have to
Review Comment:
Fixed: the opening now scopes the exhaustiveness check to the interpreted
matches in `makeCandidateFn`, with generated code relying on `methodName`
naming a real index method.
##########
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:
Fixed: the doc names the most-but-not-all case and notes full-overlap
parity; the fixed ratios are dropped since they do not hold at cluster scale.
##########
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:
Fixed: the double views now use D-suffixed literals with a setup check that
both columns are DOUBLE, and the codegen-off probe-coverage list includes a
boxed double overlap query.
##########
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:
Fixed: added `BroadcastRangeJoin metrics count only accepted, non-duplicated
output rows` to SQLMetricsSuite covering inner, outer, left semi, and left anti
under codegen on and off.
--
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]
