wangyum commented on code in PR #59027: URL: https://github.com/apache/spark/pull/59027#discussion_r4355923090
########## 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) + val (normLo, normHi) = if (ordering.lteq(lo, hi)) (lo, hi) else (hi, lo) + buffer += ((normLo, normHi, r.copy())) + } + } + + private def sortAndBuild( + keyType: DataType, + ordering: Ordering[Any], + valid: Array[(Any, Any, InternalRow)]): IntervalIndex = { + java.util.Arrays.sort(valid, Ordering.by[(Any, Any, InternalRow), Any](_._1)(ordering)) + val length = valid.length + val maxHighs = new Array[Any](length) + + def fillMaxHigh(start: Int, end: Int): Any = { + if (start >= end) return null + val mid = (start + end) >>> 1 + var max = valid(mid)._2 + val left = fillMaxHigh(start, mid) + val right = fillMaxHigh(mid + 1, end) + if (left != null) max = ordering.max(max, left) + if (right != null) max = ordering.max(max, right) + maxHighs(mid) = max + max + } + + fillMaxHigh(0, length) + val rows = new Array[InternalRow](length) + var i = 0 + while (i < length) { + rows(i) = valid(i)._3 + i += 1 + } + new IntervalIndex(keyType, valid, maxHighs, rows) + } +} + +private[joins] class OverlapIterator( + intervals: Array[(Any, Any, InternalRow)], + rows: Array[InternalRow], + maxHighs: Array[Any], + ord: Ordering[Any], + probeLow: Any, + probeHigh: Any) extends Iterator[InternalRow] { + + private[this] val depth = 32 - java.lang.Integer.numberOfLeadingZeros(intervals.length) + private[this] var stackStart = new Array[Int](math.max(8, depth * 3)) + private[this] var stackEnd = new Array[Int](math.max(8, depth * 3)) + private[this] var stackTop = 0 + private[joins] var steps = 0 + private[this] var nextRow: InternalRow = null + + push(0, intervals.length) + + private[joins] def stackCapacity: Int = stackStart.length + + private def push(s: Int, e: Int): Unit = { + if (stackTop >= stackStart.length) { + stackStart = java.util.Arrays.copyOf(stackStart, stackStart.length * 2) + stackEnd = java.util.Arrays.copyOf(stackEnd, stackEnd.length * 2) + } + stackStart(stackTop) = s + stackEnd(stackTop) = e + stackTop += 1 + } + + private def advance(): InternalRow = { + while (stackTop > 0) { + steps += 1 + stackTop -= 1 + val start = stackStart(stackTop) + val end = stackEnd(stackTop) + if (start < end) { + val mid = (start + end) >>> 1 + if (ord.gteq(maxHighs(mid), probeLow)) { + if (ord.lteq(intervals(mid)._1, probeHigh)) { + push(mid + 1, end) + } + push(mid, -1) + push(start, mid) + } + } else if (end == -1) { + val lo = intervals(start)._1 + val hi = intervals(start)._2 + if (ord.gteq(hi, probeLow) && ord.lteq(lo, probeHigh)) { + return rows(start) + } + } + } + null + } + + override def hasNext: Boolean = { + if (nextRow != null) true + else { + nextRow = advance() + nextRow != null + } + } + + override def next(): InternalRow = { + if (!hasNext) throw new NoSuchElementException("empty iterator") + val r = nextRow + nextRow = null + r + } +} + +/** + * Sorted build-side points for single cross-side inequalities (`<`, `<=`, `>`, `>=`). + */ +private[joins] class PointIndex private ( + val keyType: DataType, + private[this] val points: Array[(Any, InternalRow)], + private[this] val rows: Array[InternalRow]) + extends BoxedKeyOrdering { + + private[this] val length = points.length + + override def estimatedSize: Long = { + var size = 16L + points.length * 16L + var i = 0 + while (i < points.length) { + val (key, row) = points(i) + size += RangeIndex.valueSize(key) + RangeIndex.valueSize(row) + i += 1 + } + size + } + + /** Points whose key is less than or equal to `value`. */ + def upTo(value: Any): Iterator[InternalRow] = + if (value == null) Iterator.empty else slice(0, searchBound(value, skipEquals = true)) + + /** Points whose key is greater than or equal to `value`. */ + def from(value: Any): Iterator[InternalRow] = + if (value == null) Iterator.empty else slice(searchBound(value, skipEquals = false), length) + + private def searchBound(value: Any, skipEquals: Boolean): Int = { + val ord = ordering + var lo = 0 + var hi = length + while (lo < hi) { + val mid = (lo + hi) >>> 1 + val c = ord.compare(points(mid)._1, value) + if (c < 0 || (skipEquals && c == 0)) lo = mid + 1 else hi = mid + } + lo + } + + private def slice(from: Int, until: Int): Iterator[InternalRow] = new RowSlice(rows, from, until) +} + +private[joins] object PointIndex { + + def build( + keyType: DataType, + buildKeys: Seq[Expression], + rows: Iterator[InternalRow], + sizeHint: Option[Long]): PointIndex = { + val initialCapacity = sizeHint.map(n => math.min(n, Int.MaxValue).toInt).getOrElse(64) + val buffer = new ArrayBuffer[(Any, InternalRow)](initialCapacity) + val projection = new InterpretedProjection(buildKeys) + val lowOf = RangeIndex.getValue(buildKeys.head.dataType, 0) + val ordering = PhysicalDataType.ordering(keyType) + + while (rows.hasNext) { + val r = rows.next() + val p = projection(r) + val rawKey = lowOf(p) + if (rawKey != null) { + val key = InternalRow.copyValue(rawKey) + buffer += ((key, r.copy())) + } + } + + sortAndBuild(keyType, ordering, buffer.toArray) + } + + private def sortAndBuild( + keyType: DataType, + ordering: Ordering[Any], + present: Array[(Any, InternalRow)]): PointIndex = { + java.util.Arrays.sort(present, Ordering.by[(Any, InternalRow), Any](_._1)(ordering)) + val rows = new Array[InternalRow](present.length) + var i = 0 + while (i < present.length) { + rows(i) = present(i)._2 + i += 1 + } + new PointIndex(keyType, present, rows) + } +} + +// ===== Long-specialized indexes for fixed-width keys ======================= +// +// Indexes for fixed-width keys, stored as order-preserving longs converted at the +// boundary by `RangeIndex.longKey`. Once built, these structures carry no Spark +// types: the tree walk and the binary searches compare raw longs, and generated +// probe code passes long keys without boxing. Types that cannot be one long, or +// whose SQL equality has no bit mapping (floating point), stay on `IntervalIndex` +// and `PointIndex`. `RangeRelationBroadcastMode` picks between the two families once, +// from the build keys' data types. + +private[joins] class LongIntervalIndex private ( + val keyType: DataType, + private[this] val lows: Array[Long], + private[this] val highs: Array[Long], + private[this] val maxHighs: Array[Long], + private[this] val rows: Array[InternalRow]) + extends LongKeyConversion { + + // 1-D Helly on raw keys: the stored intervals are null-free with lo <= hi, + // so they share the intersection [coreLow, coreHigh] iff coreLow <= coreHigh. + // A window that meets the core overlaps every interval, and the walk below + // would return every row anyway. + private[this] val coreLow: Long = if (lows.isEmpty) Long.MinValue else lows(lows.length - 1) + private[this] val coreHigh: Long = { + var min = Long.MaxValue + var i = 0 + while (i < highs.length) { + if (highs(i) < min) min = highs(i) + i += 1 + } + if (highs.isEmpty) Long.MinValue else min + } + private[this] val hasCommonCore: Boolean = !lows.isEmpty && coreLow <= coreHigh + + override def estimatedSize: Long = { + var size = 24L + lows.length * 24L + var i = 0 + while (i < rows.length) { + size += RangeIndex.valueSize(rows(i)) + i += 1 + } + size + } + + def overlapping(low: Long, high: Long): Iterator[InternalRow] = { + if (lows.isEmpty) return Iterator.empty + var probeLow = low + var probeHigh = high + if (probeLow > probeHigh) { + val t = probeLow + probeLow = probeHigh + probeHigh = t + } + if (hasCommonCore && probeLow <= coreHigh && probeHigh >= coreLow) { Review Comment: Fixed: dropped hasCommonCore; the all-covering check is exact without it. ########## sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/planning/patterns.scala: ########## @@ -421,6 +422,155 @@ object ExtractSingleColumnNullAwareAntiJoin extends JoinSelectionHelper with Pre } } +/** + * Finds joins whose condition is a range predicate usable by a broadcast range join. + * + * Tries three shapes, each capturing more of the condition than the last, and keeps Review Comment: Fixed: reworded so "captures more" attaches to the two-predicate shapes, not the fallback. ########## 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 Review Comment: Fixed: doc/test now follow longKey's rule; float/double called out as boxed. ########## 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) + val (normLo, normHi) = if (ordering.lteq(lo, hi)) (lo, hi) else (hi, lo) + buffer += ((normLo, normHi, r.copy())) + } + } + + private def sortAndBuild( + keyType: DataType, + ordering: Ordering[Any], + valid: Array[(Any, Any, InternalRow)]): IntervalIndex = { + java.util.Arrays.sort(valid, Ordering.by[(Any, Any, InternalRow), Any](_._1)(ordering)) + val length = valid.length + val maxHighs = new Array[Any](length) + + def fillMaxHigh(start: Int, end: Int): Any = { + if (start >= end) return null + val mid = (start + end) >>> 1 + var max = valid(mid)._2 + val left = fillMaxHigh(start, mid) + val right = fillMaxHigh(mid + 1, end) + if (left != null) max = ordering.max(max, left) + if (right != null) max = ordering.max(max, right) + maxHighs(mid) = max + max + } + + fillMaxHigh(0, length) + val rows = new Array[InternalRow](length) + var i = 0 + while (i < length) { + rows(i) = valid(i)._3 + i += 1 + } + new IntervalIndex(keyType, valid, maxHighs, rows) + } +} + +private[joins] class OverlapIterator( + intervals: Array[(Any, Any, InternalRow)], + rows: Array[InternalRow], + maxHighs: Array[Any], + ord: Ordering[Any], + probeLow: Any, + probeHigh: Any) extends Iterator[InternalRow] { + + private[this] val depth = 32 - java.lang.Integer.numberOfLeadingZeros(intervals.length) + private[this] var stackStart = new Array[Int](math.max(8, depth * 3)) + private[this] var stackEnd = new Array[Int](math.max(8, depth * 3)) + private[this] var stackTop = 0 + private[joins] var steps = 0 + private[this] var nextRow: InternalRow = null + + push(0, intervals.length) + + private[joins] def stackCapacity: Int = stackStart.length + + private def push(s: Int, e: Int): Unit = { + if (stackTop >= stackStart.length) { + stackStart = java.util.Arrays.copyOf(stackStart, stackStart.length * 2) + stackEnd = java.util.Arrays.copyOf(stackEnd, stackEnd.length * 2) + } + stackStart(stackTop) = s + stackEnd(stackTop) = e + stackTop += 1 + } + + private def advance(): InternalRow = { + while (stackTop > 0) { + steps += 1 + stackTop -= 1 + val start = stackStart(stackTop) + val end = stackEnd(stackTop) + if (start < end) { + val mid = (start + end) >>> 1 + if (ord.gteq(maxHighs(mid), probeLow)) { + if (ord.lteq(intervals(mid)._1, probeHigh)) { + push(mid + 1, end) + } + push(mid, -1) + push(start, mid) + } + } else if (end == -1) { + val lo = intervals(start)._1 + val hi = intervals(start)._2 + if (ord.gteq(hi, probeLow) && ord.lteq(lo, probeHigh)) { + return rows(start) + } + } + } + null + } + + override def hasNext: Boolean = { + if (nextRow != null) true + else { + nextRow = advance() + nextRow != null + } + } + + override def next(): InternalRow = { + if (!hasNext) throw new NoSuchElementException("empty iterator") + val r = nextRow + nextRow = null + r + } +} + +/** + * Sorted build-side points for single cross-side inequalities (`<`, `<=`, `>`, `>=`). + */ +private[joins] class PointIndex private ( + val keyType: DataType, + private[this] val points: Array[(Any, InternalRow)], + private[this] val rows: Array[InternalRow]) + extends BoxedKeyOrdering { + + private[this] val length = points.length + + override def estimatedSize: Long = { + var size = 16L + points.length * 16L + var i = 0 + while (i < points.length) { + val (key, row) = points(i) + size += RangeIndex.valueSize(key) + RangeIndex.valueSize(row) + i += 1 + } + size + } + + /** Points whose key is less than or equal to `value`. */ + def upTo(value: Any): Iterator[InternalRow] = + if (value == null) Iterator.empty else slice(0, searchBound(value, skipEquals = true)) + + /** Points whose key is greater than or equal to `value`. */ + def from(value: Any): Iterator[InternalRow] = + if (value == null) Iterator.empty else slice(searchBound(value, skipEquals = false), length) + + private def searchBound(value: Any, skipEquals: Boolean): Int = { + val ord = ordering + var lo = 0 + var hi = length + while (lo < hi) { + val mid = (lo + hi) >>> 1 + val c = ord.compare(points(mid)._1, value) + if (c < 0 || (skipEquals && c == 0)) lo = mid + 1 else hi = mid + } + lo + } + + private def slice(from: Int, until: Int): Iterator[InternalRow] = new RowSlice(rows, from, until) +} + +private[joins] object PointIndex { + + def build( + keyType: DataType, + buildKeys: Seq[Expression], + rows: Iterator[InternalRow], + sizeHint: Option[Long]): PointIndex = { + val initialCapacity = sizeHint.map(n => math.min(n, Int.MaxValue).toInt).getOrElse(64) + val buffer = new ArrayBuffer[(Any, InternalRow)](initialCapacity) + val projection = new InterpretedProjection(buildKeys) + val lowOf = RangeIndex.getValue(buildKeys.head.dataType, 0) + val ordering = PhysicalDataType.ordering(keyType) + + while (rows.hasNext) { + val r = rows.next() + val p = projection(r) + val rawKey = lowOf(p) + if (rawKey != null) { + val key = InternalRow.copyValue(rawKey) + buffer += ((key, r.copy())) + } + } + + sortAndBuild(keyType, ordering, buffer.toArray) + } + + private def sortAndBuild( + keyType: DataType, + ordering: Ordering[Any], + present: Array[(Any, InternalRow)]): PointIndex = { + java.util.Arrays.sort(present, Ordering.by[(Any, InternalRow), Any](_._1)(ordering)) + val rows = new Array[InternalRow](present.length) + var i = 0 + while (i < present.length) { + rows(i) = present(i)._2 + i += 1 + } + new PointIndex(keyType, present, rows) + } +} + +// ===== Long-specialized indexes for fixed-width keys ======================= +// +// Indexes for fixed-width keys, stored as order-preserving longs converted at the +// boundary by `RangeIndex.longKey`. Once built, these structures carry no Spark +// types: the tree walk and the binary searches compare raw longs, and generated +// probe code passes long keys without boxing. Types that cannot be one long, or +// whose SQL equality has no bit mapping (floating point), stay on `IntervalIndex` +// and `PointIndex`. `RangeRelationBroadcastMode` picks between the two families once, +// from the build keys' data types. + +private[joins] class LongIntervalIndex private ( + val keyType: DataType, + private[this] val lows: Array[Long], + private[this] val highs: Array[Long], + private[this] val maxHighs: Array[Long], + private[this] val rows: Array[InternalRow]) + extends LongKeyConversion { + + // 1-D Helly on raw keys: the stored intervals are null-free with lo <= hi, + // so they share the intersection [coreLow, coreHigh] iff coreLow <= coreHigh. + // A window that meets the core overlaps every interval, and the walk below + // would return every row anyway. + private[this] val coreLow: Long = if (lows.isEmpty) Long.MinValue else lows(lows.length - 1) + private[this] val coreHigh: Long = { + var min = Long.MaxValue + var i = 0 + while (i < highs.length) { + if (highs(i) < min) min = highs(i) + i += 1 + } + if (highs.isEmpty) Long.MinValue else min + } + private[this] val hasCommonCore: Boolean = !lows.isEmpty && coreLow <= coreHigh + + override def estimatedSize: Long = { + var size = 24L + lows.length * 24L + var i = 0 + while (i < rows.length) { + size += RangeIndex.valueSize(rows(i)) + i += 1 + } + size + } + + def overlapping(low: Long, high: Long): Iterator[InternalRow] = { + if (lows.isEmpty) return Iterator.empty + var probeLow = low + var probeHigh = high + if (probeLow > probeHigh) { + val t = probeLow + probeLow = probeHigh + probeHigh = t + } + if (hasCommonCore && probeLow <= coreHigh && probeHigh >= coreLow) { + return new RowSlice(rows, 0, rows.length) + } + new LongOverlapIterator(lows, highs, rows, maxHighs, probeLow, probeHigh) + } +} + +private[joins] object LongIntervalIndex { + + def build( + keyType: DataType, + buildKeys: Seq[Expression], + rows: Iterator[InternalRow], + sizeHint: Option[Long]): LongIntervalIndex = { + val toLongs: Seq[(InternalRow, Int) => Long] = + buildKeys.map(k => RangeIndex.requiredLongKey(k.dataType).field) + val initialCapacity = sizeHint.map(n => math.min(n, Int.MaxValue).toInt).getOrElse(64) + val valid = new ArrayBuffer[(Long, Long, InternalRow)](initialCapacity) + val projection = new InterpretedProjection(buildKeys) + + while (rows.hasNext) { + val r = rows.next() + val p = projection(r) + if (!p.isNullAt(0) && (buildKeys.length == 1 || !p.isNullAt(1))) { + // A point build key is one key used as both ends, so it is read once. + var lo = toLongs.head(p, 0) + var hi = if (buildKeys.length == 1) lo else toLongs(1)(p, 1) + if (lo > hi) { + val t = lo + lo = hi + hi = t + } + valid += ((lo, hi, r.copy())) + } + } + + val staged = valid.toArray + java.util.Arrays.sort(staged, Ordering.by[(Long, Long, InternalRow), Long](_._1)) + val length = staged.length + val lows = new Array[Long](length) + val highs = new Array[Long](length) + val rowArr = new Array[InternalRow](length) + var i = 0 + while (i < length) { + lows(i) = staged(i)._1 + highs(i) = staged(i)._2 + rowArr(i) = staged(i)._3 + i += 1 + } + + val maxHighs = new Array[Long](length) + + def fillMaxHigh(start: Int, end: Int): Long = { + if (start >= end) return Long.MinValue + val mid = (start + end) >>> 1 + var max = highs(mid) + val left = fillMaxHigh(start, mid) + val right = fillMaxHigh(mid + 1, end) + if (left > max) max = left + if (right > max) max = right + maxHighs(mid) = max + max + } + + fillMaxHigh(0, length) + new LongIntervalIndex(keyType, lows, highs, maxHighs, rowArr) + } +} + +private[joins] class LongOverlapIterator( + lows: Array[Long], + highs: Array[Long], + rows: Array[InternalRow], + maxHighs: Array[Long], + probeLow: Long, + probeHigh: Long) extends Iterator[InternalRow] { + + private[this] val depth = 32 - java.lang.Integer.numberOfLeadingZeros(lows.length) + private[this] var stackStart = new Array[Int](math.max(8, depth * 3)) + private[this] var stackEnd = new Array[Int](math.max(8, depth * 3)) + private[this] var stackTop = 0 + private[joins] var steps = 0 Review Comment: Added a staggered LongOverlapIterator laziness test pinning steps/stackCapacity. ########## 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. + assert(empty.estimatedSize > 0) + assert(index.estimatedSize > empty.estimatedSize) + } + + test("PointIndex.build defends against a build-side iterator reusing its row buffer") { + // Same hazard as the IntervalIndex case above: a point index built straight off a + // reused row buffer must copy both the key and the row, not just reference them. + val n = 50 + val buildKeys = Seq(BoundReference(0, IntegerType, nullable = true)) + val (shared, rows) = reusingRowIterator(fields = 2, n) + val index = PointIndex.build(IntegerType, buildKeys, rows, Some(n)) + + shared.update(0, -1) + shared.update(1, -1) + + (0 until n).foreach { i => + val payload = index.from(i).next().getInt(1) + assert(payload == i, s"point $i was corrupted by later reuse of the row buffer") + } + } + + test("PointIndex.build defends against string keys sharing a reused byte buffer") { + val n = 50 + val buildKeys = Seq(BoundReference(0, StringType, nullable = true)) + val (sharedBytes, rows) = reusingByteBufferStringIterator(n) + val index = PointIndex.build(StringType, buildKeys, rows, Some(n)) + + java.util.Arrays.fill(sharedBytes, 'X'.toByte) + + (0 until n).foreach { i => + val key = UTF8String.fromString(f"$i%04d") + val matching = index.from(key).take(1).toList + assert(matching.nonEmpty && matching.head.getInt(1) == i, + s"point $i was corrupted by shared buffer reuse") + } + } + + test("point index orders NaN above +Infinity and -Infinity below every finite value") { + // Total order: -Inf < -1.0 < 0.0 < 1.0 < +Inf < NaN, with NaN the unique maximum. + val keyed = Array[(Any, InternalRow)]( + (Double.NaN, row(90)), + (Double.PositiveInfinity, row(80)), + (1.0, row(1)), + (Double.NegativeInfinity, row(-80)), + (-1.0, row(-1)), + (0.0, row(0))) + val index = buildPoints(DoubleType, keyed.toSeq) + assert(ids(index.upTo(0.0)) == List(-80, -1, 0)) + assert(ids(index.from(1.0)) == List(1, 80, 90)) + // NaN is the unique maximum: only a NaN key is at or above it. + assert(ids(index.from(Double.NaN)) == List(90)) + assert(ids(index.upTo(Double.NaN)) == List(-80, -1, 0, 1, 80, 90)) + // +Infinity is below NaN but above every finite value. + assert(ids(index.upTo(Double.PositiveInfinity)) == List(-80, -1, 0, 1, 80)) + assert(ids(index.from(Double.NegativeInfinity)) == List(-80, -1, 0, 1, 80, 90)) + } + + test("upTo/from matches a brute-force reference across randomized point sets") { + // PointIndex's counterpart to "overlapping matches a brute-force reference + // across randomized interval sets": cross-checks the binary-search bound + // scan against an independent O(n) scan, across many random key sets and + // probe values, including duplicate keys. + val rand = new scala.util.Random(42) + val bound = 50 + for (_ <- 0 until 30) { + val n = 1 + rand.nextInt(300) + val raw = (0 until n).map(i => (rand.nextInt(bound), i)) + val index = buildPoints(IntegerType, raw.map { case (k, i) => (k, row(i)) }) + + for (_ <- 0 until 20) { + val probe = rand.nextInt(bound) + val expectedUpTo = raw.collect { case (k, i) if k <= probe => i }.toSet + val expectedFrom = raw.collect { case (k, i) if k >= probe => i }.toSet + assert(ids(index.upTo(probe)).toSet == expectedUpTo, + s"n=$n probe=$probe expected=$expectedUpTo") + assert(ids(index.from(probe)).toSet == expectedFrom, + s"n=$n probe=$probe expected=$expectedFrom") + } + } + } + + test("indexes round-trip through Java and strict Kryo serialization") { + val serializers = Seq( + "Java" -> new org.apache.spark.serializer.JavaSerializer( + new org.apache.spark.SparkConf()).newInstance(), + "Kryo" -> { + val conf = new org.apache.spark.SparkConf() + .set(org.apache.spark.internal.config.Kryo.KRYO_REGISTRATION_REQUIRED, true) + new org.apache.spark.serializer.KryoSerializer(conf).newInstance() + }) + + // (key type, lo, hi, probe) covering each supported key family: fixed + // width, floating point, decimal, binary, and collated strings. The + // collated probe mixes cases, so a restored index only answers it when + // its collation ordering survived the round trip. CharType/VarcharType + // carry a StringConstraint (FixedLength/MaxLength) that Kryo must have + // registered separately from plain StringType. + val cases: Seq[(DataType, Any, Any, Any)] = Seq( + (ByteType, 0.toByte, 10.toByte, 5.toByte), + (ShortType, 0.toShort, 10.toShort, 5.toShort), + (IntegerType, 0, 10, 5), + (LongType, 0L, 10L, 5L), + (FloatType, 0.0f, 10.0f, 5.0f), + (DoubleType, 0.0, 10.0, 5.0), + (DecimalType(10, 2), Decimal(0.0), Decimal(10.0), Decimal(5.0)), + (StringType("UTF8_LCASE"), + UTF8String.fromString("abc"), UTF8String.fromString("zzz"), UTF8String.fromString("MMM")), + (CharType(3), + UTF8String.fromString("abc"), UTF8String.fromString("zzz"), UTF8String.fromString("mmm")), + (VarcharType(10), + UTF8String.fromString("abc"), UTF8String.fromString("zzz"), UTF8String.fromString("mmm")), + (BinaryType, Array[Byte](0), Array[Byte](10), Array[Byte](5)), + (DateType, 100, 200, 150), + (TimestampType, 1000L, 2000L, 1500L), + (TimestampNTZType, 1000L, 2000L, 1500L), + (BooleanType, false, true, true)) + + def roundTrip[A <: AnyRef: scala.reflect.ClassTag]( + ser: org.apache.spark.serializer.SerializerInstance, index: A): A = + ser.deserialize[A](ser.serialize(index)) + + // One reader for the sealed family, mirroring candidates() dispatch: + // the interval kinds answer a window, the point kinds answer a + // lower-bound scan, and every relation this test builds answers at + // row ordinal 0. + def answer(rel: RangeRelation, low: Any, high: Any): List[Int] = rel match { + case r: LongIntervalIndex => ids(r.overlapping(r.toLong(low), r.toLong(high))) + case r: IntervalIndex => ids(r.overlapping(low, high)) + case r: LongPointIndex => ids(r.from(r.toLong(low))) + case r: PointIndex => ids(r.from(low)) + } + + cases.foreach { case (dt, lo, hi, probe) => + // One straight span and one inverted span that normalizes to the same + // shape, plus the family the broadcast build routes the type to. Rows + // carry the asserted id at ordinal 0 and the keys at 1 and 2, so one + // reader answers every family. + val interval = buildIntervals(dt, Seq((lo, hi, row(5)), (hi, lo, row(6)))) Review Comment: Added a disjoint-interval round-trip test forcing a real walk after restore. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/LogicalQueryStageStrategy.scala: ########## @@ -58,6 +63,81 @@ object LogicalQueryStageStrategy extends Strategy { case _ => false } + /** + * Extracts the [[RangeRelationBroadcastMode]] and [[BroadcastQueryStageExec]] carried by + * a broadcast query stage, if any. + */ + private def broadcastRangeStage( + plan: LogicalPlan): Option[(RangeRelationBroadcastMode, BroadcastQueryStageExec)] = + plan match { + case LogicalQueryStage(_, bqs: BroadcastQueryStageExec) => + bqs.broadcast.outputPartitioning match { + case BroadcastPartitioning(m: RangeRelationBroadcastMode) => Some((m, bqs)) + case _ => None + } + case _ => None + } + + /** + * A range-predicate join whose left or right child is already a range-broadcast + * query stage. `unapply` looks up the mode once so `apply` does not search + * for both a guard and the constructor. + */ + private case class RangeBroadcastJoinMatch( + left: LogicalPlan, + right: LogicalPlan, + buildSide: BuildSide, + joinType: JoinType, + leftKeys: Seq[Expression], + rightKeys: Seq[Expression], + condition: Expression, + rangeJoin: RangeJoin) + + private object RangeBroadcastJoinMatch { + def unapply(plan: LogicalPlan): Option[RangeBroadcastJoinMatch] = { + if (!SQLConf.get.broadcastRangeJoinEnabled) { + None + } else { + plan match { + case j @ Join(_, _, _, Some(cond), _) => + j match { + case ExtractRangeJoinKeys(left, right, leftKeys, rightKeys, joinType, rangeJoin) => + def fromStage( + mode: RangeRelationBroadcastMode, + bqs: BroadcastQueryStageExec, + buildSide: BuildSide): Option[RangeBroadcastJoinMatch] = { + if (!mode.indexKind.matches(rangeJoin) || + !BroadcastRangeJoinExec.supports(joinType, buildSide)) { + None + } else { + val buildKeys = buildSide match { + case BuildLeft => leftKeys + case BuildRight => rightKeys + } + // Mode was bound against the physical exchange output (`bqs.output`). + // Binding against `bqs.output` ensures ordinals align even when logical + // projections above the stage were collapsed or eliminated. Review Comment: Fixed: comment now attributes the mismatch to AQE re-planning, not a wrapping projection. -- 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]
