PySpark 用户定义函数多输入

在 PySpark 中，用户定义函数（UDF）提供了自定义函数功能，可以执行各种自定义操作。对于单个输入的 UDF，定义很简单。但是，在某些情况下，我们需要一个或多个输入的 UDF。在本文中，我们将学习如何在 PySpark 中定义具有多个输入的 UDF。

多输入 UDF 的语法

在 PySpark 中定义一个多输入 UDF 非常简单。你只需要为每个输入变量定义一个参数并返回 UDF 所需的输出。

下面是 PySpark 多输入 UDF 的语法：

from pyspark.sql.functions import udf
 
multi_input_udf = udf(lambda x1, x2, ... : your_function(x1, x2, ...), outputType)

在上面的语法中，x1, x2, ... 是你想要使用的输入变量，而outputType是你定义的输出类型。

示例

让我们通过以下示例了解如何使用 PySpark 中的多输入 UDF。

假设你有一个数据框，其中包含姓名、年龄和性别列。现在，你想要创建一个 UDF，以根据年龄计算用户是否成年。

以下是一个具有两个输入变量的多输入 UDF，它将用户的年龄和性别作为输入，并将用户是否成年作为输出：

from pyspark.sql.functions import udf
from pyspark.sql.types import BooleanType
 
def is_adult(age, gender):
    """判断用户是否成年"""
    if age >= 18 and gender == 'Male':
        return True
    elif age >= 21 and gender == 'Female':
        return True
    else:
        return False
 
is_adult_udf = udf(lambda age, gender: is_adult(age, gender), BooleanType())
 
df = spark.createDataFrame([
    (1, "John", 20, "Male"),
    (2, "Alice", 23, "Female")
], ["id", "name", "age", "gender"])
 
df = df.withColumn("is_adult", is_adult_udf(df["age"], df["gender"]))
 
df.show()

输出：

+---+-----+---+------+--------+
| id| name|age|gender|is_adult|
+---+-----+---+------+--------+
|  1| John| 20|  Male|    true|
|  2|Alice| 23|Female|    true|
+---+-----+---+------+--------+

上面的代码中，我们首先定义了一个名为 is_adult 的函数，该函数获取 age 和 gender 作为输入，并返回一个布尔值。接下来，我们使用 udf() 函数封装该函数，并将其输入类型设置为 age 和 gender，输出类型设置为布尔值。最后，我们使用 withColumn() 函数将 UDF 应用于数据框中的 age 和 gender 列，并将计算的结果赋值给新列 is_adult。

这就是 PySpark 中多输入 UDF 的全部了。它为您提供了许多自定义操作的可能性，使您可以快速编写复杂的数据处理逻辑。