如何高效将Spark数组列的值替换为Pandas DataFrame中的值？

问题背景 #

现有Spark DataFrame存储购物篮数据，其中basket列为商品ID数组；另有Pandas DataFrame存储商品ID与新ID的映射关系。原使用Python UDF实现替换，但在千万级数据量下运行极慢，需更高效方案。

原始数据定义：

import pandas as pd 
import pyspark.sql.types as T
from pyspark.sql import functions as F

# Spark购物篮数据
df_baskets = spark.createDataFrame(
    [(1, ["546", "689", "946"]), (2, ["546", "799"] )],
    ("case_id","basket")
)

# Pandas映射表
product_data = pd.DataFrame({
  "product_id": ["546", "689", "946", "799"],
  "new_product_id": ["S12", "S74", "S34", "S56"]
})

高效解决方案 #

方法1：Explode + Join + 聚合（适合超大数据量） #

利用Spark分布式操作，先拆分数组为单行，关联映射表后重新聚合，全程使用内置JVM函数，避免Python UDF的性能损耗。

# 将Pandas映射表转为Spark DataFrame
df_product = spark.createDataFrame(product_data)

# 1. 拆分数组为单行记录
df_exploded = df_baskets.withColumn("product_id", F.explode(F.col("basket")))

# 2. 关联映射表获取新ID，未匹配到则保留原ID
df_joined = df_exploded.join(df_product, on="product_id", how="left") \
                       .withColumn("new_product_id", F.coalesce(F.col("new_product_id"), F.col("product_id")))

# 3. 按case_id聚合，重新生成数组
df_result = df_joined.groupBy("case_id", "basket") \
                     .agg(F.collect_list("new_product_id").alias("basket_renamed"))

df_result.show()

方法2：CreateMap + ArrayMap（简洁高效，适合映射表较小的场景） #

将映射关系转为Spark内置字典，直接对数组每个元素做映射，代码更简洁，性能同样优异。

# 构建product_id到new_product_id的映射字典
product_map = F.create_map(*[F.lit(x) for x in product_data.melt()["value"].tolist()])

# 用array_map遍历数组元素，完成映射（未匹配则保留原ID）
df_result = df_baskets.withColumn(
    "basket_renamed",
    F.array_map(F.col("basket"), lambda x: F.coalesce(product_map[x], x))
)

df_result.show()

性能优势说明 #

两种方案均使用Spark内置JVM函数，避免了Python UDF的跨进程序列化/反序列化开销，以及循环查找Pandas表的单线程低效操作。在千万级数据量下，性能提升可达数十倍甚至上百倍。

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