『How Query Engines Work』

仕事で Cloudera Impala に触れ、クエリエンジンに興味が湧き、自分で作ってみたいと思うようになって作り始めました。最初はクエリエンジンの作り方が分からなかったので、とりあえず参考書等を探して入り口を探しました。しかしそれらを読んでも実装できる気がしなかったので、まずは Apache Drill のコードを読んでいました。

今年、Apache Arrow の PMC の Andy Groove が『How Query Engines Work』という書籍をリリースしているので、興味のある方は読んでみてください。彼自身も Rust と Apache Arrow を使って DataFusion というクエリエンジンを書いています。

イチから作るのはとても時間がかかるので、色んなライブラリを使っています。今回はそれらの紹介をしたいと思います。

Parsing SQL and Building an execution plan

SQL の parse から実行計画(Logical Plan)の構築までは、Apache Calcite を使っています。Calcite については以前 blog を書いたので、そちらを参考にしてください。truffle-arrow の前は、Calcite の Arrow Adapter を書いていました。

SQL を入力すれば Calcite が Logical Plan まで作成してくれるので、あとはそこから Physical Plan を生成し、Physical Plan から先で SQL の処理を実行して結果を返すだけです。Physical Plan の生成も Calcite がサポートします。

Calcite では、実行計画の Scan / Filter / Project 等を RelNode というオブジェクトとして扱います。Physical Plan の RelNode のクラスを作り、Logical Plan の各 RelNode を Physical Plan の RelNode に変換していきます(ConverterRule)。1:1 で対応させていく必要はなく、2 つの RelNode を 1 つに纏めたり、RelNode を組み替えたりすることができます。

Physical Plan の RelNode が、入力された SQL を処理していきます。Calcite の標準実装的な位置付けである Enumerable Adapter は、これらの RelNode の中で実行計画を参照しながら Java のコードを生成し、最後にそれらをコンパイルして実行し、結果を返すようになっています。

Data Source

内部では Apache Arrow 形式のデータを処理しています。Apache Arrow については、以下を参照してください。

Apache Arrow - データ処理ツールの次世代プラットフォーム

現時点では Arrow 形式のデータを予めカラムごとに分けて出力し、実行計画を見て必要なカラムのファイルを読むようにしています。

Code Generation

以前 Calcite の Arrow Adapter を書いていた頃は、Enumerable Adapter と同じように Java のコードを生成していたのですが、Java のコードを生成する Java コードを書くのがどうにもシンドくなってしまい、別の方法を探すことにしました。

暫くして Graal / Truffle を知り、コード生成の部分に Truffle が使えないかと思うようになりました。Truffle は言語実装フレームワークです。Truffle については、以下を参照してください。

GraalVMで使われている、他言語をJVM上に実装する仕組みを学ぼう

そして Web を探してみた結果、Calcite + Truffle で実装しているプロジェクトを見つけました。

fivetran/truffle-sql: Experimental data-lake implemented using Truffle / Graal

このプロジェクトは現在 Update されていませんが、この truffle-sql と Arrow Adapter のコードを組み合わせたのが、truffle-arrow のベースになっています。

Expression Evaluation

Calcite では、SQL 文の中の式等を RexNode というオブジェクトとして扱います。truffle-sql では、この RexNode を Truffle の Node に置き換えていきます。Truffle の Node に置き換えれたオブジェクトグラフは式を表します。この式を構築する際に、評価した結果を Truffle の VirtualFrame という Truffle 上のスタックに書き込むようにしておきます。

そして最後に RelNode を下から順に実行していきます。最初の Scan では、必要なデータを load して VirtualFrame に書き込んでおきます。それ以降の RelNode では、VirtualFrame から値を取り出して評価し、結果を VirtualFrame に書き込みます。最後に、VirtualFrame から JDBC の ResultSet 経由で値を取り出せるようにしています。

Java のコード生成が Truffle に変わったことで、大分楽になりました。とはいえ Truffle は言語実装フレームワークなので、Calcite の RexNode を置き換えて評価するには色々と合わない部分もあります。Truffle に関しては、困ったら Slack で質問していました。

JDBC API

特に実装しなくても、Calcite の実行結果は JDBC API 経由で取得できるようになっています。しかし、ネットワークを介して結果を取得することはできません。そこで、Apache Avatica を使ってネットワークを介して結果を取得できるようにしています。

Avatica を使うことにより、HTTP + JSON か Protocol Buffers で DB とデータをやり取りすることができます。クライアント側は Avatica の JDBC Driver を使うことになります。

Source Code

Scan / Filter / Project / Aggregate(+ count 関数) まで実装したものを以下に置いています。

https://github.com/masayuki038/truffle-arrow

Conclusion

クエリエンジンを作るために使っているライブラリを紹介しました。Java 系の言語を使って SQL クエリエンジンを作るのであれば、その取っ掛かりとしてApache Calcite を見てみることをお勧めします。

最初は、なかなか動くところまで行くのに時間がかかりました。今も、形を変えては失敗し、また形を変えてみる、の繰り返しの日々です。ただ、その過程で学べることも多いので、暫くは続けていきたいと思います。

truffle-arrow は現在 Parallel Query を実装しようとしています。それが終わったら、一旦パフォーマンスを計測してみようと思っています。