向量搜索

你可以使用 client.create_table() 创建表，并用 VectorField 定义向量字段。

以下示例创建了一个包含四个列的 documents 表：

• id：表的主键。
• text：文档的文本内容。
• text_vec：文本内容的向量嵌入。
• meta：文档的元信息，是一个 JSON 对象。

from pytidb.schema import TableModel, Field, VectorField
from pytidb.datatype import TEXT, JSON

class Document(TableModel):
    __tablename__ = "documents"

    id: int = Field(primary_key=True)
    text: str = Field(sa_type=TEXT)
    text_vec: list[float] = VectorField(dimensions=3)
    meta: dict = Field(sa_type=JSON, default_factory=dict)

table = client.create_table(schema=Document, if_exists="overwrite")

VectorField 类接受以下参数：

• dimensions：向量的维度。指定后，该字段只能存储具有该精确维度的向量。
• index：是否为该向量字段创建向量索引。默认为 True。
• distance_metric：用于向量索引的距离度量。支持的取值：
  • DistanceMetric.COSINE（默认）：余弦距离度量，适合衡量文本相似性
  • DistanceMetric.L2：L2 距离度量，适合衡量整体差异

table.bulk_insert([
    Document(text="dog", text_vec=[1,2,1], meta={"category": "animal"}),
    Document(text="fish", text_vec=[1,2,4], meta={"category": "animal"}),
    Document(text="tree", text_vec=[1,0,0], meta={"category": "plant"}),
])

使用 table.search() 方法执行向量搜索。默认使用 search_mode="vector"。

table.search([1, 2, 3]).limit(3).to_list()

[
    {"id": 2, "text": "fish", "text_vec": [1,2,4], "_distance": 0.00853986601633272},
    {"id": 1, "text": "dog", "text_vec": [1,2,1], "_distance": 0.12712843905603044},
    {"id": 3, "text": "tree", "text_vec": [1,0,0], "_distance": 0.7327387580875756},
]

结果显示，最相关的文档是 fish，其距离为 0.00853986601633272。

table.search() API 支持以下距离度量：

要更改向量搜索使用的距离度量，可使用 .distance_metric() 方法。

示例：使用 L2 距离度量

from pytidb.schema import DistanceMetric

results = (
    table.search([1, 2, 3])
        .distance_metric(DistanceMetric.L2)
        .limit(10)
        .to_list()
)

使用 .distance_threshold() 方法为搜索结果设置最大距离。只有距离小于该阈值的记录才会被返回。

示例：仅返回距离小于 0.5 的文档

results = table.search([1, 2, 3]).distance_threshold(0.5).limit(10).to_list()

使用 .distance_range() 方法同时设置最小和最大距离值。只有距离在该范围内的记录才会被返回。

示例：仅返回距离在 0.01 到 0.05 之间的文档

results = table.search([1, 2, 3]).distance_range(0.01, 0.05).limit(10).to_list()

使用 .filter() 方法和过滤字典为向量搜索添加过滤条件。

默认情况下，table.search() API 使用后过滤模式，以最大化向量索引的搜索性能。

示例：向量搜索结合后过滤

results = (
    table.search([1, 2, 3])
        # `meta` 是 JSON 字段，其值为类似 {"category": "animal"} 的 JSON 对象
        .filter({"meta.category": "animal"})
        .num_candidate(50)
        .limit(10)
        .to_list()
)

较大的 num_candidate 值通常可以提升召回率，但可能降低查询性能。请根据你的数据集和准确性需求调整该值。

要启用预过滤，在 .filter() 方法中设置 prefilter=True。

示例：向量搜索结合预过滤

results = (
    table.search([1, 2, 3])
        .filter({"meta.category": "animal"}, prefilter=True)
        .limit(10)
        .to_list()
)

有关支持的过滤运算符，请参见 过滤。

你可以在 schema 中定义多个向量字段，并通过 .vector_column() 方法指定要搜索的向量字段。

示例：指定搜索的向量字段

# 创建包含多个向量字段的表
class RichTextDocument(TableModel):
    __tablename__ = "rich_text_documents"
    id: int = Field(primary_key=True)
    text: str = Field(sa_type=TEXT)
    text_vec: list[float] = VectorField(dimensions=3)
    image_url: str
    image_vec: list[float] = VectorField(dimensions=3)

table = client.create_table(schema=RichTextDocument, if_exists="overwrite")

# 插入示例数据 ...

# 使用图像向量字段进行搜索
results = (
    table.search([1, 2, 3])
        .vector_column("image_vec")
        .distance_metric(DistanceMetric.COSINE)
        .limit(10)
        .to_list()
)