将函数映射到 pyspark 数据框的多列

Posted 2023-04-15

【中文标题】将函数映射到 pyspark 数据框的多列

【英文标题】：Mapping a function to multiple columns of pyspark dataframe

【发布时间】：2021-09-06 07:34:54

【问题描述】：

我有一个 pyspark df，它有很多列，但一个子集如下所示：

datetime	eventid	sessionid	lat	lon	filtertype
someval	someval	someval	someval	someval	someval
someval	someval	someval	someval	someval	someval

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

有没有办法单独检索函数内的列值并将该函数映射到整个数据框？ （类似于我们可以使用 map-lambda 和 df.apply() 对 Pandas df 执行的操作）？

【问题讨论】：

这能回答你的问题吗？ Pyspark: Pass multiple columns in UDF

【参考方案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)