使用Python字典替换Spark DataFrame数组列中的值

核心思路 #

通过Spark内置高阶函数transform结合map_from_arrays创建映射关系，高效完成数组内元素的批量替换，避免自定义UDF带来的性能开销。

Python 实现代码 #

from pyspark.sql import SparkSession
from pyspark.sql.functions import transform, map_from_arrays, lit

# 初始化Spark会话
spark = SparkSession.builder.appName("ArrayElementReplace").getOrCreate()

# 构建示例DataFrame
sample_data = [("text", [0, 1, 2, 3]), ("page", [0, 1, 4, 3])]
df = spark.createDataFrame(sample_data, schema=["data", "arraydata"])

# 定义替换映射的键值对
map_keys = [0, 1, 2, 3, 4]
map_values = ["negative", "positive", "name", "sequel", "odd"]

# 创建Spark可识别的映射字典
value_mapping = map_from_arrays(lit(map_keys), lit(map_values))

# 对数组列执行替换操作
result_df = df.withColumn(
    "arraydata_replaced",
    transform("arraydata", lambda elem: value_mapping[elem])
)

# 查看结果
result_df.show(truncate=False)

Scala 实现代码 #

import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.functions.{transform, map_from_arrays, lit}

object ArrayReplaceDemo {
  def main(args: Array[String]): Unit = {
    val spark = SparkSession.builder.appName("ArrayElementReplace").getOrCreate()
    import spark.implicits._

    // 构建示例DataFrame
    val sampleData = Seq(("text", Seq(0, 1, 2, 3)), ("page", Seq(0, 1, 4, 3)))
    val df = sampleData.toDF("data", "arraydata")

    // 定义替换映射的键值对
    val mapKeys = Seq(0, 1, 2, 3, 4)
    val mapValues = Seq("negative", "positive", "name", "sequel", "odd")

    // 创建Spark可识别的映射字典
    val valueMapping = map_from_arrays(lit(mapKeys), lit(mapValues))

    // 对数组列执行替换操作
    val resultDF = df.withColumn(
      "arraydata_replaced",
      transform($"arraydata", elem => valueMapping(elem))
    )

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

输出结果 #

扩展：处理未匹配元素 #

如果需要对映射中不存在的元素设置默认值（比如替换为特定字符串），可以结合when函数：

Python 示例：

from pyspark.sql.functions import when

result_df = df.withColumn(
    "arraydata_replaced",
    transform("arraydata", lambda elem: when(value_mapping.isNotNullAt(elem), value_mapping[elem]).otherwise(lit("unknown")))
)

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