Polars Lazy API with_context过滤报错求助

问题描述 #

我搞不懂Polars惰性API里with_context方法的用法。官方文档说这个方法允许表达式访问不属于当前DataFrame的列，但我用另一个DataFrame的列过滤当前列时，会报错“其他DataFrame的列不存在”；换成select操作却能正常执行。

报错代码（main.rs） #

use polars::prelude::*;

fn main() {
    let df0 = df! {
        "id" => [1, 2, 3],
        "name" => ["foo", "bar", "baz"],
    }
    .unwrap()
    .lazy();
    let other_df = df! {
        "other_id" => [1,2,2,1],
        "name" => ["w", "x", "y", "z"],
    }
    .unwrap()
    .lazy();

    let lf = df0.with_context(&[other_df]);
    let res = lf
        .filter(col("id").is_in(col("other_id")))
        .collect()
        .unwrap();
    println!("{:?}", res);
}

Cargo.toml依赖 #

[dependencies]
polars = {git = "https://github.com/pola-rs/polars", branch = "master", features = ["lazy", "is_in"]}

可正常执行的select替代代码 #

let res = lf
    .select(&[col("id").is_in(col("other_id"))])
    .collect()
    .unwrap();

问题原因 #

with_context的核心是允许表达式访问上下文列，但**filter操作要求逐行对齐当前DataFrame与上下文DataFrame的行**。你的示例中，df0有3行，other_df有4行，行数不匹配导致Polars无法完成行对齐，因此报错。

而select中的is_in操作逻辑不同：它会自动将上下文列other_id作为一个整体集合处理，判断当前DataFrame每行的id是否属于这个集合，不需要行对齐，因此能正常执行。

解决方案 #

方案1：提取上下文列的集合后过滤 #

先将上下文DataFrame的other_id提取为一个独立的集合，再用于过滤当前DataFrame：

use polars::prelude::*;

fn main() {
    let df0 = df! {
        "id" => [1, 2, 3],
        "name" => ["foo", "bar", "baz"],
    }
    .unwrap()
    .lazy();
    let other_df = df! {
        "other_id" => [1,2,2,1],
        "name" => ["w", "x", "y", "z"],
    }
    .unwrap();

    // 获取other_id的唯一值集合
    let other_ids = other_df.column("other_id").unwrap().unique().unwrap();

    let res = df0
        .filter(col("id").is_in(other_ids))
        .collect()
        .unwrap();
    println!("{:?}", res);
}

方案2：通过Join关联后过滤 #

使用内连接（Inner Join）关联两个DataFrame，再去重得到结果：

use polars::prelude::*;

fn main() {
    let df0 = df! {
        "id" => [1, 2, 3],
        "name" => ["foo", "bar", "baz"],
    }
    .unwrap()
    .lazy();
    let other_df = df! {
        "other_id" => [1,2,2,1],
        "name" => ["w", "x", "y", "z"],
    }
    .unwrap()
    .lazy();

    let res = df0
        .join(other_df, col("id"), col("other_id"), JoinType::Inner)
        .select(&[col("id"), col("name_left").alias("name")])
        .unique()
        .collect()
        .unwrap();
    println!("{:?}", res);
}

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