Amazon RDS for PostgreSQLがpgvectorモジュールに対応しベクトル検索できるようになりました

2023.05.04

近年は機械学習のモデル・アルゴリズムを使ってオブジェクトをベクトル表現し(embedding)、テキストやメディアなど様々なオブジェクトの検索や分類が盛んに行われており、LLMの流行と共にその勢いが増しています。

RDBのPostgreSQL向けにベクトル管理・検索を行う拡張モジュール pgvector がAmazon RDS for RDSでも利用できるようになり、Pineconeのようなベクトル検索特化型のデータベースを用意することなく、RDSインスタンス単体でベクトル検索ができるようになりました。

データベースのマネージドサービスとしては、これまでも Supabase 等が pgvector に対応していましたが、RDS PostgreSQL もついに対応しました。

制限

pgvector

  • 対応する最近傍探索
    • exact
    • Approximation
  • 対応距離
    • ユークリッド(L2)距離
    • 内積
    • コサイン類似度
  • 次元
    • 2000次元まで

詳細は GitHub のpgvectorのレポジトリを参照してください。

AWS

  • RDS PostgreSQL 15.2 以降で対応
  • pgvector 0.41
  • Aurora PostgreSQL は現時点で未対応

pgvector の有効化

pgvector モジュールを利用するには、15.2以上のRDS PostgreSQL を起動し、モジュールを有効化するだけです。

test=> select version();
                                                   version
--------------------------------------------------------------------------------------------------------------
 PostgreSQL 15.2 on aarch64-unknown-linux-gnu, compiled by gcc (GCC) 7.3.1 20180712 (Red Hat 7.3.1-6), 64-bit
(1 row)

test=> SELECT * FROM pg_available_extensions where name like '%vector%';
  name  | default_version | installed_version |                  comment
--------+-----------------+-------------------+--------------------------------------------
 vector | 0.4.1           |                   | vector data type and ivfflat access method
(1 row)

test=> CREATE EXTENSION vector;
CREATE EXTENSION

test=> \dx
                        List of installed extensions
  Name   | Version |   Schema   |                Description
---------+---------+------------+--------------------------------------------
 plpgsql | 1.0     | pg_catalog | PL/pgSQL procedural language
 vector  | 0.4.1   | public     | vector data type and ivfflat access method
(2 rows)

pgvector のベクトル検索

ベクトルは vector 型のカラムで管理します。2,000次元まで対応しています。

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

test=> \d items
                               Table "public.items"
  Column   |   Type    | Collation | Nullable |              Default
-----------+-----------+-----------+----------+-----------------------------------
 id        | bigint    |           | not null | nextval('items_id_seq'::regclass)
 embedding | vector(3) |           |          |
Indexes:
    "items_pkey" PRIMARY KEY, btree (id)

実際に embedded ベクトルを挿入してみましょう。

test=> INSERT INTO items (embedding) VALUES ('[1,2,3]'), ('[4,5,6]'), ('[-1,0,1]');
INSERT 0 3

test=> select * from items;
 id | embedding
----+-----------
  1 | [1,2,3]
  2 | [4,5,6]
  3 | [-1,0,1]
(3 rows)

pgvector は以下の距離に対応し、専用の演算子が割り当てられています。

  • ユークリッド(L2) : <->
  • 内積の負 : <#>
  • コサイン : <=>

それぞれの距離を実際に計算し、ユークリッド距離でソートします。

SELECT *, 
  embedding <-> '[3,1,2]' as l2,
  embedding <#> '[3,1,2]' as inner_product,
  embedding <=> '[3,1,2]' as cos
FROM items
ORDER BY embedding <-> '[3,1,2]' asc;

 id | embedding |        l2         | inner_product |         cos
----+-----------+-------------------+---------------+---------------------
  1 | [1,2,3]   | 2.449489742783178 |           -11 |  0.2142857142857143
  3 | [-1,0,1]  | 4.242640687119285 |             1 |  1.1889822365046137
  2 | [4,5,6]   | 5.744562646538029 |           -29 | 0.11673988938389968
(3 rows)

インデックス済みの既存オブジェクト(id=1)に近いオブジェクトを探すSQL

SELECT * FROM items
WHERE id != 1
ORDER BY embedding <-> (SELECT embedding FROM items WHERE id = 1)
LIMIT 3;

 id | embedding
----+-----------
  3 | [-1,0,1]
  2 | [4,5,6]
(2 rows)

演算子とは別に距離ごとに関数も定義されています。

内積を計算する関数(inner_product)は定義通りに計算するのに対して、内積の演算子(<#>)は、計算結果を負にする点については注意が必要です。「負」にしているのは、PostgreSQLは演算子に対してASC順インデックススキャンのみをサポートするためです。

test=> select inner_product('[1,2,3]', '[3,1,2]'), '[1,2,3]' <#> '[3,1,2]';
 inner_product | ?column?
---------------+----------
            11 |      -11
(1 row)

embedding方法

ベクトル表現の管理・検索はPostgreSQLでこなせるとして、embedding 化は PostgreSQL の外部で行う必要があります。

AWSの公式ブログでは、SageMaker(sentence-transformer モデル)を使ったインテグレーションが紹介されています。

Building AI-powered search in PostgreSQL using Amazon SageMaker and pgvector | AWS Database Blog

※ 図は公式ブログから

よりお手軽な方法としては、 OpenAI の Embedding API があります。

Embeddings - OpenAI API

最後に

RDS PostgreSQL が pgvectorモジュールに対応しベクトル検索できるようになりました。

ベクトルは2000次元の制約があるため、現時点で大規模な利用は難しいかもしれませんが、既存のシステムにベクトル検索をシームレスに組み込みたい場合、慣れ親しんだ PostgreSQL を活用するのは非常に有力な選択肢の一つと思えます。

Aurora PostgreSQL の pgvector 対応も待ち遠しいです。

GitHub には pgvector に対応したマネージドサービスをトラッキングするイシューも立っています。

Hosted Providers · Issue #54 · pgvector/pgvector · GitHub

それでは。

参考