如何避免.rdd转换开销获取Spark DataFrame的分区数

我完全理解你的痛点——用df.rdd.getNumPartitions获取DataFrame分区数的开销实在太大，尤其是在需要动态调整分区的场景下，每次都要承担和count差不多的耗时，完全得不偿失。下面我给你分享几种更高效的替代方案，帮你解决这个问题：

1. 借助Spark Catalog查询已注册表格的元数据 #

如果你的DataFrame是从已注册的表（不管是临时视图还是Hive持久化表）生成的，可以直接通过Spark Catalog查询元数据来获取分区信息，完全不需要触发任何数据扫描：

// 获取临时视图的物理分区数
spark.catalog.getTable("your_temp_view_name").partitioning.length

// 如果是Hive表，若想获取物理分区的数量（即实际存储的分区目录数），可以执行：
spark.sql("SHOW PARTITIONS your_hive_table").count()

注意：SHOW PARTITIONS返回的是表的逻辑分区数（按分区列划分的），只有当表的物理存储分区和逻辑分区完全一致时，这个数值才等于DataFrame的物理分区数。如果是普通DataFrame（非分区表生成），还是用下面的方法更可靠。

2. 从DataFrame的执行计划中提取分区数 #

Spark的DataFrame内置了queryExecution属性，里面包含了完整的逻辑和物理执行计划，我们可以直接从计划里解析出物理分区数，全程不需要触发实际的计算作业，开销几乎为0：

import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec

def getPhysicalPartitionCount(df: DataFrame): Int = {
  val physicalPlan = df.queryExecution.executedPlan
  // 优先查找计划中的Shuffle节点，获取洗牌后的分区数
  physicalPlan.collectFirst { 
    case shuffle: ShuffleExchangeExec => shuffle.outputPartitioning.numPartitions 
  }.getOrElse {
    // 如果没有Shuffle，直接获取最终输出的分区数
    physicalPlan.outputPartitioning.numPartitions
  }
}

这个方法的核心是解析Spark的执行计划逻辑，不需要触碰实际数据，所以速度极快，完全不会有你之前遇到的5分钟耗时问题。

3. 重构你的动态分区调整逻辑 #

结合上面的高效分区数获取方法，我们可以优化你原有的分区调整函数，彻底去掉df.rdd.getNumPartitions这个性能瓶颈：

import org.apache.spark.sql.DataFrame
import org.apache.log4j.Logger

val logger = Logger.getLogger("DynamicPartitionAdjuster")

def getPhysicalPartitionCount(df: DataFrame): Int = {
  val physicalPlan = df.queryExecution.executedPlan
  physicalPlan.collectFirst { 
    case shuffle: ShuffleExchangeExec => shuffle.outputPartitioning.numPartitions 
  }.getOrElse(physicalPlan.outputPartitioning.numPartitions)
}

def repartition(inDf: DataFrame, minPartitions: Option[Int], maxPartitions: Option[Int]): DataFrame = {
  val inputPartitions = getPhysicalPartitionCount(inDf)
  val outDf = minPartitions.flatMap { minp =>
    if (inputPartitions < minp) {
      logger.info(s"Repartition input from $inputPartitions to $minp partitions..")
      Option(inDf.repartition(minp))
    } else {
      None
    }
  }.getOrElse(
    maxPartitions.map { maxp =>
      if (inputPartitions > maxp) {
        logger.info(s"Coalesce input from $inputPartitions to $maxp partitions..")
        inDf.coalesce(maxp)
      } else {
        inDf
      }
    }.getOrElse(inDf)
  )
  outDf
}

重构后的函数完全避免了数据扫描的开销，能瞬间获取分区数，再根据你设定的最小/最大分区阈值动态调整，性能会有质的提升。

补充说明 #

你提到Spark GUI里df.rdd操作耗时和最长SQL相当，这是因为df.rdd会触发Spark执行一个轻量作业来生成RDD的分区结构，本质上和执行一次count的开销是一样的。而用执行计划解析的方法，只是在内存里解析计划逻辑，不需要启动任何任务，所以完全不会有这个耗时问题。

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