Spark Scala：将DataFrame数组列按映射表转换为对应值

核心思路 #

优先使用Spark内置函数实现，避免自定义UDF的序列化开销，让Spark Catalyst优化器充分发挥作用，这是大数据场景下的高效方案。

实现步骤 #

1. 准备示例数据

先构建你提供的映射表和测试DataFrame：

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

val spark = SparkSession.builder().appName("ArrayMapping").master("local[*]").getOrCreate()
import spark.implicits._

// 你的Scala映射表
val map1 = Map("S" -> 1, "T" -> 2, "U" -> 3)
// 测试DataFrame
val df = Seq(Array("S", "U"), Array("U", "U")).toDF("mappedcol")

2. 推荐方案：使用Spark内置函数

将Scala Map转换为Spark的MapType字面量，结合transform和map_get函数对数组列的每个元素做映射：

// 把Scala Map转为Spark可识别的MapType表达式
val sparkMap = map(map1.map { case (k, v) => lit(k) -> lit(v) }.toSeq: _*)

// 对数组列逐元素映射
val resultDF = df.withColumn("mapped_result", transform(col("mappedcol"), x => map_get(sparkMap, x)))

// 查看结果
resultDF.show(false)

执行后输出：

+---------+-------------+
|mappedcol|mapped_result|
+---------+-------------+
|[S, U]   |[1, 3]       |
|[U, U]   |[3, 3]       |
+---------+-------------+

3. 优化：广播大映射表

如果你的映射表数据量较大，建议将其转为广播变量，避免每个Task重复加载映射数据：

val broadcastMap = broadcast(sparkMap)
val optimizedResultDF = df.withColumn("mapped_result", transform(col("mappedcol"), x => map_get(broadcastMap, x)))

4. 备选方案：自定义UDF（适合复杂逻辑）

如果映射逻辑更复杂，内置函数无法满足，可以用UDF实现，但性能不如内置函数：

// 定义UDF，可指定未匹配键的默认值（这里设为0）
val mapArrayUdf = udf((arr: Array[String]) => arr.map(map1.getOrElse(_, 0)))

val resultDFWithUDF = df.withColumn("mapped_result", mapArrayUdf(col("mappedcol")))
resultDFWithUDF.show(false)

性能说明 #

• 内置函数方案由Spark Catalyst优化器直接处理，无需Java/Scala对象的序列化/反序列化，在大数据量下性能远优于UDF。
• 广播映射表能进一步减少Executor端的内存占用和数据传输开销，适合大映射表场景。

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