如何在Polars中将多列列表数据转为结构体列表列？

问题场景 #

现有通过CsvReader读取并处理后的DataFrame，包含两列origin和destination，类型均为List(Utf8)。需求是将这两列转换为名为MyData的列，该列类型为List(Struct([origin: Utf8, destination: Utf8]))——即每一行的MyData是一个结构体列表，每个结构体由origin和destination列表中对应位置的元素配对组成，最终导出为指定结构的JSON。

解决方案代码 #

修改原有Rust代码，通过Polars的arrays::zip_with函数实现列表元素的配对与结构体构造：

use polars::prelude::*;
use polars::functions::arrays;

fn read_csv() -> Result<(), PolarsError> {
    println!("Hello, polars! 🐯");
    let df = CsvReader::from_path("./test_data/tar-data/csv/unziped/body.csv")?
        .has_header(false)
        .finish()?;

    let df = df.lazy()
        .select([
            col("column_15").str().split("|").alias("origin"),
            col("column_16").str().split("|").alias("destination"),
        ])
        .with_column(
            arrays::zip_with(
                col("origin"),
                col("destination"),
                |origin_list, dest_list| {
                    // 遍历两个列表的对应元素，构造结构体
                    let structs = origin_list.into_iter()
                        .zip(dest_list.into_iter())
                        .map(|(origin_val, dest_val)| {
                            StructValue::new(vec![
                                ("origin".into(), origin_val.unwrap_or(Utf8Value::from(""))),
                                ("destination".into(), dest_val.unwrap_or(Utf8Value::from(""))),
                            ])
                        })
                        .collect::<ListValue>();
                    
                    Ok(structs)
                },
            ).alias("MyData")
        )
        // 可选：移除原有的origin和destination列
        .drop_columns(["origin", "destination"])
        .collect()?;

    println!("Schema {:?}", df.schema());
    println!("{:?}", df);

    // 导出为JSON
    df.write_json("./output.json")?;

    Ok(())
}

关键逻辑说明 #

• arrays::zip_with：针对每一行的origin和destination列表进行逐元素配对处理，是Polars中处理列表列元素级操作的核心函数。
• 结构体构造：通过StructValue::new创建包含origin和destination字段的结构体，将两个列表对应位置的元素分别赋值给结构体字段。
• 空值处理：使用unwrap_or为可能的空元素提供默认值（此处为空字符串），可根据业务需求调整为空值保留或其他默认值。

结果验证 #

执行代码后，MyData列的Schema将符合预期：

Schema Schema:
name: MyData, data type: List(Struct([Field { name: "origin", dtype: Utf8 }, Field { name: "destination", dtype: Utf8 }]))

数据示例将变为：

shape: (10, 1)
┌───────────────────────────────────────────────────────────────────────────┐
│ MyData                                                                   │
│ ---                                                                      │
│ list[struct[2]]                                                          │
╞═══════════════════════════════════════════════════════════════════════════╡
│ [{"origin": "JOI", "destination": "VCP"}, {"origin": "GRU", "destination": "DFW"}, {"origin": "DFW", "destination": "SLC"}] │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ [{"origin": "JOI", "destination": "GRU"}, {"origin": "GRU", "destination": "ATL"}, {"origin": "ATL", "destination": "SLC"}] │
└───────────────────────────────────────────────────────────────────────────┘

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