Python 基于非零列的行式聚合

我面临的场景是，我有一个稀疏的数据集(对于每一行，同时只填充4列，其余的将是零).对于每个变量前缀，将有3个变量(例如a_qty、a_height、ma_width)，列数完全是任意的，因为随时可能添加新的前缀.

举个例子，数据帧可能类似于:

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],
    }
)
# prints
┌─────────┬───────┬──────────┬─────────┬───┬─────────┬───────┬──────────┬─────────┐
│ some_id ┆ a_qty ┆ a_height ┆ a_width ┆ … ┆ b_width ┆ c_qty ┆ c_height ┆ c_width │
│ ---     ┆ ---   ┆ ---      ┆ ---     ┆   ┆ ---     ┆ ---   ┆ ---      ┆ ---     │
│ str     ┆ i64   ┆ i64      ┆ i64     ┆   ┆ i64     ┆ i64   ┆ i64      ┆ i64     │
╞═════════╪═══════╪══════════╪═════════╪═══╪═════════╪═══════╪══════════╪═════════╡
│ x       ┆ 1     ┆ 2        ┆ 3       ┆ … ┆ 0       ┆ 0     ┆ 0        ┆ 0       │
│ y       ┆ 0     ┆ 0        ┆ 0       ┆ … ┆ 8       ┆ 0     ┆ 0        ┆ 0       │
│ z       ┆ 0     ┆ 0        ┆ 0       ┆ … ┆ 0       ┆ 10    ┆ 11       ┆ 12      │
└─────────┴───────┴──────────┴─────────┴───┴─────────┴───────┴──────────┴─────────┘

我正在寻找一种方法来缩小数据集的范围，方法是减少列数并go 掉给定行的0.目标数据帧的每一行应包含:

1. qty:给定行中包含qty的列的总和(同样适用于width和height)
2. is_prod_*根据非零列的前缀将伪变量设置为1.前缀不能混合(每行只有一个字母)

df_target = pl.DataFrame(
    {
        "some_id": ["x", "y", "z"],
        "height": [2, 5, 11],
        "width": [3, 8, 12],
        "qty": [1, 3, 10],
        "is_prod_a": [1, 0, 0],
        "is_prod_b": [0, 1, 0],
        "is_prod_c": [0, 0, 1]
    }
)
# prints
┌─────────┬────────┬───────┬─────┬────────┬────────┬────────┐
│ some_id ┆ height ┆ width ┆ qty ┆ prod_a ┆ prod_b ┆ prod_c │
│ ---     ┆ ---    ┆ ---   ┆ --- ┆ ---    ┆ ---    ┆ ---    │
│ str     ┆ i64    ┆ i64   ┆ i64 ┆ i64    ┆ i64    ┆ i64    │
╞═════════╪════════╪═══════╪═════╪════════╪════════╪════════╡
│ x       ┆ 2      ┆ 3     ┆ 1   ┆ 1      ┆ 0      ┆ 0      │
│ y       ┆ 5      ┆ 8     ┆ 3   ┆ 0      ┆ 1      ┆ 0      │
│ z       ┆ 11     ┆ 12    ┆ 10  ┆ 0      ┆ 0      ┆ 1      │
└─────────┴────────┴───────┴─────┴────────┴────────┴────────┘

我try 过按行号分组和按类型求和变量:

width_cols = [col for col in df_test.columns if "width" in col],
height_cols = [col for col in df_test.columns if "height" in col]
qty_cols = [col for col in df_test.columns if "qty" in col]

(
    df
    .with_row_count()
    .group_by("row_nr")
    .agg([
        pl.col(*width_cols).sum().alias('width')
    ])
)

# DuplicateError: column with name 'width' has more than one occurrences

这显然不起作用，因为我试图在第二行创建一个重复的列.实现这种聚合最优雅的方式是什么？

编辑1

使用pl.DataFrame.fold可以得到逐行的列和，这似乎可以做到这一点.仍然不确定如何根据条件创建名为base的伪列.

(
    df_test
    .with_columns([
        pl.fold(0, lambda acc, s: acc + s, pl.col(*width_cols)).alias("width"),
        pl.fold(0, lambda acc, s: acc + s, pl.col(*height_cols)).alias("height"),
        pl.fold(0, lambda acc, s: acc + s, pl.col(*qty_cols)).alias("qty")
    ])
)