PySpark中如何遍历DataFrame数组列并与另一DataFrame关联

问题描述 #

我有两个Spark DataFrame：
dataframe_1（some_id为数组列）：

dataframe_2：

由于dataframe_1的some_id是数组列，无法直接关联dataframe_2，我需要从dataframe_2中筛选出some_id存在于dataframe_1任意数组元素中的行。之前尝试遍历每行再遍历数组元素，但代码逻辑有问题。

尝试的错误代码：

df = < 某个查询语句 >
df_extra = (df.select("some_id","cor_name",df.propertiesMap.getItem("_value")))
df_Lineage = spark.sql("""select HLineage from someTable""")

for row in df_Lineage ():
    lis_1 = list(row)
    for element in lis_1:
      df_extra = ( df_extra.filter("terr_id").contains(element[0])
                 ) ## 我想根据每一行的每个元素过滤第二个DataFrame的值
      display(df_extra)

期望结果：

解决方案 #

Spark是分布式计算框架，禁止遍历行处理（效率极低且违背Spark设计逻辑），推荐以下两种高效实现方式：

方法1：展开数组列后关联

先将dataframe_1的数组列拆分为单行，再与dataframe_2做关联：

from pyspark.sql.functions import explode

# 展开数组列，得到所有需要匹配的some_id
df1_exploded = dataframe_1.select(explode("some_id").alias("match_id"))

# 去重避免重复关联
df1_exploded = df1_exploded.distinct()

# 内关联筛选匹配行
result_df = dataframe_2.join(df1_exploded, dataframe_2.some_id == df1_exploded.match_id, "inner")

# 移除多余的辅助列
result_df = result_df.drop("match_id")

display(result_df)

方法2：集合匹配过滤

收集所有需要匹配的ID到集合，再用isin直接过滤dataframe_2：

from pyspark.sql.functions import explode

# 分布式收集所有唯一匹配ID到集合
all_match_ids = dataframe_1.select(explode("some_id")).distinct().rdd.flatMap(lambda x: x).collect()
all_match_ids = set(all_match_ids)

# 过滤dataframe_2
result_df = dataframe_2.filter(dataframe_2.some_id.isin(all_match_ids))

display(result_df)

注意事项

• 你尝试的代码存在语法错误：filter("terr_id").contains(element[0])写法不正确，且遍历行的方式在数据量大时会直接导致性能崩溃。
• 若dataframe_1数据量极大，方法2的collect()可能触发Driver内存溢出，此时优先使用方法1（全分布式执行）。

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