如何将Polars DataFrame列合并为Parquet兼容的MAP类型（结构体列表）？

要实现将Polars DataFrame多列合并为Parquet MAP类型兼容的List(Struct([Field('key', Utf8), Field('value', Utf8)]))格式，核心思路是先为每列生成包含key（列名）和value（列值）的结构体，再将这些结构体合并为List列。以下是具体实现步骤和示例代码：

核心实现步骤 #

1. 确定目标列：列出需要转换为键值对的列名。
2. 生成Key-Value结构体列：对每个目标列，构造包含固定key（列名）和动态value（列数据）的结构体列。
3. 合并为List列：将所有结构体列组合成一个List类型的列，该结构与Polars读取Parquet MAP类型时的解析格式完全匹配。
4. 写入Parquet：直接使用Polars的write_parquet方法，该方法会自动识别此结构并以Parquet MAP类型写入。

示例代码 #

use polars::prelude::*;

fn main() -> PolarsResult<()> {
    // 创建示例DataFrame
    let df = df!(
        "name" => ["Alice", "Bob", "Charlie"],
        "age" => ["30", "25", "35"],
        "city" => ["New York", "London", "Paris"]
    )?;

    // 定义需要转换的列名
    let target_columns = ["name", "age", "city"];

    // 为每个列生成key-value结构体列
    let struct_cols: Vec<Series> = target_columns
        .iter()
        .map(|col_name| {
            // 构造key列（固定为列名）和value列（对应列数据）
            let key_series = Series::new("key", vec![col_name.to_string(); df.height()]);
            let value_series = df.column(col_name)?.clone().alias("value");
            
            // 组合为结构体列
            StructChunked::new("kv_pair", &[key_series, value_series])?.into_series()
        })
        .collect::<PolarsResult<_>>()?;

    // 将结构体列合并为List列
    let map_column = ListChunked::from_structs("map_data", struct_cols)?
        .into_series()
        .alias("map_data");

    // 合并到原DataFrame并写入Parquet
    let result_df = df.with_column(map_column)?;
    result_df.write_parquet("output.parquet", Default::default())?;

    Ok(())
}

注意事项 #

• 类型一致性：如果目标列不是Utf8类型，需先通过cast(DataType::Utf8)转换，确保key和value类型匹配（若需要其他类型，可调整结构体字段的DataType）。
• 性能优化：上述示例使用矢量化操作，避免了逐行处理的性能损耗，适合大数据量场景。如果必须逐行处理，可使用apply方法，但效率会显著降低。
• 结构验证：可以通过result_df.schema()查看生成的列结构，确认是否为List(Struct([key: Utf8, value: Utf8]))。

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