如何在Python Polars的.group_by().map_groups()结果拼接中忽略异常结果？

问题 #

我有一个需要多步运算得到最终结果的DataFrame，打算用.group_by().map_groups()封装逻辑来简化代码（不介意Python带来的性能损失，只求代码简洁）。但部分组因为不满足特定条件（比如样本不足）没有结果，这不是错误。现在遇到的技术问题：

• UDF返回的DataFrame Schema和输入不同
• 返回None会触发报错，返回符合Schema的空DataFrame又太繁琐（因为Schema还在开发中可能随时变化）
  想知道有没有办法让.group_by().map_groups()直接跳过这类无结果的组，不用处理拼接问题？

示例代码 #

测试数据集

import polars as pl

df = pl.DataFrame(
    {
        "id": [0, 1, 2, 3, 4],
        "color": ["blue", "green", "green", "red", "red"],
        "shape": ["square", "triangle", "square", "triangle", "square"],
    }
)
print(df)

输出：

shape: (5, 3)
┌─────┬───────┬──────────┐
│ id  ┆ color ┆ shape    │
│ --- ┆ ---   ┆ ---      │
│ i64 ┆ str   ┆ str      │
╞═════╪═══════╪══════════╡
│ 0   ┆ blue  ┆ square   │
│ 1   ┆ green ┆ triangle │
│ 2   ┆ green ┆ square   │
│ 3   ┆ red   ┆ triangle │
│ 4   ┆ red   ┆ square   │
└─────┴───────┴──────────┘

原UDF（触发报错）

def udf(df):
    result = df.filter(pl.col('id') < 3)
    # red组无符合条件的元素，返回None
    if result.shape[0] < 1:
        return None
    # 输出Schema与输入不同
    return result.select('id', 'color')

错误信息

df.group_by('color').map_groups(udf)

报错：

PanicException: Could not get DataFrame attribute '_df'. 
Make sure that you return a DataFrame object.: PyErr { type: <class 'AttributeError'>, value: AttributeError("'NoneType' object has no attribute '_df'"), traceback: None }

期望结果

shape: (3, 2)
┌─────┬───────┐
│ id  ┆ color │
│ --- ┆ ---   │
│ i64 ┆ str   │
╞═════╪═══════╡
│ 1   ┆ green │
│ 2   ┆ green │
│ 0   ┆ blue  │
└─────┴───────┘

解决方案 #

直接返回与正常输出Schema一致的空DataFrame，就能让map_groups自动跳过该组的拼接，同时避免手动维护Schema的麻烦。修改后的UDF如下：

def udf(df):
    # 直接生成最终格式的结果
    result = df.filter(pl.col('id') < 3).select('id', 'color')
    # 结果为空时，返回同Schema的空DataFrame
    return result if result.height > 0 else pl.DataFrame(schema=result.schema)

调用方式不变：

df.group_by('color').map_groups(udf)

核心逻辑说明 #

• Polars要求map_groups的每个返回结果必须是Schema统一的DataFrame，空DataFrame只要Schema匹配就不会被加入最终结果
• 通过result.schema动态获取输出结构，无需手动定义，完美适配开发期Schema变化的场景
• 代码简洁，同时满足跳过无结果组的需求

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