wangyum commented on code in PR #59027:
URL: https://github.com/apache/spark/pull/59027#discussion_r4273211410
##########
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:
good catch.
--
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]
