使用DataFrame - Pyspark按组累计求和

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

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

# 创建示例数据
data = [("A", 1), ("A", 2), ("B", 3), ("B", 4), ("B", 5)]

# 创建DataFrame
df = spark.createDataFrame(data, ["group", "value"])

# 定义窗口规范
windowSpec = Window.partitionBy("group").orderBy("value").rowsBetween(Window.unboundedPreceding, 0)

# 添加累计求和列
df = df.withColumn("cumulative_sum", sum(df.value).over(windowSpec))

# 显示结果
df.show()

+-----+-----+--------------+
|group|value|cumulative_sum|
+-----+-----+--------------+
|    A|    1|             1|
|    A|    2|             3|
|    B|    3|             3|
|    B|    4|             7|
|    B|    5|            12|
+-----+-----+--------------+