首页 > 解决方案 > 如果包含字符串列表,则过滤 pyspark 数据帧

问题描述

假设我们有一个 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'、...的行)

标签: python-3.xpyspark

解决方案


IIUC,您想要返回column_a“类似”(在 SQL 意义上)中的任何值的行list_a

一种方法是使用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]

推荐阅读