cloud-fan commented on code in PR #59027:
URL: https://github.com/apache/spark/pull/59027#discussion_r4434569014
##########
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:
**Nit (P3):** As written, "but inclusivity otherwise" reads as a second
object of "selects", so the sentence seems to say `strict` selects inclusivity,
and the plural "stay" has no clear subject. The intended contract, which
`PartialRangeJoin.strict` already states, is that `strict` only picks the
exclusive point-index scan in `partialRange`, while inclusivity and every
leftover conjunct stay on the original condition. Something like: "`strict`
only selects the exclusive point-index scan in `partialRange`; inclusivity and
every leftover conjunct stay on the original condition, which is evaluated for
every candidate."
This extends an [existing
discussion](https://github.com/apache/spark/pull/59027#discussion_r4346073085).
**Verification:**
- **Inspection:** The rewritten sentence matches how partialRange carries
strict into the partial-range tags and how PointScan.forJoin uses it, and it
says the original condition rechecks every candidate.
##########
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:
**Nit (P3):** Each point fixture here is one key equal to the probe, so
`from(probe)` returns row 7 after any restore that keeps a single entry,
whatever ordering it re-derives. A `PointIndex.read` or `LongPointIndex.read`
that restored keys out of step with their rows, dropped entries after the
first, or lost the collation ordering would still pass. The comment above says
the collated probe needs the restored ordering, which holds for the interval
fixtures but not for these. The SQL suites run with a local master, where tasks
get the driver's object, so nothing else deserializes a point index. Could the
point families get a multi-entry case like `restored interval indexes answer
windows that need the walk`, with duplicate and collation-equal keys, checked
through `upTo`, `from`, and both exclusive scans after each serializer?
This extends an [existing
discussion](https://github.com/apache/spark/pull/59027#discussion_r4314225488).
**Verification:**
- **Regression:** For both point families and for Java, strict Kryo, and
Kryo without reference tracking, a multi-entry point index with duplicate and
collation-equal keys answers inclusive and exclusive bound scans after the
round trip exactly as before it, and the check fails if the restore reverses or
truncates the entries.
- **Behavior:** The existing interval round-trip and restored-walk checks
still pass for every key family.
##########
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:
**Nit (P3):** Generated code never matches on `PointScan`:
`genCandidateIterator` splices `pointScan.methodName` into the Java source, as
this scaladoc itself says a few lines down. So the exhaustiveness check covers
only the two matches in `makeCandidateFn`, and a fifth case compiles fine on
the generated path and fails only when Janino compiles a query. Could the
opening say the check covers the interpreted matches, while generated code
relies on `methodName` naming a real index method?
**Verification:**
- **Inspection:** Every statement in the PointScan scaladoc about what the
compiler checks matches the matches in makeCandidateFn and the methodName
splice in the generated probe.
--
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]
