Robloxが機械学習で最適化されたブルームフィルターを用いてSparkの結合クエリのコストを削減する方法

概要

Robloxでは毎日、7,000万人のユーザーが数百万もの体験と関わり、四半期で合計160億時間に達しています。このインタラクションによりペタバイト規模のデータレイクが生成され、分析や機械学習（ML）の目的でデータが充実されています。データレイク内のファクトテーブルとディメンションテーブルを結合するには多大なリソースを要するため、これを最適化しデータのシャッフルを削減するために、MLを活用したスマートなデータ構造である「学習型ブルームフィルター」[1]を採用しました。 存在予測を行うことで、これらのフィルタは結合対象データを大幅に削減し、効率を向上させるとともにコストを削減します。その過程で、モデルアーキテクチャも改良し、処理に必要なメモリとCPU時間を削減し、運用安定性を高めるという大きなメリットを実証しました。

はじめに

当社のデータレイクでは、効率的なアクセスのためにファクトテーブルとデータキューブは時間軸でパーティション化されていますが、ディメンションテーブルにはそのようなパーティションがなく、更新時にファクトテーブルと結合するには多大なリソースを要します。結合のキー空間は、結合対象となるファクトテーブルの時間パーティションによって決定されます。 その時間パーティションに含まれるディメンションエンティティは、ディメンションデータセット全体に含まれるもののごく一部に過ぎません。その結果、これらの結合においてシャッフルされるディメンションデータの大部分は、最終的に破棄されてしまいます。このプロセスを最適化し、不要なシャッフルを削減するため、一意の結合キーに対してブルームフィルターの使用を検討しましたが、フィルターサイズとメモリ使用量の課題に直面しました。

これらの課題に対処するため、我々は学習型ブルームフィルター（Learned Bloom Filters）を検討しました。これは、低い誤検知率を維持しつつブルームフィルターのサイズを縮小する機械学習ベースのソリューションです。この革新的な手法により、計算コストを削減しシステムの安定性を向上させることで、結合操作の効率が向上します。以下の図は、我々の分散コンピューティング環境における従来の結合プロセスと最適化された結合プロセスを示しています。

学習型ブルームフィルターによる結合効率の向上

ファクトテーブルとディメンションテーブル間の結合を最適化するため、学習型ブルームフィルターの実装を採用しました。ファクトテーブルに含まれるキーからインデックスを構築し、結合操作の前にディメンションデータを事前フィルタリングするためにそのインデックスを適用しました。 

従来のブルームフィルターから学習型ブルームフィルターへの進化

従来のブルームフィルターは効率的ですが、目標とする偽陽性率を達成するために、それをロードする必要があるワーカーノードごとに15～25%の追加メモリを消費します。しかし、学習型ブルームフィルターを活用することで、同じ偽陽性率を維持しつつ、インデックスサイズを大幅に削減することに成功しました。これは、ブルームフィルターを二値分類問題へと変換したためです。陽性ラベルはインデックス内に値が存在することを示し、陰性ラベルは存在しないことを意味します。

機械学習モデルの導入により、値の初期チェックが容易になり、その後、偽陰性を排除するためにバックアップのブルームフィルターが使用されます。サイズが縮小されたのは、モデルの圧縮表現と、バックアップのブルームフィルターに必要なキー数の削減によるものです。これが、従来のブルームフィルターのアプローチとの違いです。 

本研究の一環として、学習済みブルームフィルタ手法を評価するための2つの指標を確立した。それは、インデックスの最終的なシリアライズ済みオブジェクトサイズと、結合クエリの実行中のCPU消費量である。 

実装上の課題への対処

当初の課題は、ファクトテーブルに次元テーブルのキーがほとんど含まれていない、偏りの激しいトレーニングデータセットへの対応でした。このデータセットでは、テーブル間のキーの重複率が約3分の1であることが確認されました。この問題に対処するため、我々は「サンドイッチ・ラーンド・ブルームフィルター」アプローチ[2]を採用しました。この手法では、初期段階として従来のブルームフィルターを統合し、ファクトテーブルに欠落していたキーの大部分を除去することでデータセットの分布を再調整し、事実上、データセットからネガティブサンプルを排除します。 その後、初期のブルームフィルタに含まれるキーと、誤検知（false positives）のみが、しばしば「学習済みオラクル」と呼ばれる機械学習モデルに転送されました。このアプローチにより、学習済みオラクル向けのバランスの取れたトレーニングデータセットが得られ、バイアスの問題を効果的に克服することができました。

2つ目の課題は、モデルアーキテクチャとトレーニング特徴量に焦点を当てたものでした。従来のフィッシングURLの問題[1]とは異なり、当社の結合キー（ほとんどの場合、ユーザーやエクスペリエンスを一意に識別する識別子）は、それ自体では情報価値が乏しいものでした。このため、ディメンションエンティティがファクトテーブルに存在するかどうかを予測するのに役立つ潜在的なモデル特徴量として、ディメンション属性を検討することになりました。 例えば、特定の言語でのエクスペリエンスに関するユーザーセッション情報を含むファクトテーブルを想定してください。ユーザーディメンションの地理的位置や言語設定属性は、個々のユーザーがファクトテーブルに存在するかどうかを示す良い指標となります。

3つ目の課題である推論のレイテンシに対処するには、誤検知（false negative）を最小限に抑えつつ、迅速な応答を提供できるモデルが必要でした。これらの主要な指標において、勾配ブースティングツリーモデルが最適な選択肢であり、精度と速度のバランスを取るためにその特徴量セットを剪定しました。

学習済みブルームフィルターを使用した更新後の結合クエリは、以下の通りです：

結果

以下は、当社のデータレイクにおける学習型ブルームフィルターの実験結果です。これらを、それぞれ異なるデータ特性を持つ5つの本番ワークロードに統合しました。これらのワークロードにおいて最も計算負荷が高い部分は、ファクトテーブルとディメンションテーブル間の結合処理です。 ファクトテーブルのキー空間は、ディメンションテーブルの約30%です。まず、最終的なシリアライズ済みオブジェクトのサイズにおいて、学習型ブルームフィルタが従来のブルームフィルタをどのように上回ったかについて説明します。次に、ワークロード処理パイプラインに学習型ブルームフィルタを統合することで確認されたパフォーマンスの向上について示します。

学習型ブルームフィルターのサイズ比較

以下に示すように、特定の誤検知率において、学習型ブルームフィルターの2つのバリエーションは、従来のブルームフィルターと比較して、オブジェクトの総サイズを17～42%削減しています。

さらに、勾配ブースティングツリーベースのモデルにおいて特徴量のサブセットを縮小することで、推論速度を向上させつつも、最適化性能の低下はごくわずかにとどまりました。

ブルームフィルターの使用結果 

このセクションでは、ブルームフィルターベースの結合と通常の結合のパフォーマンスを、いくつかの指標に基づいて比較します。 

以下の表は、学習済みブルームフィルターを使用した場合と使用しなかった場合のワークロードのパフォーマンスを比較したものです。両方の結合タイプで同じクラスター構成を維持しつつ、総誤検知率1%の学習済みブルームフィルターを用いた比較結果を示しています。 

まず、ブルームフィルターの実装は、CPU時間において通常の結合よりも最大60%も優れたパフォーマンスを発揮することがわかりました。学習型ブルームフィルターアプローチでは、ブルームフィルターの評価に追加の計算リソースが費やされたため、スキャンステップのCPU使用率が上昇しました。しかし、このステップで行われる事前フィルタリングにより、シャッフルされるデータのサイズが縮小され、下流のステップで消費されるCPUを削減できたため、結果として総CPU時間を削減することができました。

次に、学習済みブルームフィルターは、通常の結合と比較して、総データサイズが約80%少なく、書き込まれる総シャッフルバイト数も約80%少なくなっています。これにより、後述するように、結合パフォーマンスがより安定します。 

また、実験対象とした他の本番ワークロードにおいても、リソース使用量の削減が確認されました。5つのワークロードすべてを対象とした2週間の期間において、学習済みブルームフィルターアプローチは、モデルトレーニングやインデックス作成のコストも含め、1日あたり平均25%のコスト削減を実現しました。

結合処理中のシャッフルされるデータ量が減少したため、分析パイプラインの運用コストを大幅に削減すると同時に、その安定性も向上させることができました。以下のグラフは、実験対象とした5つのワークロードについて、2週間の期間における通常の結合ワークロードと学習済みブルームフィルターベースのワークロードの実行時間（実時間）の変動（変動係数を使用）を示しています。 学習済みブルームフィルターを使用した実行はより安定しており（実行時間がより一貫していた）、これにより、より安価で一時的かつ信頼性の低いコンピューティングリソースへの移行が可能になります。 

参考文献

[1] T. Kraska, A. Beutel, E. H. Chi, J. Dean, and N. Polyzotis. 「学習型インデックス構造の有用性」. 『https://arxiv.org/abs/1712.01208』, 2017.

[2] M. Mitzenmacher. 「サンドイッチ法による学習型ブルームフィルターの最適化」。 

https://arxiv.org/abs/1803.01474, 2018.

¹2023年6月30日終了の3ヶ月間時点

²2023年6月30日終了の3ヶ月間