PySpark 用户定义函数应用到PySpark中的窗口

在本文中，我们将介绍如何在 PySpark 中使用自定义函数（User Defined Function，简称UDF）应用到窗口函数（Window Function）中。

什么是窗口函数

窗口函数在 PySpark 中是一个非常有用和强大的功能。它允许我们在查询结果集上进行聚合操作，并根据定义的窗口范围进行分组和排序。这样我们可以根据特定的需求对查询结果进行更精细的控制和计算。

如何创建自定义函数

在 PySpark 中，我们可以使用 udf() 函数来创建自定义函数。这个函数接受一个 Python 函数作为参数，并将其封装为可在 Spark SQL 中使用的 UDF。

下面是一个示例，展示如何创建一个简单的自定义函数，将字符串转换为大写：

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

def uppercase(s):
    return s.upper()

uppercase_udf = udf(uppercase, StringType())

在上面的例子中，我们首先定义了一个名为 uppercase() 的函数，它接受一个字符串作为参数，然后返回该字符串的大写形式。接着，我们使用 udf() 函数将 uppercase() 函数封装为可在 Spark SQL 中使用的 UDF，并指定其返回类型为字符串。

如何将自定义函数应用到窗口

一旦我们创建了自定义函数，就可以将其应用到窗口函数中。为了演示这一点，我们将使用一个示例数据集，并计算每个部门每月的销售总额。假设我们有以下的示例数据：

from pyspark.sql import SparkSession
from pyspark.sql.window import Window
from pyspark.sql.functions import sum

# 创建一个 SparkSession
spark = SparkSession.builder.getOrCreate()

# 创建一个示例数据集
data = [("部门1", "2021-01-01", 100),
        ("部门1", "2021-01-02", 200),
        ("部门1", "2021-02-01", 150),
        ("部门1", "2021-02-02", 250),
        ("部门2", "2021-01-01", 300),
        ("部门2", "2021-01-02", 400),
        ("部门2", "2021-02-01", 350),
        ("部门2", "2021-02-02", 450)]

df = spark.createDataFrame(data, ["部门", "日期", "销售额"])

# 创建窗口规范
window_spec = Window.partitionBy("部门").orderBy("日期")

# 应用自定义函数到窗口
df_with_total_sales = df.withColumn("窗口销售总额", sum("销售额").over(window_spec))

df_with_total_sales.show()

上述代码中，我们首先创建了一个 SparkSession 对象。然后，我们创建了一个示例数据集，其中包含了每个部门每天的销售额。接下来，我们使用 Window.partitionBy() 方法按部门进行分组，并使用 Window.orderBy() 方法按日期进行排序，以定义窗口规范。最后，我们使用 withColumn() 方法将自定义函数 sum() 应用到窗口规范上，并创建一个名为 “窗口销售总额” 的新列。

运行以上代码，我们将得到类似下面的结果：

+----+----------+----+--------------+
|部门|日期      |销售额|窗口销售总额|
+----+----------+----+--------------+
|部门1|2021-01-01|100 |300           |
|部门1|2021-01-02|200 |300           |
|部门1|2021-02-01|150 |400           |
|部门1|2021-02-02|250 |400           |
|部门2|2021-01-01|300 |700           |
|部门2|2021-01-02|400 |700           |
|部门2|2021-02-01|350 |800           |
|部门2|2021-02-02|450 |800           |
+----+----------+----+--------------+

从上面的结果可以看出，我们成功地将自定义函数应用到了窗口函数中，并计算出了每个部门每月的销售总额。