ibis backend でPolarsを使うときにLazy実行で最適化されてるの？

## はじめに

[ibis](https://ibis-project.org/) という Python におけるバックエンドフリー（フリーとは依存しないの意味）なdata manipulation ライブラリーがある。R言語でいうdplyer(dbplyr)のPython版ともいえるだろう。 
なお、開発は2015年ごろからはじまったが、スポンサーもついたようで直近1年でのリリース頻度が非常に上がっている。

## ibis の特徴

ibisは、Pandas DataFrame, SparkやDaskなどの分散処理系,  RMDBS系(SQLエンジン)バックエンドを同一APIで操作可能とするパッケージである。なお、インストールは `pip install ibis-framework` であり `pip install ibis` とは異なるワナがあることに注意されたい。

現在Polarsバックエンドはexperimentalな位置づけである。Polarsの破壊的変更へのキャッチアップは若干遅れを取るが、込み入った操作でなければほぼ対応できているようだ。詳しくは[サポート対応表](https://ibis-project.org/support_matrix)で確認できる。

さて、ibisの利点は、今からPolars操作のためのAPIやmethodを覚えるならば、代わりにibis操作を覚えておけば、約18のバックエンドでほぼpython側のコード変更をなしで対応できることである。私はEDA（探索的データ分析）では 長年のPandas利用からduckdb, polarsの若干の利用を経て、現在は ibis を使い始めている。余談だが [ibis-dbt](https://binste.github.io/dbt-ibis/intro.html)を使うとdbtにおけるSQL dialect（方言）の違いを吸収できるので別途関心があるところだ。


## この記事で確認したいこと

Polarsのlazy APIでは[optimizations](https://pola-rs.github.io/polars/user-guide/lazy/optimizations/) が効くが、ibisのバックエンドにPolarsを指定した場合、この optimizations が効くのだろうか確認したい。

具体的にはpredicate pushdown(いわゆるfilter)とprojection pushdown(いわゆるselect)がibisの実行で効果があるか確認する。その他の最適化の確認はスコープ外とする（余力があれば別記事で検証したい）。なお、Polarsの lazy query について [query plan](https://pola-rs.github.io/polars/user-guide/lazy/query-plan/)のページを事前に読んでおくとさらに理解が深まるだろう。


## ibis の注意点

ibis は、デフォルトで lazy mode である。eagerとlazyの切り替えは `ibis.options.interactive = True # False` にておこなうことができる。

## Prerequisite

投稿時点の最新版をインストールする。

```bash
pip install ibis-framework[pandas,polars,duckdb]==7.1.0
pip install polars==0.19.19
```

なお、show_graph()での描画の実行には Graphviz が必要であるので、別途インストールしておく。

```sh
# homebrew or linuxbrew (wsl2)
brew install graphviz
```

### サンプルデータ

サンプルデータは何でもよいが、公式repositoryに存在するcsvを（2つ）使う。


```python
def get_food_csv(n: int) -> tuple[str, str]:
    # exist 1..5
    assert n in range(1, 6)
    return f"https://raw.githubusercontent.com/pola-rs/polars/main/examples/datasets/foods{n}.csv", f"foods{n}.csv"

# use only two files
for i in range(1, 3):
    url, filename = get_food_csv(i)
    urllib.request.urlretrieve(url, filename)
```

### import

ibisの特徴として `_` を直前のTableの代名詞として扱える。Polarsでいう `pl.col("col1")` は `_.col1` と記述できる。また今回は使わないが、 `select` 句での selectorで列指定が可能である。これは Polars の [polars.selectors as cs](https://pola-rs.github.io/polars/user-guide/expressions/column-selections/#using-selectors) と類似しているので馴染みがあるかもしれない。

```python
import polars as pl
import ibis
from ibis import _
# import ibis.selectors as s # 今回は未使用
```

## Polars での実行の確認

Polarsでの LazyFrame における最適化の実行計画は `print(lf.explain())` によるtext表現、 `lf.show_graph(optimized=True)` による画像で確認できる。
(N.B. `lf` は LayzFrame のオブジェクト名としてサンプルコードで使われることが多い）

さて、操作した結果自体に意味はないが、次のコードを実行してみよう。そしてその結果について次で解説する。
```python
q1 = (
    pl.scan_csv("foods*.csv")
    .filter(pl.col("category") != "meat")
    .group_by("category")
    .agg(
        # `alias()` を使わなくても書ける
        sugar_count=pl.col("sugars_g").count(),
        fat_mean=pl.col("fats_g").mean(),
    )
    .filter(pl.col("sugar_count") > 0)
    .sort(pl.col("fat_mean"), descending=True)
    .with_columns(
        fat_mean_kg=pl.col("fat_mean") / 1000
    )

)
print(q1.explain())

q1.collect()
```

#### explain() の出力の読み方。

- `explain()`の実行は下から上に読んでいく。なお結果はstr文字列で返されるが改行文字もそのままなので `print()` をかませる必要がある。
- `FROM` の区切りで処理が分けられるので分割されたブロックごとに解釈していく。
  -  最初にcsvファイルを二つ読んでUNIONしているこの時点で二つの最適化が行われている。なお、`show_graph()` の実行表記の `σ` が selection (sのギリシャ文字)、 `π` が projection（pのギリシャ文字）で表現されている。 `π 3/4` とは全4列から3列が選択されていると読む。
      1. poject: 必要な列だけの取得
      2. selectin: 条件に適する行だけ取得
  - 次に `GROUP BY` による `AGGREGATE` を実行し
  - 最後に`SORT BY`, `FILTER` および `WITH_COLUMNS` を実行している。


```sql
 WITH_COLUMNS:
 [[(col("fat_mean")) / (1000.0)].alias("sugar_mean_kg")]
  SORT BY [col("fat_mean")]
    FILTER [(col("sugar_count")) > (0)] FROM

    AGGREGATE
    	[col("sugars_g").count().alias("sugar_count"), col("fats_g").mean().alias("fat_mean")] BY [col("category")] FROM
      UNION
        PLAN 0:
          FAST_PROJECT: [sugars_g, fats_g, category]

              Csv SCAN foods1.csv
              PROJECT 3/4 COLUMNS
              SELECTION: [(col("category")) != (Utf8(meat))]
        PLAN 1:
          FAST_PROJECT: [sugars_g, fats_g, category]

              Csv SCAN foods2.csv
              PROJECT 3/4 COLUMNS
              SELECTION: [(col("category")) != (Utf8(meat))]
      END UNION
```

## ibisでの実行

次に ibis経由での操作を見てみよう

```python
# backendの指定。デフォルトは "duckdb" なので変更が必要
ibis.set_backend("polars")

q2 = (
    ibis.read_csv("foods*.csv")
    .filter(_.category != "meat")
    .group_by("category")
    .agg(
        sugar_count=_.sugars_g.count(),
        fat_mean=_.fats_g.mean()
    )
    .filter(_.sugar_count > 0)
    .order_by(_.fat_mean.desc())
    .mutate(
        fat_mean_kg=_.fat_mean / 1000
    )
)
print(q2.compile().explain())

q2.execute()

# 値は一致
q1.collect().to_pandas().compare(q2.execute())
```

最適化の確認の前に、ibisとの違いを簡単に説明しよう。

1. 各処理のメソッドを見てみると、methodの意味はほぼおなじ。違っていても推測できるレベルで読める。Polarsユーザーならばibisコードは初見でほぼ読める。書くのもdocument参照すればほぼ問題ない。なお若干 `window関数` 系の操作がibisは冗長である。
2. ibis操作での変数 `q2` は `ibis.expr.types.relations.Table` オブジェクトだが、 backend="polars" で `compile()` メソッドの実行でPolarsオブジェクトに変換される。すなわち `type(q2.compile())` の結果は `polars.lazyframe.frame.LazyFrame` となる。
3. ibis の `q2.execute()` の結果は **Pandas** Dataframe を返す。**Polars** DataFrame を返したい場合は `q2.compile().collect()` を実行する。
4. Int の cast が変わりうる。致命的ではないが pandas Dataframe.equals() では False となってしまう。
    ```python
    q1.collect().dtypes
    # [Utf8, UInt32, Float64, Float64]
    
    q2.compile().collect().dtypes
    # [Utf8, Int64, Float64, Float64]
    ```

次に `explain()` の結果を見よう。

```sql
 SELECT [col("category"), col("sugar_count"), col("fat_mean"), [(col("fat_mean")) / (1000.0)].alias("fat_mean_kg")] FROM
  SORT BY [col("fat_mean")]
    FILTER [(col("sugar_count")) > (0)] FROM

    AGGREGATE
    	[col("sugars_g").filter(col("sugars_g").is_not_null()).count().strict_cast(Int64).alias("sugar_count").alias("sugar_count"), col("fats_g").filter(col("fats_g").is_not_null()).mean().strict_cast(Float64).alias("fat_mean").alias("fat_mean")] BY [col("category")] FROM
      UNION
        PLAN 0:
          FAST_PROJECT: [category, sugars_g, fats_g]

              Csv SCAN /tmp/foods1.csv
              PROJECT 3/4 COLUMNS
              SELECTION: [(col("category")) != (Utf8(meat))]
        PLAN 1:
          FAST_PROJECT: [category, sugars_g, fats_g]

              Csv SCAN /tmp/foods2.csv
              PROJECT 3/4 COLUMNS
              SELECTION: [(col("category")) != (Utf8(meat))]
      END UNION
```

最適化である projection と selection ば同様に機能している。本記事の目的は達成できた。

一方で、次の違いがある
- ibisの mutate は `WITH_COLUMNS` に翻訳されず、 `SELECT` で直接に変換表現となる。 
- `AGGREGATE` の部分の実行文は若干変わっている。対象列に対して `is_not_null()` での filter 操作が[加わっている](https://github.com/ibis-project/ibis/blob/master/ibis/backends/polars/compiler.py#L687)こと、 `strict_cast()` が加わっていること。（aliasが二つ重なっているのはおそらく ibis 側での変換バグであろう。）最後の実行は `WITH_COLUMNS` ではなく `SELECT` となっている。意味論的には違いがないはずだが、実行速度のパフォーマンスへはプラスになるものではない。とはいっても影響は過大でがないと考えうる。

## まとめ

ibis にて Polars backend を使った際に select, filter について layz optimization が効いていることが確認できた。ただし、aggregation句にて filter(.is_not_null()) と strict_cast() が追加されてしまうぶん、生Polarsよりパフォーマンスが（若干だが）劣る可能性がある。実務上問題があるかは操作するデータ構造・量に依存するので答えられないが、個人的には問題ない範疇と考える。
 
## 免責事項

本記事の内容はの個人的見解・意見を述べるものであり、所属組織の見解ではなくまたそれらを代表するものでは一切ありません。