k 近傍法を探索して Spanner で類似度ベクトル検索を行う

このページでは、コサイン距離ベクトル関数、ユークリッド距離ベクトル関数、ドット積ベクトル関数を使用して k 近傍法を探索し、Spanner で類似度ベクトル検索を行う方法について説明します。このページを読む前に、次のコンセプトを理解しておく必要があります。

  • ユークリッド距離: 2 つのベクトル間の最短距離を測定します。
  • コサイン距離: 2 つのベクトル間の角度のコサインを測定します。
  • ドット積: 角度のコサインに、対応するベクトルの大きさを掛けます。データセット内のすべてのベクトル エンベディングが正規化されていることがわかっている場合は、DOT_PRODUCT() を距離関数として使用できます。
  • K 近傍法（KNN）: 分類や回帰の問題を解決するために使用される教師あり機械学習アルゴリズム。

ベクトル距離関数を使用して、類似度検索や検索拡張生成などのユースケースで K 近傍法（KNN）ベクトル検索を実行できます。Spanner はベクトル エンベディングを操作する COSINE_DISTANCE() 関数、EUCLIDEAN_DISTANCE() 関数、DOT_PRODUCT() 関数をサポートしています。これにより、入力エンベディングの KNN を探索できます。

たとえば、運用 Spanner データをベクトル エンベディングとして生成して保存した後、これらのベクトル エンベディングをクエリの入力パラメータとして指定して、N 次元空間内で最も近いベクトルを探索し、意味的に類似または関連するアイテムを検索できます。

3 つのすべての距離関数は、タイプ array<> の引数 vector1 と vector2 を使用し、ディメンションと長さが同じである必要があります。これらの関数の詳細については、以下をご覧ください。

例

次の例は、KNN 検索、パーティション分割データの KNN 検索、KNN でのセカンダリ インデックスの使用を示しています。

例はすべて EUCLIDEAN_DISTANCE() を使用しています。COSINE_DISTANCE() を使用することもできます。さらに、データセット内のすべてのベクトル エンベディングが正規化されている場合は、DOT_PRODUCT() を距離関数として使用できます。

DocContents バイト列の事前計算されたテキスト エンベディングの列（DocEmbedding）を含む Documents テーブルについて考えてみます。

GoogleSQL 

CREATE TABLE Documents (
UserId       INT64 NOT NULL,
DocId        INT64 NOT NULL,
Author       STRING(1024),
DocContents  BYTES,
DocEmbedding ARRAY<FLOAT32>
) PRIMARY KEY (UserId, DocId);

PostgreSQL 

CREATE TABLE Documents (
UserId       bigint primary key,
DocId        bigint primary key,
Author       varchar(1024),
DocContents  bytea,
DocEmbedding float4[]
);

「baseball, but not professional baseball」の入力エンベディングが配列 [0.3, 0.3, 0.7, 0.7] であると仮定すると、次のクエリで、一致する上位 5 つのドキュメントを検索できます。

GoogleSQL 

SELECT DocId, DocEmbedding FROM Documents
ORDER BY EUCLIDEAN_DISTANCE(DocEmbedding,
ARRAY<FLOAT32>[0.3, 0.3, 0.7, 0.8])
LIMIT 5;

PostgreSQL 

SELECT DocId, DocEmbedding FROM Documents
ORDER BY spanner.euclidean_distance(DocEmbedding,
'{0.3, 0.3, 0.7, 0.8}'::float4[])
LIMIT 5;

この例で想定される結果は、次のとおりです。

Documents
+---------------------------+-----------------+
| DocId                     | DocEmbedding    |
+---------------------------+-----------------+
| 24                        | [8, ...]        |
+---------------------------+-----------------+
| 25                        | [6, ...]        |
+---------------------------+-----------------+
| 26                        | [3.2, ...]      |
+---------------------------+-----------------+
| 27                        | [38, ...]       |
+---------------------------+-----------------+
| 14229                     | [1.6, ...]      |
+---------------------------+-----------------+

例 2: パーティション分割データの KNN 検索

前の例のクエリは、WHERE 句に条件を追加し、ベクトル検索をデータのサブセットに限定することで変更できます。この一般的な用途の 1 つは、特定の UserId に属する行などのパーティション分割データを検索することです。

GoogleSQL 

SELECT UserId, DocId, DocEmbedding FROM Documents
WHERE UserId=18
ORDER BY EUCLIDEAN_DISTANCE(DocEmbedding,
ARRAY<FLOAT32>[0.3, 0.3, 0.7, 0.8])
LIMIT 5;

PostgreSQL 

SELECT UserId, DocId, DocEmbedding FROM Documents
WHERE UserId=18
ORDER BY spanner.euclidean_distance(DocEmbedding,
'{0.3, 0.3, 0.7, 0.8}'::float4[])
LIMIT 5;

この例で想定される結果は、次のとおりです。

Documents
+-----------+-----------------+-----------------+
| UserId    | DocId           | DocEmbedding    |
+-----------+-----------------+-----------------+
| 18        | 234             | [12, ...]       |
+-----------+-----------------+-----------------+
| 18        | 12              | [1.6, ...]      |
+-----------+-----------------+-----------------+
| 18        | 321             | [22, ...]       |
+-----------+-----------------+-----------------+
| 18        | 432             | [3, ...]        |
+-----------+-----------------+-----------------+

例 3: セカンダリ インデックスの範囲に対する KNN 検索

使用している WHERE 句フィルタがテーブルの主キーに含まれていない場合は、セカンダリ インデックスを作成して、インデックスのみのスキャンによりオペレーションを高速化できます。

GoogleSQL 

CREATE INDEX DocsByAuthor
ON Documents(Author)
STORING (DocEmbedding);

SELECT Author, DocId, DocEmbedding FROM Documents
WHERE Author="Mark Twain"
ORDER BY EUCLIDEAN_DISTANCE(DocEmbedding,
   <embeddings for "book about the time traveling American">)
LIMIT 5;

PostgreSQL 

CREATE INDEX DocsByAuthor
ON Documents(Author)
INCLUDE (DocEmbedding);

SELECT Author, DocId, DocEmbedding FROM Documents
WHERE Author="Mark Twain"
ORDER BY spanner.euclidean_distance(DocEmbedding,
   <embeddings for "that book about the time traveling American">)
LIMIT 5;

この例で想定される結果は、次のとおりです。

Documents
+------------+-----------------+-----------------+
| Author     | DocId           | DocEmbedding    |
+------------+-----------------+-----------------+
| Mark Twain | 234             | [12, ...]       |
+------------+-----------------+-----------------+
| Mark Twain | 12              | [1.6, ...]      |
+------------+-----------------+-----------------+
| Mark Twain | 321             | [22, ...]       |
+------------+-----------------+-----------------+
| Mark Twain | 432             | [3, ...]        |
+------------+-----------------+-----------------+
| Mark Twain | 375             | [9, ...]        |
+------------+-----------------+-----------------+

次のステップ