如何在Polars中按复合分组获取列的众数？

问题场景 #

按复合主键（pk1、pk2）分组，为每组计算value列的众数并新增majority_value列，但使用.mode().over()窗口函数时触发如下错误：

ComputeError: the length of the window expression did not match that of the group

输入示例DataFrame #

import polars as pl

df = pl.from_repr("""
┌───────┬─────┬─────┬───────┐
│ index ┆ pk1 ┆ pk2 ┆ value │
│ ---   ┆ --- ┆ --- ┆ ---   │
│ i64   ┆ str ┆ str ┆ i64   │
╞═══════╪═════╪═════╪═══════╡
│ 0     ┆ a   ┆ x   ┆ 42    │
│ 1     ┆ b   ┆ y   ┆ 69    │
│ 2     ┆ a   ┆ x   ┆ 36    │
│ 3     ┆ b   ┆ x   ┆ 12    │
│ 4     ┆ a   ┆ x   ┆ 36    │
└───────┴─────┴─────┴───────┘
""")

出错代码 #

def get_majority_value(primary_keys: list):
    return (
        pl.col("value")
        .mode()
        .over(primary_keys)
        .alias('majority_value')
    )

df.with_columns(
    get_majority_value(primary_keys=['pk1', 'pk2'])
)

期望输出 #

shape: (5, 5)
┌───────┬─────┬─────┬───────┬────────────────┐
│ index ┆ pk1 ┆ pk2 ┆ value ┆ majority_value │
│ ---   ┆ --- ┆ --- ┆ ---   ┆ ---            │
│ i64   ┆ str ┆ str ┆ i64   ┆ i64            │
╞═══════╪═════╪═════╪═══════╪════════════════╡
│ 0     ┆ a   ┆ x   ┆ 42    ┆ 36             │
│ 1     ┆ b   ┆ y   ┆ 69    ┆ 69             │
│ 2     ┆ a   ┆ x   ┆ 36    ┆ 36             │
│ 3     ┆ b   ┆ x   ┆ 12    ┆ 12             │
│ 4     ┆ a   ┆ x   ┆ 36    ┆ 36             │
└───────┴─────┴─────┴───────┴────────────────┘

错误原因 #

Polars的mode()方法返回的是数组类型（即使分组内只有一个众数，也会返回单元素数组），而窗口函数.over()要求每个分组返回的结果长度必须与组内行数完全匹配。数组类型的结果无法直接广播到组内所有行，因此触发长度不匹配的错误。

解决方案 #

方案1：分组聚合后关联回原表 #

先通过group_by计算每组的众数（取第一个众数处理多众数场景），再通过复合主键关联回原表：

# 分组计算每组的众数
mode_df = df.group_by(['pk1', 'pk2']).agg(
    pl.col('value').mode().first().alias('majority_value')
)

# 关联回原表
result_df = df.join(mode_df, on=['pk1', 'pk2'])
print(result_df)

方案2：窗口函数中提取数组元素 #

在窗口函数内，对mode()返回的数组调用.list.first()提取第一个元素，将数组转为标量后广播到组内所有行：

def get_majority_value(primary_keys: list):
    return (
        pl.col("value")
        .mode()
        .list.first()  # 提取众数数组的第一个元素
        .over(primary_keys)
        .alias('majority_value')
    )

result_df = df.with_columns(get_majority_value(primary_keys=['pk1', 'pk2']))
print(result_df)

两种方案均可得到符合预期的输出。如果分组存在多个众数，可根据业务需求调整（比如取最大/最小众数，或保留数组）。

内容的提问来源于stack exchange，提问作者Vinícius Queiroz