cloud-fan commented on code in PR #59027:
URL: https://github.com/apache/spark/pull/59027#discussion_r4346073085
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/planning/patterns.scala:
##########
@@ -421,6 +422,197 @@ 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 RangeJoin
+
+/** A point-in-range join: a value from one side falls within a `[low, high]`
range on the other. */
+case object PointInRangeJoin extends RangeJoin
+
+/**
+ * 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 RangeJoin
+
+/**
+ * A partial-range join: one inequality, one column per side.
+ * `leftIsLower` says which child holds the lower bound.
+ */
+sealed abstract class PartialRangeJoin extends RangeJoin {
+ 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 column per side.
+ * Leftover conjuncts, including inclusivity ([[RangePredicate]] only extracts
the
Review Comment:
**Nit (P3):** Thanks for moving this block back above `object
ExtractRangeJoinKeys`, so it documents the extractor again. The parenthetical
still has the pre-strict wording, though. `RangePredicate` now returns
`(dataType, low, high, strict)`, and `partialRange` carries `strict` into
`LessPartialRangeJoin`/`GreaterPartialRangeJoin`, where `PointScan.forJoin`
uses it to pick `fromExclusive`/`upToExclusive`. The original condition still
re-checks inclusivity and every leftover conjunct, but strictness is no longer
left to it alone. So saying `RangePredicate` 'only extracts the canonical `low
<= high` shape' tells a reader that planning ignores strictness. Could the
clause mention the strict flag and its use for the exclusive scans?
**Verification:**
- **Inspection:** The extractor doc's description of RangePredicate matches
its four-field result and the strict flag's use in selecting exclusive
point-index scans.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala:
##########
@@ -1220,6 +1220,20 @@ 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 may perform worse than nested loop joins when nearly all intervals
overlap. " +
Review Comment:
**Nit (P3):** The committed benchmark results show a wider slow regime than
this sentence describes. In all three result files, the staggered
interval-overlap cases (most but not all build intervals overlap each probe, so
the tree cannot prune) run at 0.7X-0.8X on long keys and 0.5X-0.6X on double
keys, at 50% overlap as well as 90%. The fully overlapping case, which the
all-covering shortcut answers straight from the row array, runs at 1.0X-1.1X.
So 'nearly all intervals overlap' leaves out the 50% rows and the cost, and it
can read as if full overlap were the risky case. Once the flag is on,
JoinSelection picks the range join for every eligible non-equi join with no
cost check, so this sentence is the only guidance on when to leave the flag
off. The previous revision's wording ('slower when most, but not all, build
intervals overlap each probe: about 0.7X on long keys and 0.6X on boxed keys')
matched the results. Could that wording, or an equivalent, come back?
**Verification:**
- **Inspection:** Every performance statement in the config doc agrees with
the full-overlap, staggered 50% and 90% overlap, tiny-probe, and median-bound
rows of all three committed benchmark result files.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/joins/BroadcastRangeJoinExec.scala:
##########
@@ -0,0 +1,598 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution.joins
+
+import org.apache.spark.SparkException
+import org.apache.spark.broadcast.Broadcast
+import org.apache.spark.rdd.RDD
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.expressions.codegen._
+import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight,
BuildSide}
+import org.apache.spark.sql.catalyst.planning.RangeJoin
+import org.apache.spark.sql.catalyst.plans._
+import org.apache.spark.sql.catalyst.plans.physical.{BroadcastDistribution,
Distribution, Partitioning, UnspecifiedDistribution}
+import org.apache.spark.sql.execution.{CodegenSupport, ExplainUtils, SparkPlan}
+import org.apache.spark.sql.execution.metric.SQLMetrics
+
+/**
+ * Performs a range join on two tables. Supported types are inner, left outer,
right outer,
+ * left semi, and left anti. The preserved side is streamed: unmatched build
rows are not
+ * emitted, because the broadcast is copied to every partition. See
`ExtractRangeJoinKeys`
+ * for the exact set of join-condition shapes that are recognized as a range
join.
+ *
+ * The build side is an index. Point-in-range and interval overlap both build
an
+ * [[IntervalIndex]] and probe it with `overlapping`. A single inequality
builds a
+ * [[PointIndex]] and scans the side of the bound the build side holds. Keys
that
+ * map to one order-preserving long instead build [[LongIntervalIndex]] or
+ * [[LongPointIndex]], so their probes compare raw longs without boxing. Either
+ * probe returns a superset. A candidate is kept only when the original join
+ * condition is true, so inclusivity stays on the predicate. Heavy overlap can
+ * make one probe return almost every build row.
+ */
+case class BroadcastRangeJoinExec(
+ left: SparkPlan,
+ right: SparkPlan,
+ buildSide: BuildSide,
+ joinType: JoinType,
+ leftKeys: Seq[Expression],
+ rightKeys: Seq[Expression],
+ condition: Option[Expression],
+ rangeJoin: RangeJoin)
+ extends JoinCodegenSupport {
+
+ if (condition.isEmpty) {
+ throw SparkException.internalError("BroadcastRangeJoinExec requires a join
condition.")
+ }
+
+ override lazy val metrics = Map(
+ "numOutputRows" -> SQLMetrics.createMetric(sparkContext, "number of output
rows"))
+
+ override lazy val requiredChildDistribution: Seq[Distribution] = {
+ val mode = RangeRelationBroadcastMode(
+ BindReferences.bindReferences(buildKeys, buildOutput),
+ RangeIndexKind(rangeJoin))
+ buildSide match {
+ case BuildLeft =>
+ BroadcastDistribution(mode) :: UnspecifiedDistribution :: Nil
+ case BuildRight =>
+ UnspecifiedDistribution :: BroadcastDistribution(mode) :: Nil
+ }
+ }
+
+ override def output: Seq[Attribute] = {
+ joinType match {
+ case _: InnerLike =>
+ left.output ++ right.output
+ case LeftOuter =>
+ left.output ++ right.output.map(_.withNullability(true))
+ case RightOuter =>
+ left.output.map(_.withNullability(true)) ++ right.output
+ case LeftSemi | LeftAnti =>
+ left.output
+ case t =>
+ throw SparkException.internalError(s"BroadcastRangeJoin should not
take $t as the JoinType")
+ }
+ }
+
+ // The preserved side is streamed, so its partitioning and ordering pass
through.
+ override def outputPartitioning: Partitioning =
streamedPlan.outputPartitioning
+
+ override def outputOrdering: Seq[SortOrder] = streamedPlan.outputOrdering
+
+ override def simpleStringWithNodeId(): String = {
+ val opId = ExplainUtils.getOpId(this)
+ s"$nodeName $joinType $buildSide ($opId)".trim
+ }
+
+ private[this] lazy val buildPlan = if (buildSide == BuildLeft) left else
right
+ private[this] lazy val streamedPlan = if (buildSide == BuildLeft) right else
left
+
+ private[this] lazy val buildKeys = if (buildSide == BuildLeft) leftKeys else
rightKeys
+ private[this] lazy val streamedKeys = if (buildSide == BuildLeft) rightKeys
else leftKeys
+
+ @transient private[this] lazy val buildOutput = buildPlan.output
+ @transient private[this] lazy val streamedOutput = streamedPlan.output
+
+ // The point-index scan used to probe this join's build side. See
`PointScan.forJoin`.
+ private lazy val pointScan: PointScan = PointScan.forJoin(rangeJoin,
buildSide)
+
+ // Alias so `innerJoin`/`outerJoin`/`leftExistenceJoin` need no change
below; the class
+ // itself lives on the companion object (see there for why).
+ private type ProbeContext = BroadcastRangeJoinExec.ProbeContext
+
+ private def perPartition(
+ relation: Broadcast[RangeRelation])(
+ f: ProbeContext => Iterator[InternalRow]): RDD[InternalRow] = {
+ val boundStreamKeys = BindReferences.bindReferences(streamedKeys,
streamedOutput)
+ val schema = left.output ++ right.output
+ val numBuildFields = buildPlan.output.length
+ val join = rangeJoin
+ val side = buildSide
+ streamedPlan.execute().mapPartitionsInternal { stream =>
+ // relation.value is read here, inside the task closure, so each
executor fetches
+ // the index through the broadcast machinery (and shares it across the
partitions
+ // that closure runs over) instead of the driver baking the materialized
index
+ // into the closure that gets shipped out. PointScan.forJoin is passed
by-name into
+ // makeCandidateFn, so it is only evaluated - and only throws for a
non-partial-range
+ // join - inside the point-index branches that actually need it, same as
before. It
+ // reads `join`/`side` (copied above), not `rangeJoin`/`buildSide`
directly, so this
+ // closure captures only those small values instead of the enclosing
plan node.
+ val candidateFn = BroadcastRangeJoinExec.makeCandidateFn(
+ relation.value, boundStreamKeys, PointScan.forJoin(join, side))
+ f(new BroadcastRangeJoinExec.ProbeContext(
+ stream,
+ Predicate.create(condition.get, schema).eval _,
Review Comment:
**Nit (P3):** `condition` here is a field of the case class, so reading it
inside the `mapPartitionsInternal` lambda captures `this`. The whole
`BroadcastRangeJoinExec`, children included, then ships with the task. That
contradicts the comment just above, which says copying `join`/`side` into
locals keeps the closure to those small values. It also means placing
`makeCandidateFn` and `ProbeContext` on the companion object, which is
documented for the same reason, buys nothing. `doExecute`'s
`mapPartitionsWithIndexInternal` closure captures `this` the same way, through
the instance method `genResultProjection`. Answers are unaffected, and
`BroadcastNestedLoopJoinExec` captures its node the same way, so either fix
works. One is to bind `condition.get`, plus the projection's output and input
attributes, to locals before both closures so the comments hold. The other is
to drop the claims.
**Verification:**
- **Inspection:** No interpreted-path task closure in the range join
operator references a member of the plan node, so serializing it does not pull
in the node or its children.
- **Behavior:** Every supported join type still returns the nested-loop
answer with whole-stage codegen disabled.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/joins/RangeRelation.scala:
##########
@@ -0,0 +1,1249 @@
+/*
+ * 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,
InterpretedProjection}
+import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildSide}
+import org.apache.spark.sql.catalyst.planning.{IntervalOverlapJoin,
PartialRangeJoin, PointInRangeJoin, RangeJoin}
+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(rangeJoin: RangeJoin): Boolean = RangeIndexKind(rangeJoin) ==
this
+}
+
+private[execution] case object IntervalIndexKind extends RangeIndexKind
+
+private[execution] case object PointIndexKind extends RangeIndexKind
+
+private[execution] object RangeIndexKind {
+ def apply(rangeJoin: RangeJoin): RangeIndexKind = rangeJoin 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. `RangePredicate`
+ * gives both sides of a range predicate one type after analysis, 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 well
under the
Review Comment:
**Nit (P3):** A few comments here and in the tests don't match the code:
- Here: `math.min(n, 1L << 20)` equals the hint for any hint up to 2^20
rows, so 'capped well under the hint itself' holds only for larger hints, where
the cap is a fixed 2^20.
- The long-index section comment says these structures 'carry no Spark
types' once built. But `LongIntervalIndex` and `LongPointIndex` keep `_keyType:
DataType` (serialized, and used by `toLong`) and `rows: Array[InternalRow]`.
Only their comparisons avoid Spark types.
- The `RangeIndexSuite` test for the `coreLow`/`coreHigh` empty-set
identities says an all-covering window 'correctly fails both comparisons'. With
`coreHigh = Long.MaxValue` and `coreLow = Long.MinValue`, `probeLow <= coreHigh
&& probeHigh >= coreLow` holds for every window. Empty indexes are handled by
the `lows.isEmpty` early return instead.
- The probe-coverage comment in `RangeJoinSQLSuite` and the round-trip
reader in `RangeIndexSuite` refer to a `candidates()` branch or dispatch.
`ProbeContext.candidates` only delegates; the per-index dispatch is
`BroadcastRangeJoinExec.makeCandidateFn`.
Could these comments describe what the code does?
**Verification:**
- **Inspection:** The initialCapacity doc matches its min-with-2^20 bound,
the long-index comment matches the retained fields, the empty-identity comment
matches the early return for empty indexes, and the dispatch comments name the
method that dispatches by index class.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/joins/RangeIndexSuite.scala:
##########
@@ -0,0 +1,945 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution.joins
+
+import java.io.{ByteArrayInputStream, ByteArrayOutputStream,
ObjectInputStream, ObjectOutputStream}
+
+import org.apache.spark.SparkFunSuite
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.BoundReference
+import org.apache.spark.sql.catalyst.types.PhysicalDataType
+import org.apache.spark.sql.types._
+import org.apache.spark.unsafe.Platform
+import org.apache.spark.unsafe.types.UTF8String
+
+/**
+ * Unit tests for the range-join indexes: the boxed pair ([[IntervalIndex]],
+ * [[PointIndex]]), the raw-long pair ([[LongIntervalIndex]],
[[LongPointIndex]]),
+ * and the broadcast-mode routing that picks between the two families.
+ * Operator-level behavior is covered by `RangeJoinSuite`.
+ */
+class RangeIndexSuite extends SparkFunSuite {
+
+ private def row(id: Int): InternalRow = InternalRow(id)
+
+ private def span(lo: Any, hi: Any, id: Int): (Any, Any, InternalRow) = (lo,
hi, row(id))
+
+ private def buildIntervals(spans: Seq[(Any, Any, InternalRow)]):
IntervalIndex =
+ buildIntervals(IntegerType, spans)
+
+ /**
+ * Builds the boxed interval index through the production build. The row
+ * layout puts the asserted id at ordinal 0, exactly where a real build row
+ * carries its columns; the keys bind to ordinals 1 and 2.
+ */
+ private def buildIntervals(
+ keyType: DataType,
+ spans: Seq[(Any, Any, InternalRow)]): IntervalIndex = {
+ val keys = Seq(
+ BoundReference(1, keyType, nullable = true),
+ BoundReference(2, keyType, nullable = true))
+ IntervalIndex.build(keyType, keys,
+ spans.iterator.map(t => InternalRow(t._3.getInt(0), t._1, t._2)),
+ Some(spans.size.toLong))
+ }
+
+ /** The point-index build with the same row layout: id at 0, key at 1. */
+ private def buildPoints(keyType: DataType, keyed: Seq[(Any, InternalRow)]):
PointIndex = {
+ val key = Seq(BoundReference(1, keyType, nullable = true))
+ PointIndex.build(keyType, key,
+ keyed.iterator.map(t => InternalRow(t._2.getInt(0), t._1)),
+ Some(keyed.size.toLong))
+ }
+
+ private def ids(rows: Iterator[InternalRow]): List[Int] =
rows.map(_.getInt(0)).toList
+
+ private def probeIds(index: IntervalIndex, low: Any, high: Any): List[Int] =
+ ids(index.overlapping(low, high))
+
+ /**
+ * An iterator that hands back the same mutable row on every `next()` call,
each field
+ * set to the current index. Mirrors how Spark's row iterators commonly
reuse a single
+ * buffer (e.g. `UnsafeRow` deserialization) instead of allocating a fresh
row per record.
+ */
+ private def reusingRowIterator(
+ fields: Int,
+ n: Int): (org.apache.spark.sql.catalyst.expressions.GenericInternalRow,
+ Iterator[InternalRow]) = {
+ val shared = new
org.apache.spark.sql.catalyst.expressions.GenericInternalRow(fields)
+ val iter = new Iterator[InternalRow] {
+ private var i = 0
+ override def hasNext: Boolean = i < n
+ override def next(): InternalRow = {
+ var c = 0
+ while (c < fields) { shared.update(c, i); c += 1 }
+ i += 1
+ shared
+ }
+ }
+ (shared, iter)
+ }
+
+ private def reusingByteBufferStringIterator(
+ n: Int): (Array[Byte], Iterator[InternalRow]) = {
+ val sharedBytes = new Array[Byte](4)
+ val sharedRow = new
org.apache.spark.sql.catalyst.expressions.GenericInternalRow(2)
+ val iter = new Iterator[InternalRow] {
+ private var i = 0
+ override def hasNext: Boolean = i < n
+ override def next(): InternalRow = {
+ val s = f"$i%04d"
+ java.lang.System.arraycopy(s.getBytes("UTF-8"), 0, sharedBytes, 0, 4)
+ val utf8 = UTF8String.fromAddress(sharedBytes,
Platform.BYTE_ARRAY_OFFSET, 4)
+ sharedRow.update(0, utf8)
+ sharedRow.update(1, i)
+ i += 1
+ sharedRow
+ }
+ }
+ (sharedBytes, iter)
+ }
+
+ private def javaRoundTrip(index: IntervalIndex): (Array[Byte],
IntervalIndex) = {
+ val bytes = new ByteArrayOutputStream()
+ val out = new ObjectOutputStream(bytes)
+ out.writeObject(index)
+ out.close()
+ val raw = bytes.toByteArray
+ val in = new ObjectInputStream(new ByteArrayInputStream(raw))
+ try {
+ (raw, in.readObject().asInstanceOf[IntervalIndex])
+ } finally {
+ in.close()
+ }
+ }
+
+ test("build keeps a point, orders an inverted interval, and drops a null
bound") {
+ val index = buildIntervals(Seq(
+ span(5, 5, 5),
+ span(1, 3, 1),
+ span(3, 1, 3),
+ (null, 1, row(1))))
+ // (3, 1) is stored as [1, 3]. The null bound is absent.
+ assert(ids(index.overlapping(2, 2)).sorted == List(1, 3))
+ assert(ids(index.overlapping(5, 5)) == List(5))
+ assert(ids(index.overlapping(4, 4)).isEmpty)
+ }
+
+ test("IntervalIndex.build defends against a build-side iterator reusing its
row buffer") {
+ // If build() captured keys or rows by reference instead of copying them,
every stored
+ // entry would collapse to the iterator's last value once the shared
buffer is mutated
+ // again after the build-side iterator is drained -- exactly what a
broadcast's row
+ // iterator does once BroadcastExchangeExec finishes feeding it to
RangeRelationBroadcastMode.
+ val n = 50
+ val buildKeys = Seq(
+ BoundReference(0, IntegerType, nullable = true),
+ BoundReference(1, IntegerType, nullable = true))
+ val (shared, rows) = reusingRowIterator(fields = 3, n)
+ val index = IntervalIndex.build(IntegerType, buildKeys, rows, Some(n))
+
+ // Simulate the caller reusing the buffer after the build-side iterator is
exhausted.
+ shared.update(0, -1)
+ shared.update(1, -1)
+ shared.update(2, -1)
+
+ (0 until n).foreach { i =>
+ val payloads = index.overlapping(i, i).map(_.getInt(2)).toList
+ assert(payloads == List(i), s"interval $i was corrupted by later reuse
of the row buffer")
+ }
+ }
+
+ test("IntervalIndex.build defends against string keys sharing a reused byte
buffer") {
+ val n = 50
+ val buildKeys = Seq(
+ BoundReference(0, StringType, nullable = true),
+ BoundReference(0, StringType, nullable = true))
+ val (sharedBytes, rows) = reusingByteBufferStringIterator(n)
+ val index = IntervalIndex.build(StringType, buildKeys, rows, Some(n))
+
+ java.util.Arrays.fill(sharedBytes, 'X'.toByte)
+
+ (0 until n).foreach { i =>
+ val key = UTF8String.fromString(f"$i%04d")
+ val matching = index.overlapping(key, key).toList
+ assert(matching.size == 1, s"interval $i was corrupted by shared buffer
reuse")
+ assert(matching.head.getInt(1) == i)
+ }
+ }
+
+ test("a one-key interval build owns a single copy of the point key") {
+ // One build key is one key used as both ends. The build must copy it once
+ // and serve both bounds from that copy, instead of copying twice.
+ val n = 50
+ val oneKey = Seq(BoundReference(0, StringType, nullable = true))
+ val (sharedBytes, rows) = reusingByteBufferStringIterator(n)
+ val index = IntervalIndex.build(StringType, oneKey, rows, Some(n))
+
+ java.util.Arrays.fill(sharedBytes, 'X'.toByte)
+
+ (0 until n).foreach { i =>
+ val key = UTF8String.fromString(f"$i%04d")
+ val matching = index.overlapping(key, key).toList
+ assert(matching.size == 1, s"interval $i was corrupted by shared buffer
reuse")
+ assert(matching.head.getInt(1) == i)
+ }
+
+ // SizeEstimator's identity tracking charges the shared bounds once, so the
+ // estimate stays below the same build made from two independently copied
keys.
+ val twoKeys = Seq(
+ BoundReference(0, StringType, nullable = true),
+ BoundReference(0, StringType, nullable = true))
+ val (_, moreRows) = reusingByteBufferStringIterator(n)
+ val twoCopy = IntervalIndex.build(StringType, twoKeys, moreRows, Some(n))
+ assert(index.estimatedSize < twoCopy.estimatedSize)
+ }
+
+ test("an inverted build interval reaches the windows the condition accepts")
{
+ // Int keys route to the long index, so this also pins swap normalization
there.
+ val index = RangeRelationBroadcastMode(
+ Seq(BoundReference(0, IntegerType, nullable = true),
+ BoundReference(1, IntegerType, nullable = true)),
+ IntervalIndexKind)
+ .transform(Array(InternalRow(3, 1), InternalRow(0, 5)))
+ .asInstanceOf[LongIntervalIndex]
+ assert(index.overlapping(0L, 4L).map(_.getInt(0)).toSet == Set(3, 0))
+ assert(index.overlapping(4L, 9L).map(_.getInt(0)).toList == List(0))
+ }
+
+ test("interval tree returns each overlapping row once") {
+ // Spans arrive out of low order: [10, 20] then [0, 5].
+ val unsorted = buildIntervals(Seq(span(10, 20, 0), span(0, 5, 1)))
+ assert(probeIds(unsorted, 3, 3) == List(1))
+ assert(probeIds(unsorted, 15, 15) == List(0))
+
+ // Probe 50 is a build key. Row 0 spans it; row 1 starts there.
+ val spanning = buildIntervals(Seq(span(0, 100, 0), span(50, 200, 1)))
+ assert(probeIds(spanning, 50, 50).toSet == Set(0, 1))
+ assert(probeIds(spanning, 25, 25) == List(0))
+ assert(probeIds(spanning, 150, 150) == List(1))
+
+ val duplicated = buildIntervals(Seq(span(0, 100, 0), span(0, 100, 1)))
+ assert(probeIds(duplicated, 50, 50).sorted == List(0, 1))
+
+ val nested = buildIntervals(Seq(
+ span(0, 100, 0), span(10, 20, 1), span(10, 50, 2), span(90, 110, 3)))
+ def assertOnce(low: Int, expected: Set[Int]): Unit = {
+ val actual = probeIds(nested, low, low)
+ assert(actual.toSet == expected, s"probe $low -> $actual")
+ assert(actual.size == expected.size, s"probe $low duplicated $actual")
+ }
+ assertOnce(0, Set(0))
+ assertOnce(10, Set(0, 1, 2))
+ assertOnce(15, Set(0, 1, 2))
+ assertOnce(30, Set(0, 2))
+ assertOnce(95, Set(0, 3))
+ assertOnce(105, Set(3))
+ assertOnce(200, Set.empty)
+
+ assert(probeIds(buildIntervals(Seq.empty), 1, 1).isEmpty)
+
+ val withNull = buildIntervals(Seq(span(0, 10, 0), span(5, 5, 1), (null, 4,
row(9))))
+ assert(probeIds(withNull, null, 5).isEmpty)
+ assert(probeIds(withNull, 5, null).isEmpty)
+ // The null bound is dropped, so two intervals remain and the estimate
covers both.
+ val one = buildIntervals(Seq(span(0, 10, 0)))
+ assert(withNull.estimatedSize > one.estimatedSize)
+
+ // At an interior endpoint, the range that ends and the one that starts
+ // are each a candidate once.
+ val chained = buildIntervals(Seq(span(0, 1, 0), span(1, 2, 1), span(2, 3,
2)))
+ assert(probeIds(chained, 0, 0) == List(0))
+ assert(probeIds(chained, 1, 1).sorted == List(0, 1))
+ assert(probeIds(chained, 2, 2).sorted == List(1, 2))
+ assert(probeIds(chained, 3, 3) == List(2))
+
+ // A probe whose endpoints are reversed is the closed window between them.
+ // [10, 25] meets [0, 30], [15, 40], and the point at 15.
+ val inverted = buildIntervals(Seq(span(0, 30, 0), span(15, 40, 1),
span(15, 15, 2)))
+ assert(probeIds(inverted, 25, 10).sorted == List(0, 1, 2))
+ assert(probeIds(inverted, 10, 20).sorted == List(0, 1, 2))
+ }
+
+ test("interval tree is correct under heavy build-side overlap") {
+ // Every interval nests inside the one before it, so a point in the middle
+ // hits about half of the rows. The tree still stores one node per row.
+ val n = 300
+ val nested = buildIntervals((0 until n).map(i => span(i, 2 * n - i, i)))
+
+ def expected(probe: Int): Set[Int] =
+ (0 until n).filter(i => i <= probe && probe <= 2 * n - i).toSet
+
+ Seq(0, 1, n / 2, n, 2 * n - 1, 2 * n).foreach { probe =>
+ assert(probeIds(nested, probe, probe).toSet == expected(probe), s"probe
$probe")
+ }
+ }
+
+ test("broadcast of heavily overlapping intervals stays linear") {
+ // [i, i + n] for i in 0 until n. Every interval contains n. The broadcast
+ // is the flat arrays, one slot per interval, so serialization stays
linear.
+ val n = 200
+ val index = buildIntervals((0 until n).map(i => span(i, i + n, i)))
+ val (raw, restored) = javaRoundTrip(index)
+ assert(raw.length < n * 1024, s"serialized ${raw.length} bytes")
+ val probe = n
+ assert(probeIds(restored, probe, probe).toSet == probeIds(index, probe,
probe).toSet)
+ assert(probeIds(restored, probe, probe).size == n)
+ }
+
+ test("estimatedSize depends on the interval count, not the overlap depth") {
+ val n = 200
+ val one = buildIntervals(Seq(span(0, 1, 0)))
+ val disjoint = buildIntervals((0 until n).map(i => span(2 * i, 2 * i + 1,
i)))
+ val overlap = buildIntervals((0 until n).map(i => span(i, i + n, i)))
+ val again = buildIntervals((0 until n).map(i => span(2 * i, 2 * i + 1, i)))
+ assert(disjoint.estimatedSize > one.estimatedSize)
+ assert(overlap.estimatedSize > one.estimatedSize)
+ assert(disjoint.estimatedSize == again.estimatedSize)
+ // Nesting stores no extra copy of the rows. Key objects differ, so the two
+ // indexes need not be equal, but overlap stays the same order of
magnitude.
+ assert(overlap.estimatedSize < disjoint.estimatedSize * 3)
+ }
+
+ test("estimatedSize includes distinct key objects") {
+ val ints = buildIntervals(Seq(span(0, 1, 0)))
+ val wide = UTF8String.fromString("x" * 64)
+ val strings = buildIntervals(StringType, Seq((wide, wide, row(0))))
+ assert(strings.estimatedSize > ints.estimatedSize)
+ }
+
+ test("interval tree treats NaN as its own point and Infinity as a bound") {
+ // PhysicalDoubleType's ordering (SQLOrderingUtil.compareDoubles) totally
orders
+ // -Inf < ... < +Inf < NaN, with NaN comparing equal only to itself. So
[1.0, +Inf]
+ // includes +Inf under inclusive probe, and a NaN key overlaps only a
probe that
+ // also lands exactly on NaN.
+ val index = buildIntervals(DoubleType, Seq(
+ (1.0, Double.PositiveInfinity, row(1)),
+ (Double.NegativeInfinity, -1.0, row(2)),
+ (Double.NaN, Double.NaN, row(3))))
+
+ assert(ids(index.overlapping(5.0, 5.0)) == List(1))
+ assert(ids(index.overlapping(Double.PositiveInfinity,
Double.PositiveInfinity)) == List(1))
+ assert(ids(index.overlapping(-5.0, -5.0)) == List(2))
+ assert(ids(index.overlapping(Double.NegativeInfinity,
Double.NegativeInfinity)) == List(2))
+ assert(ids(index.overlapping(Double.NaN, Double.NaN)) == List(3))
+ assert(ids(index.overlapping(0.0, 0.0)).isEmpty)
+ }
+
+ test("interval index merges keys that compare equal but are not equal") {
+ // Array equals is identity, and UTF8String equals is binary, so these keys
+ // compare equal under the index ordering without being equal. The tree
+ // orders them with that comparator, so a probe on either key returns both.
+ // Signed zero is treated as one key: the index ordering treats -0.0 and
0.0 as equal.
Review Comment:
**Nit (P3):** This line says signed zero is treated as one key, but nothing
below checks it. `assertMerged` runs only for `BinaryType` and `UNICODE_CI`,
and its `assert(leftKey != rightKey)` precondition would fail for boxed `-0.0`
and `0.0`, which Scala's `==` treats as equal. I couldn't find a test in
`RangeIndexSuite`, `RangeJoinSuite`, or `RangeJoinSQLSuite` that builds or
probes `-0.0` against `0.0`, and the routing test pins only `FloatType` with
`PointIndexKind`. Signed zero is part of why `RangeIndex.longKey` keeps
floating point on the boxed indexes. If double keys were ever routed to the
long family through a bit-preserving encoding, a probe of `-0.0` against `[0.0,
0.0]` would drop a row that the nested loop join returns, and every test would
still pass. Could you add a `-0.0`/`0.0` case for both index kinds and pin
`DoubleType` (both kinds) and `FloatType` interval routing to the boxed family?
Otherwise, could the claim here be dropped?
**Verification:**
- **Regression:** Routing double keys to the long family through a
bit-preserving encoding makes a signed-zero assertion fail.
- **Behavior:** A range join that probes -0.0 against 0.0 bounds returns the
nested-loop answer.
--
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]
