如何在PySpark中为交易数据集创建CurrentBalance列？

我用PySpark创建了如下DataFrame：

+---------------+-------------+---------------+------+
|TransactionDate|AccountNumber|TransactionType|Amount|
+---------------+-------------+---------------+------+
|     2023-01-01|          100|         Credit|  1000|
|     2023-01-02|          100|         Credit|  1500|
|     2023-01-03|          100|          Debit|  1000|
|     2023-01-02|          200|         Credit|  3500|
|     2023-01-03|          200|          Debit|  2000|
|     2023-01-04|          200|         Credit|  3500|
|     2023-01-13|          300|         Credit|  4000|
|     2023-01-14|          300|          Debit|  4500|
|     2023-01-15|          300|         Credit|  5000|
+---------------+-------------+---------------+------+

需要为该DataFrame新增一列CurrentBalance，期望输出如下：

+---------------+-------------+---------------+------+--------------+
|TransactionDate|AccountNumber|TransactionType|Amount|CurrentBalance|
+---------------+-------------+---------------+------+--------------+
|     2023-01-01|          100|         Credit|  1000|          1000|
|     2023-01-02|          100|         Credit|  1500|          2500|
|     2023-01-03|          100|          Debit|  1000|          1500|
|     2023-01-02|          200|         Credit|  3500|          3500|
|     2023-01-03|          200|          Debit|  2000|          1500|
|     2023-01-04|          200|         Credit|  3500|          5000|
|     2023-01-13|          300|         Credit|  4000|          4000|
|     2023-01-14|          300|          Debit|  4500|          -500|
|     2023-01-15|          300|         Credit|  5000|          1000|
+---------------+-------------+---------------+------+--------------+

我曾尝试按AccountNumber分组查找TransactionDate的最小日期，并结合when条件计算收支，但未成功，尝试的代码如下：

# Find minimum date in TransactionDate column, grouped by AccountNumber column
df_new.groupBy('AccountNumber').agg(f.min('TransactionDate').alias('min_date'))

解决方案 #

要实现需求，需要利用窗口函数按账户分组并按交易日期排序，累计计算收支余额：

1. 导入所需函数：

from pyspark.sql import functions as f
from pyspark.sql.window import Window

2. 定义窗口规则：按账户分组，按交易日期升序排列

window_spec = Window.partitionBy("AccountNumber").orderBy("TransactionDate")

3. 调整金额正负（Credit为正，Debit为负），再计算累计余额：

df_result = df.withColumn(
    "adjusted_amount",
    f.when(f.col("TransactionType") == "Credit", f.col("Amount"))
    .when(f.col("TransactionType") == "Debit", -f.col("Amount"))
    .otherwise(0)
).withColumn(
    "CurrentBalance",
    f.sum("adjusted_amount").over(window_spec)
).drop("adjusted_amount")

4. 查看结果：

df_result.show()

说明 #

• partitionBy("AccountNumber") 确保每个账户的交易单独计算余额，不会和其他账户混淆
• orderBy("TransactionDate") 保证交易按时间顺序累加，符合实际账户余额的计算逻辑
• sum("adjusted_amount").over(window_spec) 计算从该账户第一笔交易到当前交易的累计金额，即当前余额

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