基于PySpark计算带0下限的流入滚动累计值

问题需求 #

在PySpark DataFrame中计算滚动累计值，要求累计值不得为负（低于0时取0）。因数据量过大无法转换为Pandas处理，Excel中已实现该逻辑，但PySpark常规窗口函数无法复用前一行计算结果，导致滚动求和方法失效。

模拟数据集代码 #

from pyspark.sql import SparkSession
from pyspark.sql.functions import to_date, col

spark = SparkSession.builder.appName("CumulativeFloor").getOrCreate()

df2 = spark.createDataFrame ([
    ['30-Jan-2019', -200], ['30-Jun-2019', -100], ['30-Jan-2020',100], 
    ['30-Jun-2020', 50], ['30-Jan-2021', 300], ['30-Jan-2021', -500], 
    ['30-Jul-2021', 800], ['15-Sep-2021',-300]
], ["transaction_date", "transaction_amount"])

df2 = df2.select(to_date(col("transaction_date"), "dd-MMM-yyyy").alias('transaction_date'), 'transaction_amount')

尝试过的无效代码 #

windowval = (Window.orderBy('posting_date').rangeBetween(Window.unboundedPreceding, 0))

df2.withColumn('cum_sum', when(col('transaction_amount')>0, sum(col('net_sls')).over(windowval)).otherwise(0)).show()

注：该代码存在列名不匹配（posting_date/net_sls为错误列名）及逻辑错误，未实现带下限的累计逻辑。

解决方案：递归CTE实现 #

PySpark窗口函数基于原始数据集计算，无法直接依赖前一行的计算结果，因此需要使用递归CTE来实现这种依赖前序结果的累计逻辑。

完整实现代码 #

from pyspark.sql.functions import row_number
from pyspark.sql.window import Window

# 1. 为数据添加排序后的行号，确保计算顺序正确
order_window = Window.orderBy('transaction_date')
df_with_row = df2.withColumn("row_num", row_number().over(order_window))
df_with_row.createOrReplaceTempView("temp_transactions")

# 2. 递归CTE计算带0下限的累计值
result_df = spark.sql("""
WITH RECURSIVE cumulative_cte AS (
    -- 初始化：第一行的累计值取交易金额和0的较大值
    SELECT transaction_date, transaction_amount, row_num,
           GREATEST(transaction_amount, 0) AS cumulative_value
    FROM temp_transactions
    WHERE row_num = 1
    
    UNION ALL
    
    -- 递归：每一行的累计值 = 上一行累计值 + 当前金额，与0取较大值
    SELECT curr.transaction_date, curr.transaction_amount, curr.row_num,
           GREATEST(prev.cumulative_value + curr.transaction_amount, 0) AS cumulative_value
    FROM temp_transactions curr
    JOIN cumulative_cte prev ON curr.row_num = prev.row_num + 1
)
SELECT transaction_date, transaction_amount, cumulative_value
FROM cumulative_cte
ORDER BY row_num
""")

# 查看结果
result_df.show()

输出结果 #

+----------------+------------------+---------------+
|transaction_date|transaction_amount|cumulative_value|
+----------------+------------------+---------------+
|      2019-01-30|              -200|              0|
|      2019-06-30|              -100|              0|
|      2020-01-30|               100|            100|
|      2020-06-30|                50|            150|
|      2021-01-30|               300|            450|
|      2021-01-30|              -500|              0|
|      2021-07-30|               800|            800|
|      2021-09-15|              -300|            500|
+----------------+------------------+---------------+

扩展说明 #

如果需要按分组（如用户ID）计算，只需在递归CTE中加入分组键，并调整行号窗口为分区窗口，修改递归条件为同分组内的行号递增即可。

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