HyukjinKwon commented on code in PR #59027: URL: https://github.com/apache/spark/pull/59027#discussion_r4118191678
########## sql/core/src/main/scala/org/apache/spark/sql/execution/joins/RangeIndex.scala: ########## @@ -0,0 +1,350 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.execution.joins + +import scala.collection.immutable.HashMap +import scala.collection.mutable + +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions._ +import org.apache.spark.sql.types.DataType + +/** + * Broadcast payload of a range join. An index only names candidate rows. + * [[BroadcastRangeJoinExec]] keeps a candidate when the original join condition + * is true, so inclusivity never lives in the index. + * + * - [[IntervalIndex]] answers "which ranges overlap this window?" + * - [[PointIndex]] answers "which points lie on this side of a bound?" + */ +private[execution] sealed trait RangeRelation extends Serializable { + def sizeInBytes(): Long +} + +private[execution] object RangeIndex { + + /** In-memory size of one indexed row. Non-unsafe rows are a test-only estimate. */ + def rowSize(row: InternalRow): Long = row match { + case u: UnsafeRow => u.getSizeInBytes.toLong + case _ => row.numFields * 8L + } + + /** + * One sweep-line event used to build an [[IntervalIndex]]. + * + * @param key the bound value this event occurs at + * @param flow [[RangeEvent.Start]], [[RangeEvent.Point]], or [[RangeEvent.End]] + * @param row the build-side row that produced this event + * @param index a unique id for `row` so deactivation is O(1) + */ + case class RangeEvent(key: Any, flow: Int, row: InternalRow, index: Int) + + object RangeEvent { + /** Interval start. Sorts after Point at the same key. */ + val Start = 1 + /** Degenerate interval (low == high). */ + val Point = 0 + /** Interval end. Sorts before Point at the same key. */ + val End = -1 + } + + /** + * A type-specialized accessor that reads a single field from an `InternalRow` as a boxed + * value (or `null` if the field is null). Shared by the exchange (build-side key extraction) + * and the join (stream-side key extraction) so the dispatch table lives in one place. + */ + def getValue(dt: DataType, ordinal: Int): InternalRow => Any = { + val accessor = InternalRow.getAccessor(dt) + (input: InternalRow) => accessor(input, ordinal) + } + + /** + * First index in `keys` whose key compares greater than `value` under `cmp`, or, when + * `orEqual` is set, greater than or equal to `value`. Shared by the lowerBound/upperBound + * pair on both [[IntervalIndex]] and [[PointIndex]]; `cmp` is where they differ, since + * `IntervalIndex` treats a leading null key as sorting before every value. + */ + def searchBound(keys: Array[Any], value: Any, orEqual: Boolean, cmp: (Any, Any) => Int): Int = { + var lo = 0 + var hi = keys.length + while (lo < hi) { + val mid = lo + ((hi - lo) >>> 1) + val c = cmp(keys(mid), value) + if (if (orEqual) c <= 0 else c < 0) lo = mid + 1 else hi = mid + } + lo + } + + /** + * Turn a projected `(low, high)` row into sweep events. A null bound emits none. + * + * An inverted interval (`low > high`) is normalized to `(high, low)` rather than dropped. + * The join condition still accepts a row the predicate accepts, so dropping would lose it: + * with build `(3, 1)` and stream `(0, 5)`, `a.lo < b.hi AND b.lo < a.hi` holds and the row + * has to be a candidate. Normalizing only widens the window. `start <= end` also keeps the + * sweep's invariant, so `IntervalIndex` never sees a reversed pair. + */ + def toRangeEvent( + buildSideKeyValueGetter: List[InternalRow => Any], + lowHighExtr: Projection, + cmp: Ordering[Any]): (InternalRow, Int) => Seq[RangeEvent] = { + (row: InternalRow, index: Int) => { + val lowHigh: InternalRow = lowHighExtr(row) + val low = buildSideKeyValueGetter(0)(lowHigh) + val high = buildSideKeyValueGetter(1)(lowHigh) + if (low != null && high != null) { + val result = cmp.compare(low, high) + if (result == 0) { + RangeEvent(low, RangeEvent.Point, row, index) :: Nil + } else { + val (start, end) = if (result < 0) (low, high) else (high, low) + RangeEvent(start, RangeEvent.Start, row, index) :: + RangeEvent(end, RangeEvent.End, row, index) :: Nil + } + } else { + Nil + } + } + } +} + +/** + * Sweep-line index of build-side intervals. `overlapping` returns every row whose + * interval may overlap `[low, high]`. The window is a superset: the low search is + * exclusive and the high search is inclusive, and the caller drops false positives. + */ +private[execution] object IntervalIndex { + + /** Build an interval index from unsorted events. */ + def build(ordering: Ordering[Any], events: Array[RangeIndex.RangeEvent]): IntervalIndex = { + val eventComparator = new java.util.Comparator[RangeIndex.RangeEvent] { + override def compare(a: RangeIndex.RangeEvent, b: RangeIndex.RangeEvent): Int = { + val keyCmp = ordering.compare(a.key, b.key) + if (keyCmp != 0) keyCmp else Integer.compare(a.flow, b.flow) + } + } + java.util.Arrays.sort(events, eventComparator) + buildFromSorted(ordering, events) + } + + private def buildFromSorted( + ordering: Ordering[Any], + events: Array[RangeIndex.RangeEvent]): IntervalIndex = { + // A leading null key makes the binary search uniform, including an empty index. + val keys = mutable.Buffer[Any](null) + val offsets = mutable.Buffer[Int](0) + val activatedRows = mutable.Buffer.empty[InternalRow] + val activeOld = mutable.Buffer.empty[HashMap[Int, InternalRow]] + val activeAll = mutable.Buffer.empty[HashMap[Int, InternalRow]] + + // Keyed by the event id so a row leaves the active set in O(1). Two snapshots + // are recorded per distinct key: `activeOld` (before this key's own Start + // events) and `activeAll` (after them); `overlapping` below explains why both + // are needed. A build side of heavily-overlapping intervals (e.g. n intervals + // all spanning the same window) has O(n) distinct keys each with an O(n) + // active set, so a plain map copied at every key would make the snapshots + // alone O(n^2). `HashMap` is persistent: `updated`/`removed` share structure + // with the previous version instead of copying it, so recording a snapshot is + // O(log n) and the total cost across all keys stays O(n log n). It is also + // `Serializable` for the broadcast, unlike the order-preserving persistent + // maps in the same package (`TreeSeqMap`, `VectorMap`) -- iteration order + // does not matter here since old/all are tracked as separate snapshots. + var currentKey: Any = null + var currentActiveRows = HashMap.empty[Int, InternalRow] + var currentOldActiveRows = currentActiveRows + var oldActiveRowsCaptured = false + + def writeActiveRows(): Unit = { + activeOld += (if (oldActiveRowsCaptured) currentOldActiveRows else currentActiveRows) + activeAll += currentActiveRows + } + + events.foreach { event => + // Group by the index ordering. Array equality is identity, and + // UTF8String equality is binary, so `!=` would split one sweep key. + if (currentKey == null || ordering.compare(currentKey, event.key) != 0) { + writeActiveRows() + currentKey = event.key + oldActiveRowsCaptured = false + keys += event.key + offsets += activatedRows.size + } + // The active set before any row starts at this key. Ends at this key have + // already been removed by the time the first Point/Start event is seen. + if (event.flow >= RangeIndex.RangeEvent.Point && !oldActiveRowsCaptured) { + currentOldActiveRows = currentActiveRows + oldActiveRowsCaptured = true + } + event.flow match { + case RangeIndex.RangeEvent.Start => + activatedRows += event.row + currentActiveRows = currentActiveRows.updated(event.index, event.row) + case RangeIndex.RangeEvent.Point => + activatedRows += event.row + case RangeIndex.RangeEvent.End => + currentActiveRows = currentActiveRows.removed(event.index) + } + } + writeActiveRows() + + new IntervalIndex(ordering, keys.toArray, offsets.toArray, activeOld.toArray, + activeAll.toArray, activatedRows.toArray) + } +} + +private[execution] class IntervalIndex( + private[this] val ordering: Ordering[Any], + private[this] val keys: Array[Any], + private[this] val offsets: Array[Int], + private[this] val activeOld: Array[HashMap[Int, InternalRow]], + private[this] val activeAll: Array[HashMap[Int, InternalRow]], + private[this] val activated: Array[InternalRow]) + extends RangeRelation { + + override def sizeInBytes(): Long = activated.map(RangeIndex.rowSize).sum Review Comment: I think the persistent-`HashMap` argument above (line 161) holds only on the driver, before the index is broadcast. Scala 2.13's `immutable.HashMap` serializes through `DefaultSerializationProxy`, which writes every entry, so each `activeOld`/`activeAll` snapshot is serialized in full and deserialized on each executor as an independent map. For heavily overlapping build intervals (the O(n) keys x O(n) active rows case the comment describes), the serialized broadcast and the executor-side index are then O(n^2), not O(n log n). `sizeInBytes` also counts only `activated`, so `spark.sql.maxBroadcastTableSize` doesn't see the snapshots. A build side well under the broadcast threshold could still blow up driver serialization or executor memory. Could `sizeInBytes` include the snapshot entry counts, or could the snapshots be stored in a form that stays compact on the wire (e.g. an interval tree, or per-key deltas rebuilt lazily on the executor)? A test that serializes an index of n fully overlapping intervals and checks the size would pin this down. -- 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]
