polars读取csv数据时自定义列数据类型

默认情况下，CsvReader自动匹配csv文件中的数据类型：

1	let df = CsvReader::from_path(filecsv)?.has_header(true).finish()?; Ok(df)

当csv文件中的某列出现前导0时，默认的读取会将这一列的类型设置为整数类型，且抹除前导0，这并不是我们想要的，因为前导0的列往往需要被解析为字符str类型。

有两种解决方案。

部分转换

将一些（不是全部）字段转换为 str 类型，这样做的前提是你事先知道要指定将哪些列转换。

项目的Cargo.toml依赖如下：

1
2
3

[dependencies]
polars = { version = "0.33.2", features = ["lazy","ndarray"] }
smartstring = "1.0.1"

代码如下：

use polars::datatypes::DataType::Utf8;
use polars::prelude::*;
use smartstring::SmartString;
use std::sync::Arc;
fn main() {
    let mut schema = Schema::new();
    schema.with_column(SmartString::from("some_columns"), Utf8);
    let df_csv = CsvReader::from_path("some_input.csv")
        .unwrap()
        .infer_schema(None)
        .has_header(true)
        .with_dtypes(Some(Arc::new(schema)))
        .finish()
        .unwrap();
    println!("{}", df_csv);
}

通过with_column设置schema，传入需要自定义的列和类型，这里我需要的是将其转换为polars::datatypes::DataType::Utf8类型。最后通过with_dtypes即可应用schema。

全部转换

在一些情况下，部分转换并不方便，比如：

事先不知道csv中有哪些字段
字段数量很多，使用with_column不方便
想将所有列（无论它原本是什么类型）都读取为str

经过进一步的实验，我找到了另一种方法。这需要使用 csv 库。

项目的Cargo.toml依赖如下：

1
2
3

[dependencies]
polars = { version = "0.33.2", features = ["lazy","ndarray"] }
csv = "1.3.0"

代码如下：

use polars::datatypes::DataType::Utf8;
use polars::prelude::*;
use std::sync::Arc;

fn main() {
    
    let mut rdr = csv::Reader::from_path("some_input.csv").unwrap();

    let column_names = rdr.headers().unwrap().iter().map(|item| Field::new(item,Utf8));
    let schema = Schema::from_iter(column_names);
    let df_csv = CsvReader::from_path("some_input.csv")
        .unwrap()
        .infer_schema(None)
        .has_header(true)
        .with_dtypes(Some(Arc::new(schema)))
        .finish()
        .unwrap();
    println!("{}", df_csv);
}