如何在Polars表达式函数中整合多步骤数据处理逻辑？

需求说明 #

按分组列foo（实际场景为3列）分组，执行以下逻辑：

• 若组内bar列的首个值为空（空字符串''或NULL），则用组内baz列的首个值替换该组所有bar的值
• 若组内bar首个值非空，则保留原bar数据

目标是将核心逻辑封装为单个Polars表达式函数以支持链式调用，同时寻求更简洁的实现方式。

初始数据 #

import polars as pl

df = pl.DataFrame({'foo': [1, 1, 1, 2, 2, 2, 3, 3],
                   'bar': ['a', 'a', 'a', None, None, None, 'c', 'c'],
                   'baz': ['x', None, 'q', 'z', 'r', None, 'y', 's']})

初始数据结构：

shape: (8, 3)
┌─────┬──────┬──────┐
│ foo ┆ bar  ┆ baz  │
│ --- ┆ ---  ┆ ---  │
│ i64 ┆ str  ┆ str  │
╞═════╪══════╪══════╡
│ 1   ┆ a    ┆ x    │
│ 1   ┆ a    ┆ null │
│ 1   ┆ a    ┆ q    │
│ 2   ┆ null ┆ z    │
│ 2   ┆ null ┆ r    │
│ 2   ┆ null ┆ null │
│ 3   ┆ c    ┆ y    │
│ 3   ┆ c    ┆ s    │
└─────┴──────┴──────┘

现有实现代码 #

df = (df.with_columns(pl.col('baz').first().over(['foo']).alias('temp'))
        .with_columns(pl.when((pl.col('bar') == '') | (pl.col('bar').is_null()))
                      .then(pl.col('temp'))
                      .otherwise(pl.col('bar')).alias('bar2'))
        .with_columns(pl.col('bar2').alias('bar'))
        .drop(['temp', 'bar2'])
      )

封装为单个表达式函数 #

将所有逻辑整合到一个表达式中，无需临时列，直接返回可用于链式调用的表达式：

def update_bar(group_cols: list[str]) -> pl.Expr:
    # 获取组内bar的首个值
    first_bar = pl.col('bar').first().over(group_cols)
    # 获取组内baz的首个值
    first_baz = pl.col('baz').first().over(group_cols)
    # 判断替换逻辑并返回结果
    return pl.when((first_bar == '') | first_bar.is_null())
             .then(first_baz)
             .otherwise(pl.col('bar'))
             .alias('bar')

链式调用示例 #

result_df = df.with_columns(update_bar(group_cols=['foo']))

更简洁的实现方式 #

利用pl.coalesce简化空值判断逻辑，进一步压缩代码：

def update_bar(group_cols: list[str]) -> pl.Expr:
    first_bar = pl.col('bar').first().over(group_cols)
    first_baz = pl.col('baz').first().over(group_cols)
    return pl.when(pl.coalesce(first_bar == '', first_bar.is_null()))
             .then(first_baz)
             .otherwise(pl.col('bar'))
             .alias('bar')

若无需封装函数，也可直接使用单行表达式完成转换：

result_df = df.with_columns(
    pl.when((pl.col('bar').first().over(['foo']) == '') | pl.col('bar').first().over(['foo']).is_null())
    .then(pl.col('baz').first().over(['foo']))
    .otherwise(pl.col('bar'))
    .alias('bar')
)

验证结果 #

执行后得到预期输出：

shape: (8, 3)
┌─────┬──────┬──────┐
│ foo ┆ bar  ┆ baz  │
│ --- ┆ ---  ┆ ---  │
│ i64 ┆ str  ┆ str  │
╞═════╪══════╪══════╡
│ 1   ┆ a    ┆ x    │
│ 1   ┆ a    ┆ null │
│ 1   ┆ a    ┆ q    │
│ 2   ┆ z    ┆ z    │
│ 2   ┆ z    ┆ r    │
│ 2   ┆ z    ┆ null │
│ 3   ┆ c    ┆ y    │
│ 3   ┆ c    ┆ s    │
└─────┴──────┴──────┘

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