使用向量嵌入搜索

本页面介绍如何通过以下方法使用 Cloud Firestore 执行 K 最近邻 (KNN) 向量搜索:

  • 存储向量值
  • 创建和管理 KNN 向量索引
  • 使用一个支持的向量距离函数执行 K 最近邻 (KNN) 查询

存储向量嵌入

您可以使用 Cloud Firestore 数据创建向量值(例如文本嵌入),并将它们存储在 Cloud Firestore 文档中。

通过向量嵌入执行写入操作

以下示例展示了如何将向量嵌入存储在 Cloud Firestore 文档中:

Python
from google.cloud import firestore
from google.cloud.firestore_v1.vector import Vector

firestore_client = firestore.Client()
collection = firestore_client.collection("coffee-beans")
doc = {
  "name": "Kahawa coffee beans",
  "description": "Information about the Kahawa coffee beans.",
  "embedding_field": Vector([1.0 , 2.0, 3.0])
}

collection.add(doc)
    
Node.js
import {
  Firestore,
  FieldValue,
} from "@google-cloud/firestore";

const db = new Firestore();
const coll = db.collection('coffee-beans');
await coll.add({
  name: "Kahawa coffee beans",
  description: "Information about the Kahawa coffee beans.",
  embedding_field: FieldValue.vector([1.0 , 2.0, 3.0])
});
    

使用 Cloud Functions 函数计算向量嵌入

如要在文档每次创建或更新时便计算并存储相应的向量嵌入,您可以设置一个 Cloud Functions 函数

Python
@functions_framework.cloud_event
def store_embedding(cloud_event) -> None:
  """Triggers by a change to a Firestore document.
  """
  firestore_payload = firestore.DocumentEventData()
  payload = firestore_payload._pb.ParseFromString(cloud_event.data)

  collection_id, doc_id = from_payload(payload)
  # Call a function to calculate the embedding
  embedding = calculate_embedding(payload)
  # Update the document
  doc = firestore_client.collection(collection_id).document(doc_id)
  doc.set({"embedding_field": embedding}, merge=True)
    
Node.js
/**
 * A vector embedding will be computed from the
 * value of the `content` field. The vector value
 * will be stored in the `embedding` field. The
 * field names `content` and `embedding` are arbitrary
 * field names chosen for this example.
 */
async function storeEmbedding(event: FirestoreEvent<any>): Promise<void> {
  // Get the previous value of the document's `content` field.
  const previousDocumentSnapshot = event.data.before as QueryDocumentSnapshot;
  const previousContent = previousDocumentSnapshot.get("content");

  // Get the current value of the document's `content` field.
  const currentDocumentSnapshot = event.data.after as QueryDocumentSnapshot;
  const currentContent = currentDocumentSnapshot.get("content");

  // Don't update the embedding if the content field did not change
  if (previousContent === currentContent) {
    return;
  }

  // Call a function to calculate the embedding for the value
  // of the `content` field.
  const embeddingVector = calculateEmbedding(currentContent);

  // Update the `embedding` field on the document.
  await currentDocumentSnapshot.ref.update({
    embedding: embeddingVector,
  });
}
    

创建和管理向量索引

您必须先创建相应的索引,然后才能通过向量嵌入执行最近邻搜索。以下示例展示了如何创建和管理向量索引。

创建矢量索引

要创建向量索引,请使用 gcloud alpha firestore indexes composite create

gcloud
gcloud alpha firestore indexes composite create \
--collection-group=collection-group \
--query-scope=COLLECTION \
--field-config field-path=vector-field,vector-config='vector-configuration' \
--database=database-id
    

其中:

  • collection-group 是集合组的 ID。
  • vector-field 是包含向量嵌入的字段的名称。
  • database-id 是相应数据库的 ID。
  • vector-configuration 包含向量 dimension 和索引类型。dimension 是一个不超过 2,048 的整数。索引类型必须为 flat。按如下方式设置索引配置的格式:{"dimension":"DIMENSION", "flat": "{}"}

以下示例创建了一个复合索引,其中包含字段 vector-field 的向量索引以及字段 color 的升序索引。您可以在执行最近邻搜索之前使用此类索引预先过滤数据

gcloud
gcloud alpha firestore indexes composite create \
--collection-group=collection-group \
--query-scope=COLLECTION \
--field-config=order=ASCENDING,field-path="color" \
--field-config field-path=vector-field,vector-config='{"dimension":"1024", "flat": "{}"}' \
--database=database-id
    

列出所有向量索引

gcloud
gcloud alpha firestore indexes composite list --database=database-id

database-id 替换为相应数据库的 ID。

删除矢量索引

gcloud
gcloud alpha firestore indexes composite delete index-id --database=database-id
    

其中:

  • index-id 是要删除的索引的 ID。可使用 indexes composite list 检索索引 ID。
  • database-id 是相应数据库的 ID。

描述向量索引

gcloud
gcloud alpha firestore indexes composite describe index-id --database=database-id
    

其中:

  • index-id 是要描述的索引的 ID。可使用 indexes composite list 检索索引 ID。
  • database-id 是相应数据库的 ID。

执行最近邻查询

您可以执行相似度搜索来查找向量嵌入的最近邻。相似度搜索需要使用向量索引。如果没有现成的索引,Cloud Firestore 会使用 gcloud CLI 建议一个可创建的索引。

Python
from google.cloud.firestore_v1.base_vector_query import DistanceMeasure

collection = collection("coffee-beans")

# Requires vector index
collection.find_nearest(
   vector_field="embedding_field",
   query_vector=Vector([3.0, 1.0, 2.0]),
   distance_measure=DistanceMeasure.EUCLIDEAN,
   limit=5)
    
Node.js
import {
  Firestore,
  FieldValue,
  VectorQuery,
  VectorQuerySnapshot,
} from "@google-cloud/firestore";

// Requires single-field vector index
const vectorQuery: VectorQuery = coll.findNearest('embedding_field', FieldValue.vector([3.0, 1.0, 2.0]), {
  limit: 5,
  distanceMeasure: 'EUCLIDEAN'
});

const vectorQuerySnapshot: VectorQuerySnapshot = await vectorQuery.get();
    

向量距离

最近邻查询支持下列向量距离选项:

  • EUCLIDEAN:测量向量之间的欧几里得距离。如需了解详情,请参阅欧几里得
  • COSINE:基于向量之间的角度来比较向量,这样可以测量不依赖于向量大小的相似度。对于单位归一化向量,建议使用 DOT_PRODUCT,而不是余弦距离,虽然两者在数学上是等效的,但前者性能更好。如需了解详情,请参阅余弦相似度
  • DOT_PRODUCT:与 COSINE 类似,但受向量大小影响。如需了解详情,请参阅点积

预过滤数据

如要在查找最近邻之前预过滤数据,您可以将相似度搜索与其他过滤条件(不等式过滤条件除外)结合使用。支持 andor 复合过滤条件。对于字段过滤条件,支持以下过滤条件:

  • == 等于
  • in
  • array_contains
  • array_contains_any
Python
# Similarity search with pre-filter
# Requires composite vector index
collection.where("color", "==", "red").find_nearest(
   vector_field="embedding_field",
   query_vector=Vector([3.0, 1.0, 2.0]),
   distance_measure=DistanceMeasure.EUCLIDEAN,
   limit=5)
    
Node.js
// Similarity search with pre-filter
// Requires composite vector index
const preFilteredVectorQuery: VectorQuery = coll
  .where("color", "==", "red")
  .findNearest("embedding_field", FieldValue.vector([3.0, 1.0, 2.0]), {
    limit: 5,
    distanceMeasure: "EUCLIDEAN",
  });

vectorQueryResults = await preFilteredVectorQuery.get();
    

限制

请注意,在使用向量嵌入时,有以下限制:

  • 支持的嵌入维度上限为 2,048。如要存储更大的索引,可使用降维
  • 通过最近邻查询返回的文档数量上限为 1,000。
  • 向量搜索不支持实时快照监听器
  • 不能使用不等式过滤条件来预过滤数据。
  • 只有 Python 和 Node.js 客户端库支持向量搜索。

后续步骤