Databricks Feature Storeがサブ秒単位の鮮度でフィーチャーを提供する方法

機械学習モデルの性能は、入力されるシグナルの質によって決まります。不正検知のユースケースでは、ユーザーが購入ボタンを押してから数ミリ秒以内に、取引を承認するかどうかを判断する必要があります。正しい判断を下せるかどうかは、わずか数秒前に発生した不審な取引を検知できるかどうかにかかっています。過去30日間のユーザーの平均取引実績と、過去10分間の総取引額を組み合わせることで、潜在的な不正行為を浮き彫りにすることができます。長期的な集計によってユーザーの基準となるプロファイル（通常の状態）を設定し、直近のデータによって異常な行動が発生した瞬間にそれをあぶり出すことができます。パーソナライズも同様の課題に直面しています。最新のシグナルこそがユーザーの現在の意図を捉え、エンゲージメントを向上させる鍵となります。

Sparkパイプラインは、過去の基準となる特徴量（ベースライン特徴量）を作成するために、レイクハウス内の大量データを一括処理する確立された手法です。これらのバッチジョブを定期的なスケジュールで実行する方法は広く普及していますが、数分から数時間の遅延（ラグ）が発生します。ユーザーに関するベースラインシグナルについては、インフラをシンプルに保つための許容可能なコストと言えます。しかし、モデルが最新のシグナルを必要とする場合、このインフラは破綻します。既存のフィーチャーストアプラットフォームでは、秒単位やミリ秒単位のリアルタイム性を実現することは不可能です。最新の特徴量の価値を引き出すために、データサイエンティストは、これらの集計を処理するための複雑なストリーミング専用ロジックを実装し、独自のホスト型インフラを構築することを余儀なくされています。

Databricks Feature Storeを使用すると、特徴量を一度作成すればどこでも使用できます。同じ定義を使用して、オフラインでの大規模なバッチフローと、オンラインでの極めて新鮮な特徴量パイプラインの両方を動かすことができます。このフレームワークはインフラの負担を取り除き、継続的なストリーム処理のためのSpark Real-Time Mode (RTM)、ストリーミングに最適化されたオンラインストレージであるLakebase、および大規模な取得を行うためのModel Servingをオーケストレーションします。そして、一度作成された特徴量はミリ秒単位で提供されます。Kafkaにイベントが到着してからオンラインフィーチャーストアで利用可能になるまでのエンドツーエンドのp99レイテンシは200msです。

アーキテクチャ：KafkaからFeature Storeまで200ms

それでは、Databricks Feature Storeがインフラに依存しない特徴量定義を取り込み、ミリ秒単位で一貫して計算するパイプラインをどのように構築しているのか、その内部の仕組みを見てみましょう。ストリーミング特徴量のエンドツーエンドの経路は以下のようになります。

1. イベントがKafkaに到達する - クレジットカードの取引、広告のインプレッション、クリックストリームイベントなどの生データ
2. サーバーレスのLakeflow Spark Delta Pipelines上で動作するSpark RTMパイプラインがイベントを継続的に処理し、ローリング集計をリアルタイムで計算する
3. 更新された集計データは、新しいストリーミングJDBCシンクを介してLakebaseに書き込まれ、オンラインフィーチャーストアに格納される
4. Model Servingエンドポイントが推論時にLakebaseから最新の特徴量を取得し、モデルに自動的に供給する

これを、過去10分間のユーザーの取引金額の合計という、不正検知の特徴量に当てはめて考えてみましょう。流入する各イベントには、金額、場所、ユーザーID、加盟店情報などの取引の詳細が含まれており、ステートフルなパイプラインにルーティングされます。パイプラインは、ユーザーの現在の取引合計を保持するローカルのRocksDBインスタンスを参照します。このインスタンスには、ウィンドウを過去10分間に制限するための有効期限が設定されています。パイプラインはローカルで値を読み取ってインクリメントし、更新された特徴量をLakebaseに書き込みます。そのため、新しい取引を承認するためにモデルへのクエリが発生した際、最新の取引合計額が1秒未満の鮮度でフィーチャーストアから利用可能になります。この合計値の特徴量は、承認の判断材料として、ユーザーの過去の購買ベースラインとともに取得されます。過去のベースラインを大幅に上回る合計額は、モデルにとって潜在的な不正行為を示す強力な指標となります。

このパイプラインの各コンポーネントは、流入するイベントのルーティング、集計の計算、オンラインストアへの特徴量の書き込みができる限り迅速に行われるよう最適化されています。

ローリングウィンドウ：ミリ秒単位での集計更新

インフラの詳細に入る前に、集計特徴量と、バッチ同期からリアルタイム更新へのパラダイムシフトについて説明します。

時間ウィンドウにおける集計特徴量（カウント、合計、平均など）は、リアルタイムMLにおいて強力かつ柔軟なシグナルとなります。長期的なバッチ特徴量は、一定期間におけるユーザーの過去のベースラインを設定し、モデルが各ユーザーの行動を理解して適応できるようにします。一方、短期間の新鮮な特徴量は、変化する状況に素早く反応し、ユーザーの新しい関心や不正なアクティビティを識別します。時間ウィンドウは、時間範囲（例：10分）と、それらの時間範囲が時間の経過とともにどのように変化するか（例：重複するか、重複しないか）を定義します。

Databricks Feature Storeは、3つの異なる時間ウィンドウをサポートしています。

• タンブリングウィンドウは、実時間（ウォールクロック）の間隔に合わせられ、前の間隔が終了するとすぐに開始されます。たとえば、10分間のタンブリングウィンドウは12:00〜12:10、次に12:10〜12:20をカバーします。イベントはこれらの固定された間隔にバッチ処理され、間隔の最後に特徴量が出力されます。つまり、集計値が最新になるのは間隔の境界のみです。
• スライディングウィンドウも実時間の間隔に合わせられますが、間隔の重複が許容されます。たとえば、スライド間隔が5分の10分間スライディングウィンドウは、12:00〜12:10、次に12:05〜12:15、そして12:10〜12:20をカバーします。
• ローリングウィンドウは実時間には合わせられず、各イベントのタイムスタンプからミリ秒単位の解像度で過去を遡ります。「現在の実時間における過去10分間の取引合計」は、新しいイベントが発生するたびにウィンドウが移動するため、常に最新の状態に保たれます。これにより、RollingWindowは「現在」が常に変化するリアルタイムサービングに最適です。

タンブリングウィンドウとスライディングウィンドウは、特徴量が頻繁に変化しない場合に依然として有用です。これらは更新の出力回数が少なく、維持コストが抑えられ、よりシンプルなスケジュール設定されたパイプラインに自然に適合します。ローリングウィンドウはその効率性と引き換えに最大限の鮮度を提供します。これは、新しいイベントが発生するたびに、モデルに提供される値に即座に反映されるべきシグナルにおいて最も価値を発揮します。

Feature Storeの宣言的APIを使用して、ローリングウィンドウ特徴量を定義するのがいかに簡単であるかを以下に示します。

Spark Real-Time Mode：特徴量計算のエンジン

基盤となるインフラに目を向けると、高スループットで新鮮な特徴量を実現しているのはストリーミングパイプラインです。このパイプラインは、Kafkaからオンラインフィーチャーストアまでデータを一貫して配信します。このストリーミングパイプラインは、Spark Structured Streamingの根本的に新しい実行モードであるSpark Real-Time Mode (RTM)によって稼働しています。RTMは、ミリ秒単位の鮮度を可能にする極めて重要なアーキテクチャの革新です。

同時実行ステージとステートフル処理

従来のマイクロバッチモード（MBM）では、Sparkはストリーミングデータを個別のバッチとして処理します。各バッチは設定可能な間隔でイベントを収集し、各ステージを通じて順次処理し、チェックポイントを作成してから、次のバッチを開始します。これによりレイテンシの下限が生じます。どれほど積極的にチューニングを行っても、ステートフルな集計を行うMBMパイプラインは通常、数秒から数分のオーダーで動作します。一方、RTMはステージを同時に実行します。集計オペレーターは、上流のステージがすべての行の処理を完了するのを待つことなく、行が利用可能になった瞬間に即座に処理します。

ローリング集計には、2つの重要なステージがあります。第1のステージは、データ処理、スキーマ検証、データの合体（coalescing）、型キャストです。これにより、一般的なアクションイベントを特徴量集計に適した形状に変換するビジネスロジックが実行されます。第2のステージは、エンティティごとにデータを集計してローリングウィンドウ集計を計算することです。流入する各行は、ローカルのRocksDBステートストア内の集計値を即座に更新し、新しい値を下流に出力します。ウィンドウの有効期限切れも行ごとに発生します。特定のイベントのウィンドウ期間が経過すると、パイプラインはそのイベントの寄与分を削除し、修正された集計値をLakebaseに出力します。RocksDBは各エグゼキューター上でローカルに実行されるため、クラスターのメモリ容量を超えるステートサイズを処理できます。

サーバーレスRTMにおけるパイプラインのステート管理

チェックポイント作成（チェックポインティング）は、個々のパイプラインワーカーの障害からパイプラインを回復できるようにするため、ステートフルなストリーミングにおけるフォールトトレランス（耐障害性）に不可欠です。しかし、チェックポイント作成にはコストが伴います。マイクロバッチモードでは、Sparkはバッチの境界ごとにチェックポイントを作成します。各チェックポイントはクラウドのオブジェクトストアとやり取りするため、パイプラインにレイテンシが加わります。

RTMは異なるアプローチを採用しています。プランニングとチェックポイント作成のコストを、より長い間隔に分散（アモルタイズ）させます。チェックポイント作成のコストは、各バッチの境界でパイプラインをブロックするのではなく、その間隔内に処理されたすべての行に分散されます。これにより、フォールトトレランス（耐障害性）が犠牲になることはありません。Exactly-once（正確に1回）の処理保証は維持され、障害発生時にはパイプラインがKafkaソースから最大5分間のデータを再実行（リプレイ）します。トレードオフとして、定常状態の処理レイテンシが大幅に削減される代わりに、再実行されるデータ量がわずかに増加します。

Feature Storeは、Lakeflow Spark Delta Pipelines（SDP）上でサーバーレスのRTMパイプラインを実行するため、クラスター管理やキャパシティプランニングが完全に不要になります。マシンのプロビジョニング、エグゼキューター数のチューニング、クラスターのメンテナンスについて心配する必要はありません。インフラストラクチャのアップデートによりパイプラインの再起動が必要になった場合、SDPがハンドオフを調整します。古いクラスターが停止する前に、新しいサーバーレスクラスターがプロビジョニングされ、完全に準備が整います。この調整は5分間のチェックポイント作成間隔に同期されるため、ダウンタイムを最小限に抑え、再処理のギャップを防ぐことができます。これにより、メンテナンス時間中であっても、特徴量の新しさ（フレッシュネス）への影響をほぼゼロに抑えることができます。

Lakebase：ストリーミング書き込みのオーバーヘッドを最小限に抑える

Databricks Feature Storeは、推論用のオンライン特徴量を保存するためにLakebaseを使用しています。コンピュートとストレージを分離するLakebaseのアーキテクチャにより、モデル推論の変動する負荷に対応するためのオートスケーリングが可能になります。オンラインFeature Storeはこの機能を活用し、数十ミリ秒のレイテンシで毎秒数万回の読み取りに対応できるようスケールします。

ストリーミング書き込みは、Kafkaの行を受信するたびに新しいローリングウィンドウの値が出力され、大量の小さなアップサート（upsert）が発生するため、特に困難です。標準的なPostgresでは、リカバリを容易にするためにフルページ書き込みを使用するため、このパターンによって大量のライトアヘッドログ（WAL）が生成される可能性があります。各チェックポイントの後、ページに対する最初の変更では、小さな論理的変更だけでなく、8KBのページイメージ全体がライトアヘッドログ（WAL）に書き込まれます。頻繁に更新されるホットなエンティティ行の場合、これによりWALの増幅（WAL amplification）が発生し、書き込みスループット、レプリケーション、およびリカバリのオーバーヘッドのボトルネックになります。

Lakebaseは、コンピュートと分散ストレージの分離を活用することで、標準的なPostgresと比較してストリーミング書き込みの増幅を最小限に抑えます。Lakebaseのアーキテクチャにより、PostgresはWALにフル8KBのページスナップショットを繰り返し書き込む代わりに、小さくコンパクトな変更レコードを書き込むことができます。これらのコンパクトなレコードは、分散されたセーフキーパーノードのクォーラム（合意形成）によって確認されるため、永続性は維持されます。十分な変更レコードが蓄積された後のリカバリには、依然としてフルページスナップショットが必要ですが、これらは書き込みパスを肥大化させることなく、後からストレージレイヤーで生成されます。Feature Storeにおけるこの結果、RTMはWALの増幅を大幅に抑え、追加のレイテンシを最小限に抑えながら、新しい特徴量をLakebaseに継続的にパブリッシュできるようになります。

Model Serving：大規模環境での低レイテンシな特徴量取得

最後のステップは、Lakebaseから最新の特徴量を取得し、推論時にモデルに提供することです。これは、高いQPSと低レイテンシのワークロード向けに最適化された、完全に管理されたサービングインフラストラクチャであるDatabricks Model Servingによって処理されます。

Model Servingは、リアルタイムMLのスループット要件に合わせて構築されています。

• 完全に水平スケーラブルなアーキテクチャ：推論サーバー、認証レイヤー、プロキシ、レートリミッターがすべて独立してスケールし、CPUエンドポイントで10万以上のQPSを維持します
• 高速で弾力性のあるスケーリング：過剰なプロビジョニングを行うことなく、トラフィックの急増と減少にシステムが適応し、実際の需要に合わせてコストを抑えます
• モデルのガバナンスと監視：AI Gatewayを使用して、ネットワークアクセスやモデルエンドポイントの権限を管理し、品質を監視します。

Feature Storeとの統合はシームレスです。モデルがMLflowに記録されると、その特徴量の依存関係が記録されます。推論時、Model ServingはLakebaseから必要な特徴量を自動的に検索します。カスタムの検索コードや手動での接続は不要です。RTMによって計算され、Lakebaseに保存された最新の集計値が取得され、推論リクエストと透過的に結合されます。

ストリーミングにとどまらないFeature Store

高性能なリアルタイム機能は、Feature Storeが解決できる課題の一部にすぎません。他の2つの課題についても簡単に検討する価値があります。

ストリーム特徴量のトレーニングデータ

ストリームの保持期間が短いため、ストリーミング特徴量のトレーニングデータ生成は困難であり、別途オフラインストアを維持する必要があります。Databricks Feature Storeは、取り込んだKafkaデータのオフラインコピーを保存することで、この課題を解決します。モデルのトレーニングにおいて、Feature Storeは過去の値に対してストリーミングパイプラインと同じ特徴量を計算し、ポイントインタイム（時点）で正確な結合を行います。この同じ機能は、オンラインストリーミング特徴量のバックフィルにも使用され、本番環境への迅速なローンチを可能にします。

統合

上記のように、Feature Storeはいくつかの複雑なインフラストラクチャコンポーネントをオーケストレーションします。その断片化により、ガバナンス、リネージ、および特徴量の再利用が困難になる可能性があります。また、エンジニアがシステムの境界を越えて変更を調整する必要があるため、開発スピードも低下します。

Databricksでは、特徴量はUnity Catalogの第一級オブジェクトであり、検出可能で、アクセス制御によって管理され、完全なリネージで追跡されます。特徴量の変換処理はモデルとともにパッケージ化され、MLflowがどの特徴量が使用されたかを記録し、デプロイメントリネージがモデルをその特徴量の依存関係に接続します。このプラットフォームは、MLスタック全体の開発、デプロイ、およびガバナンスを行うためのワンストップショップです。

DatabricksのFeature Storeは、Spark RTM、Lakebase、Model Servingなどの主要なビルディングブロックをオーケストレーションするため、インフラストラクチャをご自身で管理することなく、クラス最高のレイテンシとスケールを実現できます。これらの各システムはストリーミングワークロード向けに高度に最適化されており、機械学習の特徴量において200ミリ秒のフレッシュネス（新しさ）を実現します。

ストリーミング特徴量の定義方法については、ストリーミングパイプラインのドキュメントをご覧ください。既存の特徴量で実験を行い、ミリ秒レベルのフレッシュネスによってどれほど強力なシグナルが得られるかを確認してください。

基盤となるテクノロジーについてさらに詳しく知りたい場合は、書き込みを高速化するLakebaseアーキテクチャに関するブログおよびRTMアーキテクチャの解説をご覧ください。

このような課題に取り組みたい方は、採用情報をご覧ください！

(このブログ記事はAI翻訳ツールを使用して翻訳されています) 原文記事