11. Adaptive Query Execution

var result = createQueryStages(fun, currentPhysicalPlan, firstRun = true)
val events = new LinkedBlockingQueue[StageMaterializationEvent]()
while (!result.allChildStagesMaterialized) {
  currentPhysicalPlan = result.newPlan
  if (result.newStages.nonEmpty) {
    // SPARK-33933: Broadcast Stage를 먼저 제출 (Broadcast timeout 방지)
    val reorderedNewStages = result.newStages.sortWith { ... BroadcastQueryStageExec 우선 ... }
    reorderedNewStages.foreach { stage =>
      stage.materialize().onComplete { res => ... }   // 비동기 제출
    }
  }
  // 다음 Stage가 끝날 때까지 대기 → 새 런타임 통계
  val nextMsg = events.take()
  events.drainTo(rem)                                  // 비슷한 시점에 끝난 Stage도 함께 처리
  (Seq(nextMsg) ++ rem.asScala).foreach {
    case StageSuccess(stage, res) => stage.resultOption.set(Some(res))   // = MapOutputStatistics
    case StageFailure(stage, ex) => ...
  }

  val afterReOptimize = reOptimize(logicalPlan)
  if (afterReOptimize.isDefined) {
    val (newPhysicalPlan, newLogicalPlan) = afterReOptimize.get
    val origCost = costEvaluator.evaluateCost(currentPhysicalPlan)
    val newCost = costEvaluator.evaluateCost(newPhysicalPlan)
    if (newCost < origCost ||
        (newCost == origCost && currentPhysicalPlan != newPhysicalPlan)) {
      currentPhysicalPlan = newPhysicalPlan
      currentLogicalPlan = newLogicalPlan
    }
  }
  result = createQueryStages(...)
}

def computeStats(): Option[Statistics] = if (isMaterialized) {
  val runtimeStats = getRuntimeStatistics
  Some(Statistics(dataSize, numOutputRows, attributeStats, isRuntime = true))
} else { None }

override protected def doMaterialize(): Future[Any] = shuffle.submitShuffleJob()

case class SimpleCostEvaluator(forceOptimizeSkewedJoin: Boolean) extends CostEvaluator {
  override def evaluateCost(plan: SparkPlan): Cost = {
    val numShuffles = plan.collect { case s: ShuffleExchangeLike => s }.size
    if (forceOptimizeSkewedJoin) {
      val numSkewJoins = plan.collect { case j: ShuffledJoin if j.isSkewJoin => j }.size
      // Skew Join 수를 상위 32비트에 음수로 → Skew 처리가 많을수록 저비용
      SimpleCost(-numSkewJoins.toLong << 32 | numShuffles)
    } else {
      SimpleCost(numShuffles)
    }
  }
}

def numShufflePartitions: Int = {
  if (adaptiveExecutionEnabled && coalesceShufflePartitionsEnabled) {
    getConf(COALESCE_PARTITIONS_INITIAL_PARTITION_NUM).getOrElse(defaultNumShufflePartitions)
  } else { defaultNumShufflePartitions }
}

val minNumPartitions = conf.getConf(SQLConf.COALESCE_PARTITIONS_MIN_PARTITION_NUM).getOrElse {
  if (conf.getConf(SQLConf.COALESCE_PARTITIONS_PARALLELISM_FIRST)) {
    session.sparkContext.defaultParallelism   // 기본: 병렬도 우선
  } else { 1 }
}

val totalPostShuffleInputSize = mapOutputStatistics.flatMap(_.map(_.bytesByPartitionId.sum)).sum
val maxTargetSize = math.ceil(totalPostShuffleInputSize / minNumPartitions.toDouble).toLong
val targetSize = maxTargetSize.min(advisoryTargetSize).max(minPartitionSize)

while (i < end) {
  var totalSizeOfCurrentPartition = 0L
  var j = 0
  while (j < mapOutputStatistics.length) {     // 같은 그룹의 모든 Shuffle에서 i번 파티션 크기 합산
    totalSizeOfCurrentPartition += mapOutputStatistics(j).bytesByPartitionId(i)
    j += 1
  }
  if (i > latestSplitPoint && coalescedSize + totalSizeOfCurrentPartition > targetSize) {
    if (coalescedSize < minPartitionSize) {
      // 지금 묶음이 1MB 미만이면 앞뒤 중 작은 쪽에 붙임
      ...
    } else {
      createPartitionSpec()                     // 새 묶음 시작
    }
  } else {
    coalescedSize += totalSizeOfCurrentPartition
  }
  i += 1
}

/** A partition is considered as a skewed partition if its size is larger than the median
 *  partition size * SKEW_JOIN_SKEWED_PARTITION_FACTOR and also larger than
 *  SKEW_JOIN_SKEWED_PARTITION_THRESHOLD. */
def getSkewThreshold(medianSize: Long): Long = {
  conf.getConf(SQLConf.SKEW_JOIN_SKEWED_PARTITION_THRESHOLD).max(
    (medianSize * conf.getConf(SQLConf.SKEW_JOIN_SKEWED_PARTITION_FACTOR)).toLong)
}

private def targetSize(sizes: Array[Long], skewThreshold: Long): Long = {
  val nonSkewSizes = sizes.filter(_ <= skewThreshold)
  math.max(advisorySize, nonSkewSizes.sum / nonSkewSizes.length)
}

val leftParts = if (isLeftSkew) {
  ShufflePartitionsUtil.createSkewPartitionSpecs(left.mapStats.get.shuffleId, partitionIndex, leftTargetSize)
  ...
} else { leftNoSkewPartitionSpec }
...
for {
  leftSidePartition <- leftParts
  rightSidePartition <- rightParts
} {
  leftSidePartitions += leftSidePartition
  rightSidePartitions += rightSidePartition
}

def createSkewPartitionSpecs(shuffleId: Int, reducerId: Int, targetSize: Long, ...)
    : Option[Seq[PartialReducerPartitionSpec]] = {
  val mapPartitionSizes = getMapSizesForReduceId(shuffleId, reducerId)
  val mapStartIndices = splitSizeListByTargetSize(mapPartitionSizes, targetSize, ...)
  if (mapStartIndices.length > 1) {
    Some(... PartialReducerPartitionSpec(reducerId, startMapIndex, endMapIndex, dataSize) ...)
  } else { None }
}

val optimized = optimizeSkewJoin(plan)
val requirementSatisfied = ValidateRequirements.validate(optimized, ...)
if (requirementSatisfied) { optimized... }
else if (conf.getConf(SQLConf.ADAPTIVE_FORCE_OPTIMIZE_SKEWED_JOIN)) { ... }
else { plan }      // 추가 Shuffle이 필요해지면 Skew 최적화를 통째로 포기

private def hasManyEmptyPartitions(mapStats: MapOutputStatistics): Boolean = {
  val partitionCnt = mapStats.bytesByPartitionId.length
  val nonZeroCnt = mapStats.bytesByPartitionId.count(_ > 0)
  partitionCnt > 0 && nonZeroCnt > 0 &&
    (nonZeroCnt * 1.0 / partitionCnt) < conf.nonEmptyPartitionRatioForBroadcastJoin  // 0.2
}

val dataframe = Dataset.ofRows(
  sessionHolder.session,
  planner.transformRelation(request.getPlan.getRoot, cachePlan = true),
  tracker, shuffleCleanupMode)
processAsArrowBatches(dataframe, responseObserver, executeHolder)