如何用Polars实现滚动累计：上月期末余额转下月期初

问题描述 #

模拟企业初始资金的滚动累计：初始资金1000美元，每月生成随机盈亏，计算一段时间后的资金情况。用Polars创建带日期列的DataFrame后，尝试通过with_columns()堆叠生成期初/期末余额时，无法实现上月期末转下月期初的滚动累计，各月期初始终为初始值1000。循环遍历每行的方法虽能得到正确结果，但无法利用Polars的优化特性，效率低下，需用Polars原生功能实现需求。

尝试的无效代码 #

import polars as pl
import datetime as dt
from dateutil.relativedelta import relativedelta
from random import normalvariate

start_date = dt.date.today() + relativedelta(months=1, day=1)
df = pl.DataFrame(
    pl.date_range(start_date, start_date + relativedelta(months=5), '1mo', eager=True).alias('date'),
)
beginning_balance = 1000.0
df = df.with_columns(
    pl.lit(beginning_balance).alias('beginning_balance'),
    pl.lit(beginning_balance).alias('closing_balance'),
).with_columns(
    pl.when(pl.col('date') == start_date)
    .then(pl.col('beginning_balance'))
    .otherwise(pl.col('closing_balance').shift(1))
    .alias('beginning_balance'), 
    pl.Series([normalvariate(100, 80) for _ in range(len(df))]).round(2).alias('profit'),
    pl.Series([normalvariate(100, 75) for _ in range(len(df))]).round(2).alias('loss'),
).with_columns(
    (pl.col('beginning_balance') + pl.col('profit') - pl.col('loss')).alias('closing_balance'),
)
df

无效代码执行结果 #

shape: (6, 5)
┌────────────┬───────────────────┬─────────────────┬────────┬────────┐
│ date       ┆ beginning_balance ┆ closing_balance ┆ profit ┆ loss   │
│ ---        ┆ ---               ┆ ---             ┆ ---    ┆ ---    │
│ date       ┆ f64               ┆ f64             ┆ f64    ┆ f64    │
╞════════════╪═══════════════════╪═════════════════╪════════╪════════╡
│ 2024-11-01 ┆ 1000.0            ┆ 934.95          ┆ -58.53 ┆ 6.52   │
│ 2024-12-01 ┆ 1000.0            ┆ 903.15          ┆ 69.02  ┆ 165.87 │
│ 2025-01-01 ┆ 1000.0            ┆ 1007.71         ┆ 111.21 ┆ 103.5  │
│ 2025-02-01 ┆ 1000.0            ┆ 1011.97         ┆ 209.43 ┆ 197.46 │
│ 2025-03-01 ┆ 1000.0            ┆ 998.85          ┆ 32.22  ┆ 33.37  │
│ 2025-04-01 ┆ 1000.0            ┆ 1151.32         ┆ 277.49 ┆ 126.17 │
└────────────┴───────────────────┴─────────────────┴────────┴────────┘

可行但低效的循环方法 #

current_tally = beginning_balance
for t in range(len(df)):
    beginning_balance = current_tally
    current_tally = beginning_balance + df[t, 'profit'] - df[t, 'loss']
    df[t, 'beginning_balance'] = beginning_balance
    df[t, 'closing_balance'] = current_tally

df

循环方法正确结果 #

┌────────────┬───────────────────┬─────────────────┬────────┬────────┐
│ date       ┆ beginning_balance ┆ closing_balance ┆ profit ┆ loss   │
│ ---        ┆ ---               ┆ ---             ┆ ---    ┆ ---    │
│ date       ┆ f64               ┆ f64             ┆ f64    ┆ f64    │
╞════════════╪═══════════════════╪═════════════════╪════════╪════════╡
│ 2024-11-01 ┆ 1000.0            ┆ 934.95          ┆ -58.53 ┆ 6.52   │
│ 2024-12-01 ┆ 934.95            ┆ 838.1           ┆ 69.02  ┆ 165.87 │
│ 2025-01-01 ┆ 838.1             ┆ 845.81          ┆ 111.21 ┆ 103.5  │
│ 2025-02-01 ┆ 845.81            ┆ 857.78          ┆ 209.43 ┆ 197.46 │
│ 2025-03-01 ┆ 857.78            ┆ 856.63          ┆ 32.22  ┆ 33.37  │
│ 2025-04-01 ┆ 856.63            ┆ 1007.95         ┆ 277.49 ┆ 126.17 │
└────────────┴───────────────────┴─────────────────┴────────┴────────┘

Polars原生实现方案 #

原代码失效的原因是：with_columns()基于当前DataFrame的列计算，无法引用同一批次中新生成列的更新值。通过计算每月净盈亏的累加和，结合初始余额推导期末、期初余额，可实现高效的滚动累计。

正确代码 #

import polars as pl
import datetime as dt
from dateutil.relativedelta import relativedelta
from random import normalvariate

start_date = dt.date.today() + relativedelta(months=1, day=1)
df = pl.DataFrame(
    pl.date_range(start_date, start_date + relativedelta(months=5), '1mo', eager=True).alias('date'),
)
beginning_balance = 1000.0

# 生成盈亏数据并计算每月净盈亏
df = df.with_columns(
    pl.Series([normalvariate(100, 80) for _ in range(len(df))]).round(2).alias('profit'),
    pl.Series([normalvariate(100, 75) for _ in range(len(df))]).round(2).alias('loss'),
    (pl.col('profit') - pl.col('loss')).alias('net_profit'),
)

# 计算累计净盈亏，得到每月期末余额
df = df.with_columns(
    (beginning_balance + pl.col('net_profit').cumsum()).alias('closing_balance'),
)

# 推导期初余额：首月用初始值，后续取上月期末余额
df = df.with_columns(
    pl.col('closing_balance').shift(1).fill_null(beginning_balance).alias('beginning_balance'),
)

# 调整列顺序（可选）
df = df.select(['date', 'beginning_balance', 'closing_balance', 'profit', 'loss'])
print(df)

代码解释 #

1. 计算净盈亏：先得到每月profit - loss的净变化量，为后续累计做准备。
2. 累计净盈亏：用cumsum()实现向量化累加，加上初始余额直接得到每月期末余额，这一步完全利用Polars的优化特性，性能远高于循环。
3. 推导期初余额：通过shift(1)将上月期末余额下移，首月用fill_null()填充初始值，得到正确的期初余额。

此代码输出结果与循环方法完全一致，且适合处理大规模数据集。

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