diff --git a/README.md b/README.md index ee6d1b5..431a7c1 100644 --- a/README.md +++ b/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 diff --git a/daft_lance/_lance.py b/daft_lance/_lance.py index c43e496..9af80c4 100644 --- a/daft_lance/_lance.py +++ b/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" diff --git a/daft_lance/lance_vector_index.py b/daft_lance/lance_vector_index.py index fa5ec2d..f4565d5 100644 --- a/daft_lance/lance_vector_index.py +++ b/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): diff --git a/tests/io/lance/test_lance_vector_index.py b/tests/io/lance/test_lance_vector_index.py index 46753f9..6da9c09 100644 --- a/tests/io/lance/test_lance_vector_index.py +++ b/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,16 +102,20 @@ 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)) @@ -119,7 +123,7 @@ def _recall_at_5(ds, vectors: np.ndarray) -> float: @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()