Daily-Note 10-24

  /**
   * Transform a [[LogicalPlan]] into a [[SparkPlan]].
   *
   * Note that the returned physical plan still needs to be prepared for execution.
   */
  def createSparkPlan(
      sparkSession: SparkSession,
      planner: SparkPlanner,
      plan: LogicalPlan): SparkPlan = {
    // TODO: We use next(), i.e. take the first plan returned by the planner, here for now,
    //       but we will implement to choose the best plan.
    planner.plan(ReturnAnswer(plan)).next()
  }

def plan(plan: LogicalPlan): Iterator[PhysicalPlan] = {
    // Obviously a lot to do here still...

    // Collect physical plan candidates.
    val candidates = strategies.iterator.flatMap(_(plan))

    // The candidates may contain placeholders marked as [[planLater]],
    // so try to replace them by their child plans.
    val plans = candidates.flatMap { candidate =>
      val placeholders = collectPlaceholders(candidate)

      if (placeholders.isEmpty) {
        // Take the candidate as is because it does not contain placeholders.
        Iterator(candidate)
      } else {
        // Plan the logical plan marked as [[planLater]] and replace the placeholders.
        placeholders.iterator.foldLeft(Iterator(candidate)) {
          case (candidatesWithPlaceholders, (placeholder, logicalPlan)) =>
            // Plan the logical plan for the placeholder.
            val childPlans = this.plan(logicalPlan)

            candidatesWithPlaceholders.flatMap { candidateWithPlaceholders =>
              childPlans.map { childPlan =>
                // Replace the placeholder by the child plan
                candidateWithPlaceholders.transformUp {
                  case p if p.eq(placeholder) => childPlan
                }
              }
            }
        }
      }
    }

override def strategies: Seq[Strategy] =
  experimentalMethods.extraStrategies ++
    extraPlanningStrategies ++ (
    LogicalQueryStageStrategy ::
    PythonEvals ::
    new DataSourceV2Strategy(session) ::
    FileSourceStrategy ::
    DataSourceStrategy ::
    SpecialLimits ::
    Aggregation ::
    Window ::
    JoinSelection ::
    InMemoryScans ::
    BasicOperators :: Nil)

case PhysicalAggregation(groupingExpressions, aggExpressions, resultExpressions, child)
if aggExpressions.forall(expr => expr.isInstanceOf[AggregateExpression]) =>
	// ...
  val aggregateOperator =
  if (functionsWithDistinct.isEmpty) {
    AggUtils.planAggregateWithoutDistinct(
      normalizedGroupingExpressions,
      aggregateExpressions,
      resultExpressions,
      planLater(child))
  } else {
    // ...
  }
  aggregateOperator

object JoinSelection extends Strategy
  with PredicateHelper
  with JoinSelectionHelper {

  def apply(plan: LogicalPlan): Seq[SparkPlan] = plan match {
    // If it is an equi-join, we first look at the join hints w.r.t. the following order:
    //   1. broadcast hint: pick broadcast hash join if the join type is supported. If both sides
    //      have the broadcast hints, choose the smaller side (based on stats) to broadcast.
    //   2. sort merge hint: pick sort merge join if join keys are sortable.
    //   3. shuffle hash hint: We pick shuffle hash join if the join type is supported. If both
    //      sides have the shuffle hash hints, choose the smaller side (based on stats) as the
    //      build side.
    //   4. shuffle replicate NL hint: pick cartesian product if join type is inner like.
    //
    // If there is no hint or the hints are not applicable, we follow these rules one by one:
    //   1. Pick broadcast hash join if one side is small enough to broadcast, and the join type
    //      is supported. If both sides are small, choose the smaller side (based on stats)
    //      to broadcast.
    //   2. Pick shuffle hash join if one side is small enough to build local hash map, and is
    //      much smaller than the other side, and `spark.sql.join.preferSortMergeJoin` is false.
    //   3. Pick sort merge join if the join keys are sortable.
    //   4. Pick cartesian product if join type is inner like.
    //   5. Pick broadcast nested loop join as the final solution. It may OOM but we don't have
    //      other choice.
      case j @ ExtractEquiJoinKeys(joinType, leftKeys, rightKeys, nonEquiCond, left, right, hint) =>
    // ...
    
    // If it is not an equi-join, we first look at the join hints w.r.t. the following order:
    //   1. broadcast hint: pick broadcast nested loop join. If both sides have the broadcast
    //      hints, choose the smaller side (based on stats) to broadcast for inner and full joins,
    //      choose the left side for right join, and choose right side for left join.
    //   2. shuffle replicate NL hint: pick cartesian product if join type is inner like.
    //
    // If there is no hint or the hints are not applicable, we follow these rules one by one:
    //   1. Pick broadcast nested loop join if one side is small enough to broadcast. If only left
    //      side is broadcast-able and it's left join, or only right side is broadcast-able and
    //      it's right join, we skip this rule. If both sides are small, broadcasts the smaller
    //      side for inner and full joins, broadcasts the left side for right join, and broadcasts
    //      right side for left join.
    //   2. Pick cartesian product if join type is inner like.
    //   3. Pick broadcast nested loop join as the final solution. It may OOM but we don't have
    //      other choice. It broadcasts the smaller side for inner and full joins, broadcasts the
    //      left side for right join, and broadcasts right side for left join.
      case logical.Join(left, right, joinType, condition, hint) =>
    // ...
  }

createBroadcastHashJoin(true)
  .orElse { if (hintToSortMergeJoin(hint)) createSortMergeJoin() else None }
  .orElse(createShuffleHashJoin(true))
  .orElse { if (hintToShuffleReplicateNL(hint))
           createCartesianProduct() else None }
  .getOrElse(createJoinWithoutHint())

def createBroadcastHashJoin(onlyLookingAtHint: Boolean) = {
  getBroadcastBuildSide(
    left, right, joinType, hint, onlyLookingAtHint, conf).map {
    buildSide =>
      Seq(joins.BroadcastHashJoinExec(
        leftKeys,
        rightKeys,
        joinType,
        buildSide,
        nonEquiCond,
        planLater(left),
        planLater(right)))
  }
}

def createJoinWithoutHint() = {
  createBroadcastHashJoin(false) // 没有 hint 条件下创建 BroadHashJoin
    .orElse {
      if (!conf.preferSortMergeJoin) { // prefer 是否指定 SortMerge 优先
        createShuffleHashJoin(false)
      } else {
        None
      }
    }
    .orElse(createSortMergeJoin())
    .orElse(createCartesianProduct())
    .getOrElse {
      // This join could be very slow or OOM
      val buildSide = getSmallerSide(left, right)
      Seq(joins.BroadcastNestedLoopJoinExec( // NestedLoopJoin 策略为兜底
        planLater(left), planLater(right), buildSide, joinType, nonEquiCond))
    }
}

// joins.scala

val buildLeft = if (hintOnly) {
  hintToBroadcastLeft(hint)
} else {
  canBroadcastBySize(left, conf) && !hintToNotBroadcastLeft(hint)
}

def canBroadcastBySize(plan: LogicalPlan, conf: SQLConf): Boolean = {
  plan.stats.sizeInBytes >= 0 &&
  plan.stats.sizeInBytes <= conf.autoBroadcastJoinThreshold
}

private def canBuildLocalHashMapBySize(plan: LogicalPlan, conf: SQLConf): Boolean = {
  plan.stats.sizeInBytes < conf.autoBroadcastJoinThreshold * conf.numShufflePartitions
}