Rustのデータフレームcrateのpolarsとpandasの比較

published: 2021/1/9 update: 2021/1/9

TL;DR

rustにも実はpandas likeなcrateがあることを知ったのでpandasとの対応関係をまとめてた。最善である保証はありません。

これを使えば大きなファイルを素早く処理することができそうですが、さすがにrustなのでお手軽感はそんなになかった。

plotterとかはjupyter対応してるので、もしかしてそのうちjupyterで使える可能性に期待しています。間違い等があればTwitterやメールにお願いします。

polars

Apache Arrowsをベースにしたデータフレームライブラリ。なんかpy-polarsみたいなのもあって、pandasより速いらしい。polarsのgithubのREADMEにベンチマークテストがある。使い勝手としてはどちらかといえばRのtidyverseに似ている気がする。

ChunkedArray

多分特徴的なのが、Seriesから変換できるChunkedArrayという構造体を持つ点。ChunkedArrayは型があるので様々な演算ができる。また、条件をつかった列選択を行う際にはChunckedArray<BooleanType>を使う必要がある。

Install

色々featureもあって、並列実行やndarrayへの変換、ランダムサンプリングなどに対応している。あとはjsonのserdeやApache Parquet formatとかのIOとか。今回はndarrayと並列実行、ランダムサンプリングなどを試してみる。あとエラーハンドリングにanyhowを入れておく。

[dependencies]
anyhow = "1.0"
polars = { version = "0.10", features = ["ndarray", "parallel", "random"]}

また、nightlyが必要なのでOverrideしておく。

rustup override nightly

pandasはお好みのパッケージ管理ツールでインストールしてください。

rust側は下記のtodo!()部分に相当する場所を書いているつもりです。

use polars::prelude::*;
use anyhow::Result;

fn main() -> Result<()> {
    todo!();
    Ok(())
}

python側も下記のimportを行っている前提です。

import pandas as pd

SeriesとDataFrameとChunkedArrayの演算

非常に長いので畳んである。ChunkedArrayは大抵の演算ができる。Seriesの比較は条件による行選択の際に必要になってくるので見ておくとよいと思い気がする。

c1.lt(c2);

DataFrameの作成

df = pd.DataFrame({
    "A": ["a", "b", "a"],
    "B": [1, 3, 5],
    "C": [10, 11, 12],
    "D": [2, 4, 6]
})

マクロが便利

let mut df = df!("A" => &["a", "b", "a"],
             "B" => &[1, 3, 5],
             "C" => &[10, 11, 12],
             "D" => &[2, 4, 6]
    )?;

列選択

df["A"]
df[["A", "B"]]

selectで選ぶと、Result<DataFrame>が返ってくる。

df.select("A")?;
df.select(("A", "B"))?;
df.select(vec!["A", "B", "C"])?;

columnで選ぶと、Result<Series>が返ってくる。

df.column("A")?;

条件に応じた列選択

どちらもcolumnsをとってきてfilterなりなんなりをすればよい。多分strメソッドを使った方がpandasっぽい？ rustはget_columnsでcolumnsをもって来れる。もう少し何とかならないかな...

df.loc[:, [c.startswith("A") for c in df.columns]]
df.loc[:, df.columns.str.startswith("A")]

df.select(&df.get_column_names()
            .iter()
            .filter(|x| x.starts_with("A"))
            .map(|&x| x)
            .collect::<Vec<&str>>()
        )?;

列の入れ替え

df[["B", "A"]]

df.select(("B", "A"))?;

列追加

df["E"] = df["B"] * 2
# or
df["E"] = df["B"].map(lambda x: x * 2)
# or
df.assign(E = lambda df: df.B * 2)

polarsのcolumの追加はadd_column関数で行える。
assignみたいないい感じの関数が見つからなかった。四則演算や簡単な演算はSeriesにして計算すればいける。to_owned()二回やってるの解消できる気がするけどできなかった。
無名関数を使いたい際には、一端ChunkedArrayに変換してからapplyやmapを使う。Seriesは型を持たないが、ChunkedArrayは型があるので演算ができる。
DataFrame構造体にはapplyが存在しているが、&mut selfなので、本体が変わってしまう。なのでselectかcloneしてからみたいな処理になるけどどっちが早いのだろうか。

df.add_column(
    Series::new("E", &[2, 2, 2])
)?;

df.add_column(
    df.column("B")?
    .to_owned()
    .rename("E")
    .to_owned() * 2)?;
// or
   df.add_column(df.column("B")?.i32()?
    .apply(|x| x * 2 )
    .into_series()
    .rename("E")
    .to_owned())?;
// or
df.add_column(df.clone() // or df.select("B")
    .rename("B", "E")?
    .apply("E", |x| x * 2)?
    .column("E")?
    .to_owned())?;

また、新しいDataFrameを作りたければwith_columnを使う。中身の書き方はadd_columnと同じ。

let new_df = df.with_column(df.column("B")?.i32()?
    .apply(|x| x * 2 )
    .into_series()
    .rename("E").to_owned())?;

条件による行選択

単独条件

df.loc[:, df["B"] <= 4]
df.query("B <= 4")

df.filter(&df.column(B)?.lt_eq(4))?;

複数条件

df.loc[:, (df["B"] == 1) | (df["C"] == 12)]
df.query("B == 1 | C == 12")

ChunkedArrayはbit演算ができます。

df.filter(&(
    df.column("B")?.eq(1)? | df.column("C").eq(12)?
))

含まれているかなどの演算

l = [1, 3]
df.query("B in l")

たぶんChunkedArrayに変換してやるしかない、と思う。applyはSelfを返すので、ChunkedArray<Int32Type>からChunkedArray<BooleanType>に変換はできない。なので、mapを使った後collectする必要がある。

let v: Vec<i32> = vec![1, 2];
df.filter(&(
    df.column("B")?
        .i32()?
        .map(|x| v.contains(&x))?
        .collect()
))?;

重複行の抽出

df.loc[df.is_duplicates()]

df.filter(&df.is_duplicated()?)?;

重複行の削除

両方ともsubsetを選ぶことで、同じように特定の列の重複行を削除することができる。

df.drop_duplicates()

df.drop_duplicates(true, None)? // maintain_order, subset;

numpy / ndarrayへの変換

df.values()

df.to_ndarray()?;

read csv

csv以外ならsep = "\t"とかしてください。

df = pd.read_csv(path)

csv以外ならwith_delimiterの引数を好きに変えてください。なくても動きます。

あとparalellのfeatureがあると、daskみたいな感じでCPUの上限コア数を使って読み込みます。いやな場合は、.with_n_threads(Some(2))とかしてください。with_n_threadsはfrom_pathを使ってCsvReaderを作った時しか使えないようです。

let df = CsvReader::from_path(path)?
        .infer_schema(None)
        .with_delimiter(b',')
        .has_header(true)
        .finish()?

write csv

readと同様。

df.to_csv(path)

let mut f = std::fs::File::create(path)?;
CsvWriter::new(&mut f)
    .has_headers(true)
    .with_delimiter(b',')
    .finish(df)?;

TODO

• groupby
• pyvot
• melt
• vstack
• join系
• fillna系
• sample_n
• macro

気長に埋めていきます。

感想

できることは多い感じがします。Pandasみたいに柔軟に処理する分にはあまり使えなさそうですが、決まりきった処理ならこちらで記述すると生産効率が上がりそうな気がします。