wangyum commented on code in PR #59027: URL: https://github.com/apache/spark/pull/59027#discussion_r4355924300
########## 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: Fixed: strict flag now carried to the point-index scan via fromExclusive/upToExclusive. ########## 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: Fixed: lows-bound wording, Kryo reference-tracking note, and suite doc all corrected. ########## 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: Fixed: dropped the 14 extra DataType Kryo registrations, kept only the 4 index classes. ########## 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: Fixed: wide-overlap/tiny-probe comments now match the actual benchmark generators. ########## sql/core/src/test/scala/org/apache/spark/sql/execution/benchmark/RangeJoinBenchmark.scala: ########## @@ -0,0 +1,492 @@ +/* + * 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.catalyst.planning.{GreaterPartialRangeJoin, IntervalOverlapJoin, LessPartialRangeJoin, PointInRangeJoin, RangeJoin} +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. A partial range whose bound sits at the median is not selective. + * Wide interval overlap is mostly selective: a probe overlaps a build + * interval only when their starts sit within `wideSpan`, so about 2% of + * probes hit roughly `wideSpan / width` intervals each and the rest hit + * nothing. 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, 100, + * 1000, 5000, or 10000 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 concentrate 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: Fixed: wide-overlap comment matches generators; scans say fromExclusive/upToExclusive. -- 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]
