基于非零列的按行聚合：Polars稀疏数据集处理

核心思路 #

利用Polars的行方向聚合和动态列选择特性，无需分组即可完成行内求和，同时通过前缀动态生成哑变量，自动适配新增的前缀列。

完整实现代码 #

import polars as pl

df = pl.DataFrame(
    {
        "some_id": ["x", "y", "z"],
        "a_qty": [1, 0, 0],
        "a_height": [2, 0, 0],
        "a_width": [3, 0, 0],
        "b_qty": [0, 3, 0],
        "b_height": [0, 5, 0],
        "b_width": [0, 8, 0],
        "c_qty": [0, 0, 10],
        "c_height": [0, 0, 11],
        "c_width": [0, 0, 12],
    }
)

# 动态提取所有前缀（自动适配新增前缀）
prefixes = sorted({col.split("_")[0] for col in df.columns if "_" in col})

# 构建聚合操作列表
agg_operations = [
    # 行内求和生成qty/height/width列
    pl.col("*_qty").sum(axis=1).alias("qty"),
    pl.col("*_height").sum(axis=1).alias("height"),
    pl.col("*_width").sum(axis=1).alias("width"),
    # 生成每个前缀的哑变量（非零则为1，否则为0）
    *[pl.col(f"{p}_qty").ne(0).cast(pl.Int64).alias(f"is_prod_{p}") for p in prefixes]
]

# 执行聚合并保留some_id列
result = df.select(["some_id", *agg_operations])

print(result)

代码说明 #

1. 动态提取前缀：通过列名拆分提取所有前缀（如a、b、c），新增前缀时无需修改代码。
2. 行内求和：使用sum(axis=1)直接对所有同后缀列做行内求和，替代手动列选择+分组的冗余操作。
3. 哑变量生成：利用ne(0)判断对应前缀的列是否非零（因每行仅单个前缀有值），再转为整数类型得到哑变量。

解决你之前的报错问题 #

你之前的代码中width_cols = [col for col in df_test.columns if "width" in col],多了一个逗号，导致变量变成元组类型，解包时引发重复列名错误。而上述方案完全无需手动列选择，从根源避免了这类问题。

输出结果 #

┌─────────┬─────┬────────┬───────┬──────────┬──────────┬──────────┐
│ some_id ┆ qty ┆ height ┆ width ┆ is_prod_a┆ is_prod_b┆ is_prod_c│
│ ---     ┆ --- ┆ ---    ┆ ---   ┆ ---      ┆ ---      ┆ ---      │
│ str     ┆ i64 ┆ i64    ┆ i64   ┆ i64      ┆ i64      ┆ i64      │
╞═════════╪═════╪════════╪═══════╪══════════╪══════════╪══════════╡
│ x       ┆ 1   ┆ 2      ┆ 3     ┆ 1        ┆ 0        ┆ 0        │
│ y       ┆ 3   ┆ 5      ┆ 8     ┆ 0        ┆ 1        ┆ 0        │
│ z       ┆ 10  ┆ 11     ┆ 12    ┆ 0        ┆ 0        ┆ 1        │
└─────────┴─────┴────────┴───────┴──────────┴──────────┴──────────┘

内容的提问来源于stack exchange，提问作者anddt