Kaggle などのデータ分析コンペで使われるテクニックのひとつに reduce_mem_usage() 関数がある。
これは、一般に pandas の DataFrame のメモリ使用量を削減するために用いられる。
具体的には、カラムに出現する値を調べて、それを表現する上で必要最低限な型にキャストする。
たとえば、64 ビット整数のカラムを 32 ビット整数にできれば、理屈の上では必要なメモリ使用量がおよそ半分になる。
ただし、この関数は pandas が組み込みで提供しているわけではない。 そのため、各々がスニペットを秘伝のタレのように持っているか、あるいは必要に応じてウェブ上から探して利用する場合が多いはず。
一方で、Polars にはこれに相当する機能が組み込みで用意されている。
具体的には polars.Expr#shrink_dtype() という Expr オブジェクトを返すメソッドがある。
あまり知られていないようなので、今回は紹介してみる。
使った環境は次のとおり。
$ sw_vers ProductName: macOS ProductVersion: 14.6 BuildVersion: 23G80 $ python -V Python 3.11.9 $ pip list | grep polars polars 1.4.0
もくじ
下準備
まずは Polars をインストールしておく。
$ pip install polars
適当なデータセットとして Diamonds データセットをダウンロードする。
$ wget https://raw.githubusercontent.com/mwaskom/seaborn-data/master/diamonds.csv
Python のインタプリタを起動する。
$ python
CSV ファイルを読み込んで DataFrame オブジェクトを作る。 そのままだとサイズが分かりにくいので 200 個ほど連結してサイズをかさ増しする。
>>> import polars as pl >>> raw_df = pl.read_csv("diamonds.csv") >>> df = pl.concat([raw_df for _ in range(200)]) >>> df shape: (10_788_000, 10) ┌───────┬───────────┬───────┬─────────┬───┬───────┬──────┬──────┬──────┐ │ carat ┆ cut ┆ color ┆ clarity ┆ … ┆ price ┆ x ┆ y ┆ z │ │ --- ┆ --- ┆ --- ┆ --- ┆ ┆ --- ┆ --- ┆ --- ┆ --- │ │ f64 ┆ str ┆ str ┆ str ┆ ┆ i64 ┆ f64 ┆ f64 ┆ f64 │ ╞═══════╪═══════════╪═══════╪═════════╪═══╪═══════╪══════╪══════╪══════╡ │ 0.23 ┆ Ideal ┆ E ┆ SI2 ┆ … ┆ 326 ┆ 3.95 ┆ 3.98 ┆ 2.43 │ │ 0.21 ┆ Premium ┆ E ┆ SI1 ┆ … ┆ 326 ┆ 3.89 ┆ 3.84 ┆ 2.31 │ │ 0.23 ┆ Good ┆ E ┆ VS1 ┆ … ┆ 327 ┆ 4.05 ┆ 4.07 ┆ 2.31 │ │ 0.29 ┆ Premium ┆ I ┆ VS2 ┆ … ┆ 334 ┆ 4.2 ┆ 4.23 ┆ 2.63 │ │ 0.31 ┆ Good ┆ J ┆ SI2 ┆ … ┆ 335 ┆ 4.34 ┆ 4.35 ┆ 2.75 │ │ … ┆ … ┆ … ┆ … ┆ … ┆ … ┆ … ┆ … ┆ … │ │ 0.72 ┆ Ideal ┆ D ┆ SI1 ┆ … ┆ 2757 ┆ 5.75 ┆ 5.76 ┆ 3.5 │ │ 0.72 ┆ Good ┆ D ┆ SI1 ┆ … ┆ 2757 ┆ 5.69 ┆ 5.75 ┆ 3.61 │ │ 0.7 ┆ Very Good ┆ D ┆ SI1 ┆ … ┆ 2757 ┆ 5.66 ┆ 5.68 ┆ 3.56 │ │ 0.86 ┆ Premium ┆ H ┆ SI2 ┆ … ┆ 2757 ┆ 6.15 ┆ 6.12 ┆ 3.74 │ │ 0.75 ┆ Ideal ┆ D ┆ SI2 ┆ … ┆ 2757 ┆ 5.83 ┆ 5.87 ┆ 3.64 │ └───────┴───────────┴───────┴─────────┴───┴───────┴──────┴──────┴──────┘
特に指定しない場合、デフォルトでは数値のカラムが 64 ビットの型になる。
>>> df.dtypes [Float64, String, String, String, Float64, Float64, Int64, Float64, Float64, Float64]
この状況では DataFrame のサイズは約 683MB だった。
>>> df.estimated_size(unit="mb") 683.1520080566406
使用メモリを削減する
それでは、実際に polars.Expr.shrink_dtype() を使って DataFrame の使用メモリを減らしてみよう。
すべてのカラムを処理の対象とする場合には pl.all().shrink_dtype() とすれば良い。
あとは得られた Expr オブジェクトを DataFrame#select() に渡せば DataFrame が変換される。
>>> shrinked_df = df.select(pl.all().shrink_dtype()) >>> shrinked_df shape: (10_788_000, 10) ┌───────┬───────────┬───────┬─────────┬───┬───────┬──────┬──────┬──────┐ │ carat ┆ cut ┆ color ┆ clarity ┆ … ┆ price ┆ x ┆ y ┆ z │ │ --- ┆ --- ┆ --- ┆ --- ┆ ┆ --- ┆ --- ┆ --- ┆ --- │ │ f32 ┆ str ┆ str ┆ str ┆ ┆ i16 ┆ f32 ┆ f32 ┆ f32 │ ╞═══════╪═══════════╪═══════╪═════════╪═══╪═══════╪══════╪══════╪══════╡ │ 0.23 ┆ Ideal ┆ E ┆ SI2 ┆ … ┆ 326 ┆ 3.95 ┆ 3.98 ┆ 2.43 │ │ 0.21 ┆ Premium ┆ E ┆ SI1 ┆ … ┆ 326 ┆ 3.89 ┆ 3.84 ┆ 2.31 │ │ 0.23 ┆ Good ┆ E ┆ VS1 ┆ … ┆ 327 ┆ 4.05 ┆ 4.07 ┆ 2.31 │ │ 0.29 ┆ Premium ┆ I ┆ VS2 ┆ … ┆ 334 ┆ 4.2 ┆ 4.23 ┆ 2.63 │ │ 0.31 ┆ Good ┆ J ┆ SI2 ┆ … ┆ 335 ┆ 4.34 ┆ 4.35 ┆ 2.75 │ │ … ┆ … ┆ … ┆ … ┆ … ┆ … ┆ … ┆ … ┆ … │ │ 0.72 ┆ Ideal ┆ D ┆ SI1 ┆ … ┆ 2757 ┆ 5.75 ┆ 5.76 ┆ 3.5 │ │ 0.72 ┆ Good ┆ D ┆ SI1 ┆ … ┆ 2757 ┆ 5.69 ┆ 5.75 ┆ 3.61 │ │ 0.7 ┆ Very Good ┆ D ┆ SI1 ┆ … ┆ 2757 ┆ 5.66 ┆ 5.68 ┆ 3.56 │ │ 0.86 ┆ Premium ┆ H ┆ SI2 ┆ … ┆ 2757 ┆ 6.15 ┆ 6.12 ┆ 3.74 │ │ 0.75 ┆ Ideal ┆ D ┆ SI2 ┆ … ┆ 2757 ┆ 5.83 ┆ 5.87 ┆ 3.64 │ └───────┴───────────┴───────┴─────────┴───┴───────┴──────┴──────┴──────┘
上記を見ても分かるとおり、行数や列数などは何も変わっていない。 ただし、カラムの型は必要最低限なものにキャストされている。
>>> shrinked_df.dtypes [Float32, String, String, String, Float32, Float32, Int16, Float32, Float32, Float32]
処理後の DataFrame は、使用メモリが約 374MB まで減っている。
>>> shrinked_df.estimated_size(unit="mb") 374.5048522949219
いじょう。
まとめ
- Polars には
reduce_mem_usage()関数に相当する機能が組み込みで用意されている polars.Expr.shrink_dtype()というExprオブジェクトを返す関数を使う
