PySpark如何从列表选取多列并按不同值筛选数据

核心实现代码 #

首先导入依赖：

from functools import reduce
from pyspark.sql.functions import col

定义你的列名列表和对应要匹配的目标值：

# 按需求自定义列列表和匹配值
list1 = ["col1", "col3", "col4", "col11"]
target_value1 = 0
list2 = ["col2", "col6", "col9", "col10"]
target_value2 = 1

动态生成筛选条件并执行查询：

# 拼接list1所有列等于目标值的与条件
condition1 = reduce(lambda a, b: a & b, [col(c) == target_value1 for c in list1])
# 拼接list2所有列等于目标值的与条件
condition2 = reduce(lambda a, b: a & b, [col(c) == target_value2 for c in list2])
# 合并条件筛选数据
df1 = df.filter(condition1 & condition2)

兼容空列表的优化版本 #

如果存在list1或list2为空的场景，可增加兜底逻辑避免报错：

from pyspark.sql.functions import lit

def generate_condition(col_list, target):
    if not col_list:
        return lit(True)
    return reduce(lambda a, b: a & b, [col(c) == target for c in col_list])

condition1 = generate_condition(list1, 0)
condition2 = generate_condition(list2, 1)
df1 = df.filter(condition1 & condition2)

说明 #

• 方案适配任意长度的列列表，无需手动修改筛选逻辑，仅调整列名列表和目标值即可
• 全程使用PySpark原生表达式，无自定义UDF开销，性能和手动逐列编写的筛选条件完全一致，适配百万行、数千列的大表场景
• 注意列名字符串需要和DataFrame实际列名的大小写完全匹配

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