通过查找 K 个最近邻，在 Spanner 中执行向量相似度搜索

本页介绍了如何使用余弦距离、欧几里得距离和点积向量函数在 Spanner 中执行向量相似度搜索，以查找 K 最近邻。此信息适用于 GoogleSQL 方言数据库和 PostgreSQL 方言数据库。在阅读本页内容之前，请务必先了解以下概念：

  • 欧几里得距离：衡量两个向量之间的最短距离。
  • 余弦距离：计算两个向量之间夹角的余弦值。
  • 点积：计算角的余弦值与相应向量大小的乘积。如果您知道数据集中的所有向量嵌入都已归一化，那么可以使用 DOT_PRODUCT() 作为距离函数。
  • K 最近邻 (KNN)：一种监督式机器学习算法，用于解决分类或回归问题。

您可以使用向量距离函数执行 K 最近邻 (KNN) 向量搜索，以实现相似度搜索或检索增强生成等用例。Spanner 支持 COSINE_DISTANCE()、EUCLIDEAN_DISTANCE() 和 DOT_PRODUCT() 函数，这些函数可对向量嵌入进行操作，让您能够找到输入嵌入的 KNN。

例如，在您生成并保存作为向量嵌入的 Spanner 运营数据后，您可以在查询中提供这些向量嵌入作为输入参数，以在 N 维空间中查找最近的向量，从而搜索语义相似或相关的内容。

这三个距离函数都接受 vector1 和 vector2 实参，这两个实参的类型均为 array<>，并且必须具有相同的维度和长度。如需详细了解这些函数，请参阅：

示例

以下示例展示了 KNN 搜索、对分区数据进行 KNN 搜索以及将二级索引与 KNN 搭配使用。

所有示例都使用 EUCLIDEAN_DISTANCE()。您还可以使用 COSINE_DISTANCE()。此外，如果数据集中的所有向量嵌入都已归一化，您可以使用 DOT_PRODUCT() 作为距离函数。

假设有一个 Documents 表，其中包含一个列 (DocEmbedding)，该列包含来自 DocContents 字节列的预计算文本嵌入。

GoogleSQL 

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

PostgreSQL 

CREATE TABLE Documents (
UserId       bigint NOT NULL,
DocId        bigint NOT NULL,
Author       varchar(1024),
DocContents  bytea,
DocEmbedding float4[],
PRIMARY KEY  (UserId, DocId)
);

假设“棒球，但不是职业棒球”的输入嵌入是数组 [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 子句中添加条件，将向量搜索限制为仅搜索部分数据。此功能的一项常见应用是搜索分区数据，例如属于特定 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, ...]        |
+------------+-----------------+-----------------+

后续步骤