Rust Polars如何简洁整合多聚合表达式构建指定格式DataFrame

核心思路 #

之前的冗余代码问题本质是把每个聚合操作的查询、转置、重命名、拼接逻辑重复编写，只要把重复逻辑抽象为批量处理即可，不需要为每个统计量写单独逻辑。
根据性能需求可以选两种实现方式，最终输出结构完全符合要求：原始字段为行，统计量为列，非数值字段自动填充null。 #

实现方式1：逻辑直观版（和原生手写逻辑完全一致） #

这个版本逻辑和最初的手写代码完全对齐，把重复操作抽成循环，所有聚合规则统一在一个列表维护，增删统计量只需要修改列表即可，不需要改其他代码：

use polars::prelude::*;

fn main() -> PolarsResult<()> {
    let series_a = Series::new("ID", vec![1, 2, 3, 4]);
    let series_b = Series::new("Amount", vec![10.0, 22.0, 13.3, 54.1]);
    let series_c = Series::new("Name", vec!["Item 1", "Item 2", "Item 3", "Item 4"]);

    let df = DataFrame::new(vec![series_a, series_b, series_c])?;
    let lf = df.lazy();

    // 所有需要的聚合规则统一维护在这里，新增/删除统计量只需要修改这个数组
    let agg_list = [
        ("mean", |c: Expr| c.mean()),
        ("median", |c: Expr| c.median()),
        ("std", |c: Expr| c.std(1)), // 样本标准差，自由度传1
        ("min", |c: Expr| c.min()),
        ("max", |c: Expr| c.max()),
    ];

    // 批量执行每个聚合查询
    let mut agg_dfs = Vec::with_capacity(agg_list.len());
    for (_, agg_func) in agg_list.iter() {
        let one_agg_res = lf.clone().select([agg_func(col("*"))]).collect()?;
        agg_dfs.push(one_agg_res);
    }

    // 把所有聚合结果纵向拼接，列名保持和原始字段一致
    let mut combined_agg = agg_dfs[0].clone();
    for single_df in agg_dfs[1..].iter() {
        combined_agg.vstack_mut(single_df)?;
    }

    // 仅需要一次转置，统一设置列名为统计量名称
    let mut final_df = combined_agg.transpose(None, None)?;
    final_df.set_column_names(
        &agg_list.iter().map(|(name, _)| name.as_str()).collect::<Vec<_>>()
    )?;

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

实现方式2：性能优化版（单次扫描数据） #

如果数据量较大，推荐用这个版本，所有聚合计算会在一次查询中完成，仅扫描一次数据，不会重复执行计划，性能更高：

use polars::prelude::*;

fn main() -> PolarsResult<()> {
    let series_a = Series::new("ID", vec![1, 2, 3, 4]);
    let series_b = Series::new("Amount", vec![10.0, 22.0, 13.3, 54.1]);
    let series_c = Series::new("Name", vec!["Item 1", "Item 2", "Item 3", "Item 4"]);

    let df = DataFrame::new(vec![series_a, series_b, series_c])?;
    let lf = df.lazy();
    let original_cols = df.get_column_names();

    // 统一维护聚合规则
    let agg_list = [
        ("mean", |c: Expr| c.mean()),
        ("median", |c: Expr| c.median()),
        ("std", |c: Expr| c.std(1)),
        ("min", |c: Expr| c.min()),
        ("max", |c: Expr| c.max()),
    ];

    // 一次性生成所有聚合表达式，通过加前缀避免列名重名
    let mut all_agg_exprs = Vec::new();
    for (agg_name, agg_func) in agg_list.iter() {
        for col_name in original_cols.iter() {
            all_agg_exprs.push(
                agg_func(col(col_name))
                    .alias(&format!("{}__{}", agg_name, col_name))
            );
        }
    }

    // 单次collect完成所有聚合计算
    let wide_agg_res = lf.select(all_agg_exprs).collect()?;

    // 按统计量拆分列，重构为「统计量为行、原始字段为列」的结构
    let mut stat_rows = Vec::with_capacity(agg_list.len());
    for (agg_name, _) in agg_list.iter() {
        let col_prefix = format!("{}__", agg_name);
        let select_cols = original_cols.iter()
            .map(|c| format!("{}{}", col_prefix, c))
            .collect::<Vec<_>>();
        let mut single_stat_df = DataFrame::new(wide_agg_res.select_series(select_cols)?)?;
        single_stat_df.set_column_names(original_cols)?;
        stat_rows.push(single_stat_df);
    }

    // 纵向拼接后转置得到最终结果
    let mut combined = stat_rows[0].clone();
    for s_df in stat_rows[1..].iter() {
        combined.vstack_mut(s_df)?;
    }
    let mut final_df = combined.transpose(None, None)?;
    final_df.set_column_names(
        &agg_list.iter().map(|(n,_)| n.as_str()).collect::<Vec<_>>()
    )?;

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

说明 #

之前尝试的map_alias写法得到宽表，是因为直接给每个原始字段的聚合结果单独命名为「字段名+统计量」，没有做行列转置，自然无法得到「统计量为列、原始字段为行」的结构。

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