如果包含字符串列表，请过滤pyspark dataframe

Question

假设我们有一个 pyspark 数据框，其中一个列（column_a）包含一些字符串值，同时还有一个字符串列表（list_a）。

数据框：

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

列表_a：

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

我希望筛选这个数据框，只保留包含list_a中任意一项的column_a的值所在的行。

以下是筛选单个字符串的代码：

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

那么我们如何针对字符串列表获得相同的结果呢？（保留列column_a值为'string'、'third'等的行）

- Saeed Esmaili

---

你是在寻找精确匹配还是子字符串匹配？如果是前者，可以使用 isin，例如：df.where(df['column_a'].isin(list_a))。 - pault

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

Answer 1

如果我理解正确的话，您想返回那些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的比较，并将结果“OR”起来。以下是执行计划：

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]

- pault