如何在Polars的groupby apply中传递参数？

Polars的groupby.apply() API设计和Pandas不同，它的function参数仅接受单个可调用对象（该对象默认接收分组后的DataFrame作为第一个参数），无法像Pandas那样直接展开参数列表传递额外参数。你可以通过以下两种方式解决：

方法1：用lambda函数包装传递参数 #

通过匿名lambda函数接收分组DataFrame，再调用你的自定义函数并传入参数：

import polars as pl

# 修正自定义函数，适配Polars的返回格式（生成Result列）
class ApplyFunction:
    def Close(df, *args):
        return pl.DataFrame({'Result': df['Close'] * args[0]})

# 沿用你定义的Placeholder和call对象
class Placeholder:
    def __init__(self):
        self.ind = ApplyFunction.Close
        self.args = [2]
call = Placeholder()

# 转换为Polars DataFrame
df = pl.from_pandas(df.reset_index(drop=False))

# 使用lambda包装传参
df = df.with_columns(
    df.groupby("Symbol", maintain_order=True).apply(
        function=lambda group_df: call.ind(group_df, *call.args),
    )
)

print(df)

运行结果：

shape: (4, 4)
┌────────┬─────────────────────┬───────┬────────┐
│ Symbol ┆ Date                ┆ Close ┆ Result │
│ ---    ┆ ---                 ┆ ---   ┆ ---    │
│ str    ┆ datetime[ns]        ┆ i64   ┆ i64    │
╞════════╪═════════════════════╪═══════╪════════╡
│ AAA    ┆ 2020-01-01 00:00:00 ┆ 8     ┆ 16     │
│ AAA    ┆ 2020-01-02 00:00:00 ┆ 9     ┆ 18     │
│ BBB    ┆ 2020-01-01 00:00:00 ┆ 10    ┆ 20     │
│ BBB    ┆ 2020-01-02 00:00:00 ┆ 11    ┆ 22     │
└────────┴─────────────────────┴───────┴────────┘

方法2：用functools.partial绑定参数 #

通过functools.partial预先将参数绑定到自定义函数上，生成一个只需要接收分组DataFrame的新可调用对象：

import polars as pl
import functools

# 沿用之前的ApplyFunction、Placeholder和call对象
call = Placeholder()

df = pl.from_pandas(df.reset_index(drop=False))

# 使用partial绑定参数
df = df.with_columns(
    df.groupby("Symbol", maintain_order=True).apply(
        function=functools.partial(call.ind, *call.args),
    )
)

print(df)

该方法的运行结果和方法1完全一致。

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