PostgreSQLで直接機械学習:SQLで学習と予測を実行する

AI tutorial - IT technology blog
AI tutorial - IT technology blog

従来の機械学習パイプラインに潜む摩擦

アプリケーションにAI機能を追加しようとすると、まずお決まりの悩みに直面します。本番環境のPostgreSQLデータベースにはクリーンなユーザーデータが蓄積されていますが、MLモデルは別のPython環境やクラウドサービス上に存在します。これらを接続するためにETL(抽出・変換・格納)パイプラインを構築しますが、チームメンバーがテーブルスキーマを変更した瞬間に、そのパイプラインは壊れてしまうことが多々あります。

私はこれまで、トレーニングデータと本番データの同期が取れなくなったことによるパイプラインの失敗のデバッグに、週末を丸ごと費やしたことが何度もあります。ネットワーク転送を経てデータがScikit-learnやXGBoostモデルに到達する頃には、そのデータはすでに鮮度を失っています。このタイムラグにより、クレジットカードの不正利用を200ミリ秒以内にブロックするといったリアルタイム機能の維持は、ほぼ不可能になります。

なぜデータの移動が真のボトルネックなのか

核心的な問題はモデルの背後にある数学ではなく、「データグラビティ(データの重力)」にあります。大規模なデータセットをネットワーク境界を越えて移動させると、主に3つの失敗要因が生じます。

  • レイテンシ(遅延): 行をJSONにシリアライズし、ネットワーク経由で送信してPythonのDataFrameにロードする作業は、大きなオーバーヘッドを生みます。100MBのデータセットでも、数秒の遅延が簡単に発生します。
  • セキュリティリスク: 暗号化されたデータベースからデータを取り出し、S3バケットやローカルのCSVファイルに保存するたびに、攻撃対象領域(アタックサーフェス)が拡大します。
  • インフラの肥大化: データベース、フィーチャーストア、トレーニングサーバー、そしてFastAPIのようなモデルサービングAPIを管理し続けることになります。

コンピューティングとストレージを分離すると、高い「アーキテクチャ税」を支払うことになります。データベースはすでに数百万行のスキャンに最適化されています。多くの場合、データをモデルに移動させるよりも、モデルをデータのある場所に移動させる方が高速です。

アプローチの比較:Pythonスクリプト vs データベース内ML

インテリジェンスを実装する際、ほとんどのチームは以下の3つの道のいずれかを選択します。

1. 従来のPythonによる手法

エンジニアは psycopg2 を使ってデータを取得し、pickle を使ってモデルを保存します。柔軟性はありますが、このアプローチでは環境の安定性を保つために膨大な「グルーコード(接着剤となるコード)」が必要です。Pythonのバージョンが変わるだけで、モデルがロードできなくなる可能性もあります。

2. クラウドネイティブな手法(SageMaker/Vertex AI)

マネージドサービスは堅牢ですが、コストが高く、ベンダーロックインの問題も伴います。また、データベースとMLサービスを24時間365日完璧に同期させなければならないという同期の問題は依然として残ります。

3. データベース内での手法(PostgresML)

この手法では、機械学習エンジンをPostgreSQL内部に組み込みます。リモートAPIを呼び出す代わりに、SELECT pgml.predict() を実行します。これにより、スタックからETLレイヤーを完全に排除できます。

PostgresMLの実装

PostgresMLは、Scikit-learn、XGBoost、LightGBMをデータベースプロセスに直接取り込む拡張機能です。MLモデルを標準的なデータベースオブジェクトとして扱います。私の本番環境でのテストでは、この構成により、リードスコアリングのような構造化データタスクのデプロイ時間が数日から数分に短縮されました。

環境のセットアップ

事前に設定済みのDockerコンテナを使用することで、C++の依存関係をコンパイルする手間を省けます。これにはPostgreSQLとPostgresML拡張機能が最初から含まれています。

docker run \
  -it \
  -v postgresml_data:/var/lib/postgresql \
  -p 5432:5432 \
  -p 8000:8000 \
  ghcr.io/postgresml/postgresml:latest

コンテナが起動したら、任意のSQLクライアントで接続し、拡張機能を有効にします:

CREATE EXTENSION IF NOT EXISTS pgml;

SQLでのモデルトレーニング

例えば、user_activity テーブルがあるとします。ログイン頻度に基づいて、どのユーザーがサブスクリプションを解約(チャーン)するかを予測したいと考えています。5GBのCSVファイルをエクスポートする代わりに、単一のクエリを実行します:

SELECT * FROM pgml.train(
    'チャーン予測モデル',
    'classification',
    'user_activity_table',
    'churn_label',
    'xgboost'
);

PostgresMLが重い処理をすべて引き受けます。データのスナップショットを作成し、トレーニング用とテスト用に分割し、モデルの重みをシステムテーブルに直接保存します。外部ファイルは一切不要です。

リアルタイム予測の実行

トレーニングが完了すると、モデルはすぐに利用可能になります。アプリケーションの既存のクエリに予測を直接組み込むことができます。

SELECT 
    email, 
    pgml.predict('チャーン予測モデル', ARRAY[login_count, support_tickets, days_active]) AS churn_probability
FROM users
WHERE user_id = 12345;

予測はデータベース内で行われるため、ネットワーク遅延は実質ゼロです。アプリのダッシュボードにリスクスコアを表示する必要がある場合、バックエンドは標準的なSQLクエリを実行するだけです。外部のAIサービスと連携する必要はありません。

特徴量エンジニアリングの処理

生データを特徴量に変換する作業は、多くの場合、全工程の80%を占めます。PostgresMLでは、これにSQLビューを活用します。30日間の平均支出額を計算するといったロジックを標準SQLで定義し、そのビューをトレーニングジョブの対象として指定できます。

CREATE VIEW user_features AS
SELECT 
    user_id,
    extract(day from now() - created_at) as account_age, -- アカウント経過日数
    count(logs.id) as total_logins -- 合計ログイン回数
FROM users
JOIN logs ON logs.user_id = users.id
GROUP BY users.id;

本番環境でのパフォーマンスと安定性

機械学習を実行するとデータベースがクラッシュするのではないかという懸念はよく聞かれます。しかし、PostgresMLはPostgreSQLのメモリ管理を共有しつつ、別のプロセス空間で実行されます。私の経験では、標準的な4コアのデータベースインスタンスで、10ミリ秒未満のレイテンシを維持しながら、毎秒5,000件以上の予測を処理できました。

さらにスケールさせるには、リードレプリカ戦略を使用します。オフピーク時にプライマリノードで重いトレーニングを行い、モデルテーブルをリードオンリーのインスタンスに複製することで、クラスター全体で高速な推論を実現できます。

最後に

「データのある場所にコンピューティングを移動させる」ことは、バックエンドエンジニアにとって大きな転換です。PostgresMLを使用すれば、AIプロジェクトにありがちな脆弱なインフラを構築することなく、本番グレードのモデルをデプロイできます。SELECT 文が書けるのであれば、今や誰でも機械学習を構築し、デプロイすることができるのです。

Share: