Pyspark -添加一个列，对先前连续的累积值进行计数

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

spark = SparkSession.builder.appName("Counting Cumulative Values").getOrCreate()

data = [("A", 1), ("A", 2), ("A", 3), ("B", 4), ("B", 5), ("B", 6)]
df = spark.createDataFrame(data, ["col1", "col2"])
df.show()

+----+----+
|col1|col2|
+----+----+
|   A|   1|
|   A|   2|
|   A|   3|
|   B|   4|
|   B|   5|
|   B|   6|
+----+----+

window_spec = Window.partitionBy("col1").orderBy("col2")
df = df.withColumn("cumulative_count", sum(col("col2")).over(window_spec))
df.show()

+----+----+----------------+
|col1|col2|cumulative_count|
+----+----+----------------+
|   A|   1|               1|
|   A|   2|               3|
|   A|   3|               6|
|   B|   4|               4|
|   B|   5|               9|
|   B|   6|              15|
+----+----+----------------+