Question

我有数据，

数据

ID          filter
1             A
2             A
3             A
4             A
5             B
6             B
7             B
8             B

我想为数据框应用函数

def add(x):
    y = x+1
    return(y)

from pyspark.sql.functions import *                 
from pyspark.sql.functions import udf

ol_val = udf(add, StringType())

data = data.withColumn("sum",ol_val(data.ID))

这给出了一个输出，

数据

ID          filter        sum
1             A            2
2             A            3
3             A            4
4             A            5
5             B            6
6             B            7
7             B            8
8             B            9

我想仅在filter = A 时应用此功能，而用于剩余部分我希望它为NULL 。我想要的输出是，

数据

ID          filter        sum
1             A            2
2             A            3
3             A            4
4             A            5
5             B            NULL
6             B            NULL
7             B            NULL
8             B            NULL

这里的值为NULL，因为它不满足条件filter = A.我希望仅当filter = A时才应用该函数。

有人可以帮我更改代码以便在pyspark中获得此输出吗？

Answer 1

您需要使用when和otherwise。顺便说一下，您不必创建UDF。

df = sc.parallelize([
    (1, "a"),
    (1, "b"),
    (3, "c")
  ]).toDF(["id", "filter"])

df.select("*", when(col("filter") == lit("a"), col("id") + 1).otherwise(None).alias("result")).show()

如果您确实需要调用该功能，只需将col("id") + 1替换为yourUDF(col("id"))

在pyspark中为过滤的数据框调用函数

1 个答案: