编写Spark Scala UDF：手动控制Parquet读取分区并保留原分区

问题概述 #

现有一个按key重分区的DataFrame，每行对应一个key和一组Parquet文件路径列表。需要读取这些文件的col_i、col_j列，且保证同一key的所有结果行保留在原分区中，避免后续Join操作时产生Shuffle。

输入DataFrame示例：

+---+---------------------+
|key|data                 |
+---+---------------------+
|k0 |[file0, file1, file2]|
|k1 |[file0, file3, file4]|
+---+---------------------+

期望输出DataFrame（同一key的行在同一分区）：

+---+------------------+------------------+
|key|             col_i|             col_j|
+---+------------------+------------------+
| k0|val_i_0_from_file0|val_j_0_from_file0|
| k0|val_i_1_from_file0|val_j_1_from_file0|
| k0|val_i_0_from_file1|val_j_0_from_file1|
| k0|val_i_0_from_file2|val_j_0_from_file2|
| k1|val_i_0_from_file0|val_j_0_from_file0|
| k1|val_i_1_from_file0|val_j_1_from_file0|
| k1|val_i_0_from_file3|val_j_0_from_file3|
| k1|val_i_0_from_file4|val_j_0_from_file4|
+---+------------------+------------------+

解决方案 #

由于需要保留分区边界，不能使用普通行级UDF，需采用分区级处理（mapPartitions算子）。核心思路是：利用原DataFrame已按key重分区的特性，每个分区仅处理一个key对应的所有Parquet文件，读取后将key与文件数据关联，直接输出结果以维持分区结构。

代码实现 #

1. 导入依赖与定义数据结构

import org.apache.spark.sql.{SparkSession, DataFrame}
import org.apache.spark.sql.functions._

// 定义输入输出数据结构
case class InputRow(key: String, data: Array[String])
case class OutputRow(key: String, col_i: String, col_j: String)

2. 实现分区级Parquet读取逻辑

/**
 * 针对单个分区读取指定Parquet文件，并关联key返回结果
 * @param key 当前分区对应的唯一key
 * @param filePaths 需要读取的Parquet文件路径列表
 * @param spark 活跃的SparkSession
 * @return 包含key的结果行迭代器
 */
def readParquetForPartition(key: String, filePaths: Array[String], spark: SparkSession): Iterator[OutputRow] = {
  // 读取指定文件的目标列，减少数据加载量
  val parquetDF = spark.read
    .parquet(filePaths: _*)
    .select("col_i", "col_j")

  // 将key与每行数据绑定，转换为OutputRow迭代器
  parquetDF.as[(String, String)]
    .map { case (ci, cj) => OutputRow(key, ci, cj) }
    .toLocalIterator
}

3. 应用到原DataFrame

// 初始化SparkSession（驱动端执行）
val spark = SparkSession.builder()
  .appName("ControlledParquetRead")
  .getOrCreate()

// 假设原DataFrame已按key重分区（必须保证每个分区仅对应一个key）
val inputDF: DataFrame = spark.read... // 加载你的输入DataFrame
val partitionedInputDF = inputDF.repartition(col("key")).as[InputRow]

// 执行分区级读取
val resultDF = partitionedInputDF
  .mapPartitions { iter =>
    // 每个分区仅一行数据（因提前按key重分区）
    val inputRow = iter.next()
    val activeSpark = SparkSession.getActiveSession.getOrElse(
      throw new RuntimeException("No active SparkSession found in executor")
    )
    readParquetForPartition(inputRow.key, inputRow.data, activeSpark)
  }
  .toDF()

关键注意事项 #

1. 提前按key重分区：必须确保输入DataFrame执行repartition(col("key"))，保证每个分区仅包含一个key的数据，这是维持分区结构的前提。
2. 文件路径可见性：Executor节点必须能访问所有Parquet文件路径，推荐使用HDFS、S3等分布式存储，避免使用本地文件（集群环境下本地文件无法跨节点共享）。
3. 性能优化：
   • 读取时仅选择需要的列（select("col_i", "col_j")），减少IO和内存开销。
   • 可通过spark.sql.parquet.enableVectorizedReader开启向量读取优化。
4. 异常处理：可在readParquetForPartition中添加异常捕获逻辑，比如处理文件不存在、列缺失等情况，避免单个文件故障导致整个分区失败。

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