如何避免Spark在when/otherwise子句中提前执行全分支代码

问题根源 #

你的代码中，lit(spark.sql(s"SHOW Partitions $TableName").count())是在Driver端立即执行的，而非作为Column表达式延迟到数据处理阶段执行。这意味着不管when的条件是否满足，SHOW Partitions语句都会提前触发——一旦目标表未分区，就会直接抛出AnalysisException，根本无法进入后续的数据处理流程。

可行解决方案 #

核心思路是把获取分区数的逻辑移到Driver端的条件判断/异常捕获中，只有当表是分区表时才执行SHOW Partitions，否则直接返回null值。

方案1：通过异常捕获处理未分区表 #

import org.apache.spark.sql.AnalysisException

// 尝试获取表的分区数，捕获未分区表的异常
val partitionCount: Option[Long] = try {
  Some(spark.sql(s"SHOW Partitions $TableName").count())
} catch {
  case e: AnalysisException if e.getMessage.contains("SHOW PARTITIONS is not allowed on a table that is not partitioned") =>
    None
}

// 根据分区数是否存在，构建目标列
val delta = df.withColumn(
  "PC",
  when(
    col("PartitionKey") != "[]",
    // 如果表是分区表则用实际分区数，否则用null
    partitionCount.map(lit).getOrElse(lit(null))
  ).otherwise(null)
)

方案2：通过Spark Catalog直接判断表分区状态 #

// 通过Catalog获取表元数据，判断是否分区
val isPartitioned = spark.catalog.getTable(TableName).isPartitioned
val partitionCount = if (isPartitioned) Some(spark.sql(s"SHOW Partitions $TableName").count()) else None

// 构建目标列
val delta = df.withColumn(
  "PC",
  when(col("PartitionKey") != "[]", partitionCount.map(lit).getOrElse(lit(null))).otherwise(null)
)

逻辑说明 #

1. Driver端提前处理：两种方案都在Driver端完成表状态判断/分区数计算，避免了在数据处理阶段触发无效的SHOW Partitions调用。
2. 适配业务逻辑：当表未分区时，无论PartitionKey的值是什么，PC列都会返回null；当表是分区表时，仅在PartitionKey != "[]"时返回实际分区数，符合需求。

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