10. Catalyst and Whole-Stage Codegen

// SparkSession.sql (4.x에서 구현체는 classic 패키지)
val plan = tracker.measurePhase(QueryPlanningTracker.PARSING) {
  sessionState.sqlParser.parsePlanWithParameters(sqlText, paramContext)
}
Dataset.ofRows(self, plan, tracker)

private val lazyAnalyzed = LazyTry {
  val plan = executePhase(QueryPlanningTracker.ANALYSIS) {
    analyzer.executeAndCheck(sqlScriptExecuted, tracker)
  }
  tracker.setAnalyzed(plan); plan
}

private val lazyOptimizedPlan = LazyTry {
  executePhase(QueryPlanningTracker.OPTIMIZATION) {
    // clone the plan to avoid sharing the plan instance between different stages
    sparkSession.sessionState.optimizer.executeAndTrack(withCachedData.clone(), tracker)
  }
}

def createSparkPlan(planner: SparkPlanner, plan: LogicalPlan): SparkPlan = {
  // TODO: We use next(), i.e. take the first plan returned by the planner
  planner.plan(ReturnAnswer(plan)).next()
}

adaptiveExecutionRule.toSeq ++     // InsertAdaptiveSparkPlan (11장)
Seq(
  CoalesceBucketsInJoin,
  PlanDynamicPruningFilters(sparkSession),
  PlanSubqueries(sparkSession),
  RemoveRedundantProjects,
  EnsureRequirements(),            // Exchange · Sort 삽입 (10.7)
  ...
  CollapseCodegenStages()) ++      // Whole-Stage Codegen (10.8)
(if (subquery) Nil else Seq(ReuseExchangeAndSubquery))

df.explain("extended")   # Parsed / Analyzed / Optimized Logical + Physical
df.explain("formatted")  # Physical Plan을 노드별 상세와 함께
df.explain("codegen")    # 생성된 Java 코드

abstract class Rule[TreeType <: TreeNode[_]] extends SQLConfHelper with Logging {
  protected lazy val ruleId = RuleIdCollection.getRuleId(this.ruleName)
  def apply(plan: TreeType): TreeType
}

def transformDownWithPruning(cond: TreePatternBits => Boolean, ruleId: RuleId = UnknownRuleId)
    (rule: PartialFunction[BaseType, BaseType]): BaseType = {
  if (!cond.apply(this) || isRuleIneffective(ruleId)) return this   // 볼 필요 없는 서브트리 건너뜀
  val afterRule = rule.applyOrElse(this, identity[BaseType])        // 매치 안 되면 그대로
  if (this fastEquals afterRule) {
    val rewritten = mapChildren(_.transformDownWithPruning(cond, ruleId)(rule))
    if (this eq rewritten) { markRuleAsIneffective(ruleId); this }  // 효과 없던 Rule 기억
    else rewritten
  } else {
    afterRule.mapChildren(_.transformDownWithPruning(cond, ruleId)(rule))
  }
}

case object Once extends Strategy { val maxIterations = 1 }
case class FixedPoint(override val maxIterations: Int, ...) extends Strategy
case class Batch(name: String, strategy: Strategy, rules: Rule[TreeType]*)

while (continue) {
  curPlan = batch.rules.foldLeft(curPlan) { case (plan, rule) => rule(plan) }
  iteration += 1
  if (iteration > batch.strategy.maxIterations) continue = false
  if (curPlan.fastEquals(lastPlan)) continue = false     // 고정점 도달
  lastPlan = curPlan
}

val operatorOptimizationBatch: Seq[Batch] = Seq(
  Batch("Operator Optimization before Inferring Filters", fixedPoint, operatorOptimizationRuleSet: _*),
  Batch("Infer Filters", Once, InferFiltersFromGenerate, InferFiltersFromConstraints),
  Batch("Operator Optimization after Inferring Filters", fixedPoint, operatorOptimizationRuleSet: _*),
  Batch("Push extra predicate through join", fixedPoint, PushExtraPredicateThroughJoin, PushDownPredicates))

object PushDownPredicates extends Rule[LogicalPlan] {
  def apply(plan: LogicalPlan): LogicalPlan = plan.transformWithPruning(
    _.containsAnyPattern(FILTER, JOIN)) {
    CombineFilters.applyLocally
      .orElse(PushPredicateThroughNonJoin.applyLocally)
      .orElse(PushPredicateThroughJoin.applyLocally)
  }
}

val (candidates, nonDeterministic) =
  splitConjunctivePredicates(filter.condition).partition(_.deterministic)
val (pushDown, rest) = candidates.partition { cond =>
  cond.references.subsetOf(grandchild.outputSet)
}
val stayUp = rest ++ nonDeterministic

case p @ Project(_, p2: Project) if !p2.outputSet.subsetOf(p.references) =>
  p.copy(child = p2.copy(projectList = p2.projectList.filter(p.references.contains)))

/** Applies a projection only when the child is producing unnecessary attributes */
private def prunedChild(c: LogicalPlan, allReferences: AttributeSet) =
  if (!c.outputSet.subsetOf(allReferences)) {
    Project(c.output.filter(allReferences.contains), c)
  } else { c }

private def tryFold(expr: Expression, isConditionalBranch: Boolean): Expression = {
  try {
    Literal.create(expr.freshCopyIfContainsStatefulExpression().eval(EmptyRow), expr.dataType)
  } catch {
    case NonFatal(_) if isConditionalBranch =>
      // 실행되지 않을 수도 있는 CASE WHEN 분기에서 실패하면 쿼리 실패로 보지 않음
      expr.setTagValue(FAILED_TO_EVALUATE, ()); expr
  }
}

if (!conf.cboEnabled || !conf.joinReorderEnabled) { plan } else {
  if (items.size > 2 && items.size <= conf.joinReorderDPThreshold && conditions.nonEmpty &&
      items.forall(_.stats.rowCount.isDefined)) {
    JoinReorderDP.search(conf, items, conditions, output)
  }
}

Math.pow(relativeRows, conf.joinReorderCardWeight) *
  Math.pow(relativeSize, 1 - conf.joinReorderCardWeight) < 1

// 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 ...
//   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.

def createJoinWithoutHint() = {
  createBroadcastHashJoin(false)
    .orElse(createShuffleHashJoin(false))
    .orElse(createSortMergeJoin())
    .orElse(createCartesianProduct())
    .getOrElse { BroadcastNestedLoopJoinExec(...) }
}

def canBroadcastBySize(plan: LogicalPlan, conf: SQLConf): Boolean = {
  val autoBroadcastJoinThreshold = if (plan.stats.isRuntime) {
    conf.getConf(SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD)
      .getOrElse(conf.autoBroadcastJoinThreshold)
  } else { conf.autoBroadcastJoinThreshold }
  plan.stats.sizeInBytes >= 0 && plan.stats.sizeInBytes <= autoBroadcastJoinThreshold
}

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

private def muchSmaller(a: LogicalPlan, b: LogicalPlan, conf: SQLConf): Boolean =
  a.stats.sizeInBytes * conf.getConf(SQLConf.SHUFFLE_HASH_JOIN_FACTOR) <= b.stats.sizeInBytes

// ShuffledJoin (SMJ·SHJ 공통): 양쪽이 Join Key 기준으로 모여 있어야 함
override def requiredChildDistribution: Seq[Distribution] =
  ClusteredDistribution(leftKeys) :: ClusteredDistribution(rightKeys) :: Nil

// SortMergeJoinExec: 파티션 안에서 Key 순으로 정렬돼 있어야 함
override def requiredChildOrdering: Seq[Seq[SortOrder]] =
  requiredOrders(leftKeys) :: requiredOrders(rightKeys) :: Nil

distribution match {
  case BroadcastDistribution(mode) => BroadcastExchangeExec(mode, child)
  case _ => ShuffleExchangeExec(distribution.createPartitioning(numPartitions), child, shuffleOrigin)
}
...
// Now that we've performed any necessary shuffles, add sorts:
children = children.zip(requiredChildOrderings).map { case (child, requiredOrdering) =>
  if (SortOrder.orderingSatisfies(child.outputOrdering, requiredOrdering)) child
  else SortExec(requiredOrdering, global = false, child = child)
}

s"""
   | while ($limitNotReachedCond $input.hasNext()) {
   |   InternalRow $row = (InternalRow) $input.next();
   |   ${consume(ctx, outputVars, ...).trim}
   |   ${shouldStopCheckCode}
   | }
 """

s"""
   |do {
   |  $predicateCode        // 조건 불만족이면 continue
   |  $numOutput.add(1);
   |  ${consume(ctx, resultVars)}
   |} while (false);
 """

while (input.hasNext()) {
  InternalRow row = (InternalRow) input.next();
  long amount = row.getLong(2);
  if (!(amount > 100)) continue;           // Filter
  // Project: 필요한 컬럼만 결과 버퍼에 기록
  rowWriter.write(0, row.getUTF8String(1));
  rowWriter.write(1, amount);
  append(rowWriter.getRow());
}

final class $className extends BufferedRowIterator {
  private Object[] references;
  private scala.collection.Iterator[] inputs;
  ...
  protected void processNext() throws java.io.IOException { ${code.trim} }
}

val (_, compiledCodeStats) = try {
  CodeGenerator.compile(cleanedSource)
} catch {
  case NonFatal(_) if !Utils.isTesting && conf.codegenFallback =>
    return child.execute()               // 컴파일 실패 → 인터프리트 방식으로
}
if (compiledCodeStats.maxMethodCodeSize > conf.hugeMethodLimit) {
  return child.execute()                 // 메서드가 너무 크면 JIT이 안 됨 → 포기
}

private def supportCodegen(plan: SparkPlan): Boolean = plan match {
  case plan: CodegenSupport if plan.supportCodegen =>
    val willFallback = plan.expressions.exists(_.exists(e => !supportCodegen(e)))
    val hasTooManyOutputFields = WholeStageCodegenExec.isTooManyFields(conf, plan.schema)
    ...
    !willFallback && !hasTooManyOutputFields && !hasTooManyInputFields
  case _ => false
}

private Object baseObject;
private long baseOffset;

private long getFieldOffset(int ordinal) {
  return baseOffset + bitSetWidthInBytes + ordinal * 8L;
}
public long getLong(int ordinal) {
  return Platform.getLong(baseObject, getFieldOffset(ordinal));   // 주소 계산 한 번, boxing 없음
}
public UTF8String getUTF8String(int ordinal) {
  final long offsetAndSize = getLong(ordinal);
  final int offset = (int) (offsetAndSize >> 32);
  final int size = (int) offsetAndSize;
  return UTF8String.fromAddress(baseObject, baseOffset + offset, size);  // 같은 버퍼를 가리키는 view
}