cloud-fan commented on code in PR #59027: URL: https://github.com/apache/spark/pull/59027#discussion_r4333126465
########## sql/core/src/main/scala/org/apache/spark/sql/execution/joins/RangeRelation.scala: ########## @@ -0,0 +1,1131 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.execution.joins + +import java.io.{Externalizable, ObjectInput, ObjectOutput} + +import scala.collection.mutable.ArrayBuffer + +import com.esotericsoftware.kryo.{Kryo, KryoSerializable} +import com.esotericsoftware.kryo.io.{Input, Output} + +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{BoundReference, Expression, InterpretedProjection} +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.util.{KnownSizeEstimation, SizeEstimator, Utils} + +/** + * Broadcast payload of a range join. An index only names candidate rows. + * [[BroadcastRangeJoinExec]] evaluates the original join condition to filter + * candidates. Keys that map to one order-preserving long are stored raw and + * compared unboxed ([[LongIntervalIndex]], [[LongPointIndex]]); every other + * supported type is compared as boxed values ([[IntervalIndex]], + * [[PointIndex]]). [[RangeRelationBroadcastMode]] picks the family from the + * build keys' data types. + * + * - The interval indexes answer "which ranges overlap this window?" + * - The point indexes answer "which points lie on this side of a bound?" + */ +private[execution] sealed trait RangeRelation extends KnownSizeEstimation with Serializable { + def keyType: DataType +} + +/** + * Which [[RangeRelation]] a broadcast materializes. + * + * A point index has no direction, so one broadcast of a column serves both + * `<` and `>`. The join records which side holds the lower bound and scans + * accordingly. + */ +private[execution] sealed abstract class RangeIndexKind { + def matches(rangeJoin: RangeJoin): Boolean = RangeIndexKind(rangeJoin) == this +} + +private[execution] case object IntervalIndexKind extends RangeIndexKind + +private[execution] case object PointIndexKind extends RangeIndexKind + +private[execution] object RangeIndexKind { + def apply(rangeJoin: RangeJoin): RangeIndexKind = rangeJoin match { + case PointInRangeJoin | IntervalOverlapJoin => IntervalIndexKind + case _: PartialRangeJoin => PointIndexKind + } +} + +/** + * The boxed-key index family's shared ordering cache, derived from `keyType` + * on first comparison and kept out of the serialized form, so each executor + * re-derives its own instance instead of receiving one per broadcast. The + * long index family compares raw keys and has no ordering to cache. + */ +private[joins] sealed trait BoxedKeyOrdering extends RangeRelation { + @transient + @volatile private[this] var _ordering: Ordering[Any] = _ + + private[joins] def ordering: Ordering[Any] = { + var ord = _ordering + if (ord == null) { + ord = PhysicalDataType.ordering(keyType) + _ordering = ord + } + ord + } +} + +/** + * The long index family's shared key conversion, derived from `keyType` on + * first probe and kept out of the serialized form, so each executor derives + * its own converter after deserialization. The interpreted probe path calls + * [[toLong]] once per streamed key; generated probe code keeps its own + * conversion, emitted at plan time. + */ +private[joins] sealed trait LongKeyConversion extends RangeRelation { + @transient + @volatile private[this] var _toLong: Any => Long = _ + + /** + * One evaluated streamed key to its order-preserving long. `RangePredicate` + * gives both sides of a range predicate one type after analysis, so the + * streamed key type is this index's `keyType`. + */ + private[joins] def toLong(key: Any): Long = { + var f = _toLong + if (f == null) { + f = RangeIndex.requiredLongKey(keyType).value + _toLong = f + } + f(key) + } +} + +private[joins] object RangeIndex { + + /** 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 key type that has one, or None when the type + * must stay on the generic boxed index: variable-length keys, floating point, + * whose SQL equality for -0.0 and NaN has no bit-preserving order, and + * decimals wider than `Decimal.MAX_LONG_DIGITS`, whose unscaled value does + * not fit a long. This is the only place a key type is specialized; the long + * indexes and the generated probe code both consult it. + */ + def longKey(dt: DataType): Option[LongKey] = { + def key(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. + */ + def canMapKeysToLong(keys: Seq[Expression]): Boolean = + keys.forall(k => longKey(k.dataType).isDefined) + + /** longKey for a type the caller has routed to the long family. */ + def requiredLongKey(dt: DataType): LongKey = { + longKey(dt).getOrElse( + throw new IllegalStateException(s"$dt has no order-preserving long key.")) + } +} + +/** + * Array-backed candidate scan. The index hands rows to the join in array order, + * without a closure or a tuple hop: the per-pair overhead then matches a bare + * loop over a broadcast row array. One instance per probe; never serialized. + */ +private[joins] final class RowSlice( + rows: Array[InternalRow], + from: Int, + until: Int) extends Iterator[InternalRow] { + private var i = from + def hasNext: Boolean = i < until + def next(): InternalRow = { + val r = rows(i) + i += 1 + r + } +} + +private[joins] object RowSlice { + /** A shared empty scan: a null probe bound matches nothing. */ + val empty: RowSlice = new RowSlice(new Array[InternalRow](0), 0, 0) +} + +// ===== Boxed-key indexes ==================================================== +// +// Keys are the boxed values `RangeIndex.getValue` reads, compared through +// `BoxedKeyOrdering`. Floating point, strings, binary, and decimals wider +// than 18 digits stay here: these types have no order-preserving +// single-long form, which is the property that routes a key type to the +// long indexes below (`RangeIndex.longKey`). + +/** + * Augmented interval tree of build-side intervals. Lows and highs sit in + * parallel arrays sorted by low bound, with a row array in the same order + * feeding the scan path and `maxHighIdx(mid)` indexing the maximum high + * bound in the span rooted at mid. Each key and row object is owned once: + * the tree holds one reference per value plus the single row array, so a + * serialized index keeps the same object count it was built with. + * Subtrees whose max high is below the probe window, or whose low starts + * after it, are pruned during search, and a window covering every interval + * skips the walk entirely. + */ +private[joins] class IntervalIndex private ( + private[this] var _keyType: DataType, + private[this] var lows: Array[Any], + private[this] var highs: Array[Any], + private[this] var maxHighIdx: Array[Int], + private[this] var rows: Array[InternalRow]) + extends BoxedKeyOrdering with Externalizable with KryoSerializable { + + // Needed for Java/Kryo deserialization; see readExternal/read below. + def this() = this(null, null, null, null, null) + + def keyType: DataType = _keyType + + // All-covering window bounds. `lows` is sorted, so the last entry bounds Review Comment: **Nit (P3):** `lows` is sorted ascending, so its last entry is the largest low. It bounds every low from above, and the all-covering check uses it that way (`probeHigh >= coreLow`). `LongIntervalIndex` has the same wording (lines 689-690). Two nearby comments have the same kind of slip: - The `IntervalIndex` scaladoc (lines 224-226) says a serialized index "keeps the same object count it was built with". In a one-key build, `lows(i)` and `highs(i)` are the same key object. Java serialization keeps that alias, but Kryo's `MapReferenceResolver` never tracks primitive wrappers, so a float or double point build comes back with two boxed keys per entry even with `spark.kryo.referenceTracking=true`. With tracking off, every shared key is duplicated. - `RangeIndexSuite`'s scaladoc (line 31) describes "the two range-join indexes", but the suite also covers `LongIntervalIndex`, `LongPointIndex`, and the broadcast-mode routing between the two families. **Verification:** - **Inspection:** The coreLow comments describe an upper bound on every low. The interval index serialization note matches its behavior under Java serialization and under Kryo with and without reference tracking, including float and double keys. The suite description names all four index classes and the family-routing tests. ########## sql/core/src/test/scala/org/apache/spark/sql/execution/benchmark/RangeJoinBenchmark.scala: ########## @@ -0,0 +1,479 @@ +/* + * 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.benchmark + +import java.sql.Date + +import org.apache.spark.benchmark.Benchmark +import org.apache.spark.sql.{Column, DataFrame} +import org.apache.spark.sql.execution.joins.{BroadcastNestedLoopJoinExec, BroadcastRangeJoinExec} +import org.apache.spark.sql.functions.{broadcast, col, date_add, lit, lpad, when} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.StringType + +/** + * Benchmark broadcast range join against broadcast nested loop join. + * + * Long keys cross every supported shape with every supported join type. + * Inner joins broadcast each side once. Left anti probes values that match + * nothing. Two long-key cases are not selective: wide interval overlap, and a + * partial range whose bound sits at the median. Point-in-range inner is + * repeated for date, timestamp, decimal, and UNICODE_CI strings. + * + * More cases probe the cost model directly: a streamed side of 1, 10, or 20 + * rows, where range join pays to sort the whole build side instead of one + * nested-loop scan, and a fully overlapping inner join, where every pair + * matches and the index prunes nothing. + * + * Staggered high-overlap builds cover the regime between those: most, but not + * all, build intervals overlap each probe window, miss intervals are + * interleaved so no subtree's max high falls under the probe, and the build + * has no common core, so the tree is walked without pruning and evaluated per + * candidate. Long and double (boxed) keys both run at 50% and 90% overlap. + * + * To run this benchmark: + * {{{ + * 1. build/sbt "sql/Test/runMain org.apache.spark.sql.execution.benchmark.RangeJoinBenchmark" + * 2. generate result: + * SPARK_GENERATE_BENCHMARK_FILES=1 build/sbt \ + * "sql/Test/runMain org.apache.spark.sql.execution.benchmark.RangeJoinBenchmark" + * Results will be written to "benchmarks/RangeJoinBenchmark-results.txt". + * }}} + */ +object RangeJoinBenchmark extends SqlBasedBenchmark { + + // An optional regex in the environment variable RANGE_BENCH_FILTER runs only the + // sections whose name matches, e.g. RANGE_BENCH_FILTER=".*tiny probe.*|.*full overlap.*". + // A workaround for sbt splitting runMain arguments on spaces. + private var nameFilter: Option[String] = Option(System.getenv("RANGE_BENCH_FILTER")) + + private def shouldRun(name: String): Boolean = nameFilter.forall(name.matches) + + private val buildRows = 4000 + private val probeRows = 400000 + private val width = 10 + // Wide overlap and the median bound explode the output, so they stay smaller. + private val heavyBuildRows = 2000 + private val heavyProbeRows = 100000 + private val halfBuildRows = 1000 + private val halfProbeRows = 10000 + private val wideSpan = 20 * width + + private val epoch = lit(Date.valueOf("1970-01-01")) + + // [0, buildRows * width). Each range covers `width` values and none overlap. + private lazy val ranges: DataFrame = cached(spark.range(buildRows).select( + (col("id") * width).as("lo"), + (col("id") * width + (width - 1)).as("hi"))) + + // One hit per point: the value falls in exactly one range above. + private lazy val points: DataFrame = cached(spark.range(probeRows).select( + (col("id") % (buildRows.toLong * width)).as("p"))) + + // Above every range, so a point-in-range probe matches nothing. + private lazy val misses: DataFrame = cached(spark.range(probeRows).select( + (lit(buildRows.toLong * width) + (col("id") % width)).as("p"))) + + // Short intervals. A probe overlaps the build intervals whose starts are near it. + private lazy val probeIntervals: DataFrame = + cached(intervals(probeRows, "a", width + width / 2)) + private lazy val buildIntervals: DataFrame = + cached(intervals(buildRows, "b", width + width / 2)) + // Starts after every build interval. + private lazy val missIntervals: DataFrame = cached(spark.range(probeRows).select( + (lit(buildRows.toLong * width + width) + col("id") * width).as("alo"), + (lit(buildRows.toLong * width + width * 2) + col("id") * width).as("ahi"))) + + // About 40 build intervals overlap each probe. Review Comment: **Nit (P3):** This doesn't match the generators. `wideProbe` has 100,000 intervals `[10i, 10i + 200]` and `wideBuild` has 2,000 intervals `[10j, 10j + 200]`, so a probe overlaps a build interval only when `|i - j| <= 20`. About 2% of probes (`i <= 2019`) hit roughly 41 intervals, and the rest hit nothing. The case is mostly selective, then, not one of the two cases the scaladoc calls "not selective" (lines 34-35), and it doesn't explode the output as line 70 says. The PR description's "Interval overlap, wide (~40 hits per probe)" row has the same misreading. Separately, the scaladoc (line 38) and the comment at line 125 say the tiny-probe cases use 1, 10, or 20 streamed rows, but `tinyProbeRowCounts` is `Seq(1, 10, 100, 1000, 5000, 10000)`. Could these comments describe what the generators actually produce? **Verification:** - **Inspection:** The tiny-probe description lists exactly the sizes the benchmark runs, and the wide-overlap description matches the overlap implied by its generators: a probe overlaps a build interval only when their start indexes differ by at most 20, over 100,000 probes and 2,000 build intervals. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/joins/BroadcastRangeJoinExec.scala: ########## @@ -0,0 +1,591 @@ +/* + * 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.broadcast.Broadcast +import org.apache.spark.rdd.RDD +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions._ +import org.apache.spark.sql.catalyst.expressions.codegen._ +import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide, PartialRangeJoin, RangeJoin} +import org.apache.spark.sql.catalyst.plans._ +import org.apache.spark.sql.catalyst.plans.physical.{BroadcastDistribution, Distribution, Partitioning, UnspecifiedDistribution} +import org.apache.spark.sql.execution.{CodegenSupport, ExplainUtils, SparkPlan} +import org.apache.spark.sql.execution.metric.SQLMetrics + +/** + * Performs a range join on two tables. Supported types are inner, left outer, right outer, + * left semi, and left anti. The preserved side is streamed: unmatched build rows are not + * emitted, because the broadcast is copied to every partition. See `ExtractRangeJoinKeys` + * for the exact set of join-condition shapes that are recognized as a range join. + * + * The build side is an index. Point-in-range and interval overlap both build an + * [[IntervalIndex]] and probe it with `overlapping`. A single inequality builds a + * [[PointIndex]] and scans the side of the bound the build side holds. Keys that + * map to one order-preserving long instead build [[LongIntervalIndex]] or + * [[LongPointIndex]], so their probes compare raw longs without boxing. Either + * probe returns a superset. A candidate is kept only when the original join + * condition is true, so inclusivity stays on the predicate. Heavy overlap can + * make one probe return almost every build row. + */ +case class BroadcastRangeJoinExec( + left: SparkPlan, + right: SparkPlan, + buildSide: BuildSide, + joinType: JoinType, + leftKeys: Seq[Expression], + rightKeys: Seq[Expression], + condition: Option[Expression], + rangeJoin: RangeJoin) + extends JoinCodegenSupport { + + override lazy val metrics = Map( + "numOutputRows" -> SQLMetrics.createMetric(sparkContext, "number of output rows")) + + override lazy val requiredChildDistribution: Seq[Distribution] = { + val mode = RangeRelationBroadcastMode( + BindReferences.bindReferences(buildKeys, buildOutput), + RangeIndexKind(rangeJoin)) + buildSide match { + case BuildLeft => + BroadcastDistribution(mode) :: UnspecifiedDistribution :: Nil + case BuildRight => + UnspecifiedDistribution :: BroadcastDistribution(mode) :: Nil + } + } + + override def output: Seq[Attribute] = { + joinType match { + case _: InnerLike => + left.output ++ right.output + case LeftOuter => + left.output ++ right.output.map(_.withNullability(true)) + case RightOuter => + left.output.map(_.withNullability(true)) ++ right.output + case LeftSemi | LeftAnti => + left.output + case t => + throw new IllegalArgumentException( + s"BroadcastRangeJoin should not take $t as the JoinType") + } + } + + // The preserved side is streamed, so its partitioning and ordering pass through. + override def outputPartitioning: Partitioning = streamedPlan.outputPartitioning + + override def outputOrdering: Seq[SortOrder] = streamedPlan.outputOrdering + + override def simpleStringWithNodeId(): String = { + val opId = ExplainUtils.getOpId(this) + s"$nodeName $joinType $buildSide ($opId)".trim + } + + private[this] lazy val buildPlan = if (buildSide == BuildLeft) left else right + private[this] lazy val streamedPlan = if (buildSide == BuildLeft) right else left + + private[this] lazy val buildKeys = if (buildSide == BuildLeft) leftKeys else rightKeys + private[this] lazy val streamedKeys = if (buildSide == BuildLeft) rightKeys else leftKeys + + @transient private[this] lazy val buildOutput = buildPlan.output + @transient private[this] lazy val streamedOutput = streamedPlan.output + + /** Build side holds the lower bound, so points at or below the stream key match. */ + private def buildHoldsLowerBound(partial: PartialRangeJoin): Boolean = + partial.leftIsLower == (buildSide == BuildLeft) + + /** + * Candidates for one streamed row. A null bound is unknown: the index + * returns an empty iterator, and the original condition is not evaluated. + * A point probe has one key. Pass it as both ends so it is evaluated once. + */ + private def candidates( + relation: RangeRelation, + boundStreamKeys: Seq[Expression], + streamRow: InternalRow): Iterator[InternalRow] = { + val low = boundStreamKeys.head.eval(streamRow) + relation match { + case longs: LongIntervalIndex => + if (low == null) { + Iterator.empty + } else { + // A point probe has one key. Pass it as both ends so it is evaluated once. + val high = if (boundStreamKeys.length == 1) low else boundStreamKeys(1).eval(streamRow) + if (high == null) Iterator.empty else { + longs.overlapping(longs.toLong(low), longs.toLong(high)) + } + } + case points: LongPointIndex => + if (low == null) { + Iterator.empty + } else { + rangeJoin match { + case partial: PartialRangeJoin if buildHoldsLowerBound(partial) => + points.upTo(points.toLong(low)) + case _: PartialRangeJoin => + points.from(points.toLong(low)) + case _ => + throw new IllegalStateException( + s"Point index is only built for a partial range, got $rangeJoin.") + } + } + case intervals: IntervalIndex => + val high = if (boundStreamKeys.length == 1) low else boundStreamKeys(1).eval(streamRow) + intervals.overlapping(low, high) + case points: PointIndex => + rangeJoin match { + case partial: PartialRangeJoin if buildHoldsLowerBound(partial) => + points.upTo(low) + case _: PartialRangeJoin => + points.from(low) + case _ => + throw new IllegalStateException( + s"Point index is only built for a partial range, got $rangeJoin.") + } + } + } + + /** + * Broadcast relation, bound probe keys, and the original condition for one + * streamed partition. The joined row is reused; the caller copies it. + */ + private class ProbeContext( + val stream: Iterator[InternalRow], + relation: RangeRelation, + keys: Seq[Expression], + val condition: InternalRow => Boolean, + joinedRow: JoinedRow, + val nullBuild: InternalRow) { + + def candidates(row: InternalRow): Iterator[InternalRow] = + BroadcastRangeJoinExec.this.candidates(relation, keys, row) + + def join(streamRow: InternalRow, buildRow: InternalRow): InternalRow = buildSide match { + case BuildRight => joinedRow(streamRow, buildRow) + case BuildLeft => joinedRow(buildRow, streamRow) + } + } + + private def perPartition( + relation: Broadcast[RangeRelation])( + f: ProbeContext => Iterator[InternalRow]): RDD[InternalRow] = { + val boundStreamKeys = BindReferences.bindReferences(streamedKeys, streamedOutput) + val cond = condition.get + val schema = left.output ++ right.output + val numBuildFields = buildPlan.output.length + streamedPlan.execute().mapPartitionsInternal { stream => + f(new ProbeContext( + stream, + relation.value, + boundStreamKeys, + Predicate.create(cond, schema).eval _, + new JoinedRow, + new GenericInternalRow(numBuildFields))) + } + } + + private def innerJoin(relation: Broadcast[RangeRelation]): RDD[InternalRow] = { + perPartition(relation) { p => + p.stream.flatMap { row => + p.candidates(row).flatMap { buildRow => + val joined = p.join(row, buildRow) + if (p.condition(joined)) Iterator.single(joined) else Iterator.empty + } + } + } + } + + /** + * Emit each candidate the condition accepts. When none were, emit the stream + * row once with the build side padded by nulls. + * + * The loop test checks `streamRow` before `stream.hasNext`. A whole-stage + * stream reuses one `UnsafeRow`, and `hasNext` fills that buffer. It must + * not run while the current row is still being joined. + */ + private def outerJoin(relation: Broadcast[RangeRelation]): RDD[InternalRow] = { + perPartition(relation) { p => + new Iterator[InternalRow] { + private var streamRow: InternalRow = null + private var matches: Iterator[InternalRow] = Iterator.empty + private var foundMatch: Boolean = false + private var resultRow: InternalRow = null + + private def findNext(): Boolean = { + while (streamRow != null || p.stream.hasNext) { + if (streamRow == null) { + streamRow = p.stream.next() + foundMatch = false + matches = p.candidates(streamRow) + } + while (matches.hasNext) { + val joined = p.join(streamRow, matches.next()) + if (p.condition(joined)) { + foundMatch = true + resultRow = joined + return true + } + } + if (!foundMatch) { + resultRow = p.join(streamRow, p.nullBuild) + streamRow = null + return true + } + streamRow = null + } + false + } + + override def hasNext: Boolean = resultRow != null || findNext() + + override def next(): InternalRow = { + val r = resultRow + resultRow = null + r + } + } + } + } + + /** Semi keeps a stream row that matches. Anti keeps one that does not. */ + private def leftExistenceJoin( + relation: Broadcast[RangeRelation], + exists: Boolean): RDD[InternalRow] = { + perPartition(relation) { p => + p.stream.filter { row => + val matched = p.candidates(row).exists { buildRow => + p.condition(p.join(row, buildRow)) + } + matched == exists + } + } + } + + private def genResultProjection: UnsafeProjection = joinType match { + case LeftSemi | LeftAnti => + UnsafeProjection.create(output, output) + case _ => + UnsafeProjection.create( + output, (left.output ++ right.output).map(_.withNullability(true))) + } + + override def doExecute(): RDD[InternalRow] = { + require(BroadcastRangeJoinExec.supports(joinType, buildSide), + s"BroadcastRangeJoin does not support $joinType with $buildSide") + val relation = buildPlan.executeBroadcast[RangeRelation]() + val joined = joinType match { + case _: InnerLike => innerJoin(relation) + case LeftOuter | RightOuter => outerJoin(relation) + case LeftSemi => leftExistenceJoin(relation, exists = true) + case LeftAnti => leftExistenceJoin(relation, exists = false) + case t => + throw new IllegalArgumentException( + s"BroadcastRangeJoin should not take $t as the JoinType") + } + val numOutputRows = longMetric("numOutputRows") + joined.mapPartitionsWithIndexInternal { (index, iter) => + val resultProj = genResultProjection + resultProj.initialize(index) + iter.map { r => + numOutputRows += 1 + resultProj(r) + } + } + } + + override def supportCodegen: Boolean = { + BroadcastRangeJoinExec.supports(joinType, buildSide) + } + + override def inputRDDs(): Seq[RDD[InternalRow]] = { + streamedPlan.asInstanceOf[CodegenSupport].inputRDDs() + } + + private def multipleOutputForOneInput: Boolean = joinType match { + case _: InnerLike | LeftOuter | RightOuter => true + case _ => false + } + + override def needCopyResult: Boolean = { + streamedPlan.asInstanceOf[CodegenSupport].needCopyResult || multipleOutputForOneInput + } + + override def doProduce(ctx: CodegenContext): String = { + streamedPlan.asInstanceOf[CodegenSupport].produce(ctx, this) + } + + override def doConsume(ctx: CodegenContext, input: Seq[ExprCode], row: ExprCode): String = { + joinType match { + case _: InnerLike => codegenInner(ctx, input) + case LeftOuter | RightOuter => codegenOuter(ctx, input) + case LeftSemi => codegenLeftExistence(ctx, input, exists = true) + case LeftAnti => codegenLeftExistence(ctx, input, exists = false) + case t => + throw new IllegalArgumentException( + s"BroadcastRangeJoin code-gen should not take $t as the JoinType") + } + } + + private def prepareBroadcast(ctx: CodegenContext): String = { + val broadcast = buildPlan.executeBroadcast[RangeRelation]() + val broadcastTerm = ctx.addReferenceObj("broadcast", broadcast) + val cls = classOf[RangeRelation].getName + ctx.addMutableState(cls, "rangeRelation", + v => s"$v = ($cls) $broadcastTerm.value();", forceInline = true) + } + + /** + * Box each streamed range key. The index compares the same boxed values + * `RangeIndex.getValue` produces on the build side. + */ + private def genProbeKeys( + ctx: CodegenContext, + input: Seq[ExprCode]): (String, Seq[String]) = { + ctx.currentVars = input + val keys = streamedKeys.map { key => + val bound = BindReferences.bindReference(key, streamedOutput) + val ev = bound.genCode(ctx) + val name = ctx.freshName("probeKey") + val javaType = CodeGenerator.javaType(key.dataType) + val boxed = CodeGenerator.boxedType(key.dataType) + val valueExpr = if (javaType == boxed) { + s"${ev.value}" + } else { + s"$boxed.valueOf(${ev.value})" + } + val code = + s""" + |${ev.code} + |Object $name = ${ev.isNull} ? null : $valueExpr; + """.stripMargin + (code, name) + } + (keys.map(_._1).mkString("\n"), keys.map(_._2)) + } + + /** + * Probe the broadcast index. The index kind and scan direction are fixed when + * the plan is built, so generated code calls the concrete index method + * directly. A null key yields the empty scan those methods already return. + * Generated code does not allocate a key array. An interval probe owns its + * walk stack. + */ + private def genCandidateIterator( + ctx: CodegenContext, + relationTerm: String, + keyVars: Seq[String]): (String, String) = { + val matches = ctx.freshName("matches") + val iterCls = classOf[Iterator[InternalRow]].getName + val call = rangeJoin match { + case partial: PartialRangeJoin => + val points = classOf[PointIndex].getName + val scan = if (buildHoldsLowerBound(partial)) "upTo" else "from" + s"(($points) $relationTerm).$scan(${keyVars.head})" + case _ => + val intervals = classOf[IntervalIndex].getName + // A point probe has one key. Reuse it as both ends so it is evaluated once. + val high = if (keyVars.length == 1) keyVars.head else keyVars(1) + s"(($intervals) $relationTerm).overlapping(${keyVars.head}, $high)" + } + val code = s"$iterCls $matches = $call;" + (matches, code) + } + + /** + * Keys that map to one order-preserving long (`RangeIndex.longKey`) route + * to the long indexes: build and probe keys then cross as raw longs, with + * no boxing on either side, and every other supported key stays boxed. + * One test picks both the `RangeRelationBroadcastMode` index family and + * the generated probe code, so they cannot disagree. + */ + private[this] lazy val usesLongKeys: Boolean = RangeIndex.canMapKeysToLong(buildKeys) + + /** + * Evaluate each streamed key into a primitive long. Each key keeps the null + * term of its expression; the candidate call turns a null into an empty scan. + */ + private def genLongProbeKeys( + ctx: CodegenContext, + input: Seq[ExprCode]): (String, Seq[String], Seq[String]) = { + ctx.currentVars = input + val keys = streamedKeys.map { key => + val bound = BindReferences.bindReference(key, streamedOutput) + val ev = bound.genCode(ctx) + val toLong = RangeIndex.requiredLongKey(key.dataType) + val name = ctx.freshName("probeKey") + val code = + s""" + |${ev.code} + |long $name = ${ev.isNull} ? 0L : ${toLong.code(ev.value)}; + """.stripMargin + (code, name, ev.isNull) + } + (keys.map(_._1).mkString("\n"), keys.map(_._2), keys.map(_._3)) + } + + private def genLongCandidateIterator( + ctx: CodegenContext, + relationTerm: String, + keyVars: Seq[String], + nullVars: Seq[String]): (String, String) = { + val matches = ctx.freshName("matches") + val iterCls = classOf[Iterator[InternalRow]].getName + val emptyScan = s"${classOf[RowSlice].getName}$$.MODULE$$.empty()" + val call = rangeJoin match { + case partial: PartialRangeJoin => + val points = classOf[LongPointIndex].getName + val scan = if (buildHoldsLowerBound(partial)) "upTo" else "from" Review Comment: **Non-blocking (P2):** With `a.x < b.y` built on the right, this picks `from(x)`, and `PointIndex.from` / `LongPointIndex.from` start at the first key `>= x` (`searchBound(value, skipEquals = false)`). Every build row whose key equals `x` comes first, and each one fails the strict condition. `RangePredicate` maps `<`, `<=`, `>`, and `>=` to the same `(type, low, high)` shape, so the scan can't know the predicate is strict. For semi and anti joins this defeats the early exit. `leftExistenceJoin` and `codegenLeftExistence` evaluate dup(x) + 1 conditions per stream row before they can stop, while BNLJ breaks on the first matching build row. For example, take a build of 100,000 rows over 100 dates (1,000 per date, in no particular order) and `a LEFT SEMI JOIN b ON a.d < b.d`, which is what `WHERE EXISTS (... b.d > a.d)` rewrites to. A probe on the 99th date is selective, because only 1% of build rows are later, yet the range join evaluates about 1,001 conditions where BNLJ needs about 100 on average. For a probe on the first date, it is about 1,001 against about 1. That is the selective case the config doc promises to make an order of magnitude faster. The committed partial-range semi/anti benchmarks use distinct build keys, so they can't see this. The `a.x > b.y` orientation scans `upTo(x)` from the smallest key, so its equal keys come last and it is unaffected. The cost is bounded by O(N) per probe and depends on build order: an ascending-sorted build makes BNLJ the slower plan instead. **Recommended change:** Record whether the extracted single inequality is strict, and carry that on the partial-range join tag. When it is strict, have the range join operator use the exclusive form of the point-index bound, which skips build keys equal to the probe value. Do this in the interpreted path and in both generated paths, for both the boxed and the long point index. Add range join tests that cover strict and inclusive inequalities over duplicate build keys. **Why this works:** Both point indexes already locate their bounds with a binary search that can skip equal keys. Exposing the exclusive bound, and choosing it from the predicate's strictness, makes the candidate slice contain only keys that can satisfy the extracted inequality. The first candidate of a semi or anti join then already passes the range predicate, and inner and outer joins stop evaluating equal keys the predicate always rejects. **Scope:** Carry the strictness of the extracted single inequality from range-key extraction to the point-index scan, so strict predicates scan exclusive bounds in every execution path. **Compatibility:** Inclusive inequalities keep their current inclusive scans, and the broadcast index contents and adaptive-execution reuse keys are unchanged. **Risks:** If strictness is attributed to the wrong predicate, the index drops keys the condition accepts and rows are lost silently. For example, a condition can contain both a strict and an inclusive inequality over the same columns. **Constraints:** Strictness must come from the exact inequality whose keys were extracted, and the exclusive bound must use the same ordering as the condition, so that keys equal under a collation or under SQL floating-point semantics are excluded exactly when the predicate rejects them. Point-in-range and interval-overlap indexes remain inclusive supersets whose candidates are filtered by the original condition. **Exclusions:** - A cost-based or overlap-based fallback to the nested loop join, which the documented heavy-overlap slowdown leaves out of scope. **Success:** For a strict single inequality, the point-index candidates exclude build keys equal to the probe value, in both index families and in both interpreted and generated execution. A left semi or left anti join on a strict inequality over duplicate-heavy build keys evaluates a bounded number of conditions per probe instead of one per duplicate build row. Answers for strict and inclusive single-inequality joins, in both directions and on both build sides, still match the nested loop join. **Verification:** - **Regression:** With whole-stage codegen on and off, a left semi and a left anti join on a strict inequality whose build side holds many rows equal to each probe value evaluate fewer conditions per probe than the number of equal build rows, for both long and boxed keys. - **Behavior:** Strict and inclusive single-inequality joins of every supported join type, in both directions and on both build sides, return the nested-loop result when probe values coincide with build keys, including collation-equal strings and signed-zero doubles. - **Compatibility:** Under adaptive execution, a point-index broadcast built for a strict inequality is still recognized and reused by the range join strategy, and strict and inclusive partial ranges on the same key map to the same index kind. ########## sql/core/src/test/scala/org/apache/spark/sql/execution/joins/RangeIndexSuite.scala: ########## @@ -0,0 +1,966 @@ +/* + * 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("a one-key interval build owns a single copy of the point key") { + // One build key is one key used as both ends. The build must copy it once + // and serve both bounds from that copy, instead of copying twice. + val n = 50 + val oneKey = Seq(BoundReference(0, StringType, nullable = true)) + val (sharedBytes, rows) = reusingByteBufferStringIterator(n) + val index = IntervalIndex.build(StringType, oneKey, rows, Some(n)) + + java.util.Arrays.fill(sharedBytes, 'X'.toByte) + + (0 until n).foreach { i => + val key = UTF8String.fromString(f"$i%04d") + val matching = index.overlapping(key, key).toList + assert(matching.size == 1, s"interval $i was corrupted by shared buffer reuse") + assert(matching.head.getInt(1) == i) + } + + // SizeEstimator's identity tracking charges the shared bounds once, so the + // estimate stays below the same build made from two independently copied keys. + val twoKeys = Seq( + BoundReference(0, StringType, nullable = true), + BoundReference(0, StringType, nullable = true)) + val (_, moreRows) = reusingByteBufferStringIterator(n) + val twoCopy = IntervalIndex.build(StringType, twoKeys, moreRows, Some(n)) + assert(index.estimatedSize < twoCopy.estimatedSize) + } + + test("an inverted build interval reaches the windows the condition accepts") { + // Int keys route to the long index, so this also pins swap normalization there. + val index = RangeRelationBroadcastMode( + Seq(BoundReference(0, IntegerType, nullable = true), + BoundReference(1, IntegerType, nullable = true)), + IntervalIndexKind) + .transform(Array(InternalRow(3, 1), InternalRow(0, 5))) + .asInstanceOf[LongIntervalIndex] + assert(index.overlapping(0L, 4L).map(_.getInt(0)).toSet == Set(3, 0)) + assert(index.overlapping(4L, 9L).map(_.getInt(0)).toList == List(0)) + } + + test("interval tree returns each overlapping row once") { + // Spans arrive out of low order: [10, 20] then [0, 5]. + val unsorted = buildIntervals(Seq(span(10, 20, 0), span(0, 5, 1))) + assert(probeIds(unsorted, 3, 3) == List(1)) + assert(probeIds(unsorted, 15, 15) == List(0)) + + // Probe 50 is a build key. Row 0 spans it; row 1 starts there. + val spanning = buildIntervals(Seq(span(0, 100, 0), span(50, 200, 1))) + assert(probeIds(spanning, 50, 50).toSet == Set(0, 1)) + assert(probeIds(spanning, 25, 25) == List(0)) + assert(probeIds(spanning, 150, 150) == List(1)) + + val duplicated = buildIntervals(Seq(span(0, 100, 0), span(0, 100, 1))) + assert(probeIds(duplicated, 50, 50).sorted == List(0, 1)) + + val nested = buildIntervals(Seq( + span(0, 100, 0), span(10, 20, 1), span(10, 50, 2), span(90, 110, 3))) + def assertOnce(low: Int, expected: Set[Int]): Unit = { + val actual = probeIds(nested, low, low) + assert(actual.toSet == expected, s"probe $low -> $actual") + assert(actual.size == expected.size, s"probe $low duplicated $actual") + } + assertOnce(0, Set(0)) + assertOnce(10, Set(0, 1, 2)) + assertOnce(15, Set(0, 1, 2)) + assertOnce(30, Set(0, 2)) + assertOnce(95, Set(0, 3)) + assertOnce(105, Set(3)) + assertOnce(200, Set.empty) + + assert(probeIds(buildIntervals(Seq.empty), 1, 1).isEmpty) + + val withNull = buildIntervals(Seq(span(0, 10, 0), span(5, 5, 1), (null, 4, row(9)))) + assert(probeIds(withNull, null, 5).isEmpty) + assert(probeIds(withNull, 5, null).isEmpty) + // The null bound is dropped, so two intervals remain and the estimate covers both. + val one = buildIntervals(Seq(span(0, 10, 0))) + assert(withNull.estimatedSize > one.estimatedSize) + + // At an interior endpoint, the range that ends and the one that starts + // are each a candidate once. + val chained = buildIntervals(Seq(span(0, 1, 0), span(1, 2, 1), span(2, 3, 2))) + assert(probeIds(chained, 0, 0) == List(0)) + assert(probeIds(chained, 1, 1).sorted == List(0, 1)) + assert(probeIds(chained, 2, 2).sorted == List(1, 2)) + assert(probeIds(chained, 3, 3) == List(2)) + + // A probe whose endpoints are reversed is the closed window between them. + // [10, 25] meets [0, 30], [15, 40], and the point at 15. + val inverted = buildIntervals(Seq(span(0, 30, 0), span(15, 40, 1), span(15, 15, 2))) + assert(probeIds(inverted, 25, 10).sorted == List(0, 1, 2)) + assert(probeIds(inverted, 10, 20).sorted == List(0, 1, 2)) + } + + test("interval tree is correct under heavy build-side overlap") { + // Every interval nests inside the one before it, so a point in the middle + // hits about half of the rows. The tree still stores one node per row. + val n = 300 + val nested = buildIntervals((0 until n).map(i => span(i, 2 * n - i, i))) + + def expected(probe: Int): Set[Int] = + (0 until n).filter(i => i <= probe && probe <= 2 * n - i).toSet + + Seq(0, 1, n / 2, n, 2 * n - 1, 2 * n).foreach { probe => + assert(probeIds(nested, probe, probe).toSet == expected(probe), s"probe $probe") + } + } + + test("broadcast of heavily overlapping intervals stays linear") { + // [i, i + n] for i in 0 until n. Every interval contains n. The broadcast + // is the flat arrays, one slot per interval, so serialization stays linear. + val n = 200 + val index = buildIntervals((0 until n).map(i => span(i, i + n, i))) + val (raw, restored) = javaRoundTrip(index) + assert(raw.length < n * 1024, s"serialized ${raw.length} bytes") + val probe = n + assert(probeIds(restored, probe, probe).toSet == probeIds(index, probe, probe).toSet) + assert(probeIds(restored, probe, probe).size == n) + } + + test("estimatedSize depends on the interval count, not the overlap depth") { + val n = 200 + val one = buildIntervals(Seq(span(0, 1, 0))) + val disjoint = buildIntervals((0 until n).map(i => span(2 * i, 2 * i + 1, i))) + val overlap = buildIntervals((0 until n).map(i => span(i, i + n, i))) + val again = buildIntervals((0 until n).map(i => span(2 * i, 2 * i + 1, i))) + assert(disjoint.estimatedSize > one.estimatedSize) + assert(overlap.estimatedSize > one.estimatedSize) + assert(disjoint.estimatedSize == again.estimatedSize) + // Nesting stores no extra copy of the rows. Key objects differ, so the two + // indexes need not be equal, but overlap stays the same order of magnitude. + assert(overlap.estimatedSize < disjoint.estimatedSize * 3) + } + + test("estimatedSize includes distinct key objects") { + val ints = buildIntervals(Seq(span(0, 1, 0))) + val wide = UTF8String.fromString("x" * 64) + val strings = buildIntervals(StringType, Seq((wide, wide, row(0)))) + assert(strings.estimatedSize > ints.estimatedSize) + } + + test("interval tree treats NaN as its own point and Infinity as a bound") { + // PhysicalDoubleType's ordering (SQLOrderingUtil.compareDoubles) totally orders + // -Inf < ... < +Inf < NaN, with NaN comparing equal only to itself. So [1.0, +Inf] + // includes +Inf under inclusive probe, and a NaN key overlaps only a probe that + // also lands exactly on NaN. + val index = buildIntervals(DoubleType, Seq( + (1.0, Double.PositiveInfinity, row(1)), + (Double.NegativeInfinity, -1.0, row(2)), + (Double.NaN, Double.NaN, row(3)))) + + assert(ids(index.overlapping(5.0, 5.0)) == List(1)) + assert(ids(index.overlapping(Double.PositiveInfinity, Double.PositiveInfinity)) == List(1)) + assert(ids(index.overlapping(-5.0, -5.0)) == List(2)) + assert(ids(index.overlapping(Double.NegativeInfinity, Double.NegativeInfinity)) == List(2)) + assert(ids(index.overlapping(Double.NaN, Double.NaN)) == List(3)) + assert(ids(index.overlapping(0.0, 0.0)).isEmpty) + } + + test("interval index merges keys that compare equal but are not equal") { + // Array equals is identity, and UTF8String equals is binary, so these keys + // compare equal under the index ordering without being equal. The tree + // orders them with that comparator, so a probe on either key returns both. + // Signed zero is 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 two empty arrays' shells, and the five kept + // points on top of them. + 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 Review Comment: **Nit (P3):** No `DataType` reaches Kryo here. All four index classes write the key type with `writeString(keyType.json)` (`RangeRelation.scala` lines 287, 577, 733, 988) and rebuild it with `DataType.fromJson` in `read()`. The keys and rows they write are rows, `UTF8String`, `Decimal`, `byte[]`, or boxed primitives. Removing any of the 14 `DataType`/`StringConstraint` entries this PR adds to `KryoSerializer`'s default list (`BooleanType$` through `TimestampNTZType$`) would therefore leave this test green. They look left over from the round that serialized the key type and comparator as objects, and they still shift the ID of every later default registration. Could you drop them, keep the four index classes, and update this comment to say the key type crosses as JSON? **Verification:** - **Compatibility:** With Kryo registration required, interval and point indexes for every supported key family, including collated, char and varchar strings, round-trip and answer the same probes as before. - **Inspection:** Compared with the merge target, the default Kryo registration list adds only the interval and point index classes of both families, and the test comment no longer claims to exercise DataType registrations. -- 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]
