Polars中基于前序变量的向量化计算c值的最优方案

问题 #

现有三个变量a、b、c，各有n个实例，其中c除第一个实例外均为null。需要根据公式计算后续每个c的值：

c = [(1 + a) * (current_c) * (b)] + [(1 + b) * (current_c) * (a)]

要求不使用原生Python循环。

已尝试的方法均未成功：

• 使用pl.int_range(my_index_column_value, pl.len() + 1)（索引从1开始）
• 使用pl.rolling(...)但开销过高
• 结合上述方法与.over(...)、pl.select(...).item()的pl.when(...).then(...).otherwise(...)

核心问题：shift操作会一次性完成，无法实现依赖前一行计算结果的递推逻辑；尝试用索引列值作为偏移量时，Polars无法正确识别该值为数字，即使转换类型也报错。目前考虑新增行用于偏移传递值，但不确定效率。

参考代码（存在逻辑问题）：

import polars as pl

if __name__ == "__main__":
    initial_c_value = 3

    df = pl.DataFrame(((2, 3, 4, 5, 8), (3, 7, 4, 9, 2)), schema=('a', 'b'))
    df = df.with_row_index('i', 1).with_columns(pl.lit(None).alias('c'))

    df = df.with_columns(pl.when(pl.col('i') == 1)
    .then(
        (((1 + pl.col('a')) * (initial_c_value) * (pl.col('b'))) +
        ((1 + pl.col('b')) * (initial_c_value) * (pl.col('a')))).alias('c'))
    .otherwise(
        ((1 + pl.col('a')) * (pl.col('c').shift(1)) * (pl.col('b'))) +
        ((1 + pl.col('b')) * (pl.col('c').shift(1)) * (pl.col('a')))).shift(1).alias('c'))

    print(df)

最优实现方案 #

在Polars中，这类递推计算可以通过pl.fold()函数高效实现，它能在列上执行累积式的递推操作，完全避免Python循环，且性能远优于滚动窗口。

步骤说明 #

1. 先简化递推公式：原公式可合并同类项减少计算量

   c_next = c_prev * [(1+a)*b + (1+b)*a] = c_prev * (a + b + 2ab)
   

2. 使用pl.fold()从初始值开始，逐行累积计算c的值
3. 无需额外索引列，直接基于a、b列生成递推结果

完整代码 #

import polars as pl

if __name__ == "__main__":
    initial_c_value = 3

    df = pl.DataFrame(
        ((2, 3, 4, 5, 8), (3, 7, 4, 9, 2)),
        schema=('a', 'b')
    )

    # 计算递推系数：a + b + 2*a*b
    df = df.with_columns(
        coeff = pl.col('a') + pl.col('b') + 2 * pl.col('a') * pl.col('b')
    )

    # 使用fold执行递推计算
    df = df.with_columns(
        c = pl.fold(
            acc=pl.lit(initial_c_value),
            function=lambda acc, x: acc * x,
            exprs=pl.col('coeff')
        )
    )

    print(df)

代码解释 #

• pl.fold()的acc参数是初始累积值（对应第一个c的计算结果：代入原公式得(1+2)*3*3 + (1+3)*3*2 = 51，而coeff第一行值为17，3*17=51，与原公式结果一致）
• function定义每一步递推逻辑：用前一步累积值乘以当前行的系数
• exprs指定用于递推的列（预先计算的coeff）

这种实现完全利用Polars向量化计算能力，性能远优于滚动窗口或模拟偏移的方法，且代码简洁易读。

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