// 创建RDD
val rdd = sc.parallelize(1 to 100, 4)

// map：一对一转换
val mapRDD = rdd.map(x => x * 2)

// filter：过滤数据
val filterRDD = rdd.filter(x => x % 2 == 0)

// flatMap：一对多转换
val flatMapRDD = rdd.flatMap(x => 1 to x)

// groupByKey：按键分组
val kvRDD = rdd.map(x => (x % 10, x))
val groupedRDD = kvRDD.groupByKey()

// reduceByKey：按键聚合
val reducedRDD = kvRDD.reduceByKey(_ + _)

// join：连接操作
val rdd2 = sc.parallelize(Seq((1, "a"), (2, "b"), (3, "c")))
val joinedRDD = kvRDD.join(rdd2)

// collect：收集所有元素到Driver
val result = rdd.collect()

// count：计算元素数量
val cnt = rdd.count()

// first：获取第一个元素
val firstElement = rdd.first()

// take：获取前n个元素
val firstN = rdd.take(10)

// reduce：聚合所有元素
val sum = rdd.reduce(_ + _)

// foreach：遍历每个元素
rdd.foreach(println)

// saveAsTextFile：保存到文件
rdd.saveAsTextFile("hdfs://output/path")

// map, filter, union等操作产生窄依赖
val rdd1 = sc.parallelize(1 to 10, 2)
val rdd2 = rdd1.map(_ * 2)        // 窄依赖
val rdd3 = rdd2.filter(_ > 10)    // 窄依赖

// groupByKey, reduceByKey, join等操作产生宽依赖
val rdd1 = sc.parallelize(Seq((1, "a"), (2, "b"), (1, "c")), 2)
val rdd2 = rdd1.groupByKey()      // 宽依赖，需要Shuffle
val rdd3 = rdd1.reduceByKey(_ + _) // 宽依赖，需要Shuffle

import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.types._

val spark = SparkSession.builder()
  .appName("DataFrameExample")
  .getOrCreate()

import spark.implicits._

// 方式1：从RDD创建
val rdd = sc.parallelize(Seq(("Alice", 25), ("Bob", 30), ("Charlie", 35)))
val df1 = rdd.toDF("name", "age")

// 方式2：从序列创建
val df2 = Seq(("Alice", 25), ("Bob", 30)).toDF("name", "age")

// 方式3：从外部数据源创建
val df3 = spark.read.json("path/to/file.json")
val df4 = spark.read.parquet("path/to/file.parquet")

// 方式4：通过Schema创建
val schema = StructType(Seq(
  StructField("name", StringType, nullable = true),
  StructField("age", IntegerType, nullable = true)
))
val df5 = spark.createDataFrame(rdd, schema)

// 定义样例类
case class Person(name: String, age: Int, city: String)

// 方式1：从序列创建
val ds1 = Seq(
  Person("Alice", 25, "Beijing"),
  Person("Bob", 30, "Shanghai")
).toDS()

// 方式2：从DataFrame转换
val ds2 = df.as[Person]

// 方式3：从外部数据源创建
val ds3 = spark.read.json("path/to/file.json").as[Person]

// 性能测试示例
import org.apache.spark.sql.functions._

// RDD方式 - 性能较低
val rddResult = rdd.filter(_.age > 25)
  .map(p => (p.city, 1))
  .reduceByKey(_ + _)
  .collect()

// DataFrame方式 - 性能优化
val dfResult = df.filter($"age" > 25)
  .groupBy("city")
  .count()
  .collect()

// Dataset方式 - 类型安全 + 性能优化
val dsResult = ds.filter(_.age > 25)
  .groupByKey(_.city)
  .count()
  .collect()

// 创建带分区的RDD
val rdd = sc.parallelize(1 to 100, 4)  // 4个分区

// 查看分区信息
println(s"分区数: ${rdd.getNumPartitions}")
println(s"分区内容: ${rdd.glom().collect().map(_.toList).toList}")

// 重新分区
val repartitionedRDD = rdd.repartition(8)  // 增加分区数
val coalescedRDD = rdd.coalesce(2)         // 减少分区数

// Key-Value RDD分区
val kvRDD = sc.parallelize(Seq((1, "a"), (2, "b"), (3, "c"), (4, "d")), 2)

// 使用HashPartitioner
val hashPartitioned = kvRDD.partitionBy(new HashPartitioner(3))

// 使用RangePartitioner
val rangePartitioned = kvRDD.partitionBy(new RangePartitioner(3, kvRDD))

import org.apache.spark.Partitioner

// 自定义分区器：按用户ID的地区分区
class RegionPartitioner(regions: Array[String]) extends Partitioner {
  
  override def numPartitions: Int = regions.length
  
  override def getPartition(key: Any): Int = {
    val userId = key.asInstanceOf[String]
    val region = getUserRegion(userId)
    math.abs(regions.indexOf(region)) % numPartitions
  }
  
  private def getUserRegion(userId: String): String = {
    // 根据用户ID确定地区的业务逻辑
    userId.substring(0, 2) match {
      case "01" | "02" => "North"
      case "03" | "04" => "South"
      case "05" | "06" => "East"
      case _ => "West"
    }
  }
}

// 使用自定义分区器
val regions = Array("North", "South", "East", "West")
val customPartitioner = new RegionPartitioner(regions)
val customPartitioned = kvRDD.partitionBy(customPartitioner)

// 分区数设置原则
val totalCores = 16  // 集群总核心数
val optimalPartitions = totalCores * 2  // 推荐分区数为核心数的2-3倍

// 动态调整分区数
def getOptimalPartitions(dataSize: Long): Int = {
  val targetPartitionSize = 128 * 1024 * 1024  // 128MB per partition
  math.max(1, (dataSize / targetPartitionSize).toInt)
}

// 分区倾斜检测
def detectPartitionSkew(rdd: RDD[_]): Unit = {
  val partitionSizes = rdd.mapPartitionsWithIndex { (index, iter) =>
    Iterator((index, iter.size))
  }.collect()
  
  val avgSize = partitionSizes.map(_._2).sum / partitionSizes.length
  val maxSize = partitionSizes.map(_._2).max
  val skewRatio = maxSize.toDouble / avgSize
  
  if (skewRatio > 2.0) {
    println(s"警告：检测到分区倾斜，倾斜比例: $skewRatio")
    partitionSizes.foreach { case (index, size) =>
      println(s"分区 $index: $size 条记录")
    }
  }
}

// 根据数据特征预分区
val userRDD = sc.textFile("hdfs://users/*")
  .map(parseUser)
  .partitionBy(new HashPartitioner(numPartitions))
  .cache()  // 缓存预分区的数据

// Coalesce：减少分区，避免全量Shuffle
val reducedRDD = largeRDD.coalesce(10)

// Repartition：重新分区，会进行全量Shuffle
val reshuffledRDD = largeRDD.repartition(20)

// 条件分区调整
def smartRepartition[T](rdd: RDD[T], targetPartitions: Int): RDD[T] = {
  val currentPartitions = rdd.getNumPartitions
  if (targetPartitions < currentPartitions) {
    rdd.coalesce(targetPartitions)
  } else {
    rdd.repartition(targetPartitions)
  }
}

// 使用mapPartitions保持分区结构
val optimizedRDD = rdd.mapPartitions { iter =>
  // 分区内处理逻辑
  iter.map(processRecord)
}

// 避免破坏分区的操作
val goodRDD = partitionedRDD.mapValues(_ * 2)  // 保持分区
val badRDD = partitionedRDD.map(x => (x._1, x._2 * 2))  // 可能破坏分区

// SparkContext核心功能
class SparkContext(config: SparkConf) extends Logging {
  
  // 1. 初始化核心组件
  private val env = SparkEnv.createDriverEnv(conf, isLocal, listenerBus, numCores, mockOutputCommitCoordinator)
  private val statusTracker = new SparkStatusTracker(this, sparkUI)
  private val taskScheduler = createTaskScheduler(this, master, deployMode)
  private val dagScheduler = new DAGScheduler(this)
  
  // 2. 创建RDD
  def parallelize[T: ClassTag](seq: Seq[T], numSlices: Int = defaultParallelism): RDD[T] = {
    new ParallelCollectionRDD[T](this, seq, numSlices, Map[Int, Seq[String]]())
  }
  
  def textFile(path: String, minPartitions: Int = defaultMinPartitions): RDD[String] = {
    hadoopFile(path, classOf[TextInputFormat], classOf[LongWritable], classOf[Text], minPartitions)
      .map(pair => pair._2.toString)
  }
  
  // 3. 提交作业
  def runJob[T, U: ClassTag](
      rdd: RDD[T],
      func: (TaskContext, Iterator[T]) => U,
      partitions: Seq[Int],
      resultHandler: (Int, U) => Unit): Unit = {
    dagScheduler.runJob(rdd, func, partitions, callSite, resultHandler, localProperties.get)
  }
  
  // 4. 资源管理
  def stop(): Unit = {
    dagScheduler.stop()
    taskScheduler.stop()
    env.stop()
  }
}

// Driver程序示例
object WordCount {
  def main(args: Array[String]): Unit = {
    // 1. 创建SparkContext
    val conf = new SparkConf().setAppName("WordCount")
    val sc = new SparkContext(conf)
    
    try {
      // 2. 创建RDD并定义转换操作
      val lines = sc.textFile(args(0))
      val words = lines.flatMap(_.split("\\s+"))
      val wordCounts = words.map((_, 1)).reduceByKey(_ + _)
      
      // 3. 触发Action，提交作业
      wordCounts.saveAsTextFile(args(1))
      
    } finally {
      // 4. 停止SparkContext
      sc.stop()
    }
  }
}

// YARN Client模式
spark-submit \
  --master yarn \
  --deploy-mode client \
  --num-executors 10 \
  --executor-memory 4g \
  --executor-cores 2 \
  --class com.example.MyApp \
  myapp.jar

// YARN Cluster模式  
spark-submit \
  --master yarn \
  --deploy-mode cluster \
  --num-executors 10 \
  --executor-memory 4g \
  --executor-cores 2 \
  --class com.example.MyApp \
  myapp.jar

// Executor核心组件
class Executor(
    executorId: String,
    executorHostname: String,
    env: SparkEnv,
    userClassPath: Seq[URL] = Nil,
    isLocal: Boolean = false)
  extends Logging {

  // 1. 线程池管理
  private val threadPool = ThreadUtils.newDaemonCachedThreadPool(
    "Executor task launch worker", sparkConf.get(EXECUTOR_CORES), 60)
  
  // 2. 内存管理
  private val memoryManager = env.memoryManager
  
  // 3. 存储管理
  private val blockManager = env.blockManager
  
  // 4. 任务执行
  def launchTask(context: ExecutorBackend, taskDescription: TaskDescription): Unit = {
    val tr = new TaskRunner(context, taskDescription)
    runningTasks.put(taskDescription.taskId, tr)
    threadPool.execute(tr)
  }
  
  // 5. 任务运行器
  class TaskRunner(
      execBackend: ExecutorBackend,
      private val taskDescription: TaskDescription)
    extends Runnable {
    
    override def run(): Unit = {
      try {
        // 反序列化任务
        val task = ser.deserialize[Task[Any]](taskDescription.serializedTask, Thread.currentThread.getContextClassLoader)
        
        // 执行任务
        val value = task.run(
          taskAttemptId = taskDescription.taskId,
          attemptNumber = taskDescription.attemptNumber,
          metricsSystem = env.metricsSystem)
        
        // 返回结果
        execBackend.statusUpdate(taskDescription.taskId, TaskState.FINISHED, ser.serialize(value))
        
      } catch {
        case e: Exception =>
          execBackend.statusUpdate(taskDescription.taskId, TaskState.FAILED, ser.serialize(TaskFailedReason))
      }
    }
  }
}

// DAGScheduler事件类型定义
sealed trait DAGSchedulerEvent

case class JobSubmitted(
    jobId: Int,
    finalRDD: RDD[_],
    func: (TaskContext, Iterator[_]) => _,
    partitions: Array[Int],
    callSite: CallSite,
    listener: JobListener,
    properties: Properties) extends DAGSchedulerEvent

case class StageCompleted(stage: Stage) extends DAGSchedulerEvent
case class TaskCompleted(task: Task[_], reason: TaskEndReason) extends DAGSchedulerEvent
case class TaskFailed(taskId: Long, reason: TaskFailedReason) extends DAGSchedulerEvent
case class ExecutorLost(execId: String, reason: ExecutorLossReason) extends DAGSchedulerEvent

// 事件处理循环实现
class DAGSchedulerEventProcessLoop(dagScheduler: DAGScheduler)
  extends EventLoop[DAGSchedulerEvent]("dag-scheduler-event-loop") {
  
  override def onReceive(event: DAGSchedulerEvent): Unit = {
    val timerContext = dagScheduler.metricsSource.messageProcessingTimer.time()
    try {
      doOnReceive(event)
    } finally {
      timerContext.stop()
    }
  }
  
  private def doOnReceive(event: DAGSchedulerEvent): Unit = event match {
    case JobSubmitted(jobId, rdd, func, partitions, callSite, listener, properties) =>
      dagScheduler.handleJobSubmitted(jobId, rdd, func, partitions, callSite, listener, properties)
      
    case StageCompleted(stage) =>
      dagScheduler.handleStageCompleted(stage)
      
    case TaskCompleted(task, reason) =>
      dagScheduler.handleTaskCompleted(task, reason)
      
    case TaskFailed(taskId, reason) =>
      dagScheduler.handleTaskFailed(taskId, reason)
      
    case ExecutorLost(execId, reason) =>
      dagScheduler.handleExecutorLost(execId, reason)
  }
}

// DAGScheduler.scala - 核心调度逻辑
class DAGScheduler(
    private[scheduler] val sc: SparkContext,
    private[scheduler] val taskScheduler: TaskScheduler,
    listenerBus: LiveListenerBus,
    mapOutputTracker: MapOutputTrackerMaster,
    blockManagerMaster: BlockManagerMaster,
    env: SparkEnv,
    clock: Clock = new SystemClock())
  extends Logging {

  // 事件处理循环
  private val eventProcessLoop = new DAGSchedulerEventProcessLoop(this)
  
  // Stage和Job管理
  private val stageIdToStage = new HashMap[Int, Stage]
  private val shuffleIdToMapStage = new HashMap[Int, ShuffleMapStage]
  private val jobIdToActiveJob = new HashMap[Int, ActiveJob]
  private val activeJobs = new HashSet[ActiveJob]
  
  // 等待和运行中的Stage
  private val waitingStages = new HashSet[Stage]
  private val runningStages = new HashSet[Stage]
  private val failedStages = new HashSet[Stage]
  
  // 1. 提交作业的核心方法
  def runJob[T, U](
      rdd: RDD[T],
      func: (TaskContext, Iterator[T]) => U,
      partitions: Seq[Int],
      callSite: CallSite,
      resultHandler: (Int, U) => Unit,
      properties: Properties): Unit = {
    
    val start = System.nanoTime
    val waiter = submitJob(rdd, func, partitions, callSite, resultHandler, properties)
    ThreadUtils.awaitReady(waiter, Duration.Inf)
    
    waiter.value.get match {
      case scala.util.Success(_) =>
        logInfo("Job %d finished: %s, took %f s".format
          (waiter.jobId, callSite.shortForm, (System.nanoTime - start) / 1e9))
      case scala.util.Failure(exception) =>
        logInfo("Job %d failed: %s, took %f s".format
          (waiter.jobId, callSite.shortForm, (System.nanoTime - start) / 1e9))
        throw exception
    }
  }
  
  // 2. 提交作业
  def submitJob[T, U](
      rdd: RDD[T],
      func: (TaskContext, Iterator[T]) => U,
      partitions: Seq[Int],
      callSite: CallSite,
      resultHandler: (Int, U) => Unit,
      properties: Properties): JobWaiter[U] = {
    
    // 检查分区有效性
    val maxPartitions = rdd.partitions.length
    partitions.find(p => p >= maxPartitions || p < 0).foreach { p =>
      throw new IllegalArgumentException(
        "Attempting to access a non-existent partition: " + p + ". " +
          "Total number of partitions: " + maxPartitions)
    }
    
    val jobId = nextJobId.getAndIncrement()
    if (partitions.size == 0) {
      // 空分区直接返回
      return new JobWaiter[U](this, jobId, 0, resultHandler)
    }
    
    assert(partitions.size > 0)
    val func2 = func.asInstanceOf[(TaskContext, Iterator[_]) => _]
    val waiter = new JobWaiter(this, jobId, partitions.size, resultHandler)
    
    // 提交JobSubmitted事件
    eventProcessLoop.post(JobSubmitted(
      jobId, rdd, func2, partitions.toArray, callSite, waiter, properties))
    waiter
  }
  
  // 3. 处理作业提交
  private[scheduler] def handleJobSubmitted(jobId: Int,
      finalRDD: RDD[_],
      func: (TaskContext, Iterator[_]) => _,
      partitions: Array[Int],
      callSite: CallSite,
      listener: JobListener,
      properties: Properties): Unit = {
    
    var finalStage: ResultStage = null
    try {
      // 创建ResultStage
      finalStage = createResultStage(finalRDD, func, partitions, jobId, callSite)
    } catch {
      case e: Exception =>
        logWarning("Creating new stage failed due to exception - job: " + jobId, e)
        listener.jobFailed(e)
        return
    }
    
    // 创建ActiveJob
    val job = new ActiveJob(jobId, finalStage, callSite, listener, properties)
    clearCacheLocs()
    
    logInfo("Got job %s (%s) with %d output partitions".format(
      job.jobId, callSite.shortForm, partitions.length))
    logInfo("Final stage: " + finalStage + " (" + finalStage.name + ")")
    logInfo("Parents of final stage: " + finalStage.parents)
    logInfo("Missing parents: " + getMissingParentStages(finalStage))
    
    val jobSubmissionTime = clock.getTimeMillis()
    jobIdToActiveJob(jobId) = job
    activeJobs += job
    finalStage.setActiveJob(job)
    
    val stageIds = jobIdToStageIds(jobId)
    val stageInfos = stageIds.flatMap(id => stageIdToStage.get(id).map(_.latestInfo))
    listenerBus.post(
      SparkListenerJobStart(job.jobId, jobSubmissionTime, stageInfos, properties))
    
    // 提交Stage
    submitStage(finalStage)
  }
  
  // 4. Stage划分核心算法
  private def getOrCreateParentStages(rdd: RDD[_], firstJobId: Int): List[Stage] = {
    val parents = new ArrayBuffer[Stage]()
    val visited = new HashSet[RDD[_]]
    
    def visit(r: RDD[_]): Unit = {
      if (!visited(r)) {
        visited += r
        for (dep <- r.dependencies) {
          dep match {
            case shufDep: ShuffleDependency[_, _, _] =>
              // 宽依赖，创建新的ShuffleMapStage
              parents += getOrCreateShuffleMapStage(shufDep, firstJobId)
            case _ =>
              // 窄依赖，递归访问父RDD
              visit(dep.rdd)
          }
        }
      }
    }
    
    visit(rdd)
    parents.toList
}
  
  // 5. 创建或获取ShuffleMapStage
  private def getOrCreateShuffleMapStage(
      shuffleDep: ShuffleDependency[_, _, _],
      firstJobId: Int): ShuffleMapStage = {
    
    shuffleIdToMapStage.get(shuffleDep.shuffleId) match {
      case Some(stage) =>
        stage
        
      case None =>
        // 递归创建父Stage
        getMissingAncestorShuffleDependencies(shuffleDep.rdd).foreach { dep =>
          if (!shuffleIdToMapStage.contains(dep.shuffleId)) {
            createShuffleMapStage(dep, firstJobId)
          }
        }
        createShuffleMapStage(shuffleDep, firstJobId)
    }
  }
  
  // 6. 提交Stage
  private def submitStage(stage: Stage): Unit = {
    val jobId = activeJobForStage(stage)
    if (jobId.isDefined) {
      logDebug("submitStage(" + stage + ")")
      if (!waitingStages(stage) && !runningStages(stage) && !failedStages(stage)) {
        val missing = getMissingParentStages(stage).sortBy(_.id)
        logDebug("missing: " + missing)
        if (missing.isEmpty) {
          logInfo("Submitting " + stage + " (" + stage.rdd + "), which has no missing parents")
          submitMissingTasks(stage, jobId.get)
        } else {
          for (parent <- missing) {
            submitStage(parent)
          }
          waitingStages += stage
        }
      }
    } else {
      abortStage(stage, "No active job for stage " + stage.id, None)
    }
  }
  
  // 7. 提交缺失的任务
  private def submitMissingTasks(stage: Stage, jobId: Int): Unit = {
    logDebug("submitMissingTasks(" + stage + ")")
    
    // 获取缺失的分区
    val partitionsToCompute: Seq[Int] = stage.findMissingPartitions()
    
    // 添加到运行中的Stage
    runningStages += stage
    
    // 创建任务
    val tasks: Seq[Task[_]] = try {
      val serializedTaskMetrics = closureSerializer.serialize(stage.latestInfo.taskMetrics).array()
      stage match {
        case stage: ShuffleMapStage =>
          partitionsToCompute.map { id =>
            val locs = taskIdToLocations(id)
            val part = stage.rdd.partitions(id)
            new ShuffleMapTask(stage.id, stage.latestInfo.attemptId,
              taskBinary, part, locs, stage.latestInfo.taskMetrics, properties, serializedTaskMetrics)
          }
          
        case stage: ResultStage =>
          partitionsToCompute.map { id =>
            val p: Int = stage.partitions(id)
            val part = stage.rdd.partitions(p)
            val locs = taskIdToLocations(id)
            new ResultTask(stage.id, stage.latestInfo.attemptId,
              taskBinary, part, locs, id, properties, serializedTaskMetrics)
          }
      }
    } catch {
      case NonFatal(e) =>
        abortStage(stage, s"Task creation failed: $e\n${Utils.exceptionString(e)}", Some(e))
        runningStages -= stage
        return
    }
    
    if (tasks.size > 0) {
      logInfo(s"Submitting ${tasks.size} missing tasks from $stage (${stage.rdd}) (first 15 " +
        s"tasks are for partitions ${partitionsToCompute.take(15).mkString(", ")})")
      
      // 提交TaskSet给TaskScheduler
      taskScheduler.submitTasks(new TaskSet(
        tasks.toArray, stage.id, stage.latestInfo.attemptId, jobId, properties))
    } else {
      // 没有任务需要运行，标记Stage完成
      markStageAsFinished(stage, None)
      
      val debugString = stage match {
        case stage: ShuffleMapStage =>
          s"Stage ${stage} is actually done; " +
            s"(available: ${stage.isAvailable}," +
            s"available outputs: ${stage.numAvailableOutputs}," +
            s"partitions: ${stage.numPartitions})"
        case stage : ResultStage =>
          s"Stage ${stage} is actually done; (partitions: ${stage.numPartitions})"
      }
      logDebug(debugString)
      
      submitWaitingChildStages(stage)
    }
  }
  
  // 8. 处理任务完成
  def handleTaskCompletion(event: CompletionEvent): Unit = {
    val task = event.task
    val taskId = event.taskInfo.taskId
    val stageId = task.stageId
    val taskType = Utils.getFormattedClassName(task)
    
    // 更新累加器
    event.accumUpdates.foreach { case (id, partialValue) =>
      val acc = AccumulatorContext.get(id)
      if (acc != null) {
        acc.asInstanceOf[AccumulatorV2[Any, Any]].merge(partialValue.asInstanceOf[Any])
      }
    }
    
    // 处理不同的任务结果
    event.reason match {
      case Success =>
        task match {
          case rt: ResultTask[_, _] =>
            // ResultTask完成
            val resultStage = stage.asInstanceOf[ResultStage]
            resultStage.activeJob match {
              case Some(job) =>
                if (!job.finished(rt.outputId)) {
                  job.finished(rt.outputId) = true
                  job.numFinished += 1
                  
                  // 调用结果处理器
                  job.listener.taskSucceeded(rt.outputId, event.result)
                  
                  // 检查Job是否完成
                  if (job.numFinished == job.numPartitions) {
                    markStageAsFinished(resultStage)
                    cleanupStateForJobAndIndependentStages(job)
                    listenerBus.post(SparkListenerJobEnd(job.jobId, clock.getTimeMillis(), JobSucceeded))
                  }
                }
              case None =>
                logInfo("Ignoring result from " + rt + " because its job has finished")
            }
            
          case smt: ShuffleMapTask =>
            // ShuffleMapTask完成
            val shuffleStage = stage.asInstanceOf[ShuffleMapStage]
            shuffleStage.addOutputLoc(smt.partitionId, event.result.asInstanceOf[MapStatus])
            
            if (runningStages.contains(shuffleStage) && shuffleStage.pendingPartitions.isEmpty) {
              markStageAsFinished(shuffleStage)
              logInfo("looking for newly runnable stages")
              logInfo("running: " + runningStages)
              logInfo("waiting: " + waitingStages)
              logInfo("failed: " + failedStages)
              
              // 提交等待中的子Stage
              submitWaitingChildStages(shuffleStage)
            }
        }
        
      case _: TaskFailedException =>
        // 任务失败处理
        handleTaskFailure(task, event.reason.asInstanceOf[TaskFailedException])
    }
  }
}

// Stage划分核心逻辑
private def getOrCreateParentStages(rdd: RDD[_], firstJobId: Int): List[Stage] = {
  val parents = new ArrayBuffer[Stage]()
  val visited = new HashSet[RDD[_]]
  
  def visit(r: RDD[_]): Unit = {
    if (!visited(r)) {
      visited += r
      for (dep <- r.dependencies) {
        dep match {
          case shufDep: ShuffleDependency[_, _, _] =>
            // 宽依赖，创建新的ShuffleMapStage
            parents += getOrCreateShuffleMapStage(shufDep, firstJobId)
          case _ =>
            // 窄依赖，递归访问父RDD
            visit(dep.rdd)
        }
      }
    }
  }
  
  visit(rdd)
  parents.toList
}

// 查找缺失的父依赖
private def getMissingAncestorShuffleDependencies(
    rdd: RDD[_]): ArrayStack[ShuffleDependency[_, _, _]] = {
  val ancestors = new ArrayStack[ShuffleDependency[_, _, _]]
  val visited = new HashSet[RDD[_]]
  val waitingForVisit = new ArrayStack[RDD[_]]
  
  waitingForVisit.push(rdd)
  while (waitingForVisit.nonEmpty) {
    val toVisit = waitingForVisit.pop()
    if (!visited(toVisit)) {
      visited += toVisit
      toVisit.dependencies.foreach {
        case shuffleDep: ShuffleDependency[_, _, _] =>
          if (!shuffleIdToMapStage.contains(shuffleDep.shuffleId)) {
            ancestors.push(shuffleDep)
            waitingForVisit.push(shuffleDep.rdd)
          }
        case narrowDep: NarrowDependency[_] =>
          waitingForVisit.push(narrowDep.rdd)
      }
    }
  }
  ancestors
}

// 创建ResultStage
private def createResultStage(
    rdd: RDD[_],
    func: (TaskContext, Iterator[_]) => _,
    partitions: Array[Int],
    jobId: Int,
    callSite: CallSite): ResultStage = {
  
  val parents = getOrCreateParentStages(rdd, jobId)
  val id = nextStageId.getAndIncrement()
  val stage = new ResultStage(id, rdd, func, partitions, parents, jobId, callSite)
  stageIdToStage(id) = stage
  updateJobIdStageIdMaps(jobId, stage)
  stage
}

// 创建ShuffleMapStage
private def createShuffleMapStage(shuffleDep: ShuffleDependency[_, _, _], jobId: Int): ShuffleMapStage = {
  val rdd = shuffleDep.rdd
  val numTasks = rdd.partitions.length
  val parents = getOrCreateParentStages(rdd, jobId)
  val id = nextStageId.getAndIncrement()
  val stage = new ShuffleMapStage(id, rdd, numTasks, parents, jobId, rdd.creationSite, shuffleDep)
  
  stageIdToStage(id) = stage
  shuffleIdToMapStage(shuffleDep.shuffleId) = stage
  updateJobIdStageIdMaps(jobId, stage)
  
  if (!mapOutputTracker.containsShuffle(shuffleDep.shuffleId)) {
    logInfo("Registering RDD " + rdd.id + " (" + rdd.getCreationSite + ") as input to " +
      "shuffle " + shuffleDep.shuffleId)
    mapOutputTracker.registerShuffle(shuffleDep.shuffleId, rdd.partitions.length)
  }
  stage
}

// DAGScheduler事件类型
sealed trait DAGSchedulerEvent

case class JobSubmitted(
    jobId: Int,
    finalRDD: RDD[_],
    func: (TaskContext, Iterator[_]) => _,
    partitions: Array[Int],
    callSite: CallSite,
    listener: JobListener,
    properties: Properties) extends DAGSchedulerEvent

case class StageCompleted(stage: Stage) extends DAGSchedulerEvent
case class TaskCompleted(task: Task[_], reason: TaskEndReason) extends DAGSchedulerEvent

// 事件处理循环
class DAGSchedulerEventProcessLoop(dagScheduler: DAGScheduler)
  extends EventLoop[DAGSchedulerEvent]("dag-scheduler-event-loop") {
  
  override def onReceive(event: DAGSchedulerEvent): Unit = {
    event match {
      case JobSubmitted(jobId, rdd, func, partitions, callSite, listener, properties) =>
        dagScheduler.handleJobSubmitted(jobId, rdd, func, partitions, callSite, listener, properties)
      case StageCompleted(stage) =>
        dagScheduler.handleStageCompletion(stage)
      case TaskCompleted(task, reason) =>
        dagScheduler.handleTaskCompletion(task, reason)
    }
  }
}

// Stage划分核心逻辑
private def getOrCreateParentStages(rdd: RDD[_], firstJobId: Int): List[Stage] = {
  val parents = new ArrayBuffer[Stage]()
  val visited = new HashSet[RDD[_]]
  
  def visit(r: RDD[_]): Unit = {
    if (!visited(r)) {
      visited += r
      for (dep <- r.dependencies) {
        dep match {
          case shufDep: ShuffleDependency[_, _, _] =>
            // 宽依赖，创建新的ShuffleMapStage
            parents += getOrCreateShuffleMapStage(shufDep, firstJobId)
          case _ =>
            // 窄依赖，递归访问父RDD
            visit(dep.rdd)
        }
      }
    }
  }
  
  visit(rdd)
  parents.toList
}

// TaskSetManager中的本地性调度实现
class TaskSetManager(
    sched: TaskSchedulerImpl,
    val taskSet: TaskSet,
    val maxTaskFailures: Int,
    val blacklistTracker: Option[BlacklistTracker] = None,
    clock: Clock = new SystemClock()) extends Schedulable with Logging {

  // 本地性级别定义
  private val myLocalityLevels = computeValidLocalityLevels()
  private val localityWaits = myLocalityLevels.map(getLocalityWait)

  // 计算有效的本地性级别
  private def computeValidLocalityLevels(): Array[TaskLocality.TaskLocality] = {
    import TaskLocality._
    val levels = new ArrayBuffer[TaskLocality.TaskLocality]
    
    if (!pendingTasksForExecutor.isEmpty) {
      levels += PROCESS_LOCAL
    }
    if (!pendingTasksForHost.isEmpty) {
      levels += NODE_LOCAL
    }
    if (!pendingTasksWithNoPrefs.isEmpty) {
      levels += NO_PREF
    }
    if (!pendingTasksForRack.isEmpty) {
      levels += RACK_LOCAL
    }
    levels += ANY
    
    logDebug("Valid locality levels for " + taskSet + ": " + levels.mkString(", "))
    levels.toArray
  }

  // 获取本地性等待时间
  private def getLocalityWait(level: TaskLocality.TaskLocality): Long = {
    val defaultWait = sched.conf.get("spark.locality.wait", "3s")
    val localityWaitKey = level match {
      case TaskLocality.PROCESS_LOCAL => "spark.locality.wait.process"
      case TaskLocality.NODE_LOCAL => "spark.locality.wait.node"
      case TaskLocality.RACK_LOCAL => "spark.locality.wait.rack"
      case _ => null
    }
    
    if (localityWaitKey != null) {
      sched.conf.getTimeAsMs(localityWaitKey, defaultWait)
    } else {
      0L
    }
  }

  // 为特定Executor和本地性级别查找任务
  private def findTask(
      execId: String,
      host: String,
      maxLocality: TaskLocality.TaskLocality): Option[(Int, TaskLocality.TaskLocality)] = {
    
    for (locality <- TaskLocality.values) {
      if (isAllowedLocality(locality, maxLocality)) {
        val task = locality match {
          case TaskLocality.PROCESS_LOCAL =>
            findTaskFromList(pendingTasksForExecutor.getOrElse(execId, ArrayBuffer()))
            
          case TaskLocality.NODE_LOCAL =>
            findTaskFromList(pendingTasksForHost.getOrElse(host, ArrayBuffer()))
            
          case TaskLocality.NO_PREF =>
            findTaskFromList(pendingTasksWithNoPrefs)
            
          case TaskLocality.RACK_LOCAL =>
            val rack = sched.getRackForHost(host)
            if (rack != null) {
              findTaskFromList(pendingTasksForRack.getOrElse(rack, ArrayBuffer()))
            } else {
              None
            }
            
          case TaskLocality.ANY =>
            findTaskFromList(allPendingTasks)
        }
        
        if (task.isDefined) {
          return Some((task.get, locality))
        }
      }
    }
    None
  }
  
  // 检查是否允许特定本地性级别
  private def isAllowedLocality(
      taskLocality: TaskLocality.TaskLocality,
      maxLocality: TaskLocality.TaskLocality): Boolean = {
    TaskLocality.isAllowed(taskLocality, maxLocality)
  }
}

// 推测执行实现
class TaskSetManager {
  
  // 推测执行配置
  private val speculationEnabled = sched.conf.getBoolean("spark.speculation", false)
  private val speculationQuantile = sched.conf.getDouble("spark.speculation.quantile", 0.75)
  private val speculationMultiplier = sched.conf.getDouble("spark.speculation.multiplier", 1.5)
  
  // 检查是否需要推测执行
  def checkSpeculatableTasks(minFinishedForSpeculation: Int): Boolean = {
    if (!speculationEnabled || taskSet.stageId < 0 || taskSet.stageAttemptId < 0) {
      return false
    }
    
    var foundTasks = false
    val finishedTasks = runningTasksSet.size - pendingPartitions.size
    
    if (finishedTasks >= minFinishedForSpeculation && finishedTasks > 0) {
      val time = clock.getTimeMillis()
      val medianDuration = medianTaskDuration
      
      if (medianDuration > 0) {
        val threshold = speculationMultiplier * medianDuration
        
        // 检查每个正在运行的任务
        for ((taskId, info) <- runningTasks) {
          val runtime = time - info.launchTime
          if (runtime > threshold && !speculatableTasks.contains(taskId)) {
            logInfo(s"Marking task $taskId in stage ${taskSet.stageId} (on ${info.host}) " +
              s"as speculatable because it ran more than ${threshold.toDouble / 1000.0} seconds")
            speculatableTasks += taskId
            foundTasks = true
          }
        }
      }
    }
    foundTasks
  }
  
  // 获取可推测的任务
  def dequeueSpeculativeTask(
      execId: String,
      host: String,
      locality: TaskLocality.TaskLocality): Option[(Int, TaskLocality.TaskLocality)] = {
    
    speculatableTasks.retain(index => runningTasks.contains(index))
    
    if (speculatableTasks.nonEmpty) {
      // 检查是否有适合该Executor的推测任务
      for (index <- speculatableTasks) {
        val locations = tasks(index).preferredLocations
        val allowedLocality = getAllowedLocalityLevel(clock.getTimeMillis())
        
        if (TaskLocality.isAllowed(locality, allowedLocality)) {
          if (isTaskLocalToExecutor(index, execId, host, locality)) {
            speculatableTasks -= index
            return Some((index, locality))
          }
        }
      }
    }
    None
  }
  
  // 处理推测任务完成
  def handleSuccessfulTask(tid: Long, result: DirectTaskResult[_]): Unit = {
    val info = taskInfos(tid)
    val index = info.index
    
    // 如果这是推测任务的成功，取消其他副本
    if (speculatableTasks.contains(index)) {
      for ((taskId, taskInfo) <- runningTasks) {
        if (taskInfo.index == index && taskId != tid) {
          logInfo(s"Killing speculative task $taskId")
          sched.backend.killTask(taskId, taskInfo.executorId, interruptThread = true, 
            reason = "another copy of the task succeeded")
        }
      }
      speculatableTasks -= index
    }
    
    // 标记任务成功
    successful(index) = true
    removeRunningTask(tid)
    
    // 如果所有任务都完成，通知DAGScheduler
    if (tasksSuccessful == numTasks) {
      isZombie = true
    }
  }
}

// 故障处理实现
class TaskSetManager {
  
  // 处理任务失败
  def handleFailedTask(
      tid: Long,
      state: TaskState,
      reason: TaskFailedReason): Unit = {
    
    val info = taskInfos(tid)
    val index = info.index
    
    removeRunningTask(tid)
    
    reason match {
      case fetchFailed: FetchFailed =>
        // Shuffle数据获取失败，需要重新计算上游Stage
        logWarning(s"Lost task ${info.id} in stage ${taskSet.stageId} (TID $tid, ${info.host}, " +
          s"executor ${info.executorId}): ${reason}")
        
        if (!successful(index)) {
          successful(index) = true
          tasksSuccessful += 1
        }
        
        // 标记需要重新计算的Stage
        abort(s"Task $tid failed with FetchFailedException", Some(fetchFailed))
        
      case exceptionFailure: ExceptionFailure =>
        // 任务执行异常
        val key = (info.host, info.executorId)
        numFailures(index) += 1
        
        if (numFailures(index) >= maxTaskFailures) {
          // 任务失败次数超过阈值，整个TaskSet失败
          val errorMessage = s"Task $index failed $maxTaskFailures times, most recent failure: " +
            exceptionFailure.description
          abort(errorMessage, exceptionFailure.exception)
          return
        } else {
          // 重新加入待调度队列
          addPendingTask(index)
        }
        
      case taskKilled: TaskKilled =>
        // 任务被主动杀死
        logWarning(s"Task $tid killed: ${taskKilled.reason}")
        
      case executorLostFailure: ExecutorLostFailure =>
        // Executor丢失导致的任务失败
        logInfo(s"Task $tid failed because executor ${executorLostFailure.execId} was lost")
        addPendingTask(index)
        
      case _ =>
        // 其他类型失败
        logWarning(s"Lost task $tid in stage ${taskSet.stageId}: $reason")
        addPendingTask(index)
    }
    
    // 更新失败统计
    if (!isZombie) {
      taskSetBlacklistHelperOpt.foreach(_.updateBlacklistForFailedTask(
        info.host, info.executorId, index, reason))
    }
  }
  
  // 处理Executor丢失
  def executorLost(execId: String, host: String, reason: ExecutorLossReason): Unit = {
    // 重新调度该Executor上运行的所有任务
    for ((tid, info) <- runningTasks) {
      if (info.executorId == execId) {
        val reason = ExecutorLostFailure(execId, exitCausedByApp = true, 
          Some(s"Executor lost: $reason"))
        handleFailedTask(tid, TaskState.FAILED, reason)
      }
    }
    
    // 从本地性映射中移除该Executor
    for {
      (host, execIds) <- pendingTasksForHost
      if execIds.contains(execId)
    } {
      execIds -= execId
      if (execIds.isEmpty) {
        pendingTasksForHost.remove(host)
      }
    }
    
    // 重新计算有效的本地性级别
    recomputeLocality()
  }
}