PySpark 使用DataFrame计算分组累计和

在本文中，我们将介绍如何使用PySpark的DataFrame计算分组累计和。累计和是指对于每个分组，按照指定的列累计计算前面行的和。

1. 创建DataFrame

首先，我们需要创建一个DataFrame来演示如何计算分组累计和。假设我们有一个包含三个列的DataFrame，分别为group、value1和value2，其中group列表示分组的标识，value1和value2表示需要计算累计和的数值列。下面是一个示例DataFrame的代码：

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

spark = SparkSession.builder.getOrCreate()

data = [("A", 1, 10), ("A", 2, 20), ("A", 3, 30), ("B", 4, 40), ("B", 5, 50), ("B", 6, 60)]
df = spark.createDataFrame(data, ["group", "value1", "value2"])
df.show()

运行以上代码，将得到如下的输出：

+-----+------+------+ 
|group|value1|value2| 
+-----+------+------+ 
|    A|     1|    10| 
|    A|     2|    20| 
|    A|     3|    30| 
|    B|     4|    40| 
|    B|     5|    50| 
|    B|     6|    60| 
+-----+------+------+

我们的目标是计算每个分组的value1和value2的累计和。

2. 使用窗口函数计算分组累计和

PySpark提供了窗口函数和窗口规范来进行分组操作。我们可以使用窗口函数sum()来计算每个分组的累计和。以下是使用窗口函数计算分组累计和的代码：

windowSpec = Window.partitionBy("group").orderBy("value1")

df = df.withColumn("cumulative_sum_value1", sum("value1").over(windowSpec))
df = df.withColumn("cumulative_sum_value2", sum("value2").over(windowSpec))

df.show()

运行以上代码，将得到如下的输出：

+-----+------+------+---------------------+---------------------+
|group|value1|value2|cumulative_sum_value1|cumulative_sum_value2|
+-----+------+------+---------------------+---------------------+
|    A|     1|    10|                    1|                   10|
|    A|     2|    20|                    3|                   30|
|    A|     3|    30|                    6|                   60|
|    B|     4|    40|                    4|                   40|
|    B|     5|    50|                    9|                   90|
|    B|     6|    60|                   15|                  150|
+-----+------+------+---------------------+---------------------+

通过使用窗口函数和窗口规范，我们成功地计算了每个分组的累计和。

3. 自定义窗口规范计算分组累计和

除了使用默认的窗口规范外，我们还可以自定义窗口规范来计算分组累计和。自定义窗口规范可以包含多个排序列，并且可以指定窗口的范围。以下是使用自定义窗口规范计算分组累计和的代码：

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

windowSpec = Window.partitionBy("group").orderBy("value1").rowsBetween(Window.unboundedPreceding, Window.currentRow)

df = df.withColumn("cumulative_sum_value1", sum("value1").over(windowSpec))
df = df.withColumn("cumulative_sum_value2", sum("value2").over(windowSpec))

df.show()

运行以上代码，将得到如下的输出：

+-----+------+------+---------------------+---------------------+
|group|value1|value2|cumulative_sum_value1|cumulative_sum_value2|
+-----+------+------+---------------------+---------------------+
|    A|     1|    10|                    1|                   10|
|    A|     2|    20|                    3|                   30|
|    A|     3|    30|                    6|                   60|
|    B|     4|    40|                    4|                   40|
|    B|     5|    50|                    9|                   90|
|    B|     6|    60|                   15|                  150|
+-----+------+------+---------------------+---------------------+

通过自定义窗口规范，我们可以指定窗口的范围，从而获取不同范围内的累计和。