Pyspark 数据帧中的条件

Question

我有以下数据框

Column_1     Column_2       Column_3
A            1
A            2
A            3
A            4
A            5
B            1
B            4
B            5
C            1
C            2

我必须根据 Column_1 和 Column_2 中的值填充 Column_3。如果 Column_1 在 ('A','B') 和 Column_2 不在 ('1','3','5') 我必须填充 column_3用 X 否则 Y。

预期输出：

Column_1     Column_2       Column_3
A            1              Y
A            2              X
A            3              Y
A            4              X
A            5              Y
B            1              Y            
B            4              X
B            5              Y
C            1              Y
C            2              Y

我尝试了什么：

我尝试使用 when 和 otherwise 语句，但不确定如何将 Not in 与 when 语句一起使用。对此的任何帮助将不胜感激

Answer 1

您可以利用 isin 和反向 ~:

import pyspark.sql.functions as F

c = (F.when(df['Column_1'].isin(['A','B']) & 
      (~df['Column_2'].isin([1,3,5])),'X').otherwise('Y'))
df.withColumn("Column_3",c).show()

或：

expr = """CASE 
        WHEN Column_1 IN ('A','B') and Column_2 NOT IN (1,3,5) 
        THEN 'X' ELSE 'Y' 
        END as Column_3"""
df.selectExpr("*",expr).show()

+--------+--------+--------+
|Column_1|Column_2|Column_3|
+--------+--------+--------+
|       A|       1|       Y|
|       A|       2|       X|
|       A|       3|       Y|
|       A|       4|       X|
|       A|       5|       Y|
|       B|       1|       Y|
|       B|       4|       X|
|       B|       5|       Y|
|       C|       1|       Y|
|       C|       2|       Y|
+--------+--------+--------+

更多详情：

df['Column_1'].isin(['A','B'])
#Column<b'(Column_1 IN (A, B))'>
~df['Column_2'].isin([1,3,5])
#Column<b'(NOT (Column_2 IN (1, 3, 5)))'>

Pyspark 数据帧中的条件

When condition in Pyspark dataframe

dataframe

.when

pyspark