如何在Rust Polars中按索引取值并过滤DataFrame？

问题分析 #

你遇到的错误核心原因：

• USER_INFO.column("user_countries")?.get(0) 返回的是 Result<AnyValue<'_>, PolarsError> 类型，直接传入 lit() 不符合Literal trait的实现要求
• user_countries 是list[str]类型列，提取后需要转换为Polars可识别的序列类型才能用于过滤

修正步骤 #

• 解析索引值并处理错误：使用?（生产环境推荐）或unwrap()获取get(0)的实际值，避免直接传递Result对象
• 提取列表中的字符串元素：匹配AnyValue::List类型，将其转换为Vec<&str>或Vec<String>格式
• 适配过滤函数要求：将转换后的向量传入is_in()，Polars会自动处理为可用于过滤的字面量

完整修正代码 #

use polars::prelude::*;

fn main() -> Result<(), PolarsError> {
    // 假设USER_INFO和SHOPS_INFO已提前初始化完成
    // 提取第一个用户的国家列表
    let first_user_countries = match USER_INFO
        .column("user_countries")?
        .get(0)? {
            AnyValue::List(s) => s
                .iter()
                .map(|val| val.str().unwrap())
                .collect::<Vec<_>>(),
            _ => panic!("user_countries列不是list[str]类型"),
        };

    // 过滤SHOPS_INFO DataFrame
    let filtered_df = SHOPS_INFO
        .lazy()
        .filter(col("country").is_in(lit(first_user_countries)))
        .collect()?;

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

额外优化：提取所有用户的国家列表（去重后过滤） #

如果需要基于所有用户的国家（去重后）过滤SHOPS_INFO，可以用以下方式：

// 提取所有用户的国家并去重
let all_countries = USER_INFO
    .column("user_countries")?
    .explode()?
    .unique()?
    .iter()
    .map(|val| val.str().unwrap())
    .collect::<Vec<_>>();

// 执行过滤
let filtered_df = SHOPS_INFO
    .lazy()
    .filter(col("country").is_in(lit(all_countries)))
    .collect()?;

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