polarsの`.when().then().otherwise()`と`.replace()`の速度比較

# polarsの`.when().then().otherwise()`と`.replace()`の速度比較

## 前提

特定の列の値の置換を行う場合、`pl.when().then()...otherwise()`を使う方法と、`Expr.replace()`を使う方法とがある。

例えば、1から10までの整数が格納されている列を、a ~ j のアルファベットに置換する作業を行うとしよう。

1つ目の方法として、`.when().then()`を地道に続けて書くことができる。

```python
EXPR1 = (
    pl.when(column_0=1)
    .then(pl.lit("a"))
    .when(column_0=2)
    .then(pl.lit("b"))
    .when(column_0=3)
    .then(pl.lit("c"))
    .when(column_0=4)
    .then(pl.lit("d"))
    .when(column_0=5)
    .then(pl.lit("e"))
    .when(column_0=6)
    .then(pl.lit("f"))
    .when(column_0=7)
    .then(pl.lit("g"))
    .when(column_0=8)
    .then(pl.lit("h"))
    .when(column_0=9)
    .then(pl.lit("i"))
    .when(column_0=10)
    .then(pl.lit("j"))
    .otherwise(pl.lit("other"))
)
```

第2の選択肢として、辞書を用意していれば、`.replace()`を使うことができる。

```python
MY_DICT = dict(enumerate("abcdefghij", 1))
EXPR2 = pl.col("column_0").replace(MY_DICT, default="other")
```

どちらで実行したほうが速いのだろうか。

## 比較

上に挙げた例を使って比較してみよう（polarsのバージョンは0.19.19）。

ここでは1～1000万行2列のデータフレームを用意した。全てに0~10の整数が格納されている。2列あるうち1列目の1~10をa~jに変換した列を追加する。0は"other"に変換される。

```python
import benchit
import numpy as np
import polars as pl

# ~~~中略~~~ （変数 EXPR1, EXPR2 は上記参照）

def when_then_otherwise(df: pl.DataFrame):
    return df.with_columns(new_column=EXPR1)


def replace_dict(df: pl.DataFrame):
    return df.with_columns(new_column=EXPR2)

rng = np.random.default_rng(0)
funcs = [when_then_otherwise, replace_dict]
inputs = {n: pl.DataFrame(rng.integers(0, 4, (n, 2))) for n in 10 ** np.arange(8)}
t = benchit.timings(funcs, inputs)
t.plot(figsize=(8, 5), logx=True)
```

![output1.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/394333/cfcd2eaf-72b4-688a-29ac-41987db96b77.png)


これを見ると、データフレームの長さ（行数）が1万行未満では`.when().then()...otherwise()`の方が速いが、1万行以上では`.replace()`の方が速くなるようだ。

## 置換パターンを減らして比較

上の例では数字からアルファベットへの置換パターンが10種類あった。これを3種類に減らしたバージョンを作ってみる。すなわち、1~3をa~cに置換する。

```python
EXPR3 = (
    pl.when(column_0=1)
    .then(pl.lit("a"))
    .when(column_0=2)
    .then(pl.lit("b"))
    .when(column_0=3)
    .then(pl.lit("c"))
    .otherwise(pl.lit("other"))
)

MY_DICT2 = dict(enumerate("abc", 1))
EXPR4 = pl.col("column_0").replace(MY_DICT2, default="other")
```

先の10方向置換パターンに、この3方向置換パターンを加えて比較してみよう。

```python
def when_then_otherwise_v3(df: pl.DataFrame):
    return df.with_columns(new_column=EXPR3)


def replace_dict_v3(df: pl.DataFrame):
    return df.with_columns(new_column=EXPR4)


funcs = [when_then_otherwise, replace_dict, when_then_otherwise_v3, replace_dict_v3]
inputs = {n: pl.DataFrame(rng.integers(0, 10, (n, 2))) for n in 10 ** np.arange(8)}
t = benchit.timings(funcs, inputs)
t.plot(figsize=(8, 5), logx=True)
```

![output2.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/394333/ba237725-f8d9-677a-e704-2def58316f61.png)

`.when().then()...otherwise()`は速度が向上し（真下に平行移動したような形だ）、`.replace()`はほとんど変わらなかった（1万行以上で気持ち向上している）。それでも、10万行くらいをこえると`.when().then()...otherwise()`の速度を`.replace()`が上回っている。

### 2024/06/18追記

半年経過したので、バージョン0.19.19から0.20.31にあげて再検証してみた。

![output3.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/394333/b1bd50ab-e279-6c6e-1f74-f847b4bcf023.png)

全体的に速度が速くなっているが、今回の比較の観点で言えば、`.when().then()...otherwise()`がさらに有利になったようである。

## 辞書を`.when().then()...otherwise()`に渡す

そうなると、データフレームの大きさが数万行程度なら、置換パターンの辞書オブジェクトを持っていても、それを直接`.replace()`に渡すよりも、何らかの方法で`.when().then()...otherwise()`を使ったほうが処理は速いということになる。

しかし、辞書を`.when().then()...otherwise()`タイプのエクスプレッションに渡すのは結構めんどくさそうだ。

```python:例えば
expr_tmp = pl
for key, value in MY_DICT.items():
    expr_tmp = expr_tmp.when(column_0=key).then(pl.lit(value))
EXPR = expr_tmp.otherwise("other")
```

↑何をやっているのか分かりづらい。

## 結論

- 数万行以下ではたいてい`.when().then()...otherwise()`の方が`.replace()`よりも速い。
- `.when().then().when().then()....`は続けば続くほど（置換パターンが多いほど）ぐんぐん遅くなる。
- `.replace()`の速度は辞書の長さ（置換パターンの多さ）の影響をほとんど受けない。

それだけ。