如何在Polars DataFrame中高效实现跨行计算？

在使用Polars DataFrame时，需要基于其他行的值为每行计算两个新列：

• sum_lower：所有小于当前元素的数值之和
• max_other：排除当前元素后，DataFrame中的最大值

当前使用map_elements逐行处理的方式效率低下，代码如下：

import polars as pl

COL_VALUE = "value"

def fun_sum_lower(current_row, df):
    tmp_df = df.filter(pl.col(COL_VALUE) < current_row[COL_VALUE])
    sum_lower = tmp_df.select(pl.sum(COL_VALUE)).item()
    return sum_lower

def fun_max_other(current_row, df):
    tmp_df = df.filter(pl.col(COL_VALUE) != current_row[COL_VALUE])
    max_other = tmp_df.select(pl.col(COL_VALUE)).max().item()
    return max_other

if __name__ == '__main__':
    df = pl.DataFrame({COL_VALUE: [3, 7, 1, 9, 4]})

    df = df.with_columns(
        pl.struct([COL_VALUE])
        .map_elements(lambda row: fun_sum_lower(row, df), return_dtype=pl.Int64)
        .alias("sum_lower")
    )

    df = df.with_columns(
        pl.struct([COL_VALUE])
        .map_elements(lambda row: fun_max_other(row, df), return_dtype=pl.Int64)
        .alias("max_other")
    )

    print(df)

运行后输出：

shape: (5, 3)
┌───────┬───────────┬───────────┐
│ value ┆ sum_lower ┆ max_other │
│ ---   ┆ ---       ┆ ---       │
│ i64   ┆ i64       ┆ i64       │
╞═══════╪═══════════╪═══════════╡
│ 3     ┆ 1         ┆ 9         │
│ 7     ┆ 8         ┆ 9         │
│ 1     ┆ 0         ┆ 9         │
│ 9     ┆ 15        ┆ 7         │
│ 4     ┆ 4         ┆ 9         │
└───────┴───────────┴───────────┘

需要找到不需要lambda、行迭代或Python代码的Polars原生高效实现方式。

利用Polars的窗口函数、排序和聚合函数可以完全实现原生高效计算，避免逐行操作：

import polars as pl

COL_VALUE = "value"

if __name__ == '__main__':
    df = pl.DataFrame({COL_VALUE: [3, 7, 1, 9, 4]})
    
    # 计算sum_lower：先排序计算前缀和，再映射回原数据
    sorted_df = df.sort(COL_VALUE).with_columns(
        pl.col(COL_VALUE).cum_sum().shift().fill_null(0).alias("sum_lower")
    )
    
    # 计算max_other：分两种情况，当前值是全局最大值时取第二大，否则取全局最大值
    global_max = df.select(pl.max(COL_VALUE)).item()
    second_max = df.filter(pl.col(COL_VALUE) != global_max).select(pl.max(COL_VALUE)).item()
    
    df = df.join(sorted_df, on=COL_VALUE, how="left").with_columns(
        pl.when(pl.col(COL_VALUE) == global_max)
        .then(second_max)
        .otherwise(global_max)
        .alias("max_other")
    )
    
    print(df)

代码解释 #

1. sum_lower的计算：

   • 先对DataFrame按value排序，确保所有小于当前值的元素都在当前行之前
   • 使用cum_sum()计算前缀和，通过shift()将前缀和对应到下一行（即当前值的所有更小值之和），最后用fill_null(0)处理最小值的情况（没有更小值时和为0）
   • 通过join将计算好的sum_lower映射回原DataFrame的原始顺序

2. max_other的计算：

   • 先获取全局最大值global_max和排除最大值后的第二大值second_max
   • 使用when/then/otherwise条件判断：如果当前值是全局最大值，max_other取第二大值；否则直接取全局最大值

这种方式完全基于Polars的向量化操作，效率远高于逐行迭代的实现，且输出结果与原代码一致。

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