From b6f2d4367d1bdd69b073ecf43873f37497e4821e Mon Sep 17 00:00:00 2001
From: Qinren Zhou <zhouqinren.zqr@alibaba-inc.com>
Date: Thu, 2 Apr 2026 20:03:15 +0800
Subject: [PATCH] fix: fix python type annotations (#307)

---
 python/zvec/__init__.py                  |  2 +-
 python/zvec/__init__.pyi                 |  4 ++++
 python/zvec/model/collection.py          | 16 +++++++++++++---
 python/zvec/model/param/vector_query.py  |  4 ++--
 python/zvec/model/schema/field_schema.py |  9 ++++++---
 5 files changed, 26 insertions(+), 9 deletions(-)

diff --git a/python/zvec/__init__.py b/python/zvec/__init__.py
index 8958978..40db92b 100644
--- a/python/zvec/__init__.py
+++ b/python/zvec/__init__.py
@@ -112,7 +112,6 @@ __all__ = [
     "InvertIndexParam",
     "HnswIndexParam",
     "HnswRabitqIndexParam",
-    "HnswRabitqQueryParam",
     "FlatIndexParam",
     "IVFIndexParam",
     "CollectionOption",
@@ -121,6 +120,7 @@ __all__ = [
     "AddColumnOption",
     "AlterColumnOption",
     "HnswQueryParam",
+    "HnswRabitqQueryParam",
     "IVFQueryParam",
     # Extensions
     "DenseEmbeddingFunction",
diff --git a/python/zvec/__init__.pyi b/python/zvec/__init__.pyi
index efb1b2d..97a57a1 100644
--- a/python/zvec/__init__.pyi
+++ b/python/zvec/__init__.pyi
@@ -19,6 +19,8 @@ from .model.param import (
     FlatIndexParam,
     HnswIndexParam,
     HnswQueryParam,
+    HnswRabitqIndexParam,
+    HnswRabitqQueryParam,
     IndexOption,
     InvertIndexParam,
     IVFIndexParam,
@@ -54,6 +56,8 @@ __all__: list = [
     "FlatIndexParam",
     "HnswIndexParam",
     "HnswQueryParam",
+    "HnswRabitqIndexParam",
+    "HnswRabitqQueryParam",
     "IVFIndexParam",
     "IVFQueryParam",
     "IndexOption",
diff --git a/python/zvec/model/collection.py b/python/zvec/model/collection.py
index eee571d..f7a12b4 100644
--- a/python/zvec/model/collection.py
+++ b/python/zvec/model/collection.py
@@ -28,6 +28,7 @@ from .param import (
     CollectionOption,
     FlatIndexParam,
     HnswIndexParam,
+    HnswRabitqIndexParam,
     IndexOption,
     InvertIndexParam,
     IVFIndexParam,
@@ -38,7 +39,12 @@ from .schema import CollectionSchema, CollectionStats, FieldSchema
 
 __all__ = ["Collection"]
 
-_VECTOR_INDEX_TYPES = (HnswIndexParam, IVFIndexParam, FlatIndexParam)
+_VECTOR_INDEX_TYPES = (
+    HnswIndexParam,
+    HnswRabitqIndexParam,
+    IVFIndexParam,
+    FlatIndexParam,
+)
 
 
 class Collection:
@@ -107,7 +113,11 @@ class Collection:
         self,
         field_name: str,
         index_param: Union[
-            HnswIndexParam, IVFIndexParam, FlatIndexParam, InvertIndexParam
+            HnswIndexParam,
+            HnswRabitqIndexParam,
+            IVFIndexParam,
+            FlatIndexParam,
+            InvertIndexParam,
         ],
         option: IndexOption = IndexOption(),
     ) -> None:
@@ -118,7 +128,7 @@ class Collection:
 
         Args:
             field_name (str): Name of the field to index.
-            index_param (Union[HnswIndexParam, IVFIndexParam, FlatIndexParam, InvertIndexParam]):
+            index_param (Union[HnswIndexParam, HnswRabitqIndexParam, IVFIndexParam, FlatIndexParam, InvertIndexParam]):
                 Index configuration.
             option (Optional[IndexOption], optional): Index creation options.
                 Defaults to ``IndexOption()``.
diff --git a/python/zvec/model/param/vector_query.py b/python/zvec/model/param/vector_query.py
index 97d105a..312a7ca 100644
--- a/python/zvec/model/param/vector_query.py
+++ b/python/zvec/model/param/vector_query.py
@@ -17,7 +17,7 @@ from dataclasses import dataclass
 from typing import Optional, Union
 
 from ...common import VectorType
-from . import HnswQueryParam, IVFQueryParam
+from . import HnswQueryParam, HnswRabitqQueryParam, IVFQueryParam
 
 __all__ = ["VectorQuery"]
 
@@ -55,7 +55,7 @@ class VectorQuery:
     field_name: str
     id: Optional[str] = None
     vector: VectorType = None
-    param: Optional[Union[HnswQueryParam, IVFQueryParam]] = None
+    param: Optional[Union[HnswQueryParam, HnswRabitqQueryParam, IVFQueryParam]] = None
 
     def has_id(self) -> bool:
         """Check if the query is based on a document ID.
diff --git a/python/zvec/model/schema/field_schema.py b/python/zvec/model/schema/field_schema.py
index da193dd..1af00b5 100644
--- a/python/zvec/model/schema/field_schema.py
+++ b/python/zvec/model/schema/field_schema.py
@@ -21,6 +21,7 @@ from _zvec.schema import _FieldSchema
 from zvec.model.param import (
     FlatIndexParam,
     HnswIndexParam,
+    HnswRabitqIndexParam,
     InvertIndexParam,
     IVFIndexParam,
 )
@@ -209,7 +210,7 @@ class VectorSchema:
         data_type: DataType,
         dimension: Optional[int] = 0,
         index_param: Optional[
-            Union[HnswIndexParam, FlatIndexParam, IVFIndexParam]
+            Union[HnswIndexParam, HnswRabitqIndexParam, FlatIndexParam, IVFIndexParam]
         ] = None,
     ):
         if name is None or not isinstance(name, str):
@@ -263,8 +264,10 @@ class VectorSchema:
         return self._cpp_obj.dimension
 
     @property
-    def index_param(self) -> Union[HnswIndexParam, IVFIndexParam, FlatIndexParam]:
-        """Union[HnswIndexParam, IVFIndexParam, FlatIndexParam]: Index configuration for the vector."""
+    def index_param(
+        self,
+    ) -> Union[HnswIndexParam, HnswRabitqIndexParam, IVFIndexParam, FlatIndexParam]:
+        """Union[HnswIndexParam, HnswRabitqIndexParam, IVFIndexParam, FlatIndexParam]: Index configuration for the vector."""
         return self._cpp_obj.index_param
 
     def __dict__(self) -> dict[str, Any]: