使用向量嵌入搜索

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

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

准备工作

在将嵌入存储到 Cloud Firestore 中之前,您必须生成向量嵌入。Cloud Firestore 不会生成嵌入。您可以使用 Vertex AI 等服务从 Cloud Firestore 数据创建向量值(例如文本嵌入)。然后,您可以将这些嵌入重新存储在 Cloud Firestore 文档中。

如需详细了解嵌入,请参阅什么是嵌入?

如需了解如何使用 Vertex AI 获取文本嵌入,请参阅获取文本嵌入

存储向量嵌入

以下示例演示了如何在 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([0.18332680, 0.24160706, 0.3416704]),
}

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])
});
Go
import (
	"context"
	"fmt"
	"io"

	"cloud.google.com/go/firestore"
)

type CoffeeBean struct {
	Name           string             `firestore:"name,omitempty"`
	Description    string             `firestore:"description,omitempty"`
	EmbeddingField firestore.Vector32 `firestore:"embedding_field,omitempty"`
	Color          string             `firestore:"color,omitempty"`
}

func storeVectors(w io.Writer, projectID string) error {
	ctx := context.Background()

	// Create client
	client, err := firestore.NewClient(ctx, projectID)
	if err != nil {
		return fmt.Errorf("firestore.NewClient: %w", err)
	}
	defer client.Close()

	// Vector can be represented by Vector32 or Vector64
	doc := CoffeeBean{
		Name:           "Kahawa coffee beans",
		Description:    "Information about the Kahawa coffee beans.",
		EmbeddingField: []float32{1.0, 2.0, 3.0},
		Color:          "red",
	}
	ref := client.Collection("coffee-beans").NewDoc()
	if _, err = ref.Set(ctx, doc); err != nil {
		fmt.Fprintf(w, "failed to upsert: %v", err)
		return err
	}

	return nil
}
Java
import com.google.cloud.firestore.CollectionReference;
import com.google.cloud.firestore.DocumentReference;
import com.google.cloud.firestore.FieldValue;
import com.google.cloud.firestore.VectorQuery;

CollectionReference coll = firestore.collection("coffee-beans");

Map<String, Object> docData = new HashMap<>();
docData.put("name", "Kahawa coffee beans");
docData.put("description", "Information about the Kahawa coffee beans.");
docData.put("embedding_field", FieldValue.vector(new double[] {1.0, 2.0, 3.0}));

ApiFuture<DocumentReference> future = coll.add(docData);
DocumentReference documentReference = future.get();

使用 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,
  });
}
Go
  // Not yet supported in the Go client library
Java
  // Not yet supported in the Java client library

创建和管理向量索引

您必须先创建相应的索引,然后才能通过向量嵌入执行最近邻搜索。以下示例演示了如何使用 Google Cloud CLI 创建和管理向量索引。您还可以使用 Firebase CLI 和 Terraform 管理向量索引。

创建矢量索引

在创建向量索引之前,请先升级到最新版本的 Google Cloud CLI

gcloud components update

如需创建向量索引,请使用 gcloud firestore indexes composite create

gcloud
gcloud 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 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 firestore indexes composite list --database=database-id

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

删除矢量索引

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

其中:

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

描述向量索引

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

其中:

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

执行最近邻查询

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

以下示例会查找查询向量的 10 个最近邻。

Python
from google.cloud.firestore_v1.base_vector_query import DistanceMeasure
from google.cloud.firestore_v1.vector import Vector

collection = db.collection("coffee-beans")

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

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

const vectorQuerySnapshot: VectorQuerySnapshot = await vectorQuery.get();
Go
import (
	"context"
	"fmt"
	"io"

	"cloud.google.com/go/firestore"
)

func vectorSearchBasic(w io.Writer, projectID string) error {
	ctx := context.Background()

	// Create client
	client, err := firestore.NewClient(ctx, projectID)
	if err != nil {
		return fmt.Errorf("firestore.NewClient: %w", err)
	}
	defer client.Close()

	collection := client.Collection("coffee-beans")

	// Requires a vector index
	// https://firebase.google.com/docs/firestore/vector-search#create_and_manage_vector_indexes
	vectorQuery := collection.FindNearest("embedding_field",
		[]float32{3.0, 1.0, 2.0},
		5,
		// More info: https://firebase.google.com/docs/firestore/vector-search#vector_distances
		firestore.DistanceMeasureEuclidean,
		nil)

	docs, err := vectorQuery.Documents(ctx).GetAll()
	if err != nil {
		fmt.Fprintf(w, "failed to get vector query results: %v", err)
		return err
	}

	for _, doc := range docs {
		fmt.Fprintln(w, doc.Data()["name"])
	}
	return nil
}
Java
import com.google.cloud.firestore.VectorQuery;
import com.google.cloud.firestore.VectorQuerySnapshot;

VectorQuery vectorQuery = coll.findNearest(
        "embedding_field",
        new double[] {3.0, 1.0, 2.0},
        /* limit */ 10,
        VectorQuery.DistanceMeasure.EUCLIDEAN);

ApiFuture<VectorQuerySnapshot> future = vectorQuery.get();
VectorQuerySnapshot vectorQuerySnapshot = future.get();

向量距离

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

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

选择距离衡量方式

根据您的所有向量嵌入是否已归一化,您可以确定要使用哪种距离衡量方式来计算距离衡量值。已归一化向量嵌入的大小(长度)正好为 1.0。

此外,如果您知道用于训练模型的距离衡量方式,请使用该距离衡量方式来计算向量嵌入之间的距离。

已归一化数据

如果您的数据集中所有向量嵌入都已归一化,那么这三种距离衡量方式都会提供相同的语义搜索结果。从本质上讲,虽然每种距离衡量方式都会返回不同的值,但这些值的排序方式相同。如果嵌入已归一化,DOT_PRODUCT 通常具有最高的计算效率,但在大多数情况下,差异可以忽略不计。不过,如果您的应用对性能非常敏感,DOT_PRODUCT 可能会有助于进行性能调优。

未归一化数据

如果您的数据集中的向量嵌入未归一化,那么从数学角度来看,使用 DOT_PRODUCT 作为距离衡量方式是不正确的,因为点积无法衡量距离。根据生成嵌入的方式和偏好的搜索类型,COSINEEUCLIDEAN 距离衡量方式生成的搜索结果在主观上优于其他距离衡量方式。您可能需要对 COSINEEUCLIDEAN 进行实验,以确定哪个最适合您的应用场景。

不确定数据是已归一化还是未归一化

如果您不确定数据是否已归一化,但又想使用 DOT_PRODUCT,我们建议您改用 COSINECOSINE 类似于内置了归一化的 DOT_PRODUCT。使用 COSINE 衡量的距离范围为 02。结果接近 0 表示向量非常相似。

预先过滤文档

如需在查找最近邻之前预先过滤文档,您可以将相似度搜索与其他查询运算符结合使用。支持 andor 复合过滤条件。如需详细了解支持的字段过滤条件,请参阅查询运算符

Python
from google.cloud.firestore_v1.base_vector_query import DistanceMeasure
from google.cloud.firestore_v1.vector import Vector

collection = db.collection("coffee-beans")

# Similarity search with pre-filter
# Requires a composite vector index
vector_query = collection.where("color", "==", "red").find_nearest(
    vector_field="embedding_field",
    query_vector=Vector([0.3416704, 0.18332680, 0.24160706]),
    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({
      vectorField: "embedding_field",
      queryVector: [3.0, 1.0, 2.0],
      limit: 5,
      distanceMeasure: "EUCLIDEAN",
    });

const vectorQueryResults = await preFilteredVectorQuery.get();
Go
import (
	"context"
	"fmt"
	"io"

	"cloud.google.com/go/firestore"
)

func vectorSearchPrefilter(w io.Writer, projectID string) error {
	ctx := context.Background()

	// Create client
	client, err := firestore.NewClient(ctx, projectID)
	if err != nil {
		return fmt.Errorf("firestore.NewClient: %w", err)
	}
	defer client.Close()

	collection := client.Collection("coffee-beans")

	// Similarity search with pre-filter
	// Requires a composite vector index
	vectorQuery := collection.Where("color", "==", "red").
		FindNearest("embedding_field",
			[]float32{3.0, 1.0, 2.0},
			5,
			// More info: https://firebase.google.com/docs/firestore/vector-search#vector_distances
			firestore.DistanceMeasureEuclidean,
			nil)

	docs, err := vectorQuery.Documents(ctx).GetAll()
	if err != nil {
		fmt.Fprintf(w, "failed to get vector query results: %v", err)
		return err
	}

	for _, doc := range docs {
		fmt.Fprintln(w, doc.Data()["name"])
	}
	return nil
}
Java
import com.google.cloud.firestore.VectorQuery;
import com.google.cloud.firestore.VectorQuerySnapshot;

VectorQuery preFilteredVectorQuery = coll
        .whereEqualTo("color", "red")
        .findNearest(
                "embedding_field",
                new double[] {3.0, 1.0, 2.0},
                /* limit */ 10,
                VectorQuery.DistanceMeasure.EUCLIDEAN);

ApiFuture<VectorQuerySnapshot> future = preFilteredVectorQuery.get();
VectorQuerySnapshot vectorQuerySnapshot = future.get();

检索计算出的向量距离

您可以通过对 FindNearest 查询分配 distance_result_field 输出属性名称来检索计算出的向量距离,如以下示例所示:

Python
from google.cloud.firestore_v1.base_vector_query import DistanceMeasure
from google.cloud.firestore_v1.vector import Vector

collection = db.collection("coffee-beans")

vector_query = collection.find_nearest(
    vector_field="embedding_field",
    query_vector=Vector([0.3416704, 0.18332680, 0.24160706]),
    distance_measure=DistanceMeasure.EUCLIDEAN,
    limit=10,
    distance_result_field="vector_distance",
)

docs = vector_query.stream()

for doc in docs:
    print(f"{doc.id}, Distance: {doc.get('vector_distance')}")
Node.js
const vectorQuery: VectorQuery = coll.findNearest(
    {
      vectorField: 'embedding_field',
      queryVector: [3.0, 1.0, 2.0],
      limit: 10,
      distanceMeasure: 'EUCLIDEAN',
      distanceResultField: 'vector_distance'
    });

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

snapshot.forEach((doc) => {
  console.log(doc.id, ' Distance: ', doc.get('vector_distance'));
});
Go
import (
	"context"
	"fmt"
	"io"

	"cloud.google.com/go/firestore"
)

func vectorSearchDistanceResultField(w io.Writer, projectID string) error {
	ctx := context.Background()

	client, err := firestore.NewClient(ctx, projectID)
	if err != nil {
		return fmt.Errorf("firestore.NewClient: %w", err)
	}
	defer client.Close()

	collection := client.Collection("coffee-beans")

	// Requires a vector index
	// https://firebase.google.com/docs/firestore/vector-search#create_and_manage_vector_indexes
	vectorQuery := collection.FindNearest("embedding_field",
		[]float32{3.0, 1.0, 2.0},
		10,
		firestore.DistanceMeasureEuclidean,
		&firestore.FindNearestOptions{
			DistanceResultField: "vector_distance",
		})

	docs, err := vectorQuery.Documents(ctx).GetAll()
	if err != nil {
		fmt.Fprintf(w, "failed to get vector query results: %v", err)
		return err
	}

	for _, doc := range docs {
		fmt.Fprintf(w, "%v, Distance: %v\n", doc.Data()["name"], doc.Data()["vector_distance"])
	}
	return nil
}
Java
import com.google.cloud.firestore.VectorQuery;
import com.google.cloud.firestore.VectorQueryOptions;
import com.google.cloud.firestore.VectorQuerySnapshot;

VectorQuery vectorQuery = coll.findNearest(
        "embedding_field",
        new double[] {3.0, 1.0, 2.0},
        /* limit */ 10,
        VectorQuery.DistanceMeasure.EUCLIDEAN,
        VectorQueryOptions.newBuilder().setDistanceResultField("vector_distance").build());

ApiFuture<VectorQuerySnapshot> future = vectorQuery.get();
VectorQuerySnapshot vectorQuerySnapshot = future.get();

for (DocumentSnapshot document : vectorQuerySnapshot.getDocuments()) {
    System.out.println(document.getId() + " Distance: " + document.get("vector_distance"));
}

如果您想使用字段掩码返回部分文档字段以及 distanceResultField,则还必须在字段掩码中添加 distanceResultField 的值,如以下示例所示:

Python
vector_query = collection.select(["color", "vector_distance"]).find_nearest(
    vector_field="embedding_field",
    query_vector=Vector([0.3416704, 0.18332680, 0.24160706]),
    distance_measure=DistanceMeasure.EUCLIDEAN,
    limit=10,
    distance_result_field="vector_distance",
)
Node.js
const vectorQuery: VectorQuery = coll
    .select('name', 'description', 'vector_distance')
    .findNearest({
      vectorField: 'embedding_field',
      queryVector: [3.0, 1.0, 2.0],
      limit: 10,
      distanceMeasure: 'EUCLIDEAN',
      distanceResultField: 'vector_distance'
    });
Go
import (
	"context"
	"fmt"
	"io"

	"cloud.google.com/go/firestore"
)

func vectorSearchDistanceResultFieldMasked(w io.Writer, projectID string) error {
	ctx := context.Background()

	client, err := firestore.NewClient(ctx, projectID)
	if err != nil {
		return fmt.Errorf("firestore.NewClient: %w", err)
	}
	defer client.Close()

	collection := client.Collection("coffee-beans")

	// Requires a vector index
	// https://firebase.google.com/docs/firestore/vector-search#create_and_manage_vector_indexes
	vectorQuery := collection.Select("color", "vector_distance").
		FindNearest("embedding_field",
			[]float32{3.0, 1.0, 2.0},
			10,
			firestore.DistanceMeasureEuclidean,
			&firestore.FindNearestOptions{
				DistanceResultField: "vector_distance",
			})

	docs, err := vectorQuery.Documents(ctx).GetAll()
	if err != nil {
		fmt.Fprintf(w, "failed to get vector query results: %v", err)
		return err
	}

	for _, doc := range docs {
		fmt.Fprintf(w, "%v, Distance: %v\n", doc.Data()["color"], doc.Data()["vector_distance"])
	}
	return nil
}
Java
import com.google.cloud.firestore.VectorQuery;
import com.google.cloud.firestore.VectorQueryOptions;
import com.google.cloud.firestore.VectorQuerySnapshot;

VectorQuery vectorQuery = coll
        .select("name", "description", "vector_distance")
        .findNearest(
          "embedding_field",
          new double[] {3.0, 1.0, 2.0},
          /* limit */ 10,
          VectorQuery.DistanceMeasure.EUCLIDEAN,
          VectorQueryOptions.newBuilder()
            .setDistanceResultField("vector_distance")
            .build());

ApiFuture<VectorQuerySnapshot> future = vectorQuery.get();
VectorQuerySnapshot vectorQuerySnapshot = future.get();

for (DocumentSnapshot document : vectorQuerySnapshot.getDocuments()) {
    System.out.println(document.getId() + " Distance: " + document.get("vector_distance"));
}

指定距离阈值

您可以指定相似度阈值,以便仅返回相似度在阈值范围内的文档。阈值字段的行为取决于您选择的距离衡量方式:

  • EUCLIDEANCOSINE 距离会将阈值限制为距离小于或等于指定阈值的文档。随着向量的相似度增加,这些距离衡量值会减小。
  • DOT_PRODUCT 距离会将阈值限制为距离大于或等于指定阈值的文档。随着向量的相似度增加,点积距离会增加。

以下示例展示了如何指定距离阈值,以使用 EUCLIDEAN 距离指标返回最多 10 个距离不超过 4.5 个单位的最近文档:

Python
from google.cloud.firestore_v1.base_vector_query import DistanceMeasure
from google.cloud.firestore_v1.vector import Vector

collection = db.collection("coffee-beans")

vector_query = collection.find_nearest(
    vector_field="embedding_field",
    query_vector=Vector([0.3416704, 0.18332680, 0.24160706]),
    distance_measure=DistanceMeasure.EUCLIDEAN,
    limit=10,
    distance_threshold=4.5,
)

docs = vector_query.stream()

for doc in docs:
    print(f"{doc.id}")
Node.js
const vectorQuery: VectorQuery = coll.findNearest({
  vectorField: 'embedding_field',
  queryVector: [3.0, 1.0, 2.0],
  limit: 10,
  distanceMeasure: 'EUCLIDEAN',
  distanceThreshold: 4.5
});

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

snapshot.forEach((doc) => {
  console.log(doc.id);
});
Go
import (
	"context"
	"fmt"
	"io"

	"cloud.google.com/go/firestore"
)

func vectorSearchDistanceThreshold(w io.Writer, projectID string) error {
	ctx := context.Background()

	client, err := firestore.NewClient(ctx, projectID)
	if err != nil {
		return fmt.Errorf("firestore.NewClient: %w", err)
	}
	defer client.Close()

	collection := client.Collection("coffee-beans")

	// Requires a vector index
	// https://firebase.google.com/docs/firestore/vector-search#create_and_manage_vector_indexes
	vectorQuery := collection.FindNearest("embedding_field",
		[]float32{3.0, 1.0, 2.0},
		10,
		firestore.DistanceMeasureEuclidean,
		&firestore.FindNearestOptions{
			DistanceThreshold: firestore.Ptr[float64](4.5),
		})

	docs, err := vectorQuery.Documents(ctx).GetAll()
	if err != nil {
		fmt.Fprintf(w, "failed to get vector query results: %v", err)
		return err
	}

	for _, doc := range docs {
		fmt.Fprintln(w, doc.Data()["name"])
	}
	return nil
}
Java
import com.google.cloud.firestore.VectorQuery;
import com.google.cloud.firestore.VectorQueryOptions;
import com.google.cloud.firestore.VectorQuerySnapshot;

VectorQuery vectorQuery = coll.findNearest(
        "embedding_field",
        new double[] {3.0, 1.0, 2.0},
        /* limit */ 10,
        VectorQuery.DistanceMeasure.EUCLIDEAN,
        VectorQueryOptions.newBuilder()
          .setDistanceThreshold(4.5)
          .build());

ApiFuture<VectorQuerySnapshot> future = vectorQuery.get();
VectorQuerySnapshot vectorQuerySnapshot = future.get();

for (DocumentSnapshot document : vectorQuerySnapshot.getDocuments()) {
    System.out.println(document.getId());
}

限制

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

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

后续步骤