如何在Rust Polars中使用闭包实现过滤及相关数据处理需求

核心实现方案 #

Polars支持通过内置表达式或闭包两种方式实现你的需求，优先推荐内置表达式（向量化执行性能更高），复杂逻辑再用闭包补充，以下针对每个需求给出具体代码：

1. 过滤非单调递增的列值 #

本质是对比当前行与前一行的列值，用shift函数获取前一行数据，结合filter即可实现：

use polars::prelude::*;

fn main() -> Result<()> {
    let df = df! {
        "values" => [1, 3, 2, 4, 5, 3]
    }?;

    // 生成前一行的移位列，首行无数据会显示null
    let shifted_values = df.column("values")?.shift(1);
    // 过滤出当前值 >= 前一行值的行，首行直接保留
    let filtered_df = df.filter(
        col("values") >= shifted_values
            .or(shifted_values.is_null())
    )?;

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

如果一定要用闭包，可遍历序列生成掩码后过滤：

let values_series = df.column("values")?.i32()?.into_iter();
let mut prev_val = None;
let mask: Vec<bool> = values_series.map(|opt_val| {
    let keep = match (prev_val, opt_val) {
        (None, Some(_)) => true,
        (Some(p), Some(v)) => v >= p,
        _ => false,
    };
    prev_val = opt_val;
    keep
}).collect();

let filtered_df = df.filter(mask.into_series())?;

2. 根据同一行其他值过滤当前值 #

直接用列表达式组合即可，无需闭包。比如过滤col_a大于col_b两倍的行：

let df = df! {
    "col_a" => [10, 5, 20, 8],
    "col_b" => [3, 4, 7, 5]
}?;

let filtered_df = df.filter(col("col_a") > col("col_b") * 2)?;

复杂逻辑下，可通过map_rows生成掩码：

let mask: Vec<bool> = df.map_rows(|row| {
    let col_a = row.get(0).unwrap().i32().unwrap();
    let col_b = row.get(1).unwrap().i32().unwrap();
    col_a > col_b * 2
}).collect();

let filtered_df = df.filter(mask.into_series())?;

3. 根据当前行部分值修改其他值 #

用when/then/otherwise表达式是最优方案，比如当status为"error"时，将value设为0：

let df = df! {
    "status" => ["ok", "error", "ok", "error"],
    "value" => [100, 200, 300, 400]
}?;

let modified_df = df.with_column(
    when(col("status") == lit("error"))
        .then(lit(0))
        .otherwise(col("value"))
        .alias("value")
)?;

复杂逻辑可通过闭包遍历列组合修改：

let modified_values = df.column("value")?.i32()?.into_iter()
    .zip(df.column("status")?.str()?.into_iter())
    .map(|(val_opt, status_opt)| {
        match (val_opt, status_opt) {
            (Some(val), Some("error")) => Some(0),
            (v, _) => v,
        }
    }).collect::<Series>();

let modified_df = df.with_column(modified_values.alias("value"))?;

闭包使用注意事项 #

• 优先用Polars内置表达式，性能远高于闭包（内置表达式是向量化执行）。
• 闭包仅适合处理无法用内置表达式组合实现的复杂逻辑，常用函数有map_rows、map（针对Series）、filter配合自定义Series掩码。

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