Merge pull request #1745 from weaviate/rq_add_rescore_limit

RQ: Add rescoreLimit parameter

Author: dirkkul

integration/test_collection_config.py

@@ -622,7 +622,7 @@ def test_hnsw_with_rq(collection_factory: CollectionFactory) -> None:
     collection = collection_factory(
         vector_index_config=Configure.VectorIndex.hnsw(
             vector_cache_max_objects=5,
-            quantizer=Configure.VectorIndex.Quantizer.rq(bits=8),
+            quantizer=Configure.VectorIndex.Quantizer.rq(bits=8, rescore_limit=20),
         ),
     )
 
@@ -633,6 +633,7 @@ def test_hnsw_with_rq(collection_factory: CollectionFactory) -> None:
     assert isinstance(config.vector_index_config.quantizer, _RQConfig)
     assert config.vector_index_config.quantizer is not None
     assert config.vector_index_config.quantizer.bits == 8
+    assert config.vector_index_config.quantizer.rescore_limit == 20
 
 
 @pytest.mark.parametrize(

mock_tests/mock_data.py

@@ -63,7 +63,7 @@
         },
         "bq": {"enabled": False},
         "sq": {"enabled": False, "trainingLimit": 100000, "rescoreLimit": 20},
-        "rq": {"enabled": False, "bits": 8},
+        "rq": {"enabled": False, "bits": 8, "rescoreLimit": 20},
     },
     "vectorIndexType": "hnsw",
     "vectorizer": "text2vec-contextionary",

test/collection/schema.py

@@ -66,6 +66,7 @@ def multi_vector_schema(quantizer: Optional[Literal["pq", "bq", "sq", "rq"]] = N
                     "rq": {
                         "enabled": quantizer == "rq",
                         "bits": 8,
+                        "rescoreLimit": 20,
                     },
                 },
                 "vectorIndexType": "hnsw",
@@ -101,6 +102,7 @@ def multi_vector_schema(quantizer: Optional[Literal["pq", "bq", "sq", "rq"]] = N
                     "rq": {
                         "enabled": quantizer == "rq",
                         "bits": 8,
+                        "rescoreLimit": 20,
                     },
                 },
                 "vectorIndexType": "hnsw",

test/collection/test_config.py

@@ -1368,13 +1368,14 @@ def test_vector_config_hnsw_sq() -> None:
 def test_vector_config_hnsw_rq() -> None:
     vector_index = Configure.VectorIndex.hnsw(
         ef_construction=128,
-        quantizer=Configure.VectorIndex.Quantizer.rq(bits=8),
+        quantizer=Configure.VectorIndex.Quantizer.rq(bits=8, rescore_limit=123),
     )
 
     vi_dict = vector_index._to_dict()
 
     assert vi_dict["efConstruction"] == 128
     assert vi_dict["rq"]["bits"] == 8
+    assert vi_dict["rq"]["rescoreLimit"] == 123
 
 
 def test_vector_config_flat_pq() -> None:

weaviate/collections/classes/config.py

@@ -1501,6 +1501,7 @@ class _SQConfig(_ConfigBase):
 @dataclass
 class _RQConfig(_ConfigBase):
     bits: Optional[int]
+    rescore_limit: int
 
 
 BQConfig = _BQConfig
@@ -2220,6 +2221,7 @@ def sq(
 
     @staticmethod
     def rq(
+        rescore_limit: Optional[int] = None,
         enabled: bool = True,
     ) -> _RQConfigUpdate:
         """Create a `_RQConfigUpdate` object to be used when updating the Rotational quantization (RQ) configuration of Weaviate.
@@ -2229,7 +2231,7 @@ def rq(
         Arguments:
             See [the docs](https://weaviate.io/developers/weaviate/concepts/vector-index#hnsw-with-compression) for a more detailed view!
         """  # noqa: D417 (missing argument descriptions in the docstring)
-        return _RQConfigUpdate(enabled=enabled)
+        return _RQConfigUpdate(enabled=enabled, rescoreLimit=rescore_limit)
 
 
 class _VectorIndexUpdate:

weaviate/collections/classes/config_methods.py

@@ -151,6 +151,7 @@ def __get_quantizer_config(
     elif "rq" in config and config["rq"].get("enabled"):
         quantizer = _RQConfig(
             bits=config["rq"].get("bits"),
+            rescore_limit=config["rq"].get("rescoreLimit"),
         )
     return quantizer
 

weaviate/collections/classes/config_vector_index.py

@@ -273,6 +273,7 @@ def quantizer_name() -> str:
 
 class _RQConfigCreate(_QuantizerConfigCreate):
     bits: Optional[int]
+    rescoreLimit: Optional[int]
 
     @staticmethod
     def quantizer_name() -> str:
@@ -303,6 +304,7 @@ def quantizer_name() -> str:
 
 class _RQConfigUpdate(_QuantizerConfigUpdate):
     enabled: Optional[bool]
+    rescoreLimit: Optional[int]
 
     @staticmethod
     def quantizer_name() -> str:
@@ -432,6 +434,7 @@ def sq(
     @staticmethod
     def rq(
         bits: Optional[int] = None,
+        rescore_limit: Optional[int] = None,
     ) -> _RQConfigCreate:
         """Create a `_RQConfigCreate` object to be used when defining the Rotational quantization (RQ) configuration of Weaviate.
 
@@ -442,6 +445,7 @@ def rq(
         """  # noqa: D417 (missing argument descriptions in the docstring)
         return _RQConfigCreate(
             bits=bits,
+            rescoreLimit=rescore_limit,
         )