Skip to content

Commit 3e99f50

Browse files
authored
README: Add django example (#66)
1 parent d34d3b9 commit 3e99f50

File tree

1 file changed

+88
-15
lines changed

1 file changed

+88
-15
lines changed

README.md

Lines changed: 88 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -2,19 +2,13 @@
22

33
Use TiDB Vector Search with Python.
44

5-
## Installation
6-
7-
```bash
8-
pip install tidb-vector
9-
```
10-
115
## Usage
126

137
TiDB is a SQL database so that this package introduces Vector Search capability for Python ORMs:
148

159
- [#SQLAlchemy](#sqlalchemy)
16-
- [#Django](#django)
1710
- [#Peewee](#peewee)
11+
- [#Django](#django)
1812

1913
Pick one that you are familiar with to get started. If you are not using any of them, we recommend [#SQLAlchemy](#sqlalchemy).
2014

@@ -24,13 +18,18 @@ We also provide a Vector Search client for simple usage:
2418

2519
### SQLAlchemy
2620

21+
Install:
22+
2723
```bash
2824
pip install tidb-vector sqlalchemy pymysql
2925
```
3026

27+
Usage:
28+
3129
```python
32-
from sqlalchemy import Integer, Text, Column
30+
from sqlalchemy import Integer, Column
3331
from sqlalchemy import create_engine, select
32+
from sqlalchemy.dialects.mysql import LONGTEXT
3433
from sqlalchemy.orm import Session, declarative_base
3534

3635
import tidb_vector
@@ -44,15 +43,15 @@ Base = declarative_base()
4443
class Doc(Base):
4544
__tablename__ = "doc"
4645
id = Column(Integer, primary_key=True)
47-
embedding = Column(VectorType(3)) # Vector with 3 dimensions
48-
content = Column(Text)
46+
embedding = Column(VectorType(dim=3))
47+
content = Column(LONGTEXT)
4948

5049

5150
# Create empty table
5251
Base.metadata.drop_all(engine) # clean data from last run
5352
Base.metadata.create_all(engine)
5453

55-
# Create index using L2 distance
54+
# Create index for L2 distance
5655
adaptor = VectorAdaptor(engine)
5756
adaptor.create_vector_index(
5857
Doc.embedding, tidb_vector.DistanceMetric.L2, skip_existing=True
@@ -69,7 +68,7 @@ with Session(engine) as session:
6968
with Session(engine) as session:
7069
results = session.execute(
7170
select(Doc.id, Doc.content)
72-
.order_by(Doc.embedding.cosine_distance([1, 2, 3]))
71+
.order_by(Doc.embedding.l2_distance([1, 2, 3]))
7372
.limit(1)
7473
).all()
7574
print(results)
@@ -78,16 +77,90 @@ with Session(engine) as session:
7877
with Session(engine) as session:
7978
results = session.execute(
8079
select(Doc.id, Doc.content)
81-
.where(Doc.id > 2)
82-
.order_by(Doc.embedding.cosine_distance([1, 2, 3]))
80+
.where(Doc.content == "dog")
81+
.order_by(Doc.embedding.l2_distance([1, 2, 3]))
8382
.limit(1)
8483
).all()
8584
print(results)
8685
```
8786

8887
### Django
8988

90-
To use vector field in Django, you need to use [`django-tidb`](https://github.com/pingcap/django-tidb?tab=readme-ov-file#vector-beta).
89+
> [!TIP]
90+
>
91+
> Django is a full-featured web framework, not just an ORM. The following usage introducutions are provided for existing Django users.
92+
>
93+
> For new users to get started, consider using SQLAlchemy or Peewee.
94+
95+
Install:
96+
97+
```bash
98+
pip install 'django-tidb[vector]~=5.0.0' 'django~=5.0.0' mysqlclient
99+
```
100+
101+
Usage:
102+
103+
1\. Configure `django_tidb` as engine, like:
104+
105+
```python
106+
DATABASES = {
107+
'default': {
108+
'ENGINE': 'django_tidb',
109+
'NAME': 'django',
110+
'USER': 'root',
111+
'PASSWORD': '',
112+
'HOST': '127.0.0.1',
113+
'PORT': 4000,
114+
},
115+
}
116+
```
117+
118+
2\. Define a model with a vector field and vector index:
119+
120+
```python
121+
from django.db import models
122+
from django_tidb.fields.vector import VectorField, VectorIndex, L2Distance
123+
124+
class Doc(models.Model):
125+
id = models.IntegerField(primary_key=True)
126+
embedding = VectorField(dimensions=3)
127+
content = models.TextField()
128+
class Meta:
129+
indexes = [VectorIndex(L2Distance("embedding"), name="idx")]
130+
```
131+
132+
3\. Insert data:
133+
134+
```python
135+
Doc.objects.create(id=1, content="dog", embedding=[1, 2, 1])
136+
Doc.objects.create(id=2, content="fish", embedding=[1, 2, 4])
137+
Doc.objects.create(id=3, content="tree", embedding=[1, 0, 0])
138+
```
139+
140+
4\. Perform Vector Search for Top K=1:
141+
142+
```python
143+
queryset = (
144+
Doc.objects
145+
.order_by(L2Distance("embedding", [1, 2, 3]))
146+
.values("id", "content")[:1]
147+
)
148+
print(queryset)
149+
```
150+
151+
5\. Perform filtered Vector Search by adding a Where Clause:
152+
153+
```python
154+
queryset = (
155+
Doc.objects
156+
.filter(content="dog")
157+
.order_by(L2Distance("embedding", [1, 2, 3]))
158+
.values("id", "content")[:1]
159+
)
160+
print(queryset)
161+
```
162+
163+
For more details, see [django-tidb](https://github.com/pingcap/django-tidb?tab=readme-ov-file#vector-beta).
91164

92165
### Peewee
93166

0 commit comments

Comments
 (0)