Spark SQL物理计划中RDD转换逻辑的外部暴露方法问询（Spark 3.2.1）

针对你提出的从Spark SQL物理计划中暴露内部RDD转换逻辑的需求，结合TPCH Query 1的示例，以下是几种可行的实现方案：

方案一：反射调用物理算子的doExecute()方法 #

Spark SQL的每个物理算子（如FileScanExec、HashAggregateExec）都通过doExecute()方法生成对应的RDD，该方法为protected修饰，可通过反射直接调用，获取每个算子对应的RDD转换结果：

步骤示例 #

1. 获取DataFrame的查询执行实例：

val df = spark.sql(query)
val qe = df.queryExecution
val physicalPlan = qe.executedPlan

2. 遍历物理计划节点，从叶子节点（如FileScanExec）开始，反射调用doExecute()获取RDD：

import org.apache.spark.rdd.RDD
import org.apache.spark.sql.catalyst.InternalRow

// 定位FileScanExec节点
val fileScanNode = physicalPlan.find(_.isInstanceOf[org.apache.spark.sql.execution.datasources.FileScanExec])
  .get.asInstanceOf[org.apache.spark.sql.execution.datasources.FileScanExec]

// 反射调用doExecute()获取初始RDD
val doExecuteMethod = fileScanNode.getClass.getDeclaredMethod("doExecute")
doExecuteMethod.setAccessible(true)
val initialRDD: RDD[InternalRow] = doExecuteMethod.invoke(fileScanNode).asInstanceOf[RDD[InternalRow]]

3. 依次向上遍历物理计划的父节点，重复调用doExecute()，即可得到每个转换步骤对应的RDD（如ColumnarToRowExec、ProjectExec等节点的转换结果）。

方案二：手动复现物理计划的转换逻辑 #

根据物理计划的算子类型，手动调用Spark RDD/Dataset API复现相同的转换逻辑，可控性更高：

对应TPCH Query 1的转换步骤 #

1. 初始RDD（FileScanRDD）：直接读取Parquet文件生成RDD

val initialRDD = spark.read.parquet("hdfs://node13-opa:8020/user/spark_benchmark/tpch_100/dataset/lineitem")
  .where("l_shipdate <= '1998-09-02'")
  .select("l_quantity", "l_extendedprice", "l_discount", "l_tax", "l_returnflag", "l_linestatus")
  .rdd

2. ColumnarToRow转换：将列存格式转换为行格式（对应MapPartitionsRDD[4]）

import org.apache.spark.sql.execution.columnar.ColumnarBatchUtil
val rowRDD = initialRDD.mapPartitions(ColumnarBatchUtil.columnarToRow)

3. Project算子转换：提取指定字段（对应MapPartitionsRDD[5]）

val projectedRDD = rowRDD.map(row => 
  Row(
    row.getDecimal(0), row.getDecimal(1), row.getDecimal(2),
    row.getDecimal(3), row.getString(4), row.getString(5)
  )
)

4. 局部HashAggregate（partial）：实现局部聚合逻辑（对应MapPartitionsRDD[6]）

case class PartialAgg(
  returnFlag: String, lineStatus: String,
  sumQty: BigDecimal, sumExtPrice: BigDecimal,
  sumDiscPrice: BigDecimal, sumCharge: BigDecimal,
  countQty: Long, countExtPrice: Long, countDisc: Long,
  totalCount: Long
)

val partialAggRDD = projectedRDD.mapPartitions(iter => {
  val aggMap = scala.collection.mutable.Map[(String, String), PartialAgg]()
  iter.foreach(row => {
    val key = (row.getString(4), row.getString(5))
    val qty = row.getDecimal(0)
    val extPrice = row.getDecimal(1)
    val discount = row.getDecimal(2)
    val tax = row.getDecimal(3)
    val discPrice = extPrice * (BigDecimal(1.00) - discount)
    val charge = discPrice * (BigDecimal(1.00) + tax)
    
    aggMap.update(key, aggMap.get(key) match {
      case Some(agg) => PartialAgg(
        agg.returnFlag, agg.lineStatus,
        agg.sumQty + qty, agg.sumExtPrice + extPrice,
        agg.sumDiscPrice + discPrice, agg.sumCharge + charge,
        agg.countQty + 1, agg.countExtPrice + 1, agg.countDisc + 1,
        agg.totalCount + 1
      )
      case None => PartialAgg(
        key._1, key._2,
        qty, extPrice, discPrice, charge,
        1L, 1L, 1L, 1L
      )
    })
  })
  aggMap.values.toIterator
})

5. Hash Exchange（Shuffle）：按returnFlag和lineStatus进行hash分区（对应ShuffledRowRDD[7]）

import org.apache.spark.HashPartitioner
val shuffledRDD = partialAggRDD
  .keyBy(agg => (agg.returnFlag, agg.lineStatus))
  .partitionBy(new HashPartitioner(200))
  .values

后续的全局HashAggregate、Range Exchange、Sort等步骤，均可按照物理计划的逻辑继续手动实现。

方案三：自定义物理算子扩展（深度定制） #

如果需要长期复用或深度定制转换逻辑，可以继承Spark的物理算子类，重写doExecute()并将转换逻辑暴露为公共方法。例如：

class CustomFileScanExec(delegate: FileScanExec) extends FileScanExec(delegate.relation, delegate.output, delegate.requiredSchema, delegate.partitionFilters, delegate.dataFilters, delegate.tableIdentifier) {
  // 暴露公共方法获取转换后的RDD
  def getRDD: RDD[InternalRow] = doExecute()
}

该方案需要对Spark源码有一定了解，适合需要深度集成的场景。

内容的提问来源于stack exchange，提问作者user20089473