EMR Spark-SQL performance optimization reveals Native Codegen Framework

背景と動機

SparkSQL のパフォーマンス最適化は長年、オプティマイザーとランタイムの 2 つの領域に焦点を当ててきた。
前者は最適な実行計画を導出することを目的とし、後者は導出された計画を可能な限り迅速に実行することを目的とする。


オプティマイザーはランタイムと比較して、より汎用的で実装に依存しない最適化を実現する。
Java 環境 (Spark、Hive) でも C++ 環境 (Impala、MaxCompute) でも、バッチベース (Spark、Hive) でも MPP ベース (Impala、Presto) でも、さらにはビッグデータ、従来のデータベース、HTAP (HyPer、ADB) のいずれの領域でも、オプティマイザーレベルでは非常に似た問題を検討する。
統計情報の収集、コスト評価、計画選択であり、使用される最適化手法も JoinReorder、CTE、GroupKey Elimination など共通している。
インデックスの有無などコンテキストの違いによるコストモデルの構築方法の差や、遺伝的アルゴリズム対動的計画法など特定のシナリオでの探索戦略の違いはあるものの、全体的なアプローチは基本的に同じである。


長きにわたり、ランタイムの最適化作業は基本的に当時のハードウェアボトルネックの解消に注力してきた。
たとえば、MapReduce 登場時はネットワーク帯域幅がボトルネックであり、Google はローカリティ最適化を数多く実施した。
Spark 登場時はディスク IO が新たなボトルネックであり [1]、近年は CPU 性能の向上がランタイム領域の重要な最適化方向となっている。


CPU 性能向上の主流技術は 2 つある。
MonetDB/X100 [2] に代表されるベクトル化処理技術 (現在は VectorWise [3] に進化) と、HyPer [5][6] に代表されるコード生成技術 (CodeGen) である (Spark も CodeGen [9] を採用)。
簡単に説明すると、ベクトル化技術は火山モデルを踏襲しつつ、SQL 演算子に 1 レコードずつ計算させるのではなく、データをバッチで蓄積してから実行する。
バッチ単位の計算はレコード単位の計算よりも最適化の余地が大きい。
たとえば、仮想関数のオーバーヘッド削減、SIMD 最適化、キャッシュフレンドリ性の向上などである。
この技術の欠点は、演算子間で転送されるデータがレコード単位からバッチ単位に変わるため、中間データの実体化オーバーヘッドが増加することである。
CodeGen 技術は演算子融合という別の角度から仮想関数のオーバーヘッドと中間データの実体化の問題を解決する。
簡単に言えば、CodeGen フレームワークは火山モデルを「平坦化」し、演算子間の境界をなくして、元のイテレーターチェーンを 1 つの大きな for ループに圧縮し、同じセマンティクスを持つコード (Java/C++/LLVM) を生成する。
そして対応するツールチェーンで生成コードをコンパイルし、コンパイル済みのクラス (Java) や so (C++、LLVM) を使って実行することで、解釈実行をコンパイル実行に変換する。
さらに、関数呼び出しがなくなるため、1 レコードは初期演算子 (Stage 内) から最終演算子まで基本的にレジスター上に保持され、メモリに実体化されることはない。
CodeGen 技術の欠点は、SIMD などの最適化を適用しにくいことである。


両派はそれぞれの優位性を示す論文を発表した後 [4][8]、協力関係に移行し、一連のプロジェクトと論文を生み出した。
現在の学術界の主流見解は、融合が最適解であるというものであり、HyPer の進化版 [6]、Peloton [7] など、融合方式を採用したプロジェクトが時代に応じて登場している。


学術界は融合に到達したが、産業界には融合への移行を促す強い動機がない。
その理由を探ると、主に 3 つの要因がある。
第 1 に、現在の融合方式は個別最適化と比較して質的な飛躍をもたらしていないこと。
第 2 に、業界で広く認められたベストプラクティスがまだ探索段階にあること。
第 3 に、産業界は単一技術でもまだ最大限の潜在能力を発揮していないことである。
SparkSQL を例に取ると、2015 年の Expression-level コード生成の初登場から、HyPer を参考にした WholeStage コード生成の実装まで、長年にわたる改良を経て、SparkSQL のコード生成技術は成熟し、性能も 2 桁の向上を達成した。
しかし、保守性や開発者の受容性の問題からか、SparkSQL のコード生成は Java コードの生成に限定されており、ネイティブコード (C/C++、LLVM) への試みはなされていない。
Java の性能は既に優れているものの、ネイティブコードと比較すると一定のオーバーヘッドがあり、SIMD (Java はこの機能に対応中)、Prefetch などのセマンティクスが不足している。
さらに重要なのは、ネイティブコードはベアメタルを直接操作するため、ハードウェア性能を極限まで引き出しやすく、GPU などのアクセラレーターや AEP などの新しいハードウェアをサポートするのにも便利である。


上記の動機に基づき、EMR チームは SparkSQL ネイティブコード生成フレームワークを探求・開発し、SparkSQL のエンジンを変更した。
新しいエンジンは約 20% の性能向上をもたらし、EMR が世界トップクラスの性能を再び獲得することに貢献した。
本稿ではネイティブコード生成フレームワークを詳しく紹介する。


核心的な課題
ネイティブコード生成を行う際、3 つの核心的な課題がある。

1. 何を生成するか？

2. どのように生成するか？

3. Spark にどう統合するか？


何を生成するか
生成するコードについて、研究の知見と開発チームの技術スタックを踏まえ、C/C++、LLVM、Weld IR の 3 つの候補がある。
C/C++ の利点は実装が比較的簡単なことであり、Spark が生成する Java コードのロジックを書き換えればよい。
欠点はコンパイル時間が長すぎることである。
下図は HyPer の評価データであり、C++ のコンパイル時間は LLVM より 1 桁高い。
compile time.jpg
コンパイル時間が長すぎると小規模なクエリには不向きであり、極端な場合はコンパイル時間が実行時間を超えることもある。
この考慮から C/C++ オプションを除外した。
上図から LLVM のコンパイル時間は非常に良好であり、HyPer、Impala、Alibaba Cloud の独自開発ビッグデータエンジン MaxCompute、ADB など、多くのネイティブコード生成エンジンが LLVM を対象コードとして使用している。
我々にとって (あなたにとってはそうではないかもしれないが :D)、LLVM の最大の欠点は低レベルすぎることである。
構文がアセンブリに近く、SparkSQL の演算子をアセンブリで書き換える作業量を想像してほしい。
ほとんどのエンジンは LLVM でフルコードを書いていない。
たとえば HyPer は LLVM で演算子のコアロジックのみを生成し、spill や複雑なデータ構造管理などの他の汎用関数は C++ で書いて事前コンパイルしている。
LLVM+C++ でもかなりの作業量を節約できるが、それでも我々には許容できないため、第 3 のオプションに目を向けた。
Weld IR (中間表現) である。


まず Weld の簡単な紹介から始める。
Weld の作者 Shoumik Palkar は、Spark の作者としておなじみの Matei Zaharia の弟子である。
Weld が当初解決しようとした問題は、異なるライブラリ間で相互呼び出しする際のデータ転送のオーバーヘッドである。
たとえば pandas で numpy のインターフェイスを呼び出す場合、まず pandas がデータをメモリに書き込み、次に numpy がメモリを読み取って計算する。
極限まで最適化されたライブラリでは、メモリの読み書きの時間が計算自体を大幅に上回ることがある。
この問題を解決するため、Weld は Common Runtime を開発し、IR を提供した。
遅延評価の機能も備えているため、(簡単に) ライブラリを Weld 仕様に準拠させるだけで、異なるライブラリ間で Weld Runtime を共有できる。
Weld Runtime は遅延評価を利用してクロスライブラリのパイプラインを実現し、データの実体化オーバーヘッドを排除する。
さらに Weld Runtime はループ融合、ループ展開、ベクトル化、適応実行などの最適化も施している。
また、Weld は C コードの呼び出しをサポートしており、サードパーティライブラリも簡単に呼び出せる。


我々が注目しているのは Weld が提供する IR と対応する Runtime である。
Weld IR はデータ分析向けに設計されているため、セマンティクスが SQL に非常に近く、演算子をより適切に表現できる。
データ構造の面では、Weld IR のコアデータ構造は vec と struct であり、SparkSQL の UnsafeRow Batch をより適切に表現できる。
struct と vec を基に dict を構築でき、SQL で多用される Hash 構造をより適切に表現できる。
操作の面では、Weld IR は map、filter、イテレーターなどの関数型言語に似たセマンティクスを提供し、builder セマンティクスと組み合わせることで、Project、Filter、Agg、ブロードキャスト JOIN などの演算子のセマンティクスを簡潔に表現できる。
たとえば以下の IR は Filter + Project のセマンティクスを表現しており、2 列目が 10 より大きい場合に 1 列目を返すことを意味する。


これにより Weld IR の利点は明らかである。
性能 (最終的に LLVM コードを生成) と使いやすさ (LLVM や C++ よりもはるかに簡単にコード生成できる) の両方を兼ね備えている。
これらの考慮から、最終的に Weld IR を対象コードとして選択した。


この過程で 2 つの重要な課題がある。


1. 演算子間の伝達媒体は何か？

2. Weld がサポートしていない演算子をどう処理するか？


伝達媒体
Java と異なり、Weld IR はループ構造を提供しない。
代わりに vec 構造とそれに対する汎用イテレーター操作を使用する。
そのため、Weld IR は Java Codegen のように Stage 外側に大きなループを設定して各演算子が 1 レコードずつ処理する方式をそのまま取り入れにくい。
代わりに各演算子がバッチデータを処理し、IR レベルで仮の実体化を行い、Weld のループ融合最適化に依存して実体化を排除する。
たとえば前述の Filter の後に Project が続く場合、Filter 演算子が生成する IR は以下の通りであり、2 列目が 10 以下のデータをフィルタリングする。


|v:vec[{i32,i32}]| let res_fil = for(v,appender,|b,i,n| if(n.$1>10, merge(b,n), b)
Project 演算子が生成する IR は以下の通りであり、1 列目のデータを返す。


let res_proj = for(res_fil,appender,|b,i,n| merge(b,n.$0))
表面上、Filter 演算子が中間結果を実体化しているように見える。
実際には Weld のループ融合オプティマイザーがこの実体化を排除する。
最適化後のコードは以下の通りである。


|v: vec[{i32,i32}]| for(v,appender,|b,i,n| if(n.$1 > 10, merge(b,n.$0), b))
Weld のループ融合最適化に依存することで CodeGen のロジックを大幅に簡略化できるが、開発中にループ融合プロセスが非常に時間がかかることが判明した。
3 レベル以上のネストを含む複雑な SQL では、限られた時間内に結果を得られないこともある。当時、2 つの選択肢があった。Weld の実装を修正するか、CodeGen を修正してループ融合後のコードを直接生成するかであり、後者を選択した。リファクタリング後の生成コードは以下の通りである。1、2、11 行目は Scan 演算子が生成し、3、4、5、6、8、9、10 行目は Filter 演算子が生成し、7 行目は Project 演算子が生成する。

この最適化により、コンパイル時間は 1 秒未満に戻った。

フォールバック機構

Weld の現在の表現力に制約があり、SortMergeJoin、Rollup などの一部の演算子は Weld で実装できない。元の Java CodeGen でも、Outer Join などの一部の演算子はコード生成に対応していない。そのため、適切なフォールバックを確保することが正確性を保証する前提となる。我々が採用する戦略は直感的である。現在の演算子がネイティブコード生成に対応していない場合、Java CodeGen が引き継ぐ。ここでの重要な課題はフォールバックの粒度である。演算子レベルか、それとも Stage レベルか。

実装の難易度は別として、演算子単位のフォールバックは直感的にはより適切だが、実際にはより深刻な問題を引き起こす。Stage 内部のパイプラインの分断である。前述の通り、CodeGen の利点の 1 つは Stage 全体のロジックをパイプライン化し、演算子間の境界をなくすことである。単一レコードが初期演算子から最終演算子まで実行され、全体を通して実体化がない。しかし演算子粒度のフォールバックは、Stage の一部がネイティブランタイムを使い、他の部分が Java ランタイムを使うことを意味するため、両者の接続部分で中間データの実体化が必然的に発生する。このオーバーヘッドは通常、ネイティブランタイムがもたらすメリットを上回る。

上記の考慮から、Stage レベルのフォールバックを選択した。CodeGen 段階でサポートされていない演算子に遭遇したら、Stage 全体を Java CodeGen にフォールバックさせる。統計によると、TPC-DS ベンチマーク全体でネイティブコード生成 Stage の 80% が該当する。

Spark との統合

コード生成とフォールバック機構を完成させた後、最後の課題は Spark との統合である。Spark の WholeStageCodegenExec の実行はブラックボックスと見なせる。上流がテーブルスキャン、シャッフルリード、ブロードキャストのいずれでも、ブラックボックスへの入力は RowBatch (上流がテーブルスキャン) または行イテレーター (上流がテーブルスキャン以外) の 2 種類のみであり、ブラックボックスの出力は行イテレーターに固定されている。

前述の通り Stage レベルのフォールバックを選択したが、これはブラックボックスが Java ランタイムかネイティブランタイムのいずれかであることを意味し、混在することはない。そのため、RowBatch または行イテレーターを Weld のメモリレイアウトに変換する方法と、Weld の出力を行イテレーターに変換する方法にのみ集中すればよい。問題をさらに簡略化するために、Shuffle Reader/ブロードキャストの入力は行イテレーターだが、リモートシリアル化のデータ構造は本質的に RowBatch であることに着目した。Spark がデシリアライズして行イテレーターに変換してから CodeGen Module に渡しているだけで、RowBatch を行イテレーターにラップするのは非常に簡単である。したがって、ネイティブランタイムの入出力は RowBatch に統一できる。

解決策はもう目の前である。RowBatch を Weld の vec に変換する。しかしもう一歩進めて、RowBatch を直接 Weld に渡してメモリ変換を省いてはどうか。本質的に RowBatch は一定の仕様を満たすバイトストリームに過ぎない。Spark は OffHeap モードも提供しており、ヒープ外にメモリを直接保存できる (Scan Stage のみ。シャッフルデータとブロードキャストデータはヒープ外に読み出す必要がある)。Weld はそれに直接アクセスできる。Spark UnsafeRow のメモリレイアウトは概ね以下の通りである。

特定のスキーマでは、null ビットマップと固定長データの構造は固定であり、struct にマッピングできる。可変長データについては、これらのデータを連続したメモリアドレスにコピーするアプローチを取っている。これにより、可変長データがない RowBatch はメモリブロックを直接 Weld に渡し、可変長データがある場合は固定長部分と可変長部分をそれぞれ別々にコピーするだけで、列レベルの細粒度なコピー変換は不要である。

前述の Filter+Project の例を続けると、1 レコードに 2 つの int 列が含まれており、UnsafeRow のメモリレイアウトは以下の通りである (アライメントのため、Spark の固定長部分は最低 8 バイトを使用する)。

これで Input の問題が解決した。Weld Output を RowBatch に変換するのは、本質的に上記プロセスの逆変換であるため、ここでは繰り返さない。

Java とネイティブ間のデータ変換の問題が解決したら、残りは実行方法である。まず現在の Stage の Mode に基づいて、Java ランタイムとネイティブランタイムのどちらを使用するかを決定する。ネイティブブランチでは、まず StageInit を実行して Stage レベルの初期化作業を行う。Weld の初期化、コンパイル済み Weld Module の読み込み、ブロードキャストデータのプル (あれば) などが含まれる。次にループに入り、各ループで RowBatch (Scan またはシャッフルリーダーから) を読み込み、ネイティブランタイムに渡して実行し、出力を変換してシャッフルライターに渡す。

まとめ
本稿では EMR チームが Spark ネイティブコード生成の方向で探求・実践した内容を紹介した。紙面の都合上、一部の技術ポイントや最適化については詳細を展開していない。別の記事で詳しく説明できる。たとえば以下の通りである。

1. ネイティブ演算子の極限最適化
2. データ変換の詳細説明
3. Weld Dict 最適化