pgvectorを触ってみる

基本情報

• ホームページ
  • https://github.com/pgvector/pgvector
• pgvectorとは
  • ベクトルデータの類似性検索機能を提供するPostgreSQLの拡張機能
• AWS対応状況
  • Amazon Aurora Postgre の拡張機能バージョンSQL
  • Aurora PostgreSQL 16.4 -> pgvector 0.7.3

---

インストール方法(Dockerfile)

FROM postgres:16.3-bookworm

# パッケージ更新＆パッケージのビルドで必要なもの
RUN apt-get update && apt-get install -y \
    git \
    make \
    gcc \
    postgresql-server-dev-16

WORKDIR /usr/local/src
RUN git clone --branch v0.7.0 https://github.com/pgvector/pgvector.git
WORKDIR /usr/local/src/pgvector
RUN export PATH=${PATH}:/usr/local/pgsql164/bin
RUN make
RUN make install

---

インストール方法(docker-compose.yml)

    volumes:
      - ./local/postgres/initdb:/docker-entrypoint-initdb.d

• local/postgres/initdb/init_extension_vector.sql

CREATE EXTENSION IF NOT EXISTS vector;

---

基本的な使い方

• ベクトルカラムを含んでテーブル作成

CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3));

• 値をインサートする

INSERT INTO items (embedding) VALUES ('[1,2,3]'), ('[4,5,6]');

• 検索

SELECT * FROM items ORDER BY embedding <-> '[3,1,2]' LIMIT 5;

---

Indexing 1

インデックス	説明	インデックス構築時間	メモリ使用量	クエリのパフォーマンス
IVFFlat	ベクトルをリストに分割してクエリベクトルに最も近いリストのサブセットを検索する	短い	少ない	速度とリコールのトレードオフが大きい
HNSW	最下層に全てのベクトル点が含まれ上の層になるほどベクトル点が少なくなる層状のグラフ構造で上層から下層にアイテムを辿っていく	長い	多い	速度とリコールのトレードオフが少ない

• 参照: https://www.sraoss.co.jp/tech-blog/pgsql/pgvector-intro/#:~:text=pgvector(pgvector%2Fpgvector%3A%20Open,%E8%A6%81%E7%B4%A0%E3%82%92%E6%A4%9C%E7%B4%A2%E3%81%A7%E3%81%8D%E3%81%BE%E3%81%99%E3%80%82

---

Indexing 2

• 作成

CREATE INDEX ON items USING hnsw (embedding vector_l2_ops);

• sqlalchemyで作成してみる

class Rag(Base):
    __table_args__ = (
        Index(
            'embedding_index',
            'embedding',
            postgresql_using='hnsw',
            postgresql_with={'m': 16, 'ef_construction': 64},
            postgresql_ops={'embedding': 'vector_l2_ops'}
        ),
    )

    id = Column(Integer, primary_key=True, index=True)
    embedding = Column(Vector)
    text = Column(String)
    file_id = Column(Integer)

---

ベクトルサイズ変更 1

• 問題
  • embeddingする時に、モデルによってベクトルサイズが違い、例：
    • Bedrock(amazon.titan-embed-text-v2:0): 1024
    • AzureOpenAI(text-embedding-ada-002): 1536
  • embedding用のモデルを切り替えると、pgvectorのカラム(Vector)のサイズも更新しないといけない
• 対応
  • テーブル作成する時に、vectorのサイズを指定しないように
  • 検索するときにvector_dims(embedding) = 5のような条件を追加する
    • Bedrock: 1024
    • Azure: 1536

---

ベクトルサイズ変更 2

• サンプル

vec_test=# create table vecs2 (id bigserial PRIMARY KEY, embedding vector);
CREATE TABLE
vec_test=# insert into vecs2 (embedding) values ('[1,2,3,4,5]');
INSERT 0 1
vec_test=# insert into vecs2 (embedding) values ('[1,2,3]');
INSERT 0 1
vec_test=# select * from vecs2 where vector_dims(embedding) = 3;
 id | embedding
----+-----------
  2 | [1,2,3]
(1 row)

vec_test=# select * from vecs2 where vector_dims(embedding) = 5;
 id |  embedding
----+-------------
  1 | [1,2,3,4,5]
(1 row)

---

ご清聴ありがとうございます