RustのPolarsにおいてSchemaでデータ型指定してDataframeを読み込みnullを0で置換する

2023/04/23

当サイトではアフィリエイト広告を利用しています

PythonにおけるPandasに相当するのがRustにおけるPolarsです。

ただ、Pandasと比べて情報が少なく、Python版Polarsと情報が混同されていることが多いです。

今回はタイトル通りの内容を実現するのための情報集めに手間取りましたので、備忘録として残しておきます。

PythonにおけるPandasに相当するのがRustにおけるPolarsです。

ただ、PolarsはRust版とPython版があり、Rust版の情報を探しづらいうえに、公式のドキュメントにサンプルがあまりなくて使いにくい問題があります。

ここでは、Rust版PolarsにおいてSchemaを利用してDataframe読み込み時の列のデータ型を指定し、nullだった場合に0埋めをするようなよくある処理についてまとめます。

Polars自体は非常に高速でありRustで使える利点は大きいのですが、公式ドキュメントの情報が少ないのが唯一の難点だと個人的に思います。

メジャーバージョン1.0に期待したいところです。

利用したバージョン

今回はpolars = "0.28.0"を用いています。

Polarsはまだメジャーバージョン1になってないこともあり、メソッド名などがバージョンアップの際に置き換わって利するので要注意です。

今回利用するtry_apply_at_idx()関数も少し前はmay_apply_at_idx()という名前の関数だったようです。

対象のデータ

サンプルとして以下のようなデータに抜けのあるCSVを対象とします。

名前,整数1,小数1,整数2,小数2
A,1,2.0,2,1.1
B,,2.0,2,
C,1,,2,
D,,2.0,,1.1

データの型指定

infer_schema()で序盤のx行からデータ型を推定することもできるのですが、データ形式によっては事前に指定したいことがありますし、空欄混在の場合はstr列とみなされることもあります。

独自にデータ型を指定するならSchemaを使って型を指定できます。

その場合はwith_columnを使って一列毎にデータ形式を指定することになるようです。

nullの置き換え

nullの書き換えにはfill_null()関数を使うことになるのですが、Dataframeのfill_null()関数は0置換に対応していないです。

これはstr列なども混在することを考えると仕方ないのでしょう。

なので列ごとにSeriesを抽出して、Seriesに対してfill_null()関数を適用することになります。

その際、FillNullStrategy::Zeroを引数を与えることで0置換ができます。

また、列ごとに処理を適用するにはapply_at_idx()とtry_apply_at_idx()の２通りの手段があってどちらでも対応可能ですが、エラー処理回数を減らす目的でtry_apply_at_idx()を使ったほうが設計に沿っていると思われます。

サンプルコード

データの列ごとのデータ型を指定し、抜けがありnullになっているデータを0に置き換える処理をまとめると以下のようなコードになりました。

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

fn main(){
    let path = Path::new("data.csv");

    // データ型の指定
    let mut schema = Schema::new();
    schema.with_column("名前".to_string().into(), DataType::Utf8);
    schema.with_column("整数1".to_string().into(), DataType::Int64);
    schema.with_column("整数2".to_string().into(), DataType::Int64);
    schema.with_column("小数1".to_string().into(), DataType::Float64);
    schema.with_column("小数2".to_string().into(), DataType::Float64);

    // Schemaを指定したCSVデータの読み込み
    let mut df = CsvReader::from_path(path)
        .unwrap()
        .with_dtypes(Some(Arc::new(schema)))
        .has_header(true)
        .finish()
        .unwrap();
    println!("{:?}", &df);

    // データ列の型を抽出し、Float64とInt64に対してのみ0埋めを実施する
    let dtypes = df.dtypes();
    println!("{:?}", &dtypes);

    for (i, dt) in dtypes.iter().enumerate() {
        if let DataType::Float64 | DataType::Int64 = dt {
            // 以下はどちらでも可

            // df.apply_at_idx(i, |s| s.fill_null
            // apply_at_idxはクロージャの戻り値がSeries(FillNullStrategy::Zero).unwrap()).unwrap();

            // try_apply_at_idxはクロージャの戻り値がResult
            df.try_apply_at_idx(i, |s| s.fill_null(FillNullStrategy::Zero)).unwrap();
        }
    }
    println!("{:?}", &df);
}

出力

以下がnull処理前のDataframeとnull処理後のDataframeになります。

意図通りにnullが置換できていることがわかります。

shape: (4, 5)
┌──────┬───────┬───────┬───────┬───────┐
│ 名前 ┆ 整数1 ┆ 小数1 ┆ 整数2 ┆ 小数2 │
│ ---  ┆ ---   ┆ ---   ┆ ---   ┆ ---   │
│ str  ┆ i64   ┆ f64   ┆ i64   ┆ f64   │
╞══════╪═══════╪═══════╪═══════╪═══════╡
│ A    ┆ 1     ┆ 2.0   ┆ 2     ┆ 1.1   │
│ B    ┆ null  ┆ 2.0   ┆ 2     ┆ null  │
│ C    ┆ 1     ┆ null  ┆ 2     ┆ null  │
│ D    ┆ null  ┆ 2.0   ┆ null  ┆ 1.1   │
└──────┴───────┴───────┴───────┴───────┘
[Utf8, Int64, Float64, Int64, Float64]
shape: (4, 5)
┌──────┬───────┬───────┬───────┬───────┐
│ 名前 ┆ 整数1 ┆ 小数1 ┆ 整数2 ┆ 小数2 │
│ ---  ┆ ---   ┆ ---   ┆ ---   ┆ ---   │
│ str  ┆ i64   ┆ f64   ┆ i64   ┆ f64   │
╞══════╪═══════╪═══════╪═══════╪═══════╡
│ A    ┆ 1     ┆ 2.0   ┆ 2     ┆ 1.1   │
│ B    ┆ 0     ┆ 2.0   ┆ 2     ┆ 0.0   │
│ C    ┆ 1     ┆ 0.0   ┆ 2     ┆ 0.0   │
│ D    ┆ 0     ┆ 2.0   ┆ 0     ┆ 1.1   │
└──────┴───────┴───────┴───────┴───────┘

まとめ

RustのPolarsでDataframeのデータ型していにはSchemaを使う
Dataframeに対してfill_null関数で0埋めはできないのでSeriesに使う
Seriesに対してはapply_at_idx()やtry_apply_at_idx()を使うとよい