如果包含字符串列表，则过滤 pyspark 数据帧

Posted 2023-04-15

【中文标题】如果包含字符串列表，则过滤 pyspark 数据帧

【英文标题】：Filter pyspark dataframe if contains a list of strings

【发布时间】：2019-10-01 13:55:08

【问题描述】：

假设我们有一个 pyspark 数据框，其中一列 (column_a) 包含一些字符串值，并且还有一个字符串列表 (list_a)。

数据框：

column_a      | count
some_string   |  10
another_one   |  20
third_string  |  30

list_a：

['string', 'third', ...]

我想过滤此数据框，并且仅在 column_a 的值包含 list_a 的项目之一时保留行。

这是用于根据单个字符串过滤column_a 的代码：

df['column_a'].like('%string_value%')

但是我们怎样才能得到一个字符串列表的相同结果呢？ （保留column_a的值为'string', 'third', ...的行）

【问题讨论】：

您是在寻找完全匹配还是子字符串匹配？对于前者，使用isin 例如：df.where(df['column_a'].isin(list_a))

@pault 是的，我知道isin()。我想检查值是否包含列表的字符串。

【参考方案1】：

IIUC，您想返回column_a 与list_a 中的任何值“类似”（在SQL 意义上）的行。

一种方法是使用functools.reduce：

from functools import reduce

list_a = ['string', 'third']

df1 = df.where(
    reduce(lambda a, b: a|b, (df['column_a'].like('%'+pat+"%") for pat in list_a))
)
df1.show()
#+------------+-----+
#|    column_a|count|
#+------------+-----+
#| some_string|   10|
#|third_string|   30|
#+------------+-----+

基本上，您循环遍历list_a 中的所有可能字符串以在like 中进行比较并“或”结果。这是执行计划：

df1.explain()
#== Physical Plan ==
#*(1) Filter (Contains(column_a#0, string) || Contains(column_a#0, third))
#+- Scan ExistingRDD[column_a#0,count#1]

---

另一种选择是使用pyspark.sql.Column.rlike 而不是like。

df2 = df.where(
    df['column_a'].rlike("|".join(["(" + pat + ")" for pat in list_a]))
)

df2.show()
#+------------+-----+
#|    column_a|count|
#+------------+-----+
#| some_string|   10|
#|third_string|   30|
#+------------+-----+

有对应的执行计划：

df2.explain()
#== Physical Plan ==
#*(1) Filter (isnotnull(column_a#0) && column_a#0 RLIKE (string)|(third))
#+- Scan ExistingRDD[column_a#0,count#1]