PostgreSQLで実践するベクトル検索入門：AI時代のRDBMS活用術

CREATE EXTENSION IF NOT EXISTS vector;

-- サンプルテーブル作成（必要に応じてカスタマイズ）
CREATE TABLE documents (
    id SERIAL PRIMARY KEY,
    title TEXT,
    content TEXT,
    embedding VECTOR(1024)
);

CREATE INDEX ON documents USING hnsw (embedding vector_cosine_ops);

import streamlit as st
import pandas as pd
import numpy as np
from datetime import datetime
import psycopg2
from psycopg2.extras import execute_values
from sqlalchemy.sql import text
from streamlit.logger import get_logger
from embedding import Embedding

logger = get_logger(__name__)
embedding = Embedding()

# ページ設定
st.set_page_config(
    page_title="Vector Database Demo",
    page_icon="🔍",
    layout="wide"
)

def get_embedding(text):
    return embedding.get_embedding([text])[0]

# データベース接続関数
def get_connection():
    return st.connection('postgresql', type='sql')

# メインアプリケーション
def main():
    st.title("📊 Vector Database Demo")
    
    # サイドバーでの操作選択
    operation = st.sidebar.selectbox(
        "操作を選択",
        ["データ表示", "データ追加", "ベクトル検索"]
    )
    
    if operation == "データ表示":
        show_data()
    elif operation == "データ追加":
        add_data()
    elif operation == "ベクトル検索":
        vector_search()

def show_data():
    st.header("📋 登録データ一覧")
    
    conn = get_connection()

    # データ取得
    query = "SELECT id, title,content FROM documents LIMIT 100"
    df = conn.query(query, ttl=0)
    
    if not df.empty:
        st.dataframe(df)
    else:
        st.info("データが登録されていません")

def add_data():
    st.header("➕ データ追加")
    
    # 入力フォーム
    with st.form("data_form"):
        title = st.text_input("タイトルを入力")
        content = st.text_area("テキストを入力")
        submitted = st.form_submit_button("登録")
        
        if submitted and content:
            conn = get_connection()
            
            # サンプルとして、ランダムな1536次元ベクトルを生成
            # 実際のアプリケーションでは、適切なエンベッディングモデルを使用する
            embedding = get_embedding(title + " " + content)
            
            # データ登録
            query = text("""
            INSERT INTO documents (title, content, embedding)
            VALUES (:title, :content, :embedding)
            """)
            params = {"title": title, "content": content, "embedding": embedding}
            
            try:
                with conn.session as session:
                    session.execute(query, params)
                    session.commit()
                st.success("データを登録しました")
            except Exception as e:
                st.error(f"エラーが発生しました: {str(e)}")

def vector_search():
    st.header("🔍 ベクトル検索")
    
    search_text = st.text_input("検索テキストを入力")
    k = st.slider("表示件数", min_value=1, max_value=10, value=5)
    
    if st.button("検索") and search_text:
        # サンプルとして、ランダムなクエリベクトルを生成
        # 実際のアプリケーションでは、入力テキストを適切にエンベッディング
        query_vector = get_embedding(search_text)
        
        conn = get_connection()
        
        # コサイン類似度による検索
        query = """
        SELECT title,content, 1 - (embedding <-> :query_vector) as similarity
        FROM documents
        ORDER BY embedding <-> :query_vector
        LIMIT :k
        """
        params = {"query_vector": str(query_vector), "k": k}
        
        try:
            df = conn.query(query, params=params, ttl=0)
            
            if not df.empty:
                # 結果表示
                for _, row in df.iterrows():
                    with st.expander(f"{row['title']} 類似度: {row['similarity']:.4f}"):
                        st.write(row['content'])
            else:
                st.info("検索結果が見つかりませんでした")
        except Exception as e:
            st.error(f"検索中にエラーが発生しました: {str(e)}")

# アプリケーション実行
if __name__ == "__main__":
    main()

import torch.nn.functional as F

from torch import Tensor
from transformers import AutoTokenizer, AutoModel

class Embedding:
    def __init__(self, model_name: str = 'intfloat/multilingual-e5-large'):
        self.model_name = model_name
        self.load_model()

    def load_model(self):
        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
        self.model = AutoModel.from_pretrained(self.model_name)

    def average_pool(
        self,
        last_hidden_states: Tensor,
        attention_mask: Tensor
    ) -> Tensor:
        last_hidden = last_hidden_states.masked_fill(~attention_mask[..., None].bool(), 0.0)
        return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]

    def get_embedding(self, input_texts: list[str]) -> list[float]:
        batch_dict = self.tokenizer(input_texts, max_length=512, padding=True, truncation=True, return_tensors='pt')
        outputs = self.model(**batch_dict)
        embeddings = self.average_pool(outputs.last_hidden_state, batch_dict['attention_mask'])
        return F.normalize(embeddings, p=2, dim=1).tolist()

import psycopg2
import numpy as np
from psycopg2.extras import execute_values
from embedding import Embedding

test_data = [
    {
        "title": "Dockerコンテナのベストプラクティス2025年版",
        "content": "Dockerコンテナを本番環境で効率的に運用するためのベストプラクティスを解説します。イメージサイズの最適化、セキュリティ対策、ネットワーク設定、ボリューム管理など、実践的なトピックを網羅的にカバーします。マルチステージビルドの活用方法や、環境変数の適切な管理方法についても詳しく説明します。"
    },
    {
        "title": "PyTorchによる深層学習モデルの最適化手法",
        "content": "PyTorchを使用した深層学習モデルのパフォーマンス最適化について解説します。バッチサイズの調整、学習率スケジューリング、データローダーの最適化、GPUメモリの効率的な使用方法など、実践的な最適化テクニックを紹介します。"
    },
    {
        "title": "マイクロサービスアーキテクチャの設計パターン",
        "content": "マイクロサービスアーキテクチャを採用する際の主要な設計パターンについて解説します。サービス間通信、データ一貫性の確保、障害対策、モニタリング戦略など、実装時の重要なポイントを詳しく説明します。"
    },
    {
        "title": "Kubernetes運用管理の実践ガイド",
        "content": "Kubernetesクラスタの効率的な運用管理方法について解説します。リソース管理、オートスケーリング、モニタリング、セキュリティ対策など、実運用で必要となる知識を体系的に説明します。"
    },
    {
        "title": "効率的なデータベースインデックス設計",
        "content": "リレーショナルデータベースにおけるインデックス設計のベストプラクティスを解説します。クエリパフォーマンスの最適化、インデックス選択の基準、メンテナンス戦略など、実践的なアプローチを紹介します。"
    },
    {
        "title": "GraphQLによるモダンAPIの構築",
        "content": "GraphQLを使用したAPIの設計と実装について解説します。スキーマ設計、リゾルバの実装、N+1問題の解決、キャッシュ戦略など、実践的な開発手法を紹介します。"
    },
    {
        "title": "CI/CDパイプラインの自動化戦略",
        "content": "継続的インテグレーション/デリバリーパイプラインの効率的な構築方法について解説します。テスト自動化、デプロイ戦略、品質管理、モニタリングなど、実践的な自動化手法を紹介します。"
    },
    {
        "title": "セキュアなWebアプリケーション開発",
        "content": "Webアプリケーションのセキュリティ対策について包括的に解説します。XSS対策、CSRF対策、認証・認可の実装、セキュアなセッション管理など、重要なセキュリティ考慮事項を説明します。"
    },
    {
        "title": "効率的なキャッシュ戦略の実装",
        "content": "Webアプリケーションにおけるキャッシュ戦略の設計と実装について解説します。CDN、ブラウザキャッシュ、アプリケーションキャッシュ、データベースキャッシュなど、多層的なキャッシュ戦略を紹介します。"
    },
    {
        "title": "大規模データ処理のベストプラクティス",
        "content": "大規模データ処理システムの設計と実装について解説します。バッチ処理、ストリーム処理、データパイプライン、スケーラビリティ確保など、実践的なアプローチを紹介します。"
    },
    {
        "title": "ReactとTypeScriptによるフロントエンド開発",
        "content": "ReactとTypeScriptを組み合わせた最新のフロントエンド開発手法について解説します。型安全な開発、コンポーネント設計、状態管理、パフォーマンス最適化など、実践的な開発テクニックを紹介します。"
    },
    {
        "title": "AWSでのスケーラブルなインフラ構築",
        "content": "AWSを使用したスケーラブルなインフラストラクチャの構築方法について解説します。オートスケーリング、負荷分散、障害対策、コスト最適化など、クラウドインフラの設計ポイントを説明します。"
    },
    {
        "title": "効率的なログ管理とモニタリング",
        "content": "分散システムにおけるログ管理とモニタリングの実践的アプローチについて解説します。ログ収集、分析、可視化、アラート設定など、効果的な運用監視の方法を紹介します。"
    },
    {
        "title": "マイクロフロントエンドアーキテクチャの実装",
        "content": "マイクロフロントエンドアーキテクチャの設計と実装について解説します。モジュール分割、統合戦略、ルーティング、状態管理など、フロントエンド開発の新しいアプローチを紹介します。"
    },
    {
        "title": "NoSQLデータベースの設計パターン",
        "content": "NoSQLデータベースを使用する際の効果的な設計パターンについて解説します。データモデリング、クエリ最適化、スケーリング戦略など、実践的な使用方法を紹介します。"
    },
    {
        "title": "機械学習モデルの本番環境デプロイ",
        "content": "機械学習モデルを本番環境にデプロイする際の実践的アプローチについて解説します。モデルのバージョン管理、スケーリング、モニタリング、再学習戦略など、運用上の重要ポイントを説明します。"
    },
    {
        "title": "Terraformによるインフラのコード化",
        "content": "Terraformを使用したインフラストラクチャのコード化について解説します。リソース管理、モジュール設計、状態管理、チーム開発など、IaCの実践的な適用方法を紹介します。"
    },
    {
        "title": "効率的なAPIバージョニング戦略",
        "content": "WebAPIのバージョニング戦略について実践的な方法を解説します。バージョン管理手法、下位互換性の確保、マイグレーション戦略など、長期的な API 運用のポイントを説明します。"
    },
    {
        "title": "セキュアなマイクロサービス間通信",
        "content": "マイクロサービス環境におけるセキュアな通信方法について解説します。認証、認可、暗号化、証明書管理など、サービス間通信のセキュリティ確保方法を説明します。"
    },
    {
        "title": "効率的なデータベースマイグレーション",
        "content": "大規模データベースのマイグレーション戦略について解説します。ダウンタイム最小化、データ整合性確保、ロールバック計画など、安全なマイグレーションの実施方法を紹介します。"
    }
]

def insert_test_data():
    conn = psycopg2.connect(
        dbname="vectordb",
        user="postgres",
        password="postgres",
        host="postgres",
        port="5432"
    )
    cur = conn.cursor()
    embedding = Embedding()

    for data in test_data:
        # サンプルとして1536次元のランダムベクトルを生成
        emb = embedding.get_embedding([data["title"] + " " + data["content"]])[0]
        
        cur.execute(
            "INSERT INTO documents (title, content, embedding) VALUES (%s, %s, %s)",
            (data["title"], data["content"], emb)
        )

    conn.commit()
    cur.close()
    conn.close()

if __name__ == "__main__":
    insert_test_data()

[connections.postgresql]
dialect = "postgresql"
host = "postgres"
port = 5432
database = "vectordb"
username = "postgres"
password = "postgres"

SQLAlchemy==2.0.35
streamlit==1.32.0
pandas==2.2.0
numpy==1.26.0
psycopg2-binary==2.9.9
torch==2.6.0
transformers==4.48.2

FROM postgres:16.3

# pgvectorインストール
RUN apt-get update && \\
    apt-get install -y \\
    build-essential \\
    git \\
    postgresql-server-dev-16

RUN git clone <https://github.com/pgvector/pgvector.git> /tmp/pgvector && \\
    cd /tmp/pgvector && \\
    make && \\
    make install && \\
    rm -rf /tmp/pgvector

FROM python:3.11-slim

WORKDIR /app

COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

CMD ["streamlit", "run", "./streamlit/app.py", "--server.port=8501"]

version: "3.9"

services:
  postgres:
    build: ./postgres
    environment:
      POSTGRES_USER: postgres
      POSTGRES_PASSWORD: postgres
      POSTGRES_DB: vectordb
    volumes:
      - postgres_data:/var/lib/postgresql/data
      - ./postgres/initdb:/docker-entrypoint-initdb.d
    ports:
      - "5432:5432"
    networks:
      - app-network

  streamlit:
    build: ./streamlit
    volumes:
      - .:/app
    environment:
      - STREAMLIT_SERVER_PORT=8501
    ports:
      - "8501:8501"
    depends_on:
      - postgres
    networks:
      - app-network

volumes:
  postgres_data:

networks:
  app-network:
    driver: bridge

docker compose build

docker compose up -d

$ docker compose ps
NAME                   IMAGE                COMMAND                  SERVICE     CREATED       STATUS       PORTS
pgvector-postgres-1    pgvector-postgres    "docker-entrypoint.s…"   postgres    2 hours ago   Up 2 hours   0.0.0.0:5432->5432/tcp, :::5432->5432/tcp
pgvector-streamlit-1   pgvector-streamlit   "streamlit run ./str…"   streamlit   2 hours ago   Up 2 hours   0.0.0.0:8501->8501/tcp, :::8501->8501/tcp

$ docker exec -it pgvector-streamlit-1 python streamlit/seed.py