PySpark中如何计算当日与前一日的日初、日终余额？

需求说明 #

已知每个账户首日（D1）的初始余额，每日按当日交易金额扣减余额，需要按日期维度计算每个账户的：

• 当日日初余额
• 当日日终余额
• 前一日日终余额

输入数据样例：

期望输出结果：

核心计算逻辑 #

• 首日（每个账户日期排序后的第一条记录）日初余额 = 账户初始余额
• 当日日终余额 = 当日日初余额 - 当日交易扣减金额
• 非首日的日初余额 = 上一日的日终余额
• 前一日余额 = 同账户按日期排序后上一条记录的日终余额，首日该值为空
  实现时直接使用Spark窗口函数做累计计算，避免逐行迭代的性能损耗，适配大数据量场景

实现代码 #

1. 导入依赖 #

from pyspark.sql import SparkSession
from pyspark.sql.window import Window
import pyspark.sql.functions as F

2. 加载/构造输入数据 #

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

# 以下为测试数据构造，生产环境直接替换为读表/读文件的DataFrame即可
# 字段说明：账户ID、交易日期、当日交易扣减金额、账户首日初始余额
source_df = spark.createDataFrame(
    [
        ("A001", "2024-01-01", 10, 100),
        ("A001", "2024-01-02", 20, 100),
        ("A001", "2024-01-03", 15, 100),
        ("A002", "2024-01-01", 5, 200),
        ("A002", "2024-01-02", 25, 200)
    ],
    schema=["acct_id", "dt", "txn_deduct_amt", "init_balance"]
)

3. 窗口计算各余额字段 #

# 定义窗口规则：按账户分区，按交易日期升序排序
balance_window = Window.partitionBy("acct_id").orderBy("dt")

result_df = source_df.withColumn(
    # 计算截至前一日的累计扣减金额
    "cum_deduct_prev",
    F.sum("txn_deduct_amt").over(balance_window.rowsBetween(Window.unboundedPreceding, -1))
).withColumn(
    # 日初余额 = 初始余额 - 历史累计扣减，首日无历史扣减时取0
    "day_start_balance",
    F.col("init_balance") - F.coalesce(F.col("cum_deduct_prev"), F.lit(0))
).withColumn(
    # 日终余额 = 日初余额 - 当日扣减
    "day_end_balance",
    F.col("day_start_balance") - F.col("txn_deduct_amt")
).withColumn(
    # 取上一行的日终余额作为前一日余额
    "prev_day_end_balance",
    F.lag("day_end_balance", 1).over(balance_window)
).drop("cum_deduct_prev") # 删除中间计算字段

4. 结果查看 #

result_df.orderBy("acct_id", "dt").show(truncate=False)

执行后输出结果和期望样例完全对齐：

+-------+----------+---------------+--------------+-----------------+---------------+--------------------+
|acct_id|dt        |txn_deduct_amt |init_balance  |day_start_balance|day_end_balance|prev_day_end_balance|
+-------+----------+---------------+--------------+-----------------+---------------+--------------------+
|A001   |2024-01-01|10             |100           |100              |90             |null                |
|A001   |2024-01-02|20             |100           |90               |70             |90                  |
|A001   |2024-01-03|15             |100           |70               |55             |70                  |
|A002   |2024-01-01|5              |200           |200              |195            |null                |
|A002   |2024-01-02|25             |200           |195              |170            |195                 |
+-------+----------+---------------+--------------+-----------------+---------------+--------------------+

注意事项 #

• 如果账户存在无交易的日期断档，需要先按账户维度补齐连续日期序列，缺失日期的交易扣减额填0，再执行上述计算，避免余额计算断层
• 如果交易流水中同时包含入账、扣减，只需将txn_deduct_amt替换为带正负方向的交易发生额（入账为正、扣减为负），同时将余额计算的减号调整为加号即可适配
• 该实现全量使用Spark原生API，无UDF、无逐行循环，TB级数据量下也可稳定运行

内容的提问来源于stack exchange，提问作者Carlos Eduardo Bilar Rodrigues