PySpark：基于数组型订单ID列关联两个DataFrame的方法

我是PySpark及DataFrame的新手，多次搜索都没找到解决方案。需要将profiles与orders两个DataFrame关联：profiles包含存储订单ID数组的ordersId列，希望为其新增orders列，该列包含与ordersId中ID对应的订单信息。

我尝试了如下代码，但因在where上下文无法访问withColumn中的ordersId数组，导致代码无法正常运行：

profiles = profiles.withColumn('orders', lit(orders.where(array_contains(profiles.ordersId, orders.id)).collect()))

示例数据 #

from pyspark.sql.types import StructType,StructField, StringType, IntegerType, ArrayType


data = [("aaaa", ['cc', 'dd'])]

schema = StructType([ \
    StructField("id", StringType(), False), \
    StructField("ordersId", ArrayType(StringType(), False), False), \
  ])
 
df = spark.createDataFrame(data=data, schema=schema)
df.printSchema()
df.show(truncate=False)


data2 = [("cc", "titi"), ("dd", "toto")]

schema2 = StructType([ \
    StructField("id", StringType(), False), \
    StructField("alias", StringType(), False), \
  ])
 
df2 = spark.createDataFrame(data=data2, schema=schema2)
df2.printSchema()
df2.show(truncate=False)

示例输出 #

root
 |-- id: string (nullable = false)
 |-- ordersId: array (nullable = false)
 |    |-- element: string (containsNull = false)

+----+--------+
|id  |ordersId|
+----+--------+
|aaaa|[cc, dd]|
+----+--------+

root
 |-- id: string (nullable = false)
 |-- alias: string (nullable = false)

+---+-----+
|id |alias|
+---+-----+
|cc |titi |
|dd |toto |
+---+-----+

期望结果 #

+----+--------+-----------------------------------------------+
|id  |ordersId|orders                                         |
+----+--------+-----------------------------------------------+
|aaaa|[cc, dd]|[{id: cc, alias: titi}, {id: dd, alias: toto}] |
+----+--------+-----------------------------------------------+

解决方案 #

你的问题在于直接在withColumn里调用collect()会把orders数据拉到Driver端，而且array_contains无法跨DataFrame直接关联。正确的做法是通过展开数组、关联、再聚合的方式实现：

方法一：使用explode+join+groupBy（通用版本） #

from pyspark.sql import functions as F

# 1. 展开profiles中的ordersId数组，生成每行对应一个订单ID的临时表
exploded_profiles = df.withColumn("order_id", F.explode(F.col("ordersId")))

# 2. 和orders表关联，获取对应的订单信息
joined_df = exploded_profiles.join(df2, exploded_profiles.order_id == df2.id, "left")

# 3. 按profiles的id分组，将订单信息收集成数组
result = joined_df.groupBy("id", "ordersId")\
                  .agg(F.collect_list(F.struct(F.col("id").alias("id"), F.col("alias").alias("alias"))).alias("orders"))

result.show(truncate=False)

方法二：使用transform+映射字典（Spark 3.0+ 简洁版） #

如果你的Spark版本是3.0及以上，可以用更简洁的方式，先把orders转成键值对映射，再匹配数组元素生成嵌套结构：

from pyspark.sql import functions as F

# 将orders表转成(id -> struct(id, alias))的映射字典
order_map = df2.groupBy().agg(F.map_from_entries(F.collect_list(F.struct(F.col("id"), F.struct(F.col("id"), F.col("alias"))))).alias("order_map")).first()["order_map"]

# 遍历ordersId数组，从映射中取出对应订单信息
result = df.withColumn("orders", F.transform(F.col("ordersId"), lambda x: order_map[x]))

result.show(truncate=False)

说明 #

• 方法一适用于所有Spark版本，是分布式场景下的通用方案，不会把全量数据拉到Driver端，适合大数据量场景。
• 方法二更简洁，但依赖Spark 3.0+，且如果orders数据量过大，first()会把映射字典拉到Driver端，更适合小体量的orders表。

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