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

本文介绍如何使用 Databricks 矢量搜索创建和查询矢量搜索索引。

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

要求

  • 已启用 Unity Catalog 的工作区。
  • 已启用无服务器计算。
  • 源表必须启用“更改数据馈送”。
  • 若要创建索引,必须具有目录架构中支持创建索引的 CREATE TABLE 权限。 若要查询其他用户拥有的索引,必须具有其他权限。 请参阅查询矢量搜索终结点
  • 如果要使用个人访问令牌(不建议用于生产工作负载),请检查是否已启用个人访问令牌。 若要改用服务主体令牌,请使用 SDK 或 API 调用显式传递它。

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

%pip install databricks-vectorsearch

dbutils.library.restartPython()

from databricks.vector_search.client import VectorSearchClient

创建矢量搜索终结点

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

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

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

  1. 单击左边栏中的“计算”。

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

    创建终结点窗体

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

  4. 单击“确认” 。

使用 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 = VectorSearch(service_principal_client_id=<CLIENT_ID>,service_principal_client_secret=<CLIENT_SECRET>)

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

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

请参阅 POST /api/2.0/vector-search/endpoints

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

如果选择让 Databricks 计算嵌入,则必须设置一个模型服务终结点为嵌入模型提供服务。 请参阅创建基础模型服务终结点中的说明。 有关示例笔记本,请参阅用于调用嵌入模型的笔记本示例

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

注意

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

创建矢量搜索索引

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

有两种类型的索引:

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

使用 UI 创建索引

  1. 单击左边栏中的“目录”打开目录资源管理器 UI

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

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

    “创建索引”按钮

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

    “创建索引”对话框

    名称:用于 Unity 目录中联机表的名称。 该名称需要三级命名空间 <catalog>.<schema>.<name>。 只允许使用字母数字字符和下划线。

    主键:用作主键的列。

    终结点:选择要使用的模型服务终结点。

    嵌入源:指示是否希望 Databricks 计算 Delta 表中文本列的嵌入(计算嵌入),或者 Delta 表是否包含预先计算的嵌入(使用现有嵌入列)

    • 如果选择了“计算嵌入”,请选择要为其计算嵌入的列以及为嵌入模型提供服务的终结点。 仅支持文本列。
    • 如果选择了“使用现有嵌入列”,请选择包含预先计算的嵌入和嵌入维度的列。 预计算的嵌入列的格式应为 array[float]

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

    同步模式:“连续”模式使索引与延迟秒数保持同步。 但是,由于预配了计算群集以运行连续同步流式处理管道,因此它的成本更高。 “被触发”模式更具成本效益,但必须使用 API 手动启动。 对于“连续”和“触发”,更新都是增量更新,即只处理自上次同步以来发生更改的数据

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

使用 Python SDK 创建索引

以下示例使用 Databricks 计算的嵌入创建 Delta 同步索引。

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"
)

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


client = VectorSearchClient()

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

使用 REST API 创建索引

请参阅 POST /api/2.0/vector-search/indexes

保存生成的嵌入内容表

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

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

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

更新矢量搜索索引

更新 Delta 同步索引

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

Python sdk

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, 2.0, 3.0]
       },
       {"id": 2,
        "field2": "value2",
        "field3": 3.0,
        "text_vector": [1.1, 2.1, 3.0]
        }
        ])

REST API

请参阅 REST API (POST /api/2.0/vector-search/indexes)。

查询矢量搜索终结点

只能使用 Python SDK 或 REST API 查询矢量搜索终结点。

注意

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

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

Python sdk

results = index.similarity_search(
    query_text="Greek myths",
    columns=["id", "text"],
    num_results=2
    )

results

REST API

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

对查询使用筛选器

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

筛选器运算符 行为 示例
NOT 对筛选器求反。 键必须以“NOT”结尾。 例如,值为“red”的“color NOT”与颜色不为红色的文档匹配。 {"id NOT": 2} {“color NOT”: “red”}
< 检查字段值是否小于筛选器值。 键必须以“<”结尾。 例如,值为 100 的“price <”与价格小于 100 的文档匹配。 {"id <": 200}
<= 检查字段值是否小于或等于筛选器值。 键必须以“<=”结尾。 例如,值为 100 的“price <=”与价格小于或等于 100 的文档匹配。 {"id <=": 200}
> 检查字段值是否大于筛选器值。 键必须以“>”结尾。 例如,值为 100 的“price >”与价格大于 100 的文档匹配。 {"id >": 200}
>= 检查字段值是否大于或等于筛选器值。 键必须以“>=”结尾。 例如,值为 100 的“price >=”与价格大于或等于 100 的文档匹配。 {"id >=": 200}
OR 检查字段值是否与任何一个筛选器值匹配。 键必须包含 OR 以分隔多个子键。 例如,值为 ["red", "blue"]color1 OR color2color1redcolor2blue 的文档匹配。 {"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

示例笔记本

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

LangChain 示例

请参阅如何将 LangChain 与 Databricks 矢量搜索配合使用,以将 Databricks 矢量搜索用于与 LangChain 包的集成。

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

矢量搜索与 Python SDK 笔记本

获取笔记本

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

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

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

获取笔记本

使用 Databricks 模型服务笔记本调用 BGE 嵌入模型

获取笔记本

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

获取笔记本