rotate: add an optional random rotation feature in INT8/INT4 quantization method (#483)

Co-authored-by: rayx <rui.xing@alibaba-inc.com>
Co-authored-by: Jalin Wang <wangjianning.wjn@alibaba-inc.com>
Co-authored-by: Copilot Autofix powered by AI <223894421+github-code-quality[bot]@users.noreply.github.com>

python/tests/test_params.py

@@ -36,6 +36,7 @@ from zvec import (
     IndexType,
     MetricType,
     QuantizeType,
+    QuantizerParam,
     DataType,
     VectorSchema,
 )
@@ -452,3 +453,88 @@ class TestVectorQueryDeprecated:
             warnings.simplefilter("always")
             vq = VectorQuery(field_name="embedding", id="doc123")
         assert isinstance(vq, Query)
+
+
+# ----------------------------
+# QuantizerParam Test Case
+# ----------------------------
+
+
+class TestQuantizerParam:
+    def test_default(self):
+        qp = QuantizerParam()
+        assert qp.enable_rotate is False
+
+    def test_enable_rotate_true(self):
+        qp = QuantizerParam(enable_rotate=True)
+        assert qp.enable_rotate is True
+
+    def test_enable_rotate_false(self):
+        qp = QuantizerParam(enable_rotate=False)
+        assert qp.enable_rotate is False
+
+    def test_equality(self):
+        qp1 = QuantizerParam(enable_rotate=True)
+        qp2 = QuantizerParam(enable_rotate=True)
+        qp3 = QuantizerParam(enable_rotate=False)
+        assert qp1 == qp2
+        assert qp1 != qp3
+
+    def test_to_dict(self):
+        qp = QuantizerParam(enable_rotate=True)
+        d = qp.to_dict()
+        assert isinstance(d, dict)
+        assert d.get("enable_rotate") is True
+
+    def test_repr(self):
+        qp = QuantizerParam(enable_rotate=True)
+        r = repr(qp)
+        assert "enable_rotate" in r or "QuantizerParam" in r
+
+    def test_pickle_roundtrip(self):
+        import pickle
+
+        qp = QuantizerParam(enable_rotate=True)
+        data = pickle.dumps(qp)
+        qp2 = pickle.loads(data)
+        assert qp2.enable_rotate is True
+        assert qp == qp2
+
+
+# ----------------------------
+# HnswIndexParam with QuantizerParam
+# ----------------------------
+
+
+class TestHnswIndexParamQuantizer:
+    def test_default_quantizer_param(self):
+        param = HnswIndexParam()
+        assert param.quantizer_param is not None
+        assert param.quantizer_param.enable_rotate is False
+
+    def test_with_quantizer_param(self):
+        qp = QuantizerParam(enable_rotate=True)
+        param = HnswIndexParam(
+            metric_type=MetricType.L2,
+            quantize_type=QuantizeType.INT8,
+            quantizer_param=qp,
+        )
+        assert param.quantizer_param.enable_rotate is True
+        assert param.quantize_type == QuantizeType.INT8
+
+
+# ----------------------------
+# FlatIndexParam with QuantizerParam
+# ----------------------------
+
+
+class TestFlatIndexParamQuantizer:
+    def test_with_quantizer_param(self):
+        qp = QuantizerParam(enable_rotate=True)
+        param = FlatIndexParam(
+            metric_type=MetricType.L2,
+            quantize_type=QuantizeType.INT8,
+            quantizer_param=qp,
+        )
+        assert param.quantizer_param.enable_rotate is True
+        assert param.quantize_type == QuantizeType.INT8

python/zvec/__init__.py

@@ -108,6 +108,7 @@ from .model.param import (
     IVFIndexParam,
     IVFQueryParam,
     OptimizeOption,
+    QuantizerParam,
     VamanaIndexParam,
     VamanaQueryParam,
 )
@@ -171,6 +172,7 @@ __all__ = [
     "HnswQueryParam",
     "HnswRabitqQueryParam",
     "IVFQueryParam",
+    "QuantizerParam",
     "VamanaIndexParam",
     "VamanaQueryParam",
     # Extensions

python/zvec/__init__.pyi

@@ -30,6 +30,7 @@ from .model.param import (
     IVFIndexParam,
     IVFQueryParam,
     OptimizeOption,
+    QuantizerParam,
     VamanaIndexParam,
     VamanaQueryParam,
 )
@@ -79,6 +80,7 @@ __all__: list = [
     "MetricType",
     "OptimizeOption",
     "QuantizeType",
+    "QuantizerParam",
     "Query",
     "ReRanker",
     "RrfReRanker",

python/zvec/model/param/__init__.py

@@ -31,6 +31,7 @@ from zvec._zvec.param import (
     IVFIndexParam,
     IVFQueryParam,
     OptimizeOption,
+    QuantizerParam,
     VamanaIndexParam,
     VamanaQueryParam,
 )
@@ -53,6 +54,7 @@ __all__ = [
     "IndexOption",
     "InvertIndexParam",
     "OptimizeOption",
+    "QuantizerParam",
     "VamanaIndexParam",
     "VamanaQueryParam",
 ]

python/zvec/model/param/__init__.pyi

@@ -26,6 +26,7 @@ __all__: list[str] = [
     "IndexParam",
     "InvertIndexParam",
     "OptimizeOption",
+    "QuantizerParam",
     "QueryParam",
     "SegmentOption",
     "VectorIndexParam",
@@ -147,6 +148,8 @@ class FlatIndexParam(VectorIndexParam):
         quantize_type (QuantizeType): Optional quantization type for vector
             compression (e.g., FP16, INT8). Use ``QuantizeType.UNDEFINED`` to
             disable quantization. Default is ``QuantizeType.UNDEFINED``.
+        quantizer_param (QuantizerParam): Optional quantizer parameters. See
+            ``QuantizerParam`` for available options. Default is ``QuantizerParam()``.
 
     Examples:
         >>> from zvec.typing import MetricType, QuantizeType
@@ -163,6 +166,7 @@ class FlatIndexParam(VectorIndexParam):
         self,
         metric_type: zvec._zvec.typing.MetricType = ...,
         quantize_type: zvec._zvec.typing.QuantizeType = ...,
+        quantizer_param: QuantizerParam = ...,
     ) -> None:
         """
         Constructs a FlatIndexParam instance.
@@ -171,6 +175,8 @@ class FlatIndexParam(VectorIndexParam):
             metric_type (MetricType, optional): Distance metric. Defaults to MetricType.IP.
             quantize_type (QuantizeType, optional): Vector quantization type.
                 Defaults to QuantizeType.UNDEFINED (no quantization).
+            quantizer_param (QuantizerParam, optional): Quantizer configuration.
+                Defaults to QuantizerParam().
         """
 
     def __repr__(self) -> str: ...
@@ -226,6 +232,7 @@ class HnswIndexParam(VectorIndexParam):
         ef_construction: typing.SupportsInt = 500,
         quantize_type: zvec._zvec.typing.QuantizeType = ...,
         use_contiguous_memory: bool = False,
+        quantizer_param: QuantizerParam = ...,
     ) -> None: ...
     def __repr__(self) -> str: ...
     def __setstate__(self, arg0: tuple) -> None: ...
@@ -283,6 +290,7 @@ class HnswQueryParam(QueryParam):
         >>> print(params.to_dict() if hasattr(params, 'to_dict') else params)
         {"type":"HNSW", "ef":300}
     """
+
     def __getstate__(self) -> tuple: ...
     def __init__(
         self,
@@ -307,6 +315,7 @@ class HnswQueryParam(QueryParam):
                 - ``prefetch_lines`` (int): Number of 64B cache lines to prefetch
                   per neighbour vector (PL). ``0`` (default) means auto-derive from vector size.
         """
+
     def __repr__(self) -> str: ...
     def __setstate__(self, arg0: tuple) -> None: ...
     @property
@@ -430,6 +439,7 @@ class HnswRabitqQueryParam(QueryParam):
         >>> print(params.ef)
         300
     """
+
     def __getstate__(self) -> tuple: ...
     def __init__(
         self,
@@ -448,6 +458,7 @@ class HnswRabitqQueryParam(QueryParam):
             is_linear (bool, optional): Force linear search. Default is False.
             is_using_refiner (bool, optional): Whether to use refiner for the query. Default is False.
         """
+
     def __repr__(self) -> str: ...
     def __setstate__(self, arg0: tuple) -> None: ...
     @property
@@ -498,6 +509,7 @@ class IVFIndexParam(VectorIndexParam):
         n_iters: typing.SupportsInt = 10,
         use_soar: bool = False,
         quantize_type: zvec._zvec.typing.QuantizeType = ...,
+        quantizer_param: QuantizerParam = ...,
     ) -> None:
         """
         Constructs an IVFIndexParam instance.
@@ -511,6 +523,8 @@ class IVFIndexParam(VectorIndexParam):
             use_soar (bool, optional): Enable SOAR optimization. Defaults to False.
             quantize_type (QuantizeType, optional): Vector quantization type.
                 Defaults to QuantizeType.UNDEFINED.
+            quantizer_param (QuantizerParam, optional): Quantizer configuration.
+                Defaults to QuantizerParam().
         """
 
     def __repr__(self) -> str: ...
@@ -558,6 +572,7 @@ class IVFQueryParam(QueryParam):
         >>> print(params.nprobe)
         20
     """
+
     def __getstate__(self) -> tuple: ...
     def __init__(self, nprobe: typing.SupportsInt = 10) -> None:
         """
@@ -567,6 +582,7 @@ class IVFQueryParam(QueryParam):
             nprobe (int, optional): Number of inverted lists to probe during search.
                 Higher values improve accuracy. Defaults to 10.
         """
+
     def __repr__(self) -> str: ...
     def __setstate__(self, arg0: tuple) -> None: ...
     @property
@@ -592,6 +608,7 @@ class VamanaIndexParam(VectorIndexParam):
     Examples:
         >>> params = VamanaIndexParam(metric_type=MetricType.COSINE, max_degree=64)
     """
+
     def __getstate__(self) -> tuple: ...
     def __init__(
         self,
@@ -610,18 +627,23 @@ class VamanaIndexParam(VectorIndexParam):
     @property
     def max_degree(self) -> int:
         """int: Maximum out-degree (R) of every node in the Vamana graph."""
+
     @property
     def search_list_size(self) -> int:
         """int: Candidate list size during Vamana graph construction."""
+
     @property
     def alpha(self) -> float:
         """float: Vamana RobustPrune alpha factor."""
+
     @property
     def saturate_graph(self) -> bool:
         """bool: Whether to saturate every node to max_degree neighbors."""
+
     @property
     def use_contiguous_memory(self) -> bool:
         """bool: Whether to allocate a single contiguous memory arena."""
+
     @property
     def use_id_map(self) -> bool:
         """bool: Reserved flag for engine-level id remapping."""
@@ -644,6 +666,7 @@ class VamanaQueryParam(QueryParam):
         >>> print(params.ef_search)
         200
     """
+
     def __getstate__(self) -> tuple: ...
     def __init__(
         self,
@@ -667,14 +690,17 @@ class VamanaQueryParam(QueryParam):
                 - ``prefetch_lines`` (int): Cache lines to prefetch per vector (PL).
                   ``0`` (default) means auto-derive from vector size.
         """
+
     def __repr__(self) -> str: ...
     def __setstate__(self, arg0: tuple) -> None: ...
     @property
     def ef_search(self) -> int:
         """int: Size of the dynamic candidate list during Vamana search."""
+
     @property
     def prefetch_offset(self) -> int:
         """int: Graph prefetch offset used by the Vamana fast path."""
+
     @property
     def prefetch_lines(self) -> int:
         """int: Override of prefetch cache lines per vector (0=auto)."""
@@ -760,6 +786,7 @@ class FtsQueryParam(QueryParam):
         >>> print(params.default_operator)
         AND
     """
+
     def __getstate__(self) -> tuple: ...
     def __init__(
         self,
@@ -772,6 +799,7 @@ class FtsQueryParam(QueryParam):
             default_operator (str, optional): Default boolean operator for adjacent
                 bare terms. Supported: "OR", "AND". Defaults to "" (uses engine default).
         """
+
     def __repr__(self) -> str: ...
     def __setstate__(self, arg0: tuple) -> None: ...
     @property
@@ -865,6 +893,7 @@ class InvertIndexParam(IndexParam):
         >>> print(config)
         {'enable_range_optimization': True, 'enable_extended_wildcard': False}
     """
+
     def __getstate__(self) -> tuple: ...
     def __init__(
         self,
@@ -880,18 +909,21 @@ class InvertIndexParam(IndexParam):
             enable_extended_wildcard (bool, optional): If True, enables extended wildcard
                 search including suffix and infix patterns. Defaults to False.
         """
+
     def __repr__(self) -> str: ...
     def __setstate__(self, arg0: tuple) -> None: ...
     def to_dict(self) -> dict:
         """
         Convert to dictionary with all fields
         """
+
     @property
     def enable_extended_wildcard(self) -> bool:
         """
         bool: Whether extended wildcard (suffix and infix) search is enabled.
         Note: Prefix search is always enabled regardless of this setting.
         """
+
     @property
     def enable_range_optimization(self) -> bool:
         """
@@ -948,6 +980,7 @@ class QueryParam:
             using the index. Useful for debugging or small datasets. Default is False.
         is_using_refiner (bool, optional): Whether to use refiner for the query. Default is False.
     """
+
     def __getstate__(self) -> tuple: ...
     def __setstate__(self, arg0: tuple) -> None: ...
     @property
@@ -955,16 +988,19 @@ class QueryParam:
         """
         bool: Whether to bypass the index and use brute-force linear search.
         """
+
     @property
     def is_using_refiner(self) -> bool:
         """
         bool: Whether to use refiner for the query.
         """
+
     @property
     def radius(self) -> float:
         """
         IndexType: The type of index this query targets.
         """
+
     @property
     def type(self) -> zvec._zvec.typing.IndexType:
         """
@@ -1015,6 +1051,49 @@ class SegmentOption:
         bool: Whether the segment is read-only.
         """
 
+class QuantizerParam:
+    """
+
+    Optional parameters for quantizer configuration.
+
+    This class is only needed when customizing quantizer behavior (e.g., enabling
+    random rotation). It can be omitted for default quantization settings.
+
+    Attributes:
+        enable_rotate (bool): Whether to apply random rotation before INT8/INT4
+            quantization to reduce quantization error.
+            Only effective with quantize_type=INT8 or INT4. Defaults to False.
+
+    Examples:
+        >>> qp = QuantizerParam(enable_rotate=True)
+        >>> print(qp.enable_rotate)
+        True
+    """
+
+    def __getstate__(self) -> tuple: ...
+    def __init__(self, enable_rotate: bool = False) -> None:
+        """
+        Constructs a QuantizerParam instance.
+
+        Args:
+            enable_rotate (bool, optional): Whether to apply random rotation
+                before INT8/INT4 quantization. Defaults to False.
+        """
+
+    def __repr__(self) -> str: ...
+    def __setstate__(self, arg0: tuple) -> None: ...
+    def __eq__(self, arg0: typing.Any) -> bool: ...
+    def to_dict(self) -> dict:
+        """
+        Convert to dictionary with all fields
+        """
+
+    @property
+    def enable_rotate(self) -> bool:
+        """
+        bool: Whether random rotation is enabled before INT8/INT4 quantization.
+        """
+
 class VectorIndexParam(IndexParam):
     """
 
@@ -1026,6 +1105,7 @@ class VectorIndexParam(IndexParam):
         type (IndexType): The specific vector index type (e.g., HNSW, FLAT).
         metric_type (MetricType): Distance metric used for similarity search.
         quantize_type (QuantizeType): Optional vector quantization type.
+        quantizer_param (QuantizerParam): Optional quantizer parameters.
     """
 
     def __getstate__(self) -> tuple: ...
@@ -1047,6 +1127,12 @@ class VectorIndexParam(IndexParam):
         QuantizeType: Vector quantization type (e.g., FP16, INT8).
         """
 
+    @property
+    def quantizer_param(self) -> QuantizerParam:
+        """
+        QuantizerParam: Quantizer configuration including enable_rotate.
+        """
+
 class _SearchQuery:
     field_name: str
     filter: str