NEC デジタルテクノロジー開発研究所の荒木 拓也です。こちらは NEC デジタルテクノロジー開発研究所 Advent Calendar 2023 18日目の記事です。
FireDucksとは？
すでに本Advent Calendar内で何人も関連記事を書いていますが、pandasと互換性のあるAPIを持つ高速なデータフレームライブラリです。現在、ベータ版を公開中で無償でお試し頂けます。
NEC 研究開発/新技術 プレスリリース（2023年10月19日）：
FireDucks 公式ウェブサイト：
をご覧ください。
FireDucksは
$ pip install fireducks
でインストールできます。また、
# import pandas as pd
import fireducks.pandas as pd
のようにimportするモジュールを切り替えるか、
$ python3 -m fireducks.imhook your_program.py
のようにimport hookという機能を用いることで、pandasを使ったプログラムを高速に実行することができるようになります。
FireDucksにおける最適化処理
FireDucksはpandasより高速に動作しますが、その理由は
- CPUが持つ複数のコアを使って並列に動作すること
- 内部で独自の最適化を行っていること
にあります。本記事では後者の最適化について紹介します。
FireDucksではpandasのAPI呼び出しを行う際、実際にはそのAPIの実行はせず、中間言語に変換します（中間言語はIR：Intermediate Representationと呼ばれることもあります）。そして、実行結果が必要になった時点で、中間言語のコンパイルと実行を行います。遅延評価と呼ばれる仕組みですね。
例えば、
import fireducks.pandas as pd
df = pd.read_csv("sample.csv")
sorted = df.sort_values("b")
result = sorted[["a"]]
print(result)
というプログラムがあったとします。このプログラムを実行すると、FireDucksは、read_csv("sample.csv"), sort_values("b"), sorted[["a"]]の呼び出し時点では関数の実行を行わず、下記のような中間言語に変換します。（説明のため、実際の中間言語から変更・省略しています。）
%v1 = "fireducks.read_csv"("sample.csv")
%v2 = "fireducks.sort_values"(%v1,"b")
%v3 = "fireducks.project"(%v2,["a"])
そして、実行結果が必要となるprint(result)という文を呼び出した時点で、今まで生成した中間言語をコンパイルし、実行するという仕組みです。ちなみにこの中間言語は、LLVMという著名なコンパイラフレームワークを利用しており、その中のMLIRとよばれる中間言語を拡張する仕組みを使っています。
さて、サンプルにしたPythonプログラムをもう一度見てみましょう。読み込んだデータをb列でソートして、a列を抽出、表示しています。ここでもし読み込んだデータに多量の列があったとしたらどうでしょうか。実際に必要なのはa列だけなのに、他の多量の列もb列でソートした後の順序入れ替えの対象となり、結局捨てられてしまいます。これは無駄な処理ですので、
df = pd.read_csv("sample.csv")
df2 = df[["a","b"]]
sorted = df2.sort_values("b")
result = sorted[["a"]]
print(result)
のようにソートの前に必要な列だけ抽出しておくと高速に実行できると期待されます。（pandasやFireDucksでは、データフレームはおおむね列の集合として実現されているので、列を抽出する処理は非常に軽いものになっています。）
FireDucksでは、これを中間言語のレベルの最適化により自動的に行います。動作としては、
- 列の抽出（
project）を中間言語の中から見つけます - projectの入力（
%v2）を定義している中間言語を見つけます。この場合はsort_valuesです。これはLLVMの機能を使うことで容易に実現できます。 - 対象となる中間言語の前に
projectを挿入します。挿入されるprojectは対象となる中間言語によって変わります。sort_valuesの場合は、ソートに使うキーも必要になるため、元々抽出する列であるaとソートキーであるbを抽出するprojectが挿入されます。 - この操作を対象とできない中間言語にたどり着くまで再帰的に繰り返します。すなわちこの場合では
sort_valuesの入力%v1を定義している中間言語を探して同じことを繰り返します。この場合はread_csvで、この前にprojectは挿入できないので、ここで動作を停止します。
結果として、この場合は
%v1 = "fireducks.read_csv"("sample.csv")
%v11 = "fireducks.project"(%v1,["a","b"])
%v2 = "fireducks.sort_values"(%v11,"b")
%v3 = "fireducks.project"(%v2,["a"])
のような中間言語に最適化されます。
mergeがあると若干複雑になったり、不要になったprojectは削除したりしますが、おおむね上記のように動作します。
上記では、私が実装を担当したということもあり、列の抽出を前もって行うという最適化を紹介しましたが、他にもいくつかの最適化が中間言語レベルで行われています。
最適化の評価
ではこれらの最適化がどれくらい効果があるのか、TPC-Hというベンチマークで評価してみました。FireDucksは実は環境変数でコンパイラの動作を制御できます。FIREDUCKS_FLAGSという環境変数に"-O0"という文字列を設定するとすべての最適化がスキップされますので、これを利用しました。データサイズはSF-10（全体で約10GB）、Xeon Gold 6226 2ソケット（24コア）マシンで評価しています。FireDucksのバージョンは0.9.0、pandasのバージョンは1.5.3です。3回実行した実行時間の最小値を使っています。
グラフの縦軸はpandasに対する速度向上率です。青が最適化無しの場合、オレンジが最適化有りの場合です。最適化の効果がある場合は最大で2倍以上の高速化が達成できていることが分かります。TPC-Hは最適化がかかりやすいプログラムだったということもあるかも知れませんが、最適化による高速化としては、悪くないのではないでしょうか。
おわりに
今回紹介した最適化手法は、データベースのクエリ最適化で行われている内容と類似しており、特に目新しいものではありません。一方、SQLのようなドメイン特化型の言語では無く、Pythonという一般的な言語を対象として、実行時に中間言語を作成し、それを対象に最適化を行うという仕組みは興味深いものでは無いでしょうか。
例えば、この仕組みで作成された中間言語には、条件分岐や関数呼び出しは含まれていないため（pandasのAPIを呼び出す際に中間言語に変換されるという仕組みのため）、大規模な処理でも容易に最適化することが可能になります。（fallbackされなければ、ですが…）
FireDucksはこれからも発展していきます。今後も新たな最適化処理を追加することで、高速化を進めていきます。