使用 Java 在 Azure Cosmos DB for NoSQL 中编制矢量索引和查询矢量

项目
12/20/2024

在使用矢量索引和搜索功能之前，必须先启用该功能。本文涵盖以下步骤：

启用 Azure Cosmos DB for NoSQL 中的矢量搜索功能
为矢量搜索设置 Azure Cosmos DB 容器
创建矢量嵌入策略
将矢量索引添加到容器索引策略
使用矢量索引和矢量嵌入策略创建容器
对存储的数据执行矢量搜索

本指南逐步讲解创建矢量数据、为数据编制索引，然后查询容器中的数据的过程。

先决条件

一个现有的 Azure Cosmos DB for NoSQL 帐户。
- 如果你没有 Azure 订阅，请免费试用 Azure Cosmos DB for NoSQL。
- 如果你当前有 Azure 订阅，请创建新的 Azure Cosmos DB for NoSQL 帐户。
最新版本的 Azure Cosmos DB Java SDK。

启用功能

Azure Cosmos DB for NoSQL 的矢量搜索需要功能启用。按照以下步骤注册：

导航到 Azure Cosmos DB for NoSQL 资源页。
选择“设置”菜单项下的“功能”窗格。
选择“Azure Cosmos DB for NoSQL 中的矢量搜索”。
阅读有关该功能的说明，确认你想要启用该功能。
选择“启用”以启用 Azure Cosmos DB for NoSQL 中的矢量搜索功能。

提示

或使用 Azure CLI 更新帐户的功能以支持 NoSQL 矢量搜索。

az cosmosdb update \
     --resource-group <resource-group-name> \
     --name <account-name> \
     --capabilities EnableNoSQLVectorSearch

注意

注册请求将自动获得批准，但可能需要 15 分钟才能生效。

了解矢量搜索所涉及的步骤

以下步骤假定你知道如何设置 Cosmos DB NoSQL 帐户并创建数据库。现有容器当前不支持矢量搜索功能，因此需要创建一个新容器并指定容器级矢量嵌入策略以及创建容器时的矢量索引策略。

让我们以为 Internet 书店创建数据库为例，为每本书存储“书名”、“作者”、ISBN 和“描述”。我们还定义了两个属性来包含矢量嵌入。第一个是“contentVector”属性，它包含从书籍文本内容生成的文本嵌入（例如，在创建嵌入之前将“书名”“作者”“isbn”和“描述”属性连接起来）。第二个是“coverImageVector”，基于书籍封面图像生成。

为要执行矢量搜索的字段创建并存储矢量嵌入。
在矢量嵌入策略中指定矢量嵌入路径。
在容器的索引策略中包含任何所需的矢量索引。

在本文的后续部分中，我们将考虑对存储在容器中的项采用以下结构：

{
"title": "book-title", 
"author": "book-author", 
"isbn": "book-isbn", 
"description": "book-description", 
"contentVector": [2, -1, 4, 3, 5, -2, 5, -7, 3, 1], 
"coverImageVector": [0.33, -0.52, 0.45, -0.67, 0.89, -0.34, 0.86, -0.78] 
}

首先，创建 CosmosContainerProperties 对象。

CosmosContainerProperties collectionDefinition = new CosmosContainerProperties(UUID.randomUUID().toString(), "Partition_Key_Def");

为容器创建矢量嵌入策略

接下来，需要定义容器矢量策略。该策略提供的信息用于告知 Azure Cosmos DB 查询引擎如何处理 VectorDistance 系统函数中的矢量属性。这也为矢量索引策略提供了必要的信息，如果你选择指定一个的话。包含的矢量策略中包含以下信息：

“path”：包含矢量的属性路径
“datatype”：矢量元素的类型（默认为 Float32）
“dimensions”：路径中每个矢量的长度（默认为 1536）
“distanceFunction”：用于计算距离/相似性的指标（默认余弦）

对于我们的书籍详细信息示例，矢量策略可能如 JSON 示例所示：

// Creating vector embedding policy
CosmosVectorEmbeddingPolicy cosmosVectorEmbeddingPolicy = new CosmosVectorEmbeddingPolicy();

CosmosVectorEmbedding embedding1 = new CosmosVectorEmbedding();
embedding1.setPath("/coverImageVector");
embedding1.setDataType(CosmosVectorDataType.FLOAT32);
embedding1.setDimensions(8L);
embedding1.setDistanceFunction(CosmosVectorDistanceFunction.COSINE);

CosmosVectorEmbedding embedding2 = new CosmosVectorEmbedding();
embedding2.setPath("/contentVector");
embedding2.setDataType(CosmosVectorDataType.FLOAT32);
embedding2.setDimensions(10L);
embedding2.setDistanceFunction(CosmosVectorDistanceFunction.DOT_PRODUCT);

cosmosVectorEmbeddingPolicy.setCosmosVectorEmbeddings(Arrays.asList(embedding1, embedding2, embedding3));

collectionDefinition.setVectorEmbeddingPolicy(cosmosVectorEmbeddingPolicy);

在索引策略中创建矢量索引

确定矢量嵌入路径后，需要将矢量索引添加到索引策略。目前，Azure Cosmos DB for NoSQL 的矢量搜索功能只支持新容器，因此你需要在创建容器时应用矢量策略，以后无法修改。对于此示例，索引策略如下所示：

IndexingPolicy indexingPolicy = new IndexingPolicy();
indexingPolicy.setIndexingMode(IndexingMode.CONSISTENT);
ExcludedPath excludedPath1 = new ExcludedPath("/coverImageVector/*");
ExcludedPath excludedPath2 = new ExcludedPath("/contentVector/*");
indexingPolicy.setExcludedPaths(ImmutableList.of(excludedPath1, excludedPath2));

IncludedPath includedPath1 = new IncludedPath("/*");
indexingPolicy.setIncludedPaths(Collections.singletonList(includedPath1));

// Creating vector indexes
CosmosVectorIndexSpec cosmosVectorIndexSpec1 = new CosmosVectorIndexSpec();
cosmosVectorIndexSpec1.setPath("/coverImageVector");
cosmosVectorIndexSpec1.setType(CosmosVectorIndexType.QUANTIZED_FLAT.toString());

CosmosVectorIndexSpec cosmosVectorIndexSpec2 = new CosmosVectorIndexSpec();
cosmosVectorIndexSpec2.setPath("/contentVector");
cosmosVectorIndexSpec2.setType(CosmosVectorIndexType.DISK_ANN.toString());

indexingPolicy.setVectorIndexes(Arrays.asList(cosmosVectorIndexSpec1, cosmosVectorIndexSpec2, cosmosVectorIndexSpec3));

collectionDefinition.setIndexingPolicy(indexingPolicy);

最后，使用容器索引策略和矢量索引策略创建容器。

database.createContainer(collectionDefinition).block();

重要

将矢量路径添加到索引策略的“excludedPaths”部分可确保优化插入性能。不将矢量路径添加到“excludedPaths”会导致矢量插入的 RU 费用和延迟较高。

运行矢量相似性搜索查询

使用所需矢量策略创建容器并将矢量数据插入容器后，可以在查询中使用 Vector Distance 系统函数进行矢量搜索。假设你想通过查看描述来搜索有关美食食谱的书籍，首先需要获取查询文本的嵌入。在这种情况下，可能需要为查询文本“食谱”生成嵌入。获得搜索查询的嵌入后，可以在矢量搜索查询中的 VectorDistance 函数中使用它，获得与查询相似的所有项，如下所示：

SELECT TOP 10 c.title, VectorDistance(c.contentVector, [1,2,3,4,5,6,7,8,9,10]) AS SimilarityScore   
FROM c  
ORDER BY VectorDistance(c.contentVector, [1,2,3,4,5,6,7,8,9,10])

此查询检索书名以及与你的查询相对应的相似度分数。下面是一个使用 Java 的示例：

float[] embedding = new float[10];
for (int i = 0; i < 10; i++) {
    array[i] = i + 1;
}
ArrayList<SqlParameter> paramList = new ArrayList<SqlParameter>();
  paramList.add(new SqlParameter("@embedding", embedding));
  SqlQuerySpec querySpec = new SqlQuerySpec("SELECT c.title, VectorDistance(c.contentVector,@embedding) AS SimilarityScore  FROM c ORDER BY VectorDistance(c.contentVector,@embedding)", paramList);
  CosmosPagedIterable<Family> filteredFamilies = container.queryItems(querySpec, new CosmosQueryRequestOptions(), Family.class);

  if (filteredFamilies.iterator().hasNext()) {
      Family family = filteredFamilies.iterator().next();
      logger.info(String.format("First query result: Family with (/id, partition key) = (%s,%s)",family.getId(),family.getLastName()));
  }

通过