Python: scikit-learn は v1.4 から Polars をサポートした

scikit-learn に組み込みで用意されている Transformer は、長らく入出力として NumPy 配列にしか対応していなかった。その状況が変わったのは v1.2 以降で、Pandas の DataFrame を扱えるようになった。

blog.amedama.jp

そして v1.4 からは、ついに Polars の DataFrame もサポートされた。今回は、実際にその機能を試してみよう。

scikit-learn.org

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

$ sw_vers    
ProductName:        macOS
ProductVersion:     14.5
BuildVersion:       23F79
$ python -V                      
Python 3.11.9
$ pip list | egrep "(scikit-learn|polars)"
polars        0.20.27
scikit-learn  1.5.0

下準備

まずは scikit-learn と Polars をインストールする。

$ pip install scikit-learn polars

インストールできたら、次に Python のインタプリタを起動する。

$ python

続いて、動作確認用のデータを用意する。今回は Iris データセットにした。ひとまず NumPy 配列として読み込む。

>>> from sklearn.datasets import load_iris
>>> X, _ = load_iris(as_frame=False, return_X_y=True)
>>> X[:5]
array([[5.1, 3.5, 1.4, 0.2],
       [4.9, 3. , 1.4, 0.2],
       [4.7, 3.2, 1.3, 0.2],
       [4.6, 3.1, 1.5, 0.2],
       [5. , 3.6, 1.4, 0.2]])

続いて、適当な Transformer を用意する。今回は StandardScaler にした。

>>> from sklearn.preprocessing import StandardScaler
>>> scaler = StandardScaler()

ひとまず、用意した StandardScaler のインスタンスを使ってデータを変換する。すると、デフォルトでは NumPy 配列として結果が返ってくる。

>>> X_scaled = scaler.fit_transform(X)
>>> X_scaled[:5]
array([[-0.90068117,  1.01900435, -1.34022653, -1.3154443 ],
       [-1.14301691, -0.13197948, -1.34022653, -1.3154443 ],
       [-1.38535265,  0.32841405, -1.39706395, -1.3154443 ],
       [-1.50652052,  0.09821729, -1.2833891 , -1.3154443 ],
       [-1.02184904,  1.24920112, -1.34022653, -1.3154443 ]])

この振る舞いは従来から変わっていない。

出力を Polars の DataFrame にする

それでは、次に出力を Polars の DataFrame にしてみよう。

具体的には、先ほど用意した StandardScaler のインスタンスの set_output() メソッドを呼び出す。このとき、引数の transform に "polars" を指定する。

>>> scaler.set_output(transform="polars")

この状況で、もう一度データを変換してみよう。

>>> X_scaled = scaler.fit_transform(X)

すると、返ってくるのが Polars の DataFrame になる。

>>> type(X_scaled)
<class 'polars.dataframe.frame.DataFrame'>
>>> X_scaled
shape: (150, 4)
┌───────────┬───────────┬───────────┬───────────┐
│ x0        ┆ x1        ┆ x2        ┆ x3        │
│ ---       ┆ ---       ┆ ---       ┆ ---       │
│ f64       ┆ f64       ┆ f64       ┆ f64       │
╞═══════════╪═══════════╪═══════════╪═══════════╡
│ -0.900681 ┆ 1.019004  ┆ -1.340227 ┆ -1.315444 │
│ -1.143017 ┆ -0.131979 ┆ -1.340227 ┆ -1.315444 │
│ -1.385353 ┆ 0.328414  ┆ -1.397064 ┆ -1.315444 │
│ -1.506521 ┆ 0.098217  ┆ -1.283389 ┆ -1.315444 │
│ -1.021849 ┆ 1.249201  ┆ -1.340227 ┆ -1.315444 │
│ …         ┆ …         ┆ …         ┆ …         │
│ 1.038005  ┆ -0.131979 ┆ 0.819596  ┆ 1.448832  │
│ 0.553333  ┆ -1.282963 ┆ 0.705921  ┆ 0.922303  │
│ 0.795669  ┆ -0.131979 ┆ 0.819596  ┆ 1.053935  │
│ 0.432165  ┆ 0.788808  ┆ 0.933271  ┆ 1.448832  │
│ 0.068662  ┆ -0.131979 ┆ 0.762758  ┆ 0.790671  │
└───────────┴───────────┴───────────┴───────────┘

なお、入力が NumPy 配列なので、カラム名は x0, x1, x2, ... という命名規則で自動的に与えられている。

入力として Polars の DataFrame を与えてみる

次は入力として Polars の DataFrame を与えてみるパターンも試しておこう。

NumPy 配列を元にして Polars の DataFrame を作成する。

>>> import polars as pl
>>> col_names = {
...     "sepal_length",
...     "sepal_width",
...     "petal_length",
...     "petal_width",
... }
>>> df = pl.DataFrame(data=X, schema=col_names)

作った DataFrame を StandardScaler のインスタンスに渡す。すると、結果が先ほどと同じように Polars の DataFrame として得られる。

>>> scaler.fit_transform(df)
shape: (150, 4)
┌──────────────┬─────────────┬──────────────┬─────────────┐
│ sepal_length ┆ petal_width ┆ petal_length ┆ sepal_width │
│ ---          ┆ ---         ┆ ---          ┆ ---         │
│ f64          ┆ f64         ┆ f64          ┆ f64         │
╞══════════════╪═════════════╪══════════════╪═════════════╡
│ -0.900681    ┆ 1.019004    ┆ -1.340227    ┆ -1.315444   │
│ -1.143017    ┆ -0.131979   ┆ -1.340227    ┆ -1.315444   │
│ -1.385353    ┆ 0.328414    ┆ -1.397064    ┆ -1.315444   │
│ -1.506521    ┆ 0.098217    ┆ -1.283389    ┆ -1.315444   │
│ -1.021849    ┆ 1.249201    ┆ -1.340227    ┆ -1.315444   │
│ …            ┆ …           ┆ …            ┆ …           │
│ 1.038005     ┆ -0.131979   ┆ 0.819596     ┆ 1.448832    │
│ 0.553333     ┆ -1.282963   ┆ 0.705921     ┆ 0.922303    │
│ 0.795669     ┆ -0.131979   ┆ 0.819596     ┆ 1.053935    │
│ 0.432165     ┆ 0.788808    ┆ 0.933271     ┆ 1.448832    │
│ 0.068662     ┆ -0.131979   ┆ 0.762758     ┆ 0.790671    │
└──────────────┴─────────────┴──────────────┴─────────────┘

ただし、今回は元の DataFrame にカラム名があるので、それがそのまま使われている。

常に出力を Polars の DataFrame にしたい

ここまでの例は、個別のインスタンスにおいて出力の形式を指定するやり方だった。一方で、実際に使う場合にはデフォルトの形式を変更したいことは多いはず。

そんなときは sklearn.set_config() 関数が使える。引数の transform_output に "polars" を指定しておけばデフォルトの形式が Polars の DataFrame になる。

>>> from sklearn import set_config
>>> set_config(transform_output="polars")

上記を実行した後で、あらためて StandardScaler のインスタンスを作成しよう。

>>> scaler = StandardScaler()

そしてデータを変換させてみると、今度は最初から Polars の DataFrame が返ってくる。

>>> scaler.fit_transform(X)
shape: (150, 4)
┌───────────┬───────────┬───────────┬───────────┐
│ x0        ┆ x1        ┆ x2        ┆ x3        │
│ ---       ┆ ---       ┆ ---       ┆ ---       │
│ f64       ┆ f64       ┆ f64       ┆ f64       │
╞═══════════╪═══════════╪═══════════╪═══════════╡
│ -0.900681 ┆ 1.019004  ┆ -1.340227 ┆ -1.315444 │
│ -1.143017 ┆ -0.131979 ┆ -1.340227 ┆ -1.315444 │
│ -1.385353 ┆ 0.328414  ┆ -1.397064 ┆ -1.315444 │
│ -1.506521 ┆ 0.098217  ┆ -1.283389 ┆ -1.315444 │
│ -1.021849 ┆ 1.249201  ┆ -1.340227 ┆ -1.315444 │
│ …         ┆ …         ┆ …         ┆ …         │
│ 1.038005  ┆ -0.131979 ┆ 0.819596  ┆ 1.448832  │
│ 0.553333  ┆ -1.282963 ┆ 0.705921  ┆ 0.922303  │
│ 0.795669  ┆ -0.131979 ┆ 0.819596  ┆ 1.053935  │
│ 0.432165  ┆ 0.788808  ┆ 0.933271  ┆ 1.448832  │
│ 0.068662  ┆ -0.131979 ┆ 0.762758  ┆ 0.790671  │
└───────────┴───────────┴───────────┴───────────┘