feat(turbo): support product quantization end to end

python/tests/test_ivf_pq.py

@@ -0,0 +1,309 @@
+"""Public HNSW/IVF PQ configuration, persistence and coarse/refine behavior."""
+
+import pickle
+
+import numpy as np
+import pytest
+import zvec
+
+
+def index_param(kind, **kwargs):
+    if kind == "ivf":
+        return zvec.IVFIndexParam(n_list=4, n_iters=3, **kwargs)
+    return zvec.HnswIndexParam(m=16, ef_construction=100, **kwargs)
+
+
+def query_param(kind, **kwargs):
+    if kind == "ivf":
+        return zvec.IVFQueryParam(nprobe=4, **kwargs)
+    return zvec.HnswQueryParam(ef=2000, **kwargs)
+
+
+@pytest.mark.parametrize(
+    "kind,bits,fast_scan",
+    [
+        ("ivf", 8, False),
+        ("ivf", 4, False),
+        ("ivf", 4, True),
+        ("hnsw", 8, False),
+        ("hnsw", 4, False),
+    ],
+)
+@pytest.mark.parametrize("rotate", [False, True])
+def test_pq_parameter_round_trip(kind, bits, fast_scan, rotate):
+    qp = zvec.QuantizerParam(
+        enable_rotate=rotate,
+        num_chunk=4,
+        num_bits=bits,
+        fast_scan=fast_scan,
+        opq_iter=3,
+        opq_pq_iter=2,
+    )
+    param = index_param(
+        kind,
+        metric_type=zvec.MetricType.L2,
+        quantize_type=zvec.QuantizeType.PQ,
+        quantizer_param=qp,
+    )
+    assert pickle.loads(pickle.dumps(qp)) == qp
+    restored = pickle.loads(pickle.dumps(param))
+    assert restored.quantizer_param == qp
+    assert restored.to_dict() == param.to_dict()
+    assert param.to_dict()["quantizer_param"] == qp.to_dict()
+    assert '"fast_scan"' in repr(param)
+
+
+def test_legacy_quantizer_and_ivf_pickle():
+    qp = zvec.QuantizerParam.__new__(zvec.QuantizerParam)
+    qp.__setstate__((True,))
+    assert qp.enable_rotate
+    assert qp.num_chunk == 8 and qp.num_bits == 8
+    assert not qp.fast_scan
+    for state in [
+        (zvec.MetricType.L2, 4, 3, False, zvec.QuantizeType.INT8),
+        (zvec.MetricType.L2, 4, 3, False, zvec.QuantizeType.INT8, True),
+    ]:
+        param = zvec.IVFIndexParam.__new__(zvec.IVFIndexParam)
+        param.__setstate__(state)
+        assert param.quantizer_param.enable_rotate == (len(state) == 6)
+        assert param.quantizer_param.num_bits == 8
+
+
+@pytest.mark.parametrize(
+    "kind,bits,fast_scan",
+    [
+        ("ivf", 8, False),
+        ("ivf", 4, False),
+        ("ivf", 4, True),
+        ("hnsw", 8, False),
+        ("hnsw", 4, False),
+    ],
+)
+@pytest.mark.parametrize("rotate", [False, True])
+@pytest.mark.parametrize(
+    "metric", [zvec.MetricType.L2, zvec.MetricType.IP, zvec.MetricType.COSINE]
+)
+def test_pq_collection_reopen_and_refine(
+    tmp_path, kind, bits, fast_scan, rotate, metric
+):
+    rng = np.random.default_rng(903)
+    vectors = rng.normal(size=(1100, 16)).astype(np.float32)
+    query = rng.normal(size=16).astype(np.float32)
+    qp = zvec.QuantizerParam(
+        enable_rotate=rotate,
+        num_chunk=4,
+        num_bits=bits,
+        fast_scan=fast_scan,
+        opq_iter=2,
+        opq_pq_iter=2,
+    )
+    schema = zvec.CollectionSchema(
+        name="ivf_pq",
+        vectors=[
+            zvec.VectorSchema(
+                "vector",
+                zvec.DataType.VECTOR_FP32,
+                dimension=16,
+                index_param=index_param(
+                    kind,
+                    metric_type=metric,
+                    quantize_type=zvec.QuantizeType.PQ,
+                    quantizer_param=qp,
+                ),
+            )
+        ],
+    )
+    path = str(tmp_path / "index")
+    coll = zvec.create_and_open(path, schema)
+    try:
+        for start in range(0, len(vectors), 200):
+            statuses = coll.insert(
+                [
+                    zvec.Doc(id=str(i), vectors={"vector": vectors[i].tolist()})
+                    for i in range(start, min(start + 200, len(vectors)))
+                ]
+            )
+            assert all(status.ok() for status in statuses)
+        # Untrained writer blocks remain searchable through raw Flat.
+        assert coll.query(
+            zvec.Query(field_name="vector", vector=query.tolist()), topk=5
+        )
+        coll.optimize()
+        before = coll.query(
+            zvec.Query(
+                field_name="vector",
+                vector=query.tolist(),
+                param=query_param(kind),
+            ),
+            topk=40,
+            output_fields=[],
+        )
+    finally:
+        coll.close()
+    coll = zvec.open(path, zvec.CollectionOption(read_only=True, enable_mmap=True))
+    try:
+        restored = coll.schema.vector("vector").index_param
+        assert restored.quantize_type == zvec.QuantizeType.PQ
+        assert restored.quantizer_param == qp
+        coarse = coll.query(
+            zvec.Query(
+                field_name="vector",
+                vector=query.tolist(),
+                param=query_param(kind),
+            ),
+            topk=40,
+            output_fields=[],
+        )
+        assert [doc.id for doc in coarse] == [doc.id for doc in before]
+        np.testing.assert_allclose(
+            [doc.score for doc in coarse], [doc.score for doc in before]
+        )
+        refined = coll.query(
+            zvec.Query(
+                field_name="vector",
+                vector=query.tolist(),
+                param=query_param(kind, is_using_refiner=True, scale_factor=4),
+            ),
+            topk=10,
+            output_fields=[],
+        )
+        ids = np.array([int(doc.id) for doc in coarse])
+        candidates = vectors[ids]
+        if metric == zvec.MetricType.L2:
+            scores = np.sum((candidates - query) ** 2, axis=1)
+            order = np.argsort(scores)[:10]
+        else:
+            scores = candidates @ query
+            if metric == zvec.MetricType.COSINE:
+                scores /= np.linalg.norm(candidates, axis=1) * np.linalg.norm(query)
+            order = np.argsort(-scores)[:10]
+            if metric == zvec.MetricType.COSINE:
+                scores = 1.0 - scores
+        # Verify optimize built compressed postings rather than leaving the raw Flat fallback.
+        assert np.max(np.abs(np.array([doc.score for doc in coarse]) - scores)) > 1e-4
+        assert [int(doc.id) for doc in refined] == ids[order].tolist()
+        np.testing.assert_allclose(
+            [doc.score for doc in refined], scores[order], rtol=1e-4, atol=1e-5
+        )
+    finally:
+        coll.close()
+
+
+def test_legacy_hnsw_pickle():
+    for state in [
+        (zvec.MetricType.L2, 16, 100, zvec.QuantizeType.INT8, False),
+        (
+            zvec.MetricType.L2,
+            16,
+            100,
+            zvec.QuantizeType.INT8,
+            False,
+            True,
+            True,
+            zvec.DataType.VECTOR_FP16,
+        ),
+    ]:
+        param = zvec.HnswIndexParam.__new__(zvec.HnswIndexParam)
+        param.__setstate__(state)
+        assert param.quantizer_param.enable_rotate == (len(state) == 8)
+        assert param.quantizer_param.num_bits == 8
+
+
+@pytest.mark.parametrize("bits", [4, 8])
+@pytest.mark.parametrize("enable_mmap", [False, True])
+def test_hnsw_pq_incremental_writes_and_retraining(tmp_path, bits, enable_mmap):
+    zvec.init()
+    rng = np.random.default_rng(137)
+    vectors = rng.normal(size=(303, 16)).astype(np.float32)
+    qp = zvec.QuantizerParam(num_chunk=4, num_bits=bits)
+    param = zvec.HnswIndexParam(
+        metric_type=zvec.MetricType.L2,
+        m=16,
+        ef_construction=100,
+        quantize_type=zvec.QuantizeType.PQ,
+        quantizer_param=qp,
+        flat_data_type=zvec.DataType.VECTOR_FP16,
+    )
+    schema = zvec.CollectionSchema(
+        name="pq_writes",
+        vectors=[
+            zvec.VectorSchema(
+                "vector", zvec.DataType.VECTOR_FP32, 16, index_param=param
+            )
+        ],
+    )
+    path = str(tmp_path / "index")
+    option = zvec.CollectionOption(enable_mmap=enable_mmap)
+    coll = zvec.create_and_open(path, schema, option=option)
+
+    def check(ids):
+        hits = coll.query(
+            zvec.Query(
+                "vector",
+                vector=vectors[302].tolist(),
+                param=zvec.HnswQueryParam(is_linear=True, is_using_refiner=True),
+            ),
+            topk=400,
+        )
+        assert {hit.id for hit in hits} == ids
+        assert coll.schema.vector("vector").index_param.quantizer_param == qp
+
+    try:
+        for start in range(0, 300, 100):
+            statuses = coll.insert(
+                [
+                    zvec.Doc(id=str(i), vectors={"vector": vectors[i].tolist()})
+                    for i in range(start, start + 100)
+                ]
+            )
+            assert all(status.ok() for status in statuses)
+        coll.optimize()
+        coll.close()
+        coll = zvec.open(path, option)
+        assert coll.insert(
+            zvec.Doc(id="300", vectors={"vector": vectors[300].tolist()})
+        ).ok()
+        assert coll.update(
+            zvec.Doc(id="1", vectors={"vector": vectors[301].tolist()})
+        ).ok()
+        assert coll.delete("0").ok()
+        expected = {str(i) for i in range(1, 301)}
+        check(expected)
+        coll.optimize()
+        check(expected)
+        coll.close()
+        coll = zvec.open(path, option)
+        check(expected)
+    finally:
+        coll.close()
+
+
+@pytest.mark.parametrize(
+    "kwargs",
+    [
+        {"fast_scan": True, "num_bits": 4},
+        {"num_chunk": 0},
+        {"num_chunk": 17},
+        {"num_bits": 6},
+        {"num_chunk": 3, "num_bits": 4},
+        {"enable_rotate": True, "opq_iter": 0},
+    ],
+)
+def test_hnsw_pq_rejects_invalid_configuration(tmp_path, kwargs):
+    schema = zvec.CollectionSchema(
+        name="invalid_pq",
+        vectors=[
+            zvec.VectorSchema(
+                "vector",
+                zvec.DataType.VECTOR_FP32,
+                16,
+                index_param=zvec.HnswIndexParam(
+                    quantize_type=zvec.QuantizeType.PQ,
+                    quantizer_param=zvec.QuantizerParam(**kwargs),
+                ),
+            )
+        ],
+    )
+    with pytest.raises(ValueError, match="PQ"):
+        coll = zvec.create_and_open(str(tmp_path / "index"), schema)
+        coll.close()

python/zvec/model/param/__init__.pyi

@@ -335,7 +335,7 @@ class HnswIndexParam(VectorIndexParam):
             neighbors during index construction. Larger values yield better
             graph quality at the cost of slower build time. Default is 500.
         quantize_type (QuantizeType): Optional quantization type for vector
-            compression (e.g., FP16, INT8). Default is `QuantizeType.UNDEFINED` to
+            compression (FP16, INT8, INT4, or PQ). Default is `QuantizeType.UNDEFINED` to
             disable quantization.
         use_contiguous_memory (bool): If True, the HNSW streamer allocates a
             single contiguous memory arena for all graph nodes, improving cache
@@ -347,6 +347,10 @@ class HnswIndexParam(VectorIndexParam):
             reference index. Defaults to ``VECTOR_FP32``; ``VECTOR_FP16`` and
             ``VECTOR_UINT8`` are also supported.
 
+    PQ learns its codebook during Collection.optimize(). Before optimization,
+    writes remain searchable through Flat. Configure PQ4/PQ8 and optional OPQ
+    with QuantizerParam; fast_scan is supported only by IVF.
+
     Examples:
         >>> from zvec.typing import MetricType, QuantizeType
         >>> params = HnswIndexParam(
@@ -1401,7 +1405,7 @@ class QuantizerParam:
     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.
+            For PQ, enables learned OPQ. Defaults to False.
 
     Examples:
         >>> qp = QuantizerParam(enable_rotate=True)
@@ -1410,13 +1414,26 @@ class QuantizerParam:
     """
 
     def __getstate__(self) -> tuple: ...
-    def __init__(self, enable_rotate: bool = False) -> None:
+    def __init__(
+        self,
+        enable_rotate: bool = False,
+        num_chunk: int = 8,
+        num_bits: int = 8,
+        fast_scan: bool = False,
+        opq_iter: int = 5,
+        opq_pq_iter: int = 4,
+    ) -> None:
         """
         Constructs a QuantizerParam instance.
 
         Args:
+            num_chunk (int): PQ subquantizers; default 8.
+            num_bits (int): PQ bits per code, 4 or 8; default 8.
+            fast_scan (bool): Packed IVF PQ4 scan; default False.
+            opq_iter (int): OPQ iterations; default 5.
+            opq_pq_iter (int): PQ iterations per OPQ iteration; default 4.
             enable_rotate (bool, optional): Whether to apply random rotation
-                before INT8/INT4 quantization. Defaults to False.
+                before INT8/INT4 quantization, or learned OPQ for PQ. Defaults to False.
         """
 
     def __repr__(self) -> str: ...
@@ -1430,9 +1447,20 @@ class QuantizerParam:
     @property
     def enable_rotate(self) -> bool:
         """
-        bool: Whether random rotation is enabled before INT8/INT4 quantization.
+        bool: Enable random rotation for INT8/INT4 or learned OPQ for PQ.
         """
 
+    @property
+    def num_chunk(self) -> int: ...
+    @property
+    def num_bits(self) -> int: ...
+    @property
+    def fast_scan(self) -> bool: ...
+    @property
+    def opq_iter(self) -> int: ...
+    @property
+    def opq_pq_iter(self) -> int: ...
+
 class VectorIndexParam(IndexParam):
     """