在 polars 中加载数据框时如何定义列的类型?

How to define types of columns while loading dataframe in polars?

我正在使用 polars,我想在加载数据框时定义列的类型。在 pandas 中,我可以使用 dtype:

df=pd.read_csv("iris.csv", dtype={'petal_length':str})

我正在尝试在极地做同样的事情,但直到现在都没有成功。这是我尝试过的:

use polars::prelude::*;
use std::fs::File;
use std::collections::HashMap;


fn main() {
    let df = example();
    println!("{:?}", df.expect("Cannot find dataframe").head(Some(10)))
}

fn example() -> Result<DataFrame> {
    let file = File::open("iris.csv")
                    .expect("could not read file");
    let mut myschema = HashMap::new();
    myschema.insert("sepal_length", f64);
    myschema.insert("sepal_width", f64); 
    myschema.insert("petal_length",String); 
    myschema.insert("petal_width", f64); 
    myschema.insert("species", String); 

    CsvReader::new(file)
            .with_schema(myschema)
            .has_header(true)
            .finish()
}

我的疑问是实施 with_schema 需要什么类型的数据?我打印了使用 infer_schema(None) 加载的 DataFrame 的模式。这打印了一个看起来像字典的对象:

Schema { fields: [Field { name: "sepal_length", data_type: Float64 }, Field { name: "sepal_width", data_type: Float64 }, Field { name: "petal_length", data_type: Float64 }, Field { name: "petal_width", data_type: Float64 }, Field { name: "species", data_type: Utf8 }] }

但我不知道应该使用什么对象来实现我的模式。

此外,有一种方法可以指定一个变量的类型,而不是指定所有变量的类型吗?

with_schema 方法需要 Arc<Schema> 类型,而不是 Hashmap

以下代码有效:

use polars::prelude::*;
use std::sync::Arc;

fn example() -> Result<DataFrame> {
    let file = "iris.csv";

    let myschema = Schema::new(
        vec![
            Field::new("sepal_length", DataType::Float64),
            Field::new("sepal_width", DataType::Float64),
            Field::new("petal_length", DataType::Utf8),
            Field::new("petal_width", DataType::Float64),
            Field::new("species", DataType::Utf8),
        ]
    );

    CsvReader::from_path(file)?
        .with_schema(Arc::new(myschema))
        .has_header(true)
        .finish()
}

请注意,通常建议使用 CsvReader::from_path 而不是 CsvReader::new,因为这样性能更高。

Also, there is a way to specify the type of one variable, instead of all of them?

是的,您可以使用 with_dtype_overwrite。需要部分架构。