本文介绍如何在 Polars 中高效实现“按分组（如 A 列）查找某列（x）极值所在行，进而提取另一列（B）的值”，并生成新列 y；核心方法是组合使用 arg_max()、get() 和 over()。
本文介绍如何在 polars 中高效实现“按分组（如 a 列）查找某列（x）极值所在行，进而提取另一列（b）的值”，并生成新列 y；核心方法是组合使用 `arg_max()`、`get()` 和 `over()`。
在 Polars 数据处理中，常需基于分组聚合逻辑提取“极值位置对应的其他字段值”——例如：按列 A 分组后，在每组内找到 x 最大值所在的行，并将该行的 B 值广播到整组作为新列 y。这不同于常规聚合（如 max('x')），而是典型的“索引对齐式取值”，需借助位置索引能力。
关键在于理解三个表达式的作用与协同：
-
pl.col("x").arg_max()：返回当前分组内x列最大值的行索引（0-based），类型为UInt32； -
pl.col("B").get(...)：根据给定索引从B列中提取标量值（支持单个索引或表达式）； -
.over("A")：将上述组合操作以窗口函数形式按A分组独立执行，确保每组内独立计算索引并取值。
完整代码如下：
import polars as pl
df = pl.DataFrame({
'A': ['a0', 'a0', 'a1', 'a1'],
'B': ['b1', 'b2', 'b1', 'b2'],
'x': [0, 10, 5, 1]
})
result = df.with_columns(
pl.col("B").get(pl.col("x").arg_max()).over("A").alias("y")
)
print(result)
输出结果符合预期：
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 │ └─────┴─────┴─────┴─────┘
✅ 注意事项：
- 若
x在组内存在多个相同最大值，arg_max()返回第一个匹配位置（稳定行为）； - 若需获取最小值对应列，将
arg_max()替换为arg_min()即可； -
get()要求索引有效（非空组），若存在空组或全 null 的x，建议先用filter或fill_null()预处理； - 该方案纯表达式驱动，无 Python 循环或 UDF，性能优异，适用于大规模数据。
此模式可泛化至任意“按组取极值行中指定列”的场景，是 Polars 窗口计算与位置索引能力的典型实践。
