本文介绍使用 polars 对 dataframe 每行中数值最大的指定数量（如前 3）列进行求和，并同步生成二进制标识矩阵，精确指出哪些原始列参与了该行的 top-k 聚合。
本文介绍使用 polars 对 dataframe 每行中数值最大的指定数量（如前 3）列进行求和，并同步生成二进制标识矩阵，精确指出哪些原始列参与了该行的 top-k 聚合。
在数据处理中，常需对宽格式（wide-format）数值列执行“行内 top-k 聚合”——例如：从 val1 至 val6 六列中取每行最大的 3 个值求和，并记录哪几列被选中（用于可解释性或后续逻辑）。Polars 原生不支持直接的行级 top_k 操作，但可通过 unpivot → group_by → top_k_by + top_k 的组合策略高效实现，无需 Python 循环或 apply，全程保持惰性计算与高性能。
✅ 核心步骤解析
重塑结构（Unpivot）
将所有目标数值列（如val*）转为长格式：新增variable（列名）和value（对应值）两列，保留id_col作为分组键。-
按 ID 分组聚合
- 使用
pl.col("value").top_k(3).sum()计算每组（即每行原始数据）最大 3 个值的和； - 同时用
pl.col("variable").top_k_by("value", k=3)获取对应列名列表（自动处理并列：按原始顺序稳定截断）。
- 使用
重建宽格式标识列
利用pl.lit(col).is_in("variable")对每个原始列名col进行逐行判断：若其出现在该行 top-3 列名列表中，则返回True，再经.cast(pl.Int8)转为0/1整型标志位。
? 完整实现代码
import polars as pl
import polars.selectors as cs
df = pl.DataFrame({
"id_col": [0, 1, 2, 3, 4],
"val1": [10, 0, 0, 20, 5],
"val2": [5, 1, 2, 3, 10],
"val3": [8, 2, 2, 2, 5],
"val4": [1, 7, 7, 4, 1],
"val5": [3, 0, 0, 6, 0],
"val6": [2, 7, 5, 5, 4],
})
# 动态获取所有 val* 列名
val_cols = cs.expand_selector(df, cs.starts_with("val"))
result = (
df.unpivot(
index="id_col",
on=val_cols, # 显式指定列更清晰（等价于 cs.starts_with("val")）
)
.group_by("id_col")
.agg(
pl.col("value").top_k(3).sum().alias("agg_col"),
pl.col("variable").top_k_by("value", k=3).alias("top_vars"),
)
.select(
"id_col",
# 为每个 val* 列生成 0/1 标志列
*[pl.lit(col).is_in("top_vars").cast(pl.Int8).alias(col) for col in val_cols],
"agg_col",
)
.sort("id_col") # 确保输出顺序与输入一致
)
print(result)
⚠️ 注意事项与最佳实践
-
并列处理：当第 3 大值存在多个时（如第 2 行中
val2==val3==val6==2），top_k_by默认按 输入行序 选取（即最早出现的列优先），结果确定且可复现；如需其他策略（如字典序最小列名），可先对variable排序再取 top-k。 -
性能提示：
unpivot+group_by是 Polars 高效向量化操作，远优于row_wise().apply(...)；对于超宽表（数百列），建议预过滤无关列以减少内存占用。 -
类型安全：
.cast(pl.Int8)显式指定小整型，节省内存；若需布尔型，可改用.cast(pl.Boolean)。 -
扩展性：只需修改
k=3和cs.starts_with("val")即可适配任意列前缀与 top-k 数量。
该方案兼顾表达力、性能与可维护性，是 Polars 生态中处理“行内 top-k 聚合+溯源”任务的标准范式。
