Add distributed IVF_RQ builds with a shared rotation model

README.md

@@ -107,18 +107,20 @@ returns the updated dataset; for a distributed rebuild use
 ### Vector Indexing
 
 Vector indexes (`IVF_FLAT`, `IVF_PQ`, `IVF_SQ`, `IVF_HNSW_FLAT`,
-`IVF_HNSW_PQ`, `IVF_HNSW_SQ`) are built distributed in three phases: the
+`IVF_HNSW_PQ`, `IVF_HNSW_SQ`, `IVF_RQ`) are built distributed in three phases: the
 driver trains the global model once — IVF centroids (and the PQ codebook for
 PQ variants) via pylance's `IndicesBuilder` — then Daft workers each build one
 uncommitted index segment for their fragment batch with the shared model, and
 the coordinator commits all segments atomically. Segments built against the
-same centroids and codebook commit as one logical index.
+same shared model commit as one logical index. IVF_RQ also generates one shared
+RaBitQ rotation model on the coordinator and uses it in every segment.
 
 ```python
 from daft_lance import create_vector_index
 
 create_vector_index("s3://bucket/my_dataset", column="vector", index_type="IVF_PQ", num_partitions=16)
 create_vector_index("s3://bucket/my_dataset", column="embedding", index_type="IVF_FLAT", metric="cosine")
+create_vector_index("s3://bucket/my_dataset", column="vector", index_type="IVF_RQ", num_partitions=16)
 ```
 
 `replace` defaults to `False`, matching pylance's `create_index` — which,
@@ -161,6 +163,13 @@ The coordinator always trains the shared IVF centroids and, for PQ types,
 the shared codebook; user-supplied models and `ivf_centroids_file` are not supported.
 PQ builds retain `num_bits` (default 8; 4 and 8 are supported by Lance). The same
 bit width is used for shared codebook training and every segment build.
+For IVF_RQ, `num_bits` defaults to 1 bit per vector dimension. The vector dimension
+must be divisible by 8; Lance validates supported configurations. The shared
+rotation model lets independently built RQ segments use consistent encoding and
+be merged by `optimize_indices`. `rabitq_model` is managed internally and cannot
+be supplied by callers. Omitting `num_bits` preserves the existing 8-bit PQ default.
+RQ is lossy: query recall depends on bit width and data. Use Lance nearest-query
+`refine_factor` to fetch more candidates and rerank them using original vectors.
 
 Training `sample_rate` is a multiplier, not a percentage: for example,
 4 IVF partitions and `sample_rate=32` require 128 rows; 8-bit PQ requires

daft_lance/_lance.py

@@ -698,7 +698,7 @@ def create_vector_index(
     metric: str = "L2",
     num_partitions: int | None = None,
     num_sub_vectors: int | None = None,
-    num_bits: int = 8,
+    num_bits: int | None = None,
     sample_rate: int = 256,
     storage_options: dict[str, Any] | None = None,
     version: int | str | None = None,
@@ -716,7 +716,8 @@ def create_vector_index(
 
     The build runs in three phases: the driver trains the global model once
     (IVF centroids via ``IndicesBuilder.train_ivf``, and a PQ codebook via
-    ``IndicesBuilder.train_pq`` for PQ variants), fragment batches are
+    ``IndicesBuilder.train_pq`` for PQ variants, or a shared RaBitQ rotation
+    model for IVF_RQ), fragment batches are
     distributed across Daft workers where each builds one uncommitted index
     segment with Lance's ``create_index_uncommitted`` using the shared model,
     and the coordinator commits all segments atomically with
@@ -726,7 +727,7 @@ def create_vector_index(
 
     Incremental indexing of appended data belongs to ``optimize_indices``.
     This API does not accept ``fragment_ids``, ``ivf_centroids`` or
-    ``pq_codebook``. Models are always trained by the driver.
+    ``pq_codebook`` or ``rabitq_model``. Models are always generated by the driver.
 
     Args:
         uri: The URI of the Lance table (supports remote URLs to object stores such as `s3://` or `gs://`)
@@ -738,7 +739,9 @@ def create_vector_index(
             {"root": "/data"} for "dir" or {"uri": "http://host:port"} for "rest".
         column: Vector column name to index.
         index_type: Vector index type to build: "IVF_FLAT", "IVF_PQ", "IVF_SQ",
-            "IVF_HNSW_FLAT", "IVF_HNSW_PQ", or "IVF_HNSW_SQ". Independently
+            "IVF_HNSW_FLAT", "IVF_HNSW_PQ", "IVF_HNSW_SQ", or "IVF_RQ".
+            IVF_RQ uses a shared RaBitQ rotation model and requires a vector
+            dimension divisible by 8. Independently
             built SQ segments may have different quantization metadata;
             ``optimize_indices`` cannot merge segments with incompatible metadata.
         name: Name of the index (generated if None).
@@ -759,9 +762,11 @@ def create_vector_index(
             parallelism here comes from ``fragment_group_size``/``max_concurrency``.
         num_sub_vectors: Number of PQ sub-vectors (PQ variants only). If None,
             Lance derives it from the vector dimension during training.
-        num_bits: Bits per PQ code, default 8; Lance supports 4 and 8. Configurable
-            only for ``IVF_PQ`` and ``IVF_HNSW_PQ``. Passed to shared codebook
-            training and every segment build. Lance validates supported configurations.
+        num_bits: Quantization bit width. None selects 8 for PQ
+            (4 and 8 supported), 1 for RQ. For RQ this is the number of bits per
+            vector dimension. Configurable for ``IVF_PQ``, ``IVF_HNSW_PQ`` and
+            ``IVF_RQ``; passed consistently to the shared model and every segment
+            build. Lance validates supported configurations.
         sample_rate: Training sampling multiplier, not a percentage. Lance samples
             up to this many rows per IVF partition or PQ centroid. For example,
             4 IVF partitions and sample_rate=32 require 128 rows; 8-bit PQ
@@ -809,6 +814,11 @@ def create_vector_index(
         ...     "s3://my-bucket/dataset/", column="vector", index_type="IVF_PQ", num_partitions=16
         ... )
 
+        Build a distributed IVF_RQ index with the default 1-bit encoding:
+        >>> daft_lance.create_vector_index(
+        ...     "s3://my-bucket/dataset/", column="vector", index_type="IVF_RQ", num_partitions=16
+        ... )
+
         Build an IVF_FLAT index with cosine distance:
         >>> daft_lance.create_vector_index(
         ...     "s3://my-bucket/dataset/", column="embedding", index_type="IVF_FLAT", metric="cosine"

daft_lance/lance_vector_index.py

@@ -13,6 +13,7 @@
     from daft_lance.namespace import DatasetOpenContext
 
 import lance
+from lance.lance import indices as lance_indices
 
 from daft_lance.lance_scalar_index import (
     _validate_segments_against_manifest,
@@ -22,8 +23,10 @@
 logger = logging.getLogger(__name__)
 
 # Vector index types built with the distributed segment workflow. IVF and
-# IVF_HNSW families over FLAT/PQ/SQ storage.
-VECTOR_INDEX_TYPES = frozenset({"IVF_FLAT", "IVF_PQ", "IVF_SQ", "IVF_HNSW_FLAT", "IVF_HNSW_PQ", "IVF_HNSW_SQ"})
+# IVF_HNSW families over FLAT/PQ/SQ storage, plus IVF_RQ.
+VECTOR_INDEX_TYPES = frozenset(
+    {"IVF_FLAT", "IVF_PQ", "IVF_SQ", "IVF_HNSW_FLAT", "IVF_HNSW_PQ", "IVF_HNSW_SQ", "IVF_RQ"}
+)
 
 # PQ variants need a codebook trained alongside the IVF centroids.
 _PQ_INDEX_TYPES = frozenset({"IVF_PQ", "IVF_HNSW_PQ"})
@@ -44,6 +47,7 @@
         "num_bits",
         "ivf_centroids",
         "pq_codebook",
+        "rabitq_model",
         "storage_options",
         "index_uuid",
         "ivf_centroids_file",
@@ -71,7 +75,7 @@ def _validate_worker_kwargs(kwargs: dict[str, Any]) -> None:
             "create_vector_index no longer accepts fragment_ids; it builds all fragments. "
             "Use optimize_indices to index appended data."
         )
-    removed_models = sorted(set(kwargs) & {"ivf_centroids", "pq_codebook", "ivf_centroids_file"})
+    removed_models = sorted(set(kwargs) & {"ivf_centroids", "pq_codebook", "ivf_centroids_file", "rabitq_model"})
     if removed_models:
         raise TypeError(
             f"create_vector_index does not accept {removed_models}; models are trained by the build workflow."
@@ -109,7 +113,7 @@ class VectorFragmentIndexHandler:
     Each Daft worker receives a subset of Lance fragment IDs and builds one
     uncommitted vector index segment for just those fragments with Lance's
     public ``create_index_uncommitted`` API, using the IVF centroids (and PQ
-    codebook, for PQ variants) trained once on the driver so every segment
+    codebook for PQ variants, or rotation model for RQ) generated once on the driver so every segment
     shares the same model. The worker pickles the returned ``lance.Index``
     metadata back to the coordinator, which commits all segments into the
     dataset manifest atomically with ``commit_existing_index_segments``.
@@ -170,7 +174,7 @@ def create_vector_index_internal(
     metric: str = "L2",
     num_partitions: int | None = None,
     num_sub_vectors: int | None = None,
-    num_bits: int = 8,
+    num_bits: int | None = None,
     sample_rate: int = 256,
     fragment_group_size: int | None = None,
     max_concurrency: int | None = None,
@@ -188,6 +192,7 @@ def create_vector_index_internal(
        ``IndicesBuilder.train_ivf`` and, for PQ variants, the PQ codebook via
        ``IndicesBuilder.train_pq``. The user's ``sample_rate`` is passed
        unchanged; Lance validates that enough training data is available.
+       IVF_RQ builds one shared RaBitQ rotation model for all segments.
     2. Fragment batches are distributed across Daft workers, one partition per
        batch; each worker calls ``create_index_uncommitted`` with the shared
        model for its fragments and pickles the segment metadata back.
@@ -215,8 +220,10 @@ def create_vector_index_internal(
 
     if sample_rate <= 0:
         raise ValueError(f"sample_rate must be positive, got {sample_rate}")
-    if index_type not in _PQ_INDEX_TYPES and num_bits != 8:
-        raise ValueError("num_bits is configurable only for IVF_PQ and IVF_HNSW_PQ.")
+    if num_bits is None:
+        num_bits = 1 if index_type == "IVF_RQ" else 8
+    if index_type not in _PQ_INDEX_TYPES and index_type != "IVF_RQ" and num_bits != 8:
+        raise ValueError("num_bits is configurable only for IVF_PQ, IVF_HNSW_PQ and IVF_RQ.")
 
     # Reject kwargs Lance's segment build would silently swallow (misspelled
     # or unsupported parameters must not build a misconfigured index).
@@ -293,6 +300,9 @@ def create_vector_index_internal(
         "pq_codebook": pq_codebook,
     }
 
+    if index_type == "IVF_RQ":
+        model_kwargs["rabitq_model"] = lance_indices.build_rq_model(dimension=builder.dimension, num_bits=num_bits)
+
     # Phase 2: distribute fragment batches across Daft workers, one partition
     # per batch so distributed runners build the segments in parallel.
     if fragment_group_size > len(fragment_ids_to_use):

tests/io/lance/test_lance_vector_index.py

@@ -13,7 +13,7 @@
 
 warnings.filterwarnings("ignore", category=DeprecationWarning, module="lance")
 
-ALL_VECTOR_INDEX_TYPES = ["IVF_FLAT", "IVF_PQ", "IVF_SQ", "IVF_HNSW_FLAT", "IVF_HNSW_PQ", "IVF_HNSW_SQ"]
+ALL_VECTOR_INDEX_TYPES = ["IVF_FLAT", "IVF_PQ", "IVF_SQ", "IVF_HNSW_FLAT", "IVF_HNSW_PQ", "IVF_HNSW_SQ", "IVF_RQ"]
 
 
 def _make_vector_dataset(
@@ -102,24 +102,28 @@ def test_ivf_pq_shared_model_multi_segment(tmp_path: Path) -> None:
     "IVF_HNSW_SQ": 0.8,
     "IVF_PQ": 0.2,
     "IVF_HNSW_PQ": 0.2,
+    "IVF_RQ": 0.8,
 }
 
 
-def _recall_at_5(ds, vectors: np.ndarray) -> float:
+def _recall_at_5(ds, vectors: np.ndarray, refine_factor: int | None = None) -> float:
     """Recall@5 of the index against brute-force ground truth, over fixed queries."""
     queries = list(range(0, 400, 40)) + [7]
     hits = 0
     for qi in queries:
         truth = set(np.argsort(((vectors - vectors[qi]) ** 2).sum(axis=1))[:5].tolist())
-        got = set(ds.to_table(nearest={"column": "vector", "q": vectors[qi], "k": 5})["id"].to_pylist())
+        nearest = {"column": "vector", "q": vectors[qi], "k": 5}
+        if refine_factor is not None:
+            nearest["refine_factor"] = refine_factor
+        got = set(ds.to_table(nearest=nearest)["id"].to_pylist())
         assert len(got) == 5
         hits += len(truth & got)
     return hits / (5 * len(queries))
 
 
 @pytest.mark.parametrize("index_type", ALL_VECTOR_INDEX_TYPES)
 def test_every_index_type_builds_and_answers_queries(tmp_path: Path, index_type: str) -> None:
-    """All six advertised vector index types build multiple segments and return results."""
+    """All advertised vector index types build multiple segments and return results."""
     uri, vectors = _make_vector_dataset(tmp_path / f"t_{index_type}.lance", seed=13)
 
     create_vector_index(
@@ -133,7 +137,8 @@ def test_every_index_type_builds_and_answers_queries(tmp_path: Path, index_type:
     assert len(desc.segments) == 2
     assert _covered_fragments(uri, "vector_idx") == {f.fragment_id for f in ds.get_fragments()}
 
-    recall = _recall_at_5(ds, vectors)
+    # One-bit RQ needs exact refinement, particularly for this small 8-dimensional fixture.
+    recall = _recall_at_5(ds, vectors, refine_factor=100 if index_type == "IVF_RQ" else None)
     assert recall >= _MIN_RECALL[index_type], f"{index_type} recall@5 {recall:.2f} < {_MIN_RECALL[index_type]}"
 
 
@@ -239,19 +244,21 @@ def test_cosine_metric_orders_by_angle_not_distance(tmp_path: Path) -> None:
     assert got.isdisjoint(l2_truth)
 
 
-def test_replace_default_false_rejects_existing_name(tmp_path: Path) -> None:
+@pytest.mark.parametrize("index_type", ["IVF_FLAT", "IVF_RQ"])
+def test_replace_default_false_rejects_existing_name(tmp_path: Path, index_type: str) -> None:
     uri, _ = _make_vector_dataset(tmp_path / "repl.lance", num_rows=240, rows_per_file=40, seed=5)
-    create_vector_index(uri, column="vector", index_type="IVF_FLAT", num_partitions=4, sample_rate=32, name="v_idx")
+    create_vector_index(uri, column="vector", index_type=index_type, num_partitions=4, sample_rate=32, name="v_idx")
 
     with pytest.raises(ValueError, match="already exists. Set replace=True"):
-        create_vector_index(uri, column="vector", index_type="IVF_FLAT", num_partitions=4, sample_rate=32, name="v_idx")
+        create_vector_index(uri, column="vector", index_type=index_type, num_partitions=4, sample_rate=32, name="v_idx")
 
 
-def test_replace_true_rebuilds_atomically_in_one_version(tmp_path: Path) -> None:
+@pytest.mark.parametrize("index_type", ["IVF_FLAT", "IVF_RQ"])
+def test_replace_true_rebuilds_atomically_in_one_version(tmp_path: Path, index_type: str) -> None:
     uri, vectors = _make_vector_dataset(tmp_path / "rebuild.lance", num_rows=240, rows_per_file=40, seed=6)
     build = dict(
         column="vector",
-        index_type="IVF_FLAT",
+        index_type=index_type,
         num_partitions=4,
         sample_rate=32,
         name="v_idx",
@@ -268,17 +275,20 @@ def test_replace_true_rebuilds_atomically_in_one_version(tmp_path: Path) -> None
     assert len(described) == 1
     assert described[0].name == "v_idx"
     assert len(described[0].segments) == 3
-    results = ds.to_table(nearest={"column": "vector", "q": vectors[7], "k": 5})
+    results = ds.to_table(nearest={"column": "vector", "q": vectors[7], "k": 5, "nprobes": 4, "refine_factor": 20})
     assert 7 in results["id"].to_pylist()
 
 
-def test_worker_failure_keeps_version_and_old_index(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
+@pytest.mark.parametrize("index_type", ["IVF_FLAT", "IVF_RQ"])
+def test_worker_failure_keeps_version_and_old_index(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch, index_type: str
+) -> None:
     """A failing segment build aborts before the commit; the old index stays usable."""
     uri, vectors = _make_vector_dataset(tmp_path / "fail.lance", num_rows=240, rows_per_file=40, seed=9)
     create_vector_index(
         uri,
         column="vector",
-        index_type="IVF_FLAT",
+        index_type=index_type,
         num_partitions=4,
         sample_rate=32,
         name="v_idx",
@@ -298,7 +308,7 @@ def __call__(self, fragment_ids: list[int]) -> bytes:
         create_vector_index(
             uri,
             column="vector",
-            index_type="IVF_FLAT",
+            index_type=index_type,
             num_partitions=4,
             sample_rate=32,
             name="v_idx",
@@ -312,7 +322,7 @@ def __call__(self, fragment_ids: list[int]) -> bytes:
     assert len(described) == 1
     assert described[0].name == "v_idx"
     assert len(described[0].segments) == 3
-    results = ds.to_table(nearest={"column": "vector", "q": vectors[7], "k": 5})
+    results = ds.to_table(nearest={"column": "vector", "q": vectors[7], "k": 5, "nprobes": 4, "refine_factor": 20})
     assert 7 in results["id"].to_pylist()
 
 
@@ -523,7 +533,7 @@ def test_optimize_indices_covers_appended_vectors(tmp_path: Path, index_type: st
     assert np.isfinite(result["_distance"].to_numpy()).all()
 
 
-@pytest.mark.parametrize("argument", ["ivf_centroids", "pq_codebook"])
+@pytest.mark.parametrize("argument", ["ivf_centroids", "pq_codebook", "rabitq_model"])
 @pytest.mark.parametrize("value", [None, pa.FixedSizeListArray.from_arrays(pa.array([0.0] * 8, type=pa.float32()), 8)])
 def test_external_model_arguments_are_rejected(
     tmp_path: Path, argument: str, value, monkeypatch: pytest.MonkeyPatch
@@ -596,3 +606,57 @@ def test_insufficient_training_rows_raise_lance_error(
         )
     assert lance.dataset(uri).version == before.version
     assert not lance.dataset(uri).describe_indices()
+
+
+@pytest.mark.parametrize("num_bits", [None, 2, 4])
+def test_rq_shared_rotation_survives_segment_merge(tmp_path: Path, num_bits: int | None) -> None:
+    """Shared RaBitQ rotations allow independently built segments to merge and remain searchable."""
+    uri, vectors = _make_vector_dataset(tmp_path / "rq.lance", dim=16, rows_per_file=512)
+    create_vector_index(
+        uri,
+        column="vector",
+        index_type="IVF_RQ",
+        num_partitions=4,
+        num_bits=num_bits,
+        sample_rate=8,
+        fragment_group_size=2,
+    )
+    ds = lance.dataset(uri)
+    segments = ds.stats.index_stats("vector_idx")["indices"]
+    assert len(segments) == 2
+    model = segments[0]["sub_index"]
+    assert model["num_bits"] == (1 if num_bits is None else num_bits)
+    assert model["code_dim"] == 16
+    assert model["fast_rotation_signs"]
+    assert all(segment["sub_index"]["fast_rotation_signs"] == model["fast_rotation_signs"] for segment in segments)
+
+    def check_query(dataset):
+        truth = set(np.argsort(((vectors - vectors[7]) ** 2).sum(axis=1))[:5].tolist())
+        result = dataset.to_table(
+            nearest={
+                "column": "vector",
+                "q": vectors[7],
+                "k": 5,
+                "nprobes": 4,
+                "refine_factor": 100,
+            }
+        )
+        assert len(set(result["id"].to_pylist()) & truth) >= 4
+        assert 7 in result["id"].to_pylist()
+
+    check_query(ds)
+    optimize_indices(uri, indices=["vector_idx"], num_indices_to_merge=2)
+    merged = lance.dataset(uri)
+    assert merged.version == ds.version + 1
+    assert len(merged.describe_indices()[0].segments) == 1
+    assert merged.describe_indices()[0].num_rows_indexed == len(vectors)
+    check_query(merged)
+
+
+def test_rq_dimension_error_preserves_dataset(tmp_path: Path) -> None:
+    uri, _ = _make_vector_dataset(tmp_path / "rq_dimension.lance", dim=10)
+    before = lance.dataset(uri).version
+    with pytest.raises(ValueError, match="divisible by 8"):
+        create_vector_index(uri, column="vector", index_type="IVF_RQ", num_partitions=4, sample_rate=8)
+    assert lance.dataset(uri).version == before
+    assert not lance.dataset(uri).describe_indices()