Repository: tensorchord/vechord

Issue:
feat: allow explicit configuration of the vector/keyword/maxsim index

Discussion:

Solution patch:
diff --git a/README.md b/README.md
index d5cfd03..f66961b 100644
--- a/README.md
+++ b/README.md
@@ -126,6 +126,17 @@ vr.run("https://paulgraham.com/best.html")  # only accept the arguments for the
 print(vr.search_by_vector(Chunk, emb.vectorize_query("startup")))
 ```
 
+### Customized Index Configuration
+
+```python
+from vechord.spec import VectorIndex
+
+class Chunk(Table, kw_only=True):
+    uid: Optional[PrimaryKeyAutoIncrease] = None
+    vector: Annotated[DenseVector, VectorIndex(distance="cos", lists=128)]
+    text: str
+```
+
 ### HTTP Service
 
 This creates a WSGI application that can be served by any WSGI server.
diff --git a/docs/source/api.md b/docs/source/api.md
index 59f87cc..f124f33 100644
--- a/docs/source/api.md
+++ b/docs/source/api.md
@@ -11,7 +11,7 @@
 
 ```{eval-rst}
 .. automodule:: vechord.spec
-   :members: Vector,ForeignKey,PrimaryKeyAutoIncrease,Table,Keyword
+   :members: Vector,ForeignKey,PrimaryKeyAutoIncrease,Table,Keyword,VectorIndex,MultiVectorIndex,KeywordIndex,IndexColumn
 ```
 
 ## Augment
diff --git a/docs/source/conf.py b/docs/source/conf.py
index f65899c..5ef8639 100644
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -46,7 +46,6 @@
 html_extra_path = ["robots.txt"]
 # myst
 myst_enable_extensions = ["tasklist"]
-myst_enable_checkboxes = True
 
 # -- Options for HTML output -------------------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#options-for-html-output
diff --git a/vechord/__init__.py b/vechord/__init__.py
index efad51e..ef284c5 100644
--- a/vechord/__init__.py
+++ b/vechord/__init__.py
@@ -13,7 +13,17 @@
 from vechord.registry import VechordRegistry
 from vechord.rerank import CohereReranker
 from vechord.service import create_web_app
-from vechord.spec import ForeignKey, PrimaryKeyAutoIncrease, Table, Vector
+from vechord.spec import (
+    ForeignKey,
+    IndexColumn,
+    Keyword,
+    KeywordIndex,
+    MultiVectorIndex,
+    PrimaryKeyAutoIncrease,
+    Table,
+    Vector,
+    VectorIndex,
+)
 
 __all__ = [
     "CohereReranker",
@@ -24,8 +34,12 @@
     "GeminiDenseEmbedding",
     "GeminiEvaluator",
     "GeminiExtractor",
+    "IndexColumn",
+    "Keyword",
+    "KeywordIndex",
     "LocalLoader",
     "Memory",
+    "MultiVectorIndex",
     "OpenAIDenseEmbedding",
     "PrimaryKeyAutoIncrease",
     "RegexChunker",
@@ -36,6 +50,7 @@
     "VechordRegistry",
     "Vector",
     "VectorChordClient",
+    "VectorIndex",
     "WordLlamaChunker",
     "create_web_app",
 ]
diff --git a/vechord/client.py b/vechord/client.py
index e82c222..6cc388a 100644
--- a/vechord/client.py
+++ b/vechord/client.py
@@ -7,7 +7,9 @@
 from pgvector.psycopg import register_vector
 from psycopg import sql
 
-from vechord.spec import Keyword
+from vechord.spec import IndexColumn, Keyword
+
+DEFAULT_TOKENIZER = ("bert_base_uncased", "wiki_tocken")
 
 active_cursor = contextvars.ContextVar("active_cursor", default=None)
 select_transaction_buffer = contextvars.ContextVar(
@@ -30,9 +32,14 @@ def __init__(self, namespace: str, url: str):
         self.ns = namespace
         self.url = url
         self.conn = psycopg.connect(url, autocommit=True)
-        self.conn.execute("CREATE EXTENSION IF NOT EXISTS vchord CASCADE")
-        self.conn.execute("CREATE EXTENSION IF NOT EXISTS vchord_bm25 CASCADE")
-        self.conn.execute('SET search_path TO "$user", public, bm25_catalog')
+        with self.transaction():
+            cursor = self.get_cursor()
+            cursor.execute("CREATE EXTENSION IF NOT EXISTS vchord CASCADE")
+            cursor.execute("CREATE EXTENSION IF NOT EXISTS vchord_bm25")
+            cursor.execute("CREATE EXTENSION IF NOT EXISTS pg_tokenizer")
+            cursor.execute(
+                'SET search_path TO "$user", public, bm25_catalog, tokenizer_catalog'
+            )
         register_vector(self.conn)
 
     @contextlib.contextmanager
@@ -69,61 +76,44 @@ def create_table_if_not_exists(self, name: str, schema: Sequence[tuple[str, str]
                 )
             )
 
-    def create_vector_index(self, name: str, column: str):
-        config = """
-        residual_quantization = true
-        [build.internal]
-        lists = [1]
-        spherical_centroids = false
-        """
+    def create_tokenizer(self):
         with self.transaction():
             cursor = self.get_cursor()
-            cursor.execute(
-                sql.SQL(
-                    "CREATE INDEX IF NOT EXISTS {index} ON "
-                    "{table} USING vchordrq ({column} vector_l2_ops) WITH "
-                    "(options = $${config}$$);"
-                ).format(
-                    table=sql.Identifier(f"{self.ns}_{name}"),
-                    index=sql.Identifier(f"{self.ns}_{name}_{column}_vec_idx"),
-                    column=sql.Identifier(column),
-                    config=sql.SQL(config),
-                )
-            )
+            try:
+                for name in DEFAULT_TOKENIZER:
+                    cursor.execute(
+                        sql.SQL(
+                            "SELECT create_tokenizer({name}, $$model={model}$$)"
+                        ).format(
+                            name=sql.Literal(name),
+                            model=sql.Identifier(name),
+                        )
+                    )
+            except psycopg.errors.DatabaseError as err:
+                if "already exists" not in str(err):
+                    raise
 
-    def create_multivec_index(self, name: str, column: str):
-        config = "build.internal.lists = []"
+    def create_index_if_not_exists(self, name: str, column: IndexColumn):
         with self.transaction():
             cursor = self.get_cursor()
-            cursor.execute(
-                sql.SQL(
-                    "CREATE INDEX IF NOT EXISTS {index} ON "
-                    "{table} USING vchordrq ({column} vector_maxsim_ops) WITH "
-                    "(options = $${config}$$);"
-                ).format(
-                    table=sql.Identifier(f"{self.ns}_{name}"),
-                    index=sql.Identifier(f"{self.ns}_{name}_{column}_multivec_idx"),
-                    column=sql.Identifier(column),
-                    config=sql.SQL(config),
-                )
+            query = sql.SQL(
+                "CREATE INDEX IF NOT EXISTS {index_name} ON {table} "
+                "USING {index} ({column} {op_name})"
+            ).format(
+                index_name=sql.Identifier(self._index_name(name, column)),
+                table=sql.Identifier(f"{self.ns}_{name}"),
+                index=sql.SQL(column.index.index),
+                column=sql.Identifier(column.name),
+                op_name=sql.SQL(column.index.op_name),
             )
-
-    def _keyword_index_name(self, name: str, column: str):
-        return f"{self.ns}_{name}_{column}_bm25_idx"
-
-    def create_keyword_index(self, name: str, column: str):
-        with self.transaction():
-            cursor = self.get_cursor()
-            cursor.execute(
-                sql.SQL(
-                    "CREATE INDEX IF NOT EXISTS {index} ON "
-                    "{table} USING bm25 ({column} bm25_ops);"
-                ).format(
-                    table=sql.Identifier(f"{self.ns}_{name}"),
-                    index=sql.Identifier(self._keyword_index_name(name, column)),
-                    column=sql.Identifier(column),
+            if config := column.index.config():
+                query += sql.SQL(" WITH (options = $${config}$$)").format(
+                    config=sql.SQL(config)
                 )
-            )
+            cursor.execute(query)
+
+    def _index_name(self, name: str, column: IndexColumn):
+        return f"{self.ns}_{name}_{column.name}_{column.index.name}"
 
     def select(
         self,
@@ -167,7 +157,7 @@ def _to_placeholder(kv: tuple[str, Any]):
         key, value = kv
         if isinstance(value, Keyword):
             return sql.SQL("tokenize({}, {})").format(
-                sql.Placeholder(key), sql.Literal(value._tokenizer)
+                sql.Placeholder(key), sql.Literal(value._model)
             )
         return sql.Placeholder(key)
 
@@ -205,7 +195,7 @@ def delete(self, name: str, kvs: dict):
     def query_vec(
         self,
         name: str,
-        vec_col: str,
+        vec_col: IndexColumn,
         vec: np.ndarray,
         return_fields: list[str],
         topk: int = 10,
@@ -213,11 +203,12 @@ def query_vec(
         columns = sql.SQL(", ").join(map(sql.Identifier, return_fields))
         cursor = self.conn.execute(
             sql.SQL(
-                "SELECT {columns} FROM {table} ORDER BY {vec_col} <-> %s LIMIT %s;"
+                "SELECT {columns} FROM {table} ORDER BY {vec_col} {op} %s LIMIT %s;"
             ).format(
                 table=sql.Identifier(f"{self.ns}_{name}"),
                 columns=columns,
-                vec_col=sql.Identifier(vec_col),
+                op=sql.SQL(vec_col.index.op_symbol),
+                vec_col=sql.Identifier(vec_col.name),
             ),
             (vec, topk),
         )
@@ -226,16 +217,19 @@ def query_vec(
     def query_multivec(  # noqa: PLR0913
         self,
         name: str,
-        multivec_col: str,
+        multivec_col: IndexColumn,
         vec: np.ndarray,
         max_maxsim_tuples: int,
+        probe: Optional[int],
         return_fields: list[str],
         topk: int = 10,
     ):
         columns = sql.SQL(", ").join(map(sql.Identifier, return_fields))
         with self.transaction():
             cursor = self.get_cursor()
-            cursor.execute("SET vchordrq.probes = '';")
+            cursor.execute(
+                sql.SQL("SET vchordrq.probes = {};").format(sql.Literal(probe or ""))
+            )
             cursor.execute(
                 sql.SQL("SET vchordrq.max_maxsim_tuples = {};").format(
                     sql.Literal(max_maxsim_tuples)
@@ -247,7 +241,7 @@ def query_multivec(  # noqa: PLR0913
                 ).format(
                     table=sql.Identifier(f"{self.ns}_{name}"),
                     columns=columns,
-                    multivec_col=sql.Identifier(multivec_col),
+                    multivec_col=sql.Identifier(multivec_col.name),
                 ),
                 (vec, topk),
             )
@@ -256,7 +250,7 @@ def query_multivec(  # noqa: PLR0913
     def query_keyword(  # noqa: PLR0913
         self,
         name: str,
-        keyword_col: str,
+        keyword_col: IndexColumn,
         keyword: str,
         return_fields: Sequence[str],
         tokenizer: str,
@@ -266,13 +260,13 @@ def query_keyword(  # noqa: PLR0913
         cursor = self.conn.execute(
             sql.SQL(
                 "SELECT {columns} FROM {table} ORDER BY {keyword_col} <&> "
-                "to_bm25query({index}, %s, {tokenizer}) LIMIT %s;"
+                "to_bm25query({index}, tokenize(%s, {tokenizer})) LIMIT %s;"
             ).format(
                 table=sql.Identifier(f"{self.ns}_{name}"),
                 columns=columns,
-                index=sql.Literal(self._keyword_index_name(name, keyword_col)),
+                index=sql.Literal(self._index_name(name, keyword_col)),
                 tokenizer=sql.Literal(tokenizer),
-                keyword_col=sql.Identifier(keyword_col),
+                keyword_col=sql.Identifier(keyword_col.name),
             ),
             (keyword, topk),
         )
diff --git a/vechord/registry.py b/vechord/registry.py
index cce919e..33f2d27 100644
--- a/vechord/registry.py
+++ b/vechord/registry.py
@@ -44,13 +44,14 @@ def __init__(self, namespace: str, url: str):
         self.tables: list[type[Table]] = []
         self.pipeline: list[Callable] = []
 
-    def register(self, tables: list[type[Table]]):
+    def register(self, tables: list[type[Table]], create_index: bool = True):
         """Register the given tables to the registry.
 
         This will create the tables in the database if not exists.
 
         Args:
             tables: a list of Table classes to be registered.
+            create_index: whether or not to create the index if not exists.
         """
         for table in tables:
             if not issubclass(table, Table):
@@ -58,28 +59,28 @@ def register(self, tables: list[type[Table]]):
 
             self.client.create_table_if_not_exists(table.name(), table.table_schema())
             logger.debug("create table %s if not exists", table.name())
-            if vector_column := table.vector_column():
-                self.client.create_vector_index(table.name(), vector_column)
-                logger.debug(
-                    "create vector index for %s.%s if not exists",
-                    table.name(),
-                    vector_column,
-                )
-            if multivec_column := table.multivec_column():
-                self.client.create_multivec_index(table.name(), multivec_column)
-                logger.debug(
-                    "create multivec index for %s.%s if not exists",
-                    table.name(),
-                    multivec_column,
-                )
-            if keyword_column := table.keyword_column():
-                self.client.create_keyword_index(table.name(), keyword_column)
+            self.tables.append(table)
+
+            if table.keyword_column() is not None:
+                self.client.create_tokenizer()
+
+            if not create_index:
+                continue
+
+            # create index
+            for index_column in (
+                table.vector_column(),
+                table.multivec_column(),
+                table.keyword_column(),
+            ):
+                if index_column is None:
+                    continue
+                self.client.create_index_if_not_exists(table.name(), index_column)
                 logger.debug(
-                    "create keyword index for %s.%s if not exists",
+                    "create index for %s.%s if not exists",
                     table.name(),
-                    keyword_column,
+                    index_column.name,
                 )
-            self.tables.append(table)
 
     def set_pipeline(self, pipeline: list[Callable]):
         """Set the pipeline to be executed in the `run` method."""
@@ -172,13 +173,14 @@ def search_by_vector(
             for r in res
         ]
 
-    def search_by_multivec(
+    def search_by_multivec(  # noqa: PLR0913
         self,
         cls: type[Table],
         multivec: np.ndarray,
         topk: int = 10,
         return_fields: Optional[Sequence[str]] = None,
         max_maxsim_tuples: int = 1000,
+        probe: Optional[int] = None,
     ) -> list[Table]:
         """Search the multivec for the given `Table` class.
 
@@ -188,6 +190,7 @@ def search_by_multivec(
             topk: the number of results to be returned.
             max_maxsim_tuples: the maximum number of tuples to be considered for
                 the each vector in the multivec.
+            probe: TODO
             return_fields: the fields to be returned, if not set, all the
                 non-[vector,keyword] fields will be returned.
         """
@@ -202,6 +205,7 @@ def search_by_multivec(
             multivec_col,
             multivec,
             max_maxsim_tuples=max_maxsim_tuples,
+            probe=probe,
             topk=topk,
             return_fields=fields,
         )
diff --git a/vechord/spec.py b/vechord/spec.py
index bd6ea93..dfb88aa 100644
--- a/vechord/spec.py
+++ b/vechord/spec.py
@@ -1,3 +1,7 @@
+import dataclasses
+import enum
+import inspect
+from abc import ABC, abstractmethod
 from datetime import datetime
 from types import UnionType
 from typing import (
@@ -142,17 +146,16 @@ class Keyword(str):
     `bm25vector` in PostgreSQL.
     """
 
-    _tokenizer: Literal["Unicode", "Bert", "Tocken"] = "Bert"
+    _model: Literal["bert_base_uncased", "wiki_tocken"] = "bert_base_uncased"
 
     @classmethod
     def schema(cls) -> str:
         return "bm25vector"
 
     @classmethod
-    def with_tokenizer(
-        cls, tokenizer: Literal["Unicode", "Bert", "Tocken"]
-    ) -> Type["Keyword"]:
-        cls._tokenizer = tokenizer
+    def with_model(cls, model: Literal["bert_base_uncased", "wiki_tocken"]) -> Type:
+        """TODO: test this"""
+        cls._model = model
         return cls
 
 
@@ -168,19 +171,23 @@ def with_tokenizer(
 }
 
 
-def is_optional_type(typ) -> bool:
+def is_optional_type(typ: Type) -> bool:
     return (get_origin(typ) is Union or get_origin(typ) is UnionType) and type(
         None
     ) in get_args(typ)
 
 
-def get_first_type_from_optional(typ) -> Type:
+def get_first_type_from_optional(typ: Type) -> Type:
     for arg in get_args(typ):
         if arg is not type(None):
             return arg
     raise ValueError(f"no non-None type in {typ}")
 
 
+def is_list_of_vector_type(typ: Type) -> bool:
+    return get_origin(typ) is list and issubclass(typ.__args__[0].__class__, VectorMeta)
+
+
 def type_to_psql(typ) -> str:
     if is_optional_type(typ):
         typ = get_first_type_from_optional(typ)
@@ -190,7 +197,7 @@ def type_to_psql(typ) -> str:
         origin = typ.__origin__
         schema = [type_to_psql(origin)]
         for m in meta:
-            if issubclass(m, ForeignKey):
+            if inspect.isclass(m) and issubclass(m, ForeignKey):
                 schema.append(m.schema())
         return " ".join(schema)
     elif get_origin(typ) is list:
@@ -202,6 +209,108 @@ def type_to_psql(typ) -> str:
     raise ValueError(f"unsupported type {typ}")
 
 
+class VectorDistance(enum.Enum):
+    L2 = "l2"
+    COS = "cos"
+    DOT = "dot"
+
+
+class BaseIndex(ABC):
+    name: str = ""
+    index: str = ""
+    op_name: str = ""
+    op_symbol: str = ""
+
+    def verify(self):
+        """This will construct a new data-only struct to verify all the fields.
+
+        This is to ensure that the `__post_init__` will not be called recursively.
+        """
+        fields = [
+            (
+                f.name,
+                f.type,
+                f.default
+                if f.default is not f.default_factory
+                else msgspec.field(default_factory=f.default_factory),
+            )
+            for f in dataclasses.fields(self)
+        ]
+        spec = msgspec.defstruct(f"Spec{self.__class__.__name__}", fields)
+        instance = msgspec.convert(dataclasses.asdict(self), spec)
+        for field in instance.__struct_fields__:
+            setattr(self, field, getattr(instance, field))
+
+    @abstractmethod
+    def config(self) -> str:
+        raise NotImplementedError
+
+
+class IndexColumn(msgspec.Struct, frozen=True, order=True):
+    name: str
+    index: BaseIndex
+
+
+@dataclasses.dataclass
+class VectorIndex(BaseIndex):
+    distance: VectorDistance = VectorDistance.L2
+    lists: int = 1
+
+    def __post_init__(self):
+        self.verify()
+        self.name = "vec_idx"
+        self.index = "vchordrq"
+        match self.distance:
+            case VectorDistance.L2:
+                self.op_symbol = "<->"
+                self.op_name = "vector_l2_ops"
+            case VectorDistance.COS:
+                self.op_symbol = "<=>"
+                self.op_name = "vector_cosine_ops"
+            case VectorDistance.DOT:
+                self.op_symbol = "<#>"
+                self.op_name = "vector_ip_ops"
+
+    def config(self):
+        is_l2 = self.distance == VectorDistance.L2
+        return f"""
+residual_quantization = {"true" if is_l2 else "false"}
+[build.internal]
+lists = [{self.lists}]
+spherical_centroids = {"false" if is_l2 else "true"}
+"""
+
+
+@dataclasses.dataclass
+class MultiVectorIndex(BaseIndex):
+    lists: Optional[int] = None
+
+    def __post_init__(self):
+        self.verify()
+        self.name = "multivec_idx"
+        self.index = "vchordrq"
+        self.op_name = "vector_maxsim_ops"
+        self.op_symbol = "@#"
+
+    def config(self):
+        return f"build.internal.lists = [{self.lists or ''}]"
+
+
+@dataclasses.dataclass
+class KeywordIndex(BaseIndex):
+    model: Literal["bert_base_uncased", "wiki_tocken"] = "bert_base_uncased"
+
+    def __post_init__(self):
+        self.verify()
+        self.name = "keyword_idx"
+        self.index = "bm25"
+        self.op_name = "bm25_ops"
+        self.op_symbol = "<&>"
+
+    def config(self):
+        return ""
+
+
 class Storage(msgspec.Struct):
     @classmethod
     def name(cls) -> str:
@@ -228,35 +337,56 @@ def table_schema(cls) -> Sequence[tuple[str, str]]:
         return tuple((name, type_to_psql(typ)) for name, typ in hints.items())
 
     @classmethod
-    def vector_column(cls) -> Optional[str]:
+    def vector_column(cls) -> Optional[IndexColumn]:
         """Get the vector column name."""
         for name, typ in get_type_hints(cls, include_extras=True).items():
             if issubclass(typ.__class__, VectorMeta):
-                return name
+                return IndexColumn(name, VectorIndex())
+            elif get_origin(typ) is Annotated and issubclass(
+                typ.__origin__.__class__, VectorMeta
+            ):
+                for m in typ.__metadata__:
+                    if isinstance(m, VectorIndex):
+                        return IndexColumn(name, m)
         return None
 
     @classmethod
-    def multivec_column(cls) -> Optional[str]:
+    def multivec_column(cls) -> Optional[IndexColumn]:
         """Get the multivec column name."""
         for name, typ in get_type_hints(cls, include_extras=True).items():
-            if get_origin(typ) is list and issubclass(
-                typ.__args__[0].__class__, VectorMeta
-            ):
-                return name
+            if is_list_of_vector_type(typ):
+                return IndexColumn(name, MultiVectorIndex())
+            elif get_origin(typ) is Annotated:
+                inner_type = typ.__origin__
+                if is_list_of_vector_type(inner_type):
+                    for m in typ.__metadata__:
+                        if isinstance(m, MultiVectorIndex):
+                            return IndexColumn(name, m)
         return None
 
     @classmethod
-    def keyword_column(cls) -> Optional[str]:
+    def keyword_column(cls) -> Optional[IndexColumn]:
         """Get the keyword column name."""
         for name, typ in get_type_hints(cls, include_extras=True).items():
             if typ is Keyword:
-                return name
+                return IndexColumn(name, KeywordIndex())
+            elif get_origin(typ) is Annotated and typ.__origin__ is Keyword:
+                for m in typ.__metadata__:
+                    if isinstance(m, KeywordIndex):
+                        return IndexColumn(name, m)
         return None
 
     @classmethod
     def non_vec_columns(cls) -> Sequence[str]:
         """Get the column names that are not vector or keyword."""
-        exclude = (cls.vector_column(), cls.keyword_column(), cls.multivec_column())
+        exclude = (
+            col.name if col else col
+            for col in (
+                cls.vector_column(),
+                cls.keyword_column(),
+                cls.multivec_column(),
+            )
+        )
         return tuple(field for field in cls.fields() if field not in exclude)
 
     @classmethod
@@ -264,7 +394,9 @@ def keyword_tokenizer(cls) -> Optional[str]:
         """Get the keyword tokenizer."""
         for _, typ in get_type_hints(cls, include_extras=True).items():
             if typ is Keyword:
-                return typ._tokenizer
+                return typ._model
+            elif get_origin(typ) is Annotated and typ.__origin__ is Keyword:
+                return typ.__origin__._model
         return None
 
     @classmethod


Test patch:
diff --git a/tests/test_spec.py b/tests/test_spec.py
index a94b013..04eb7b4 100644
--- a/tests/test_spec.py
+++ b/tests/test_spec.py
@@ -5,7 +5,16 @@
 import numpy as np
 import pytest
 
-from vechord.spec import ForeignKey, Keyword, PrimaryKeyAutoIncrease, Table, Vector
+from vechord.spec import (
+    ForeignKey,
+    Keyword,
+    MultiVectorIndex,
+    PrimaryKeyAutoIncrease,
+    Table,
+    Vector,
+    VectorDistance,
+    VectorIndex,
+)
 
 
 class Document(Table, kw_only=True):
@@ -47,9 +56,9 @@ def test_table_cls_methods():
     assert Chunk.primary_key() == "uid", Chunk
 
     assert Document.vector_column() is None
-    assert Chunk.vector_column() == "vec"
-    assert Chunk.multivec_column() == "multivec"
-    assert Chunk.keyword_column() == "keyword"
+    assert Chunk.vector_column().name == "vec"
+    assert Chunk.multivec_column().name == "multivec"
+    assert Chunk.keyword_column().name == "keyword"
 
     def find_schema_by_name(schema, name):
         for n, t in schema:
@@ -77,3 +86,43 @@ def test_vector_type():
         Dense(np.random.rand(123))
 
     assert np.equal(Dense(np.ones(128)), Dense([1.0] * 128)).all()
+
+
+def test_index():
+    with pytest.raises(msgspec.ValidationError):
+        VectorIndex(distance="bug")
+
+    with pytest.raises(msgspec.ValidationError):
+        VectorIndex(lists="bug")
+
+    lists = 128
+    vec_idx = VectorIndex(distance="l2", lists=lists)
+    assert vec_idx.distance is VectorDistance.L2
+    assert vec_idx.lists == lists
+    assert vec_idx.op_symbol == "<->"
+    assert str(vec_idx.lists) in vec_idx.config()
+
+    multivec_idx = MultiVectorIndex(lists=1)
+    assert multivec_idx.lists == 1
+    assert MultiVectorIndex().config() == "build.internal.lists = []"
+
+
+def test_annotated_index():
+    lists = 8
+
+    class Sentence(Table, kw_only=True):
+        uid: PrimaryKeyAutoIncrease | None = None
+        text: str
+        vec: Annotated[Vector[128], VectorIndex(distance="cos", lists=lists)]
+        vecs: Annotated[list[Vector[128]], MultiVectorIndex(lists=lists)]
+
+    vec_col = Sentence.vector_column()
+    assert vec_col
+    assert vec_col.name == "vec"
+    assert vec_col.index.op_name == "vector_cosine_ops"
+    assert vec_col.index.lists == lists
+
+    multivec_col = Sentence.multivec_column()
+    assert multivec_col
+    assert multivec_col.name == "vecs"
+    assert multivec_col.index.lists == lists
diff --git a/tests/test_table.py b/tests/test_table.py
index 1177907..688889f 100644
--- a/tests/test_table.py
+++ b/tests/test_table.py
@@ -1,6 +1,6 @@
 from datetime import datetime
 from os import environ
-from typing import Annotated
+from typing import Annotated, Optional
 
 import msgspec
 import numpy as np
@@ -8,7 +8,14 @@
 
 from vechord.log import logger
 from vechord.registry import VechordRegistry
-from vechord.spec import ForeignKey, Keyword, PrimaryKeyAutoIncrease, Table, Vector
+from vechord.spec import (
+    ForeignKey,
+    Keyword,
+    PrimaryKeyAutoIncrease,
+    Table,
+    Vector,
+    VectorIndex,
+)
 
 URL = "127.0.0.1"
 # for local container development environment, use the host machine's IP
@@ -38,6 +45,12 @@ class Chunk(Table, kw_only=True):
     keyword: Keyword
 
 
+class AnnotatedChunk(Table, kw_only=True):
+    uid: Optional[PrimaryKeyAutoIncrease] = None
+    text: str
+    vector: Annotated[DenseVector, VectorIndex(distance="cos", lists=2)]
+
+
 class Sentence(Table, kw_only=True):
     uid: PrimaryKeyAutoIncrease | None = None
     text: str
@@ -80,6 +93,21 @@ def test_insert_select_remove(registry):
     assert len(registry.select_by(Document.partial_init())) == 1
 
 
+@pytest.mark.db
+def test_annotated_index(registry):
+    registry.register([AnnotatedChunk])
+    num = 100
+    topk = 5
+    for text in (f"hello {i}" for i in range(num)):
+        registry.insert(AnnotatedChunk(text=text, vector=gen_vector()))
+
+    inserted = registry.select_by(AnnotatedChunk.partial_init(), fields=["text"])
+    assert len(inserted) == num
+
+    res = registry.search_by_vector(AnnotatedChunk, gen_vector(), topk=topk)
+    assert len(res) == topk
+
+
 @pytest.mark.db
 def test_foreign_key(registry):
     docs = [