Integrate TiDB Vector Search with Django ORM

This tutorial walks you through how to use Django ORM to interact with the TiDB Vector Search, store embeddings, and perform vector search queries.

Prerequisites

To complete this tutorial, you need:

• Python 3.8 or higher installed.
• Git installed.
• A TiDB cluster.

If you don't have a TiDB cluster, you can create one as follows:

• Follow Deploy a local test TiDB cluster or Deploy a production TiDB cluster to create a local cluster.
• Follow Creating a TiDB Cloud Serverless cluster to create your own TiDB Cloud cluster.

Run the sample app

You can quickly learn about how to integrate TiDB Vector Search with Django ORM by following the steps below.

Step 1. Clone the repository

Clone the tidb-vector-python repository to your local machine:

git clone https://github.com/pingcap/tidb-vector-python.git

Step 2. Create a virtual environment

Create a virtual environment for your project:

cd tidb-vector-python/examples/orm-django-quickstart
python3 -m venv .venv
source .venv/bin/activate

Step 3. Install required dependencies

Install the required dependencies for the demo project:

pip install -r requirements.txt

Alternatively, you can install the following packages for your project:

pip install Django django-tidb mysqlclient numpy python-dotenv

If you encounter installation issues with mysqlclient, refer to the mysqlclient official documentation.

What is django-tidb

django-tidb is a TiDB dialect for Django, which enhances the Django ORM to support TiDB-specific features (for example, Vector Search) and resolves compatibility issues between TiDB and Django.

To install django-tidb, choose a version that matches your Django version. For example, if you are using django==4.2.*, install django-tidb==4.2.*. The minor version does not need to be the same. It is recommended to use the latest minor version.

For more information, refer to django-tidb repository.

Step 4. Configure the environment variables

Configure the environment variables depending on the TiDB deployment option you've selected.

• TiDB Cloud Serverless
• TiDB Self-Managed

TIDB_HOST=127.0.0.1
TIDB_PORT=4000
TIDB_USERNAME=root
TIDB_PASSWORD=
TIDB_DATABASE=test

Step 5. Run the demo

Migrate the database schema:

python manage.py migrate

Run the Django development server:

python manage.py runserver

Open your browser and visit http://127.0.0.1:8000 to try the demo application. Here are the available API paths:

Sample code snippets

You can refer to the following sample code snippets to complete your own application development.

Connect to the TiDB cluster

In the file sample_project/settings.py, add the following configurations:

dotenv.load_dotenv()

DATABASES = {
    "default": {
        # https://github.com/pingcap/django-tidb
        "ENGINE": "django_tidb",
        "HOST": os.environ.get("TIDB_HOST", "127.0.0.1"),
        "PORT": int(os.environ.get("TIDB_PORT", 4000)),
        "USER": os.environ.get("TIDB_USERNAME", "root"),
        "PASSWORD": os.environ.get("TIDB_PASSWORD", ""),
        "NAME": os.environ.get("TIDB_DATABASE", "test"),
        "OPTIONS": {
            "charset": "utf8mb4",
        },
    }
}

TIDB_CA_PATH = os.environ.get("TIDB_CA_PATH", "")
if TIDB_CA_PATH:
    DATABASES["default"]["OPTIONS"]["ssl_mode"] = "VERIFY_IDENTITY"
    DATABASES["default"]["OPTIONS"]["ssl"] = {
        "ca": TIDB_CA_PATH,
    }

You can create a .env file in the root directory of your project and set up the environment variables TIDB_HOST, TIDB_PORT, TIDB_USERNAME, TIDB_PASSWORD, TIDB_DATABASE, and TIDB_CA_PATH with the actual values of your TiDB cluster.

Create vector tables

Define a vector column

tidb-django provides a VectorField to store vector embeddings in a table.

Create a table with a column named embedding that stores a 3-dimensional vector.

class Document(models.Model):
   content = models.TextField()
   embedding = VectorField(dimensions=3)

Store documents with embeddings

Document.objects.create(content="dog", embedding=[1, 2, 1])
Document.objects.create(content="fish", embedding=[1, 2, 4])
Document.objects.create(content="tree", embedding=[1, 0, 0])

Search the nearest neighbor documents

TiDB Vector support the following distance functions:

• L1Distance
• L2Distance
• CosineDistance
• NegativeInnerProduct

Search for the top-3 documents that are semantically closest to the query vector [1, 2, 3] based on the cosine distance function.

results = Document.objects.annotate(
   distance=CosineDistance('embedding', [1, 2, 3])
).order_by('distance')[:3]

Search documents within a certain distance

Search for the documents whose cosine distance from the query vector [1, 2, 3] is less than 0.2.

results = Document.objects.annotate(
   distance=CosineDistance('embedding', [1, 2, 3])
).filter(distance__lt=0.2).order_by('distance')[:3]

See also

• Vector Data Types
• Vector Search Index