Polars：基于多列计算滚动众数（Rolling Mode）

现有如下Polars DataFrame：

import polars as pl

df = pl.from_repr("""
┌──────┬──────┬──────┬──────┬──────┐
│ col1 ┆ col2 ┆ col3 ┆ col4 ┆ col5 │
│ ---  ┆ ---  ┆ ---  ┆ ---  ┆ ---  │
│ i64  ┆ i64  ┆ i64  ┆ i64  ┆ f32  │
╞══════╪══════╪══════╪══════╪══════╡
│ 3    ┆ 3    ┆ 3    ┆ null ┆ null │
│ 2    ┆ 4    ┆ 1    ┆ 5    ┆ null │
│ 4    ┆ null ┆ null ┆ null ┆ null │
│ 7    ┆ 1    ┆ null ┆ null ┆ null │
│ 1    ┆ null ┆ null ┆ null ┆ null │
│ 10   ┆ 1    ┆ null ┆ null ┆ null │
│ 7    ┆ 9    ┆ 4    ┆ null ┆ null │
└──────┴──────┴──────┴──────┴──────┘
""")

需要创建一个新列res，存储基于窗口内所有列行值计算的滚动众数（窗口大小为2，最小周期1，非中心窗口），要求：

• 忽略原始数据中的null值
• 结果中不能将null作为众数
• 若无众数（所有值都不同）则返回None

预期结果如下：

┌──────┐
│ res  │
│ ---  │
│ i64  │
╞══════╡
│ 3    │
│ 3    │
│ 4    │
│ None │ # <- all values different
│ 1    │
│ 1    │
│ None │ # <- all values different
└──────┘

通过polars.rolling_apply结合自定义函数实现，具体步骤如下：

1. 定义众数计算函数：先过滤null值，统计每个值的出现次数，若存在唯一的最高频值则返回该值，否则返回None；
2. 利用索引窗口获取对应行的全列数据，传入自定义函数完成滚动计算。

完整代码：

import polars as pl
from collections import Counter

def calculate_mode(window_data):
    # 展平窗口内的多行多列数据，过滤null值
    flattened = [x for row in window_data for x in row if x is not None]
    if not flattened:
        return None
    # 统计值的出现频次
    count = Counter(flattened)
    max_count = max(count.values())
    # 筛选出所有频次等于最大值的候选众数
    modes = [k for k, v in count.items() if v == max_count]
    # 仅当存在唯一众数时返回，否则返回None
    return modes[0] if len(modes) == 1 else None

df = pl.from_repr("""
┌──────┬──────┬──────┬──────┬──────┐
│ col1 ┆ col2 ┆ col3 ┆ col4 ┆ col5 │
│ ---  ┆ ---  ┆ ---  ┆ ---  ┆ ---  │
│ i64  ┆ i64  ┆ i64  ┆ i64  ┆ f32  │
╞══════╪══════╪══════╪══════╪══════╡
│ 3    ┆ 3    ┆ 3    ┆ null ┆ null │
│ 2    ┆ 4    ┆ 1    ┆ 5    ┆ null │
│ 4    ┆ null ┆ null ┆ null ┆ null │
│ 7    ┆ 1    ┆ null ┆ null ┆ null │
│ 1    ┆ null ┆ null ┆ null ┆ null │
│ 10   ┆ 1    ┆ null ┆ null ┆ null │
│ 7    ┆ 9    ┆ 4    ┆ null ┆ null │
└──────┴──────┴──────┴──────┴──────┘
""")

# 新增滚动众数列
result_df = df.with_columns(
    pl.int_range(0, pl.count()).rolling_apply(
        function=lambda idx: calculate_mode(df[idx, :].to_numpy()),
        window_size=2,
        min_periods=1,
        center=False
    ).alias("res")
)

# 查看结果
print(result_df.select("res"))

运行后输出的res列与预期完全一致。

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