Когда условие в фреймворке 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_3 на основе значений в Column_1 и Column_2. Если 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 и else, но не уверен, как использовать 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)))'>