Polars：如何最优地将数值列截断至指定分位数区间？

问题背景 #

我想要基于DataFrame的表达式结果截断数值，但Polars的clip函数仅接受浮点数或整数，不支持传入表达式。给定如下DataFrame：

df = pl.DataFrame({'x': [0, 1,2,3,4,5,6,7,8,9,10]})

我尝试用内置的clip函数将值截断到20%到80%分位数之间：

df.with_column(
    pl.col("x").clip(
        min_val = pl.col("x").quantile(0.20),
        max_val = pl.col("x").quantile(0.80)
   )
   .alias("clipped")
)

执行后报错：

thread '<unnamed>' panicked at 'called `Result::unwrap()` on an `Err` value: PyErr { type: <class 'RuntimeError'>, value: RuntimeError('BindingsError: "row type not supported <polars.internals.expr.expr.Expr object at 0x0000016F4B3053C0>"'), traceback: None }', src\lazy\dsl.rs:351:53
Traceback (most recent call last):
  File "C:\Users\BWT\Anaconda3\envs\tca_ml\lib\site-packages\IPython\core\interactiveshell.py", line 3398, in run_code
    exec(code_obj, self.user_global_ns, self.user_ns)
  File "<ipython-input-19-240f333898af>", line 2, in <cell line: 1>
    pl.col("x").clip(
  File "C:\Users\BWT\Anaconda3\envs\tca_ml\lib\site-packages\polars\internals\expr\expr.py", line 4840, in clip
    return wrap_expr(self._pyexpr.clip(min_val, max_val))
pyo3_runtime.PanicException: called `Result::unwrap()` on an `Err` value: PyErr { type: <class 'RuntimeError'>, value: RuntimeError('BindingsError: "row type not supported <polars.internals.expr.expr.Expr object at 0x0000016F4B3053C0>"'), traceback: None }

我找到了一种可行但过于冗长的实现方式：

>>> lower = pl.col("x").quantile(0.20)
>>> upper = pl.col("x").quantile(0.80)
>>> df.with_columns(
    [
        pl.when(pl.col("x") < lower)
        .then(lower)
        .when(pl.col("x") > upper)
        .then(upper)
        .otherwise(pl.col("x"))
        .alias("clipped")
    ]
)

输出结果符合预期：

Out[31]: 
shape: (11, 2)
┌─────┬─────────┐
│ x   ┆ clipped │
│ --- ┆ ---     │
│ i64 ┆ f64     │
╞═════╪═════════╡
│ 0   ┆ 2.0     │
├╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┤
│ 1   ┆ 2.0     │
├╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┤
│ 2   ┆ 2.0     │
├╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┤
│ 3   ┆ 3.0     │
├╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┤
│ ... ┆ ...     │
├╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┤
│ 7   ┆ 7.0     │
├╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┤
│ 8   ┆ 8.0     │
├╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┤
│ 9   ┆ 8.0     │
├╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┤
│ 10  ┆ 8.0     │
└─────┴─────────┘

请问有没有更简洁的实现方式？

简洁实现方案 #

有两种更简洁的方式可以实现需求：

1. 利用max()+min()链式调用
   通过先取列值与下分位数的最大值，再取该结果与上分位数的最小值，等价于截断操作，代码更紧凑：

df.with_columns(
    pl.col("x")
    .max(pl.col("x").quantile(0.2))
    .min(pl.col("x").quantile(0.8))
    .alias("clipped")
)

2. 升级Polars版本后直接使用clip表达式参数
   如果你的Polars版本在0.17.0及以上，clip函数已经支持传入表达式作为min_val和max_val，可以直接复用最初的思路：

df.with_columns(
    pl.col("x")
    .clip(min_val=pl.col("x").quantile(0.2), max_val=pl.col("x").quantile(0.8))
    .alias("clipped")
)

两种方式都能得到符合预期的截断结果，代码量远少于when-then的写法。

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