本文介绍使用 Polars 实现按行选取最大 K 个数值列求和，并同步生成二值化标识矩阵（标记哪些列被选入 Top-K）的完整方案，核心依赖 unpivot + top_k / top_k_by + 行级布尔广播。
本文介绍使用 polars 实现按行选取最大 k 个数值列求和，并同步生成二值化标识矩阵（标记哪些列被选入 top-k）的完整方案，核心依赖 `unpivot` + `top_k` / `top_k_by` + 行级布尔广播。
在数据处理中，常需对宽表中多个数值列进行行内 Top-K 聚合（如“每行取最大 3 列求和”），同时保留可解释性——即明确记录哪些原始列参与了聚合。Polars 原生不支持直接对行内多列做动态 Top-K 操作，但可通过「长表转换 → 分组聚合 → 逆向广播」三步高效实现，且全程惰性执行、零 Python 循环。
以下以示例数据为例，目标为：对 val1–val6 六列，每行取最大 3 个值求和，并生成对应 val1–val6 的 0/1 标识列（1 表示该列值进入 Top-3），最终输出含 id_col、6 个标识列和 agg_col 的结果 DataFrame。
✅ 核心步骤解析
-
宽转长（Unpivot）：将所有
val*列熔合为两列 ——variable（列名）和value（对应值），以id_col为索引键。这使后续按id_col分组后，可自然对每个 ID 的全部 val 值排序。 -
分组聚合（GroupBy + Top-K）：
-
pl.col("value").top_k(3).sum()：对每个id_col组内value取最大 3 个并求和； -
pl.col("variable").top_k_by("value", k=3)：同步获取对应的最大 3 个variable名（即列名），返回字符串列表（如["val1", "val3", "val2"]）。
-
-
重建宽表标识列：利用
pl.lit(col).is_in("variable")对每个原始val*列名字面量，判断其是否出现在上一步生成的 Top-3 列名列表中；再.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"))
res = (
df.unpivot(
index="id_col",
on=val_cols,
variable_name="variable",
value_name="value"
)
.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(res)
✅ 输出说明：
agg_col值与预期一致（如第 0 行：10+8+5=23）；标识列严格对应被选中的列（如val1,val3,val2为 1）。存在并列时（如第 2 行val2==val3==val6==2），top_k_by按内部稳定排序任意取 3 个，符合题目要求。
⚠️ 注意事项
-
性能提示：
unpivot+group_by是 Polars 推荐的行级 Top-K 模式，比apply或map_rows快一个数量级以上，且支持流式处理。 -
列名一致性：
cs.starts_with("val")使用 selector 自动捕获列，避免硬编码；若列名不规则，可显式传入["val1","val2",...]。 -
类型安全：
.cast(pl.Int8)显式指定小整型，节省内存；若需布尔类型，可改用.cast(pl.Boolean)。 -
扩展性：只需修改
k=3和val_cols即可适配任意列数与 Top-K 大小。
该方法兼具简洁性、可读性与高性能，是 Polars 生态中处理“行内 Top-K 聚合 + 解释性标识”任务的标准范式。
