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

Я работаю со 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.

Ответ 1

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

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

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

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|
+------+------+------------+------------+------------+

Ответ 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']))

Ответ 3

функция withColumn в pyspark позволяет вам создать новую переменную с условиями, добавить функции when и else, и у вас будет правильно работающая структура if тогда else. Для всего этого вам нужно будет импортировать функции sparrsql, как вы увидите что следующий фрагмент кода не будет работать без функции col(). В первом бите мы объявляем новый столбец - "новый столбец", а затем задаем условие, заключенное в функцию "когда" (т.е. fruit1 == fruit2), затем даем 1, если условие истинно, если не соответствует элементу управления, иначе затем выполняет второе условие (fruit1 или fruit2 равно Null) с помощью функции isNull(), и если возвращается значение 3, а если значение false, в противном случае проверяется снова, давая 0 в качестве ответа

 from pyspark.sql import functions as F        
   df=df.withColumn('new_column', F.when(F.col('fruit1')==F.col('fruit2),1)
                          .otherwise(F.when((F.col('fruit1').isNull()) |(F.col('fruit2').isNull()),3))
                                    .otherwise(0))