Pyspark 从列表中选择多个列并过滤不同的值

Question

我有一个大约 5k 列和大约 1 M 行的表格，如下所示：

ID	Col1	Col2	Col3	Col4	Col5	Col6	Col7	Col8	Col9	Col10	Col11
ID1	0	1	0	1	0	2	1	1	2	2	0
ID2	1	0	1	0	1	2	0	0	0	2	1
ID3	1	0	0	0	0	2	1	0	1	1	1
ID4	0	1	1	1	1	1	1	1	1	0	0
ID5	0	1	1	1	1	1	0	1	0	0	0

我想从不同的列表中选择与列名匹配的不同列，并根据不同的条件对行进行子集化。例如，如果我的 list1 有 col1、col3、col4、col11，而 list2 有 col2、col6、col9、col10。我想将行过滤为 list1 == 0 AND list2 == 1。例如 df1 = df.filter((df.col1 == 0) & (df.col3 == 0) & (df.col4 == 0) & (df.col6== 1) & (df.col9 == 1) & (df.col10 == 1)) 。我希望从两个不同的列表中选择这些列，而不是每次都添加列名。如何使用 PySpark 实现这一目标？

Answer 1

如果您需要像这样比较很多列，请考虑像这样进行元组比较：

from pyspark.sql.functions import lit, struct
source_tuple = struct(col("col1"), col("col3"), col("col4"), col("col6"), col("col9"), col("col10"))
target_tuple1 = struct([lit(0), lit(0), lit(0), lit(1), lit(1), lit(1)])
df1 = df.where(source_tuple.isin([target_tuple1]))

您可以动态构建左侧和右侧，并在右侧添加多个元组。如果您需要多个左侧，请使用 UNION ALL 组合生成的 DataFrame。

Answer 2

我找到了我正在寻找的解决方案： 我能够使用以下步骤过滤大量列：

1. 使用列 ID 创建了单独的列表。我的列 ID 在数据框中。所以我将该列转换为过滤过程列表。

list1 = ped.filter((ped.pheno == 2)).select("IID")
list1 = list1.select('IID').rdd.map(lambda row : row [0]).collect()`
list2 = ped.filter((ped.pheno == 1)).select("IID")
list2 = list2.select('IID').rdd.map(lambda row : row [0]).collect()

## 2. Using these lists I filtered the columns as follows:

df1 = df.where ("AND".join([(%s ==1)"%(col) for col in list1]))
df1 = df.where ("AND".join([(%s ==0)"%(col) for col in list2]))

也感谢其他有用的解决方案！