LuciferYang commented on code in PR #59027:
URL: https://github.com/apache/spark/pull/59027#discussion_r4270808591
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/LogicalQueryStageStrategy.scala:
##########
@@ -58,6 +62,76 @@ object LogicalQueryStageStrategy extends Strategy {
case _ => false
}
+ /**
+ * Extracts the [[RangeBroadcastMode]] carried by a broadcast query stage,
if any.
+ * Keys and the join condition come from the logical join. The mode
identifies
+ * the already-built range index.
+ */
+ private def broadcastRangeMode(plan: LogicalPlan):
Option[RangeBroadcastMode] = plan match {
+ case LogicalQueryStage(_, bqs: BroadcastQueryStageExec) =>
+ bqs.broadcast.outputPartitioning match {
+ case BroadcastPartitioning(m: RangeBroadcastMode) => Some(m)
+ case _ => None
+ }
+ case _ => None
+ }
+
+ /**
+ * A range-predicate join whose left or right child is already a
range-broadcast
+ * query stage. `unapply` looks up the mode once so `apply` does not search
+ * for both a guard and the constructor.
+ */
+ private case class RangeBroadcastJoinStage(
+ left: LogicalPlan,
+ right: LogicalPlan,
+ joinType: JoinType,
+ buildSide: BuildSide,
+ leftKeys: Seq[Expression],
+ rightKeys: Seq[Expression],
+ condition: Expression,
+ rangeJoin: RangeJoin)
+
+ private object RangeBroadcastJoinStage {
+ def unapply(plan: LogicalPlan): Option[RangeBroadcastJoinStage] = plan
match {
+ case j @ Join(_, _, _, Some(cond), _) =>
+ j match {
+ case ExtractRangeJoinKeys(left, right, leftKeys, rightKeys,
joinType, rangeJoin) =>
+ def fromMode(
+ mode: RangeBroadcastMode,
+ buildSide: BuildSide): Option[RangeBroadcastJoinStage] = {
+ if (!mode.indexKind.matches(rangeJoin) ||
+ !BroadcastRangeJoinExec.supports(joinType, buildSide)) {
+ None
+ } else {
+ val buildPlan = buildSide match {
+ case BuildLeft => left
+ case BuildRight => right
+ }
+ val buildKeys = buildSide match {
+ case BuildLeft => leftKeys
+ case BuildRight => rightKeys
+ }
+ // Same index kind is not enough: a point index of column y
must not
+ // serve a join on column z. Mode equality ignores
nullability, and
+ // BroadcastPartitioning.satisfies uses the same ==, so a stage
+ // whose nullability AQE rewrote still satisfies the
distribution.
+ val bound = BindReferences.bindReferences(buildKeys,
buildPlan.output)
Review Comment:
On AQE replanning, `RangeBroadcastJoinStage` binds the build keys against
the logical output of the `LogicalQueryStage`, but the stage's mode was bound
against the physical exchange output. If the build side is a projection that
only reorders columns and `RemoveRedundantProjects` drops it, the ordinals no
longer line up, the strategy doesn't recognize its own stage, and the join goes
back to `JoinSelection`. When the index's runtime size is over the threshold at
that point, `ValidateSparkPlan` rejects the new plan and the re-optimization is
thrown away. By default that only shows up in DEBUG logs.
I ran `JOIN (SELECT hi, lo, name FROM t) r ON p.x BETWEEN r.lo AND r.hi` on
a parquet table, with the threshold set between the logical estimate and the
index size. The answer was correct, but AQE logged 2 `Invalid broadcast query
stage` failures. Matching `LogicalQueryStage(_, bqs)` in `fromMode` and binding
against `bqs.output` fixes it: the failures went away and the existing AQE
range tests pass.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/planning/patterns.scala:
##########
@@ -421,6 +422,235 @@ object ExtractSingleColumnNullAwareAntiJoin extends
JoinSelectionHelper with Pre
}
}
+/**
+ * Finds joins whose condition is a range predicate usable by a broadcast
range join.
+ *
+ * Recognizes point-in-range (`a.ip` against `[b.lo, b.hi]`), interval overlap
+ * (`a.lo < b.hi AND b.lo < a.hi`), and a cross-side inequality (`a.start <
b.end`).
+ * Point-in-range keys are the point once and `(low, high)`. The index probes
+ * that one value as both ends of the window. Overlap keys are `(low, high)`
+ * on each side. A partial range carries the one column from each side.
+ * Extra conjuncts stay on the original condition and are evaluated per
candidate.
+ * A point-in-range pair wins over overlap, and overlap wins over a partial
range.
+ * [[RangePredicate]] supplies the canonical `low <= high` form; inclusivity
stays
+ * on the original condition. Keys whose type the range index cannot compare,
and
+ * keys that are not deterministic, stay unrecognized, so the join stays a
nested
+ * loop. The index stores one evaluation of each key and the original condition
+ * evaluates it again; a second `rand()` would drop rows the condition accepts.
+ * A non-deterministic conjunct that is not a key stays on that condition.
+ * A deterministic [[With]] (`BETWEEN` lowers to one) is inlined first, so the
+ * keys are the common expressions. A non-deterministic definition is left in
+ * place and is not used as a range key.
+ */
+object ExtractRangeJoinKeys extends PredicateHelper {
+ type ReturnType = (LogicalPlan, LogicalPlan, Seq[Expression],
Seq[Expression],
+ JoinType, RangeJoin)
+
+ def unapply(plan: LogicalPlan): Option[ReturnType] = plan match {
+ case Join(left, right, joinType, Some(condition), _) =>
+ val predicates =
splitConjunctivePredicates(inlineCommonExpressions(condition))
+ val rangePreds = predicates.collect {
+ case RangePredicate(dt, low, high) if RangePredicate.supportedType(dt)
=>
+ (dt, low, high)
+ }
+ if (rangePreds.isEmpty) {
+ None
+ } else {
+ pointInRange(rangePreds, left, right, joinType)
+ .orElse(intervalOverlap(rangePreds, left, right, joinType))
+ .orElse(partialRange(rangePreds, left, right, joinType))
+ }
+ case _ => None
+ }
+
+ /**
+ * Replaces each deterministic [[With]] with its child, substituting every
+ * [[CommonExpressionRef]] by the common expression. A ref has no attributes,
+ * so `canEvaluate` would accept it on either side. A cycle or a ref with no
+ * definition in this [[With]] is left as a ref, and [[RangePredicate]] then
+ * rejects the predicate.
+ */
+ private def inlineCommonExpressions(expression: Expression): Expression = {
+ expression.transformUp {
+ case w: With if w.defs.forall(_.deterministic) =>
+ val defsById = w.defs.iterator.map(d => d.id -> d.child).toMap
+ def substitute(e: Expression, stack: Set[Long]): Expression = {
+ e.transform {
+ case r: CommonExpressionRef =>
+ defsById.get(r.id) match {
+ case Some(definition) if !stack.contains(r.id.id) =>
+ substitute(definition, stack + r.id.id)
+ case _ => r
+ }
+ }
+ }
+ substitute(w.child, Set.empty)
+ }
+ }
+
+ private def pointInRange(
+ rangePreds: Seq[(DataType, Expression, Expression)],
+ left: LogicalPlan,
+ right: LogicalPlan,
+ joinType: JoinType): Option[ReturnType] = {
+ rangePreds.combinations(2).flatMap {
+ case Seq((d1, l1, h1), (d2, l2, h2)) if d1 == d2 =>
+ val pointCandidate =
+ if (h1.semanticEquals(l2) && !l1.semanticEquals(h2) && l1.dataType
== h2.dataType) {
+ place(Seq(l2), Seq(l1, h2), left, right)
+ } else if (l1.semanticEquals(h2) && !h1.semanticEquals(l2) &&
+ l2.dataType == h1.dataType) {
+ place(Seq(l1), Seq(l2, h1), left, right)
+ } else {
+ None
+ }
+ pointCandidate.collect {
+ case (leftKeys, rightKeys) if supportedKeys(leftKeys, rightKeys) =>
+ (left, right, leftKeys, rightKeys, joinType, PointInRangeJoin)
+ }
+ case _ => None
+ }.nextOption()
+ }
+
+ /**
+ * Puts `leftExprs` on the left child when that child can evaluate all of
them
+ * and the right child can evaluate `rightExprs`. Otherwise swaps, when the
+ * opposite assignment fits. A key is a column of one input, so at most one
+ * assignment succeeds.
+ */
+ private def place(
+ leftExprs: Seq[Expression],
+ rightExprs: Seq[Expression],
+ left: LogicalPlan,
+ right: LogicalPlan): Option[(Seq[Expression], Seq[Expression])] = {
+ if (leftExprs.forall(canEvaluate(_, left)) &&
rightExprs.forall(canEvaluate(_, right))) {
+ Some((leftExprs, rightExprs))
+ } else if (leftExprs.forall(canEvaluate(_, right)) &&
+ rightExprs.forall(canEvaluate(_, left))) {
+ Some((rightExprs, leftExprs))
+ } else {
+ None
+ }
+ }
+
+ private def intervalOverlap(
+ rangePreds: Seq[(DataType, Expression, Expression)],
+ left: LogicalPlan,
+ right: LogicalPlan,
+ joinType: JoinType): Option[ReturnType] = {
+ rangePreds.combinations(2).flatMap {
+ case Seq((d1, l1, h1), (d2, l2, h2)) if d1 == d2 =>
+ // `(l1, h2)` is one interval and `(l2, h1)` is the other.
+ val overlapCandidate = place(Seq(l1, h2), Seq(l2, h1), left, right)
+ overlapCandidate.collect {
+ case (leftKeys, rightKeys)
+ if !leftKeys.head.semanticEquals(leftKeys(1)) &&
+ !rightKeys.head.semanticEquals(rightKeys(1)) &&
+ leftKeys.head.dataType == leftKeys(1).dataType &&
+ rightKeys.head.dataType == rightKeys(1).dataType &&
+ supportedKeys(leftKeys, rightKeys) =>
+ (left, right, leftKeys, rightKeys, joinType, IntervalOverlapJoin)
+ }
+ case _ => None
+ }.nextOption()
+ }
+
+ private def supportedKeys(keys: Seq[Expression]*): Boolean = {
+ keys.flatten.forall { key =>
+ key.references.nonEmpty && key.deterministic &&
RangePredicate.supportedType(key.dataType)
Review Comment:
With the flag on, a range key that contains an uncorrelated scalar subquery
makes the query fail. Take `ON p.x >= r.lo + (SELECT min(m) FROM t) AND p.x <=
r.hi`: `supportedKeys` accepts `r.lo + subquery` as a build key, so it goes
into `RangeBroadcastMode` and the driver evaluates it while building the index.
The subquery result is only set in the join's `waitForSubqueries`, and the
exchange can't see expressions inside a mode, so this throws `Subquery ... has
not finished`. I ran it and it fails with AQE both on and off. With the flag
off the same query runs as BNLJ and works.
Adding `!SubqueryExpression.hasSubquery(key)` to `supportedKeys` fixes it. I
tried that locally and the range suites pass. BHJ equi keys have the same gap
on master, which is unrelated to this PR and can go to a separate JIRA.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/joins/RangeIndex.scala:
##########
@@ -0,0 +1,329 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution.joins
+
+import scala.collection.mutable.ArrayBuffer
+
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{Expression,
InterpretedProjection}
+import org.apache.spark.sql.catalyst.types.PhysicalDataType
+import org.apache.spark.sql.types.DataType
+import org.apache.spark.util.SizeEstimator
+
+/**
+ * Broadcast payload of a range join. An index only names candidate rows.
+ * [[BroadcastRangeJoinExec]] evaluates the original join condition to filter
candidates.
+ *
+ * - [[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 estimatedSize(): Long
+ def keyType: DataType
+}
+
+private[execution] 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)
+ }
+}
+
+/**
+ * Augmented interval tree of build-side intervals, stored as flat arrays.
+ *
+ * Intervals are sorted by low bound. For any span `[start, end)`, its midpoint
+ * is the subtree root, and `maxHighs(mid)` records the maximum high bound in
+ * that span. Subtrees whose max high is below the probe window, or whose low
+ * starts after it, are pruned during search.
+ */
+private[execution] object IntervalIndex {
+
+ def build(
+ keyType: DataType,
+ buildKeys: Seq[Expression],
+ rows: Iterator[InternalRow],
+ sizeHint: Option[Long]): IntervalIndex = {
+ val initialCapacity = sizeHint.map(n => math.min(n,
Int.MaxValue).toInt).getOrElse(64)
+ val buffer = new ArrayBuffer[(Any, Any, InternalRow)](initialCapacity)
+ val projection = new InterpretedProjection(buildKeys)
+ val lowOf = RangeIndex.getValue(buildKeys.head.dataType, 0)
+ val highOf =
+ if (buildKeys.length == 1) lowOf else
RangeIndex.getValue(buildKeys(1).dataType, 1)
+ val ordering = PhysicalDataType.ordering(keyType)
+
+ while (rows.hasNext) {
+ val r = rows.next()
+ val p = projection(r)
+ normalizedInterval(ordering, lowOf(p), highOf(p), r, buffer)
+ }
+
+ sortAndBuild(keyType, ordering, buffer.toArray)
+ }
+
+ def build(
+ keyType: DataType,
+ intervals: Array[(Any, Any, InternalRow)]): IntervalIndex = {
+ val ordering = PhysicalDataType.ordering(keyType)
+ val buffer = new ArrayBuffer[(Any, Any, InternalRow)](intervals.length)
+ var i = 0
+ while (i < intervals.length) {
+ val (rawLo, rawHi, r) = intervals(i)
+ normalizedInterval(ordering, rawLo, rawHi, r, buffer)
+ i += 1
+ }
+ sortAndBuild(keyType, ordering, buffer.toArray)
+ }
+
+ /** Drops a null-bounded interval; otherwise orders (lo, hi) and copies the
key and row. */
+ private def normalizedInterval(
+ ordering: Ordering[Any],
+ rawLo: Any,
+ rawHi: Any,
+ r: InternalRow,
+ buffer: ArrayBuffer[(Any, Any, InternalRow)]): Unit = {
+ if (rawLo != null && rawHi != null) {
+ val lo = InternalRow.copyValue(rawLo)
Review Comment:
The `InternalRow.copyValue` fix from the last round doesn't have a
regression test that would catch it being dropped. I replaced the four
`copyValue` calls in `RangeIndex.scala` with the raw values, and all 44 tests
in `RangeIndexSuite`, `RangeJoinSuite` and `RangeJoinSQLSuite` still pass. The
two existing "reused row buffer" tests use `IntegerType` keys, and `copyValue`
is a no-op for those.
I wrote a throwaway test where the `UTF8String` keys point into a byte array
that the next row overwrites, which is what the typed ScalaUDF case from the
last round produces. Without `copyValue`, both the interval and the point index
miss the row they should return; with it they return it. Could you add a
string-key test like that to `RangeIndexSuite`, one per index kind?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala:
##########
@@ -410,8 +440,26 @@ abstract class SparkStrategies extends
QueryPlanner[SparkPlan] {
}
}
+ def createRangeJoin(hintOnly: Boolean) = {
+ if (!conf.broadcastRangeJoinEnabled) {
+ None
+ } else {
+ plan match {
+ case ExtractRangeJoinKeys(_, _, leftKeys, rightKeys, _,
rangeJoin) =>
+ val allowed = BroadcastRangeJoinExec.supports(joinType, _:
BuildSide)
+ broadcastSide(hintOnly, allowed).map { buildSide =>
+ Seq(BroadcastRangeJoinExec(
+ planLater(left), planLater(right), buildSide, joinType,
+ leftKeys, rightKeys, condition, rangeJoin))
+ }
+ case _ => None
+ }
+ }
+ }
+
def createJoinWithoutHint() = {
- createBroadcastNLJoin(false)
+ createRangeJoin(hintOnly = false)
Review Comment:
When neither side is under the threshold and there is no hint, left outer,
left semi and left anti joins fall through to the last-resort BNLJ in
`createJoinWithoutHint`. That branch broadcasts the right side without checking
the threshold, and the right side is exactly what range join supports, but it
never tries range join first. So with the flag on, these joins still run as
BNLJ when the build side is large. I checked: with the threshold at -1 and no
hint all three plan as BNLJ, and adding `BROADCAST(r)` turns them into range
joins.
Since that branch already broadcasts the same side and goes through the same
`maxBroadcastTableSize` check, it could try range join first for the
combinations `supports` allows. The catch is that the index estimate is larger
than the raw rows, so near the 8GB limit a BNLJ that runs today could fail with
`tooLargeToBroadcast`. I'd like to hear your view on that trade-off. If keeping
BNLJ there is intentional, a sentence in the PR description explaining why
would be enough.
--
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]
