如何在Rust Polars中按分组对数据集执行demean操作？

在Rust Polars中，完全可以实现类似Python Pandas里groupby.transform(lambda x: x - x.mean())的链式调用逻辑，核心是利用Polars的**窗口函数（over）**结合列表达式运算，无需单独计算分组均值再做减法。

优化后的链式调用实现 #

直接在Lazy API中通过with_columns或select批量处理目标列，全程链式完成分组去均值：

use ndarray::{Array, s};
use ndarray_rand::RandomExt;
use ndarray_rand::rand_distr::Uniform;
use polars::prelude::*;
use random_choice::random_choice;

fn main() -> Result<(), PolarsError> {
    // 生成1000行10列的随机均匀分布数据
    let X = Array::random((1000, 10), Uniform::new(0, 100));

    // 生成随机分组
    let samples = vec!["group 1", "group 2", "group 3"];
    let weights = vec![1.0, 1.0, 1.0];
    let groups: Vec<&&str> = random_choice().random_choice_f64(&samples, &weights, 1000);
    let groups_clone: Vec<&str> = groups.into_iter().copied().collect();

    // 构建初始DataFrame
    let X_df = df!(
        "shock_1" => X.slice(s![.., 0]).to_vec(),
        "shock_2" => X.slice(s![.., 1]).to_vec(),
        "shock_3" => X.slice(s![.., 2]).to_vec(),
        "groups" => &groups_clone
    )?;

    // 链式调用实现分组去均值：保留原分组列，对其他列做去均值处理
    let demeaned_df = X_df.lazy()
        .with_columns(
            // 排除分组列，对剩余所有列执行：列值 - 分组均值
            all().exclude("groups")
                .map(|col| col - col.mean().over(["groups"]), GetOutput::from_type(DataType::Float64))
        )
        .collect()?;

    println!("{:?}", demeaned_df);
    Ok(())
}

关键说明 #

1. 窗口函数over的作用：col.mean().over(["groups"])会计算每个分组内该列的均值，并将结果广播到对应分组的每一行，完全等价于Pandas的transform逻辑。
2. 批量处理列：all().exclude("groups")可以一次性选中所有需要去均值的列，避免逐个列编写重复表达式。
3. Lazy API的优势：全程使用Lazy模式可以让Polars优化查询计划，比先计算均值再做减法的拆分方式更高效。
4. 错误处理：用Result<(), PolarsError>替代expect，符合Rust的错误处理规范，避免程序panic。

另一种实现：直接生成去均值后的数据集 #

如果不需要保留原始列，也可以用select直接构建只包含分组列和去均值列的DataFrame：

let demeaned_df = X_df.lazy()
    .select([
        "groups",
        all().exclude("groups").map(|col| col - col.mean().over(["groups"]), GetOutput::from_type(DataType::Float64))
    ])
    .collect()?;

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