Support vector index for Peewee (#67)

A VectorAdaptor is introduced similar to how we did for SQLAlchemy.

It is possible to utilize Peewee's native index, when we only need one
SQL statement to add index. Currently we must execute two SQL
statements, so that an adaptor is needed.

Author: breezewish

README.md

@@ -52,9 +52,9 @@ Base.metadata.drop_all(engine)  # clean data from last run
 Base.metadata.create_all(engine)
 
 # Create index for L2 distance
-adaptor = VectorAdaptor(engine)
-adaptor.create_vector_index(
+VectorAdaptor(engine).create_vector_index(
     Doc.embedding, tidb_vector.DistanceMetric.L2, skip_existing=True
+    # For cosine distance, use tidb_vector.DistanceMetric.COSINE
 )
 
 # Insert content with vectors
@@ -69,6 +69,7 @@ with Session(engine) as session:
     results = session.execute(
         select(Doc.id, Doc.content)
         .order_by(Doc.embedding.l2_distance([1, 2, 3]))
+        # For cosine distance, use Doc.embedding.cosine_distance(...)
         .limit(1)
     ).all()
     print(results)
@@ -84,6 +85,78 @@ with Session(engine) as session:
     print(results)
 ```
 
+### Peewee
+
+Install:
+
+```bash
+pip install tidb-vector peewee pymysql
+```
+
+Usage:
+
+```python
+import tidb_vector
+from peewee import Model, MySQLDatabase, IntegerField, TextField
+from tidb_vector.peewee import VectorField, VectorAdaptor
+
+db = MySQLDatabase(
+    database="test",
+    user="root",
+    password="",
+    host="127.0.0.1",
+    port=4000,
+)
+
+
+# Define table schema
+class Doc(Model):
+    class Meta:
+        database = db
+        table_name = "peewee_test"
+
+    id = IntegerField(primary_key=True)
+    embedding = VectorField(3)
+    content = TextField()
+
+
+# Create empty table and index for L2 distance
+db.drop_tables([Doc])  # clean data from last run
+db.create_tables([Doc])
+# For cosine distance, use tidb_vector.DistanceMetric.COSINE
+VectorAdaptor(db).create_vector_index(Doc.embedding, tidb_vector.DistanceMetric.L2)
+
+# Insert content with vectors
+Doc.insert_many(
+    [
+        {"id": 1, "content": "dog", "embedding": [1, 2, 1]},
+        {"id": 2, "content": "fish", "embedding": [1, 2, 4]},
+        {"id": 3, "content": "tree", "embedding": [1, 0, 0]},
+    ]
+).execute()
+
+# Perform Vector Search for Top K=1
+cursor = (
+    Doc.select(Doc.id, Doc.content)
+    # For cosine distance, use Doc.embedding.cosine_distance(...)
+    .order_by(Doc.embedding.l2_distance([1, 2, 3]))
+    .limit(1)
+)
+for row in cursor:
+    print(row.id, row.content)
+
+
+# Perform filtered Vector Search by adding a Where Clause:
+cursor = (
+    Doc.select(Doc.id, Doc.content)
+    .where(Doc.content == "dog")
+    .order_by(Doc.embedding.l2_distance([1, 2, 3]))
+    .limit(1)
+)
+for row in cursor:
+    print(row.id, row.content)
+```
+
 ### Django
 
 > [!TIP]
@@ -162,83 +235,6 @@ print(queryset)
 
 For more details, see [django-tidb](https://github.com/pingcap/django-tidb?tab=readme-ov-file#vector-beta).
 
-### Peewee
-
-Define peewee table with vector field
-
-```python
-from peewee import Model, MySQLDatabase
-from tidb_vector.peewee import VectorField
-
-# Using `pymysql` as the driver
-connect_kwargs = {
-    'ssl_verify_cert': True,
-    'ssl_verify_identity': True,
-}
-
-# Using `mysqlclient` as the driver
-connect_kwargs = {
-    'ssl_mode': 'VERIFY_IDENTITY',
-    'ssl': {
-        # Root certificate default path
-        # https://docs.pingcap.com/tidbcloud/secure-connections-to-serverless-clusters/#root-certificate-default-path
-        'ca': '/etc/ssl/cert.pem'  # MacOS
-    },
-}
-
-db = MySQLDatabase(
-    'peewee_test',
-    user='xxxxxxxx.root',
-    password='xxxxxxxx',
-    host='xxxxxxxx.shared.aws.tidbcloud.com',
-    port=4000,
-    **connect_kwargs,
-)
-
-class TestModel(Model):
-    class Meta:
-        database = db
-        table_name = 'test'
-
-    embedding = VectorField(3)
-
-# or add hnsw index when creating table
-class TestModelWithIndex(Model):
-    class Meta:
-        database = db
-        table_name = 'test_with_index'
-
-    embedding = VectorField(3, constraints=[SQL("COMMENT 'hnsw(distance=l2)'")])
-
-
-db.connect()
-db.create_tables([TestModel, TestModelWithIndex])
-```
-
-Insert vector data
-
-```python
-TestModel.create(embedding=[1, 2, 3])
-```
-
-Get the nearest neighbors
-
-```python
-TestModel.select().order_by(TestModel.embedding.l2_distance([1, 2, 3.1])).limit(5)
-```
-
-Get the distance
-
-```python
-TestModel.select(TestModel.embedding.cosine_distance([1, 2, 3.1]).alias('distance'))
-```
-
-Get within a certain distance
-
-```python
-TestModel.select().where(TestModel.embedding.l2_distance([1, 2, 3.1]) < 0.5)
-```
-
 ### TiDB Vector Client
 
 Within the framework, you can directly utilize the built-in `TiDBVectorClient`, as demonstrated by integrations like [Langchain](https://python.langchain.com/docs/integrations/vectorstores/tidb_vector) and [Llama index](https://docs.llamaindex.ai/en/stable/community/integrations/vector_stores.html#using-a-vector-store-as-an-index), to seamlessly interact with TiDB Vector. This approach abstracts away the need to manage the underlying ORM, simplifying your interaction with the vector store.

tests/peewee/test_peewee.py

@@ -1,7 +1,8 @@
 import numpy as np
 import pytest
 from peewee import MySQLDatabase, Model, OperationalError
-from tidb_vector.peewee import VectorField
+import tidb_vector
+from tidb_vector.peewee import VectorField, VectorAdaptor
 from ..config import TestConfig
 
 
@@ -273,3 +274,67 @@ def test_negative_inner_product(self):
         assert items.count() == 1
         assert items.get().id == item.id
         assert items[0].distance == -14
+
+
+class TestPeeweeAdaptor:
+    def setup_method(self):
+        db.drop_tables([Item1Model, Item2Model])
+        db.create_tables([Item1Model, Item2Model])
+
+    def teardown_method(self):
+        db.drop_tables([Item1Model, Item2Model])
+
+    def test_create_index_on_dyn_vector(self):
+        adaptor = VectorAdaptor(db)
+        with pytest.raises(ValueError):
+            adaptor.create_vector_index(
+                Item1Model.embedding, distance_metric=tidb_vector.DistanceMetric.L2
+            )
+        assert adaptor.has_vector_index(Item1Model.embedding) is False
+
+    def test_create_index_on_fixed_vector(self):
+        adaptor = VectorAdaptor(db)
+        adaptor.create_vector_index(
+            Item2Model.embedding, distance_metric=tidb_vector.DistanceMetric.L2
+        )
+        assert adaptor.has_vector_index(Item2Model.embedding) is True
+
+        with pytest.raises(Exception):
+            adaptor.create_vector_index(
+                Item2Model.embedding, distance_metric=tidb_vector.DistanceMetric.L2
+            )
+
+        assert adaptor.has_vector_index(Item2Model.embedding) is True
+
+        adaptor.create_vector_index(
+            Item2Model.embedding,
+            distance_metric=tidb_vector.DistanceMetric.L2,
+            skip_existing=True,
+        )
+
+        adaptor.create_vector_index(
+            Item2Model.embedding,
+            distance_metric=tidb_vector.DistanceMetric.COSINE,
+            skip_existing=True,
+        )
+
+    def test_index_and_search(self):
+        adaptor = VectorAdaptor(db)
+        adaptor.create_vector_index(
+            Item2Model.embedding, distance_metric=tidb_vector.DistanceMetric.L2
+        )
+        assert adaptor.has_vector_index(Item2Model.embedding) is True
+
+        Item2Model.insert_many(
+            [
+                {"embedding": [1, 2, 3]},
+                {"embedding": [1, 2, 3.2]},
+            ]
+        ).execute()
+
+        distance = Item2Model.embedding.cosine_distance([1, 2, 3])
+        items = (
+            Item2Model.select(distance.alias("distance")).order_by(distance).limit(5)
+        )
+        assert items.count() == 2
+        assert items[0].distance == 0.0

tidb_vector/__init__.py

@@ -1,4 +1,4 @@
-from .constants import MAX_DIM, MIN_DIM, DistanceMetric
+from .constants import MAX_DIM, MIN_DIM, DistanceMetric, VectorDataType
 
 __version__ = "0.0.13"
-__all__ = ["MAX_DIM", "MIN_DIM", "DistanceMetric"]
+__all__ = ["MAX_DIM", "MIN_DIM", "DistanceMetric", "VectorDataType"]

tidb_vector/constants.py

@@ -1,15 +1,37 @@
 import enum
+import typing
+
+import numpy
 
 # TiDB Vector has a limitation on the dimension length
 MAX_DIM = 16000
 MIN_DIM = 1
 
 
+VectorDataType = typing.Union[numpy.ndarray, typing.List[float]]
+
+
 class DistanceMetric(enum.Enum):
+    """
+    An enumeration representing different types of distance metrics.
+
+    - `DistanceMetric.L2`: L2 (Euclidean) distance metric.
+    - `DistanceMetric.COSINE`: Cosine distance metric.
+    """
+
     L2 = "L2"
     COSINE = "COSINE"
 
     def to_sql_func(self):
+        """
+        Converts the DistanceMetric to its corresponding SQL function name.
+
+        Returns:
+            str: The SQL function name.
+
+        Raises:
+            ValueError: If the DistanceMetric enum member is not supported.
+        """
         if self == DistanceMetric.L2:
             return "VEC_L2_DISTANCE"
         elif self == DistanceMetric.COSINE:

tidb_vector/peewee/__init__.py

@@ -1,32 +1,4 @@
-from peewee import Field, fn
+from .vector_type import VectorField
+from .adaptor import VectorAdaptor
 
-from tidb_vector.utils import decode_vector, encode_vector
-
-
-class VectorField(Field):
-    field_type = "VECTOR"
-
-    def __init__(self, dimensions=None, *args, **kwargs):
-        self.dimensions = dimensions
-        super(VectorField, self).__init__(*args, **kwargs)
-
-    def get_modifiers(self):
-        return self.dimensions and [self.dimensions] or None
-
-    def db_value(self, value):
-        return encode_vector(value)
-
-    def python_value(self, value):
-        return decode_vector(value)
-
-    def l1_distance(self, vector):
-        return fn.VEC_L1_DISTANCE(self, self.to_value(vector))
-
-    def l2_distance(self, vector):
-        return fn.VEC_L2_DISTANCE(self, self.to_value(vector))
-
-    def cosine_distance(self, vector):
-        return fn.VEC_COSINE_DISTANCE(self, self.to_value(vector))
-
-    def negative_inner_product(self, vector):
-        return fn.VEC_NEGATIVE_INNER_PRODUCT(self, self.to_value(vector))
+__all__ = ["VectorField", "VectorAdaptor"]