Rust使用Polars CsvReader读取无表头CSV时如何为DataFrame添加列名

前置准备 #

先创建数据存储目录并下载测试用葡萄酒数据集，执行如下bash命令：

mkdir -p datastore
wget https://archive.ics.uci.edu/ml/machine-learning-databases/wine/wine.data -O datastore/wine.data

方案1：读取CSV阶段直接指定列名（推荐） #

Polars的CsvReader提供了with_column_names方法，可以在文件解析阶段直接给无表头CSV绑定自定义列名，不需要读取后二次修改，性能更好。
完整可运行代码如下：

use polars::prelude::*;

// 预设列名，这里定义为静态常量避免重复内存分配，也可直接传入Vec<String>类型的列名向量
const COLUMN_NAMES: &[&str] = &[
    "Class label", "Alcohol",
    "Malic acid", "Ash",
    "Alcalinity of ash", "Magnesium",
    "Total phenols", "Flavanoids",
    "Nonflavanoid phenols",
    "Proanthocyanins",
    "Color intensity", "Hue",
    "OD280/OD315 of diluted wines",
    "Proline"
];

fn read_wine_data() -> Result<DataFrame> {
    let file = "datastore/wine.data";
    CsvReader::from_path(file)?
        .has_header(false) // 声明原文件不包含表头行
        .with_column_names(COLUMN_NAMES) // 读取时直接绑定自定义列名
        .finish()
}

fn main() {
    let df = read_wine_data();
    match df {
        Ok(content) => println!("{:?}", content.head(Some(10))),
        Err(error) => panic!("Problem reading file: {:?}", error)
    }
}

方案2：读取完成后重命名DataFrame列 #

如果你已经完成CSV读取拿到了DataFrame实例，可以用set_column_names方法批量修改所有列名，需要保证传入的列名数量和DataFrame实际列数完全一致，否则会抛出错误。
示例代码如下：

use polars::prelude::*;

fn read_wine_data() -> Result<DataFrame> {
    let file = "datastore/wine.data";
    let mut df = CsvReader::from_path(file)?
        .has_header(false)
        .finish()?;
    
    // 读取完成后全量设置列名
    let column_names = vec![
        "Class label", "Alcohol",
        "Malic acid", "Ash",
        "Alcalinity of ash", "Magnesium",
        "Total phenols", "Flavanoids",
        "Nonflavanoid phenols",
        "Proanthocyanins",
        "Color intensity", "Hue",
        "OD280/OD315 of diluted wines",
        "Proline"
    ];
    df.set_column_names(&column_names)?;
    Ok(df)
}

fn main() {
    let df = read_wine_data();
    match df {
        Ok(content) => println!("{:?}", content.head(Some(10))),
        Err(error) => panic!("Problem reading file: {:?}", error)
    }
}

注意事项 #

• 两种方案最终效果完全一致，优先选择读取时指定列名的方案，减少一次DataFrame修改操作，性能更优
• 传入的列名数量必须和CSV文件实际列数匹配，本次使用的葡萄酒数据集共14列，给出的预设列名刚好14个，不会出现匹配错误
• 如果只需要修改部分列名，可以调用DataFrame的rename方法单独修改指定列名称，不需要全量传入所有列名

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