将Pandas基于时间窗口的GroupBy/Apply/EWM计算迁移至Polars

出于性能考量，需将部分Pandas脚本转为Polars实现，核心需求是按用户分组，基于日期时间计算30天半衰期的EWM均值。参考相关内容后写出近似实现，但存在两个明显缺陷：

• 用设置极长时间窗口的rolling()模拟Pandas的expanding()，写法冗余笨拙
• 以30行数据作为半衰期，而非基于date列的30天时间窗口，导致计算结果与Pandas版本存在差异

原始代码实现 #

import pandas as pd
import polars as pl
import random
from datetime import datetime, timedelta

# 定义用户列表
persons = ['Person A', 'Person B', 'Person C', 'Person D']

# 生成随机测试数据
df = pd.DataFrame(
    {'person': [random.choice(persons) for _ in range(50)],
     'rating': [random.randint(75, 110) for _ in range(50)],
     'date'  : [datetime(2022, 6, 1, 0, 0, 0)
                + timedelta(days=random.randint(0, 365))
                for _ in range(50)]}
)

df.sort_values(['date'], inplace=True)
dl = pl.from_dataframe(df)

# Pandas 实现逻辑
df['EWM_30d'] = df.groupby(
    by='person', sort=False).apply(
    lambda x: x['rating'].ewm(halflife=('30d'), times=x['date']
                              ).mean().shift(1, fill_value=80).round(2)
).to_numpy()

# 初始Polars近似实现
dl = dl.rolling(
    'date', by='person', period="100000d").agg(
    pl.col('rating').ewm_mean(half_life=30).shift(1, fill_value=80).last().alias('EWM_30d'))

优化后的Polars实现 #

针对上述缺陷，优化后的实现完全匹配Pandas的时间半衰期逻辑，同时写法更简洁：

import polars as pl
import pandas as pd
import random
from datetime import datetime, timedelta

# 生成数据部分同原始代码，此处省略
# ...

# 先按用户和日期排序，确保分组内时间顺序正确
dl = dl.sort(['person', 'date'])

# 按用户分组计算基于30天半衰期的EWM均值
result = dl.group_by('person').agg(
    pl.col('date'),
    pl.col('rating'),
    # 计算当前行与上一行的天数差，第一行填充0
    pl.col('date').diff().dt.total_days().fill_null(0).alias('days_since_last'),
).with_columns(
    # 计算30天半衰期对应的衰减系数：alpha = 1 - exp(-ln(2)/30)
    pl.lit(1 - (-pl.math.ln(2)/30).exp()).alias('alpha'),
    # 计算每个数据点的权重：基于时间差的指数衰减，累积乘积得到相对权重
    pl.col('days_since_last').map_batches(
        lambda x: (-pl.math.ln(2)/30 * x).exp().cum_prod()
    ).alias('weights'),
).with_columns(
    # 计算加权和与权重总和
    (pl.col('rating') * pl.col('weights')).cum_sum().alias('weighted_sum'),
    pl.col('weights').cum_sum().alias('weight_sum'),
    # 计算EWM均值，shift(1)后用80填充空值，保留两位小数
    (pl.col('weighted_sum') / pl.col('weight_sum')).shift(1).fill_null(80).round(2).alias('EWM_30d'),
).select(['person', 'date', 'rating', 'EWM_30d'])

优化说明 #

1. 替代expanding()的优雅方案：通过group_by+agg结合cum_sum/cum_prod实现累积计算，完全模拟Pandasexpanding()的效果，无需设置超大时间窗口
2. 匹配时间半衰期逻辑：通过计算日期差，基于指数衰减公式生成权重，完全对齐Pandasewm(halflife='30d', times=x['date'])的计算逻辑，确保结果一致

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