Polars带条件窗口函数问题：分组计算符合条件的actions极值

场景与需求 #

给定如下DataFrame：

import polars as pl

df = pl.DataFrame(
    {'id':['1']*3 + ['2']*3,
     'actions': [0,1,2, 0, 1, 2],
     'campaign_name': [None,None, '2', '1', '2', None] ,
     'event_time': [1,2,3,0,1,2],
     'session_id':['session_1']*6}
)

数据展示：

┌─────┬─────────┬───────────────┬────────────┬────────────┐
│ id  ┆ actions ┆ campaign_name ┆ event_time ┆ session_id │
│ --- ┆ ---     ┆ ---           ┆ ---        ┆ ---        │
│ str ┆ i64     ┆ str           ┆ i64        ┆ str        │
╞═════╪═════════╪═══════════════╪════════════╪════════════╡
│ 1   ┆ 0       ┆ null          ┆ 1          ┆ session_1  │
│ 1   ┆ 1       ┆ null          ┆ 2          ┆ session_1  │
│ 1   ┆ 2       ┆ 2             ┆ 3          ┆ session_1  │
│ 2   ┆ 0       ┆ 1             ┆ 0          ┆ session_1  │
│ 2   ┆ 1       ┆ 2             ┆ 1          ┆ session_1  │
│ 2   ┆ 2       ┆ null          ┆ 2          ┆ session_1  │
└─────┴─────────┴───────────────┴────────────┴────────────┘

需求：按session_id和id的组合分组，仅针对campaign_name属于['1','2']的行，计算actions的最小值（min_action）和最大值（max_action），并将结果填充至分组内所有行。

原代码问题 #

原尝试代码如下：

df.with_columns(
    pl.when(pl.col('campaign_name').is_in(['1','2']))
    .then(
        pl.col('actions').min().over('session_id', 'id').alias('min_action')
    ),
    pl.when(pl.col('campaign_name').is_in(['1','2']))
    .then(
        pl.col('actions').max().over('session_id', 'id').alias('max_action')
    )  
)

存在两个核心问题：

1. 窗口函数min()/max()未筛选行，计算的是分组内所有行的actions极值，而非仅符合campaign_name条件的行；
2. pl.when仅在符合条件的行赋值，其他行留空，无法实现全组填充的需求。

修正代码 #

通过在窗口函数内部先筛选目标行，再计算极值，即可实现需求：

df.with_columns(
    # 计算分组内符合条件的actions最小值，填充全组
    pl.col('actions')
      .filter(pl.col('campaign_name').is_in(['1','2']))
      .min()
      .over('session_id', 'id')
      .alias('min_action'),
    # 计算分组内符合条件的actions最大值，填充全组
    pl.col('actions')
      .filter(pl.col('campaign_name').is_in(['1','2']))
      .max()
      .over('session_id', 'id')
      .alias('max_action')
)

结果验证 #

执行后得到预期结果：

shape: (6, 7)
┌─────┬─────────┬───────────────┬────────────┬────────────┬────────────┬────────────┐
│ id  ┆ actions ┆ campaign_name ┆ event_time ┆ session_id ┆ min_action ┆ max_action │
│ --- ┆ ---     ┆ ---           ┆ ---        ┆ ---        ┆ ---        ┆ ---        │
│ str ┆ i64     ┆ str           ┆ i64        ┆ str        ┆ i64        ┆ i64        │
╞═════╪═════════╪═══════════════╪════════════╪════════════╪════════════╪════════════╡
│ 1   ┆ 0       ┆ null          ┆ 1          ┆ session_1  ┆ 2          ┆ 2          │
│ 1   ┆ 1       ┆ null          ┆ 2          ┆ session_1  ┆ 2          ┆ 2          │
│ 1   ┆ 2       ┆ 2             ┆ 3          ┆ session_1  ┆ 2          ┆ 2          │
│ 2   ┆ 0       ┆ 1             ┆ 0          ┆ session_1  ┆ 0          ┆ 1          │
│ 2   ┆ 1       ┆ 2             ┆ 1          ┆ session_1  ┆ 0          ┆ 1          │
│ 2   ┆ 2       ┆ null          ┆ 2          ┆ session_1  ┆ 0          ┆ 1          │
└─────┴─────────┴───────────────┴────────────┴────────────┴────────────┴────────────┘

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