如何创建和查询矢量搜索索引

本文介绍如何使用 马赛克 AI 矢量搜索创建和查询矢量搜索索引。

可以使用 UI、 Python SDKREST API 创建和管理矢量搜索组件,例如矢量搜索终结点和矢量搜索索引。

要求

使用访问控制列表配置创建和管理矢量搜索终结点的权限。 请参阅矢量搜索终结点 ACL

安装

若要使用矢量搜索 SDK,必须在笔记本中安装它。 使用以下代码安装包:

%pip install databricks-vectorsearch
dbutils.library.restartPython()

然后使用以下命令导入 VectorSearchClient

from databricks.vector_search.client import VectorSearchClient

身份验证

请参阅 数据保护和身份验证

创建矢量搜索终结点

可以使用 Databricks UI、Python SDK 或 API 创建矢量搜索终结点。

使用 UI 创建矢量搜索终结点

按照以下步骤使用 UI 创建矢量搜索终结点。

  1. 在左侧栏中,单击“ 计算”。

  2. 单击 “矢量搜索 ”选项卡,然后单击“ 创建”。

    创建终结点窗体

  3. 此时会打开 “创建终结点”窗体 。 输入此终结点的名称。

    创建矢量搜索终结点对话框。

  4. “类型” 字段中,选择“ 标准 ”或“ 存储优化”。 请参阅 终结点选项

  5. (可选)在 “高级设置”下,选择预算策略。 请参阅 马赛克 AI 矢量搜索:预算策略

  6. 单击“确认”。

使用 Python SDK 创建矢量搜索终结点

以下示例使用 create_endpoint() SDK 函数创建矢量搜索终结点。

# The following line automatically generates a PAT Token for authentication
client = VectorSearchClient()

# The following line uses the service principal token for authentication
# client = VectorSearchClient(service_principal_client_id=<CLIENT_ID>,service_principal_client_secret=<CLIENT_SECRET>)

client.create_endpoint(
    name="vector_search_endpoint_name",
    endpoint_type="STANDARD" # or "STORAGE_OPTIMIZED"
)

使用 REST API 创建矢量搜索终结点

请参阅 REST API 参考文档: POST /api/2.0/vector-search/endpoints

(可选)创建并配置终结点来为嵌入模型提供服务

如果选择让 Databricks 计算嵌入内容,则可以使用预配置的基础模型 API 终结点或创建提供终结点的模型来为所选的嵌入模型提供服务。 有关说明,请参阅 按令牌付费的基础模型 API,或 创建基础模型服务终结点。 有关示例笔记本,请参阅用于调用嵌入模型的笔记本示例

配置嵌入终结点时,Databricks 建议删除 缩放到零的默认选择。 服务终结点可能需要几分钟时间来预热,因此在终结点纵向缩减的情况下对索引执行初始查询可能会超时。

注释

如果未为数据集适当配置嵌入终结点,矢量搜索索引初始化可能会超时。 应仅对小型数据集和测试使用 CPU 终结点。 对于较大的数据集,请使用 GPU 终结点实现最佳性能。

创建矢量搜索索引

可以使用 UI、Python SDK 或 REST API 创建矢量搜索索引。 UI 是最简单的方法。

有两种类型的索引:

  • 增量同步索引 会自动与源 Delta 表同步,并在 Delta 表中的基础数据发生变化时,自动且增量地更新索引。
  • 直接矢量访问索引 支持直接读取和写入矢量和元数据。 用户负责使用 REST API 或 Python SDK 更新此表。 无法使用 UI 创建这种类型的索引。 必须使用 REST API 或 SDK。

注释

保留列名 _id。 如果源表有一个名为 _id列,请在创建矢量搜索索引之前对其进行重命名。

使用 UI 创建索引

  1. 在左侧栏中,单击 目录 打开目录资源管理器用户界面。

  2. 导航到要使用的 Delta 表。

  3. 单击右上角的“ 创建 ”按钮,然后从下拉菜单中选择 矢量搜索索引

    “创建索引”按钮

  4. 使用对话框中的选择器配置索引。

    “创建索引”对话框

    名称:用于 Unity Catalog 中在线表使用的名称。 名称需要三级命名空间。 <catalog>.<schema>.<name> 仅允许字母数字字符和下划线。

    主键:要用作主键的列。

    终结点:选择要使用的矢量搜索终结点。

    要同步的列:(仅标准终结点支持)。选择要与矢量索引同步的列。 如果将此字段留空,则源表中的所有列都与索引同步。 主键列和嵌入源列或嵌入向量列始终同步。 对于存储优化终结点,源表中的所有列始终同步。

    嵌入源:指示您是希望 Databricks 为 Delta 表的文本列计算嵌入(计算嵌入),还是您的 Delta 表中已经包含预计算的嵌入(使用现有嵌入列)。

    • 如果选择了“计算嵌入”,请选择要计算嵌入的列,以及负责嵌入模型的服务终结点。 仅支持文本列。 对于大规模嵌入生成,Databricks 建议使用按令牌付费的基础模型 databricks-gte-large-en 提高吞吐量。
    • 如果选择了“使用现有嵌入列”,请选择包含预先计算的嵌入和嵌入维度的列。 预先计算的嵌入列格式应为 array[float]。 对于存储优化终结点,嵌入维度必须均匀地被 16 分割。

    同步计算嵌入内容:切换此设置,将生成的嵌入内容保存到 Unity 目录表。 有关详细信息,请参阅 保存生成的嵌入表

    同步模式连续模式让索引以秒级延迟保持同步。 但是,由于预配了计算群集以运行连续同步流式处理管道,因此它的成本更高。 对于标准终结点,连续式触发式都执行增量更新,因此仅处理自上次同步后发生变化的数据。 对于存储优化终结点,每次同步都会完全重新生成矢量搜索索引。 请参阅 存储优化终结点限制

    使用“触发”同步模式时,使用 Python SDK 或 REST API 启动同步。请参阅更新 Delta 同步索引

    对于存储优化终结点,仅支持 触发的 同步模式。

  5. 完成索引配置后,单击“ 创建”。

使用 Python SDK 创建索引

以下示例将创建 Delta 同步索引,其中包含由 Databricks 计算的嵌入。

client = VectorSearchClient()

index = client.create_delta_sync_index(
  endpoint_name="vector_search_demo_endpoint",
  source_table_name="vector_search_demo.vector_search.en_wiki",
  index_name="vector_search_demo.vector_search.en_wiki_index",
  pipeline_type="TRIGGERED",
  primary_key="id",
  embedding_source_column="text",
  embedding_model_endpoint_name="e5-small-v2"
)

以下示例将创建 Delta 同步索引,其中包含自托管的嵌入。 此示例还演示了如何使用可选参数 columns_to_sync 仅选择要在索引中使用的列子集。

client = VectorSearchClient()

index = client.create_delta_sync_index(
  endpoint_name="vector_search_demo_endpoint",
  source_table_name="vector_search_demo.vector_search.en_wiki",
  index_name="vector_search_demo.vector_search.en_wiki_index",
  pipeline_type="TRIGGERED",
  primary_key="id",
  embedding_dimension=1024,
  embedding_vector_column="text_vector"
)

默认情况下,源表中的所有列都与索引同步。

在标准端点上,可以选择用于同步的列子集 columns_to_sync。 主键和嵌入列始终包含在索引中。

若要仅同步主键和嵌入列,必须按如下所示在 中指定它们:columns_to_sync

index = client.create_delta_sync_index(
  ...
  columns_to_sync=["id", "text_vector"] # to sync only the primary key and the embedding column
)

若要同步其他列,请指定它们,如下所示。 无需包括主键和嵌入列,因为它们始终同步。

index = client.create_delta_sync_index(
  ...
  columns_to_sync=["revisionId", "text"] # to sync the `revisionId` and `text` columns in addition to the primary key and embedding column.
)

以下示例创建直接矢量访问索引。


client = VectorSearchClient()

index = client.create_direct_access_index(
  endpoint_name="storage_endpoint",
  index_name=f"{catalog_name}.{schema_name}.{index_name}",
  primary_key="id",
  embedding_dimension=1024,
  embedding_vector_column="text_vector",
  schema={
    "id": "int",
    "field2": "string",
    "field3": "float",
    "text_vector": "array<float>"}
)

使用 REST API 创建索引

请参阅 REST API 参考文档: POST /api/2.0/vector-search/indexes

保存生成的嵌入表

如果 Databricks 生成嵌入,可以将生成的嵌入内容保存到 Unity 目录中的表中。 此表是在与矢量索引相同的架构中创建的,并且是从向量索引页链接的。

表的名称是矢量搜索索引的名称,后面追加了 _writeback_table。 该名称不可编辑。

你可以像 Unity 目录中的任何其他表一样访问和查询表。 但是,不应删除或修改表,因为它不打算手动更新。 如果删除索引,则会自动删除该表。

更新矢量搜索索引

更新 Delta 同步索引

源 Delta 表更改时,使用 连续 同步模式创建的索引会自动更新。 如果使用 触发同步 模式,则可以使用 UI、Python SDK 或 REST API 启动同步。

Databricks 用户界面

  1. 在目录资源管理器中,导航到矢量搜索索引。
  2. 在“概述”选项卡上的“数据引入”部分中,单击“立即同步”按钮,从目录资源管理器同步矢量搜索索引。

Python SDK

client = VectorSearchClient()
index = client.get_index(index_name="vector_search_demo.vector_search.en_wiki_index")

index.sync()

REST API

请参阅 REST API 参考文档: POST /api/2.0/vector-search/indexes/{index_name}/sync

更新直接矢量访问索引

可以使用 Python SDK 或 REST API 从直接矢量访问索引插入、更新或删除数据。

Python SDK

index.upsert([
    {
        "id": 1,
        "field2": "value2",
        "field3": 3.0,
        "text_vector": [1.0] * 1024
    },
    {
        "id": 2,
        "field2": "value2",
        "field3": 3.0,
        "text_vector": [1.1] * 1024
    }
])

REST API

请参阅 REST API 参考文档: POST /api/2.0/vector-search/indexes

对于生产应用程序,Databricks 建议使用服务主体而不是个人访问令牌。 每次查询的性能最多可提升 100 毫秒。

下面的代码示例演示如何使用服务主体更新索引。

export SP_CLIENT_ID=...
export SP_CLIENT_SECRET=...
export INDEX_NAME=...
export WORKSPACE_URL=https://...
export WORKSPACE_ID=...

# Set authorization details to generate OAuth token
export AUTHORIZATION_DETAILS='{"type":"unity_catalog_permission","securable_type":"table","securable_object_name":"'"$INDEX_NAME"'","operation": "WriteVectorIndex"}'

# Generate OAuth token
export TOKEN=$(curl -X POST --url $WORKSPACE_URL/oidc/v1/token -u "$SP_CLIENT_ID:$SP_CLIENT_SECRET" --data 'grant_type=client_credentials' --data 'scope=all-apis' --data-urlencode 'authorization_details=['"$AUTHORIZATION_DETAILS"']' | jq .access_token | tr -d '"')

# Get index URL
export INDEX_URL=$(curl -X GET -H 'Content-Type: application/json' -H "Authorization: Bearer $TOKEN" --url $WORKSPACE_URL/api/2.0/vector-search/indexes/$INDEX_NAME | jq -r '.status.index_url' | tr -d '"')

# Upsert data into vector search index.
curl -X POST -H 'Content-Type: application/json' -H "Authorization: Bearer $TOKEN" --url https://$INDEX_URL/upsert-data --data '{"inputs_json": "[...]"}'

# Delete data from vector search index
curl -X DELETE -H 'Content-Type: application/json' -H "Authorization: Bearer $TOKEN" --url https://$INDEX_URL/delete-data --data '{"primary_keys": [...]}'

下面的代码示例演示如何使用个人访问令牌(PAT)更新索引。

export TOKEN=...
export INDEX_NAME=...
export WORKSPACE_URL=https://...

# Upsert data into vector search index.
curl -X POST -H 'Content-Type: application/json' -H "Authorization: Bearer $TOKEN" --url $WORKSPACE_URL/api/2.0/vector-search/indexes/$INDEX_NAME/upsert-data --data '{"inputs_json": "..."}'

# Delete data from vector search index
curl -X DELETE -H 'Content-Type: application/json' -H "Authorization: Bearer $TOKEN" --url $WORKSPACE_URL/api/2.0/vector-search/indexes/$INDEX_NAME/delete-data --data '{"primary_keys": [...]}'

查询矢量搜索终结点

只能使用 Python SDK、REST API 或 SQL vector_search() AI 函数查询矢量搜索终结点。

注释

如果查询终结点的用户不是矢量搜索索引的所有者,则用户必须具有以下 UC 权限:

  • 对包含矢量搜索索引的目录具有 USE CATALOG 权限。
  • 对包含矢量搜索索引的架构具有 USE SCHEMA 权限。
  • 对矢量搜索索引具有 SELECT 权限。

默认查询类型为 ann (近似近邻)。 若要执行混合关键字相似性搜索,请将参数 query_type 设置为 hybrid。 使用混合搜索时,将包含所有文本元数据列,最多返回 200 个结果。

Python SDK 标准终结点

# Delta Sync Index with embeddings computed by Databricks
results = index.similarity_search(
    query_text="Greek myths",
    columns=["id", "field2"],
    num_results=2
    )

# Delta Sync Index using hybrid search, with embeddings computed by Databricks
results3 = index.similarity_search(
    query_text="Greek myths",
    columns=["id", "field2"],
    num_results=2,
    query_type="hybrid"
    )

# Delta Sync Index with pre-calculated embeddings
results2 = index.similarity_search(
    query_vector=[0.9] * 1024,
    columns=["id", "text"],
    num_results=2
    )

Python SDK 存储优化的端点

现有筛选器接口已重新设计为存储优化矢量搜索索引,以采用更类似于 SQL 的筛选器字符串,而不是标准矢量搜索终结点中使用的筛选器字典。

client = VectorSearchClient()
index = client.get_index(index_name="vector_search_demo.vector_search.en_wiki_index")

# similarity search with query vector
results = index.similarity_search(
    query_vector=[0.2, 0.33, 0.19, 0.52],
    columns=["id", "text"],
    num_results=2
)

# similarity search with query vector and filter string
results = index.similarity_search(
    query_vector=[0.2, 0.33, 0.19, 0.52],
    columns=["id", "text"],
    # this is a single filter string similar to SQL WHERE clause syntax
    filters="language = 'en' AND country = 'us'",
    num_results=2
)

REST API

请参阅 REST API 参考文档: POST /api/2.0/vector-search/indexes/{index_name}/query

对于生产应用程序,Databricks 建议使用服务主体而不是个人访问令牌。 除了增强安全性和访问管理之外,使用服务主体还可以将每个查询的性能提高最多至 100 毫秒。

下面的代码示例演示如何使用服务主体查询索引。

export SP_CLIENT_ID=...
export SP_CLIENT_SECRET=...
export INDEX_NAME=...
export WORKSPACE_URL=https://...
export WORKSPACE_ID=...

# Set authorization details to generate OAuth token
export AUTHORIZATION_DETAILS='{"type":"unity_catalog_permission","securable_type":"table","securable_object_name":"'"$INDEX_NAME"'","operation": "ReadVectorIndex"}'
# If you are using an route_optimized embedding model endpoint (TODO: link), then you need to have additional authorization details to invoke the serving endpoint
# export EMBEDDING_MODEL_SERVING_ENDPOINT_ID=...
# export AUTHORIZATION_DETAILS="$AUTHORIZATION_DETAILS"',{"type":"workspace_permission","object_type":"serving-endpoints","object_path":"/serving-endpoints/'"$EMBEDDING_MODEL_SERVING_ENDPOINT_ID"'","actions": ["query_inference_endpoint"]}'

# Generate OAuth token
export TOKEN=$(curl -X POST  --url $WORKSPACE_URL/oidc/v1/token -u "$SP_CLIENT_ID:$SP_CLIENT_SECRET" --data 'grant_type=client_credentials' --data 'scope=all-apis' --data-urlencode 'authorization_details=['"$AUTHORIZATION_DETAILS"']' | jq .access_token | tr -d '"')

# Get index URL
export INDEX_URL=$(curl -X GET -H 'Content-Type: application/json' -H "Authorization: Bearer $TOKEN" --url $WORKSPACE_URL/api/2.0/vector-search/indexes/$INDEX_NAME | jq -r '.status.index_url' | tr -d '"')

# Query vector search index.
curl -X GET -H 'Content-Type: application/json' -H "Authorization: Bearer $TOKEN" --url https://$INDEX_URL/query --data '{"num_results": 3, "query_vector": [...], "columns": [...], "debug_level": 1}'

# Query vector search index.
curl -X GET -H 'Content-Type: application/json' -H "Authorization: Bearer $TOKEN" --url https://$INDEX_URL/query --data '{"num_results": 3, "query_text": "...", "columns": [...], "debug_level": 1}'

下面的代码示例演示如何使用个人访问令牌(PAT)查询索引。

export TOKEN=...
export INDEX_NAME=...
export WORKSPACE_URL=https://...

# Query vector search index with `query_vector`
curl -X GET -H 'Content-Type: application/json' -H "Authorization: Bearer $TOKEN" --url $WORKSPACE_URL/api/2.0/vector-search/indexes/$INDEX_NAME/query --data '{"num_results": 3, "query_vector": [...], "columns": [...], "debug_level": 1}'

# Query vector search index with `query_text`
curl -X GET -H 'Content-Type: application/json' -H "Authorization: Bearer $TOKEN" --url $WORKSPACE_URL/api/2.0/vector-search/indexes/$INDEX_NAME/query --data '{"num_results": 3, "query_text": "...", "columns": [...], "debug_level": 1}'

SQL

重要

vector_search() AI 函数处于公共预览阶段

若要使用此 AI 函数,请参阅 vector_search 函数

对查询使用筛选器

查询可以根据 Delta 表中的任何列定义筛选器。 similarity_search 仅返回与指定筛选器匹配的行。

下表列出了支持的筛选器。

注释

存储优化矢量搜索索引支持表中列出的所有筛选器运算符。 对于存储优化索引,筛选器的行为与 SQL 筛选器行为匹配,而不是表中显示的标准终结点筛选器使用的自定义行为。

除了表中的筛选器运算符外,存储优化矢量搜索索引还支持 to_timestamp 运算符。

筛选器运算符 行为 例子
NOT 对筛选器求反。 密钥必须以“NOT”结尾。 例如,值为“red”的“color NOT”与颜色不为红色的文档匹配。 {"id NOT": 2} {“color NOT”: “red”}
< 检查字段值是否小于筛选器值。 键必须以“ <”结尾。 例如,值为 200 的“price <”与价格小于 200 的文档匹配。 {"id <": 200}
<= 检查字段值是否小于或等于筛选器值。 密钥必须以“ <=”结尾。 例如,值为 200 的“price <=”与价格小于或等于 200 的文档匹配。 {"id <=": 200}
> 检查字段值是否大于筛选器值。 键必须以“ >”结尾。 例如,值为 200 的“price >”与价格大于 200 的文档匹配。 {"id >": 200}
>= 检查字段值是否大于或等于筛选器值。 密钥必须以“ >=”结尾。 例如,值为 200 的“price >=”与价格大于或等于 200 的文档匹配。 {"id >=": 200}
OR 检查字段值是否与任何筛选器值匹配。 密钥必须包含 OR 以分隔多个子项。 例如,值为 color1 OR color2["red", "blue"]color1redcolor2blue 的文档匹配。 {"color1 OR color2": ["red", "blue"]}
LIKE 匹配字符串中空格分隔的标记。 请参阅下面的代码示例。 {"column LIKE": "hello"}
未指定筛选器运算符 筛选器检查是否完全匹配。 如果指定了多个值,则它与任何值匹配。 {"id": 200} {"id": [200, 300]}

请参阅以下代码示例:

Python SDK

# Match rows where `title` exactly matches `Athena` or `Ares`
results = index.similarity_search(
    query_text="Greek myths",
    columns=["id", "text"],
    filters={"title": ["Ares", "Athena"]},
    num_results=2
    )

# Match rows where `title` or `id` exactly matches `Athena` or `Ares`
results = index.similarity_search(
    query_text="Greek myths",
    columns=["id", "text"],
    filters={"title OR id": ["Ares", "Athena"]},
    num_results=2
    )

# Match only rows where `title` is not `Hercules`
results = index.similarity_search(
    query_text="Greek myths",
    columns=["id", "text"],
    filters={"title NOT": "Hercules"},
    num_results=2
    )

REST API

请参阅 POST /api/2.0/vector-search/indexes/{index_name}/query

喜欢

LIKE 例子

{"column LIKE": "apple"}:与字符串“apple”和“apple pear”匹配,但与“pineapple”或“pear”不匹配。 请注意,即使包含“apple”子字符串,也与“pineapple”不匹配,它查找的是包含空格分隔标记的完全匹配内容,例如“apple pear”。

{"column NOT LIKE": "apple"} 则相反。 它匹配“菠萝”和“梨”,但不匹配“苹果”或“苹果梨”。

示例笔记本

本节中的示例演示了矢量搜索 Python SDK 的用法。

LangChain 示例

请参阅如何将 LangChain 与 Mosaic AI 矢量搜索配合使用,以按与 LangChain 包集成的方式使用 Mosaic AI 矢量搜索。

以下笔记本演示如何将相似性搜索结果转换为 LangChain 文档。

使用 Python SDK 笔记本进行矢量搜索

获取笔记本

用于调用嵌入模型的笔记本示例

以下笔记本演示如何配置用于生成嵌入的Mosaic AI模型服务端点。

使用 Mosaic AI 模型服务笔记本调用 OpenAI 嵌入模型

获取笔记本

使用 Mosaic AI 模型服务笔记本调用 GTE 嵌入模型

获取笔记本

注册并提供 OSS 嵌入模型笔记本

获取笔记本