如何在Polars中按条件统计当前行上方的行数？

问题描述 #

现有如下已按date列排序的Polars DataFrame：

df = pl.DataFrame(
    {
    'date': ['2022-01-01', '2022-01-02', '2022-01-07', '2022-01-17', '2022-03-02', '2022-06-05', '2022-06-07', '2022-07-02'],
    'col1': [4, 4, 2, 2, 2, 3, 2, 1],
    'col2': [1, 2, 3, 4, 1, 3, 3, 4],
    'col3': [2, 3, 4, 4, 3, 2, 2, 1]
    }
)

需要新增一列ge，统计所有早于当前行（日期更早）且col1、col2、col3的值均大于等于当前行对应列值的行数，即满足：

Count rows
    where row_index < current_row_index &
          col1[row_index] >= col1[current_row_index] &
          col2[row_index] >= col2[current_row_index] &
          col3[row_index] >= col3[current_row_index]

预期结果：

解决方案 #

方法一：结合NumPy的高效实现（适合大数据集） #

利用广播比较和三角掩码快速计算符合条件的行数：

import numpy as np
import polars as pl

# 提取列数据为NumPy数组
col1 = df['col1'].to_numpy()
col2 = df['col2'].to_numpy()
col3 = df['col3'].to_numpy()

# 生成三角掩码：仅保留当前行之前的行（row_index < current_row_index）
mask = np.triu(np.ones((len(df), len(df)), dtype=bool), k=1).T

# 计算三列均满足>=当前行的条件矩阵
cond_matrix = (col1[:, None] >= col1) & (col2[:, None] >= col2) & (col3[:, None] >= col3)

# 结合掩码统计每行符合条件的行数
ge_counts = (cond_matrix & mask).sum(axis=1)

# 将结果添加到原DataFrame
result_df = df.with_columns(pl.Series('ge', ge_counts))

方法二：纯Polars API实现（语法更贴合Polars） #

通过int_range生成索引，逐行切片统计历史行：

import polars as pl

result_df = df.with_columns(
    pl.int_range(0, pl.count()).map(
        lambda idx: (
            df.slice(0, idx)
            .select(
                (pl.col('col1') >= df['col1'][idx]) &
                (pl.col('col2') >= df['col2'][idx]) &
                (pl.col('col3') >= df['col3'][idx])
            )
            .sum()
        )
    ).alias('ge')
)

两种方法对比 #

• 方法一借助NumPy的广播机制，计算效率更高，适合处理十万级以上的大数据集；
• 方法二完全使用Polars原生API，逻辑直观易懂，适合小数据集或偏好纯Polars语法的场景。

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