如何按类别时序累加填充PySpark DataFrame的Stock-level列

要实现你需要的按Category分组、Time升序，用前一行Stock-level加当前行Stock-change填充当前行Stock-level的需求，直接用窗口函数的累积求和就能搞定，具体实现如下：

步骤1：导入依赖并创建示例DataFrame #

首先导入PySpark所需的函数，然后创建和你输入一致的测试DataFrame：

from pyspark.sql import SparkSession
from pyspark.sql.window import Window
from pyspark.sql.functions import col, first, sum as spark_sum, coalesce

# 初始化SparkSession
spark = SparkSession.builder.appName("StockLevelFill").getOrCreate()

# 创建输入DataFrame
data = [
    ("apple", 1, 4, None),
    ("apple", 2, None, -2),
    ("apple", 3, None, 5),
    ("banana", 1, 12, None),
    ("banana", 2, None, 4),
    ("orange", 1, 1, None),
    ("orange", 2, None, -7)
]

df = spark.createDataFrame(data, ["Category", "Time", "Stock-level", "Stock-change"])

步骤2：定义窗口规则 #

按Category分组，Time升序排序，这个窗口是后续计算的基础：

window_spec = Window.partitionBy("Category").orderBy("Time").rowsBetween(Window.unboundedPreceding, Window.currentRow)

步骤3：计算填充后的Stock-level #

核心逻辑是：

• 取每个分组的第一个Stock-level（也就是初始库存）
• 计算从第一行到当前行的Stock-change累积和
• 用coalesce保留初始行的原始Stock-level，后续行用「初始库存 + 累积变动」填充

filled_df = df.withColumn(
    "Stock-level",
    coalesce(
        col("Stock-level"),
        first(col("Stock-level")).over(window_spec) + spark_sum(col("Stock-change")).over(window_spec)
    )
)

步骤4：查看结果 #

执行filled_df.show()就能得到你期望的输出：

+--------+----+-----------+------------+
|Category|Time|Stock-level|Stock-change|
+--------+----+-----------+------------+
|   apple|   1|          4|        null|
|   apple|   2|          2|          -2|
|   apple|   3|          7|           5|
|  banana|   1|         12|        null|
|  banana|   2|         16|           4|
|  orange|   1|          1|        null|
|  orange|   2|         -6|          -7|
+--------+----+-----------+------------+

关键逻辑说明 #

• first(col("Stock-level")).over(window_spec)：获取每个分组里第一行的Stock-level，也就是初始库存值
• spark_sum(col("Stock-change")).over(window_spec)：计算从分组第一行到当前行的Stock-change累积和，比如apple组第三行的累积和是-2+5=3，加上初始的4就得到7
• coalesce函数：优先使用原始的Stock-level（也就是第一行的非null值），只有当原始值为null时，才用计算后的填充值

内容的提问来源于stack exchange，提问作者Peter