如何在Polars-Rust中计算多列中位数并生成新DataFrame

依赖准备 #

确保你的Cargo.toml中包含带lazy特性的Polars依赖：

polars = { version = "0.35", features = ["lazy", "fmt"] }

实现步骤 #

核心思路是将目标列打包为每行的数组，再对数组计算中位数。以下是完整示例，包含测试数据生成、外连接及中位数计算：

use polars::prelude::*;

fn main() -> Result<()> {
    // 生成测试用的三个初始DataFrame
    let df1 = df!["date" => ["2023-01-01", "2023-01-02"], "a" => [1, 4]]?;
    let df2 = df!["date" => ["2023-01-01", "2023-01-03"], "a" => [2, 5]]?;
    let df3 = df!["date" => ["2023-01-01", "2023-01-02"], "a" => [3, 6]]?;

    // 按date列外连接，并重命名重复的a列为a1、a2
    let joined_lf = df1.lazy()
        .join(df2.lazy(), [col("date")], [col("date")], JoinArgs::new(JoinType::Outer))
        .rename("a_right" => "a1")
        .join(df3.lazy(), [col("date")], [col("date")], JoinArgs::new(JoinType::Outer))
        .rename("a_right" => "a2");

    // 计算每行a、a1、a2的中位数，生成目标DataFrame
    let result_lf = joined_lf
        .select([
            col("date"),
            // 打包三列为数组，计算中位数并命名为a_median
            arr_make(&[col("a"), col("a1"), col("a2")])
                .map(|arr| arr.median(), GetOutput::from_type(DataType::Float64))
                .alias("a_median")
        ]);

    // 执行LazyFrame并打印结果
    let result_df = result_lf.collect()?;
    println!("{}", result_df);

    Ok(())
}

关键细节说明 #

• arr_make(&[col("a"), col("a1"), col("a2")])：将指定列合并为每行的数组，外连接产生的空值会被保留在数组中。
• .map(|arr| arr.median(), ...)：对每行数组调用中位数计算，Polars的数组median方法会自动忽略空值；若所有值为空，结果也会是空值。
• GetOutput::from_type(DataType::Float64)：指定输出类型为浮点数——因为当非空值为偶数个时，中位数是中间两数的平均值，会产生小数结果。若需要整数类型，可在后续添加.cast(DataType::Int64)（需注意空值和小数转整数的截断问题）。

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