Python: Polars と NumPy / SciPy の関数を組み合わせて使う

Polars を使って数値を加工しようとすると、数学に関する API がさほど多くないことに気づく。そうしたときに、最初に思いつくのは Series オブジェクトを NumPy 配列に変換した上で処理する方法かもしれない。しかし、実際には Polars の Expr オブジェクトと NumPy の ufunc (ユニバーサル関数) にはインテグレーションが提供されている。今回は Polars と NumPy や SciPy を組み合わせて使う方法について書いてみる。

使った環境は次のとおり。

$ sw_vers    
ProductName:        macOS
ProductVersion:     13.2.1
BuildVersion:       22D68
$ python -V            
Python 3.9.16
$ pip list | egrep "(polars|numpy|scipy)"  
numpy             1.24.2
polars            0.16.14
scipy             1.10.1

下準備

あらかじめ Polars と、オプションで利用できる NumPy をインストールする。また、ついでに SciPy も入れておく。

$ pip install "polars[numpy]" scipy

Python のインタプリタを起動する。

$ python

サンプルとして x という名前のカラムを持った Polars の DataFrame を用意する。

>>> import polars as pl
>>> df = pl.DataFrame({"x": [-1, 0, 1, 2]})
>>> df
shape: (4, 1)
┌─────┐
│ x   │
│ --- │
│ i64 │
╞═════╡
│ -1  │
│ 0   │
│ 1   │
│ 2   │
└─────┘

Polars と NumPy を組み合わせて使う

たとえば x カラムの値について、絶対値を取って log1p したいとする。これはあくまで例であって、処理自体に何か意味があるわけではない。

実のところ、この処理は pl.col("x") から得られる Expr オブジェクトを np.abs() と np.log1p() の引数にするだけで実現できる。なんと、これで返り値として Polars の Expr オブジェクトが得られる。

>>> import numpy as np
>>> np.log1p(np.abs(pl.col("x")))
<polars.expr.expr.Expr object at 0x1018c5ee0>

あとは、この Expr オブジェクトを DataFrame#select() とかに突っ込めばいい。

>>> df.select(np.log1p(np.abs(pl.col("x"))))
shape: (4, 1)
┌──────────┐
│ x        │
│ ---      │
│ f64      │
╞══════════╡
│ 0.693147 │
│ 0.0      │
│ 0.693147 │
│ 1.098612 │
└──────────┘

Polars と SciPy を組み合わせて使う

また、公式にサポートは名言されていないようだけど、実は SciPy の ufunc でも同じことができるようだ。

たとえば SciPy の sp.special.expit() で試してみる。

>>> import scipy as sp
>>> sp.special.expit(pl.col("x"))
<polars.expr.expr.Expr object at 0x103180fa0>

ちゃんと Expr オブジェクトが返ってくる。

先ほどと同じように DataFrame#select() に突っ込んでみる。

>>> df.select(sp.special.expit(pl.col("x")))
shape: (4, 1)
┌──────────┐
│ x        │
│ ---      │
│ f64      │
╞══════════╡
│ 0.268941 │
│ 0.5      │
│ 0.731059 │
│ 0.880797 │
└──────────┘

うまくいっているようだ。