Polars自定义表达式构造函数

Question

我想在Polars中更改前面n行的某些值。一般情况下，这个问题可以在之前的帖子中找到解决方案。

不过，我想通过构建一个自定义函数来实现这个目标：

from polars.type_aliases import IntoExpr, IntoExprColumn
import polars as pl
import numpy as np

df = pl.DataFrame({"test": np.arange(1, 11)})

def _func(x: IntoExpr) -> pl.Expr:
    return pl.when((x+1) < 5).then(None).otherwise(x+1)

df.with_columns(
    _func(pl.col("test")).alias("test+1")
    )

1. 我该如何使用自定义函数创建索引列？（不把Polars DataFrame作为输入参数）
2. 有没有办法在不传入DataFrame作为输入参数的情况下，访问用于Polars表达式的Polars DataFrame？

Answer 1

看起来我们不能直接从 Polars Expressions 访问 Polars.DataFrame 对象，必须把它们作为输入参数传进去。

不过，我们可以使用以下查询来：

1. 把第 n 个值替换成前 n 个值的平均值
2. 把第 1-n 个值（包括左边的那个）设置为 Null
3. 其他情况下保持原值不变

def query(
    target_var: IntoExpr = pl.col("var"),
    index_col: IntoExpr | pl.Expr | None = None,
    n: int = 1,
    **kwargs,
    ) -> pl.Expr:
    
    if "index_col" in kwargs:
        index_col = kwargs.pop("index_col") # IntoExpr
    else:
        index_col = pl.int_range(0, pl.len())

    mean_nth = pl.head(target_var.meta.output_name(), n).mean()
    expr = (
        pl.when(index_col == (n-1))
            .then(mean_nth)
            .when(index_col < (n-1))
            .then(None)
            .otherwise(target_var)
            )
    
    return expr