web-dev-qa-db-fra.com

PySpark: withColumn () avec deux conditions et trois résultats

Je travaille avec Spark et PySpark. J'essaie d'obtenir le résultat équivalent au pseudocode suivant:

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

J'essaie de faire cela dans PySpark mais je ne suis pas sûr de la syntaxe. Des pointeurs? J'ai regardé dans expr() mais je n'ai pas réussi à le faire fonctionner.

Notez que df est un pyspark.sql.dataframe.DataFrame.

apache-sparkhivepysparkapache-spark-sqlhiveql
27
user2205916

Il existe quelques moyens efficaces pour mettre cela en œuvre. Commençons par les importations requises:

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

Vous pouvez utiliser la fonction Hive IF dans expr:

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

ou when + otherwise:

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

Enfin, vous pouvez utiliser le truc suivant:

from pyspark.sql.functions import coalesce, lit

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

Avec des données d'exemple:

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

vous pouvez utiliser ceci comme suit:

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

et le résultat est:

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

Vous voudrez utiliser un fichier UDF comme ci-dessous

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']))
15
David