如何获取Polars中分区上最频繁的值？

Question

我正在尝试使用 Polars 获取 DataFrame 列中每个类别中最常见的值，并且该类别由多个其他列给出（即使用复合主键）。

例如，我有一个看起来像这样的 DataFrame：

输入

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    │
└───────┴─────┴─────┴───────┘
""")

我希望输出是：

输出

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             │
└───────┴─────┴─────┴───────┴────────────────┘

我尝试将

.mode()

与

.over()

一起使用，但它返回异常：

计算错误：窗口表达式的长度与组的长度不匹配

我尝试使用的代码：

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'])
)

Answer 1

事实证明，如果存在平局（多个模式），Polars 的

.mode()

表达式可以返回多个值。

所以我们必须添加

.first()

def get_majority_value(primary_keys: list):
    return (
        pl.col("value")
        .drop_nulls() # Use this if you don't want to consider "None" as a possible mode
        .mode()
        .first()
        .over(primary_keys)
        .alias('majority_value')
    )

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