Polars 在所有内容上为 group_by 添加额外的行

Question

import polars as pl

df_test = pl.DataFrame({
    'item_num': [1, 2, 1, 2, 3, 1, 2, 3, 1, 2],
    'rate': [100, 200, 150, 250, 300, 120, 220, 350, 180, 270],
    'consider': [1, 1, 0, 1, 1, 0, 1, 1, 0, 1]
})

filter_for_agg = pl.all_horizontal(
    pl.col('consider') == 1
)

groups = (
    df_test.group_by('item_num')
    .agg(
        pl.col("rate").filter(filter_for_agg).mean().alias("mean"),
        pl.col("rate").filter(filter_for_agg).median().alias("median"),  
        pl.col("rate").filter(filter_for_agg).quantile(.75).alias("p75"),          
    )
)

上面是我的代码片段，我需要组数据框中的另一行，它将读取 item_num 中的“all”，并且整个数据集上的所有 kpi 都满足“filter_for_agg”。最好的方法是什么？

预期输出：

shape: (4, 4)
┌──────────┬────────────┬────────┬───────┐
│ item_num ┆ mean       ┆ median ┆ p75   │
│ ---      ┆ ---        ┆ ---    ┆ ---   │
│ str      ┆ f64        ┆ f64    ┆ f64   │
╞══════════╪════════════╪════════╪═══════╡
│ 2        ┆ 235.0      ┆ 235.0  ┆ 250.0 │
│ 1        ┆ 100.0      ┆ 100.0  ┆ 100.0 │
│ 3        ┆ 325.0      ┆ 325.0  ┆ 350.0 │
│ all      ┆ 241.428571 ┆ 250.0  ┆ 300.0 │
└──────────┴────────────┴────────┴───────┘

Answer 1

你可以做

total = df_test.group_by(pl.lit("all")).agg(
    pl.col("rate").filter(filter_for_agg).mean().alias("mean"),
    pl.col("rate").filter(filter_for_agg).median().alias("median"),
    pl.col("rate").filter(filter_for_agg).quantile(0.75).alias("p75"),
)
pl.concat(
    [
        groups.with_columns(pl.col("item_num").cast(pl.Utf8)),
        total.rename({"literal": "item_num"}),
    ]
)

这给出了

shape: (4, 4)
┌──────────┬────────────┬────────┬───────┐
│ item_num ┆ mean       ┆ median ┆ p75   │
│ ---      ┆ ---        ┆ ---    ┆ ---   │
│ str      ┆ f64        ┆ f64    ┆ f64   │
╞══════════╪════════════╪════════╪═══════╡
│ 3        ┆ 325.0      ┆ 325.0  ┆ 350.0 │
│ 2        ┆ 235.0      ┆ 235.0  ┆ 270.0 │
│ 1        ┆ 100.0      ┆ 100.0  ┆ 100.0 │
│ all      ┆ 241.428571 ┆ 250.0  ┆ 300.0 │
└──────────┴────────────┴────────┴───────┘