Rust Polars嵌套groupby时apply_multiple内创建DF聚合报错求解

问题场景 #

使用Rust版本Polars实现复杂聚合逻辑时，需在外层GroupBy聚合上下文的apply_multiple回调中，基于传入列数据新建DataFrame，复用DataFrame原生的filter、GroupBy、Aggregate能力完成计算（复现完整逻辑需开启ndarray特性）。
核心需求：

• 原始DataFrame先按Region列做稳定分组
• 每个Region分组内先过滤掉Scalar列为空的行，再按Month列二次分组
• 聚合计算WeightScaled（计算逻辑为col("Weight") * col("Scalar")）的月度总和：EU分组预期得到长度为2的序列[0.3, 0.7]（分别对应NOV、APR月度加权和），US分组预期得到[1.7, 2.4]（分别对应JUL、SEP月度加权和）
• 基于月度聚合结果调用ndarray完成复杂数学运算，最终每个Region分组返回一个f64类型标量作为区域健康指数

异常表现 #

原实现运行时存在两个问题：

• 对region_health_index()返回的表达式调用sum()等聚合函数会触发panic，报错信息为assertion failed: idx.len() <= self.len()
• 仅调用first()时返回空值，无法得到每个分组对应一个f64标量的预期结果

原问题复现代码如下：

use polars::prelude::*;
use polars::df;
use ndarray::prelude::*;
use ndarray::stack;

pub fn main() {
    let df = df! [
        "Region" => ["EU", "EU", "EU", "EU",                "US", "US", "US", "US"],
        "Month" => ["NOV", "DEC", "APR", "APR",             "JUL", "JAN", "JUL", "SEP"],
        "Weight" => [1, 2, 3, 4,                             5,6,7,8],
        "Scalar" => [Some(0.3), None, Some(0.1), Some(0.1),   Some(0.2), None, Some(0.1), Some(0.3)]
    ].unwrap();

    let df1 = df.clone().lazy()
        .groupby_stable([col("Region")])
        .agg( [
            // 调用sum()时报错，仅调用first()时返回空值
            region_health_index().first().alias("RegionHealth"),
            //region_health_index().sum().alias("RegionHealth2"),
        ]
        )
        .collect()
        .unwrap();
    dbg!(df1);
}

pub fn weight_scaled() -> Expr {
    col("Weight") * col("Scalar")
}

pub fn region_health_index() -> Expr {
    apply_multiple(|columns| {
        // 此处基于传入列创建新DataFrame做二次分组聚合
        let df = DataFrame::new(vec![
            columns[0].clone(), columns[1].clone()
        ])? 
        .lazy()
        .groupby_stable([col("Month")])
        .agg([
            col("WeightScaled").sum().alias("WeightScaledSumed")
        ])
        .collect()?;

        // 后续ndarray复杂计算逻辑省略
        Ok( Series::new("result", &[res]) )
        
    }, 
        &[col("Month"), weight_scaled().alias("WeightScaled")], 
        GetOutput::from_type(DataType::Float64))
}

问题根因 #

• apply_multiple在分组上下文中执行时，传入的列是当前分组对应的切片序列，原实现中回调返回的二次聚合结果长度与原分组输入长度不匹配：原分组输入长度为该Region下的原始行数，二次聚合后返回的是月度聚合后的短序列，外层聚合器校验序列长度与分组索引长度不一致时，直接触发长度断言panic。
• 原实现遗漏了分组内过滤Scalar空值的步骤，且未明确输出为单组单标量的映射规则，导致调用first()时空值匹配异常返回空值。
• 外层聚合额外套first()/sum()属于多余操作：自定义回调本身已经完成组内全量计算，再套一层聚合会重复做长度校验，触发异常。

正确实现代码 #

use polars::prelude::*;
use polars::df;
use ndarray::prelude::*;
use ndarray::stack;

pub fn main() {
    let df = df! [
        "Region" => ["EU", "EU", "EU", "EU",                "US", "US", "US", "US"],
        "Month" => ["NOV", "DEC", "APR", "APR",             "JUL", "JAN", "JUL", "SEP"],
        "Weight" => [1, 2, 3, 4,                             5,6,7,8],
        "Scalar" => [Some(0.3), None, Some(0.1), Some(0.1),   Some(0.2), None, Some(0.1), Some(0.3)]
    ].unwrap();

    let df1 = df.clone().lazy()
        .groupby_stable([col("Region")])
        .agg([
            // 直接调用自定义表达式即可，无需额外套first/sum，回调已返回单组单标量
            region_health_index().alias("RegionHealth"),
        ])
        .collect()
        .unwrap();
    dbg!(df1);
}

pub fn weight_scaled() -> Expr {
    col("Weight") * col("Scalar")
}

pub fn region_health_index() -> Expr {
    apply_multiple(|columns| {
        let month_series = columns[0].str()?;
        let ws_series = columns[1].f64()?;

        // 过滤Scalar为空（即WeightScaled为null）的无效行
        let valid_mask = ws_series.is_not_null();
        let month_valid = month_series.filter(&valid_mask)?;
        let ws_valid = ws_series.filter(&valid_mask)?;

        // 按Month二次分组计算月度加权和
        let sub_df = DataFrame::new(vec![
            month_valid.into_series(),
            ws_valid.into_series().with_name("WeightScaled")
        ])?;
        let monthly_agg = sub_df.lazy()
            .groupby_stable([col("Month")])
            .agg([col("WeightScaled").sum().alias("WeightScaledSumed")])
            .collect()?;
        
        // 提取月度聚合结果转为ndarray数组
        let monthly_sum = monthly_agg.column("WeightScaledSumed")?.f64()?;
        let arr = Array1::from_iter(monthly_sum.into_no_null_iter());

        // 替换为实际ndarray复杂数学运算逻辑，示例为直接求和
        let res: f64 = arr.sum();

        // 返回长度为1的标量Series，匹配外层分组单组单值的长度要求
        Ok(Series::new("RegionHealth", &[res]))
    }, 
        &[col("Month"), weight_scaled().alias("WeightScaled")], 
        // 明确声明输出为Float64标量类型，与分组一一对应
        GetOutput::from_type(DataType::Float64)
    )
}

关键修正点 #

• 回调入口先做空值过滤，剔除Scalar为空的无效行，避免空值影响二次聚合结果
• 二次聚合完成后，基于月度聚合结果做ndarray计算，最终返回长度固定为1的标量Series，完全匹配外层GroupBy每个分组返回一个值的长度要求，不会触发长度断言panic
• 移除外层聚合对自定义表达式额外套的first()/sum()等多余聚合函数，回调本身已经完成组内全量计算，直接返回单组结果即可
• 运行后EU分组月度聚合值为[0.3, 0.7]，US分组月度聚合值为[1.7, 2.4]，与预期完全一致，替换示例的ndarray计算逻辑即可得到目标区域健康指数。

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