Polars中top_k_by结合over窗口且k=1时报错是否为Bug？

你遇到的这个情况其实不是Bug哦，是top_k_by函数的返回特性和窗口函数over的要求不匹配导致的，咱们来一步步拆解清楚：

首先先还原你的初始DataFrame：

import polars as pl

df = pl.DataFrame({
    'A': ['a0', 'a0', 'a1', 'a1'],
    'B': ['b1', 'b2', 'b1', 'b2'],
    'x': [0, 10, 5, 1]
})

你的需求是在每个A分组里，找到x值最大对应的B列内容，尝试用top_k_by(k=1)结合窗口函数实现时触发了ComputeError，提示窗口表达式长度和分组不匹配；但设置k=2时却能正常运行。

原因分析 #

• 当k=2时，top_k_by会返回一个长度为2的列表，Polars会自动把这个列表广播到当前分组的每一行（比如a0分组有2行，就把列表['b2', 'b1']分别赋值给这两行），所以不会报错。
• 但当k=1时，Polars做了优化，直接返回一个标量值（比如a0分组返回'b2'），而窗口函数over要求每个分组的输出结果长度必须和该分组的行数一致（a0分组有2行，就需要返回2个值），这就导致了长度不匹配的报错。

解决方法 #

如果一定要用top_k_by实现需求，只需要把结果强制转换成列表，让Polars可以正常广播到分组的每一行即可：

df.with_columns(
    pl.col.B.top_k_by("x", 1).list().over("A").alias("y")
)

运行后得到的结果：

shape: (4, 4)
┌─────┬─────┬─────┬───────┐
│ A   ┆ B   ┆ x   ┆ y     │
│ --- ┆ --- ┆ --- ┆ ---   │
│ str ┆ str ┆ i64 ┆ list[str] │
╞═════╪═════╪═════╪═══════╡
│ a0  ┆ b1  ┆ 0   ┆ ["b2"] │
│ a0  ┆ b2  ┆ 10  ┆ ["b2"] │
│ a1  ┆ b1  ┆ 5   ┆ ["b1"] │
│ a1  ┆ b2  ┆ 1   ┆ ["b1"] │
└─────┴─────┴─────┴───────┘

如果想要直接得到标量而不是列表，也可以用你提到的arg_max方案，这也是更常规的写法：

df.with_columns(
    pl.col.B[pl.col.x.arg_max()].over("A").alias("y")
)

结果如下：

shape: (4, 4)
┌─────┬─────┬─────┬─────┐
│ A   ┆ B   ┆ x   ┆ y   │
│ --- ┆ --- ┆ --- ┆ --- │
│ str ┆ str ┆ i64 ┆ str │
╞═════╪═════╪═════╪═════╡
│ a0  ┆ b1  ┆ 0   ┆ b2  │
│ a0  ┆ b2  ┆ 10  ┆ b2  │
│ a1  ┆ b1  ┆ 5   ┆ b1  │
│ a1  ┆ b2  ┆ 1   ┆ b1  │
└─────┴─────┴─────┴─────┘

备注：内容来源于stack exchange，提问作者roman