PySpark自定义窗口函数

在本文中，我们将介绍如何在PySpark中自定义窗口函数。窗口函数是对数据集中某个特定范围内的数据进行计算的函数。PySpark提供了许多内置的窗口函数，但有时我们可能需要根据自己的需求定制一个特定的窗口函数。

阅读更多：PySpark 教程

什么是窗口函数？

窗口函数是一种特殊的函数，它可以在数据集中的特定窗口（或范围）内计算结果。这个窗口可以是一段时间、一组行或者其他一些自定义的条件。与普通函数不同，窗口函数可以访问其他行的数据，因此可以进行复杂的聚合计算。

如何定义窗口函数？

在PySpark中定义一个窗口函数需要以下几个步骤：

定义一个窗口规范（WindowSpec）：窗口规范包含了窗口函数需要的各种信息，如窗口的大小、排序方式等。可以使用PySpark中的Window类来创建窗口规范。例如：windowSpec = Window.partitionBy("col1").orderBy("col2").rowsBetween(-1, 1)表示按照”col1″进行分区，并按照”col2″排序，在当前行的前一行和后一行之间的范围内进行计算。
定义一个窗口函数：窗口函数是一个可以应用于窗口中的数据的函数。可以使用PySpark中的pyspark.sql.functions模块定义内置的窗口函数，或者使用Python的lambda表达式定义自定义的窗口函数。
应用窗口函数：使用pyspark.sql.functions模块中的over方法将窗口函数应用到数据集。例如：df.withColumn("result", func(col("col1")).over(windowSpec))表示将窗口函数func应用到数据集的”col1″列上，并将计算结果保存在名为”result”的新列中。

下面我们将通过一个示例来演示如何使用自定义窗口函数。

示例

假设我们有一个包含学生姓名、科目和分数的数据集。我们想要计算每个学生在每个科目上的总分和平均分，以及每个科目上的最高分和最低分。通过定义一个自定义的窗口函数，我们可以轻松实现这些计算。

首先，让我们创建一个示例数据集：

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, sum, avg, max, min
from pyspark.sql.window import Window

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

# 创建示例数据集
data = [("Alice", "Math", 90),
        ("Alice", "Science", 80),
        ("Bob", "Math", 70),
        ("Bob", "Science", 75),
        ("Bob", "English", 85),
        ("Charlie", "Science", 95),
        ("Charlie", "English", 90)]

df = spark.createDataFrame(data, ["name", "subject", "score"])
df.show()

输出结果如下：

+-------+-------+-----+
|   name|subject|score|
+-------+-------+-----+
|  Alice|   Math|   90|
|  Alice|Science|   80|
|    Bob|   Math|   70|
|    Bob|Science|   75|
|    Bob|English|   85|
|Charlie|Science|   95|
|Charlie|English|   90|
+-------+-------+-----+

接下来，我们定义一个窗口规范，并使用自定义的窗口函数计算每个学生在每个科目上的总分、平均分、最高分和最低分。

windowSpec = Window.partitionBy("name", "subject")

total_score = sum(col("score")).over(windowSpec)
average_score = avg(col("score")).over(windowSpec)
max_score = max(col("score")).over(windowSpec)
min_score = min(col("score")).over(windowSpec)

result = df.withColumn("total_score", total_score)\
           .withColumn("average_score", average_score)\
           .withColumn("max_score", max_score)\
           .withColumn("min_score", min_score)\
           .orderBy("name", "subject")
result.show()

输出结果如下：

+-------+-------+-----+-----------+------------------+---------+---------+
|   name|subject|score|total_score|     average_score|max_score|min_score|
+-------+-------+-----+-----------+------------------+---------+---------+
|  Alice|   Math|   90|       90.0|              90.0|     90.0|     90.0|
|  Alice|Science|   80|       80.0|              80.0|     80.0|     80.0|
|    Bob|English|   85|       85.0|              85.0|     85.0|     85.0|
|    Bob|   Math|   70|       70.0|              70.0|     70.0|     70.0|
|    Bob|Science|   75|75.0       |              75.0|     75.0|     75.0|
|Charlie|English|   90|       90.0|              90.0|     90.0|     90.0|
|Charlie|Science|   95|       95.0|              95.0|     95.0|     95.0|
+-------+-------+-----+-----------+------------------+---------+---------+