PySpark：传递多列到UDF

在本文中，我们将介绍如何在PySpark中将多列传递给用户定义函数（UDF）。UDF允许我们对数据应用自定义的转换逻辑，让我们先来了解一下如何使用UDF来处理多列数据。

什么是UDF？

UDF是用户定义函数（User-Defined Function）的缩写，可以让我们在Spark中自定义转换逻辑。UDF接受输入数据并返回转换后的结果。通过使用UDF，我们可以扩展Spark的功能，以满足特定的需求。

在PySpark中创建UDF

在PySpark中，我们可以使用udf函数来创建UDF。udf函数接受一个函数和返回类型作为参数，并返回一个UDF函数对象。我们可以将这个函数对象应用到DataFrame的列上，以实现自定义的转换操作。

让我们以一个简单的例子来说明如何创建一个接受多列输入的UDF。假设我们有一个包含学生信息的DataFrame，其中包括学生的姓名和年龄两列。我们想要创建一个UDF，用于将年龄大于18岁的学生的姓名转换为大写。

from pyspark.sql.functions import udf
from pyspark.sql.types import StringType

# 创建一个接受两列输入的UDF
def uppercase_name(name, age):
    if age > 18:
        return name.upper()
    else:
        return name

# 将Python函数转换为UDF函数对象
udf_uppercase_name = udf(uppercase_name, StringType())

# 在DataFrame上应用UDF
df = df.withColumn("uppercase_name", udf_uppercase_name(df.name, df.age))

在上面的代码中，我们定义了一个名为uppercase_name的Python函数，该函数接受两个参数（姓名和年龄）。如果年龄大于18岁，函数返回大写的姓名；否则，返回原始姓名。然后，我们使用udf函数将这个Python函数转换为UDF函数对象。最后，我们使用withColumn方法将应用了UDF的新列添加到DataFrame中。

传递多列到UDF

为了传递多列到UDF中，我们需要做一些额外的工作。PySpark提供了struct函数，可以将多个列打包成一个复合列。我们可以使用struct函数将多个列合并为一个列，并将该列用作传递给UDF的参数。

让我们继续上面的例子，假设我们希望将姓名和年龄两列作为参数传递给UDF，并且我们希望将年龄乘以2，然后将结果作为新列添加到DataFrame中。

from pyspark.sql.functions import col, udf
from pyspark.sql.types import IntegerType, StringType, StructType, StructField

# 创建一个接受结构化输入的UDF
def double_age(struct):
    name = struct["name"]
    age = struct["age"]
    return age * 2

# 定义结构化输入的字段类型
input_struct_type = StructType([
    StructField("name", StringType(), True),
    StructField("age", IntegerType(), True)
])

# 将Python函数转换为UDF函数对象
udf_double_age = udf(double_age, IntegerType())

# 使用struct函数将两列合并为一个复合列
df = df.withColumn("input_struct", struct(col("name"), col("age")))

# 在DataFrame上应用UDF
df = df.withColumn("double_age", udf_double_age(df.input_struct))

在上面的代码中，我们定义了一个名为double_age的Python函数，该函数接受一个结构化输入，其中包括name和age两个字段。函数将年龄乘以2，并返回结果。

我们通过定义input_struct_type来指定结构化输入的字段类型，包括一个name字段和一个age字段。然后，我们使用struct函数将name和age两个列合并为一个复合列，并将该复合列命名为input_struct。最后，我们使用withColumn方法将应用了UDF的新列添加到DataFrame中。

总结

在本文中，我们介绍了如何在PySpark中将多列传递给UDF。首先，我们创建一个接受多列输入的UDF函数，并将其应用到DataFrame的列上。然后，我们使用struct函数将多个列合并为一个复合列，并将该列作为参数传递给UDF。通过使用UDF，我们可以根据自己的需求对数据进行自定义转换操作，从而扩展PySpark的功能。