PySpark: withColumn() con dos condiciones y tres resultados

Presta atención ya que en esta noticia vas a encontrar la respuesta que buscas.Esta sección ha sido evaluado por nuestros especialistas para garantizar la calidad y exactitud de nuestro contenido.

Solución:

Hay algunas formas eficientes de implementar esto. Comencemos con las importaciones requeridas:

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

Puedes usar Colmena IF función dentro de expr:

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

o when + otherwise:

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

Finalmente podrías usar el siguiente truco:

from pyspark.sql.functions import coalesce, lit

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

Con datos de ejemplo:

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

puedes usar esto de la siguiente manera:

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

y el resultado es:

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

Querrás usar un udf como se muestra a continuación

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

La función withColumn en pyspark le permite crear una nueva variable con condiciones, agregar en el Cuándo y de lo contrario funciona correctamente y tiene una estructura if then else que funciona correctamente. Para todo esto, necesitaría importar las funciones de sparksql, ya que verá que el siguiente fragmento de código no funcionará sin la función col(). En el primer bit, declaramos una nueva columna -‘nueva columna’, y luego damos la condición incluida en la función when (es decir, fruit1==fruit2) y luego damos 1 si la condición es true, si no es cierto, el control pasa a lo contrario, que luego se encarga de la segunda condición (fruta1 o fruta2 es Nulo) con la función isNull() y si true 3 se devuelve y si false, en caso contrario se vuelve a comprobar dando 0 como respuesta.

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

Nos puedes añadir valor a nuestro contenido informacional añadiendo tu experiencia en las reseñas.

---

---

Utiliza Nuestro Buscador