PySpark:withColumn() 有两个条件和三个结果 [英] PySpark: withColumn() with two conditions and three outcomes
问题描述
我正在使用 Spark 和 PySpark.我正在尝试实现等效于以下伪代码的结果:
I am working with Spark and PySpark. I am trying to achieve the result equivalent to the following pseudocode:
df = df.withColumn('new_column',
IF fruit1 == fruit2 THEN 1, ELSE 0. IF fruit1 IS NULL OR fruit2 IS NULL 3.)
我正在尝试在 PySpark 中执行此操作,但我不确定语法.任何指针?我查看了 expr()
但无法让它工作.
I am trying to do this in PySpark but I'm not sure about the syntax. Any pointers? I looked into expr()
but couldn't get it to work.
注意 df
是一个 pyspark.sql.dataframe.DataFrame
.
推荐答案
有几种有效的方法可以实现这一点.让我们从必需的导入开始:
There are a few efficient ways to implement this. Let's start with required imports:
from pyspark.sql.functions import col, expr, when
你可以在 expr 中使用 Hive IF
函数:
You can use Hive IF
function inside 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)
最后你可以使用以下技巧:
Finally you could use following trick:
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"])
您可以按如下方式使用它:
you can use this as follows:
(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|
+------+------+------------+------------+------------+
这篇关于PySpark:withColumn() 有两个条件和三个结果的文章就介绍到这了,希望我们推荐的答案对大家有所帮助,也希望大家多多支持IT屋!