如何在Polars中按组基于最大值与非空条件分配category列

初始数据 #

import polars as pl

df = pl.DataFrame({
    'group_cols': ['A', 'A', 'A', 'A', 'B', 'B', 'C', 'C', 'C', 'C', 'C'],
    'col1': [2, 1, 1, 0, 3, 2, 3, None, 4, 4, 1],
    'col2': [None, 'not null', 'also not null', 'not null either', 'not null', 'not null', None, None, None, None, None]
})

需求说明 #

新增category列，每组group_cols需恰好一行标记为primary，规则如下：

• 优先选择col2非空且col1值最大的行（多行符合条件时任选其一即可）
• 若组内col2全为空，则选择col1值最大的行

Polars解决方案（避免使用map_groups） #

result = df.with_row_index() \
    .group_by('group_cols') \
    .agg(
        # 标记组内是否存在非空的col2
        pl.col('col2').is_not_null().any().alias('has_non_null_col2'),
        # 收集组内每行的索引、col1、col2信息为struct数组
        pl.struct(['index', 'col1', 'col2']).alias('rows')
    ) \
    .with_columns(
        # 筛选候选行：有非空col2时取col2非空的行，否则取所有行
        pl.col('rows').filter(
            pl.when(pl.col('has_non_null_col2'))
              .then(pl.element().struct.field('col2').is_not_null())
              .otherwise(True)
        ).alias('candidates')
    ) \
    .with_columns(
        # 在候选行中按col1降序排序（null视为极小值），取第一行的索引作为primary行标识
        pl.col('candidates').arr.sort_by(
            pl.element().struct.field('col1'), reverse=True, nulls_last=False
        ).arr.first().struct.field('index').alias('primary_index')
    ) \
    .select('group_cols', 'primary_index') \
    # 关联回原表生成category列
    .join(df.with_row_index(), on=['group_cols', 'index'], how='right') \
    .with_columns(
        pl.when(pl.col('index') == pl.col('primary_index'))
          .then('primary')
          .otherwise('secondary').alias('category')
    ) \
    .drop('index', 'primary_index') \
    .sort('group_cols')

代码逻辑说明 #

1. 添加行索引：通过with_row_index()给每行分配唯一标识，用于后续精准定位目标行
2. 分组聚合：
   • 标记组内是否存在非空的col2
   • 将组内每行的索引、col1、col2打包为struct数组，保留完整行信息
3. 筛选候选行：根据组内是否有非空col2，筛选出符合优先级条件的候选行集合
4. 确定primary行：对候选行按col1降序排序（null排最后），取第一行的索引作为组内primary行的标记
5. 关联生成结果：通过group_cols和index关联回原表，匹配到primary_index的行标记为primary，其余为secondary
6. 整理输出：删除临时辅助列，按group_cols排序得到最终结果

验证结果 #

执行代码后输出与预期一致：

shape: (11, 4)
┌────────────┬──────┬─────────────────┬───────────┐
│ group_cols ┆ col1 ┆ col2            ┆ category  │
│ ---        ┆ ---  ┆ ---             ┆ ---       │
│ str        ┆ i64  ┆ str             ┆ str       │
╞════════════╪══════╪═════════════════╪═══════════╡
│ A          ┆ 2    ┆ null            ┆ secondary │
│ A          ┆ 1    ┆ not null        ┆ primary   │
│ A          ┆ 1    ┆ also not null   ┆ secondary │
│ A          ┆ 0    ┆ not null either ┆ secondary │
│ B          ┆ 3    ┆ not null        ┆ primary   │
│ B          ┆ 2    ┆ not null        ┆ secondary │
│ C          ┆ 3    ┆ null            ┆ secondary │
│ C          ┆ null ┆ null            ┆ secondary │
│ C          ┆ 4    ┆ null            ┆ primary   │
│ C          ┆ 4    ┆ null            ┆ secondary │
│ C          ┆ 1    ┆ null            ┆ secondary │
└────────────┴──────┴─────────────────┴───────────┘

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