将 pyspark DataFrame 的每一行相交，这是一个字符串列表与一个主字符串列表？

Posted 2023-04-13

【中文标题】将 pyspark DataFrame 的每一行相交，这是一个字符串列表与一个主字符串列表？

【英文标题】：Intersect each row of a pyspark DataFrame which is a list of strings with a master list of strings?

【发布时间】：2018-12-05 11:14:48

【问题描述】：

假设我有一个这样的 DataFrame。

[Row(case_number='5307793179', word_list=['n', 'b', 'c']),
 Row(case_number='5307793171', word_list=['w', 'e', 'c']),
 Row(case_number='5307793172', word_list=['1', 'f', 'c']),
 Row(case_number='5307793173', word_list=['a', 'k', 'c']),
 Row(case_number='5307793174', word_list=['z', 'l', 'c']),
 Row(case_number='5307793175', word_list=['b', 'r', 'c'])]

还有一个像这样的主词表：

master_word_list = ['b', 'c']

是否有一种简洁的方法可以根据 master_word_list 过滤 word_list，因此生成的 pyspark 数据框看起来像这样。 （我的意思是不使用 UDF，如果 UDF 是最好/唯一的方法，我也会接受它作为一种解决方案）

[Row(case_number='5307793179', word_list=['b', 'c']),
 Row(case_number='5307793171', word_list=['c']),
 Row(case_number='5307793172', word_list=['c']),
 Row(case_number='5307793173', word_list=['c']),
 Row(case_number='5307793174', word_list=['c']),
 Row(case_number='5307793175', word_list=['b', 'c'])]

【参考方案1】：

array_intersect 自 Spark 2.4 起可用：

pyspark.sql.functions.array_intersect(col1, col2)

collection函数：返回col1和col2交集的元素数组，不重复。

参数：

col1 - 包含数组的列的名称 col2 - 包含数组的列的名称

from pyspark.sql.functions import array, array_intersect, lit

master_word_list_col = array(*[lit(x) for x in master_word_list])

df = spark.createDataFrame(
    [("5307793179", ["n", "b", "c"])], 
    ("case_number", "word_list")
)

df.withColumn("word_list", array_intersect("word_list", master_word_list_col)).show()

+-----------+---------+
|case_number|word_list|
+-----------+---------+
| 5307793179|   [b, c]|
+-----------+---------+