pyspark:withcolumn()有两个条件和三个结果

jdg4fx2g  于 2021-05-27  发布在  Spark
关注(0)|答案(3)|浏览(1769)

我和spark和Pypark一起工作。我正在尝试实现与以下伪代码等效的结果:

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 .

34gzjxbg

34gzjxbg1#

有几种有效的方法来实现这一点。让我们从所需的导入开始:

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

n3h0vuf22#

pyspark中的withcolumn函数允许您创建一个带有条件的新变量,添加when和others函数,这样您就拥有了一个正常工作的if-then-else结构。对于所有这些,您需要导入sparksql函数,因为您将看到,如果没有col()函数,下面的代码将无法工作。在第一位中,我们声明了一个新列-“new column”,然后给出when函数中包含的条件(即,fruit1==fruit2),如果条件为真,则给出1,如果为true,则控件将转到otherwise,后者使用isnull()函数处理第二个条件(水果1或水果2为null),如果返回true 3,如果为false,则再次检查otherwise,并给出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))
bmvo0sr5

bmvo0sr53#

您需要使用如下自定义项

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

相关问题