如何在Polars Rust中使用with_column方法创建计算列

前置配置 #

• 表达式计算、col列引用方法默认不包含在基础功能中，需要在Cargo.toml中为polars依赖开启对应feature，否则无法找到col方法：

[dependencies]
polars = { version = "0.22.8", features = ["lazy", "dtype-i32", "dtype-str", "dtype-f64"] }

• 你提供的示例代码存在基础语法错误：main函数定义缺少参数括号，example函数和main函数均缺少闭合花括号，且DataFrame相关操作均返回PolarsResult类型，需要做错误处理。

具体实现 #

with_column方法接收一个Polars表达式作为入参，使用方式和PySpark一致：用col("列名")引用已有列做算术/逻辑运算，最后通过.alias("新列名")指定生成列的列名即可。
如果需要计算浮点数比值，需要先把整数列转换为浮点类型，否则会执行整数除法丢弃小数部分。
完整可运行代码如下：

use polars::prelude::*;

fn example() -> PolarsResult<()> {
    let df = df![
        "foo" => ["A", "A", "B", "B", "C"],
        "val1" => [1, 2, 2, 4, 2],
        "val2" => [1, 2, 2, 4, 2],
    ]?;

    // 新增名为ration的计算列，存储val1/val2的浮点比值
    let df_res = df.with_column(
        (col("val1").cast(DataType::Float64) / col("val2").cast(DataType::Float64))
            .alias("ration")
    )?;

    println!("{}", df_res);
    Ok(())
}

fn main() -> PolarsResult<()> {
    example()
}

运行后输出结果会包含三列原始列+新增的ration列，你提供的测试数据下比值均为1.0。替换val2的测试值即可看到不同的计算结果，不需要浮点精度时删除.cast(DataType::Float64)转换逻辑即可，Polars会自动执行整数除法。

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