Pyspark. Create A new column based on 多列多行

Question

我下面有一个 dataframe：

ID	特点1	特征2	特点3	特征4
1个	A	0	0	0
1个	b	0	0	0
1个	c	0	0	0
2个	A	1个	0	0
2个	b	0	0	0
2个	c	0	0	0
3个	b	0	0	0
3个	c	0	1个	0

我需要一个新的专栏并获得下面的 dataframe：

ID	特点1	特征2	特点3	特征4	新专栏
1个	A	0	0	0	1个
1个	b	0	0	0	1个
1个	c	0	0	0	1个
2个	A	1个	0	0	0
2个	b	0	0	0	0
2个	c	0	0	0	0
3个	b	0	0	0	0
3个	c	0	1个	0	0

对于id=1的情况，newColumn 应该为 1，因为所有具有feature2 、 feature3和feature4的行都是 0。但是，对于id==2的情况，NewColumn 应该为 0，因为在这种情况下，当feature1=a时feature2为 1。 id==3时相同，NewColumn 应该为 0，因为feature3包含 1。基本上，我需要单独考虑所有 id，并检查feature2 、 feature3和feature4列是否包含 1。如何在 Pyspark 中有效地实现它？

Answer 1

由于它需要跨行检查值，您可以使用collect_list 、 feature3和feature4 array的collect_list超过 window 按列id分区， flatten数组，使用exists检查数组是否包含1的任何值，然后cast否定的 boolean 值到 integer。（对于 Spark 版本 >= 2.4）

from pyspark.sql.functions import expr

df = df.withColumn('NewColumn', 
                   expr("""cast(
                             not exists(
                               flatten(
                                 collect_list(array(feature2,feature3,feature4)) over (partition by id)),
                               v -> v = 1)
                             as int)"""))
df.show()

# +---+--------+--------+--------+--------+---------+
# | id|feature1|feature2|feature3|feature4|NewColumn|
# +---+--------+--------+--------+--------+---------+
# |  1|       a|       0|       0|       0|        1|
# |  1|       b|       0|       0|       0|        1|
# |  1|       c|       0|       0|       0|        1|
# |  3|       b|       0|       0|       0|        0|
# |  3|       c|       0|       1|       0|        0|
# |  2|       a|       1|       0|       0|        0|
# |  2|       b|       0|       0|       0|        0|
# |  2|       c|       0|       0|       0|        0|
# +---+--------+--------+--------+--------+---------+