Vector Search

Vector search uses semantic similarity to help you find the most relevant records, even if your query does not explicitly include all the keywords.

Note

For a complete example of vector search, see Vector Search Example.

Basic usage

This section shows how to use vector search in your application in just a few steps. Before you start, you need to connect to the database.

Step 1. Create a table with a vector field

You can use client.create_table() to create a table and VectorField to define a vector field.

The following example creates a documents table with four columns:

id: The primary key of the table.
text: The text content of the document.
text_vec: The vector embedding of the text content.
meta: The metadata of the document, which is a JSON object.

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

The VectorField class accepts the following parameters:

dimensions: The vector dimension. Once specified, only vectors with this exact dimension can be stored in this field.
index: Whether to create a vector index for the vector field. Defaults to True.
distance_metric: The distance metric to use for the vector index. Supported values:
- DistanceMetric.COSINE (default): Cosine distance metric, suitable for measuring text similarity
- DistanceMetric.L2: L2 distance metric, suitable for capturing overall difference

Use the CREATE TABLE statement to create a table and use the VECTOR type to define a vector column.

CREATE TABLE documents (
    id INT PRIMARY KEY,
    text TEXT,
    text_vec VECTOR(3),
    VECTOR INDEX `vec_idx_text_vec`((VEC_COSINE_DISTANCE(`text_vec`)))
);

In this example:

The text_vec column is defined as VECTOR(3), so vectors stored in this column must have 3 dimensions.
A vector index is created using the VEC_COSINE_DISTANCE function to optimize vector search performance.

TiDB supports two distance functions for vector indexes:

VEC_COSINE_DISTANCE: Calculates the cosine distance between two vectors
VEC_L2_DISTANCE: Calculates L2 distance (Euclidean distance) between two vectors

Step 2. Insert vector data into the table

For demonstration, insert some text and their corresponding embeddings into the table.

The following example inserts three documents, each with a simple 3-dimensional vector embedding:

dog with the vector embedding [1, 2, 1]
fish with the vector embedding [1, 2, 4]
tree with the vector embedding [1, 0, 0]

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"}),
])

INSERT INTO documents (id, text, text_vec, meta)
VALUES
    (1, 'dog', '[1,2,1]', '{"category": "animal"}'),
    (2, 'fish', '[1,2,4]', '{"category": "animal"}'),
    (3, 'tree', '[1,0,0]', '{"category": "plant"}');

Note

In real-world applications, embeddings are usually generated by an embedding model.

For convenience, pytidb provides an auto embedding feature that can automatically generate vector embeddings for your text fields when you insert, update, or search—no manual processing needed.

For details, see the Auto Embedding guide.

Step 3. Perform vector search

Vector search uses vector distance metrics to measure the similarity and relevance between vectors. The closer the distance, the more relevant the record. To find the most relevant documents in the table, you need to specify a query vector.

The following example assumes the query is A swimming animal and its vector embedding is [1, 2, 3].

Use the table.search() method to perform vector search. It uses search_mode="vector" by default.

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},
]

The result shows that the most relevant document is fish with a distance of 0.00853986601633272.

Use the ORDER BY <distance_function>(<column_name>, <query_vector>) LIMIT <n> clause in a SELECT statement to get the n nearest neighbors of a query vector.

The following example uses the vec_cosine_distance function to calculate the cosine distance between the vectors stored in the text_vec column and the provided query vector [1, 2, 3].

SELECT id, text, vec_cosine_distance(text_vec, '[1,2,3]') AS distance
FROM documents
ORDER BY distance
LIMIT 3;

+----+----------+---------------------+
| id | text     | distance            |
+----+----------+---------------------+
|  2 | fish     | 0.00853986601633272 |
|  1 | dog      | 0.12712843905603044 |
|  3 | tree     |  0.7327387580875756 |
+----+----------+---------------------+
3 rows in set (0.15 sec)

The result shows that the most relevant document is fish with a distance of 0.00853986601633272.

Distance metrics

Distance metrics are a measure of the similarity between a pair of vectors. Currently, TiDB supports the following distance metrics:

The table.search() API supports the following distance metrics:

Metric Name	Description	Best For
`DistanceMetric.COSINE`	Calculates the cosine distance between two vectors (default). Measures the angle between vectors.	Text embeddings, semantic search
`DistanceMetric.L2`	Calculates the L2 distance (Euclidean distance) between two vectors. Measures the straight-line distance.	Image features

To change the distance metric used for vector search, use the .distance_metric() method.

Example: Use the L2 distance metric

from pytidb.schema import DistanceMetric

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

In SQL, you can use the following built-in functions to calculate vector distances directly in your queries:

Function Name	Description
`VEC_L2_DISTANCE`	Calculates L2 distance (Euclidean distance) between two vectors
`VEC_COSINE_DISTANCE`	Calculates the cosine distance between two vectors
`VEC_NEGATIVE_INNER_PRODUCT`	Calculates the negative of the inner product between two vectors
`VEC_L1_DISTANCE`	Calculates L1 distance (Manhattan distance) between two vectors

Distance threshold

The table.search() API allows you to set a distance threshold to control the similarity of the returned results. By specifying this threshold, you can exclude less similar vectors and return only those that meet your relevance criteria.

Use the .distance_threshold() method to set a maximum distance for search results. Only records with a distance less than the threshold are returned.

Example: Only return documents with a distance less than 0.5

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

In SQL, use the HAVING clause with a distance function to filter results by distance:

Example: Only return documents with a distance less than 0.1

SELECT id, text, vec_cosine_distance(text_vec, '[1,2,3]') AS distance
FROM documents
HAVING distance < 0.1
ORDER BY distance
LIMIT 10;

Distance range

The table.search() API also supports specifying a distance range to further refine the results.

Use the .distance_range() method to set both minimum and maximum distance values. Only records with a distance within this range are returned.

Example: Only return documents with a distance between 0.01 and 0.05

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

To specify a distance range in SQL, use BETWEEN or other comparison operators in the HAVING clause:

Example: Only return documents with a distance between 0.01 and 0.05

SELECT id, text, vec_l2_distance(text_vec, '[1,2,3]') AS distance
FROM documents
HAVING distance BETWEEN 0.01 AND 0.05
ORDER BY distance
LIMIT 10;

Metadata filtering

As a relational database, TiDB supports a rich set of SQL operators and allows flexible combinations of filtering conditions.

For vector search in TiDB, you can apply metadata filtering on scalar fields (for example, integers and strings) or JSON fields.

Typically, there are two modes for vector search combined with metadata filtering:

Post-filtering: TiDB first performs vector search to retrieve the top-k candidates from the entire vector space, then applies filters to that candidate set. The vector search stage typically uses a vector index for efficiency.
Pre-filtering: TiDB applies filters before vector search. If the filter is highly selective and the filtered field has a scalar index, this mode can reduce the search space and improve performance.

Post-filtering

Use the .filter() method with a filter dictionary to apply filtering to vector search.

By default, the table.search() API uses post-filtering mode to maximize search performance with the vector index.

Example: Vector search with post-filtering

results = (
    table.search([1, 2, 3])
        # The `meta` is a JSON field, and its value is a JSON object
        # like {"category": "animal"}
        .filter({"meta.category": "animal"})
        .num_candidate(50)
        .limit(10)
        .to_list()
)

Note

When using a vector index, if the final limit is very small, the accuracy of the results might decrease. You can use the .num_candidate() method to control how many candidates to retrieve from the vector index during the vector search phase, without changing the limit parameter.

A higher num_candidate value generally improves recall but might reduce query performance. Adjust this value based on your dataset and accuracy requirements.

Currently, vector indexes are only effective in strict ANN (Approximate Nearest Neighbor) queries, such as:

SELECT * FROM <table> ORDER BY <distance_func>(<column>) LIMIT <n>

In other words, you cannot use a WHERE clause together with a vector index in the same query.

If you need to combine vector search with additional filtering conditions, you can use the post-filtering pattern. In this approach, the ANN query will be divided into two parts:

The inner query performs the vector search using the vector index.
The outer query applies the WHERE condition to filter the results.

SELECT *
FROM (
    SELECT id, text, meta, vec_cosine_distance(text_vec, '[1,2,3]') AS distance
    FROM documents
    ORDER BY distance
    LIMIT 50
) candidates
WHERE meta->>'$.category' = 'animal'
ORDER BY distance
LIMIT 10;

Note

The post-filtering pattern might lead to empty results. For example, the inner query might retrieve the top 50 most similar records, but none of them match the WHERE condition.

To mitigate this, you can increase the LIMIT value (e.g., 50) in the inner query to fetch more candidates, improving the chances of returning enough valid results after filtering.

For supported SQL operators, see Operators in the TiDB Cloud documentation.

Pre-filtering

To enable pre-filtering, set prefilter=True in the .filter() method.

Example: Vector search with pre-filtering

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

For supported filter operators, see Filtering.

In SQL, use the ->> operator or JSON_EXTRACT to access JSON fields in the WHERE clause:

SELECT id, text, meta, vec_cosine_distance(text_vec, '[1,2,3]') AS distance
FROM documents
WHERE meta->>'$.category' = 'animal'
ORDER BY distance
LIMIT 10;

For supported SQL operators, see Operators in the TiDB Cloud documentation.

Multiple vector fields

TiDB supports defining multiple vector columns in a single table, allowing you to store and search different types of vector embeddings.

For example, you can store both text embeddings and image embeddings in the same table, which is convenient for managing multimodal data.

You can define multiple vector fields in the schema and perform vector search on the specified vector field by using the .vector_column() method.

Example: Specify the vector field to search on

# Create a table with multiple vector fields
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")

# Insert sample data ...

# Search using image vector field
results = (
    table.search([1, 2, 3])
        .vector_column("image_vec")
        .distance_metric(DistanceMetric.COSINE)
        .limit(10)
        .to_list()
)

You can create multiple vector columns in a table and search them using suitable distance functions:

-- Create a table with multiple vector fields
CREATE TABLE rich_text_documents (
    id BIGINT PRIMARY KEY,
    text TEXT,
    text_vec VECTOR(3),
    image_url VARCHAR(255),
    image_vec VECTOR(3)
);

-- Insert sample data ...

-- Search using text vector
SELECT id, image_url, vec_l2_distance(image_vec, '[4,5,6]') AS image_distance
FROM rich_text_documents
ORDER BY image_distance
LIMIT 10;

Output search results

The table.search() API lets you convert search results into several common data processing formats:

As SQLAlchemy result rows

To work with raw SQLAlchemy result rows, use:

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

As a list of Python dictionaries

For easier manipulation in Python, convert the results to a list of dictionaries:

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

As a pandas DataFrame

To display results in a user-friendly table—especially useful in Jupyter notebooks—convert them to a pandas DataFrame:

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

As a list of Pydantic model instances

The TableModel class can also be used as a Pydantic model to represent data entities. To work with results as Pydantic model instances, use:

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

Vector Search

Basic usage

Step 1. Create a table with a vector field

Step 2. Insert vector data into the table

Step 3. Perform vector search

Distance metrics

Distance threshold

Distance range

Metadata filtering

Post-filtering

Pre-filtering

Multiple vector fields

Output search results

As SQLAlchemy result rows

As a list of Python dictionaries

As a pandas DataFrame

As a list of Pydantic model instances

Was this page helpful?