PySpark：withColumn（）具有两个条件和三个结果

user2205916 发表于 Dev

user2205916：

我正在使用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。

0323：

有几种有效的方法可以实现此目的。让我们从必需的导入开始：

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

您可以IF在expr中使用Hive 函数：

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

本文收集自互联网，转载请注明来源。

如有侵权，请联系 [email protected] 删除。