wangyum commented on code in PR #59027: URL: https://github.com/apache/spark/pull/59027#discussion_r4355921941
########## sql/core/src/main/scala/org/apache/spark/sql/execution/joins/RangeRelation.scala: ########## @@ -0,0 +1,827 @@ +/* + * 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 scala.collection.mutable.ArrayBuffer + +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{BoundReference, Expression, InterpretedProjection, UnsafeRow} +import org.apache.spark.sql.catalyst.optimizer.{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.unsafe.types.UTF8String +import org.apache.spark.util.KnownSizeEstimation + +/** + * 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 { + + /** Accessor reading a single field from `InternalRow` as a boxed value (or null). */ + def getValue(dt: DataType, ordinal: Int): InternalRow => Any = { + val accessor = InternalRow.getAccessor(dt) + (input: InternalRow) => accessor(input, ordinal) + } + + /** + * One key type's bridge to an order-preserving long: reading a projected build + * field, converting an evaluated streamed value, and turning generated Java for + * a typed key into generated Java for a long. The three agree, so a build key + * and a probe key of the same type become the same long. + */ + final case class LongKey( + field: (InternalRow, Int) => Long, + value: Any => Long, + code: String => String) + + /** + * An order-preserving long for a fixed-width key, or None when the type must + * stay on the generic boxed index: variable-length keys, and floating point, + * whose SQL equality for -0.0 and NaN has no bit-preserving order. 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(field: (InternalRow, Int) => Long, value: Any => Long): LongKey = + LongKey(field, value, v => s"((long) $v)") + dt match { + case _: BooleanType => + Some(LongKey( + (r, o) => if (r.getBoolean(o)) 1L else 0L, + v => if (v.asInstanceOf[Boolean]) 1L else 0L, + v => s"($v ? 1L : 0L)")) + case _: ByteType => + Some(key((r, o) => r.getByte(o).toLong, v => v.asInstanceOf[Byte].toLong)) + case _: ShortType => + Some(key((r, o) => r.getShort(o).toLong, v => v.asInstanceOf[Short].toLong)) + case _: IntegerType | _: DateType => + Some(key((r, o) => r.getInt(o).toLong, v => v.asInstanceOf[Int].toLong)) + case _: LongType | _: TimestampType | _: TimestampNTZType => + Some(key((r, o) => r.getLong(o), v => v.asInstanceOf[Long])) + case d: DecimalType if d.precision <= Decimal.MAX_LONG_DIGITS => + Some(LongKey( + (r, o) => r.getDecimal(o, d.precision, d.scale).toUnscaledLong, + 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. `RangePredicate` gives + * both sides of a range predicate one type after analysis, so testing the + * build keys also settles the streamed keys. + */ + private[execution] 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. */ + private[execution] def requiredLongKey(dt: DataType): LongKey = { + longKey(dt).getOrElse( + throw new IllegalStateException(s"$dt has no order-preserving long key.")) + } + + /** Estimated bytes of a row or key value, without reflection or SizeEstimator. */ + private[joins] def valueSize(v: Any): Long = v match { Review Comment: Fixed: estimatedSize is now one SizeEstimator pass over the real retained arrays. ########## sql/core/src/test/scala/org/apache/spark/sql/execution/joins/RangeIndexSuite.scala: ########## @@ -0,0 +1,825 @@ +/* + * 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 two range-join indexes. [[IntervalIndex]] is the interval tree + * used by point-in-range and interval overlap. [[PointIndex]] is the sorted array + * used by a single inequality. 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("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 already one key: Scala equality treats -0.0 and 0.0 as equal. + 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("a nested probe does not drop rows from the outer iterator") { + val index = buildIntervals((0 until 32).map(i => span(i, i + 10, i))) + val outer = index.overlapping(0, 100) + val seen = scala.collection.mutable.ArrayBuffer.empty[Int] + // The inner probe runs while the outer iterator still has rows, the same + // shape as two interval joins fused in one whole-stage. + while (outer.hasNext) { + seen += outer.next().getInt(0) + assert(ids(index.overlapping(0, 5)).sorted == (0 to 5).toList) + } + assert(seen.sorted == (0 until 32).toList) + } + + test("a point probe of disjoint intervals returns only the containing row") { + val n = 1000 + val index = buildIntervals((0 until n).map(i => span(i * 2, i * 2 + 1, i))) + assert(probeIds(index, 2, 2) == List(1)) + assert(probeIds(index, 3, 3) == List(1)) + assert(probeIds(index, 4, 4) == List(2)) + assert(probeIds(index, -1, -1).isEmpty) + assert(probeIds(index, 100000, 100000).isEmpty) + } + + test("point index answers both sides of a bound, including equals") { + // Unsorted, with two rows on key 5. Equals stay, in input order. + val keyed = Array[(Any, InternalRow)]( + (8, row(8)), + (1, row(1)), + (5, row(5)), + (5, row(50)), + (3, row(3)), + (null, row(-1))) + val index = buildPoints(IntegerType, keyed.toSeq) + assert(ids(index.upTo(5)) == List(1, 3, 5, 50)) + assert(ids(index.upTo(4)) == List(1, 3)) + assert(ids(index.upTo(0)).isEmpty) + assert(ids(index.upTo(8)) == List(1, 3, 5, 50, 8)) + assert(ids(index.from(5)) == List(5, 50, 8)) + assert(ids(index.from(6)) == List(8)) + assert(ids(index.from(9)).isEmpty) + assert(ids(index.from(1)) == List(1, 3, 5, 50, 8)) + assert(index.upTo(null).isEmpty) + assert(index.from(null).isEmpty) + + val empty = buildPoints(IntegerType, Seq.empty) + assert(empty.upTo(1).isEmpty) + assert(empty.from(1).isEmpty) + // SizeEstimator charges the empty array shell, and the five kept points on top of it. Review Comment: Fixed: PointIndex.estimatedSize also uses one SizeEstimator pass; comment updated. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/joins/RangeRelation.scala: ########## @@ -0,0 +1,827 @@ +/* + * 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 scala.collection.mutable.ArrayBuffer + +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{BoundReference, Expression, InterpretedProjection, UnsafeRow} +import org.apache.spark.sql.catalyst.optimizer.{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.unsafe.types.UTF8String +import org.apache.spark.util.KnownSizeEstimation + +/** + * 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 { + + /** Accessor reading a single field from `InternalRow` as a boxed value (or null). */ + def getValue(dt: DataType, ordinal: Int): InternalRow => Any = { + val accessor = InternalRow.getAccessor(dt) + (input: InternalRow) => accessor(input, ordinal) + } + + /** + * One key type's bridge to an order-preserving long: reading a projected build + * field, converting an evaluated streamed value, and turning generated Java for + * a typed key into generated Java for a long. The three agree, so a build key + * and a probe key of the same type become the same long. + */ + final case class LongKey( + field: (InternalRow, Int) => Long, + value: Any => Long, + code: String => String) + + /** + * An order-preserving long for a fixed-width key, or None when the type must + * stay on the generic boxed index: variable-length keys, and floating point, + * whose SQL equality for -0.0 and NaN has no bit-preserving order. 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(field: (InternalRow, Int) => Long, value: Any => Long): LongKey = + LongKey(field, value, v => s"((long) $v)") + dt match { + case _: BooleanType => + Some(LongKey( + (r, o) => if (r.getBoolean(o)) 1L else 0L, + v => if (v.asInstanceOf[Boolean]) 1L else 0L, + v => s"($v ? 1L : 0L)")) + case _: ByteType => + Some(key((r, o) => r.getByte(o).toLong, v => v.asInstanceOf[Byte].toLong)) + case _: ShortType => + Some(key((r, o) => r.getShort(o).toLong, v => v.asInstanceOf[Short].toLong)) + case _: IntegerType | _: DateType => + Some(key((r, o) => r.getInt(o).toLong, v => v.asInstanceOf[Int].toLong)) + case _: LongType | _: TimestampType | _: TimestampNTZType => + Some(key((r, o) => r.getLong(o), v => v.asInstanceOf[Long])) + case d: DecimalType if d.precision <= Decimal.MAX_LONG_DIGITS => + Some(LongKey( + (r, o) => r.getDecimal(o, d.precision, d.scale).toUnscaledLong, + 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. `RangePredicate` gives + * both sides of a range predicate one type after analysis, so testing the + * build keys also settles the streamed keys. + */ + private[execution] 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. */ + private[execution] def requiredLongKey(dt: DataType): LongKey = { + longKey(dt).getOrElse( + throw new IllegalStateException(s"$dt has no order-preserving long key.")) + } + + /** Estimated bytes of a row or key value, without reflection or SizeEstimator. */ + private[joins] def valueSize(v: Any): Long = v match { + case u: UnsafeRow => u.getSizeInBytes.toLong + case s: UTF8String => s.numBytes().toLong + case b: Array[Byte] => b.length.toLong + case d: Decimal if d.precision > Decimal.MAX_LONG_DIGITS => 24L + d.precision.toLong + case _ => 8L + } +} + +/** + * 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 `RangeIndex.getValue` reads, compared through +// `BoxedKeyOrdering`. Floating point, strings, binary, and decimal keys stay +// here: their comparison has no order-preserving single-long form. + +/** + * Augmented interval tree of build-side intervals. Entries are (low, high, row) + * tuples sorted by low bound, with a parallel row array feeding the scan path + * and `maxHighs(mid)` recording the maximum high bound in the span rooted at + * mid. Subtrees whose max high is below the probe window, or whose low starts + * after it, are pruned during search, and a window meeting the common + * intersection skips the walk entirely. + */ +private[joins] class IntervalIndex private ( + val keyType: DataType, + private[this] val intervals: Array[(Any, Any, InternalRow)], Review Comment: Fixed: PointIndex keeps keys and rows in separate arrays, no duplicated aliasing. ########## sql/core/src/test/scala/org/apache/spark/sql/execution/joins/RangeJoinSQLSuite.scala: ########## @@ -0,0 +1,660 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.execution.joins + +import org.apache.spark.SparkException +import org.apache.spark.sql.QueryTest +import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, IntervalOverlapJoin, LessPartialRangeJoin} +import org.apache.spark.sql.catalyst.plans.{Inner, LeftAnti, LeftOuter, LeftSemi} +import org.apache.spark.sql.execution._ +import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.test.SharedSparkSession + +/** + * Planner and end-to-end coverage for broadcast range join. + * + * Which condition is recognized: `ExtractRangeJoinKeysSuite`. + * What the index returns: `RangeIndexSuite`. + * What the operator returns: `RangeJoinSuite`. + * This suite checks that the planner chooses the operator, and that the SQL + * result matches broadcast nested loop join where that comparison is the point. + */ +class RangeJoinSQLSuite extends QueryTest with SharedSparkSession with AdaptiveSparkPlanHelper { + + private val pointCond = + "s.ip_add_int >= l.BEGIN_IP2LONG AND s.ip_add_int <= l.END_IP2LONG" + + private def findRangeJoin(plan: SparkPlan): Seq[BroadcastRangeJoinExec] = { + collect(plan) { case j: BroadcastRangeJoinExec => j } + } + + /** `n` range joins sit in one whole-stage, not in nested stages or input adapters. */ + private def oneCodegenStageOwns(plan: SparkPlan, n: Int): Boolean = { + exists(plan) { + case w: WholeStageCodegenExec => + def count(p: SparkPlan): Int = p match { + case _: InputAdapter | _: WholeStageCodegenExec => 0 + case j: BroadcastRangeJoinExec => 1 + j.children.map(count).sum + case other => other.children.map(count).sum + } + count(w.child) == n + case _ => false + } + } + + /** The range join is compiled into a whole-stage, not hidden under an input adapter. */ + private def rangeJoinIsCodegen(plan: SparkPlan): Boolean = { + exists(plan) { + case w: WholeStageCodegenExec => + def contains(p: SparkPlan): Boolean = p match { + case _: InputAdapter => false + case _: BroadcastRangeJoinExec => true + case other => other.children.exists(contains) + } + contains(w.child) + case _ => false + } + } + + private def findBroadcastNestedLoopJoin(plan: SparkPlan): Seq[BroadcastNestedLoopJoinExec] = { + collect(plan) { case j: BroadcastNestedLoopJoinExec => j } + } + + private def findCartesianProduct(plan: SparkPlan): Seq[CartesianProductExec] = { + collect(plan) { case c: CartesianProductExec => c } + } + + private def withRangeJoin(extra: (String, String)*)(f: => Unit): Unit = { + withSQLConf((Seq(SQLConf.BROADCAST_RANGE_JOIN_ENABLED.key -> "true") ++ extra): _*)(f) + } + + private def setupIpLookupViews(): Unit = { + sql( + """ + |CREATE OR REPLACE TEMP VIEW signin_ip(id, ip_add_int) AS VALUES + | ('1', 1467630666L), + | ('2', 1186929717L), + | ('3', 3501544352L), + | ('31', 3501544320L), + | ('32', 3501544383L), + | ('33', 3501544330L), + | ('34', 3501544390L), + | ('4', 1626243295L), + | ('41', 1626243072L), + | ('42', 1626243327L), + | ('5', 1206355995L), + | ('5', 1206355995L), + | ('6', CAST(NULL AS BIGINT)), + | ('7', 0L) + |""".stripMargin) + sql( + """ + |CREATE OR REPLACE TEMP VIEW ip_lookup( + | BEGIN_IP2LONG, END_IP2LONG, TWO_CHAR_CNTRY_CD, zip_id) AS VALUES + | (1467630592L, 1467630847L, 'de', 2380520L), + | (1186929664L, 1186930175L, 'us', 6723936L), + | (1186929664L, 1186930175L, 'us', 6723936L), + | (3501544320L, 3501544383L, 'ca', 12804631L), + | (3501544330L, 3501544390L, 'ca', 12804631001L), + | (1626243072L, 1626243327L, 'us', 8849753L), + | (1206355968L, 1206356223L, 'us', 10138801L), + | (CAST(NULL AS BIGINT), 1206356223L, 'us', 10138801L), + | (1206355968L, CAST(NULL AS BIGINT), 'us', 10138801L), + | (CAST(NULL AS BIGINT), CAST(NULL AS BIGINT), 'us', 10138801L) + |""".stripMargin) + } + + private def setupIntervalViews(): Unit = { + sql( + """ + |CREATE OR REPLACE TEMP VIEW intervals_a(lo, hi) AS VALUES + | (-1, 0), (0, 1), (0, 2), (1, 5) + |""".stripMargin) + sql( + """ + |CREATE OR REPLACE TEMP VIEW intervals_b(lo, hi) AS VALUES + | (-2, -1), (-4, -2), (1, 3), (5, 7) + |""".stripMargin) + } + + private def setupDecimalViews(): Unit = { + sql( + """ + |CREATE OR REPLACE TEMP VIEW dec_ranges(lo, hi) AS VALUES + | (CAST(0.00 AS DECIMAL(10, 2)), CAST(2.50 AS DECIMAL(10, 2))), + | (CAST(3.00 AS DECIMAL(10, 2)), CAST(5.00 AS DECIMAL(10, 2))) + |""".stripMargin) + sql( + """ + |CREATE OR REPLACE TEMP VIEW dec_points(p) AS VALUES + | (CAST(1.25 AS DECIMAL(10, 2))), + | (CAST(3.00 AS DECIMAL(10, 2))), + | (CAST(9.00 AS DECIMAL(10, 2))) + |""".stripMargin) + } + + private def setupCollatedStrings(): Unit = { + // 'M' and 'm' compare equal under UNICODE_CI and differ as bytes. + sql( + """ + |CREATE OR REPLACE TEMP VIEW str_ranges(lo, hi) AS VALUES + | (CAST('a' AS STRING COLLATE UNICODE_CI), CAST('a' AS STRING COLLATE UNICODE_CI)), + | (CAST('b' AS STRING COLLATE UNICODE_CI), CAST('b' AS STRING COLLATE UNICODE_CI)), + | (CAST('c' AS STRING COLLATE UNICODE_CI), CAST('M' AS STRING COLLATE UNICODE_CI)), + | (CAST('m' AS STRING COLLATE UNICODE_CI), CAST('z' AS STRING COLLATE UNICODE_CI)) + |""".stripMargin) + sql( + """ + |CREATE OR REPLACE TEMP VIEW str_points(s) AS VALUES + | (CAST('M' AS STRING COLLATE UNICODE_CI)), + | (CAST('m' AS STRING COLLATE UNICODE_CI)) + |""".stripMargin) + } + + private def setupDateAndTimestampViews(): Unit = { + sql( + """ + |CREATE OR REPLACE TEMP VIEW dt_ranges(lo, hi) AS VALUES + | (DATE'2024-01-01', DATE'2024-01-31'), + | (DATE'2024-06-01', DATE'2024-06-30') + |""".stripMargin) + sql( + """ + |CREATE OR REPLACE TEMP VIEW dt_points(d) AS VALUES + | (DATE'2024-01-15'), + | (DATE'2024-05-01'), + | (DATE'2024-06-15') + |""".stripMargin) + sql( + """ + |CREATE OR REPLACE TEMP VIEW ts_ranges(lo, hi) AS VALUES + | (TIMESTAMP'2024-01-01 00:00:00', TIMESTAMP'2024-01-01 12:00:00') + |""".stripMargin) + sql( + """ + |CREATE OR REPLACE TEMP VIEW ts_points(t) AS VALUES + | (TIMESTAMP'2024-01-01 06:00:00'), + | (TIMESTAMP'2024-01-01 18:00:00') + |""".stripMargin) + } + + test("range join is disabled by default and falls back to nested loop join") { + setupIpLookupViews() + val df = sql(s"SELECT s.id, l.zip_id FROM signin_ip s JOIN ip_lookup l ON $pointCond") + assert(findRangeJoin(df.queryExecution.executedPlan).isEmpty) + assert(findBroadcastNestedLoopJoin(df.queryExecution.executedPlan).size == 1) + } + + test("point-in-range matches the nested loop result") { + setupIpLookupViews() + val query = s"SELECT s.id, l.zip_id FROM signin_ip s JOIN ip_lookup l ON $pointCond" + val expected = sql(query).collect() + withRangeJoin() { + val df = sql(query) + assert(findRangeJoin(df.queryExecution.executedPlan).size == 1) + checkAnswer(df, expected) + } + } + + test("BETWEEN is planned as a point-in-range join") { + setupIpLookupViews() + val query = + "SELECT s.id, l.zip_id FROM signin_ip s JOIN ip_lookup l " + + "ON s.ip_add_int BETWEEN l.BEGIN_IP2LONG AND l.END_IP2LONG" + val expected = sql(query).collect() + withRangeJoin(SQLConf.ALWAYS_INLINE_COMMON_EXPR.key -> "false") { + val df = sql(query) + val planned = findRangeJoin(df.queryExecution.executedPlan) + assert(planned.size == 1) + // BETWEEN lowers to one point expression. The other side keeps (low, high). + val keyCounts = Set(planned.head.leftKeys.length, planned.head.rightKeys.length) + assert(keyCounts == Set(1, 2)) + checkAnswer(df, expected) + } + } + + test("hints select range join, nested loop, or cartesian") { + setupIpLookupViews() + setupIntervalViews() + val expected = sql( + s"SELECT s.id, l.zip_id FROM signin_ip s JOIN ip_lookup l ON $pointCond").collect() + withRangeJoin() { + val broadcast = sql( + "SELECT /*+ BROADCAST(l) */ s.id, l.zip_id FROM signin_ip s JOIN ip_lookup l " + + s"ON $pointCond") + val joins = findRangeJoin(broadcast.queryExecution.executedPlan) + assert(joins.size == 1 && joins.head.buildSide == BuildRight) + checkAnswer(broadcast, expected) + + val dualHintInner = sql( + "SELECT /*+ BROADCAST(s), BROADCAST(l) */ s.id, l.zip_id " + + s"FROM signin_ip s JOIN ip_lookup l ON $pointCond") + val dualJoins = findRangeJoin(dualHintInner.queryExecution.executedPlan) + assert(dualJoins.size == 1 && dualJoins.head.buildSide == BuildRight) + checkAnswer(dualHintInner, expected) + + val dualHintLeftOuter = sql( + "SELECT /*+ BROADCAST(s), BROADCAST(l) */ s.id, l.zip_id " + + s"FROM signin_ip s LEFT JOIN ip_lookup l ON $pointCond") + val dualOuterJoins = findRangeJoin(dualHintLeftOuter.queryExecution.executedPlan) + assert(dualOuterJoins.size == 1 && dualOuterJoins.head.buildSide == BuildRight) + + val unsupported = sql( + "SELECT /*+ BROADCAST(a) */ a.lo, b.hi FROM intervals_a a LEFT JOIN intervals_b b " + + "ON a.lo < b.hi") + val loops = findBroadcastNestedLoopJoin(unsupported.queryExecution.executedPlan) + assert(findRangeJoin(unsupported.queryExecution.executedPlan).isEmpty) + assert(loops.size == 1 && loops.head.buildSide == BuildLeft) + + val cartesian = sql( + "SELECT /*+ SHUFFLE_REPLICATE_NL(s) */ s.id, l.zip_id " + + s"FROM signin_ip s JOIN ip_lookup l ON $pointCond") + val plan = cartesian.queryExecution.executedPlan + assert(findRangeJoin(plan).isEmpty) + assert(findCartesianProduct(plan).size == 1) + } + } + + test("supported join types match nested loop; full outer does not") { + setupIpLookupViews() + Seq( + s"SELECT s.id, l.zip_id FROM signin_ip s LEFT JOIN ip_lookup l ON $pointCond", + s"SELECT s.id, l.zip_id FROM signin_ip s RIGHT JOIN ip_lookup l ON $pointCond", + s"SELECT s.id FROM signin_ip s LEFT SEMI JOIN ip_lookup l ON $pointCond", + s"SELECT s.id FROM signin_ip s LEFT ANTI JOIN ip_lookup l ON $pointCond" + ).foreach { query => + val expected = sql(query).collect() + withRangeJoin() { + val df = sql(query) + assert(findRangeJoin(df.queryExecution.executedPlan).size == 1, query) + checkAnswer(df, expected) + } + } + val fullOuter = + s"SELECT s.id, l.zip_id FROM signin_ip s FULL OUTER JOIN ip_lookup l ON $pointCond" + withRangeJoin() { + assert(findRangeJoin(sql(fullOuter).queryExecution.executedPlan).isEmpty) + } + } + + test("partial range left outer broadcasts the right side") { + setupIntervalViews() + // No b.lo is strictly greater than a.hi (5), so the left row (1, 5) is kept with nulls. + val query = + "SELECT a.lo, a.hi, b.lo, b.hi FROM intervals_a a LEFT JOIN intervals_b b " + + "ON a.hi < b.lo" + val expected = sql(query).collect() + assert(expected.exists(_.anyNull)) + withRangeJoin() { + val df = sql(query) + val planned = findRangeJoin(df.queryExecution.executedPlan) + assert(planned.size == 1) + assert(planned.head.joinType == LeftOuter) + assert(planned.head.buildSide == BuildRight) + assert(planned.head.rangeJoin == LessPartialRangeJoin) + checkAnswer(df, expected) + } + } + + test("a residual predicate is still evaluated") { + sql("CREATE OR REPLACE TEMP VIEW complex_a(id, lo, hi) AS VALUES (1, 0, 10), (2, 0, 10)") + sql("CREATE OR REPLACE TEMP VIEW complex_b(id, point) AS VALUES (1, 5), (2, 5)") + // Both ranges cover both points, so the range predicate matches all 4 pairs. + // a.id <> b.id is not an equi-join key and filters that down to 2. + val query = + "SELECT a.id, b.id FROM complex_a a JOIN complex_b b " + + "ON b.point >= a.lo AND b.point <= a.hi AND a.id <> b.id" + val expected = sql(query).collect() + assert(expected.length == 2) + withRangeJoin() { + val df = sql(query) + assert(findRangeJoin(df.queryExecution.executedPlan).size == 1) + checkAnswer(df, expected) + } + } + + test("interval overlap is planned with two keys on each side") { + setupIntervalViews() + val query = + "SELECT a.lo, a.hi, b.lo, b.hi FROM intervals_a a JOIN intervals_b b " + + "ON a.lo < b.hi AND b.lo < a.hi" + val expected = sql(query).collect() + withRangeJoin() { + val df = sql(query) + val planned = findRangeJoin(df.queryExecution.executedPlan) + assert(planned.size == 1) + assert(planned.head.rangeJoin == IntervalOverlapJoin) + assert(planned.head.leftKeys.length == 2) + assert(planned.head.rightKeys.length == 2) + checkAnswer(df, expected) + } + } + + test("interval overlap with an inverted interval matches the nested loop result") { + sql("CREATE OR REPLACE TEMP VIEW inverted_a(lo, hi) AS VALUES (0, 5), (1, 2), (-4, 10)") + sql("CREATE OR REPLACE TEMP VIEW inverted_b(lo, hi) AS VALUES (3, 1), (7, 9)") + val query = + "SELECT a.lo, a.hi, b.lo, b.hi FROM inverted_a a " + + "JOIN inverted_b b ON a.lo < b.hi AND b.lo < a.hi" + // (3, 1) accepts rows with a.lo < 1 and a.hi > 3, so it has to reach the result. + val expected = sql(query).collect() + assert(expected.exists(r => r.getInt(2) == 3 && r.getInt(3) == 1), expected.toSeq) + assert(expected.length == 3, expected.toSeq) + // Broadcast each side in turn so the inverted interval lands on the index once. + Seq("/*+ BROADCAST(a) */", "/*+ BROADCAST(b) */").foreach { hint => + val hinted = query.replace("SELECT ", s"SELECT $hint ") + withRangeJoin() { + val df = sql(hinted) + assert(findRangeJoin(df.queryExecution.executedPlan).size == 1, hinted) + checkAnswer(df, expected) + } + } + } + + test("a scalar subquery in the condition is planned") { + setupIpLookupViews() + // The extra conjunct names both sides, so it stays on the join. The subquery + // has to be planned from that condition. + val query = + "SELECT s.id, l.zip_id FROM signin_ip s JOIN ip_lookup l " + + s"ON $pointCond " + + "AND s.ip_add_int + l.zip_id > (SELECT min(ip_add_int) FROM signin_ip)" + val expected = sql(query).collect() + withRangeJoin() { + val df = sql(query) + val planned = findRangeJoin(df.queryExecution.executedPlan) + assert(planned.size == 1) + val cond = planned.head.condition.get + val subquery = cond.collectFirst { case s: ScalarSubquery => s.plan } + assert(subquery.exists(plan => + plan.isInstanceOf[SubqueryExec] || plan.isInstanceOf[ReusedSubqueryExec]), cond) + checkAnswer(df, expected) + } + } + + test("collated strings match the nested loop result on either build side") { + setupCollatedStrings() + val cond = "p.s >= r.lo AND p.s <= r.hi" + Seq( + (s"SELECT /*+ BROADCAST(r) */ p.s, r.lo, r.hi FROM str_points p JOIN str_ranges r ON $cond", + BuildRight), + (s"SELECT /*+ BROADCAST(p) */ p.s, r.lo, r.hi FROM str_points p JOIN str_ranges r ON $cond", + BuildLeft) + ).foreach { case (query, buildSide) => + val expected = sql(query).collect() + // Each point matches the interval ending at 'M' and the one starting at 'm'. + assert(expected.length == 4, query) + withRangeJoin() { + val df = sql(query) + val joins = findRangeJoin(df.queryExecution.executedPlan) + assert(joins.size == 1 && joins.head.buildSide == buildSide, query) + checkAnswer(df, expected) + } + } + } + + test("date and timestamp range joins match the nested loop result") { + setupDateAndTimestampViews() + val dateQuery = + "SELECT p.d, r.lo, r.hi FROM dt_points p JOIN dt_ranges r ON p.d BETWEEN r.lo AND r.hi" + val tsQuery = + "SELECT p.t, r.lo, r.hi FROM ts_points p JOIN ts_ranges r ON p.t BETWEEN r.lo AND r.hi" + Seq(dateQuery, tsQuery).foreach { query => + val expected = sql(query).collect() + withRangeJoin() { + val df = sql(query) + assert(findRangeJoin(df.queryExecution.executedPlan).size == 1, query) + checkAnswer(df, expected) + } + } + } + + test("boolean range joins match the nested loop result, including a null point") { + sql( + """ + |CREATE OR REPLACE TEMP VIEW bool_ranges(lo, hi) AS VALUES + | (false, true), + | (true, true) + |""".stripMargin) + sql( + """ + |CREATE OR REPLACE TEMP VIEW bool_points(p) AS VALUES + | (false), + | (true), + | (CAST(NULL AS BOOLEAN)) + |""".stripMargin) + val query = "SELECT p.p, r.lo, r.hi FROM bool_points p JOIN bool_ranges r " + + "ON p.p >= r.lo AND p.p <= r.hi" + val expected = sql(query).collect() + withRangeJoin() { + val df = sql(query) + assert(findRangeJoin(df.queryExecution.executedPlan).size == 1, query) + checkAnswer(df, expected) + } + } + + test("whole-stage codegen covers each range-join probe") { Review Comment: Added a boxed DOUBLE/STRING two-key interval test exercising the PointIndex branch. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/joins/RangeRelation.scala: ########## @@ -0,0 +1,827 @@ +/* + * 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 scala.collection.mutable.ArrayBuffer + +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{BoundReference, Expression, InterpretedProjection, UnsafeRow} +import org.apache.spark.sql.catalyst.optimizer.{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.unsafe.types.UTF8String +import org.apache.spark.util.KnownSizeEstimation + +/** + * 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 { + + /** Accessor reading a single field from `InternalRow` as a boxed value (or null). */ + def getValue(dt: DataType, ordinal: Int): InternalRow => Any = { + val accessor = InternalRow.getAccessor(dt) + (input: InternalRow) => accessor(input, ordinal) + } + + /** + * One key type's bridge to an order-preserving long: reading a projected build + * field, converting an evaluated streamed value, and turning generated Java for + * a typed key into generated Java for a long. The three agree, so a build key + * and a probe key of the same type become the same long. + */ + final case class LongKey( + field: (InternalRow, Int) => Long, + value: Any => Long, + code: String => String) + + /** + * An order-preserving long for a fixed-width key, or None when the type must + * stay on the generic boxed index: variable-length keys, and floating point, + * whose SQL equality for -0.0 and NaN has no bit-preserving order. 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(field: (InternalRow, Int) => Long, value: Any => Long): LongKey = + LongKey(field, value, v => s"((long) $v)") + dt match { + case _: BooleanType => + Some(LongKey( + (r, o) => if (r.getBoolean(o)) 1L else 0L, + v => if (v.asInstanceOf[Boolean]) 1L else 0L, + v => s"($v ? 1L : 0L)")) + case _: ByteType => + Some(key((r, o) => r.getByte(o).toLong, v => v.asInstanceOf[Byte].toLong)) + case _: ShortType => + Some(key((r, o) => r.getShort(o).toLong, v => v.asInstanceOf[Short].toLong)) + case _: IntegerType | _: DateType => + Some(key((r, o) => r.getInt(o).toLong, v => v.asInstanceOf[Int].toLong)) + case _: LongType | _: TimestampType | _: TimestampNTZType => + Some(key((r, o) => r.getLong(o), v => v.asInstanceOf[Long])) + case d: DecimalType if d.precision <= Decimal.MAX_LONG_DIGITS => + Some(LongKey( + (r, o) => r.getDecimal(o, d.precision, d.scale).toUnscaledLong, + 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. `RangePredicate` gives + * both sides of a range predicate one type after analysis, so testing the + * build keys also settles the streamed keys. + */ + private[execution] 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. */ + private[execution] def requiredLongKey(dt: DataType): LongKey = { + longKey(dt).getOrElse( + throw new IllegalStateException(s"$dt has no order-preserving long key.")) + } + + /** Estimated bytes of a row or key value, without reflection or SizeEstimator. */ + private[joins] def valueSize(v: Any): Long = v match { + case u: UnsafeRow => u.getSizeInBytes.toLong + case s: UTF8String => s.numBytes().toLong + case b: Array[Byte] => b.length.toLong + case d: Decimal if d.precision > Decimal.MAX_LONG_DIGITS => 24L + d.precision.toLong + case _ => 8L + } +} + +/** + * 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 `RangeIndex.getValue` reads, compared through +// `BoxedKeyOrdering`. Floating point, strings, binary, and decimal keys stay +// here: their comparison has no order-preserving single-long form. + +/** + * Augmented interval tree of build-side intervals. Entries are (low, high, row) + * tuples sorted by low bound, with a parallel row array feeding the scan path + * and `maxHighs(mid)` recording the maximum high bound in the span rooted at + * mid. Subtrees whose max high is below the probe window, or whose low starts + * after it, are pruned during search, and a window meeting the common + * intersection skips the walk entirely. + */ +private[joins] class IntervalIndex private ( + val keyType: DataType, + private[this] val intervals: Array[(Any, Any, InternalRow)], + private[this] val maxHighs: Array[Any], + private[this] val rows: Array[InternalRow]) + extends BoxedKeyOrdering { + + // 1-D Helly: the stored intervals are null-free with lo <= hi, so they share + // the intersection [coreLow, coreHigh] iff coreLow <= coreHigh. `intervals` + // is sorted by low, so the last low bounds the core from below. + private[this] val coreLow: Any = + if (intervals.isEmpty) null else intervals(intervals.length - 1)._1 + private[this] val coreHigh: Any = + if (intervals.isEmpty) null + else intervals.iterator.map(_._2).reduce[Any](ordering.min(_, _)) + private[this] val hasCommonCore: Boolean = + coreLow != null && ordering.lteq(coreLow, coreHigh) + + override def estimatedSize: Long = { + var size = 24L + intervals.length * 24L + var i = 0 + while (i < intervals.length) { + val (lo, hi, row) = intervals(i) + size += RangeIndex.valueSize(lo) + RangeIndex.valueSize(hi) + RangeIndex.valueSize(row) + i += 1 + } + size + } + + def overlapping(low: Any, high: Any): Iterator[InternalRow] = { + if (intervals.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 meets the core overlaps every interval, and the walk below + // would return every row anyway: return them without visiting the tree. + if (hasCommonCore && ord.lteq(probeLow, coreHigh) && ord.gteq(probeHigh, coreLow)) { + return new RowSlice(rows, 0, rows.length) + } + new OverlapIterator(intervals, rows, maxHighs, ord, probeLow, probeHigh) + } +} + +private[joins] object IntervalIndex { + + def build( + keyType: DataType, + buildKeys: Seq[Expression], + rows: Iterator[InternalRow], + sizeHint: Option[Long]): IntervalIndex = { + val initialCapacity = sizeHint.map(n => math.min(n, Int.MaxValue).toInt).getOrElse(64) + val buffer = new ArrayBuffer[(Any, Any, InternalRow)](initialCapacity) + val projection = new InterpretedProjection(buildKeys) + val lowOf = RangeIndex.getValue(buildKeys.head.dataType, 0) + val highOf = + if (buildKeys.length == 1) lowOf else RangeIndex.getValue(buildKeys(1).dataType, 1) + val ordering = PhysicalDataType.ordering(keyType) + + while (rows.hasNext) { + val r = rows.next() + val p = projection(r) + normalizedInterval(ordering, lowOf(p), highOf(p), r, buffer) + } + + sortAndBuild(keyType, ordering, buffer.toArray) + } + + /** Drops a null-bounded interval; otherwise orders (lo, hi) and copies the key and row. */ + private def normalizedInterval( + ordering: Ordering[Any], + rawLo: Any, + rawHi: Any, + r: InternalRow, + buffer: ArrayBuffer[(Any, Any, InternalRow)]): Unit = { + if (rawLo != null && rawHi != null) { + val lo = InternalRow.copyValue(rawLo) + val hi = InternalRow.copyValue(rawHi) Review Comment: Fixed: a single build key is copied once and reused for both bounds, like LongIntervalIndex. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala: ########## @@ -378,30 +396,42 @@ abstract class SparkStrategies extends QueryPlanner[SparkPlan] { if (canBuildBroadcastLeft(joinType)) BuildLeft else BuildRight } - def createBroadcastNLJoin(onlyLookingAtHint: Boolean) = { - val buildLeft = if (onlyLookingAtHint) { - hintToBroadcastLeft(hint) - } else { - canBroadcastBySize(left, conf) && !hintToNotBroadcastAndReplicateLeft(hint) - } - - val buildRight = if (onlyLookingAtHint) { - hintToBroadcastRight(hint) - } else { - canBroadcastBySize(right, conf) && !hintToNotBroadcastAndReplicateRight(hint) + /** + * Which side to broadcast. A hint-only pass honors a BROADCAST hint. + * Otherwise a side qualifies when it is small enough and is not marked + * NO_BROADCAST_AND_REPLICATION. `allowed` drops a side this operator + * cannot build. Both sides qualify: the smaller one. Review Comment: Fixed: doc now says both-qualify returns desiredBuildSide, matching the code. -- 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]
