wangyum commented on code in PR #59027:
URL: https://github.com/apache/spark/pull/59027#discussion_r4273213742
##########
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:
fixed and add a test
--
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]
