Polars中如何为列表字符串序列添加元素计数列？

问题背景 #

拆分字符串为单词列表后，尝试添加count列统计每个列表的元素数量时出现类型不匹配错误：

mismatched types
expected struct `polars::prelude::Series`, found `usize`

初始DataFrame代码：

let df = df! [
        "keys" => ["a ab", "a ab abc", "b ba abc abcd", "b ba bbc abcd bbcd"],
        "groups" => ["A", "A", "B", "C"],
    ]?;

拆分字符串代码：

let out = df.lazy().with_column(col("keys").str().split(lit(" ")));

错误的计数代码：

let out_2 = out.with_columns([col("keys")
      .apply(|s| Ok(s.len()), GetOutput::from_type(DataType::Int32))
      .alias("count")]).collect().unwrap();

错误原因 #

apply的闭包参数接收的是整个Series对象，而非单个行的列表元素。直接返回s.len()得到的是整个Series的行数（usize类型），但Polars要求返回一个Series，因此触发类型不匹配。

正确解决方案 #

Polars针对列表列提供了内置的list.len()方法，无需手动使用apply，既简洁又高效：

完整正确代码 #

use polars::prelude::*;

fn main() -> PolarsResult<()> {
    let df = df! [
        "keys" => ["a ab", "a ab abc", "b ba abc abcd", "b ba bbc abcd bbcd"],
        "groups" => ["A", "A", "B", "C"],
    ]?;

    let result_df = df.lazy()
        // 拆分字符串为列表列，保留原keys列
        .with_column(col("keys").str().split(lit(" ")).alias("split_keys"))
        // 统计每个列表的元素数量，生成count列
        .with_column(col("split_keys").list().len().alias("count"))
        .collect()?;

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

代码说明 #

1. str().split(lit(" "))将keys列的字符串按空格拆分，生成列表类型的列split_keys（避免覆盖原keys列，若不需要原列可直接覆盖）。
2. list().len()方法会遍历列表列的每个元素，统计每个列表的元素个数，返回整数类型的Series，完美匹配Polars的类型要求。

若坚持使用apply（不推荐） #

如果一定要用apply实现，需遍历Series中的每个列表元素，手动构造返回的Series：

let result_df = out.with_columns([
    col("keys")
        .apply(|s: Series| {
            // 遍历Series中的每个列表元素，统计长度并转为i32
            let len_series = s
                .iter()
                .map(|opt_list| opt_list.map(|list_val: &ListValue| list_val.len() as i32))
                .collect::<Series>();
            Ok(len_series)
        }, GetOutput::from_type(DataType::Int32))
        .alias("count")
]).collect()?;

但内置方法list.len()性能更优，代码更简洁，优先推荐使用。

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