PySpark: withColumn() с двумя условиями и тремя результатами

Question

Я работаю с Spark и PySpark. Я пытаюсь добиться результата, эквивалентного следующего псевдокод:

df = df.withColumn('new_column', 
    IF fruit1 == fruit2 THEN 1, ELSE 0. IF fruit1 IS NULL OR fruit2 IS NULL 3.) 

Я пытаюсь сделать это в PySpark, но я не уверен в синтаксисе. Любые указатели? Я просмотрел expr(), но не смог заставить его работать.

Отметьте, что df является pyspark.sql.dataframe.DataFrame.

Answer 1

Есть несколько эффективных способов реализации этого. Начнем с необходимым импортом:

from pyspark.sql.functions import col, expr, when 

Вы можете использовать улую IF функции внутри выража:

new_column_1 = expr(
    """IF(fruit1 IS NULL OR fruit2 IS NULL, 3, IF(fruit1 = fruit2, 1, 0))""" 
) 

или when + otherwise:

new_column_2 = when(
    col("fruit1").isNull() | col("fruit2").isNull(), 3 
).when(col("fruit1") == col("fruit2"), 1).otherwise(0) 

Наконец, вы можете использовать следующий трюк:

from pyspark.sql.functions import coalesce, lit 

new_column_3 = coalesce((col("fruit1") == col("fruit2")).cast("int"), lit(3)) 

с примерными данными:

df = sc.parallelize([ 
    ("orange", "apple"), ("kiwi", None), (None, "banana"), 
    ("mango", "mango"), (None, None) 
]).toDF(["fruit1", "fruit2"]) 

вы можете использовать это следующим образом:

(df 
    .withColumn("new_column_1", new_column_1) 
    .withColumn("new_column_2", new_column_2) 
    .withColumn("new_column_3", new_column_3)) 

и результат:

+------+------+------------+------------+------------+ 
|fruit1|fruit2|new_column_1|new_column_2|new_column_3| 
+------+------+------------+------------+------------+ 
|orange| apple|   0|   0|   0| 
| kiwi| null|   3|   3|   3| 
| null|banana|   3|   3|   3| 
| mango| mango|   1|   1|   1| 
| null| null|   3|   3|   3| 
+------+------+------------+------------+------------+ 

Answer 2

Вы хотите использовать UDF, как показано ниже

from pyspark.sql.types import IntegerType 
from pyspark.sql.functions import udf 

def func(fruit1, fruit2): 
    if fruit1 == None or fruit2 == None: 
     return 3 
    if fruit1 == fruit2: 
     return 1 
    return 0 

func_udf = udf(func, IntegerType()) 
df = df.withColumn('new_column',func_udf(df['fruit1'], df['fruit2']))