Spark 2.x Scala 2.11：修改partitionBy值生成自定义分区目录

嗨，这个需求其实很好实现，核心思路就是先对sourcefilename字段做值的映射转换，再用转换后的新字段来做分区保存就行，结合Spark 2.x和Scala 2.11的特性，给你两种常用的实现方式：

方法一：用withColumn+when/otherwise实现简单映射 #

如果你的映射规则是固定的键值对，直接用Spark内置的条件函数就能搞定，不用写自定义UDF：

import org.apache.spark.sql.functions._

// 第一步：定义你的文件名->分区值映射规则
val filenameToPartition = Map(
  "ABC" -> "123",
  "DEF" -> "345",
  "GHI" -> "678" // 可以继续添加更多映射关系
)

// 第二步：给DataFrame新增一个映射后的分区字段
val transformedDf = df.withColumn(
  "partition_key", // 新字段名，后续用它来分区
  // 匹配映射规则，找不到匹配的话可以保留原文件名或者设默认值
  when(col("sourcefilename").isin(filenameToPartition.keys.toSeq: _*),
       lit(filenameToPartition(col("sourcefilename").as[String]))
  ).otherwise(col("sourcefilename"))
)

// 第三步：用新字段分区保存ORC文件
transformedDf.write.mode("overwrite")
  .partitionBy("partition_key")
  .format("orc")
  .save("s3a://my/dir/path/output-data")

这样生成的分区目录会是partition_key=123、partition_key=345这种格式，完全符合你的映射需求。如果不需要保留原sourcefilename字段，可以在保存前加.drop("sourcefilename")删掉它。

方法二：自定义UDF处理复杂映射逻辑 #

如果你的映射规则比较复杂（比如需要从文件名提取子串、做正则匹配或者其他业务逻辑），就适合用自定义UDF来实现：

import org.apache.spark.sql.functions.udf
import org.apache.spark.sql.types.StringType

// 第一步：定义自定义UDF，实现文件名到分区值的转换逻辑
val mapFilenameToPartition = udf((filename: String) => filename match {
  case "ABC" => "123"
  case "DEF" => "345"
  case name if name.startsWith("XYZ") => "999" // 示例：处理前缀匹配的情况
  case _ => "unknown" // 未匹配时的默认值，也可以返回原文件名
})

// 第二步：应用UDF生成新的分区字段
val transformedDf = df.withColumn(
  "partition_key",
  mapFilenameToPartition(col("sourcefilename").cast(StringType)) // 确保字段是String类型
)

// 第三步：分区保存，逻辑和方法一一致
transformedDf.write.mode("overwrite")
  .partitionBy("partition_key")
  .format("orc")
  .save("s3a://my/dir/path/output-data")

额外注意事项 #

• Spark 2.x中，partitionBy生成的目录默认是字段名=值的格式，这是Spark识别分区列的标准格式，不建议修改成纯值目录（会导致后续读取时无法自动识别分区）；
• 如果你的sourcefilename字段不是String类型，一定要先通过.cast(StringType)转换后再做映射；
• 保存模式用overwrite时要注意，会覆盖目标路径下的所有内容，包括之前的分区目录。

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