LoginSignup
5
4
PolarsAdvent Calendar 2023Day 9
@hkzm(KZM H)in株式会社justInCase

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

Last updated at Posted at 2023-12-09

はじめに

ibis という 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の破壊的変更へのキャッチアップは若干遅れを取るが、込み入った操作でなければほぼ対応できているようだ。詳しくはサポート対応表で確認できる。

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

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

Polarsのlazy APIではoptimizations が効くが、ibisのバックエンドにPolarsを指定した場合、この optimizations が効くのだろうか確認したい。

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

ibis の注意点

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

Prerequisite

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

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

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

# homebrew or linuxbrew (wsl2)
brew install graphviz

サンプルデータ

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

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 と類似しているので馴染みがあるかもしれない。

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 のオブジェクト名としてサンプルコードで使われることが多い）

さて、操作した結果自体に意味はないが、次のコードを実行してみよう。そしてその結果について次で解説する。

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 を実行している。
 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経由での操作を見てみよう

# 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 となってしまう。 
    q1.collect().dtypes
# [Utf8, UInt32, Float64, Float64]

q2.compile().collect().dtypes
# [Utf8, Int64, Float64, Float64]

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

 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 操作が加わっていること、 strict_cast() が加わっていること。（aliasが二つ重なっているのはおそらく ibis 側での変換バグであろう。）最後の実行は WITH_COLUMNS ではなく SELECT となっている。意味論的には違いがないはずだが、実行速度のパフォーマンスへはプラスになるものではない。とはいっても影響は過大でがないと考えうる。

まとめ

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

免責事項

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

5
4
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
Sign upLogin
5
4