PySpark如何Pythonic地组合不定长列表的AND/OR过滤条件

前置测试环境准备 #

首先导入PySpark SQL函数，构造测试用DataFrame，代码如下：

import pyspark.sql.functions as func

row_data = [(1, 1, 1), (1, 1, 2), (1, 1, 3),
           (1, 2, 1), (1, 2, 2), (1, 2, 3),
           (2, 1, 1), (2, 1, 2), (2, 1, 3),
           (2, 2, 1), (2, 2, 2), (2, 2, 3),
           (2, 2, 4), (2, 2, 5), (2, 2, 6)]

test_df = spark.createDataFrame(row_data, ["A", "B", "C"])

test_df.show()

执行代码后返回的DataFrame内容如下：

+---+---+---+
|  A|  B|  C|
+---+---+---+
|  1|  1|  1|
|  1|  1|  2|
|  1|  1|  3|
|  1|  2|  1|
|  1|  2|  2|
|  1|  2|  3|
|  2|  1|  1|
|  2|  1|  2|
|  2|  1|  3|
|  2|  2|  1|
|  2|  2|  2|
|  2|  2|  3|
|  2|  2|  4|
|  2|  2|  5|
|  2|  2|  6|
+---+---+---+

固定长度条件的常规写法 #

当过滤条件列表长度固定时，可以直接用&（逻辑与）、|（逻辑或）运算符手动拼接条件。例如要求列A、列B均等于1的条件列表：

l = [func.col("A") == 1, func.col("B") == 1]

拼接两个条件后过滤DataFrame的代码如下：

t = l[0] & l[1]
test_df.filter(t).show()

执行后返回结果：

+---+---+---+
|  A|  B|  C|
+---+---+---+
|  1|  1|  1|
|  1|  1|  2|
|  1|  1|  3|
+---+---+---+

待解决问题 #

如果存储过滤条件的列表l长度为未知值n（元素均为PySpark Column类型的过滤条件），最符合Pythonic风格的单行实现方式是什么，能够按照逻辑与（&）或者逻辑或（|）的规则组合所有条件？

说明：Python内置的all()和any()函数无法实现该需求，这两个函数仅支持处理由True/False布尔值组成的普通列表，无法处理PySpark Column类型的条件表达式，无法适配l = [func.col("A") == 1, func.col("B") == 1, func.col("C") == 2]这类动态长度条件场景。

实现方案 #

使用Python标准库functools.reduce配合operator模块的逻辑运算函数即可实现，写法简洁无冗余，完全适配任意长度的条件列表：

• 逻辑与组合（所有条件同时满足）：传入operator.and_作为归并函数
• 逻辑或组合（任意条件满足即可）：传入operator.or_作为归并函数

完整示例代码：

from functools import reduce
from operator import and_, or_

# 动态长度条件列表示例
l = [func.col("A") == 1, func.col("B") == 1, func.col("C") == 2]

# 单行组合逻辑与条件
and_filter = reduce(and_, l)
test_df.filter(and_filter).show()

# 单行组合逻辑或条件
or_filter = reduce(or_, l)
test_df.filter(or_filter).show()

实现原理：reduce会从左到右依次对列表元素应用传入的二元运算函数，对长度为n的条件列表，reduce(and_, l)的计算逻辑等价于((...(l[0] & l[1]) & l[2]) & ... ) & l[n-1]，和手动拼接固定长度条件的逻辑完全一致，不需要手动编写循环处理。
以上述3个条件的逻辑与过滤为例，执行后返回结果为：

+---+---+---+
|  A|  B|  C|
+---+---+---+
|  1|  1|  2|
+---+---+---+

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