如何在Rust的Polars中为DataFrame添加条件计算列？

依赖说明 #

首先确保你的Cargo.toml引入了正确的Polars依赖：

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

• lazy feature 是延迟执行模式的必需依赖
• fmt feature 用于打印DataFrame，不需要可删除

Eager（立即执行）模式实现 #

直接对内存中的DataFrame执行计算，调用后立即生成结果：

use polars::prelude::*;

fn main() -> PolarsResult<()> {
    // 示例构造df1，可替换为实际数据源读取逻辑
    let df1 = df!(
        "a" => [10, -10, 0, 6, -6],
        "b" => [1, 2, 3, 4, 5],
        "c" => [6, 7, 8, 9, 10]
    )?;

    // 按规则计算d列
    let d_col = when(col("a").gt(5))
        .then(col("b") + col("c"))
        .when(col("a").lt(-5))
        .then(col("c") - col("b") + col("a"))
        .otherwise(lit(3.0) * col("a"))?
        .alias("d");
    
    // 合并得到df2
    let df2 = df1.with_column(d_col)?;

    println!("Eager模式执行结果：\n{}", df2);
    Ok(())
}

Lazy（延迟执行）模式实现 #

先构建执行计划，调用collect时才触发实际计算，支持查询优化，适合大数据量场景：

use polars::prelude::*;

fn main() -> PolarsResult<()> {
    // 示例构造df1，可替换为实际数据源读取逻辑
    let df1 = df!(
        "a" => [10, -10, 0, 6, -6],
        "b" => [1, 2, 3, 4, 5],
        "c" => [6, 7, 8, 9, 10]
    )?;

    // 构建执行计划并计算df2
    let df2 = df1
        .lazy()
        .with_column(
            when(col("a").gt(5))
                .then(col("b") + col("c"))
                .when(col("a").lt(-5))
                .then(col("c") - col("b") + col("a"))
                .otherwise(lit(3.0) * col("a"))
                .alias("d")
        )
        .collect()?;

    println!("Lazy模式执行结果：\n{}", df2);
    Ok(())
}

逻辑说明 #

代码中的when/then/otherwise链式结构和给出的原生Rust判断逻辑完全对应，无需自己实现逐行遍历，Polars内部会做向量化执行，性能远高于手动迭代行。

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