将函数映射到pyspark dataframe的多列

Question

我有一个 pyspark df，它有很多列，但子集看起来像这样：

我想映射一个函数some_func（），它只使用列“lat”、“lon”和“event_id”来返回一个布尔值，该值将作为名为“验证”的单独列添加到df中。基本上，我需要单独检索函数中感兴趣的列并对它们进行操作。我知道我可以使用UDF或df.withColumn（），但它们用于映射到单列。为此，我需要将感兴趣的列连接为一列，这会使代码有点混乱。

有没有办法分别检索函数中的列值，并将该函数映射到整个数据帧？(类似于我们使用map-lambda对Pandas df所做的操作

Answer 1

您可以创建一个udf，它可以将多个列作为参数

前任：

from pyspark.sql.functions as f
from pyspark.sql.types import BooleanType

def your_function(p1, p2, p3):
    # your logic goes here
    # return a bool

udf_func = f.udf(your_function, BooleanType())


df = spark.read.....

df2 = df.withColumn("verified", udf_func(f.col("lat"), f.col("lon"), f.col("event_id")))

df2.show(truncate=False)