使用arrow-rs加载Mongo数据实现独热编码遇字典类型支持问题求助

1. 修复Schema适配JSON解码器 #

当前报错是因为Arrow的JSON解码器暂不支持Dictionary(UInt16, Utf8)类型的反序列化，先将Schema中的Country字段改为Utf8类型，确保能成功生成初始RecordBatch：

let schema = Schema::new(vec![
    Field::new("Amount", DataType::Float32, false),
    Field::new("Country", DataType::Utf8, false),
]);

修改后重新运行解码器代码，就能正常生成包含Amount和Country（字符串类型）的RecordBatch。

2. 对Country字段执行独热编码 #

得到初始RecordBatch后，在Rust中完成独热编码操作，步骤如下：

提取唯一国家值 #

从Country列中提取所有不重复的国家名称，用于生成后续的独热列：

// 假设初始batch是已成功生成的RecordBatch
let country_array = batch.column(1).as_any().downcast_ref::<StringArray>().unwrap();
// 使用BTreeSet保证顺序一致，也可用HashSet（无序）
let unique_countries: Vec<&str> = country_array
    .iter()
    .flatten()
    .collect::<std::collections::BTreeSet<_>>()
    .into_iter()
    .collect();

生成独热编码列 #

为每个唯一国家创建对应的布尔数组，当前行属于该国家则值为true，否则为false：

use arrow::array::{BooleanArray, BooleanBuilder, StringArray};
use arrow::datatypes::{Field, Schema};
use arrow::record_batch::RecordBatch;
use std::sync::Arc;

// 保留原Amount列
let mut new_columns = vec![batch.column(0).clone()];

for country in &unique_countries {
    let mut builder = BooleanBuilder::new();
    for c in country_array.iter() {
        // 逐个判断当前行的国家是否匹配
        builder.append_value(c == Some(*country)).unwrap();
    }
    let bool_array = builder.finish();
    new_columns.push(Arc::new(bool_array) as Arc<dyn arrow::array::Array>);
}

构建新的Schema和RecordBatch #

基于生成的独热列创建新Schema，组合成最终要发送的RecordBatch：

// 构建新字段列表：保留Amount，加上每个国家对应的布尔字段
let mut new_fields = vec![schema.field(0).clone()];
for country in &unique_countries {
    new_fields.push(Field::new(country, DataType::Boolean, false));
}
let new_schema = Schema::new(new_fields);

// 生成最终的RecordBatch
let final_batch = RecordBatch::try_new(Arc::new(new_schema), new_columns).unwrap();

3. 通过Arrow Flight发送数据 #

将final_batch通过Arrow Flight发送给客户端即可，pyarrow客户端接收后，直接调用to_pandas()就能转换为包含独热编码列的DataFrame，无需额外处理。

注意事项 #

• 实际代码中建议替换unwrap()为错误处理逻辑（如match或?运算符），避免panic；
• 若数据集较大，可考虑分批处理独热编码，避免内存占用过高；
• 若需要固定的国家映射（而非动态提取唯一值），可预先定义映射表，跳过提取唯一值的步骤。

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