如何在Polars中高效实现分组条件式自定义函数应用？

给定如下Polars DataFrame：

import polars as pl

df = pl.DataFrame(
    {
        "a": [1, 4, 3, 2, 8, 4, 5, 6],
        "b": [2, 3, 1, 3, 9, 7, 6, 8],
        "c": [1, 1, 1, 1, 2, 2, 2, 2],
    }
)

需要完成以下任务：

1. 按列c分组
2. 对每个分组执行逻辑：
   • 若分组内所有a列的值均小于对应b列的值，直接保留原a列数据
   • 否则，调用第三方Python函数convert（接收两个numpy数组，返回同尺寸numpy数组），将分组的a、b转为numpy数组传入，再将返回结果转为Polars Series返回

已知convert函数示例：

import numpy as np
def convert(a_arr: np.ndarray, b_arr: np.ndarray) -> np.ndarray:
    # 示例逻辑，实际为第三方实现
    return np.where(a_arr > b_arr, b_arr, a_arr)

# 测试：c=1分组的a=[1,4,3,2]，b=[2,3,1,3]，调用后返回[1,3,1,2]

期望输出（展开分组后）：

shape: (8, 2)
┌─────┬─────┐
│ c   ┆ a   │
│ --- ┆ --- │
│ i64 ┆ i64 │
╞═════╪═════╡
│ 1   ┆ 1   │
│ 1   ┆ 3   │
│ 1   ┆ 1   │
│ 1   ┆ 2   │
│ 2   ┆ 8   │
│ 2   ┆ 4   │
│ 2   ┆ 5   │
│ 2   ┆ 6   │
└─────┴─────┘

核心诉求：在无需调用第三方函数的分组中充分利用Polars的Rust原生性能，仅在必要时调用Python自定义函数，避免全局apply带来的性能损耗

可以通过「先标记需要处理的分组，再拆分处理后合并」的方式实现，兼顾性能与需求：

步骤1：标记需要调用convert的分组 #

先通过Polars的原生聚合函数快速判断每个分组是否满足all(a < b)，生成分组标记：

# 生成分组判断结果：True表示需要调用convert，False表示直接保留a
group_flags = df.group_by("c").agg(
    need_convert=~pl.all(pl.col("a") < pl.col("b"))
)

步骤2：拆分DataFrame为两个子集 #

• 无需处理的分组：直接保留c和a列
• 需要处理的分组：调用convert函数处理后保留结果

# 拆分数据集
df_no_convert = df.join(group_flags.filter(~pl.col("need_convert")), on="c").select("c", "a")

df_need_convert = df.join(group_flags.filter(pl.col("need_convert")), on="c")

# 对需要处理的分组应用convert函数
def process_group(group: pl.DataFrame) -> pl.DataFrame:
    a_np = group["a"].to_numpy()
    b_np = group["b"].to_numpy()
    converted_a = convert(a_np, b_np)
    return pl.DataFrame({
        "c": group["c"],
        "a": converted_a
    })

df_converted = df_need_convert.group_by("c", maintain_order=True).apply(process_group)

步骤3：合并两个子集得到最终结果 #

final_df = pl.concat([df_no_convert, df_converted], maintain_order=True)
print(final_df)

方案优势 #

• 性能优化：无需处理的分组完全使用Polars原生操作，享受Rust的并行化和速度优势，仅对需要处理的分组调用Python函数，将性能损耗降到最低
• 逻辑清晰：拆分处理的方式避免了在分组apply中做分支判断，符合Polars的惯用写法
• 可扩展性：后续若有其他分支逻辑，可继续拆分对应子集处理

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