如何用Rust DataFusion DataFrame API读取空格分隔CSV文件

问题背景 #

使用Rust DataFusion读取UCI auto-mpg空格分隔数据集时，返回空Schema和空数据，而相同数据集用Pandas可正常读取。

关键问题点 #

原代码未处理以下核心配置：

• 未指定数据集的Schema与列名
• 未处理多连续空格分隔符（对应Pandas的skipinitialspace=True）
• 未开启带引号字段的解析（最后一列为带引号的字符串，包含空格）
• 未配置缺失值标识（对应Pandas的na_values="?"）

解决方案 #

1. 定义匹配数据集的Schema #

根据数据集字段，创建包含列名和对应数据类型的Schema：

use datafusion::arrow::datatypes::{Schema, Field, DataType};

fn get_auto_mpg_schema() -> Schema {
    Schema::new(vec![
        Field::new("MPG", DataType::Float64, true),
        Field::new("Cylinders", DataType::Int64, true),
        Field::new("Displacement", DataType::Float64, true),
        Field::new("Horsepower", DataType::Float64, true),
        Field::new("Weight", DataType::Float64, true),
        Field::new("Acceleration", DataType::Float64, true),
        Field::new("Model Year", DataType::Int64, true),
        Field::new("Origin", DataType::Int64, true),
        Field::new("Car Name", DataType::Utf8, true), // 数据集实际包含的汽车名字段
    ])
}

2. 配置正确的CSV读取选项 #

修改读取配置，添加必要的解析规则：

use datafusion::{prelude::*, error::Result};

fn get_csv_option(schema: &Schema) -> CsvReadOptions {
    let mut csv_opt = CsvReadOptions::new()
        .schema(schema)
        .has_header(false)
        .delimiter(b' ')
        .skip_initial_space(true) // 跳过字段前的空格，适配多连续分隔符
        .allow_quoted_records(true) // 解析带引号的含空格字段
        .na_values(vec!["?".to_string()]); // 将"?"识别为缺失值
    csv_opt
}

3. 完整运行代码 #

整合上述部分，修改主函数：

use datafusion::arrow::datatypes::{Schema, Field, DataType};
use datafusion::{prelude::*, error::Result};

fn get_auto_mpg_schema() -> Schema {
    Schema::new(vec![
        Field::new("MPG", DataType::Float64, true),
        Field::new("Cylinders", DataType::Int64, true),
        Field::new("Displacement", DataType::Float64, true),
        Field::new("Horsepower", DataType::Float64, true),
        Field::new("Weight", DataType::Float64, true),
        Field::new("Acceleration", DataType::Float64, true),
        Field::new("Model Year", DataType::Int64, true),
        Field::new("Origin", DataType::Int64, true),
        Field::new("Car Name", DataType::Utf8, true),
    ])
}

fn get_csv_option(schema: &Schema) -> CsvReadOptions {
    let mut csv_opt = CsvReadOptions::new()
        .schema(schema)
        .has_header(false)
        .delimiter(b' ')
        .skip_initial_space(true)
        .allow_quoted_records(true)
        .na_values(vec!["?".to_string()]);
    csv_opt
}

#[tokio::main]
async fn main() -> Result<()> {
    let schema = get_auto_mpg_schema();
    let read_option = get_csv_option(&schema);
    let ctx = SessionContext::new();
    
    // 替换为你的数据集实际路径
    let df = ctx.read_csv("data/landing/auto-mpg.data", read_option).await?;
    
    println!("{}", df.schema());
    df.show().await?;
    Ok(())
}

代码说明 #

• skip_initial_space(true)：自动跳过每个字段开头的空格，解决多连续空格作为分隔符的问题
• allow_quoted_records(true)：正确解析包含空格的带引号字段（如汽车名称）
• na_values(vec!["?".to_string()])：将数据集中的"?"识别为NULL值，与Pandas行为一致
• 显式定义Schema：让DataFusion明确每个字段的名称和数据类型，避免自动推断失败

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