From 03f39a9ccb7ea3b3a4a7c832ddb79b3a1499ea2d Mon Sep 17 00:00:00 2001 From: wangzheyan Date: Mon, 28 Sep 2026 18:17:13 +0800 Subject: [PATCH 01/11] feat: add distributed vector index creation Add create_vector_index, building IVF_* / IVF_HNSW_* indexes with the same distributed segment workflow as the scalar path: - the driver trains the global model once (IVF centroids, and the PQ codebook for PQ variants, via pylance's IndicesBuilder) unless pre-trained ivf_centroids/pq_codebook are supplied - Daft workers build one uncommitted segment per fragment batch with the shared model via create_index_uncommitted - the coordinator validates coverage and commits all segments in one commit_existing_index_segments transaction (atomic overlap replacement; replace defaults to False matching pylance's create_index) - fragment_ids partial builds and incremental backfill share the scalar workflow's semantics Multi-segment IVF_SQ builds were verified to match single-segment recall against brute force, so no per-type downgrade is needed. --- README.md | 28 ++ daft_lance/__init__.py | 2 + daft_lance/_lance.py | 162 +++++++++++ daft_lance/lance_vector_index.py | 323 ++++++++++++++++++++++ tests/io/lance/test_lance_vector_index.py | 254 +++++++++++++++++ 5 files changed, 769 insertions(+) create mode 100644 daft_lance/lance_vector_index.py create mode 100644 tests/io/lance/test_lance_vector_index.py diff --git a/README.md b/README.md index 8a4b690..84ec74b 100644 --- a/README.md +++ b/README.md @@ -104,6 +104,34 @@ returns the updated dataset; for a distributed rebuild use `ValueError`. +### 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 +driver trains the global model once — IVF centroids (and the PQ codebook for +PQ variants) via pylance's `IndicesBuilder`, unless pre-trained +`ivf_centroids`/`pq_codebook` are supplied — 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. + +```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") +``` + +`replace` defaults to `False`, matching pylance's `create_index`: an existing +index name is refused unless `replace=True`, which rebuilds atomically (one +manifest transaction retires the old segments and lands the new ones). +`fragment_ids` partial builds and incremental backfill work the same way as +for scalar indexes. Training samples `sample_rate` rows per IVF partition +(and per PQ centroid), so the dataset needs at least +`num_partitions * sample_rate` rows (`256 * sample_rate` for the 8-bit PQ +codebook). + + ### Column Merging ```python diff --git a/daft_lance/__init__.py b/daft_lance/__init__.py index 25154ae..52effec 100644 --- a/daft_lance/__init__.py +++ b/daft_lance/__init__.py @@ -7,6 +7,7 @@ from ._lance import ( compact_files, create_scalar_index, + create_vector_index, merge_columns, merge_columns_df, optimize_indices, @@ -20,6 +21,7 @@ "UpdateColumnsResult", "compact_files", "create_scalar_index", + "create_vector_index", "merge_columns", "merge_columns_df", "optimize_indices", diff --git a/daft_lance/_lance.py b/daft_lance/_lance.py index 6b8eb32..c5b9834 100644 --- a/daft_lance/_lance.py +++ b/daft_lance/_lance.py @@ -19,6 +19,7 @@ from .lance_scalar_index import create_scalar_index_internal, optimize_indices_internal from .lance_scan import LanceScanOperator from .lance_update_column import UpdateColumnsResult, update_columns_from_df, validate_update_arguments +from .lance_vector_index import create_vector_index_internal from .namespace import validate_uri_or_namespace from .utils import construct_lance_dataset_handle @@ -682,6 +683,167 @@ def create_scalar_index( ) +@PublicAPI +def create_vector_index( + uri: str | pathlib.Path | None = None, + io_config: IOConfig | None = None, + *, + table_id: list[str] | None = None, + namespace_impl: str | None = None, + namespace_properties: dict[str, Any] | None = None, + column: str, + index_type: str = "IVF_PQ", + name: str | None = None, + replace: bool = False, + metric: str = "L2", + num_partitions: int | None = None, + num_sub_vectors: int | None = None, + sample_rate: int = 256, + ivf_centroids: Any | None = None, + pq_codebook: Any | None = None, + storage_options: dict[str, Any] | None = None, + version: int | str | None = None, + asof: str | None = None, + block_size: int | None = None, + commit_lock: Any | None = None, + index_cache_size: int | None = None, + default_scan_options: dict[str, Any] | None = None, + metadata_cache_size_bytes: int | None = None, + fragment_group_size: int | None = None, + max_concurrency: int | None = None, + fragment_ids: list[int] | None = None, + **kwargs: Any, +) -> None: + """Build a distributed vector index using Daft's distributed execution. + + 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, unless pre-trained + ``ivf_centroids``/``pq_codebook`` are supplied), 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 + ``commit_existing_index_segments``. Segments built against the same + centroids and codebook commit as one logical index, so no central build + step is needed. + + Args: + uri: The URI of the Lance table (supports remote URLs to object stores such as `s3://` or `gs://`) + io_config: A custom IOConfig to use when accessing LanceDB data. Defaults to None. + table_id: Table identifier within the namespace, e.g. ["catalog", "schema", "table"]. + Mutually exclusive with ``uri``. + namespace_impl: Lance Namespace implementation, e.g. "dir" or "rest". + namespace_properties: Properties for connecting to the namespace, e.g. + {"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". + name: Name of the index (generated if None). + replace: Whether to replace an existing index with the same name. Defaults + to False, matching ``LanceDataset.create_index``; an existing name is + refused with ``ValueError`` unless ``replace=True``. Replacement is + atomic: all new segments commit in one transaction that retires the + overlapped old segments, so readers see either the old or the new index. + metric: Distance metric ("L2", "cosine", or "dot"). + num_partitions: Number of IVF partitions. If None, Lance derives it from + the number of rows during training. + num_sub_vectors: Number of PQ sub-vectors (PQ variants only). If None, + Lance derives it from the vector dimension during training. + sample_rate: Rows sampled per IVF partition (and per PQ centroid) during + training. Training requires at least ``num_partitions * sample_rate`` + rows (256 * sample_rate for the 8-bit PQ codebook). + ivf_centroids: Pre-trained IVF centroids (a pyarrow array); skips IVF + training. Supplying the same centroids to every build is how + independently built segments share one model. + pq_codebook: Pre-trained PQ codebook (PQ variants only); skips PQ training. + storage_options: Storage options for the dataset. + version: Version of the dataset to use. + asof: Timestamp to use for time travel queries. + block_size: Block size for the dataset. + commit_lock: Commit lock for the dataset. + index_cache_size: Size of the index cache. + default_scan_options: Default scan options. + metadata_cache_size_bytes: Size of the metadata cache in bytes. + fragment_group_size: Number of fragments grouped per worker task. If None, + defaults to 10. Each group becomes one independently built index segment. + max_concurrency: Maximum number of concurrent tasks for fragment batches. + If None, Daft uses its default concurrency setting. + fragment_ids: Optional list of fragment IDs to index. Fragments already + covered by committed segments of the same index are skipped and the + remaining ones appended (incremental backfill); pass ``replace=True`` + without ``fragment_ids`` for a full rebuild. + **kwargs: Additional keyword arguments forwarded to Lance's index segment + creation API (e.g. ``target_partition_size``). + + Returns: + None + + Raises: + ValueError: If input parameters are invalid (e.g. empty column name, + non-existent column, unsupported index type, or an existing index + name with ``replace=False``) + RuntimeError: If index building fails (e.g. training sample too small, + commit failures) + + Note: + This function writes Lance datasets via the Lance Python package. + To ensure that this is installed with Daft, you may install: `pip install daft[lance]` + + Examples: + Build a distributed IVF_PQ index: + >>> daft_lance.create_vector_index( + ... "s3://my-bucket/dataset/", column="vector", index_type="IVF_PQ", 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" + ... ) + + Refuse to overwrite an existing index: + >>> daft_lance.create_vector_index("s3://my-bucket/dataset/", column="vector", replace=False) + + Index appended fragments incrementally: + >>> daft_lance.create_vector_index("s3://my-bucket/dataset/", column="vector", fragment_ids=[10, 11]) + """ + io_config = context.get_context().daft_planning_config.default_io_config if io_config is None else io_config + + dataset_handle = construct_lance_dataset_handle( + uri, + storage_options=storage_options, + io_config=io_config, + namespace_impl=namespace_impl, + namespace_properties=namespace_properties, + table_id=table_id, + version=version, + asof=asof, + block_size=block_size, + commit_lock=commit_lock, + index_cache_size=index_cache_size, + default_scan_options=default_scan_options, + metadata_cache_size_bytes=metadata_cache_size_bytes, + ) + + create_vector_index_internal( + lance_ds=dataset_handle.dataset, + open_context=dataset_handle.worker_open_context(), + column=column, + index_type=index_type, + name=name, + replace=replace, + metric=metric, + num_partitions=num_partitions, + num_sub_vectors=num_sub_vectors, + sample_rate=sample_rate, + ivf_centroids=ivf_centroids, + pq_codebook=pq_codebook, + fragment_group_size=fragment_group_size, + max_concurrency=max_concurrency, + fragment_ids=fragment_ids, + **kwargs, + ) + + @PublicAPI def optimize_indices( uri: str | pathlib.Path | None = None, diff --git a/daft_lance/lance_vector_index.py b/daft_lance/lance_vector_index.py new file mode 100644 index 0000000..ecad50f --- /dev/null +++ b/daft_lance/lance_vector_index.py @@ -0,0 +1,323 @@ +from __future__ import annotations + +import logging +import pickle +from typing import TYPE_CHECKING, Any + +import daft +from daft import execution_config_ctx, from_pylist + +if TYPE_CHECKING: + from daft_lance.namespace import DatasetOpenContext + +import lance + +from daft.dependencies import pa +from daft_lance.lance_scalar_index import ( + _existing_index_coverage, + _validate_segments_against_manifest, +) +from daft_lance.utils import distribute_fragments_balanced + +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"}) + +# PQ variants need a codebook trained alongside the IVF centroids. +_PQ_INDEX_TYPES = frozenset({"IVF_PQ", "IVF_HNSW_PQ"}) + + +class VectorFragmentIndexHandler: + """Handler for distributed vector index segment creation on fragment batches. + + 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 + 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``. + """ + + def __init__( + self, + open_context: DatasetOpenContext, + column: str, + index_type: str, + name: str, + replace: bool = False, + **kwargs: Any, + ) -> None: + self.open_context = open_context + self.column = column + self.index_type = index_type + self.name = name + self.replace = replace + self.kwargs = kwargs + self._lance_ds: lance.LanceDataset | None = None + + def _dataset(self) -> lance.LanceDataset: + if self._lance_ds is None: + self._lance_ds = self.open_context.open_pinned() + return self._lance_ds + + def __call__(self, fragment_ids: list[int]) -> bytes: + """Build one uncommitted vector index segment and return pickled metadata.""" + logger.info( + "Building vector index segment for fragments %s (column=%s, type=%s)", + fragment_ids, + self.column, + self.index_type, + ) + # Segment creation uses ``replace=False`` unless the driver opts workers + # into ``replace=True`` for a same-name rebuild: the worker's pinned + # snapshot still contains the index, and Lance rejects building against + # an existing name with ``replace=False``. Replacement itself happens + # in the coordinator's single manifest commit, which retires overlapped + # old segments atomically. + index_meta = self._dataset().create_index_uncommitted( + column=self.column, + index_type=self.index_type, + name=self.name, + replace=self.replace, + train=True, + fragment_ids=fragment_ids, + **self.kwargs, + ) + return pickle.dumps(index_meta) + + +def create_vector_index_internal( + lance_ds: lance.LanceDataset, + open_context: DatasetOpenContext, + *, + column: str, + index_type: str = "IVF_PQ", + name: str | None = None, + replace: bool = False, + metric: str = "L2", + num_partitions: int | None = None, + num_sub_vectors: int | None = None, + sample_rate: int = 256, + ivf_centroids: pa.Array | None = None, + pq_codebook: pa.Array | None = None, + fragment_group_size: int | None = None, + max_concurrency: int | None = None, + fragment_ids: list[int] | None = None, + **kwargs: Any, +) -> None: + """Internal implementation of distributed vector index creation. + + ``lance_ds`` is the driver's live dataset (planning, training, commits); + ``open_context`` is the serializable handle workers reopen from and the + single source of uri, storage options and namespace kwargs. + + The build has three phases: + + 1. The driver trains the global model once — IVF centroids via + ``IndicesBuilder.train_ivf`` and, for PQ variants, the PQ codebook via + ``IndicesBuilder.train_pq`` — unless the caller supplies pre-trained + ``ivf_centroids`` / ``pq_codebook``. One shared model across segments is + what lets independently built segments commit as one logical index. + 2. Fragment batches are distributed across Daft workers; each worker calls + ``create_index_uncommitted`` with the shared model for its fragments and + pickles the segment metadata back. + 3. The coordinator validates the collected segments against the live + manifest (dead/overlapping/missing coverage) and commits them all in one + ``commit_existing_index_segments`` transaction, which retires overlapped + old segments atomically when replacing. + + ``replace`` defaults to ``False``, matching pylance's ``create_index``: + an existing index name is refused unless ``replace=True``. Column type + compatibility is validated by Lance's training and build APIs, not + duplicated here. ``fragment_ids`` restricts the build to a subset of + fragments; already-covered fragments are skipped, so a partial build + followed by a backfill is the incremental path (same semantics as the + scalar index workflow). + """ + if not column: + raise ValueError("Column name cannot be empty") + + index_type = index_type.upper() + if index_type not in VECTOR_INDEX_TYPES: + raise ValueError( + f"Unsupported distributed vector index type '{index_type}'. Supported types: " + f"{sorted(VECTOR_INDEX_TYPES)}. For scalar index types use " + f"daft_lance.create_scalar_index." + ) + + if sample_rate <= 0: + raise ValueError(f"sample_rate must be positive, got {sample_rate}") + + # Validate column exists; whether it is a vector column is Lance's rule, + # enforced by the training and build APIs with clear errors. + try: + lance_ds.schema.field(column) + except KeyError as e: + available_columns = [field.name for field in lance_ds.schema] + raise ValueError(f"Column '{column}' not found. Available: {available_columns}") from e + + # Generate index name if not provided (matches pylance's convention) + if name is None: + name = f"{column}_idx" + + fragments = lance_ds.get_fragments() + available_fragment_ids = {fragment.fragment_id for fragment in fragments} + + # Validate and normalize the requested fragment subset, if any. + requested_fragment_ids: set[int] | None = None + if fragment_ids is not None: + if len(fragment_ids) == 0: + raise ValueError("fragment_ids must be a non-empty list of fragment IDs; pass None to index all fragments.") + unique_ids = list(dict.fromkeys(fragment_ids)) + duplicates = sorted({fid for fid in unique_ids if fragment_ids.count(fid) > 1}) + if duplicates: + logger.warning("Duplicate fragment_ids %s were given; each fragment is scheduled once.", duplicates) + unknown_ids = sorted(fid for fid in unique_ids if fid not in available_fragment_ids) + if unknown_ids: + raise ValueError( + f"fragment_ids {unknown_ids} do not exist in the dataset. " + f"Available fragment IDs: {sorted(available_fragment_ids)}" + ) + requested_fragment_ids = set(unique_ids) + + existing_coverage = _existing_index_coverage(lance_ds, name) + if existing_coverage is not None: + # Column/model compatibility of a same-name index is validated by + # Lance's build and commit APIs, not duplicated here. + if not replace and requested_fragment_ids is None: + raise ValueError(f"Index with name '{name}' already exists. Set replace=True to replace it.") + + # Workers open the pinned snapshot where a same-name index still exists; + # building against that name always requires replace=True. The actual + # replacement happens in the coordinator's atomic commit. + handler_replace = existing_coverage is not None + if existing_coverage is not None and requested_fragment_ids is not None: + # Incremental backfill: skip fragments already covered by committed + # segments; only the remainder is built and appended. + covered = existing_coverage & available_fragment_ids + already_covered = requested_fragment_ids & covered + to_build = requested_fragment_ids - covered + if already_covered: + logger.info( + "Fragments %s are already covered by index '%s'; skipping them.", + sorted(already_covered), + name, + ) + if not to_build: + logger.info("All requested fragments are already covered by index '%s'; nothing to build.", name) + return + requested_fragment_ids = to_build + + if requested_fragment_ids is not None: + fragments = [fragment for fragment in fragments if fragment.fragment_id in requested_fragment_ids] + fragment_ids_to_use = sorted( + requested_fragment_ids if requested_fragment_ids is not None else (f.fragment_id for f in fragments) + ) + if not fragment_ids_to_use: + raise ValueError(f"Dataset at {open_context.uri} contains no fragments") + + # Phase 1: train the global model once on the driver so all segments + # share the same centroids (and codebook) and commit as one logical index. + builder = lance.indices.IndicesBuilder(lance_ds, column) + ivf_model: lance.indices.IvfModel | None = None + if ivf_centroids is None: + logger.info( + "Phase 1: training IVF centroids (index_type=%s, metric=%s, num_partitions=%s, sample_rate=%s)", + index_type, + metric, + num_partitions, + sample_rate, + ) + ivf_model = builder.train_ivf( + num_partitions=num_partitions, + distance_type=metric.lower(), + sample_rate=sample_rate, + ) + ivf_centroids = ivf_model.centroids + num_partitions = ivf_model.num_partitions + logger.info("IVF training completed: num_partitions=%d", num_partitions) + + if index_type in _PQ_INDEX_TYPES and pq_codebook is None: + logger.info("Phase 1: training PQ codebook (num_sub_vectors=%s, sample_rate=%s)", num_sub_vectors, sample_rate) + if ivf_model is None: + # Caller-supplied centroids: wrap them so train_pq can partition + # its samples with the same model the segments will be built with. + ivf_model = lance.indices.IvfModel(ivf_centroids, metric.lower()) + pq_model = builder.train_pq( + ivf_model, + num_subvectors=num_sub_vectors, + sample_rate=sample_rate, + ) + pq_codebook = pq_model.codebook + num_sub_vectors = pq_model.num_subvectors + logger.info("PQ training completed: num_sub_vectors=%d", num_sub_vectors) + + if ivf_centroids is None: + raise ValueError("ivf_centroids must be provided or trainable for IVF-based distributed vector indices") + if index_type in _PQ_INDEX_TYPES and pq_codebook is None: + raise ValueError("pq_codebook must be provided or trainable for PQ-based distributed vector indices") + + model_kwargs: dict[str, Any] = {"metric": metric} + if num_partitions is not None: + model_kwargs["num_partitions"] = num_partitions + if num_sub_vectors is not None: + model_kwargs["num_sub_vectors"] = num_sub_vectors + if ivf_centroids is not None: + model_kwargs["ivf_centroids"] = ivf_centroids + if pq_codebook is not None: + model_kwargs["pq_codebook"] = pq_codebook + + # Phase 2: distribute fragment batches across Daft workers. + if fragment_group_size is None: + fragment_group_size = 10 + elif fragment_group_size <= 0: + raise ValueError("fragment_group_size must be positive") + + if fragment_group_size > len(fragment_ids_to_use) and fragment_ids_to_use: + fragment_group_size = len(fragment_ids_to_use) + logger.info("Adjusted fragment_group_size to %d to match fragment count", fragment_group_size) + + fragment_data = distribute_fragments_balanced(fragments, fragment_group_size) + if not fragment_data: + raise ValueError(f"Dataset at {open_context.uri} contains no fragments") + + logger.info( + "Phase 2: building vector index segments across %d fragment batches " + "(column=%s, type=%s, name=%s, fragments=%d)", + len(fragment_data), + column, + index_type, + name, + len(fragment_ids_to_use), + ) + handler_cls = daft.cls( + VectorFragmentIndexHandler, + max_concurrency=max_concurrency, + ) + handler = handler_cls( + open_context=open_context, + column=column, + index_type=index_type, + name=name, + replace=handler_replace, + **model_kwargs, + **kwargs, + ) + + with execution_config_ctx(maintain_order=False): + df = from_pylist(fragment_data) + df = df.select(handler(df["fragment_ids"]).alias("index_meta")) + collected = df.collect() + + index_metas = [pickle.loads(raw) for raw in collected.to_pydict()["index_meta"]] + + # Phase 3: validate against the live manifest and commit atomically. + lance_ds = open_context.open_latest() + _validate_segments_against_manifest(lance_ds, index_metas, fragment_ids_to_use) + + logger.info("Collected %d vector index segments; committing as segmented index %s", len(index_metas), name) + lance_ds.commit_existing_index_segments(name, column, index_metas) + logger.info("Vector index %s committed successfully", name) diff --git a/tests/io/lance/test_lance_vector_index.py b/tests/io/lance/test_lance_vector_index.py new file mode 100644 index 0000000..1097e49 --- /dev/null +++ b/tests/io/lance/test_lance_vector_index.py @@ -0,0 +1,254 @@ +from __future__ import annotations + +import warnings +from pathlib import Path + +import lance +import numpy as np +import pytest +from lance.indices import IndicesBuilder + +from daft.dependencies import pa +from daft_lance import create_vector_index + +warnings.filterwarnings("ignore", category=DeprecationWarning, module="lance") + + +def _make_vector_dataset( + path: Path, + num_rows: int = 2048, + dim: int = 8, + rows_per_file: int = 256, + seed: int = 42, +) -> tuple[str, np.ndarray]: + """Create a vector dataset and return (uri, vectors) for ground-truth checks.""" + rng = np.random.default_rng(seed) + vectors = rng.standard_normal((num_rows, dim)).astype(np.float32) + table = pa.table( + { + "id": list(range(num_rows)), + "vector": pa.FixedSizeListArray.from_arrays(vectors.reshape(-1), dim), + } + ) + lance.write_dataset(table, str(path), mode="create", max_rows_per_file=rows_per_file) + return str(path), vectors + + +def _covered_fragments(uri: str, index_name: str) -> set[int]: + for desc in lance.dataset(uri).describe_indices(): + if desc.name == index_name: + covered: set[int] = set() + for segment in desc.segments: + covered.update(segment.fragment_ids) + return covered + return set() + + +def test_ivf_flat_multi_segment_build_and_search(tmp_path: Path) -> None: + """IVF_FLAT builds one segment per fragment group; search returns the query point.""" + uri, vectors = _make_vector_dataset(tmp_path / "flat.lance", num_rows=240, rows_per_file=40, seed=1) + + # 6 fragments in groups of 2 -> 3 segments built by separate tasks. + create_vector_index( + uri, column="vector", index_type="IVF_FLAT", num_partitions=4, sample_rate=32, fragment_group_size=2 + ) + + ds = lance.dataset(uri) + described = ds.describe_indices() + assert len(described) == 1 + desc = described[0] + assert desc.name == "vector_idx" + assert desc.index_type == "IVF_FLAT" + assert desc.num_rows_indexed == 240 + assert len(desc.segments) == 3 + + results = ds.to_table(nearest={"column": "vector", "q": vectors[7], "k": 5}) + assert results.num_rows == 5 + assert 7 in results["id"].to_pylist() + + +def test_ivf_pq_shared_model_multi_segment(tmp_path: Path) -> None: + """IVF_PQ trains one shared centroid+codebook pair; segments commit as one index.""" + uri, vectors = _make_vector_dataset(tmp_path / "pq.lance") + + create_vector_index( + uri, + column="vector", + index_type="IVF_PQ", + num_partitions=8, + num_sub_vectors=2, + sample_rate=8, + fragment_group_size=4, + ) + + ds = lance.dataset(uri) + desc = ds.describe_indices()[0] + assert desc.index_type == "IVF_PQ" + assert desc.num_rows_indexed == 2048 + assert len(desc.segments) == 2 + + results = ds.to_table(nearest={"column": "vector", "q": vectors[7], "k": 5}) + assert results.num_rows == 5 + assert 7 in results["id"].to_pylist() + + +def test_ivf_sq_multi_segment_queryable(tmp_path: Path) -> None: + """IVF_SQ also builds multi-segment; recall stays comparable to a single-segment build.""" + uri, vectors = _make_vector_dataset(tmp_path / "sq.lance", seed=7) + + create_vector_index( + uri, column="vector", index_type="IVF_SQ", num_partitions=8, sample_rate=8, fragment_group_size=4 + ) + + ds = lance.dataset(uri) + desc = ds.describe_indices()[0] + assert desc.index_type == "IVF_SQ" + assert len(desc.segments) == 2 + + hits = 0 + queries = range(0, 400, 40) + 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()) + hits += len(truth & got) + # SQ is a lossy compression; a healthy build keeps most of the true top-5. + assert hits >= 0.8 * 5 * len(list(queries)) + + +def test_pretrained_centroids_are_used(tmp_path: Path) -> None: + """Caller-supplied ivf_centroids skip driver-side training and still build.""" + uri, _ = _make_vector_dataset(tmp_path / "pre.lance", num_rows=240, rows_per_file=40, seed=3) + + centroids = ( + IndicesBuilder(lance.dataset(uri), "vector") + .train_ivf(num_partitions=4, distance_type="l2", sample_rate=32) + .centroids + ) + + create_vector_index( + uri, + column="vector", + index_type="IVF_FLAT", + ivf_centroids=centroids, + fragment_group_size=2, + ) + + desc = lance.dataset(uri).describe_indices()[0] + assert desc.index_type == "IVF_FLAT" + assert desc.num_rows_indexed == 240 + + +def test_replace_default_false_rejects_existing_name(tmp_path: Path) -> 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") + + 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") + + +def test_replace_true_rebuilds_atomically_in_one_version(tmp_path: Path) -> 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", + num_partitions=4, + sample_rate=32, + name="v_idx", + fragment_group_size=2, + ) + create_vector_index(uri, **build) + version_before = lance.dataset(uri).version + + create_vector_index(uri, replace=True, **build) + + ds = lance.dataset(uri) + assert ds.version == version_before + 1, "a successful replace must commit exactly one new version" + described = ds.describe_indices() + 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}) + assert 7 in results["id"].to_pylist() + + +def test_fragment_ids_partial_build_then_backfill(tmp_path: Path) -> None: + uri, _ = _make_vector_dataset(tmp_path / "partial.lance", num_rows=240, rows_per_file=40, seed=8) + build = dict(column="vector", index_type="IVF_FLAT", num_partitions=4, sample_rate=32, name="v_idx") + + create_vector_index(uri, fragment_ids=[0, 1, 2], **build) + assert _covered_fragments(uri, "v_idx") == {0, 1, 2} + + # Backfill appends coverage for the remaining fragments. + create_vector_index(uri, fragment_ids=[1, 3, 4, 5], **build) + assert _covered_fragments(uri, "v_idx") == {0, 1, 2, 3, 4, 5} + + # A fully covered request is a no-op: version and coverage stay unchanged. + version_before = lance.dataset(uri).version + create_vector_index(uri, fragment_ids=[2, 3], **build) + assert lance.dataset(uri).version == version_before + assert _covered_fragments(uri, "v_idx") == {0, 1, 2, 3, 4, 5} + + +def test_worker_failure_keeps_version_and_old_index(tmp_path: Path) -> 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", + num_partitions=4, + sample_rate=32, + name="v_idx", + fragment_group_size=2, + ) + ds = lance.dataset(uri) + version_before = ds.version + + original = lance.LanceDataset.create_index_uncommitted + + def fail_one_batch(self, *args, **kwargs): + fragment_ids = kwargs.get("fragment_ids") + if fragment_ids and 4 in fragment_ids: + raise RuntimeError("injected segment build failure") + return original(self, *args, **kwargs) + + lance.LanceDataset.create_index_uncommitted = fail_one_batch + try: + with pytest.raises(Exception, match="injected segment build failure|Error processing"): + create_vector_index( + uri, + column="vector", + index_type="IVF_FLAT", + num_partitions=4, + sample_rate=32, + name="v_idx", + fragment_group_size=2, + replace=True, + ) + finally: + lance.LanceDataset.create_index_uncommitted = original + + ds = lance.dataset(uri) + assert ds.version == version_before + described = ds.describe_indices() + 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}) + assert 7 in results["id"].to_pylist() + + +def test_invalid_arguments_raise(tmp_path: Path) -> None: + uri, _ = _make_vector_dataset(tmp_path / "bad.lance", num_rows=240, rows_per_file=40, seed=11) + + with pytest.raises(ValueError, match="Unsupported distributed vector index type 'BTREE'"): + create_vector_index(uri, column="vector", index_type="BTREE") + + with pytest.raises(ValueError, match="Column 'missing' not found"): + create_vector_index(uri, column="missing", index_type="IVF_FLAT") + + with pytest.raises(ValueError, match="sample_rate must be positive"): + create_vector_index(uri, column="vector", index_type="IVF_FLAT", sample_rate=0) + + with pytest.raises(ValueError, match="fragment_ids"): + create_vector_index(uri, column="vector", index_type="IVF_FLAT", fragment_ids=[]) From 118b5286b51474277bbbf510fda7dfd043dc940d Mon Sep 17 00:00:00 2001 From: wangzheyan Date: Mon, 28 Sep 2026 19:33:46 +0800 Subject: [PATCH 02/11] fix: address review findings on distributed vector index creation High severity: - Backfill (fragment_ids append to an existing index) now requires the original ivf_centroids/pq_codebook: retraining would produce a divergent IVF model that Lance's segment merge refuses (verified: optimize_indices raises 'vector index segments do not share IVF centroids'), so appending without the shared model raises instead. - One Daft partition per fragment batch (into_partitions), so distributed runners build segments in parallel instead of serializing on a single partition. - sample_rate is clamped down (with a warning) to what the dataset supports; the default 256 previously made the README example fail below ~65k rows. Medium severity: - num_sub_vectors is derived from a supplied pq_codebook (and num_partitions from supplied centroids), instead of failing deep in a worker. - Unknown **kwargs are rejected with TypeError: Lance's segment build silently drops them, so 'distance_type=' built an L2 index and any misspelling passed silently. distance_type/train/segmented get targeted messages. Low severity: fragment_group_size validated before training; Raises docstring corrected (ValueError, not RuntimeError, for training size); dead always-true guards removed; replace-under-fragment_ids skip is logged; num_partitions' different meaning vs create_scalar_index documented; README notes the driver-side training memory profile. Tests: six-type build/query matrix, cosine ground truth (differs from L2 truth), clamp test, shared-model backfill + optimize_indices merge regression, kwargs validation, tightened failure-injection regex. Suite: 466 passed. --- README.md | 25 ++- daft_lance/_lance.py | 42 +++-- daft_lance/lance_vector_index.py | 208 ++++++++++++++++++---- tests/io/lance/test_lance_vector_index.py | 195 +++++++++++++++++--- 4 files changed, 392 insertions(+), 78 deletions(-) diff --git a/README.md b/README.md index 84ec74b..9140e7c 100644 --- a/README.md +++ b/README.md @@ -122,13 +122,24 @@ create_vector_index("s3://bucket/my_dataset", column="vector", index_type="IVF_P create_vector_index("s3://bucket/my_dataset", column="embedding", index_type="IVF_FLAT", metric="cosine") ``` -`replace` defaults to `False`, matching pylance's `create_index`: an existing -index name is refused unless `replace=True`, which rebuilds atomically (one -manifest transaction retires the old segments and lands the new ones). -`fragment_ids` partial builds and incremental backfill work the same way as -for scalar indexes. Training samples `sample_rate` rows per IVF partition -(and per PQ centroid), so the dataset needs at least -`num_partitions * sample_rate` rows (`256 * sample_rate` for the 8-bit PQ +`replace` defaults to `False`, matching pylance's `create_index` — which, +unlike its scalar counterpart `create_scalar_index` (default `True` above), +defaults `replace` to `False`. An existing index name is refused unless +`replace=True`, which rebuilds atomically (one manifest transaction retires +the old segments and lands the new ones). + +`fragment_ids` partial builds and incremental backfill work like the scalar +workflow, with one vector-specific requirement: appending segments requires +the same `ivf_centroids` (and `pq_codebook` for PQ variants) the existing +segments were built with, since every segment of a logical vector index must +share one IVF model — a backfill without the original model raises instead of +silently training a divergent one. + +Training samples `sample_rate` rows per IVF partition (and per PQ centroid) +and runs in the coordinator process, so its memory footprint grows with +`num_partitions * sample_rate * dimension`; `sample_rate` is clamped down +automatically (with a warning) to what the dataset size supports +(`num_partitions * sample_rate` rows, `256 * sample_rate` for the 8-bit PQ codebook). diff --git a/daft_lance/_lance.py b/daft_lance/_lance.py index c5b9834..3dbaded 100644 --- a/daft_lance/_lance.py +++ b/daft_lance/_lance.py @@ -746,16 +746,27 @@ def create_vector_index( overlapped old segments, so readers see either the old or the new index. metric: Distance metric ("L2", "cosine", or "dot"). num_partitions: Number of IVF partitions. If None, Lance derives it from - the number of rows during training. + the number of rows during training. Note this is the IVF partition + count of the trained model — unlike ``create_scalar_index``, where + ``num_partitions`` controls Daft's fragment repartitioning; build + 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. + Lance derives it from the vector dimension during training, or from + the supplied ``pq_codebook`` when one is given. sample_rate: Rows sampled per IVF partition (and per PQ centroid) during training. Training requires at least ``num_partitions * sample_rate`` - rows (256 * sample_rate for the 8-bit PQ codebook). + rows (256 * sample_rate for the 8-bit PQ codebook); the value is + automatically clamped down (with a warning) to what the dataset + size supports, so the default works on datasets of any size, at the + cost of a smaller training sample on small datasets. ivf_centroids: Pre-trained IVF centroids (a pyarrow array); skips IVF training. Supplying the same centroids to every build is how - independently built segments share one model. - pq_codebook: Pre-trained PQ codebook (PQ variants only); skips PQ training. + independently built segments share one model. Required when + appending to an existing index (``fragment_ids`` backfill): all + segments of a logical vector index must share one IVF model. + pq_codebook: Pre-trained PQ codebook (PQ variants only); skips PQ + training. Required when appending to an existing PQ index, for the + same shared-model reason as ``ivf_centroids``. storage_options: Storage options for the dataset. version: Version of the dataset to use. asof: Timestamp to use for time travel queries. @@ -771,19 +782,28 @@ def create_vector_index( fragment_ids: Optional list of fragment IDs to index. Fragments already covered by committed segments of the same index are skipped and the remaining ones appended (incremental backfill); pass ``replace=True`` - without ``fragment_ids`` for a full rebuild. + without ``fragment_ids`` for a full rebuild. Appending requires the + same ``ivf_centroids`` (and ``pq_codebook`` for PQ variants) the + existing segments were built with — retraining would produce a + divergent model that Lance cannot merge, so a backfill without the + original model raises ``ValueError``. **kwargs: Additional keyword arguments forwarded to Lance's index segment - creation API (e.g. ``target_partition_size``). + creation API (e.g. ``target_partition_size``). Unknown arguments are + rejected with ``TypeError`` — Lance's segment build silently ignores + them, so misspellings must not pass silently here. Returns: None Raises: ValueError: If input parameters are invalid (e.g. empty column name, - non-existent column, unsupported index type, or an existing index - name with ``replace=False``) - RuntimeError: If index building fails (e.g. training sample too small, - commit failures) + non-existent column, unsupported index type, an existing index + name with ``replace=False``, a backfill without the original + shared model, or a dataset too small to train an IVF model even + at ``sample_rate=1`` — Lance raises its training ``ValueError``) + TypeError: If column type is incompatible with the chosen ``index_type``, + or an unknown keyword argument is passed + RuntimeError: If index building fails (e.g. version compatibility issues, commit failures) Note: This function writes Lance datasets via the Lance Python package. diff --git a/daft_lance/lance_vector_index.py b/daft_lance/lance_vector_index.py index ecad50f..7eda37c 100644 --- a/daft_lance/lance_vector_index.py +++ b/daft_lance/lance_vector_index.py @@ -1,7 +1,10 @@ from __future__ import annotations +import inspect import logging +import math import pickle +from functools import lru_cache from typing import TYPE_CHECKING, Any import daft @@ -28,6 +31,65 @@ # PQ variants need a codebook trained alongside the IVF centroids. _PQ_INDEX_TYPES = frozenset({"IVF_PQ", "IVF_HNSW_PQ"}) +# The 8-bit PQ codebook trains one centroid per 2^8 codes, so its sample +# requirement is 256 rows per sampled codebook entry. +_PQ_CODEBOOK_SIZE = 256 + +# Keyword arguments this workflow sets explicitly on the worker's +# ``create_index_uncommitted`` call; they are not valid user kwargs here. +_HANDLER_MANAGED_KWARGS = frozenset( + { + "column", + "index_type", + "name", + "replace", + "train", + "fragment_ids", + "metric", + "num_partitions", + "num_sub_vectors", + "ivf_centroids", + "pq_codebook", + "storage_options", + } +) + + +@lru_cache(maxsize=1) +def _accepted_worker_kwargs() -> frozenset[str]: + """Keyword arguments Lance's segment-build API actually understands. + + ``create_index_uncommitted`` ends in ``**kwargs`` and silently drops + unknown keys (verified: a misspelled or unsupported parameter builds an + index with the wrong configuration and no warning), so the driver + validates user kwargs against its signature instead of forwarding blind. + """ + params = inspect.signature(lance.LanceDataset.create_index_uncommitted).parameters + return frozenset(params) - _HANDLER_MANAGED_KWARGS + + +def _validate_worker_kwargs(kwargs: dict[str, Any]) -> None: + """Reject kwargs Lance's segment build would silently ignore.""" + if "segmented" in kwargs: + raise TypeError( + "The 'segmented' parameter was removed: the distributed segment-index " + "workflow is the only code path. Remove the argument." + ) + unknown = sorted(set(kwargs) - _accepted_worker_kwargs()) + if not unknown: + return + hints = { + "distance_type": " (create_vector_index uses 'metric'; pylance's training " + "API calls it distance_type, but this API does not accept that spelling)", + "train": " (the distributed workflow always trains segment builds)", + } + detail = "; ".join(f"'{name}'{hints.get(name, '')}" for name in unknown) + raise TypeError( + f"Unknown keyword argument(s) for create_vector_index: {detail}. Lance's " + f"index segment build silently ignores unknown arguments, so they are " + f"rejected here. Accepted: {sorted(_accepted_worker_kwargs())}." + ) + class VectorFragmentIndexHandler: """Handler for distributed vector index segment creation on fragment batches. @@ -121,9 +183,14 @@ def create_vector_index_internal( ``IndicesBuilder.train_pq`` — unless the caller supplies pre-trained ``ivf_centroids`` / ``pq_codebook``. One shared model across segments is what lets independently built segments commit as one logical index. - 2. Fragment batches are distributed across Daft workers; each worker calls - ``create_index_uncommitted`` with the shared model for its fragments and - pickles the segment metadata back. + ``sample_rate`` is clamped down to what the dataset size supports + (training needs ``num_partitions * sample_rate`` rows, and the 8-bit PQ + codebook needs ``256 * sample_rate``); ``num_sub_vectors`` and + ``num_partitions`` are derived from a supplied codebook / centroids when + not given explicitly. + 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. 3. The coordinator validates the collected segments against the live manifest (dead/overlapping/missing coverage) and commits them all in one ``commit_existing_index_segments`` transaction, which retires overlapped @@ -133,9 +200,12 @@ def create_vector_index_internal( an existing index name is refused unless ``replace=True``. Column type compatibility is validated by Lance's training and build APIs, not duplicated here. ``fragment_ids`` restricts the build to a subset of - fragments; already-covered fragments are skipped, so a partial build - followed by a backfill is the incremental path (same semantics as the - scalar index workflow). + fragments; already-covered fragments are skipped and the remainder + appended. Appending requires the same ``ivf_centroids`` (and + ``pq_codebook`` for PQ variants) the existing segments were built with: + every segment of a logical vector index must share one IVF model, or + Lance cannot merge the segments later, so a backfill without the original + model raises instead of silently training a divergent one. """ if not column: raise ValueError("Column name cannot be empty") @@ -151,10 +221,14 @@ def create_vector_index_internal( if sample_rate <= 0: raise ValueError(f"sample_rate must be positive, got {sample_rate}") + # Reject kwargs Lance's segment build would silently swallow (misspelled + # or unsupported parameters must not build a misconfigured index). + _validate_worker_kwargs(kwargs) + # Validate column exists; whether it is a vector column is Lance's rule, # enforced by the training and build APIs with clear errors. try: - lance_ds.schema.field(column) + field = lance_ds.schema.field(column) except KeyError as e: available_columns = [field.name for field in lance_ds.schema] raise ValueError(f"Column '{column}' not found. Available: {available_columns}") from e @@ -202,13 +276,36 @@ def create_vector_index_internal( to_build = requested_fragment_ids - covered if already_covered: logger.info( - "Fragments %s are already covered by index '%s'; skipping them.", + "Fragments %s are already covered by index '%s'; skipping them%s", sorted(already_covered), name, + " (replace does not rebuild covered fragments; use replace=True without " + "fragment_ids for a full rebuild)" + if replace + else "", ) if not to_build: logger.info("All requested fragments are already covered by index '%s'; nothing to build.", name) return + # Appended segments join a live logical index, and every segment of a + # logical vector index must share the same IVF model: Lance's segment + # merge refuses segments trained on different centroids (the error + # ``optimize_indices`` would later raise). Training here would sample + # non-deterministically and silently produce a divergent model, so the + # original model must be supplied. + if ivf_centroids is None: + raise ValueError( + f"Cannot append to index '{name}': all segments of a vector index must share " + "one IVF model, and retraining would produce a divergent one. Pass the same " + "ivf_centroids (and pq_codebook, for PQ variants) the existing segments were " + "built with, or rebuild the whole index with replace=True and no fragment_ids." + ) + if index_type in _PQ_INDEX_TYPES and pq_codebook is None: + raise ValueError( + f"Cannot append to PQ index '{name}' without the original pq_codebook: all " + "segments must share one model. Pass the codebook the existing segments were " + "built with, or rebuild the whole index with replace=True and no fragment_ids." + ) requested_fragment_ids = to_build if requested_fragment_ids is not None: @@ -219,29 +316,60 @@ def create_vector_index_internal( if not fragment_ids_to_use: raise ValueError(f"Dataset at {open_context.uri} contains no fragments") + # Validate grouping before any training work so argument errors surface first. + if fragment_group_size is None: + fragment_group_size = 10 + elif fragment_group_size <= 0: + raise ValueError("fragment_group_size must be positive") + # Phase 1: train the global model once on the driver so all segments # share the same centroids (and codebook) and commit as one logical index. + num_rows = lance_ds.count_rows() + effective_sample_rate = sample_rate + needs_ivf_training = ivf_centroids is None + needs_pq_training = index_type in _PQ_INDEX_TYPES and pq_codebook is None + if needs_ivf_training or needs_pq_training: + # Clamp the sample rate to what this dataset can support, mirroring + # Lance's own training requirements: num_partitions * sample_rate rows + # for IVF (Lance derives num_partitions as sqrt(num_rows) when None) + # and 256 * sample_rate rows for the 8-bit PQ codebook. + effective_partitions = num_partitions if num_partitions is not None else max(1, round(math.sqrt(num_rows))) + caps = [max(1, num_rows // effective_partitions)] + if needs_pq_training: + caps.append(max(1, num_rows // _PQ_CODEBOOK_SIZE)) + clamped = min([sample_rate, *caps]) + if clamped < sample_rate: + logger.warning( + "sample_rate %d exceeds what the dataset supports (%d rows); clamping to %d", + sample_rate, + num_rows, + clamped, + ) + effective_sample_rate = clamped + builder = lance.indices.IndicesBuilder(lance_ds, column) ivf_model: lance.indices.IvfModel | None = None - if ivf_centroids is None: + if needs_ivf_training: logger.info( - "Phase 1: training IVF centroids (index_type=%s, metric=%s, num_partitions=%s, sample_rate=%s)", + "Phase 1: training IVF centroids (index_type=%s, metric=%s, num_partitions=%s, sample_rate=%d)", index_type, metric, num_partitions, - sample_rate, + effective_sample_rate, ) ivf_model = builder.train_ivf( num_partitions=num_partitions, distance_type=metric.lower(), - sample_rate=sample_rate, + sample_rate=effective_sample_rate, ) ivf_centroids = ivf_model.centroids num_partitions = ivf_model.num_partitions logger.info("IVF training completed: num_partitions=%d", num_partitions) - if index_type in _PQ_INDEX_TYPES and pq_codebook is None: - logger.info("Phase 1: training PQ codebook (num_sub_vectors=%s, sample_rate=%s)", num_sub_vectors, sample_rate) + if needs_pq_training: + logger.info( + "Phase 1: training PQ codebook (num_sub_vectors=%s, sample_rate=%d)", num_sub_vectors, effective_sample_rate + ) if ivf_model is None: # Caller-supplied centroids: wrap them so train_pq can partition # its samples with the same model the segments will be built with. @@ -249,34 +377,40 @@ def create_vector_index_internal( pq_model = builder.train_pq( ivf_model, num_subvectors=num_sub_vectors, - sample_rate=sample_rate, + sample_rate=effective_sample_rate, ) pq_codebook = pq_model.codebook num_sub_vectors = pq_model.num_subvectors logger.info("PQ training completed: num_sub_vectors=%d", num_sub_vectors) - if ivf_centroids is None: - raise ValueError("ivf_centroids must be provided or trainable for IVF-based distributed vector indices") - if index_type in _PQ_INDEX_TYPES and pq_codebook is None: - raise ValueError("pq_codebook must be provided or trainable for PQ-based distributed vector indices") - - model_kwargs: dict[str, Any] = {"metric": metric} - if num_partitions is not None: - model_kwargs["num_partitions"] = num_partitions - if num_sub_vectors is not None: - model_kwargs["num_sub_vectors"] = num_sub_vectors - if ivf_centroids is not None: - model_kwargs["ivf_centroids"] = ivf_centroids - if pq_codebook is not None: - model_kwargs["pq_codebook"] = pq_codebook - - # Phase 2: distribute fragment batches across Daft workers. - if fragment_group_size is None: - fragment_group_size = 10 - elif fragment_group_size <= 0: - raise ValueError("fragment_group_size must be positive") + # Derive the model shape from supplied artifacts when not given: the + # partition count is the number of centroids, and the sub-vector count + # follows from the codebook's entry size versus the column dimension. + if num_partitions is None and ivf_centroids is not None: + num_partitions = len(ivf_centroids) + if num_sub_vectors is None and pq_codebook is not None: + dimension = getattr(field.type, "list_size", None) + subvector_size = getattr(pq_codebook.type, "list_size", None) + if dimension and subvector_size and dimension % subvector_size == 0: + num_sub_vectors = dimension // subvector_size + else: + raise ValueError( + f"Cannot derive num_sub_vectors from the supplied pq_codebook " + f"(column dimension {dimension}, codebook entry size {subvector_size}); " + "pass num_sub_vectors explicitly." + ) - if fragment_group_size > len(fragment_ids_to_use) and fragment_ids_to_use: + model_kwargs: dict[str, Any] = { + "metric": metric, + "ivf_centroids": ivf_centroids, + "num_partitions": num_partitions, + "num_sub_vectors": num_sub_vectors, + "pq_codebook": pq_codebook, + } + + # 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): fragment_group_size = len(fragment_ids_to_use) logger.info("Adjusted fragment_group_size to %d to match fragment count", fragment_group_size) @@ -308,7 +442,7 @@ def create_vector_index_internal( ) with execution_config_ctx(maintain_order=False): - df = from_pylist(fragment_data) + df = from_pylist(fragment_data).into_partitions(len(fragment_data)) df = df.select(handler(df["fragment_ids"]).alias("index_meta")) collected = df.collect() diff --git a/tests/io/lance/test_lance_vector_index.py b/tests/io/lance/test_lance_vector_index.py index 1097e49..54774c1 100644 --- a/tests/io/lance/test_lance_vector_index.py +++ b/tests/io/lance/test_lance_vector_index.py @@ -9,10 +9,12 @@ from lance.indices import IndicesBuilder from daft.dependencies import pa -from daft_lance import create_vector_index +from daft_lance import create_vector_index, optimize_indices 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"] + def _make_vector_dataset( path: Path, @@ -92,6 +94,26 @@ def test_ivf_pq_shared_model_multi_segment(tmp_path: Path) -> None: assert 7 in results["id"].to_pylist() +@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 multi-segment and return results.""" + uri, vectors = _make_vector_dataset(tmp_path / f"t_{index_type}.lance", seed=13) + + create_vector_index( + uri, column="vector", index_type=index_type, num_partitions=4, sample_rate=2, fragment_group_size=4 + ) + + ds = lance.dataset(uri) + desc = ds.describe_indices()[0] + assert desc.index_type == index_type + assert desc.num_rows_indexed == 2048 + assert len(desc.segments) == 2 + + results = ds.to_table(nearest={"column": "vector", "q": vectors[7], "k": 5}) + assert results.num_rows == 5 + assert 7 in results["id"].to_pylist() + + def test_ivf_sq_multi_segment_queryable(tmp_path: Path) -> None: """IVF_SQ also builds multi-segment; recall stays comparable to a single-segment build.""" uri, vectors = _make_vector_dataset(tmp_path / "sq.lance", seed=7) @@ -106,13 +128,63 @@ def test_ivf_sq_multi_segment_queryable(tmp_path: Path) -> None: assert len(desc.segments) == 2 hits = 0 - queries = range(0, 400, 40) + queries = list(range(0, 400, 40)) 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()) hits += len(truth & got) # SQ is a lossy compression; a healthy build keeps most of the true top-5. - assert hits >= 0.8 * 5 * len(list(queries)) + assert hits >= 0.8 * 5 * len(queries) + + +def test_cosine_metric_orders_by_angle_not_distance(tmp_path: Path) -> None: + """metric='cosine' must order by cosine similarity, provably not by L2 distance.""" + rng = np.random.default_rng(5) + dim = 8 + u = rng.standard_normal(dim).astype(np.float32) + u /= np.linalg.norm(u) + v = rng.standard_normal(dim).astype(np.float32) + v -= v.dot(u) * u # orthogonal direction to u + v /= np.linalg.norm(v) + + vectors = np.stack( + [u * (9.0 + (i % 10)) if i % 2 == 0 else v * (0.5 + (i % 10) * 0.01) for i in range(400)] + ).astype(np.float32) + table = pa.table( + { + "id": list(range(400)), + "vector": pa.FixedSizeListArray.from_arrays(vectors.reshape(-1), dim), + } + ) + uri = str(tmp_path / "cosine.lance") + lance.write_dataset(table, uri, mode="create", max_rows_per_file=100) + + create_vector_index(uri, column="vector", index_type="IVF_FLAT", metric="cosine", num_partitions=2, sample_rate=32) + + query = u.astype(np.float32) + got = set(lance.dataset(uri).to_table(nearest={"column": "vector", "q": query, "k": 5})["id"].to_pylist()) + # Even ids are u-direction rows (cosine similarity 1.0); odd ids are v-direction. + cosine_truth = {i for i in range(400) if i % 2 == 0} + l2_truth = set(np.argsort(((vectors - query) ** 2).sum(axis=1))[:5].tolist()) + # Sanity: on this dataset the two ground truths genuinely differ. + assert all(i % 2 == 1 for i in l2_truth) + # The index must follow the cosine ground truth, not the L2 one. + assert got <= cosine_truth + assert got.isdisjoint(l2_truth) + + +def test_small_dataset_default_sample_rate_is_clamped(tmp_path: Path) -> None: + """Default sample_rate=256 would demand 65k rows; it must clamp instead of failing.""" + uri, vectors = _make_vector_dataset(tmp_path / "clamp.lance", num_rows=8192, rows_per_file=512, seed=17) + + # No explicit sample_rate: training needs 16*256 (IVF) and 256*256 (PQ) rows + # unclamped — far more than the dataset has. + create_vector_index(uri, column="vector", index_type="IVF_PQ", num_partitions=16, fragment_group_size=2) + + desc = lance.dataset(uri).describe_indices()[0] + assert desc.index_type == "IVF_PQ" + assert desc.num_rows_indexed == 8192 + assert 7 in lance.dataset(uri).to_table(nearest={"column": "vector", "q": vectors[7], "k": 5})["id"].to_pylist() def test_pretrained_centroids_are_used(tmp_path: Path) -> None: @@ -138,6 +210,68 @@ def test_pretrained_centroids_are_used(tmp_path: Path) -> None: assert desc.num_rows_indexed == 240 +def test_pq_codebook_without_num_sub_vectors_derives_it(tmp_path: Path) -> None: + """A supplied codebook implies num_sub_vectors; no deep worker error.""" + uri, _ = _make_vector_dataset(tmp_path / "codebook.lance", seed=19) + ds = lance.dataset(uri) + builder = IndicesBuilder(ds, "vector") + ivf = builder.train_ivf(num_partitions=4, distance_type="l2", sample_rate=32) + pq = builder.train_pq(ivf, num_subvectors=2, sample_rate=8) + + create_vector_index( + uri, + column="vector", + index_type="IVF_PQ", + ivf_centroids=ivf.centroids, + pq_codebook=pq.codebook, + fragment_group_size=4, + ) + + desc = lance.dataset(uri).describe_indices()[0] + assert desc.index_type == "IVF_PQ" + assert desc.num_rows_indexed == 2048 + + +def test_backfill_requires_the_original_shared_model(tmp_path: Path) -> None: + """Appending segments without the original centroids would fork the IVF model.""" + uri, _ = _make_vector_dataset(tmp_path / "guard.lance") + build = dict(column="vector", index_type="IVF_FLAT", num_partitions=4, sample_rate=8, name="v_idx") + + create_vector_index(uri, fragment_ids=[0, 1, 2, 3], **build) + + # Retraining on the backfill would produce a divergent model — refused. + with pytest.raises(ValueError, match="ivf_centroids"): + create_vector_index(uri, fragment_ids=[4, 5], **build) + + +def test_backfill_with_shared_model_then_optimize_indices_merges(tmp_path: Path) -> None: + """A same-model backfill stays mergeable by optimize_indices (shared-centroid invariant).""" + uri, vectors = _make_vector_dataset(tmp_path / "backfill.lance") + centroids = ( + IndicesBuilder(lance.dataset(uri), "vector") + .train_ivf(num_partitions=4, distance_type="l2", sample_rate=8) + .centroids + ) + build = dict(column="vector", index_type="IVF_FLAT", ivf_centroids=centroids, name="v_idx") + + create_vector_index(uri, fragment_ids=[0, 1, 2, 3], **build) + assert _covered_fragments(uri, "v_idx") == {0, 1, 2, 3} + + # Backfill appends coverage for the remaining fragments. + create_vector_index(uri, fragment_ids=[1, 4, 5, 6, 7], **build) + assert _covered_fragments(uri, "v_idx") == {0, 1, 2, 3, 4, 5, 6, 7} + + # Segments built from one model must remain mergeable. + optimize_indices(uri, indices=["v_idx"], num_indices_to_merge=2) + assert _covered_fragments(uri, "v_idx") == {0, 1, 2, 3, 4, 5, 6, 7} + assert 7 in lance.dataset(uri).to_table(nearest={"column": "vector", "q": vectors[7], "k": 5})["id"].to_pylist() + + # A fully covered request is a no-op: version stays unchanged. + version_before = lance.dataset(uri).version + create_vector_index(uri, fragment_ids=[2, 3], **build) + assert lance.dataset(uri).version == version_before + + def test_replace_default_false_rejects_existing_name(tmp_path: Path) -> 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") @@ -171,24 +305,6 @@ def test_replace_true_rebuilds_atomically_in_one_version(tmp_path: Path) -> None assert 7 in results["id"].to_pylist() -def test_fragment_ids_partial_build_then_backfill(tmp_path: Path) -> None: - uri, _ = _make_vector_dataset(tmp_path / "partial.lance", num_rows=240, rows_per_file=40, seed=8) - build = dict(column="vector", index_type="IVF_FLAT", num_partitions=4, sample_rate=32, name="v_idx") - - create_vector_index(uri, fragment_ids=[0, 1, 2], **build) - assert _covered_fragments(uri, "v_idx") == {0, 1, 2} - - # Backfill appends coverage for the remaining fragments. - create_vector_index(uri, fragment_ids=[1, 3, 4, 5], **build) - assert _covered_fragments(uri, "v_idx") == {0, 1, 2, 3, 4, 5} - - # A fully covered request is a no-op: version and coverage stay unchanged. - version_before = lance.dataset(uri).version - create_vector_index(uri, fragment_ids=[2, 3], **build) - assert lance.dataset(uri).version == version_before - assert _covered_fragments(uri, "v_idx") == {0, 1, 2, 3, 4, 5} - - def test_worker_failure_keeps_version_and_old_index(tmp_path: Path) -> 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) @@ -214,7 +330,7 @@ def fail_one_batch(self, *args, **kwargs): lance.LanceDataset.create_index_uncommitted = fail_one_batch try: - with pytest.raises(Exception, match="injected segment build failure|Error processing"): + with pytest.raises(Exception, match="injected segment build failure"): create_vector_index( uri, column="vector", @@ -238,17 +354,50 @@ def fail_one_batch(self, *args, **kwargs): assert 7 in results["id"].to_pylist() +def test_unknown_kwargs_are_rejected(tmp_path: Path) -> None: + """Lance's segment build silently ignores unknown kwargs; this API must not.""" + uri, _ = _make_vector_dataset(tmp_path / "kwargs.lance", num_rows=240, rows_per_file=40, seed=21) + + with pytest.raises(TypeError, match="totally_bogus_param"): + create_vector_index(uri, column="vector", index_type="IVF_FLAT", totally_bogus_param=1) + + # pylance's training API spells it distance_type; here it must be 'metric' + # or the build would silently be L2. + with pytest.raises(TypeError, match="'distance_type'.*'metric'"): + create_vector_index(uri, column="vector", index_type="IVF_FLAT", distance_type="cosine") + + with pytest.raises(TypeError, match="'train'"): + create_vector_index(uri, column="vector", index_type="IVF_FLAT", train=False) + + with pytest.raises(TypeError, match="'segmented' parameter was removed"): + create_vector_index(uri, column="vector", index_type="IVF_FLAT", segmented=True) + + # A real build kwarg still passes through. + create_vector_index(uri, column="vector", index_type="IVF_FLAT", num_partitions=4, target_partition_size=1024) + + def test_invalid_arguments_raise(tmp_path: Path) -> None: uri, _ = _make_vector_dataset(tmp_path / "bad.lance", num_rows=240, rows_per_file=40, seed=11) with pytest.raises(ValueError, match="Unsupported distributed vector index type 'BTREE'"): create_vector_index(uri, column="vector", index_type="BTREE") + with pytest.raises(ValueError, match="Column name cannot be empty"): + create_vector_index(uri, column="", index_type="IVF_FLAT") + with pytest.raises(ValueError, match="Column 'missing' not found"): create_vector_index(uri, column="missing", index_type="IVF_FLAT") with pytest.raises(ValueError, match="sample_rate must be positive"): create_vector_index(uri, column="vector", index_type="IVF_FLAT", sample_rate=0) - with pytest.raises(ValueError, match="fragment_ids"): + with pytest.raises(ValueError, match="fragment_ids must be a non-empty list"): create_vector_index(uri, column="vector", index_type="IVF_FLAT", fragment_ids=[]) + + with pytest.raises(ValueError, match="do not exist in the dataset"): + create_vector_index(uri, column="vector", index_type="IVF_FLAT", fragment_ids=[99]) + + # Argument validation must fire before training work, even when training + # would also fail on this small dataset with a large sample_rate. + with pytest.raises(ValueError, match="fragment_group_size must be positive"): + create_vector_index(uri, column="vector", index_type="IVF_FLAT", fragment_group_size=0) From a33ee63ddeeb59430523f7387bce97adc17db3b6 Mon Sep 17 00:00:00 2001 From: wangzheyan Date: Tue, 29 Sep 2026 15:06:58 +0800 Subject: [PATCH 03/11] fix: CI failures on PR #77 (mypy strict, flaky PQ recall assertion) - mypy --strict (pre-commit hook): parameterize pa.Array[Any] in the internal signature and add an explicit narrowing guard before the IvfModel wrap so pq_codebook-without-centroids is typed, not just logically impossible. - The six-type matrix test asserted exact top-5 membership of the query point on lossy indexes: with sample_rate=2 the PQ codebook trains on the minimum sample and recall is low enough that the assertion flips across platforms (failed on Python 3.11). Lossy families now get recall floors calibrated with headroom over measured values (FLAT 1.00 -> 0.9, SQ 0.95 -> 0.8, PQ 0.42+ -> 0.2), the matrix builds at sample_rate=8, and the dedicated IVF_PQ test moves to a 0.5 recall floor (measured 0.73). Exact-membership assertions remain only on IVF_FLAT, where recall is 1.00. --- daft_lance/lance_vector_index.py | 8 +++-- tests/io/lance/test_lance_vector_index.py | 36 ++++++++++++++++++----- 2 files changed, 35 insertions(+), 9 deletions(-) diff --git a/daft_lance/lance_vector_index.py b/daft_lance/lance_vector_index.py index 7eda37c..c1b3667 100644 --- a/daft_lance/lance_vector_index.py +++ b/daft_lance/lance_vector_index.py @@ -163,8 +163,8 @@ def create_vector_index_internal( num_partitions: int | None = None, num_sub_vectors: int | None = None, sample_rate: int = 256, - ivf_centroids: pa.Array | None = None, - pq_codebook: pa.Array | None = None, + ivf_centroids: pa.Array[Any] | None = None, + pq_codebook: pa.Array[Any] | None = None, fragment_group_size: int | None = None, max_concurrency: int | None = None, fragment_ids: list[int] | None = None, @@ -373,6 +373,10 @@ def create_vector_index_internal( if ivf_model is None: # Caller-supplied centroids: wrap them so train_pq can partition # its samples with the same model the segments will be built with. + # needs_ivf_training is False exactly when ivf_centroids was + # supplied, so the None case cannot occur in practice. + if ivf_centroids is None: + raise ValueError("PQ codebook training requires IVF centroids") ivf_model = lance.indices.IvfModel(ivf_centroids, metric.lower()) pq_model = builder.train_pq( ivf_model, diff --git a/tests/io/lance/test_lance_vector_index.py b/tests/io/lance/test_lance_vector_index.py index 54774c1..8e89604 100644 --- a/tests/io/lance/test_lance_vector_index.py +++ b/tests/io/lance/test_lance_vector_index.py @@ -89,9 +89,32 @@ def test_ivf_pq_shared_model_multi_segment(tmp_path: Path) -> None: assert desc.num_rows_indexed == 2048 assert len(desc.segments) == 2 - results = ds.to_table(nearest={"column": "vector", "q": vectors[7], "k": 5}) - assert results.num_rows == 5 - assert 7 in results["id"].to_pylist() + assert _recall_at_5(ds, vectors) >= 0.5 + + +# Lossy compression keeps less of the true top-5; exact-membership assertions +# on a lossy index flake across platforms, so each family gets a recall floor +# calibrated with headroom over measured values (FLAT 1.00, SQ 0.95, PQ 0.42+). +_MIN_RECALL = { + "IVF_FLAT": 0.9, + "IVF_HNSW_FLAT": 0.9, + "IVF_SQ": 0.8, + "IVF_HNSW_SQ": 0.8, + "IVF_PQ": 0.2, + "IVF_HNSW_PQ": 0.2, +} + + +def _recall_at_5(ds, vectors: np.ndarray) -> 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()) + assert len(got) == 5 + hits += len(truth & got) + return hits / (5 * len(queries)) @pytest.mark.parametrize("index_type", ALL_VECTOR_INDEX_TYPES) @@ -100,7 +123,7 @@ def test_every_index_type_builds_and_answers_queries(tmp_path: Path, index_type: uri, vectors = _make_vector_dataset(tmp_path / f"t_{index_type}.lance", seed=13) create_vector_index( - uri, column="vector", index_type=index_type, num_partitions=4, sample_rate=2, fragment_group_size=4 + uri, column="vector", index_type=index_type, num_partitions=4, sample_rate=8, fragment_group_size=4 ) ds = lance.dataset(uri) @@ -109,9 +132,8 @@ def test_every_index_type_builds_and_answers_queries(tmp_path: Path, index_type: assert desc.num_rows_indexed == 2048 assert len(desc.segments) == 2 - results = ds.to_table(nearest={"column": "vector", "q": vectors[7], "k": 5}) - assert results.num_rows == 5 - assert 7 in results["id"].to_pylist() + recall = _recall_at_5(ds, vectors) + assert recall >= _MIN_RECALL[index_type], f"{index_type} recall@5 {recall:.2f} < {_MIN_RECALL[index_type]}" def test_ivf_sq_multi_segment_queryable(tmp_path: Path) -> None: From c4adeb3aaaf49adfbd197c5c834344fc3f917cfc Mon Sep 17 00:00:00 2001 From: wangzheyan Date: Tue, 29 Sep 2026 15:11:05 +0800 Subject: [PATCH 04/11] fix: clamp test no longer asserts exact top-5 membership on a lossy index The same flake pattern as the matrix test: with a heavily clamped sample rate the IVF_PQ recall is low enough that the query point's top-5 membership flips across platforms (failed on all four Python versions in CI). The test's subject is that default sample_rate clamps instead of erroring, so it now asserts the index answers queries (result row count); recall quality stays guarded by the PQ-specific recall-floor tests. --- tests/io/lance/test_lance_vector_index.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/io/lance/test_lance_vector_index.py b/tests/io/lance/test_lance_vector_index.py index 8e89604..f47ca9c 100644 --- a/tests/io/lance/test_lance_vector_index.py +++ b/tests/io/lance/test_lance_vector_index.py @@ -206,7 +206,11 @@ def test_small_dataset_default_sample_rate_is_clamped(tmp_path: Path) -> None: desc = lance.dataset(uri).describe_indices()[0] assert desc.index_type == "IVF_PQ" assert desc.num_rows_indexed == 8192 - assert 7 in lance.dataset(uri).to_table(nearest={"column": "vector", "q": vectors[7], "k": 5})["id"].to_pylist() + # The clamp makes the build succeed; recall quality under a heavily clamped + # sample is the lossy-compression tradeoff, guarded by the PQ-specific + # recall tests, so only assert the index answers queries here. + results = lance.dataset(uri).to_table(nearest={"column": "vector", "q": vectors[7], "k": 5}) + assert results.num_rows == 5 def test_pretrained_centroids_are_used(tmp_path: Path) -> None: From 2881dddf1ddea152ef9cc9ccbe68630e5bc94e7e Mon Sep 17 00:00:00 2001 From: wangzheyan Date: Thu, 8 Oct 2026 11:41:09 +0800 Subject: [PATCH 05/11] fix: restrict SQ vector index types to single-segment builds Multi-segment IVF_SQ / IVF_HNSW_SQ indexes query correctly (each segment is read with its own quantizer metadata) but can never be merged: optimize_indices raises 'vector index segments do not share quantizer metadata' (reproduced on both SQ variants; IVF_FLAT with shared centroids merges fine). pylance exposes no driver-side SQ quantizer trainer, so the metadata cannot be shared at build time. - IVF_SQ / IVF_HNSW_SQ now build exactly one segment covering all selected fragments regardless of fragment_group_size - fragment_ids appends are rejected for SQ variants with guidance (replace=True rebuild, or optimize_indices for appended fragments) - verified the end state stays maintainable: single-segment SQ + append + optimize_indices(num_indices_to_merge=2) stays at one segment and keeps answering queries Tests: matrix now expects 1 segment for SQ types; new regressions for the single-segment collapse under fragment_group_size, the backfill rejection, and the append+optimize maintenance path. Suite: 468 passed. --- README.md | 8 +++ daft_lance/_lance.py | 10 ++- daft_lance/lance_vector_index.py | 31 +++++++++ tests/io/lance/test_lance_vector_index.py | 79 +++++++++++++++++++---- 4 files changed, 113 insertions(+), 15 deletions(-) diff --git a/README.md b/README.md index 9140e7c..7400c23 100644 --- a/README.md +++ b/README.md @@ -135,6 +135,14 @@ segments were built with, since every segment of a logical vector index must share one IVF model — a backfill without the original model raises instead of silently training a divergent one. +The SQ variants (`IVF_SQ`, `IVF_HNSW_SQ`) build a **single segment** covering +all selected fragments and reject `fragment_ids` appends: pylance exposes no +shared SQ quantizer, so segments built separately query correctly but can +never be merged — `optimize_indices` refuses with "vector index segments do +not share quantizer metadata". Rebuild with `replace=True` to retrain, or let +`optimize_indices` pick up appended fragments (verified to keep a +single-segment SQ index at one segment). + Training samples `sample_rate` rows per IVF partition (and per PQ centroid) and runs in the coordinator process, so its memory footprint grows with `num_partitions * sample_rate * dimension`; `sample_rate` is clamped down diff --git a/daft_lance/_lance.py b/daft_lance/_lance.py index 3dbaded..ce54f4f 100644 --- a/daft_lance/_lance.py +++ b/daft_lance/_lance.py @@ -737,7 +737,10 @@ 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". + "IVF_HNSW_FLAT", "IVF_HNSW_PQ", or "IVF_HNSW_SQ". The SQ variants + always build a single segment covering all selected fragments: + pylance exposes no shared SQ quantizer, so segments built + separately could never be merged by ``optimize_indices``. name: Name of the index (generated if None). replace: Whether to replace an existing index with the same name. Defaults to False, matching ``LanceDataset.create_index``; an existing name is @@ -786,7 +789,10 @@ def create_vector_index( same ``ivf_centroids`` (and ``pq_codebook`` for PQ variants) the existing segments were built with — retraining would produce a divergent model that Lance cannot merge, so a backfill without the - original model raises ``ValueError``. + original model raises ``ValueError``. SQ variants reject appends + entirely (no shareable SQ quantizer exists): rebuild with + ``replace=True`` or let ``optimize_indices`` pick up appended + fragments. **kwargs: Additional keyword arguments forwarded to Lance's index segment creation API (e.g. ``target_partition_size``). Unknown arguments are rejected with ``TypeError`` — Lance's segment build silently ignores diff --git a/daft_lance/lance_vector_index.py b/daft_lance/lance_vector_index.py index c1b3667..427e5ee 100644 --- a/daft_lance/lance_vector_index.py +++ b/daft_lance/lance_vector_index.py @@ -31,6 +31,15 @@ # PQ variants need a codebook trained alongside the IVF centroids. _PQ_INDEX_TYPES = frozenset({"IVF_PQ", "IVF_HNSW_PQ"}) +# SQ variants quantize per-dimension bounds per segment, and pylance exposes +# no driver-side quantizer trainer to share them. Multi-segment SQ indexes +# query correctly (each segment is read with its own metadata) but can never +# be merged — ``optimize_indices`` raises "vector index segments do not share +# quantizer metadata" — so these types are restricted to a single segment +# (verified: a single-segment SQ index stays healthy through append + +# optimize_indices, which keeps it at one segment). +_SINGLE_SEGMENT_VECTOR_INDEX_TYPES = frozenset({"IVF_SQ", "IVF_HNSW_SQ"}) + # The 8-bit PQ codebook trains one centroid per 2^8 codes, so its sample # requirement is 256 rows per sampled codebook entry. _PQ_CODEBOOK_SIZE = 256 @@ -293,6 +302,16 @@ def create_vector_index_internal( # ``optimize_indices`` would later raise). Training here would sample # non-deterministically and silently produce a divergent model, so the # original model must be supplied. + if index_type in _SINGLE_SEGMENT_VECTOR_INDEX_TYPES: + # SQ has no shareable quantizer at all (see the type-set comment), + # so an appended SQ segment could never merge with the existing + # one. Full rebuilds and optimize_indices stay available. + raise ValueError( + f"Cannot append to SQ index '{name}': pylance exposes no shared SQ quantizer, " + "so segments built separately can never be merged (optimize_indices would " + "refuse). Rebuild the whole index with replace=True and no fragment_ids, or " + "let optimize_indices pick up newly appended fragments." + ) if ivf_centroids is None: raise ValueError( f"Cannot append to index '{name}': all segments of a vector index must share " @@ -418,6 +437,18 @@ def create_vector_index_internal( fragment_group_size = len(fragment_ids_to_use) logger.info("Adjusted fragment_group_size to %d to match fragment count", fragment_group_size) + # SQ variants build exactly one segment covering all selected fragments: + # per-segment quantizer metadata cannot be shared or merged, so a + # multi-segment SQ index would be permanently unmergeable. + if index_type in _SINGLE_SEGMENT_VECTOR_INDEX_TYPES and fragment_group_size < len(fragment_ids_to_use): + fragment_group_size = len(fragment_ids_to_use) + logger.info( + "index_type %s builds one segment covering all %d fragments: pylance has no " + "shared SQ quantizer, and multi-segment SQ indexes cannot be merged", + index_type, + len(fragment_ids_to_use), + ) + fragment_data = distribute_fragments_balanced(fragments, fragment_group_size) if not fragment_data: raise ValueError(f"Dataset at {open_context.uri} contains no fragments") diff --git a/tests/io/lance/test_lance_vector_index.py b/tests/io/lance/test_lance_vector_index.py index f47ca9c..7c84f80 100644 --- a/tests/io/lance/test_lance_vector_index.py +++ b/tests/io/lance/test_lance_vector_index.py @@ -104,6 +104,18 @@ def test_ivf_pq_shared_model_multi_segment(tmp_path: Path) -> None: "IVF_HNSW_PQ": 0.2, } +# SQ variants build a single segment covering all fragments (per-segment +# quantizer metadata cannot be merged); the others fan out one segment per +# fragment group (the matrix build uses fragment_group_size=4 over 8 fragments). +_EXPECT_SEGMENTS = { + "IVF_FLAT": 2, + "IVF_HNSW_FLAT": 2, + "IVF_PQ": 2, + "IVF_HNSW_PQ": 2, + "IVF_SQ": 1, + "IVF_HNSW_SQ": 1, +} + def _recall_at_5(ds, vectors: np.ndarray) -> float: """Recall@5 of the index against brute-force ground truth, over fixed queries.""" @@ -119,7 +131,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 multi-segment and return results.""" + """All six advertised vector index types build and return results (SQ single-segment).""" uri, vectors = _make_vector_dataset(tmp_path / f"t_{index_type}.lance", seed=13) create_vector_index( @@ -130,16 +142,23 @@ def test_every_index_type_builds_and_answers_queries(tmp_path: Path, index_type: desc = ds.describe_indices()[0] assert desc.index_type == index_type assert desc.num_rows_indexed == 2048 - assert len(desc.segments) == 2 + assert len(desc.segments) == _EXPECT_SEGMENTS[index_type] recall = _recall_at_5(ds, vectors) assert recall >= _MIN_RECALL[index_type], f"{index_type} recall@5 {recall:.2f} < {_MIN_RECALL[index_type]}" -def test_ivf_sq_multi_segment_queryable(tmp_path: Path) -> None: - """IVF_SQ also builds multi-segment; recall stays comparable to a single-segment build.""" +def test_ivf_sq_is_single_segment_by_design(tmp_path: Path) -> None: + """IVF_SQ builds one segment even when fragment_group_size would split it. + + SQ quantizer metadata cannot be shared across segments, so a multi-segment + SQ index would be permanently unmergeable; recall stays high. + """ uri, vectors = _make_vector_dataset(tmp_path / "sq.lance", seed=7) + # fragment_group_size=4 over 8 fragments would produce 2 segments for + # non-SQ types; SQ must collapse to 1 so optimize_indices can never hit + # "vector index segments do not share quantizer metadata". create_vector_index( uri, column="vector", index_type="IVF_SQ", num_partitions=8, sample_rate=8, fragment_group_size=4 ) @@ -147,16 +166,50 @@ def test_ivf_sq_multi_segment_queryable(tmp_path: Path) -> None: ds = lance.dataset(uri) desc = ds.describe_indices()[0] assert desc.index_type == "IVF_SQ" - assert len(desc.segments) == 2 + assert desc.num_rows_indexed == 2048 + assert len(desc.segments) == 1 - hits = 0 - queries = list(range(0, 400, 40)) - 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()) - hits += len(truth & got) - # SQ is a lossy compression; a healthy build keeps most of the true top-5. - assert hits >= 0.8 * 5 * len(queries) + assert _recall_at_5(ds, vectors) >= 0.8 + + +def test_sq_backfill_is_rejected(tmp_path: Path) -> None: + """Appending SQ segments would create an unmergeable pair; it must raise instead.""" + uri, _ = _make_vector_dataset(tmp_path / "sq_guard.lance") + build = dict(column="vector", index_type="IVF_HNSW_SQ", num_partitions=4, sample_rate=8, name="v_idx") + + create_vector_index(uri, fragment_ids=[0, 1, 2, 3], **build) + + with pytest.raises(ValueError, match="no shared SQ quantizer"): + create_vector_index(uri, fragment_ids=[4, 5], **build) + + +def test_sq_single_segment_stays_maintainable(tmp_path: Path) -> None: + """Regression: single-segment SQ survives append + optimize_indices (merge included). + + This is the healthy end state multi-segment SQ builds could never reach. + """ + uri, vectors = _make_vector_dataset(tmp_path / "sq_life.lance") + + create_vector_index(uri, column="vector", index_type="IVF_SQ", num_partitions=8, sample_rate=8, name="v_idx") + assert len(lance.dataset(uri).describe_indices()[0].segments) == 1 + + more = np.random.default_rng(99).standard_normal((512, 8)).astype(np.float32) + lance.write_dataset( + pa.table( + { + "id": list(range(2048, 2560)), + "vector": pa.FixedSizeListArray.from_arrays(more.reshape(-1), 8), + } + ), + uri, + mode="append", + max_rows_per_file=256, + ) + + optimize_indices(uri, indices=["v_idx"], num_indices_to_merge=2) + ds = lance.dataset(uri) + assert len(ds.describe_indices()[0].segments) == 1 + assert 7 in ds.to_table(nearest={"column": "vector", "q": vectors[7], "k": 5})["id"].to_pylist() def test_cosine_metric_orders_by_angle_not_distance(tmp_path: Path) -> None: From 42348ccd8ed684a7a241674f1e7432e23cf0a094 Mon Sep 17 00:00:00 2001 From: wangzheyan Date: Fri, 9 Oct 2026 11:52:10 +0800 Subject: [PATCH 06/11] fix: allow segmented SQ vector index builds and backfill --- README.md | 16 ++-- daft_lance/_lance.py | 14 ++-- daft_lance/lance_vector_index.py | 31 -------- tests/io/lance/test_lance_vector_index.py | 96 +++++++++++------------ 4 files changed, 61 insertions(+), 96 deletions(-) diff --git a/README.md b/README.md index 7400c23..f29a9be 100644 --- a/README.md +++ b/README.md @@ -135,13 +135,15 @@ segments were built with, since every segment of a logical vector index must share one IVF model — a backfill without the original model raises instead of silently training a divergent one. -The SQ variants (`IVF_SQ`, `IVF_HNSW_SQ`) build a **single segment** covering -all selected fragments and reject `fragment_ids` appends: pylance exposes no -shared SQ quantizer, so segments built separately query correctly but can -never be merged — `optimize_indices` refuses with "vector index segments do -not share quantizer metadata". Rebuild with `replace=True` to retrain, or let -`optimize_indices` pick up appended fragments (verified to keep a -single-segment SQ index at one segment). +The SQ variants (`IVF_SQ`, `IVF_HNSW_SQ`) support multi-segment builds and +`fragment_ids` backfill. Each segment uses its own SQ quantization metadata +for queries. Independently built segments may have different quantization +metadata even when they share IVF centroids; Lance rejects attempts to merge +incompatible segments through `optimize_indices` with "vector index segments +do not share quantizer metadata". Use `fragment_group_size` large enough to +cover all selected fragments when a single segment is preferred. To consolidate +an existing index, rebuild with that grouping, `replace=True`, and no +`fragment_ids`. Training samples `sample_rate` rows per IVF partition (and per PQ centroid) and runs in the coordinator process, so its memory footprint grows with diff --git a/daft_lance/_lance.py b/daft_lance/_lance.py index ce54f4f..fb385aa 100644 --- a/daft_lance/_lance.py +++ b/daft_lance/_lance.py @@ -737,10 +737,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". The SQ variants - always build a single segment covering all selected fragments: - pylance exposes no shared SQ quantizer, so segments built - separately could never be merged by ``optimize_indices``. + "IVF_HNSW_FLAT", "IVF_HNSW_PQ", or "IVF_HNSW_SQ". 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). replace: Whether to replace an existing index with the same name. Defaults to False, matching ``LanceDataset.create_index``; an existing name is @@ -789,10 +788,9 @@ def create_vector_index( same ``ivf_centroids`` (and ``pq_codebook`` for PQ variants) the existing segments were built with — retraining would produce a divergent model that Lance cannot merge, so a backfill without the - original model raises ``ValueError``. SQ variants reject appends - entirely (no shareable SQ quantizer exists): rebuild with - ``replace=True`` or let ``optimize_indices`` pick up appended - fragments. + original model raises ``ValueError``. SQ variants also support + backfill, but sharing IVF centroids does not guarantee compatible + SQ quantization metadata for subsequent segment merging. **kwargs: Additional keyword arguments forwarded to Lance's index segment creation API (e.g. ``target_partition_size``). Unknown arguments are rejected with ``TypeError`` — Lance's segment build silently ignores diff --git a/daft_lance/lance_vector_index.py b/daft_lance/lance_vector_index.py index 427e5ee..c1b3667 100644 --- a/daft_lance/lance_vector_index.py +++ b/daft_lance/lance_vector_index.py @@ -31,15 +31,6 @@ # PQ variants need a codebook trained alongside the IVF centroids. _PQ_INDEX_TYPES = frozenset({"IVF_PQ", "IVF_HNSW_PQ"}) -# SQ variants quantize per-dimension bounds per segment, and pylance exposes -# no driver-side quantizer trainer to share them. Multi-segment SQ indexes -# query correctly (each segment is read with its own metadata) but can never -# be merged — ``optimize_indices`` raises "vector index segments do not share -# quantizer metadata" — so these types are restricted to a single segment -# (verified: a single-segment SQ index stays healthy through append + -# optimize_indices, which keeps it at one segment). -_SINGLE_SEGMENT_VECTOR_INDEX_TYPES = frozenset({"IVF_SQ", "IVF_HNSW_SQ"}) - # The 8-bit PQ codebook trains one centroid per 2^8 codes, so its sample # requirement is 256 rows per sampled codebook entry. _PQ_CODEBOOK_SIZE = 256 @@ -302,16 +293,6 @@ def create_vector_index_internal( # ``optimize_indices`` would later raise). Training here would sample # non-deterministically and silently produce a divergent model, so the # original model must be supplied. - if index_type in _SINGLE_SEGMENT_VECTOR_INDEX_TYPES: - # SQ has no shareable quantizer at all (see the type-set comment), - # so an appended SQ segment could never merge with the existing - # one. Full rebuilds and optimize_indices stay available. - raise ValueError( - f"Cannot append to SQ index '{name}': pylance exposes no shared SQ quantizer, " - "so segments built separately can never be merged (optimize_indices would " - "refuse). Rebuild the whole index with replace=True and no fragment_ids, or " - "let optimize_indices pick up newly appended fragments." - ) if ivf_centroids is None: raise ValueError( f"Cannot append to index '{name}': all segments of a vector index must share " @@ -437,18 +418,6 @@ def create_vector_index_internal( fragment_group_size = len(fragment_ids_to_use) logger.info("Adjusted fragment_group_size to %d to match fragment count", fragment_group_size) - # SQ variants build exactly one segment covering all selected fragments: - # per-segment quantizer metadata cannot be shared or merged, so a - # multi-segment SQ index would be permanently unmergeable. - if index_type in _SINGLE_SEGMENT_VECTOR_INDEX_TYPES and fragment_group_size < len(fragment_ids_to_use): - fragment_group_size = len(fragment_ids_to_use) - logger.info( - "index_type %s builds one segment covering all %d fragments: pylance has no " - "shared SQ quantizer, and multi-segment SQ indexes cannot be merged", - index_type, - len(fragment_ids_to_use), - ) - fragment_data = distribute_fragments_balanced(fragments, fragment_group_size) if not fragment_data: raise ValueError(f"Dataset at {open_context.uri} contains no fragments") diff --git a/tests/io/lance/test_lance_vector_index.py b/tests/io/lance/test_lance_vector_index.py index 7c84f80..89150fb 100644 --- a/tests/io/lance/test_lance_vector_index.py +++ b/tests/io/lance/test_lance_vector_index.py @@ -104,18 +104,6 @@ def test_ivf_pq_shared_model_multi_segment(tmp_path: Path) -> None: "IVF_HNSW_PQ": 0.2, } -# SQ variants build a single segment covering all fragments (per-segment -# quantizer metadata cannot be merged); the others fan out one segment per -# fragment group (the matrix build uses fragment_group_size=4 over 8 fragments). -_EXPECT_SEGMENTS = { - "IVF_FLAT": 2, - "IVF_HNSW_FLAT": 2, - "IVF_PQ": 2, - "IVF_HNSW_PQ": 2, - "IVF_SQ": 1, - "IVF_HNSW_SQ": 1, -} - def _recall_at_5(ds, vectors: np.ndarray) -> float: """Recall@5 of the index against brute-force ground truth, over fixed queries.""" @@ -131,7 +119,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 and return results (SQ single-segment).""" + """All six 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( @@ -142,55 +130,63 @@ def test_every_index_type_builds_and_answers_queries(tmp_path: Path, index_type: desc = ds.describe_indices()[0] assert desc.index_type == index_type assert desc.num_rows_indexed == 2048 - assert len(desc.segments) == _EXPECT_SEGMENTS[index_type] + assert len(desc.segments) == 2 recall = _recall_at_5(ds, vectors) assert recall >= _MIN_RECALL[index_type], f"{index_type} recall@5 {recall:.2f} < {_MIN_RECALL[index_type]}" -def test_ivf_sq_is_single_segment_by_design(tmp_path: Path) -> None: - """IVF_SQ builds one segment even when fragment_group_size would split it. - - SQ quantizer metadata cannot be shared across segments, so a multi-segment - SQ index would be permanently unmergeable; recall stays high. - """ - uri, vectors = _make_vector_dataset(tmp_path / "sq.lance", seed=7) - - # fragment_group_size=4 over 8 fragments would produce 2 segments for - # non-SQ types; SQ must collapse to 1 so optimize_indices can never hit - # "vector index segments do not share quantizer metadata". - create_vector_index( - uri, column="vector", index_type="IVF_SQ", num_partitions=8, sample_rate=8, fragment_group_size=4 +@pytest.mark.parametrize("index_type", ["IVF_SQ", "IVF_HNSW_SQ"]) +def test_sq_backfill_and_failed_merge_preserve_index(tmp_path: Path, index_type: str) -> None: + """SQ backfill is searchable; incompatible merges leave the committed index intact.""" + rng = np.random.default_rng(7) + vectors = rng.uniform(0, 1, (2048, 8)).astype(np.float32) + vectors[1024:] += 10 + uri = str(tmp_path / "sq_backfill.lance") + lance.write_dataset( + pa.table({"id": list(range(2048)), "vector": pa.FixedSizeListArray.from_arrays(vectors.reshape(-1), 8)}), + uri, + max_rows_per_file=256, ) - - ds = lance.dataset(uri) - desc = ds.describe_indices()[0] - assert desc.index_type == "IVF_SQ" - assert desc.num_rows_indexed == 2048 - assert len(desc.segments) == 1 - - assert _recall_at_5(ds, vectors) >= 0.8 - - -def test_sq_backfill_is_rejected(tmp_path: Path) -> None: - """Appending SQ segments would create an unmergeable pair; it must raise instead.""" - uri, _ = _make_vector_dataset(tmp_path / "sq_guard.lance") - build = dict(column="vector", index_type="IVF_HNSW_SQ", num_partitions=4, sample_rate=8, name="v_idx") - + centroids = ( + IndicesBuilder(lance.dataset(uri), "vector") + .train_ivf(num_partitions=4, distance_type="l2", sample_rate=8) + .centroids + ) + build = dict(column="vector", index_type=index_type, ivf_centroids=centroids, name="v_idx") create_vector_index(uri, fragment_ids=[0, 1, 2, 3], **build) - - with pytest.raises(ValueError, match="no shared SQ quantizer"): - create_vector_index(uri, fragment_ids=[4, 5], **build) + create_vector_index(uri, fragment_ids=[4, 5, 6, 7], **build) + + before = lance.dataset(uri) + assert len(before.describe_indices()[0].segments) == 2 + assert _covered_fragments(uri, "v_idx") == set(range(8)) + segment_ids_before = [segment.uuid for segment in before.describe_indices()[0].segments] + results_before = [] + for qi in [7, 1300]: + result = before.to_table(nearest={"column": "vector", "q": vectors[qi], "k": 5, "nprobes": 4}) + assert qi in result["id"].to_pylist() + results_before.append(result.to_pydict()) + + with pytest.raises(OSError, match="do not share quantizer metadata"): + optimize_indices(uri, indices=["v_idx"], num_indices_to_merge=2) + + after = lance.dataset(uri) + assert after.version == before.version + assert [segment.uuid for segment in after.describe_indices()[0].segments] == segment_ids_before + assert _covered_fragments(uri, "v_idx") == set(range(8)) + for qi, expected in zip([7, 1300], results_before): + assert ( + after.to_table(nearest={"column": "vector", "q": vectors[qi], "k": 5, "nprobes": 4}).to_pydict() == expected + ) def test_sq_single_segment_stays_maintainable(tmp_path: Path) -> None: - """Regression: single-segment SQ survives append + optimize_indices (merge included). - - This is the healthy end state multi-segment SQ builds could never reach. - """ + """An explicitly requested single-segment SQ index supports incremental maintenance.""" uri, vectors = _make_vector_dataset(tmp_path / "sq_life.lance") - create_vector_index(uri, column="vector", index_type="IVF_SQ", num_partitions=8, sample_rate=8, name="v_idx") + create_vector_index( + uri, column="vector", index_type="IVF_SQ", num_partitions=8, sample_rate=8, name="v_idx", fragment_group_size=8 + ) assert len(lance.dataset(uri).describe_indices()[0].segments) == 1 more = np.random.default_rng(99).standard_normal((512, 8)).astype(np.float32) From 4c7c119dbb98c9b32adf9067b86382ede4668f25 Mon Sep 17 00:00:00 2001 From: wangzheyan Date: Fri, 9 Oct 2026 12:00:20 +0800 Subject: [PATCH 07/11] fix: isolate Lance handles across concurrent vector builds --- daft_lance/lance_vector_index.py | 10 ++--- tests/io/lance/test_lance_vector_index.py | 45 ++++++++++------------- 2 files changed, 23 insertions(+), 32 deletions(-) diff --git a/daft_lance/lance_vector_index.py b/daft_lance/lance_vector_index.py index c1b3667..abe81bd 100644 --- a/daft_lance/lance_vector_index.py +++ b/daft_lance/lance_vector_index.py @@ -118,12 +118,6 @@ def __init__( self.name = name self.replace = replace self.kwargs = kwargs - self._lance_ds: lance.LanceDataset | None = None - - def _dataset(self) -> lance.LanceDataset: - if self._lance_ds is None: - self._lance_ds = self.open_context.open_pinned() - return self._lance_ds def __call__(self, fragment_ids: list[int]) -> bytes: """Build one uncommitted vector index segment and return pickled metadata.""" @@ -139,7 +133,9 @@ def __call__(self, fragment_ids: list[int]) -> bytes: # an existing name with ``replace=False``. Replacement itself happens # in the coordinator's single manifest commit, which retires overlapped # old segments atomically. - index_meta = self._dataset().create_index_uncommitted( + # Vector builds mutably borrow the Lance handle; concurrent invocations + # must open independent handles to the same pinned snapshot. + index_meta = self.open_context.open_pinned().create_index_uncommitted( column=self.column, index_type=self.index_type, name=self.name, diff --git a/tests/io/lance/test_lance_vector_index.py b/tests/io/lance/test_lance_vector_index.py index 89150fb..5ab167f 100644 --- a/tests/io/lance/test_lance_vector_index.py +++ b/tests/io/lance/test_lance_vector_index.py @@ -9,7 +9,7 @@ from lance.indices import IndicesBuilder from daft.dependencies import pa -from daft_lance import create_vector_index, optimize_indices +from daft_lance import create_vector_index, lance_vector_index, optimize_indices warnings.filterwarnings("ignore", category=DeprecationWarning, module="lance") @@ -380,7 +380,7 @@ def test_replace_true_rebuilds_atomically_in_one_version(tmp_path: Path) -> None assert 7 in results["id"].to_pylist() -def test_worker_failure_keeps_version_and_old_index(tmp_path: Path) -> None: +def test_worker_failure_keeps_version_and_old_index(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> 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( @@ -395,29 +395,24 @@ def test_worker_failure_keeps_version_and_old_index(tmp_path: Path) -> None: ds = lance.dataset(uri) version_before = ds.version - original = lance.LanceDataset.create_index_uncommitted - - def fail_one_batch(self, *args, **kwargs): - fragment_ids = kwargs.get("fragment_ids") - if fragment_ids and 4 in fragment_ids: - raise RuntimeError("injected segment build failure") - return original(self, *args, **kwargs) - - lance.LanceDataset.create_index_uncommitted = fail_one_batch - try: - with pytest.raises(Exception, match="injected segment build failure"): - create_vector_index( - uri, - column="vector", - index_type="IVF_FLAT", - num_partitions=4, - sample_rate=32, - name="v_idx", - fragment_group_size=2, - replace=True, - ) - finally: - lance.LanceDataset.create_index_uncommitted = original + class FailingVectorFragmentIndexHandler(lance_vector_index.VectorFragmentIndexHandler): + def __call__(self, fragment_ids: list[int]) -> bytes: + if 4 in fragment_ids: + raise RuntimeError("injected segment build failure") + return super().__call__(fragment_ids) + + monkeypatch.setattr(lance_vector_index, "VectorFragmentIndexHandler", FailingVectorFragmentIndexHandler) + with pytest.raises(Exception, match="injected segment build failure"): + create_vector_index( + uri, + column="vector", + index_type="IVF_FLAT", + num_partitions=4, + sample_rate=32, + name="v_idx", + fragment_group_size=2, + replace=True, + ) ds = lance.dataset(uri) assert ds.version == version_before From 3d9efbc4cb906ecc1704f4723d1f0d41296dfc29 Mon Sep 17 00:00:00 2001 From: wangzheyan Date: Fri, 9 Oct 2026 17:51:57 +0800 Subject: [PATCH 08/11] fix: preserve vector index models and concurrent commits --- README.md | 9 ++ daft_lance/_lance.py | 13 +- daft_lance/lance_vector_index.py | 40 +++++- tests/io/lance/test_lance_vector_index.py | 108 ++++++++++++++ .../test_lance_vector_index_concurrency.py | 133 ++++++++++++++++++ 5 files changed, 294 insertions(+), 9 deletions(-) create mode 100644 tests/io/lance/test_lance_vector_index_concurrency.py diff --git a/README.md b/README.md index f29a9be..7808376 100644 --- a/README.md +++ b/README.md @@ -134,6 +134,15 @@ the same `ivf_centroids` (and `pq_codebook` for PQ variants) the existing segments were built with, since every segment of a logical vector index must share one IVF model — a backfill without the original model raises instead of silently training a divergent one. +Backfill must also use the existing index's `metric`; a different metric is +rejected because distances from different metrics cannot be ranked together. +If another task changes the same index during a new build or backfill, the +operation fails rather than replacing that task's result; retry against the +latest version. Full rebuilds with `replace=True` retain replacement semantics. + +HNSW builds accept `m`, `max_level`, and `ef_construction`. Each worker generates +its own segment UUID, so `index_uuid` cannot be supplied. For pre-trained IVF +models, pass `ivf_centroids`; `ivf_centroids_file` is not supported by this API. The SQ variants (`IVF_SQ`, `IVF_HNSW_SQ`) support multi-segment builds and `fragment_ids` backfill. Each segment uses its own SQ quantization metadata diff --git a/daft_lance/_lance.py b/daft_lance/_lance.py index fb385aa..54e3010 100644 --- a/daft_lance/_lance.py +++ b/daft_lance/_lance.py @@ -746,7 +746,10 @@ def create_vector_index( refused with ``ValueError`` unless ``replace=True``. Replacement is atomic: all new segments commit in one transaction that retires the overlapped old segments, so readers see either the old or the new index. - metric: Distance metric ("L2", "cosine", or "dot"). + New builds and backfill fail if another task changes the same index + during the build; retry the operation against the latest version. + metric: Distance metric ("L2", "cosine", or "dot"). Backfill must use + the same metric as the existing index; mismatches raise ``ValueError``. num_partitions: Number of IVF partitions. If None, Lance derives it from the number of rows during training. Note this is the IVF partition count of the trained model — unlike ``create_scalar_index``, where @@ -792,9 +795,11 @@ def create_vector_index( backfill, but sharing IVF centroids does not guarantee compatible SQ quantization metadata for subsequent segment merging. **kwargs: Additional keyword arguments forwarded to Lance's index segment - creation API (e.g. ``target_partition_size``). Unknown arguments are - rejected with ``TypeError`` — Lance's segment build silently ignores - them, so misspellings must not pass silently here. + creation API, including HNSW ``m``, ``max_level`` and ``ef_construction``. + ``index_uuid`` is managed per segment. ``ivf_centroids_file`` is not + supported; supply shared centroids through ``ivf_centroids``. + Unknown arguments raise ``TypeError`` — Lance's segment build silently + ignores them, so misspellings must not pass silently here. Returns: None diff --git a/daft_lance/lance_vector_index.py b/daft_lance/lance_vector_index.py index abe81bd..5999ec2 100644 --- a/daft_lance/lance_vector_index.py +++ b/daft_lance/lance_vector_index.py @@ -35,8 +35,8 @@ # requirement is 256 rows per sampled codebook entry. _PQ_CODEBOOK_SIZE = 256 -# Keyword arguments this workflow sets explicitly on the worker's -# ``create_index_uncommitted`` call; they are not valid user kwargs here. +# Arguments managed by this workflow or incompatible with its shared-model +# builds; they are not valid user kwargs on the worker's segment-build call. _HANDLER_MANAGED_KWARGS = frozenset( { "column", @@ -51,6 +51,8 @@ "ivf_centroids", "pq_codebook", "storage_options", + "index_uuid", + "ivf_centroids_file", } ) @@ -62,10 +64,10 @@ def _accepted_worker_kwargs() -> frozenset[str]: ``create_index_uncommitted`` ends in ``**kwargs`` and silently drops unknown keys (verified: a misspelled or unsupported parameter builds an index with the wrong configuration and no warning), so the driver - validates user kwargs against its signature instead of forwarding blind. + validates explicit parameters and the HNSW options parsed from ``**kwargs``. """ params = inspect.signature(lance.LanceDataset.create_index_uncommitted).parameters - return frozenset(params) - _HANDLER_MANAGED_KWARGS + return (frozenset(params) | {"m", "max_level", "ef_construction"}) - _HANDLER_MANAGED_KWARGS def _validate_worker_kwargs(kwargs: dict[str, Any]) -> None: @@ -75,6 +77,10 @@ def _validate_worker_kwargs(kwargs: dict[str, Any]) -> None: "The 'segmented' parameter was removed: the distributed segment-index " "workflow is the only code path. Remove the argument." ) + if "index_uuid" in kwargs: + raise TypeError("index_uuid is managed per segment; each worker must generate a unique UUID.") + if "ivf_centroids_file" in kwargs: + raise TypeError("ivf_centroids_file is not supported; pass shared centroids through ivf_centroids.") unknown = sorted(set(kwargs) - _accepted_worker_kwargs()) if not unknown: return @@ -91,6 +97,10 @@ def _validate_worker_kwargs(kwargs: dict[str, Any]) -> None: ) +def _index_segment_ids(dataset: lance.LanceDataset, name: str) -> set[str]: + return {segment.uuid for index in dataset.describe_indices() if index.name == name for segment in index.segments} + + class VectorFragmentIndexHandler: """Handler for distributed vector index segment creation on fragment batches. @@ -254,8 +264,9 @@ def create_vector_index_internal( requested_fragment_ids = set(unique_ids) existing_coverage = _existing_index_coverage(lance_ds, name) + planned_segment_ids = _index_segment_ids(lance_ds, name) if existing_coverage is not None: - # Column/model compatibility of a same-name index is validated by + # Column/index-type compatibility of a same-name index is validated by # Lance's build and commit APIs, not duplicated here. if not replace and requested_fragment_ids is None: raise ValueError(f"Index with name '{name}' already exists. Set replace=True to replace it.") @@ -283,6 +294,19 @@ def create_vector_index_internal( if not to_build: logger.info("All requested fragments are already covered by index '%s'; nothing to build.", name) return + # Lance permits mixed metrics at commit, but their distances cannot be + # compared when searching the logical index. + stats = lance_ds.stats.index_stats(name) + if stats["index_type"].startswith("IVF"): + requested_metric = metric.lower() + if requested_metric == "euclidean": + requested_metric = "l2" + existing_metrics = {segment["metric_type"] for segment in stats["indices"]} + if existing_metrics != {requested_metric}: + raise ValueError( + f"Cannot append to index '{name}' with metric '{metric}': " + f"existing segments use {sorted(existing_metrics)}. Use the same metric." + ) # Appended segments join a live logical index, and every segment of a # logical vector index must share the same IVF model: Lance's segment # merge refuses segments trained on different centroids (the error @@ -450,8 +474,14 @@ def create_vector_index_internal( # Phase 3: validate against the live manifest and commit atomically. lance_ds = open_context.open_latest() + if (not replace or requested_fragment_ids is not None) and _index_segment_ids( + lance_ds, name + ) != planned_segment_ids: + raise ValueError(f"Index '{name}' changed during the build; retry against the latest dataset version.") _validate_segments_against_manifest(lance_ds, index_metas, fragment_ids_to_use) + # Keep the checked handle: a later same-name CreateIndex transaction then + # conflicts in Lance instead of being treated as an index to replace. logger.info("Collected %d vector index segments; committing as segmented index %s", len(index_metas), name) lance_ds.commit_existing_index_segments(name, column, index_metas) logger.info("Vector index %s committed successfully", name) diff --git a/tests/io/lance/test_lance_vector_index.py b/tests/io/lance/test_lance_vector_index.py index 5ab167f..cb2fe89 100644 --- a/tests/io/lance/test_lance_vector_index.py +++ b/tests/io/lance/test_lance_vector_index.py @@ -471,3 +471,111 @@ def test_invalid_arguments_raise(tmp_path: Path) -> None: # would also fail on this small dataset with a large sample_rate. with pytest.raises(ValueError, match="fragment_group_size must be positive"): create_vector_index(uri, column="vector", index_type="IVF_FLAT", fragment_group_size=0) + + +def test_existing_segment_uuid_is_rejected_without_overwriting_index(tmp_path: Path) -> None: + """A caller cannot overwrite published segment files during a failed rebuild.""" + uri, vectors = _make_vector_dataset(tmp_path / "uuid.lance", num_rows=240, rows_per_file=60) + build = dict(column="vector", index_type="IVF_FLAT", num_partitions=4, sample_rate=8) + create_vector_index(uri, fragment_group_size=4, **build) + before = lance.dataset(uri) + segment_uuid = str(before.describe_indices()[0].segments[0].uuid) + query = {"column": "vector", "q": vectors[7], "k": 5, "nprobes": 4} + expected = before.to_table(nearest=query).to_pydict() + + with pytest.raises(TypeError, match="index_uuid"): + create_vector_index( + uri, + replace=True, + fragment_group_size=1, + max_concurrency=1, + index_uuid=segment_uuid, + **build, + ) + + after = lance.dataset(uri) + assert after.version == before.version + assert str(after.describe_indices()[0].segments[0].uuid) == segment_uuid + assert _covered_fragments(uri, "vector_idx") == {0, 1, 2, 3} + assert after.to_table(nearest=query).to_pydict() == expected + + +@pytest.mark.parametrize("initial_metric,append_metric", [("L2", "cosine"), ("cosine", None)]) +def test_backfill_rejects_a_different_metric(tmp_path: Path, initial_metric: str, append_metric: str | None) -> None: + """Segments must use comparable distances, including when backfill defaults to L2.""" + uri = str(tmp_path / "mixed_metric.lance") + vectors = np.array([[0, 1], [100, 0]], dtype=np.float32) + lance.write_dataset( + pa.table({"id": [0, 1], "vector": pa.FixedSizeListArray.from_arrays(vectors.reshape(-1), 2)}), + uri, + max_rows_per_file=1, + ) + centroids = pa.FixedSizeListArray.from_arrays(pa.array([1, 0], type=pa.float32()), 2) + build = dict(column="vector", index_type="IVF_FLAT", ivf_centroids=centroids) + create_vector_index(uri, metric=initial_metric, fragment_ids=[0], **build) + before = lance.dataset(uri) + query = {"column": "vector", "q": [1, 0], "k": 2, "nprobes": 1, "metric": initial_metric} + expected = before.to_table(nearest=query).to_pydict() + + append_kwargs = {} if append_metric is None else {"metric": append_metric} + with pytest.raises(ValueError, match="[Mm]etric"): + create_vector_index(uri, fragment_ids=[1], **build, **append_kwargs) + + after = lance.dataset(uri) + assert after.version == before.version + assert _covered_fragments(uri, "vector_idx") == {0} + assert after.to_table(nearest=query).to_pydict() == expected + + +@pytest.mark.parametrize("append_metric", ["l2", "L2", "euclidean", "EUCLIDEAN"]) +def test_backfill_accepts_equivalent_metric_spellings(tmp_path: Path, append_metric: str) -> None: + """Lance's Euclidean alias and case variations do not imply a changed metric.""" + uri, _ = _make_vector_dataset(tmp_path / "metric_alias.lance", num_rows=240, rows_per_file=120) + centroids = ( + IndicesBuilder(lance.dataset(uri), "vector") + .train_ivf(num_partitions=4, distance_type="l2", sample_rate=8) + .centroids + ) + build = dict(column="vector", index_type="IVF_FLAT", ivf_centroids=centroids) + create_vector_index(uri, metric="L2", fragment_ids=[0], **build) + create_vector_index(uri, metric=append_metric, fragment_ids=[1], **build) + assert _covered_fragments(uri, "vector_idx") == {0, 1} + + +def test_hnsw_build_parameters_are_forwarded(tmp_path: Path) -> None: + """Real HNSW kwargs are implemented by Lance even though its Python signature uses **kwargs.""" + uri, vectors = _make_vector_dataset(tmp_path / "hnsw_kwargs.lance", num_rows=240, rows_per_file=60) + create_vector_index( + uri, + column="vector", + index_type="IVF_HNSW_FLAT", + num_partitions=4, + sample_rate=8, + fragment_group_size=2, + m=4, + max_level=3, + ef_construction=24, + ) + ds = lance.dataset(uri) + assert len(ds.describe_indices()[0].segments) == 2 + for segment in ds.index_statistics("vector_idx")["segments"]: + params = segment["sub_index"]["params"] + assert params["m"] == 4 + assert params["max_level"] == 3 + assert params["ef_construction"] == 24 + assert 7 in ds.to_table(nearest={"column": "vector", "q": vectors[7], "k": 5, "nprobes": 4})["id"].to_pylist() + + +def test_centroids_file_is_rejected_before_segment_build(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + """Unsupported model-file input must fail before training or writing segment files.""" + uri, _ = _make_vector_dataset(tmp_path / "centroids_file.lance", num_rows=240, rows_per_file=60) + before = lance.dataset(uri) + + def unexpected_training(*args, **kwargs): + raise AssertionError("unsupported model-file argument reached training") + + monkeypatch.setattr(IndicesBuilder, "train_ivf", unexpected_training) + with pytest.raises(TypeError, match="ivf_centroids_file"): + create_vector_index(uri, column="vector", index_type="IVF_FLAT", ivf_centroids_file="centroids.npy") + assert lance.dataset(uri).version == before.version + assert not lance.dataset(uri).describe_indices() diff --git a/tests/io/lance/test_lance_vector_index_concurrency.py b/tests/io/lance/test_lance_vector_index_concurrency.py new file mode 100644 index 0000000..275ebd6 --- /dev/null +++ b/tests/io/lance/test_lance_vector_index_concurrency.py @@ -0,0 +1,133 @@ +from __future__ import annotations + +from pathlib import Path + +import lance +import numpy as np +import pytest +from lance.commit import CommitConflictError + +from daft.dependencies import pa +from daft_lance import create_vector_index, lance_vector_index +from daft_lance.namespace import DatasetOpenContext + + +def _dataset(tmp_path: Path): + vectors = np.random.default_rng(23).standard_normal((240, 8)).astype(np.float32) + uri = str(tmp_path / "concurrent.lance") + lance.write_dataset( + pa.table({"id": range(240), "vector": pa.FixedSizeListArray.from_arrays(vectors.ravel(), 8)}), + uri, + max_rows_per_file=40, + ) + centroids = ( + lance.indices.IndicesBuilder(lance.dataset(uri), "vector").train_ivf(num_partitions=4, sample_rate=8).centroids + ) + return uri, vectors, centroids + + +def _segments(dataset): + return {segment.uuid for segment in dataset.describe_indices()[0].segments} + + +def test_same_name_created_before_live_check_is_preserved(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + uri, vectors, centroids = _dataset(tmp_path) + original_open = DatasetOpenContext.open_latest + competing = {} + + def open_after_competing_commit(self): + dataset = lance.dataset(uri) + dataset.create_index("vector", "IVF_FLAT", name="vector_idx", ivf_centroids=centroids, num_partitions=4) + competing["version"] = dataset.version + competing["segments"] = _segments(dataset) + return original_open(self) + + monkeypatch.setattr(DatasetOpenContext, "open_latest", open_after_competing_commit) + with pytest.raises(ValueError, match="changed during the build"): + create_vector_index(uri, column="vector", index_type="IVF_FLAT", ivf_centroids=centroids, fragment_group_size=2) + + dataset = lance.dataset(uri) + assert dataset.version == competing["version"] + assert _segments(dataset) == competing["segments"] + assert 7 in dataset.to_table(nearest={"column": "vector", "q": vectors[7], "k": 5})["id"].to_pylist() + + +def test_same_name_created_after_live_check_causes_commit_conflict( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + uri, vectors, centroids = _dataset(tmp_path) + original_validate = lance_vector_index._validate_segments_against_manifest + competing = {} + + def validate_after_competing_commit(dataset, index_metas, expected_fragment_ids): + other = lance.dataset(uri) + other.create_index("vector", "IVF_FLAT", name="vector_idx", ivf_centroids=centroids, num_partitions=4) + competing["version"] = other.version + competing["segments"] = _segments(other) + return original_validate(dataset, index_metas, expected_fragment_ids) + + monkeypatch.setattr(lance_vector_index, "_validate_segments_against_manifest", validate_after_competing_commit) + with pytest.raises(CommitConflictError, match="CreateIndex"): + create_vector_index(uri, column="vector", index_type="IVF_FLAT", ivf_centroids=centroids, fragment_group_size=2) + + dataset = lance.dataset(uri) + assert dataset.version == competing["version"] + assert _segments(dataset) == competing["segments"] + assert 7 in dataset.to_table(nearest={"column": "vector", "q": vectors[7], "k": 5})["id"].to_pylist() + + +def test_backfill_rejects_concurrent_rebuild_with_changed_metric( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + uri, vectors, centroids = _dataset(tmp_path) + create_vector_index(uri, column="vector", index_type="IVF_FLAT", ivf_centroids=centroids, fragment_ids=[0, 1, 2]) + original_open = DatasetOpenContext.open_latest + competing = {} + + def open_after_rebuild(self): + dataset = lance.dataset(uri) + dataset.create_index( + "vector", + "IVF_FLAT", + name="vector_idx", + metric="cosine", + replace=True, + ivf_centroids=centroids, + num_partitions=4, + ) + competing["version"] = dataset.version + competing["segments"] = _segments(dataset) + return original_open(self) + + monkeypatch.setattr(DatasetOpenContext, "open_latest", open_after_rebuild) + with pytest.raises(ValueError, match="changed during the build"): + create_vector_index( + uri, column="vector", index_type="IVF_FLAT", ivf_centroids=centroids, fragment_ids=[3, 4, 5] + ) + + dataset = lance.dataset(uri) + assert dataset.version == competing["version"] + assert _segments(dataset) == competing["segments"] + assert dataset.index_statistics("vector_idx")["segments"][0]["metric_type"] == "cosine" + assert 7 in dataset.to_table(nearest={"column": "vector", "q": vectors[7], "k": 5})["id"].to_pylist() + + +def test_unrelated_append_after_live_check_allows_index_commit(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + uri, vectors, centroids = _dataset(tmp_path) + original_validate = lance_vector_index._validate_segments_against_manifest + + def validate_after_append(dataset, index_metas, expected_fragment_ids): + lance.write_dataset( + pa.table({"id": [240], "vector": pa.FixedSizeListArray.from_arrays(vectors[7], 8)}), uri, mode="append" + ) + return original_validate(dataset, index_metas, expected_fragment_ids) + + monkeypatch.setattr(lance_vector_index, "_validate_segments_against_manifest", validate_after_append) + create_vector_index(uri, column="vector", index_type="IVF_FLAT", ivf_centroids=centroids, fragment_group_size=2) + + dataset = lance.dataset(uri) + assert dataset.version == 3 + assert dataset.count_rows() == 241 + assert dataset.describe_indices()[0].num_rows_indexed == 240 + assert len(_segments(dataset)) == 3 + assert 7 in dataset.to_table(nearest={"column": "vector", "q": vectors[7], "k": 5})["id"].to_pylist() From 7e9a2546a124ec18be24f8583eaa53e5cf7dd0c6 Mon Sep 17 00:00:00 2001 From: wangzheyan Date: Sat, 10 Oct 2026 15:19:54 +0800 Subject: [PATCH 09/11] feat: support configurable PQ bit width in distributed builds --- .pre-commit-config.yaml | 2 +- README.md | 6 +- daft_lance/_lance.py | 20 +++-- daft_lance/lance_vector_index.py | 49 +++++------ tests/io/lance/test_lance_vector_index.py | 100 +++++++++++++++++++++- 5 files changed, 139 insertions(+), 38 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 0bab98e..635673d 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -9,7 +9,7 @@ repos: files: daft_lance additional_dependencies: - 'daft>=0.7.0' - - 'pylance>=6.0.0' + - 'pylance>=11.0.0' - 'pyarrow-stubs>=20.0.0' - repo: https://github.com/pre-commit/pre-commit-hooks diff --git a/README.md b/README.md index 7808376..b616652 100644 --- a/README.md +++ b/README.md @@ -143,6 +143,10 @@ latest version. Full rebuilds with `replace=True` retain replacement semantics. HNSW builds accept `m`, `max_level`, and `ef_construction`. Each worker generates its own segment UUID, so `index_uuid` cannot be supplied. For pre-trained IVF models, pass `ivf_centroids`; `ivf_centroids_file` is not supported by this API. +PQ builds accept `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. When +supplying a `pq_codebook`, also pass its original `num_sub_vectors` and +`num_bits`; a bare codebook array does not carry the sub-vector count. The SQ variants (`IVF_SQ`, `IVF_HNSW_SQ`) support multi-segment builds and `fragment_ids` backfill. Each segment uses its own SQ quantization metadata @@ -158,7 +162,7 @@ Training samples `sample_rate` rows per IVF partition (and per PQ centroid) and runs in the coordinator process, so its memory footprint grows with `num_partitions * sample_rate * dimension`; `sample_rate` is clamped down automatically (with a warning) to what the dataset size supports -(`num_partitions * sample_rate` rows, `256 * sample_rate` for the 8-bit PQ +(`num_partitions * sample_rate` rows, `2**num_bits * sample_rate` for the PQ codebook). diff --git a/daft_lance/_lance.py b/daft_lance/_lance.py index 54e3010..1e6dddb 100644 --- a/daft_lance/_lance.py +++ b/daft_lance/_lance.py @@ -698,6 +698,7 @@ def create_vector_index( metric: str = "L2", num_partitions: int | None = None, num_sub_vectors: int | None = None, + num_bits: int = 8, sample_rate: int = 256, ivf_centroids: Any | None = None, pq_codebook: Any | None = None, @@ -756,21 +757,27 @@ def create_vector_index( ``num_partitions`` controls Daft's fragment repartitioning; build 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, or from - the supplied ``pq_codebook`` when one is given. + Lance derives it from the vector dimension during training. Required + when supplying a ``pq_codebook``; the array does not carry this metadata. + 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. + Supplied PQ codebooks and backfill must use the matching bit width. sample_rate: Rows sampled per IVF partition (and per PQ centroid) during training. Training requires at least ``num_partitions * sample_rate`` - rows (256 * sample_rate for the 8-bit PQ codebook); the value is + rows (2**num_bits * sample_rate for the PQ codebook); the value is automatically clamped down (with a warning) to what the dataset - size supports, so the default works on datasets of any size, at the - cost of a smaller training sample on small datasets. + size supports, at the cost of a smaller training sample on small + datasets. Lance still rejects data below the minimum training size + (2**num_bits rows for PQ); clamping does not remove that requirement. ivf_centroids: Pre-trained IVF centroids (a pyarrow array); skips IVF training. Supplying the same centroids to every build is how independently built segments share one model. Required when appending to an existing index (``fragment_ids`` backfill): all segments of a logical vector index must share one IVF model. pq_codebook: Pre-trained PQ codebook (PQ variants only); skips PQ - training. Required when appending to an existing PQ index, for the + training. Supply its original ``num_sub_vectors`` and ``num_bits``. + Required when appending to an existing PQ index, for the same shared-model reason as ``ivf_centroids``. storage_options: Storage options for the dataset. version: Version of the dataset to use. @@ -863,6 +870,7 @@ def create_vector_index( metric=metric, num_partitions=num_partitions, num_sub_vectors=num_sub_vectors, + num_bits=num_bits, sample_rate=sample_rate, ivf_centroids=ivf_centroids, pq_codebook=pq_codebook, diff --git a/daft_lance/lance_vector_index.py b/daft_lance/lance_vector_index.py index 5999ec2..2fe7ee3 100644 --- a/daft_lance/lance_vector_index.py +++ b/daft_lance/lance_vector_index.py @@ -31,10 +31,6 @@ # PQ variants need a codebook trained alongside the IVF centroids. _PQ_INDEX_TYPES = frozenset({"IVF_PQ", "IVF_HNSW_PQ"}) -# The 8-bit PQ codebook trains one centroid per 2^8 codes, so its sample -# requirement is 256 rows per sampled codebook entry. -_PQ_CODEBOOK_SIZE = 256 - # Arguments managed by this workflow or incompatible with its shared-model # builds; they are not valid user kwargs on the worker's segment-build call. _HANDLER_MANAGED_KWARGS = frozenset( @@ -48,6 +44,7 @@ "metric", "num_partitions", "num_sub_vectors", + "num_bits", "ivf_centroids", "pq_codebook", "storage_options", @@ -168,6 +165,7 @@ def create_vector_index_internal( metric: str = "L2", num_partitions: int | None = None, num_sub_vectors: int | None = None, + num_bits: int = 8, sample_rate: int = 256, ivf_centroids: pa.Array[Any] | None = None, pq_codebook: pa.Array[Any] | None = None, @@ -190,10 +188,10 @@ def create_vector_index_internal( ``ivf_centroids`` / ``pq_codebook``. One shared model across segments is what lets independently built segments commit as one logical index. ``sample_rate`` is clamped down to what the dataset size supports - (training needs ``num_partitions * sample_rate`` rows, and the 8-bit PQ - codebook needs ``256 * sample_rate``); ``num_sub_vectors`` and - ``num_partitions`` are derived from a supplied codebook / centroids when - not given explicitly. + (training needs ``num_partitions * sample_rate`` rows, and the PQ + codebook needs ``2**num_bits * sample_rate``). ``num_partitions`` is + derived from supplied centroids. A supplied codebook requires explicit + ``num_sub_vectors`` because the array does not carry that metadata. 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. @@ -226,15 +224,19 @@ 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.") # Reject kwargs Lance's segment build would silently swallow (misspelled # or unsupported parameters must not build a misconfigured index). _validate_worker_kwargs(kwargs) + if index_type in _PQ_INDEX_TYPES and pq_codebook is not None and num_sub_vectors is None: + raise ValueError("A supplied pq_codebook requires its original num_sub_vectors; it cannot be inferred.") # Validate column exists; whether it is a vector column is Lance's rule, # enforced by the training and build APIs with clear errors. try: - field = lance_ds.schema.field(column) + lance_ds.schema.field(column) except KeyError as e: available_columns = [field.name for field in lance_ds.schema] raise ValueError(f"Column '{column}' not found. Available: {available_columns}") from e @@ -352,12 +354,14 @@ def create_vector_index_internal( # Clamp the sample rate to what this dataset can support, mirroring # Lance's own training requirements: num_partitions * sample_rate rows # for IVF (Lance derives num_partitions as sqrt(num_rows) when None) - # and 256 * sample_rate rows for the 8-bit PQ codebook. - effective_partitions = num_partitions if num_partitions is not None else max(1, round(math.sqrt(num_rows))) - caps = [max(1, num_rows // effective_partitions)] + # and 2**num_bits * sample_rate rows for the PQ codebook. + caps = [sample_rate] + if needs_ivf_training: + effective_partitions = num_partitions if num_partitions is not None else max(1, round(math.sqrt(num_rows))) + caps.append(max(1, num_rows // effective_partitions)) if needs_pq_training: - caps.append(max(1, num_rows // _PQ_CODEBOOK_SIZE)) - clamped = min([sample_rate, *caps]) + caps.append(max(1, num_rows // (2**num_bits))) + clamped = min(caps) if clamped < sample_rate: logger.warning( "sample_rate %d exceeds what the dataset supports (%d rows); clamping to %d", @@ -401,34 +405,23 @@ def create_vector_index_internal( pq_model = builder.train_pq( ivf_model, num_subvectors=num_sub_vectors, + num_bits=num_bits, sample_rate=effective_sample_rate, ) pq_codebook = pq_model.codebook num_sub_vectors = pq_model.num_subvectors logger.info("PQ training completed: num_sub_vectors=%d", num_sub_vectors) - # Derive the model shape from supplied artifacts when not given: the - # partition count is the number of centroids, and the sub-vector count - # follows from the codebook's entry size versus the column dimension. + # Derive the partition count from supplied centroids when not given. if num_partitions is None and ivf_centroids is not None: num_partitions = len(ivf_centroids) - if num_sub_vectors is None and pq_codebook is not None: - dimension = getattr(field.type, "list_size", None) - subvector_size = getattr(pq_codebook.type, "list_size", None) - if dimension and subvector_size and dimension % subvector_size == 0: - num_sub_vectors = dimension // subvector_size - else: - raise ValueError( - f"Cannot derive num_sub_vectors from the supplied pq_codebook " - f"(column dimension {dimension}, codebook entry size {subvector_size}); " - "pass num_sub_vectors explicitly." - ) model_kwargs: dict[str, Any] = { "metric": metric, "ivf_centroids": ivf_centroids, "num_partitions": num_partitions, "num_sub_vectors": num_sub_vectors, + "num_bits": num_bits, "pq_codebook": pq_codebook, } diff --git a/tests/io/lance/test_lance_vector_index.py b/tests/io/lance/test_lance_vector_index.py index cb2fe89..1050a73 100644 --- a/tests/io/lance/test_lance_vector_index.py +++ b/tests/io/lance/test_lance_vector_index.py @@ -285,26 +285,36 @@ def test_pretrained_centroids_are_used(tmp_path: Path) -> None: assert desc.num_rows_indexed == 240 -def test_pq_codebook_without_num_sub_vectors_derives_it(tmp_path: Path) -> None: - """A supplied codebook implies num_sub_vectors; no deep worker error.""" +def test_pq_codebook_requires_original_num_sub_vectors(tmp_path: Path) -> None: + """A bare codebook cannot reveal how the original vectors were split.""" uri, _ = _make_vector_dataset(tmp_path / "codebook.lance", seed=19) ds = lance.dataset(uri) builder = IndicesBuilder(ds, "vector") ivf = builder.train_ivf(num_partitions=4, distance_type="l2", sample_rate=32) pq = builder.train_pq(ivf, num_subvectors=2, sample_rate=8) + with pytest.raises(ValueError, match="original num_sub_vectors"): + create_vector_index( + uri, column="vector", index_type="IVF_PQ", ivf_centroids=ivf.centroids, pq_codebook=pq.codebook + ) + create_vector_index( uri, column="vector", index_type="IVF_PQ", ivf_centroids=ivf.centroids, pq_codebook=pq.codebook, + num_sub_vectors=2, fragment_group_size=4, ) desc = lance.dataset(uri).describe_indices()[0] assert desc.index_type == "IVF_PQ" assert desc.num_rows_indexed == 2048 + assert all( + segment["sub_index"]["num_sub_vectors"] == 2 + for segment in lance.dataset(uri).stats.index_stats("vector_idx")["indices"] + ) def test_backfill_requires_the_original_shared_model(tmp_path: Path) -> None: @@ -579,3 +589,89 @@ def unexpected_training(*args, **kwargs): create_vector_index(uri, column="vector", index_type="IVF_FLAT", ivf_centroids_file="centroids.npy") assert lance.dataset(uri).version == before.version assert not lance.dataset(uri).describe_indices() + + +@pytest.mark.parametrize("index_type", ["IVF_PQ", "IVF_HNSW_PQ"]) +@pytest.mark.parametrize("num_bits", [4, 8]) +def test_pq_bit_width_is_shared_by_all_segments(tmp_path: Path, index_type: str, num_bits: int) -> None: + uri, vectors = _make_vector_dataset(tmp_path / "bits.lance", num_rows=512, rows_per_file=128) + create_vector_index( + uri, + column="vector", + index_type=index_type, + num_partitions=2, + num_sub_vectors=2, + 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 + assert all(segment["sub_index"]["nbits"] == num_bits for segment in segments) + assert all(segment["sub_index"]["num_sub_vectors"] == 2 for segment in segments) + result = ds.to_table(nearest={"column": "vector", "q": vectors[7], "k": 5, "nprobes": 2}) + assert result.num_rows == 5 + assert np.isfinite(result["_distance"].to_numpy()).all() + + +def test_four_bit_pq_sampling_uses_sixteen_centroids(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + uri, _ = _make_vector_dataset(tmp_path / "sampling_bits.lance", num_rows=128, rows_per_file=64) + centroids = pa.FixedSizeListArray.from_arrays(pa.array([0.0] * 8, type=pa.float32()), 8) + original_train = IndicesBuilder.train_pq + calls = [] + + def train_with_recorded_sample(builder, *args, **kwargs): + calls.append((kwargs["num_bits"], kwargs["sample_rate"])) + return original_train(builder, *args, **kwargs) + + monkeypatch.setattr(IndicesBuilder, "train_pq", train_with_recorded_sample) + create_vector_index( + uri, + column="vector", + index_type="IVF_PQ", + ivf_centroids=centroids, + num_sub_vectors=2, + num_bits=4, + sample_rate=8, + fragment_group_size=1, + ) + # 16 * 8 samples fit in 128 rows; no IVF training or 256-entry cap applies. + assert calls == [(4, 8)] + assert all( + segment["sub_index"]["nbits"] == 4 for segment in lance.dataset(uri).stats.index_stats("vector_idx")["indices"] + ) + + +def test_supplied_four_bit_pq_model_backfills_with_original_configuration(tmp_path: Path) -> None: + uri, vectors = _make_vector_dataset(tmp_path / "supplied_bits.lance", num_rows=256, rows_per_file=128) + builder = IndicesBuilder(lance.dataset(uri), "vector") + ivf = builder.train_ivf(num_partitions=2, sample_rate=8) + pq = builder.train_pq(ivf, num_subvectors=2, num_bits=4, sample_rate=8) + build = dict( + column="vector", + index_type="IVF_PQ", + ivf_centroids=ivf.centroids, + pq_codebook=pq.codebook, + num_sub_vectors=2, + num_bits=4, + ) + create_vector_index(uri, fragment_ids=[0], **build) + create_vector_index(uri, fragment_ids=[1], **build) + ds = lance.dataset(uri) + segments = ds.stats.index_stats("vector_idx")["indices"] + assert len(segments) == 2 + assert all(segment["sub_index"]["nbits"] == 4 for segment in segments) + assert all(segment["sub_index"]["num_sub_vectors"] == 2 for segment in segments) + assert _covered_fragments(uri, "vector_idx") == {0, 1} + assert ds.to_table(nearest={"column": "vector", "q": vectors[7], "k": 5, "nprobes": 2}).num_rows == 5 + + +@pytest.mark.parametrize("index_type", ["IVF_FLAT", "IVF_SQ", "IVF_HNSW_FLAT", "IVF_HNSW_SQ"]) +def test_non_pq_bit_width_is_not_silently_ignored(tmp_path: Path, index_type: str) -> None: + uri, _ = _make_vector_dataset(tmp_path / "non_pq_bits.lance", num_rows=128, rows_per_file=64) + before = lance.dataset(uri) + with pytest.raises(ValueError, match="configurable only for IVF_PQ"): + create_vector_index(uri, column="vector", index_type=index_type, num_bits=4) + assert lance.dataset(uri).version == before.version + assert not lance.dataset(uri).describe_indices() From b12dd4b35f3df3cb80cb51b88ee2a7e5b4cf4ea1 Mon Sep 17 00:00:00 2001 From: wangzheyan Date: Sat, 10 Oct 2026 15:42:01 +0800 Subject: [PATCH 10/11] refactor: separate full vector index builds from incremental maintenance --- README.md | 49 ++++-- daft_lance/_lance.py | 42 ++--- daft_lance/lance_vector_index.py | 134 ++++---------- tests/io/lance/test_lance_vector_index.py | 165 ++++++++---------- .../test_lance_vector_index_concurrency.py | 73 ++++++-- 5 files changed, 208 insertions(+), 255 deletions(-) diff --git a/README.md b/README.md index b616652..7acc93d 100644 --- a/README.md +++ b/README.md @@ -128,17 +128,33 @@ defaults `replace` to `False`. An existing index name is refused unless `replace=True`, which rebuilds atomically (one manifest transaction retires the old segments and lands the new ones). -`fragment_ids` partial builds and incremental backfill work like the scalar -workflow, with one vector-specific requirement: appending segments requires -the same `ivf_centroids` (and `pq_codebook` for PQ variants) the existing -segments were built with, since every segment of a logical vector index must -share one IVF model — a backfill without the original model raises instead of -silently training a divergent one. -Backfill must also use the existing index's `metric`; a different metric is -rejected because distances from different metrics cannot be ranked together. -If another task changes the same index during a new build or backfill, the -operation fails rather than replacing that task's result; retry against the -latest version. Full rebuilds with `replace=True` retain replacement semantics. +`create_vector_index` always builds every fragment in the opened dataset version. +Use `optimize_indices` after appending data: Lance reuses the existing index model, +indexes uncovered fragments, and merges existing segments according to +`num_indices_to_merge`. Setting that option to `0` builds a new segment without +merging old segments. Full rebuilds remain distributed and use `replace=True`. +With `replace=False`, a concurrent same-name index creation is rejected; ordinary +data appends do not prevent the planned index commit. +If another task indexes appended fragments outside the rebuild's planned +snapshot, full replacement fails and must be retried to include those fragments. +This prevents old model segments from surviving a full rebuild. + +**Breaking change — vector index migration:** `create_vector_index` no longer +accepts `fragment_ids`, partial builds, or backfill/no-op requests. For an initial +build, remove `fragment_ids` to create coverage for the full dataset. Once the +index exists, replace calls that used `create_vector_index(..., fragment_ids=[...])` +to cover appended data with: + +```python +from daft_lance import optimize_indices + +optimize_indices("s3://bucket/my_dataset", indices=["vector_idx"], num_indices_to_merge=0) +``` + +Maintenance discovers uncovered fragments automatically; it needs neither a +fragment list nor caller-supplied models. To retrain or replace the index, use +`create_vector_index(..., replace=True)`. Scalar `fragment_ids` and backfill +semantics are unchanged. HNSW builds accept `m`, `max_level`, and `ef_construction`. Each worker generates its own segment UUID, so `index_uuid` cannot be supplied. For pre-trained IVF @@ -148,15 +164,16 @@ bit width is used for shared codebook training and every segment build. When supplying a `pq_codebook`, also pass its original `num_sub_vectors` and `num_bits`; a bare codebook array does not carry the sub-vector count. -The SQ variants (`IVF_SQ`, `IVF_HNSW_SQ`) support multi-segment builds and -`fragment_ids` backfill. Each segment uses its own SQ quantization metadata +The SQ variants (`IVF_SQ`, `IVF_HNSW_SQ`) support multi-segment builds. +Each segment uses its own SQ quantization metadata for queries. Independently built segments may have different quantization metadata even when they share IVF centroids; Lance rejects attempts to merge incompatible segments through `optimize_indices` with "vector index segments do not share quantizer metadata". Use `fragment_group_size` large enough to -cover all selected fragments when a single segment is preferred. To consolidate -an existing index, rebuild with that grouping, `replace=True`, and no -`fragment_ids`. +cover all fragments when a single segment is preferred. To consolidate +an existing index, rebuild with that grouping and `replace=True`. Use +`optimize_indices(..., num_indices_to_merge=0)` to add coverage without merging +incompatible old SQ segments. Training samples `sample_rate` rows per IVF partition (and per PQ centroid) and runs in the coordinator process, so its memory footprint grows with diff --git a/daft_lance/_lance.py b/daft_lance/_lance.py index 1e6dddb..63352ab 100644 --- a/daft_lance/_lance.py +++ b/daft_lance/_lance.py @@ -712,10 +712,9 @@ def create_vector_index( metadata_cache_size_bytes: int | None = None, fragment_group_size: int | None = None, max_concurrency: int | None = None, - fragment_ids: list[int] | None = None, **kwargs: Any, ) -> None: - """Build a distributed vector index using Daft's distributed execution. + """Build a distributed vector index over every fragment in the opened version. The build runs in three phases: the driver trains the global model once (IVF centroids via ``IndicesBuilder.train_ivf``, and a PQ codebook via @@ -728,6 +727,9 @@ def create_vector_index( centroids and codebook commit as one logical index, so no central build step is needed. + Incremental indexing of appended data belongs to ``optimize_indices``. + This API does not accept ``fragment_ids`` and does not skip covered data. + Args: uri: The URI of the Lance table (supports remote URLs to object stores such as `s3://` or `gs://`) io_config: A custom IOConfig to use when accessing LanceDB data. Defaults to None. @@ -747,10 +749,11 @@ def create_vector_index( refused with ``ValueError`` unless ``replace=True``. Replacement is atomic: all new segments commit in one transaction that retires the overlapped old segments, so readers see either the old or the new index. - New builds and backfill fail if another task changes the same index - during the build; retry the operation against the latest version. - metric: Distance metric ("L2", "cosine", or "dot"). Backfill must use - the same metric as the existing index; mismatches raise ``ValueError``. + With replace=False, a concurrent same-name creation fails this build + instead of being overwritten; retry against the latest version. + Full replacement also fails if concurrent index maintenance adds + coverage outside the planned snapshot; retry to include that data. + metric: Distance metric ("L2", "cosine", or "dot"). num_partitions: Number of IVF partitions. If None, Lance derives it from the number of rows during training. Note this is the IVF partition count of the trained model — unlike ``create_scalar_index``, where @@ -762,7 +765,7 @@ def create_vector_index( 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. - Supplied PQ codebooks and backfill must use the matching bit width. + Supplied PQ codebooks must use the matching bit width. sample_rate: Rows sampled per IVF partition (and per PQ centroid) during training. Training requires at least ``num_partitions * sample_rate`` rows (2**num_bits * sample_rate for the PQ codebook); the value is @@ -772,13 +775,9 @@ def create_vector_index( (2**num_bits rows for PQ); clamping does not remove that requirement. ivf_centroids: Pre-trained IVF centroids (a pyarrow array); skips IVF training. Supplying the same centroids to every build is how - independently built segments share one model. Required when - appending to an existing index (``fragment_ids`` backfill): all - segments of a logical vector index must share one IVF model. + independently built segments share one model. pq_codebook: Pre-trained PQ codebook (PQ variants only); skips PQ training. Supply its original ``num_sub_vectors`` and ``num_bits``. - Required when appending to an existing PQ index, for the - same shared-model reason as ``ivf_centroids``. storage_options: Storage options for the dataset. version: Version of the dataset to use. asof: Timestamp to use for time travel queries. @@ -791,16 +790,6 @@ def create_vector_index( defaults to 10. Each group becomes one independently built index segment. max_concurrency: Maximum number of concurrent tasks for fragment batches. If None, Daft uses its default concurrency setting. - fragment_ids: Optional list of fragment IDs to index. Fragments already - covered by committed segments of the same index are skipped and the - remaining ones appended (incremental backfill); pass ``replace=True`` - without ``fragment_ids`` for a full rebuild. Appending requires the - same ``ivf_centroids`` (and ``pq_codebook`` for PQ variants) the - existing segments were built with — retraining would produce a - divergent model that Lance cannot merge, so a backfill without the - original model raises ``ValueError``. SQ variants also support - backfill, but sharing IVF centroids does not guarantee compatible - SQ quantization metadata for subsequent segment merging. **kwargs: Additional keyword arguments forwarded to Lance's index segment creation API, including HNSW ``m``, ``max_level`` and ``ef_construction``. ``index_uuid`` is managed per segment. ``ivf_centroids_file`` is not @@ -814,8 +803,8 @@ def create_vector_index( Raises: ValueError: If input parameters are invalid (e.g. empty column name, non-existent column, unsupported index type, an existing index - name with ``replace=False``, a backfill without the original - shared model, or a dataset too small to train an IVF model even + name with ``replace=False``, a supplied codebook without its original + sub-vector count, or a dataset too small to train an IVF model even at ``sample_rate=1`` — Lance raises its training ``ValueError``) TypeError: If column type is incompatible with the chosen ``index_type``, or an unknown keyword argument is passed @@ -839,8 +828,8 @@ def create_vector_index( Refuse to overwrite an existing index: >>> daft_lance.create_vector_index("s3://my-bucket/dataset/", column="vector", replace=False) - Index appended fragments incrementally: - >>> daft_lance.create_vector_index("s3://my-bucket/dataset/", column="vector", fragment_ids=[10, 11]) + Maintain coverage after appending data: + >>> daft_lance.optimize_indices("s3://my-bucket/dataset/", indices=["vector_idx"]) """ io_config = context.get_context().daft_planning_config.default_io_config if io_config is None else io_config @@ -876,7 +865,6 @@ def create_vector_index( pq_codebook=pq_codebook, fragment_group_size=fragment_group_size, max_concurrency=max_concurrency, - fragment_ids=fragment_ids, **kwargs, ) diff --git a/daft_lance/lance_vector_index.py b/daft_lance/lance_vector_index.py index 2fe7ee3..821bf1f 100644 --- a/daft_lance/lance_vector_index.py +++ b/daft_lance/lance_vector_index.py @@ -17,7 +17,6 @@ from daft.dependencies import pa from daft_lance.lance_scalar_index import ( - _existing_index_coverage, _validate_segments_against_manifest, ) from daft_lance.utils import distribute_fragments_balanced @@ -69,6 +68,11 @@ def _accepted_worker_kwargs() -> frozenset[str]: def _validate_worker_kwargs(kwargs: dict[str, Any]) -> None: """Reject kwargs Lance's segment build would silently ignore.""" + if "fragment_ids" in kwargs: + raise TypeError( + "create_vector_index no longer accepts fragment_ids; it builds all fragments. " + "Use optimize_indices to index appended data." + ) if "segmented" in kwargs: raise TypeError( "The 'segmented' parameter was removed: the distributed segment-index " @@ -94,8 +98,8 @@ def _validate_worker_kwargs(kwargs: dict[str, Any]) -> None: ) -def _index_segment_ids(dataset: lance.LanceDataset, name: str) -> set[str]: - return {segment.uuid for index in dataset.describe_indices() if index.name == name for segment in index.segments} +def _index_exists(dataset: lance.LanceDataset, name: str) -> bool: + return any(index.name == name for index in dataset.describe_indices()) class VectorFragmentIndexHandler: @@ -171,7 +175,6 @@ def create_vector_index_internal( pq_codebook: pa.Array[Any] | None = None, fragment_group_size: int | None = None, max_concurrency: int | None = None, - fragment_ids: list[int] | None = None, **kwargs: Any, ) -> None: """Internal implementation of distributed vector index creation. @@ -203,13 +206,8 @@ def create_vector_index_internal( ``replace`` defaults to ``False``, matching pylance's ``create_index``: an existing index name is refused unless ``replace=True``. Column type compatibility is validated by Lance's training and build APIs, not - duplicated here. ``fragment_ids`` restricts the build to a subset of - fragments; already-covered fragments are skipped and the remainder - appended. Appending requires the same ``ivf_centroids`` (and - ``pq_codebook`` for PQ variants) the existing segments were built with: - every segment of a logical vector index must share one IVF model, or - Lance cannot merge the segments later, so a backfill without the original - model raises instead of silently training a divergent one. + duplicated here. Every fragment in the opened dataset snapshot is indexed. + Incremental coverage of appended data belongs to ``optimize_indices``. """ if not column: raise ValueError("Column name cannot be empty") @@ -245,96 +243,15 @@ def create_vector_index_internal( if name is None: name = f"{column}_idx" - fragments = lance_ds.get_fragments() - available_fragment_ids = {fragment.fragment_id for fragment in fragments} - - # Validate and normalize the requested fragment subset, if any. - requested_fragment_ids: set[int] | None = None - if fragment_ids is not None: - if len(fragment_ids) == 0: - raise ValueError("fragment_ids must be a non-empty list of fragment IDs; pass None to index all fragments.") - unique_ids = list(dict.fromkeys(fragment_ids)) - duplicates = sorted({fid for fid in unique_ids if fragment_ids.count(fid) > 1}) - if duplicates: - logger.warning("Duplicate fragment_ids %s were given; each fragment is scheduled once.", duplicates) - unknown_ids = sorted(fid for fid in unique_ids if fid not in available_fragment_ids) - if unknown_ids: - raise ValueError( - f"fragment_ids {unknown_ids} do not exist in the dataset. " - f"Available fragment IDs: {sorted(available_fragment_ids)}" - ) - requested_fragment_ids = set(unique_ids) - - existing_coverage = _existing_index_coverage(lance_ds, name) - planned_segment_ids = _index_segment_ids(lance_ds, name) - if existing_coverage is not None: - # Column/index-type compatibility of a same-name index is validated by - # Lance's build and commit APIs, not duplicated here. - if not replace and requested_fragment_ids is None: - raise ValueError(f"Index with name '{name}' already exists. Set replace=True to replace it.") - - # Workers open the pinned snapshot where a same-name index still exists; - # building against that name always requires replace=True. The actual - # replacement happens in the coordinator's atomic commit. - handler_replace = existing_coverage is not None - if existing_coverage is not None and requested_fragment_ids is not None: - # Incremental backfill: skip fragments already covered by committed - # segments; only the remainder is built and appended. - covered = existing_coverage & available_fragment_ids - already_covered = requested_fragment_ids & covered - to_build = requested_fragment_ids - covered - if already_covered: - logger.info( - "Fragments %s are already covered by index '%s'; skipping them%s", - sorted(already_covered), - name, - " (replace does not rebuild covered fragments; use replace=True without " - "fragment_ids for a full rebuild)" - if replace - else "", - ) - if not to_build: - logger.info("All requested fragments are already covered by index '%s'; nothing to build.", name) - return - # Lance permits mixed metrics at commit, but their distances cannot be - # compared when searching the logical index. - stats = lance_ds.stats.index_stats(name) - if stats["index_type"].startswith("IVF"): - requested_metric = metric.lower() - if requested_metric == "euclidean": - requested_metric = "l2" - existing_metrics = {segment["metric_type"] for segment in stats["indices"]} - if existing_metrics != {requested_metric}: - raise ValueError( - f"Cannot append to index '{name}' with metric '{metric}': " - f"existing segments use {sorted(existing_metrics)}. Use the same metric." - ) - # Appended segments join a live logical index, and every segment of a - # logical vector index must share the same IVF model: Lance's segment - # merge refuses segments trained on different centroids (the error - # ``optimize_indices`` would later raise). Training here would sample - # non-deterministically and silently produce a divergent model, so the - # original model must be supplied. - if ivf_centroids is None: - raise ValueError( - f"Cannot append to index '{name}': all segments of a vector index must share " - "one IVF model, and retraining would produce a divergent one. Pass the same " - "ivf_centroids (and pq_codebook, for PQ variants) the existing segments were " - "built with, or rebuild the whole index with replace=True and no fragment_ids." - ) - if index_type in _PQ_INDEX_TYPES and pq_codebook is None: - raise ValueError( - f"Cannot append to PQ index '{name}' without the original pq_codebook: all " - "segments must share one model. Pass the codebook the existing segments were " - "built with, or rebuild the whole index with replace=True and no fragment_ids." - ) - requested_fragment_ids = to_build + index_exists = _index_exists(lance_ds, name) + if index_exists and not replace: + raise ValueError(f"Index with name '{name}' already exists. Set replace=True to replace it.") - if requested_fragment_ids is not None: - fragments = [fragment for fragment in fragments if fragment.fragment_id in requested_fragment_ids] - fragment_ids_to_use = sorted( - requested_fragment_ids if requested_fragment_ids is not None else (f.fragment_id for f in fragments) - ) + # Each worker builds against the pinned snapshot. Replacement is committed + # once, after every fragment batch has succeeded. + handler_replace = index_exists + fragments = lance_ds.get_fragments() + fragment_ids_to_use = sorted(fragment.fragment_id for fragment in fragments) if not fragment_ids_to_use: raise ValueError(f"Dataset at {open_context.uri} contains no fragments") @@ -467,10 +384,21 @@ def create_vector_index_internal( # Phase 3: validate against the live manifest and commit atomically. lance_ds = open_context.open_latest() - if (not replace or requested_fragment_ids is not None) and _index_segment_ids( - lance_ds, name - ) != planned_segment_ids: + if not replace and _index_exists(lance_ds, name): raise ValueError(f"Index '{name}' changed during the build; retry against the latest dataset version.") + if replace: + indexed_fragments = { + fid + for index in lance_ds.describe_indices() + if index.name == name + for segment in index.segments + for fid in segment.fragment_ids + } + live_fragments = {fragment.fragment_id for fragment in lance_ds.get_fragments()} + if (indexed_fragments & live_fragments).difference(fragment_ids_to_use): + raise ValueError( + f"Index '{name}' gained coverage outside the build plan; retry against the latest dataset version." + ) _validate_segments_against_manifest(lance_ds, index_metas, fragment_ids_to_use) # Keep the checked handle: a later same-name CreateIndex transaction then diff --git a/tests/io/lance/test_lance_vector_index.py b/tests/io/lance/test_lance_vector_index.py index 1050a73..92105f4 100644 --- a/tests/io/lance/test_lance_vector_index.py +++ b/tests/io/lance/test_lance_vector_index.py @@ -131,18 +131,19 @@ def test_every_index_type_builds_and_answers_queries(tmp_path: Path, index_type: assert desc.index_type == index_type assert desc.num_rows_indexed == 2048 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) assert recall >= _MIN_RECALL[index_type], f"{index_type} recall@5 {recall:.2f} < {_MIN_RECALL[index_type]}" @pytest.mark.parametrize("index_type", ["IVF_SQ", "IVF_HNSW_SQ"]) -def test_sq_backfill_and_failed_merge_preserve_index(tmp_path: Path, index_type: str) -> None: - """SQ backfill is searchable; incompatible merges leave the committed index intact.""" +def test_sq_multi_segment_build_and_failed_merge_preserve_index(tmp_path: Path, index_type: str) -> None: + """A full SQ build is searchable; incompatible merges preserve the committed index.""" rng = np.random.default_rng(7) vectors = rng.uniform(0, 1, (2048, 8)).astype(np.float32) - vectors[1024:] += 10 - uri = str(tmp_path / "sq_backfill.lance") + vectors.reshape(8, 256, 8)[1::2] += 10 + uri = str(tmp_path / "sq_segments.lance") lance.write_dataset( pa.table({"id": list(range(2048)), "vector": pa.FixedSizeListArray.from_arrays(vectors.reshape(-1), 8)}), uri, @@ -154,8 +155,7 @@ def test_sq_backfill_and_failed_merge_preserve_index(tmp_path: Path, index_type: .centroids ) build = dict(column="vector", index_type=index_type, ivf_centroids=centroids, name="v_idx") - create_vector_index(uri, fragment_ids=[0, 1, 2, 3], **build) - create_vector_index(uri, fragment_ids=[4, 5, 6, 7], **build) + create_vector_index(uri, fragment_group_size=4, **build) before = lance.dataset(uri) assert len(before.describe_indices()[0].segments) == 2 @@ -317,46 +317,6 @@ def test_pq_codebook_requires_original_num_sub_vectors(tmp_path: Path) -> None: ) -def test_backfill_requires_the_original_shared_model(tmp_path: Path) -> None: - """Appending segments without the original centroids would fork the IVF model.""" - uri, _ = _make_vector_dataset(tmp_path / "guard.lance") - build = dict(column="vector", index_type="IVF_FLAT", num_partitions=4, sample_rate=8, name="v_idx") - - create_vector_index(uri, fragment_ids=[0, 1, 2, 3], **build) - - # Retraining on the backfill would produce a divergent model — refused. - with pytest.raises(ValueError, match="ivf_centroids"): - create_vector_index(uri, fragment_ids=[4, 5], **build) - - -def test_backfill_with_shared_model_then_optimize_indices_merges(tmp_path: Path) -> None: - """A same-model backfill stays mergeable by optimize_indices (shared-centroid invariant).""" - uri, vectors = _make_vector_dataset(tmp_path / "backfill.lance") - centroids = ( - IndicesBuilder(lance.dataset(uri), "vector") - .train_ivf(num_partitions=4, distance_type="l2", sample_rate=8) - .centroids - ) - build = dict(column="vector", index_type="IVF_FLAT", ivf_centroids=centroids, name="v_idx") - - create_vector_index(uri, fragment_ids=[0, 1, 2, 3], **build) - assert _covered_fragments(uri, "v_idx") == {0, 1, 2, 3} - - # Backfill appends coverage for the remaining fragments. - create_vector_index(uri, fragment_ids=[1, 4, 5, 6, 7], **build) - assert _covered_fragments(uri, "v_idx") == {0, 1, 2, 3, 4, 5, 6, 7} - - # Segments built from one model must remain mergeable. - optimize_indices(uri, indices=["v_idx"], num_indices_to_merge=2) - assert _covered_fragments(uri, "v_idx") == {0, 1, 2, 3, 4, 5, 6, 7} - assert 7 in lance.dataset(uri).to_table(nearest={"column": "vector", "q": vectors[7], "k": 5})["id"].to_pylist() - - # A fully covered request is a no-op: version stays unchanged. - version_before = lance.dataset(uri).version - create_vector_index(uri, fragment_ids=[2, 3], **build) - assert lance.dataset(uri).version == version_before - - def test_replace_default_false_rejects_existing_name(tmp_path: Path) -> 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") @@ -471,12 +431,6 @@ def test_invalid_arguments_raise(tmp_path: Path) -> None: with pytest.raises(ValueError, match="sample_rate must be positive"): create_vector_index(uri, column="vector", index_type="IVF_FLAT", sample_rate=0) - with pytest.raises(ValueError, match="fragment_ids must be a non-empty list"): - create_vector_index(uri, column="vector", index_type="IVF_FLAT", fragment_ids=[]) - - with pytest.raises(ValueError, match="do not exist in the dataset"): - create_vector_index(uri, column="vector", index_type="IVF_FLAT", fragment_ids=[99]) - # Argument validation must fire before training work, even when training # would also fail on this small dataset with a large sample_rate. with pytest.raises(ValueError, match="fragment_group_size must be positive"): @@ -510,48 +464,6 @@ def test_existing_segment_uuid_is_rejected_without_overwriting_index(tmp_path: P assert after.to_table(nearest=query).to_pydict() == expected -@pytest.mark.parametrize("initial_metric,append_metric", [("L2", "cosine"), ("cosine", None)]) -def test_backfill_rejects_a_different_metric(tmp_path: Path, initial_metric: str, append_metric: str | None) -> None: - """Segments must use comparable distances, including when backfill defaults to L2.""" - uri = str(tmp_path / "mixed_metric.lance") - vectors = np.array([[0, 1], [100, 0]], dtype=np.float32) - lance.write_dataset( - pa.table({"id": [0, 1], "vector": pa.FixedSizeListArray.from_arrays(vectors.reshape(-1), 2)}), - uri, - max_rows_per_file=1, - ) - centroids = pa.FixedSizeListArray.from_arrays(pa.array([1, 0], type=pa.float32()), 2) - build = dict(column="vector", index_type="IVF_FLAT", ivf_centroids=centroids) - create_vector_index(uri, metric=initial_metric, fragment_ids=[0], **build) - before = lance.dataset(uri) - query = {"column": "vector", "q": [1, 0], "k": 2, "nprobes": 1, "metric": initial_metric} - expected = before.to_table(nearest=query).to_pydict() - - append_kwargs = {} if append_metric is None else {"metric": append_metric} - with pytest.raises(ValueError, match="[Mm]etric"): - create_vector_index(uri, fragment_ids=[1], **build, **append_kwargs) - - after = lance.dataset(uri) - assert after.version == before.version - assert _covered_fragments(uri, "vector_idx") == {0} - assert after.to_table(nearest=query).to_pydict() == expected - - -@pytest.mark.parametrize("append_metric", ["l2", "L2", "euclidean", "EUCLIDEAN"]) -def test_backfill_accepts_equivalent_metric_spellings(tmp_path: Path, append_metric: str) -> None: - """Lance's Euclidean alias and case variations do not imply a changed metric.""" - uri, _ = _make_vector_dataset(tmp_path / "metric_alias.lance", num_rows=240, rows_per_file=120) - centroids = ( - IndicesBuilder(lance.dataset(uri), "vector") - .train_ivf(num_partitions=4, distance_type="l2", sample_rate=8) - .centroids - ) - build = dict(column="vector", index_type="IVF_FLAT", ivf_centroids=centroids) - create_vector_index(uri, metric="L2", fragment_ids=[0], **build) - create_vector_index(uri, metric=append_metric, fragment_ids=[1], **build) - assert _covered_fragments(uri, "vector_idx") == {0, 1} - - def test_hnsw_build_parameters_are_forwarded(tmp_path: Path) -> None: """Real HNSW kwargs are implemented by Lance even though its Python signature uses **kwargs.""" uri, vectors = _make_vector_dataset(tmp_path / "hnsw_kwargs.lance", num_rows=240, rows_per_file=60) @@ -643,7 +555,7 @@ def train_with_recorded_sample(builder, *args, **kwargs): ) -def test_supplied_four_bit_pq_model_backfills_with_original_configuration(tmp_path: Path) -> None: +def test_supplied_four_bit_pq_model_builds_with_original_configuration(tmp_path: Path) -> None: uri, vectors = _make_vector_dataset(tmp_path / "supplied_bits.lance", num_rows=256, rows_per_file=128) builder = IndicesBuilder(lance.dataset(uri), "vector") ivf = builder.train_ivf(num_partitions=2, sample_rate=8) @@ -656,8 +568,7 @@ def test_supplied_four_bit_pq_model_backfills_with_original_configuration(tmp_pa num_sub_vectors=2, num_bits=4, ) - create_vector_index(uri, fragment_ids=[0], **build) - create_vector_index(uri, fragment_ids=[1], **build) + create_vector_index(uri, fragment_group_size=1, **build) ds = lance.dataset(uri) segments = ds.stats.index_stats("vector_idx")["indices"] assert len(segments) == 2 @@ -675,3 +586,63 @@ def test_non_pq_bit_width_is_not_silently_ignored(tmp_path: Path, index_type: st create_vector_index(uri, column="vector", index_type=index_type, num_bits=4) assert lance.dataset(uri).version == before.version assert not lance.dataset(uri).describe_indices() + + +@pytest.mark.parametrize("fragment_ids", [None, [], [0]]) +def test_fragment_subset_argument_is_rejected( + tmp_path: Path, fragment_ids: list[int] | None, monkeypatch: pytest.MonkeyPatch +) -> None: + """Creation always covers the dataset; even explicitly passing None is a removed argument.""" + uri, _ = _make_vector_dataset(tmp_path / "removed_subset.lance", num_rows=240, rows_per_file=60) + before = lance.dataset(uri) + + def unexpected_training(*args, **kwargs): + raise AssertionError("removed fragment_ids argument reached training") + + monkeypatch.setattr(IndicesBuilder, "train_ivf", unexpected_training) + with pytest.raises(TypeError, match="fragment_ids.*optimize_indices"): + create_vector_index(uri, column="vector", index_type="IVF_FLAT", fragment_ids=fragment_ids) + assert lance.dataset(uri).version == before.version + assert not lance.dataset(uri).describe_indices() + + +@pytest.mark.parametrize("index_type", ALL_VECTOR_INDEX_TYPES) +def test_optimize_indices_covers_appended_vectors(tmp_path: Path, index_type: str) -> None: + """Incremental coverage belongs to maintenance, while creation covers every existing fragment.""" + uri, _ = _make_vector_dataset(tmp_path / "incremental.lance", num_rows=512, rows_per_file=128) + create_vector_index( + uri, + column="vector", + index_type=index_type, + num_partitions=2, + num_sub_vectors=2, + sample_rate=2, + fragment_group_size=2, + ) + original = lance.dataset(uri) + old_fragment_ids = {f.fragment_id for f in original.get_fragments()} + old_segment_ids = {segment.uuid for segment in original.describe_indices()[0].segments} + assert _covered_fragments(uri, "vector_idx") == old_fragment_ids + + more = np.random.default_rng(99).standard_normal((256, 8)).astype(np.float32) + lance.write_dataset( + pa.table({"id": list(range(512, 768)), "vector": pa.FixedSizeListArray.from_arrays(more.reshape(-1), 8)}), + uri, + mode="append", + max_rows_per_file=128, + ) + appended = lance.dataset(uri) + assert _covered_fragments(uri, "vector_idx") == old_fragment_ids + assert len(appended.get_fragments()) == 6 + + # Keep old segments separate, including SQ segments whose quantizers may differ. + optimize_indices(uri, indices=["vector_idx"], num_indices_to_merge=0) + maintained = lance.dataset(uri) + assert maintained.version == appended.version + 1 + assert _covered_fragments(uri, "vector_idx") == {f.fragment_id for f in maintained.get_fragments()} + assert maintained.describe_indices()[0].num_rows_indexed == 768 + assert old_segment_ids <= {segment.uuid for segment in maintained.describe_indices()[0].segments} + result = maintained.to_table(nearest={"column": "vector", "q": more[7], "k": 5, "nprobes": 2, "refine_factor": 20}) + assert result.num_rows == 5 + assert 519 in result["id"].to_pylist() + assert np.isfinite(result["_distance"].to_numpy()).all() diff --git a/tests/io/lance/test_lance_vector_index_concurrency.py b/tests/io/lance/test_lance_vector_index_concurrency.py index 275ebd6..654a6a4 100644 --- a/tests/io/lance/test_lance_vector_index_concurrency.py +++ b/tests/io/lance/test_lance_vector_index_concurrency.py @@ -8,7 +8,7 @@ from lance.commit import CommitConflictError from daft.dependencies import pa -from daft_lance import create_vector_index, lance_vector_index +from daft_lance import create_vector_index, lance_vector_index, optimize_indices from daft_lance.namespace import DatasetOpenContext @@ -76,11 +76,9 @@ def validate_after_competing_commit(dataset, index_metas, expected_fragment_ids) assert 7 in dataset.to_table(nearest={"column": "vector", "q": vectors[7], "k": 5})["id"].to_pylist() -def test_backfill_rejects_concurrent_rebuild_with_changed_metric( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch -) -> None: +def test_replace_true_replaces_concurrently_rebuilt_index(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: uri, vectors, centroids = _dataset(tmp_path) - create_vector_index(uri, column="vector", index_type="IVF_FLAT", ivf_centroids=centroids, fragment_ids=[0, 1, 2]) + create_vector_index(uri, column="vector", index_type="IVF_FLAT", ivf_centroids=centroids) original_open = DatasetOpenContext.open_latest competing = {} @@ -100,15 +98,15 @@ def open_after_rebuild(self): return original_open(self) monkeypatch.setattr(DatasetOpenContext, "open_latest", open_after_rebuild) - with pytest.raises(ValueError, match="changed during the build"): - create_vector_index( - uri, column="vector", index_type="IVF_FLAT", ivf_centroids=centroids, fragment_ids=[3, 4, 5] - ) + create_vector_index( + uri, column="vector", index_type="IVF_FLAT", ivf_centroids=centroids, replace=True, fragment_group_size=2 + ) dataset = lance.dataset(uri) - assert dataset.version == competing["version"] - assert _segments(dataset) == competing["segments"] - assert dataset.index_statistics("vector_idx")["segments"][0]["metric_type"] == "cosine" + assert dataset.version == competing["version"] + 1 + assert _segments(dataset).isdisjoint(competing["segments"]) + assert len(_segments(dataset)) == 3 + assert all(segment["metric_type"] == "l2" for segment in dataset.index_statistics("vector_idx")["segments"]) assert 7 in dataset.to_table(nearest={"column": "vector", "q": vectors[7], "k": 5})["id"].to_pylist() @@ -131,3 +129,54 @@ def validate_after_append(dataset, index_metas, expected_fragment_ids): assert dataset.describe_indices()[0].num_rows_indexed == 240 assert len(_segments(dataset)) == 3 assert 7 in dataset.to_table(nearest={"column": "vector", "q": vectors[7], "k": 5})["id"].to_pylist() + + +def test_replace_rejects_concurrently_indexed_appended_fragments( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """A replacement must not leave an old-metric segment outside its pinned build plan.""" + uri, vectors, centroids = _dataset(tmp_path) + create_vector_index(uri, column="vector", index_type="IVF_FLAT", ivf_centroids=centroids) + original_open = DatasetOpenContext.open_latest + more = np.random.default_rng(31).standard_normal((40, 8)).astype(np.float32) + 10 + competing = {} + queries = [vectors[7], more[7]] + + def open_after_incremental_maintenance(self): + lance.write_dataset( + pa.table({"id": range(240, 280), "vector": pa.FixedSizeListArray.from_arrays(more.ravel(), 8)}), + uri, + mode="append", + ) + monkeypatch.setattr(DatasetOpenContext, "open_latest", original_open) + optimize_indices(uri, indices=["vector_idx"], num_indices_to_merge=0) + dataset = lance.dataset(uri) + competing["version"] = dataset.version + competing["segments"] = _segments(dataset) + competing["results"] = [ + dataset.to_table(nearest={"column": "vector", "q": query, "k": 5, "nprobes": 4}).to_pydict() + for query in queries + ] + return original_open(self) + + monkeypatch.setattr(DatasetOpenContext, "open_latest", open_after_incremental_maintenance) + with pytest.raises(ValueError, match="coverage outside the build plan"): + create_vector_index( + uri, + column="vector", + index_type="IVF_FLAT", + ivf_centroids=centroids, + metric="cosine", + replace=True, + fragment_group_size=2, + ) + + dataset = lance.dataset(uri) + assert dataset.version == competing["version"] + assert _segments(dataset) == competing["segments"] + assert dataset.describe_indices()[0].num_rows_indexed == 280 + assert all(segment["metric_type"] == "l2" for segment in dataset.index_statistics("vector_idx")["segments"]) + for query, expected in zip(queries, competing["results"]): + assert dataset.to_table(nearest={"column": "vector", "q": query, "k": 5, "nprobes": 4}).to_pydict() == expected + assert 7 in competing["results"][0]["id"] + assert 247 in competing["results"][1]["id"] From 754f9a467c8b5703cce47f69c451722c7ea146b3 Mon Sep 17 00:00:00 2001 From: wangzheyan Date: Sat, 10 Oct 2026 15:58:38 +0800 Subject: [PATCH 11/11] refactor: train vector index models internally without sample adjustment --- README.md | 36 +-- daft_lance/_lance.py | 39 ++-- daft_lance/lance_vector_index.py | 101 +++----- tests/io/lance/test_lance_vector_index.py | 216 +++++++----------- .../test_lance_vector_index_concurrency.py | 47 ++-- 5 files changed, 169 insertions(+), 270 deletions(-) diff --git a/README.md b/README.md index 7acc93d..45b573c 100644 --- a/README.md +++ b/README.md @@ -109,8 +109,7 @@ returns the updated dataset; for a distributed rebuild use Vector indexes (`IVF_FLAT`, `IVF_PQ`, `IVF_SQ`, `IVF_HNSW_FLAT`, `IVF_HNSW_PQ`, `IVF_HNSW_SQ`) 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`, unless pre-trained -`ivf_centroids`/`pq_codebook` are supplied — then Daft workers each build one +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. @@ -140,9 +139,9 @@ snapshot, full replacement fails and must be retried to include those fragments. This prevents old model segments from surviving a full rebuild. **Breaking change — vector index migration:** `create_vector_index` no longer -accepts `fragment_ids`, partial builds, or backfill/no-op requests. For an initial -build, remove `fragment_ids` to create coverage for the full dataset. Once the -index exists, replace calls that used `create_vector_index(..., fragment_ids=[...])` +accepts `fragment_ids`, `ivf_centroids`, or `pq_codebook`, partial builds, or +backfill/no-op requests. For an initial build, remove `fragment_ids` to create +coverage for the full dataset. Once the index exists, replace calls that used `create_vector_index(..., fragment_ids=[...])` to cover appended data with: ```python @@ -157,12 +156,18 @@ fragment list nor caller-supplied models. To retrain or replace the index, use semantics are unchanged. HNSW builds accept `m`, `max_level`, and `ef_construction`. Each worker generates -its own segment UUID, so `index_uuid` cannot be supplied. For pre-trained IVF -models, pass `ivf_centroids`; `ivf_centroids_file` is not supported by this API. -PQ builds accept `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. When -supplying a `pq_codebook`, also pass its original `num_sub_vectors` and -`num_bits`; a bare codebook array does not carry the sub-vector count. +its own segment UUID, so `index_uuid` cannot be supplied. +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. + +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 +`256 * sample_rate` rows. The requested value is passed unchanged to Lance. +Automatic sample reduction has been removed. If training reports insufficient +data, lower the training parameters or provide more data. Calls that previously +provided pre-trained models should omit them and let the build train its models. The SQ variants (`IVF_SQ`, `IVF_HNSW_SQ`) support multi-segment builds. Each segment uses its own SQ quantization metadata @@ -175,12 +180,9 @@ an existing index, rebuild with that grouping and `replace=True`. Use `optimize_indices(..., num_indices_to_merge=0)` to add coverage without merging incompatible old SQ segments. -Training samples `sample_rate` rows per IVF partition (and per PQ centroid) -and runs in the coordinator process, so its memory footprint grows with -`num_partitions * sample_rate * dimension`; `sample_rate` is clamped down -automatically (with a warning) to what the dataset size supports -(`num_partitions * sample_rate` rows, `2**num_bits * sample_rate` for the PQ -codebook). +Training runs in the coordinator process. Its memory footprint grows with +`num_partitions * sample_rate * dimension` for IVF and +`2**num_bits * sample_rate * dimension` for PQ. ### Column Merging diff --git a/daft_lance/_lance.py b/daft_lance/_lance.py index 63352ab..f383c65 100644 --- a/daft_lance/_lance.py +++ b/daft_lance/_lance.py @@ -700,8 +700,6 @@ def create_vector_index( num_sub_vectors: int | None = None, num_bits: int = 8, sample_rate: int = 256, - ivf_centroids: Any | None = None, - pq_codebook: Any | None = None, storage_options: dict[str, Any] | None = None, version: int | str | None = None, asof: str | None = None, @@ -718,8 +716,7 @@ 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, unless pre-trained - ``ivf_centroids``/``pq_codebook`` are supplied), fragment batches are + ``IndicesBuilder.train_pq`` for PQ variants), 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 @@ -728,7 +725,8 @@ def create_vector_index( step is needed. Incremental indexing of appended data belongs to ``optimize_indices``. - This API does not accept ``fragment_ids`` and does not skip covered data. + This API does not accept ``fragment_ids``, ``ivf_centroids`` or + ``pq_codebook``. Models are always trained by the driver. Args: uri: The URI of the Lance table (supports remote URLs to object stores such as `s3://` or `gs://`) @@ -760,24 +758,16 @@ def create_vector_index( ``num_partitions`` controls Daft's fragment repartitioning; build 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. Required - when supplying a ``pq_codebook``; the array does not carry this metadata. + 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. - Supplied PQ codebooks must use the matching bit width. - sample_rate: Rows sampled per IVF partition (and per PQ centroid) during - training. Training requires at least ``num_partitions * sample_rate`` - rows (2**num_bits * sample_rate for the PQ codebook); the value is - automatically clamped down (with a warning) to what the dataset - size supports, at the cost of a smaller training sample on small - datasets. Lance still rejects data below the minimum training size - (2**num_bits rows for PQ); clamping does not remove that requirement. - ivf_centroids: Pre-trained IVF centroids (a pyarrow array); skips IVF - training. Supplying the same centroids to every build is how - independently built segments share one model. - pq_codebook: Pre-trained PQ codebook (PQ variants only); skips PQ - training. Supply its original ``num_sub_vectors`` and ``num_bits``. + 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 + requires 256 * sample_rate rows. The value is passed unchanged to + Lance, which rejects insufficient training data. Adjust the training + parameters or provide more data if training fails. storage_options: Storage options for the dataset. version: Version of the dataset to use. asof: Timestamp to use for time travel queries. @@ -793,7 +783,7 @@ def create_vector_index( **kwargs: Additional keyword arguments forwarded to Lance's index segment creation API, including HNSW ``m``, ``max_level`` and ``ef_construction``. ``index_uuid`` is managed per segment. ``ivf_centroids_file`` is not - supported; supply shared centroids through ``ivf_centroids``. + supported; the driver trains the shared models internally. Unknown arguments raise ``TypeError`` — Lance's segment build silently ignores them, so misspellings must not pass silently here. @@ -803,9 +793,8 @@ def create_vector_index( Raises: ValueError: If input parameters are invalid (e.g. empty column name, non-existent column, unsupported index type, an existing index - name with ``replace=False``, a supplied codebook without its original - sub-vector count, or a dataset too small to train an IVF model even - at ``sample_rate=1`` — Lance raises its training ``ValueError``) + name with ``replace=False``, or insufficient data for the requested + training parameters — Lance raises its training ``ValueError``) TypeError: If column type is incompatible with the chosen ``index_type``, or an unknown keyword argument is passed RuntimeError: If index building fails (e.g. version compatibility issues, commit failures) @@ -861,8 +850,6 @@ def create_vector_index( num_sub_vectors=num_sub_vectors, num_bits=num_bits, sample_rate=sample_rate, - ivf_centroids=ivf_centroids, - pq_codebook=pq_codebook, fragment_group_size=fragment_group_size, max_concurrency=max_concurrency, **kwargs, diff --git a/daft_lance/lance_vector_index.py b/daft_lance/lance_vector_index.py index 821bf1f..fa5ec2d 100644 --- a/daft_lance/lance_vector_index.py +++ b/daft_lance/lance_vector_index.py @@ -2,7 +2,6 @@ import inspect import logging -import math import pickle from functools import lru_cache from typing import TYPE_CHECKING, Any @@ -15,7 +14,6 @@ import lance -from daft.dependencies import pa from daft_lance.lance_scalar_index import ( _validate_segments_against_manifest, ) @@ -73,6 +71,11 @@ 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"}) + if removed_models: + raise TypeError( + f"create_vector_index does not accept {removed_models}; models are trained by the build workflow." + ) if "segmented" in kwargs: raise TypeError( "The 'segmented' parameter was removed: the distributed segment-index " @@ -80,8 +83,6 @@ def _validate_worker_kwargs(kwargs: dict[str, Any]) -> None: ) if "index_uuid" in kwargs: raise TypeError("index_uuid is managed per segment; each worker must generate a unique UUID.") - if "ivf_centroids_file" in kwargs: - raise TypeError("ivf_centroids_file is not supported; pass shared centroids through ivf_centroids.") unknown = sorted(set(kwargs) - _accepted_worker_kwargs()) if not unknown: return @@ -171,8 +172,6 @@ def create_vector_index_internal( num_sub_vectors: int | None = None, num_bits: int = 8, sample_rate: int = 256, - ivf_centroids: pa.Array[Any] | None = None, - pq_codebook: pa.Array[Any] | None = None, fragment_group_size: int | None = None, max_concurrency: int | None = None, **kwargs: Any, @@ -187,14 +186,8 @@ def create_vector_index_internal( 1. The driver trains the global model once — IVF centroids via ``IndicesBuilder.train_ivf`` and, for PQ variants, the PQ codebook via - ``IndicesBuilder.train_pq`` — unless the caller supplies pre-trained - ``ivf_centroids`` / ``pq_codebook``. One shared model across segments is - what lets independently built segments commit as one logical index. - ``sample_rate`` is clamped down to what the dataset size supports - (training needs ``num_partitions * sample_rate`` rows, and the PQ - codebook needs ``2**num_bits * sample_rate``). ``num_partitions`` is - derived from supplied centroids. A supplied codebook requires explicit - ``num_sub_vectors`` because the array does not carry that metadata. + ``IndicesBuilder.train_pq``. The user's ``sample_rate`` is passed + unchanged; Lance validates that enough training data is available. 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. @@ -228,8 +221,6 @@ def create_vector_index_internal( # Reject kwargs Lance's segment build would silently swallow (misspelled # or unsupported parameters must not build a misconfigured index). _validate_worker_kwargs(kwargs) - if index_type in _PQ_INDEX_TYPES and pq_codebook is not None and num_sub_vectors is None: - raise ValueError("A supplied pq_codebook requires its original num_sub_vectors; it cannot be inferred.") # Validate column exists; whether it is a vector column is Lance's rule, # enforced by the training and build APIs with clear errors. @@ -263,76 +254,36 @@ def create_vector_index_internal( # Phase 1: train the global model once on the driver so all segments # share the same centroids (and codebook) and commit as one logical index. - num_rows = lance_ds.count_rows() - effective_sample_rate = sample_rate - needs_ivf_training = ivf_centroids is None - needs_pq_training = index_type in _PQ_INDEX_TYPES and pq_codebook is None - if needs_ivf_training or needs_pq_training: - # Clamp the sample rate to what this dataset can support, mirroring - # Lance's own training requirements: num_partitions * sample_rate rows - # for IVF (Lance derives num_partitions as sqrt(num_rows) when None) - # and 2**num_bits * sample_rate rows for the PQ codebook. - caps = [sample_rate] - if needs_ivf_training: - effective_partitions = num_partitions if num_partitions is not None else max(1, round(math.sqrt(num_rows))) - caps.append(max(1, num_rows // effective_partitions)) - if needs_pq_training: - caps.append(max(1, num_rows // (2**num_bits))) - clamped = min(caps) - if clamped < sample_rate: - logger.warning( - "sample_rate %d exceeds what the dataset supports (%d rows); clamping to %d", - sample_rate, - num_rows, - clamped, - ) - effective_sample_rate = clamped - builder = lance.indices.IndicesBuilder(lance_ds, column) - ivf_model: lance.indices.IvfModel | None = None - if needs_ivf_training: - logger.info( - "Phase 1: training IVF centroids (index_type=%s, metric=%s, num_partitions=%s, sample_rate=%d)", - index_type, - metric, - num_partitions, - effective_sample_rate, - ) - ivf_model = builder.train_ivf( - num_partitions=num_partitions, - distance_type=metric.lower(), - sample_rate=effective_sample_rate, - ) - ivf_centroids = ivf_model.centroids - num_partitions = ivf_model.num_partitions - logger.info("IVF training completed: num_partitions=%d", num_partitions) + logger.info( + "Phase 1: training IVF centroids (index_type=%s, metric=%s, num_partitions=%s, sample_rate=%d)", + index_type, + metric, + num_partitions, + sample_rate, + ) + ivf_model = builder.train_ivf( + num_partitions=num_partitions, + distance_type=metric.lower(), + sample_rate=sample_rate, + ) + ivf_centroids = ivf_model.centroids + num_partitions = ivf_model.num_partitions + logger.info("IVF training completed: num_partitions=%d", num_partitions) - if needs_pq_training: - logger.info( - "Phase 1: training PQ codebook (num_sub_vectors=%s, sample_rate=%d)", num_sub_vectors, effective_sample_rate - ) - if ivf_model is None: - # Caller-supplied centroids: wrap them so train_pq can partition - # its samples with the same model the segments will be built with. - # needs_ivf_training is False exactly when ivf_centroids was - # supplied, so the None case cannot occur in practice. - if ivf_centroids is None: - raise ValueError("PQ codebook training requires IVF centroids") - ivf_model = lance.indices.IvfModel(ivf_centroids, metric.lower()) + pq_codebook = None + if index_type in _PQ_INDEX_TYPES: + logger.info("Phase 1: training PQ codebook (num_sub_vectors=%s, sample_rate=%d)", num_sub_vectors, sample_rate) pq_model = builder.train_pq( ivf_model, num_subvectors=num_sub_vectors, num_bits=num_bits, - sample_rate=effective_sample_rate, + sample_rate=sample_rate, ) pq_codebook = pq_model.codebook num_sub_vectors = pq_model.num_subvectors logger.info("PQ training completed: num_sub_vectors=%d", num_sub_vectors) - # Derive the partition count from supplied centroids when not given. - if num_partitions is None and ivf_centroids is not None: - num_partitions = len(ivf_centroids) - model_kwargs: dict[str, Any] = { "metric": metric, "ivf_centroids": ivf_centroids, diff --git a/tests/io/lance/test_lance_vector_index.py b/tests/io/lance/test_lance_vector_index.py index 92105f4..46753f9 100644 --- a/tests/io/lance/test_lance_vector_index.py +++ b/tests/io/lance/test_lance_vector_index.py @@ -149,12 +149,7 @@ def test_sq_multi_segment_build_and_failed_merge_preserve_index(tmp_path: Path, uri, max_rows_per_file=256, ) - centroids = ( - IndicesBuilder(lance.dataset(uri), "vector") - .train_ivf(num_partitions=4, distance_type="l2", sample_rate=8) - .centroids - ) - build = dict(column="vector", index_type=index_type, ivf_centroids=centroids, name="v_idx") + build = dict(column="vector", index_type=index_type, num_partitions=4, sample_rate=8, name="v_idx") create_vector_index(uri, fragment_group_size=4, **build) before = lance.dataset(uri) @@ -244,79 +239,6 @@ def test_cosine_metric_orders_by_angle_not_distance(tmp_path: Path) -> None: assert got.isdisjoint(l2_truth) -def test_small_dataset_default_sample_rate_is_clamped(tmp_path: Path) -> None: - """Default sample_rate=256 would demand 65k rows; it must clamp instead of failing.""" - uri, vectors = _make_vector_dataset(tmp_path / "clamp.lance", num_rows=8192, rows_per_file=512, seed=17) - - # No explicit sample_rate: training needs 16*256 (IVF) and 256*256 (PQ) rows - # unclamped — far more than the dataset has. - create_vector_index(uri, column="vector", index_type="IVF_PQ", num_partitions=16, fragment_group_size=2) - - desc = lance.dataset(uri).describe_indices()[0] - assert desc.index_type == "IVF_PQ" - assert desc.num_rows_indexed == 8192 - # The clamp makes the build succeed; recall quality under a heavily clamped - # sample is the lossy-compression tradeoff, guarded by the PQ-specific - # recall tests, so only assert the index answers queries here. - results = lance.dataset(uri).to_table(nearest={"column": "vector", "q": vectors[7], "k": 5}) - assert results.num_rows == 5 - - -def test_pretrained_centroids_are_used(tmp_path: Path) -> None: - """Caller-supplied ivf_centroids skip driver-side training and still build.""" - uri, _ = _make_vector_dataset(tmp_path / "pre.lance", num_rows=240, rows_per_file=40, seed=3) - - centroids = ( - IndicesBuilder(lance.dataset(uri), "vector") - .train_ivf(num_partitions=4, distance_type="l2", sample_rate=32) - .centroids - ) - - create_vector_index( - uri, - column="vector", - index_type="IVF_FLAT", - ivf_centroids=centroids, - fragment_group_size=2, - ) - - desc = lance.dataset(uri).describe_indices()[0] - assert desc.index_type == "IVF_FLAT" - assert desc.num_rows_indexed == 240 - - -def test_pq_codebook_requires_original_num_sub_vectors(tmp_path: Path) -> None: - """A bare codebook cannot reveal how the original vectors were split.""" - uri, _ = _make_vector_dataset(tmp_path / "codebook.lance", seed=19) - ds = lance.dataset(uri) - builder = IndicesBuilder(ds, "vector") - ivf = builder.train_ivf(num_partitions=4, distance_type="l2", sample_rate=32) - pq = builder.train_pq(ivf, num_subvectors=2, sample_rate=8) - - with pytest.raises(ValueError, match="original num_sub_vectors"): - create_vector_index( - uri, column="vector", index_type="IVF_PQ", ivf_centroids=ivf.centroids, pq_codebook=pq.codebook - ) - - create_vector_index( - uri, - column="vector", - index_type="IVF_PQ", - ivf_centroids=ivf.centroids, - pq_codebook=pq.codebook, - num_sub_vectors=2, - fragment_group_size=4, - ) - - desc = lance.dataset(uri).describe_indices()[0] - assert desc.index_type == "IVF_PQ" - assert desc.num_rows_indexed == 2048 - assert all( - segment["sub_index"]["num_sub_vectors"] == 2 - for segment in lance.dataset(uri).stats.index_stats("vector_idx")["indices"] - ) - - def test_replace_default_false_rejects_existing_name(tmp_path: Path) -> 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") @@ -413,7 +335,9 @@ def test_unknown_kwargs_are_rejected(tmp_path: Path) -> None: create_vector_index(uri, column="vector", index_type="IVF_FLAT", segmented=True) # A real build kwarg still passes through. - create_vector_index(uri, column="vector", index_type="IVF_FLAT", num_partitions=4, target_partition_size=1024) + create_vector_index( + uri, column="vector", index_type="IVF_FLAT", num_partitions=4, sample_rate=32, target_partition_size=1024 + ) def test_invalid_arguments_raise(tmp_path: Path) -> None: @@ -485,7 +409,9 @@ def test_hnsw_build_parameters_are_forwarded(tmp_path: Path) -> None: assert params["m"] == 4 assert params["max_level"] == 3 assert params["ef_construction"] == 24 - assert 7 in ds.to_table(nearest={"column": "vector", "q": vectors[7], "k": 5, "nprobes": 4})["id"].to_pylist() + result = ds.to_table(nearest={"column": "vector", "q": vectors[7], "k": 5, "nprobes": 4}) + assert result.num_rows == 5 + assert np.isfinite(result["_distance"].to_numpy()).all() def test_centroids_file_is_rejected_before_segment_build(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: @@ -506,7 +432,7 @@ def unexpected_training(*args, **kwargs): @pytest.mark.parametrize("index_type", ["IVF_PQ", "IVF_HNSW_PQ"]) @pytest.mark.parametrize("num_bits", [4, 8]) def test_pq_bit_width_is_shared_by_all_segments(tmp_path: Path, index_type: str, num_bits: int) -> None: - uri, vectors = _make_vector_dataset(tmp_path / "bits.lance", num_rows=512, rows_per_file=128) + uri, vectors = _make_vector_dataset(tmp_path / "bits.lance", num_rows=2048, rows_per_file=512) create_vector_index( uri, column="vector", @@ -527,57 +453,6 @@ def test_pq_bit_width_is_shared_by_all_segments(tmp_path: Path, index_type: str, assert np.isfinite(result["_distance"].to_numpy()).all() -def test_four_bit_pq_sampling_uses_sixteen_centroids(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: - uri, _ = _make_vector_dataset(tmp_path / "sampling_bits.lance", num_rows=128, rows_per_file=64) - centroids = pa.FixedSizeListArray.from_arrays(pa.array([0.0] * 8, type=pa.float32()), 8) - original_train = IndicesBuilder.train_pq - calls = [] - - def train_with_recorded_sample(builder, *args, **kwargs): - calls.append((kwargs["num_bits"], kwargs["sample_rate"])) - return original_train(builder, *args, **kwargs) - - monkeypatch.setattr(IndicesBuilder, "train_pq", train_with_recorded_sample) - create_vector_index( - uri, - column="vector", - index_type="IVF_PQ", - ivf_centroids=centroids, - num_sub_vectors=2, - num_bits=4, - sample_rate=8, - fragment_group_size=1, - ) - # 16 * 8 samples fit in 128 rows; no IVF training or 256-entry cap applies. - assert calls == [(4, 8)] - assert all( - segment["sub_index"]["nbits"] == 4 for segment in lance.dataset(uri).stats.index_stats("vector_idx")["indices"] - ) - - -def test_supplied_four_bit_pq_model_builds_with_original_configuration(tmp_path: Path) -> None: - uri, vectors = _make_vector_dataset(tmp_path / "supplied_bits.lance", num_rows=256, rows_per_file=128) - builder = IndicesBuilder(lance.dataset(uri), "vector") - ivf = builder.train_ivf(num_partitions=2, sample_rate=8) - pq = builder.train_pq(ivf, num_subvectors=2, num_bits=4, sample_rate=8) - build = dict( - column="vector", - index_type="IVF_PQ", - ivf_centroids=ivf.centroids, - pq_codebook=pq.codebook, - num_sub_vectors=2, - num_bits=4, - ) - create_vector_index(uri, fragment_group_size=1, **build) - ds = lance.dataset(uri) - segments = ds.stats.index_stats("vector_idx")["indices"] - assert len(segments) == 2 - assert all(segment["sub_index"]["nbits"] == 4 for segment in segments) - assert all(segment["sub_index"]["num_sub_vectors"] == 2 for segment in segments) - assert _covered_fragments(uri, "vector_idx") == {0, 1} - assert ds.to_table(nearest={"column": "vector", "q": vectors[7], "k": 5, "nprobes": 2}).num_rows == 5 - - @pytest.mark.parametrize("index_type", ["IVF_FLAT", "IVF_SQ", "IVF_HNSW_FLAT", "IVF_HNSW_SQ"]) def test_non_pq_bit_width_is_not_silently_ignored(tmp_path: Path, index_type: str) -> None: uri, _ = _make_vector_dataset(tmp_path / "non_pq_bits.lance", num_rows=128, rows_per_file=64) @@ -646,3 +521,78 @@ def test_optimize_indices_covers_appended_vectors(tmp_path: Path, index_type: st assert result.num_rows == 5 assert 519 in result["id"].to_pylist() assert np.isfinite(result["_distance"].to_numpy()).all() + + +@pytest.mark.parametrize("argument", ["ivf_centroids", "pq_codebook"]) +@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 +) -> None: + """Every build trains its model; removed inputs must fail even when their value is None.""" + uri, _ = _make_vector_dataset(tmp_path / "removed_model.lance", num_rows=240, rows_per_file=60) + before = lance.dataset(uri) + + def unexpected_training(*args, **kwargs): + raise AssertionError("removed model argument reached training") + + monkeypatch.setattr(IndicesBuilder, "train_ivf", unexpected_training) + with pytest.raises(TypeError, match=argument): + create_vector_index(uri, column="vector", index_type="IVF_PQ", **{argument: value}) + assert lance.dataset(uri).version == before.version + assert not lance.dataset(uri).describe_indices() + + +def test_requested_sample_rate_reaches_real_ivf_and_pq_training( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """A complete PQ build uses the caller's sampling rate in both training stages.""" + uri, vectors = _make_vector_dataset(tmp_path / "sample_rate.lance", num_rows=2048, rows_per_file=512) + train_ivf = IndicesBuilder.train_ivf + train_pq = IndicesBuilder.train_pq + calls = [] + + def record_ivf(builder, *args, **kwargs): + calls.append(("IVF", kwargs["sample_rate"])) + return train_ivf(builder, *args, **kwargs) + + def record_pq(builder, *args, **kwargs): + calls.append(("PQ", kwargs["sample_rate"])) + return train_pq(builder, *args, **kwargs) + + monkeypatch.setattr(IndicesBuilder, "train_ivf", record_ivf) + monkeypatch.setattr(IndicesBuilder, "train_pq", record_pq) + create_vector_index( + uri, + column="vector", + index_type="IVF_PQ", + num_partitions=4, + num_sub_vectors=2, + sample_rate=8, + fragment_group_size=2, + ) + assert calls == [("IVF", 8), ("PQ", 8)] + ds = lance.dataset(uri) + assert len(ds.describe_indices()[0].segments) == 2 + assert ds.describe_indices()[0].num_rows_indexed == 2048 + assert ds.to_table(nearest={"column": "vector", "q": vectors[7], "k": 5, "nprobes": 4}).num_rows == 5 + + +@pytest.mark.parametrize("index_type,num_rows,sample_rate", [("IVF_FLAT", 240, 256), ("IVF_PQ", 512, 8)]) +def test_insufficient_training_rows_raise_lance_error( + tmp_path: Path, index_type: str, num_rows: int, sample_rate: int +) -> None: + """Insufficient samples fail instead of silently reducing the requested training quality.""" + uri, _ = _make_vector_dataset(tmp_path / "too_small.lance", num_rows=num_rows, rows_per_file=128) + before = lance.dataset(uri) + training_stage = "IVF centroids" if index_type == "IVF_FLAT" else "PQ" + with pytest.raises(ValueError, match=f"not enough rows.*{training_stage}"): + create_vector_index( + uri, + column="vector", + index_type=index_type, + num_partitions=2, + num_sub_vectors=2, + sample_rate=sample_rate, + ) + assert lance.dataset(uri).version == before.version + assert not lance.dataset(uri).describe_indices() diff --git a/tests/io/lance/test_lance_vector_index_concurrency.py b/tests/io/lance/test_lance_vector_index_concurrency.py index 654a6a4..eeeea45 100644 --- a/tests/io/lance/test_lance_vector_index_concurrency.py +++ b/tests/io/lance/test_lance_vector_index_concurrency.py @@ -20,10 +20,7 @@ def _dataset(tmp_path: Path): uri, max_rows_per_file=40, ) - centroids = ( - lance.indices.IndicesBuilder(lance.dataset(uri), "vector").train_ivf(num_partitions=4, sample_rate=8).centroids - ) - return uri, vectors, centroids + return uri, vectors def _segments(dataset): @@ -31,20 +28,22 @@ def _segments(dataset): def test_same_name_created_before_live_check_is_preserved(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: - uri, vectors, centroids = _dataset(tmp_path) + uri, vectors = _dataset(tmp_path) original_open = DatasetOpenContext.open_latest competing = {} def open_after_competing_commit(self): dataset = lance.dataset(uri) - dataset.create_index("vector", "IVF_FLAT", name="vector_idx", ivf_centroids=centroids, num_partitions=4) + dataset.create_index("vector", "IVF_FLAT", name="vector_idx", num_partitions=4) competing["version"] = dataset.version competing["segments"] = _segments(dataset) return original_open(self) monkeypatch.setattr(DatasetOpenContext, "open_latest", open_after_competing_commit) with pytest.raises(ValueError, match="changed during the build"): - create_vector_index(uri, column="vector", index_type="IVF_FLAT", ivf_centroids=centroids, fragment_group_size=2) + create_vector_index( + uri, column="vector", index_type="IVF_FLAT", num_partitions=4, sample_rate=8, fragment_group_size=2 + ) dataset = lance.dataset(uri) assert dataset.version == competing["version"] @@ -55,20 +54,22 @@ def open_after_competing_commit(self): def test_same_name_created_after_live_check_causes_commit_conflict( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: - uri, vectors, centroids = _dataset(tmp_path) + uri, vectors = _dataset(tmp_path) original_validate = lance_vector_index._validate_segments_against_manifest competing = {} def validate_after_competing_commit(dataset, index_metas, expected_fragment_ids): other = lance.dataset(uri) - other.create_index("vector", "IVF_FLAT", name="vector_idx", ivf_centroids=centroids, num_partitions=4) + other.create_index("vector", "IVF_FLAT", name="vector_idx", num_partitions=4) competing["version"] = other.version competing["segments"] = _segments(other) return original_validate(dataset, index_metas, expected_fragment_ids) monkeypatch.setattr(lance_vector_index, "_validate_segments_against_manifest", validate_after_competing_commit) with pytest.raises(CommitConflictError, match="CreateIndex"): - create_vector_index(uri, column="vector", index_type="IVF_FLAT", ivf_centroids=centroids, fragment_group_size=2) + create_vector_index( + uri, column="vector", index_type="IVF_FLAT", num_partitions=4, sample_rate=8, fragment_group_size=2 + ) dataset = lance.dataset(uri) assert dataset.version == competing["version"] @@ -77,8 +78,8 @@ def validate_after_competing_commit(dataset, index_metas, expected_fragment_ids) def test_replace_true_replaces_concurrently_rebuilt_index(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: - uri, vectors, centroids = _dataset(tmp_path) - create_vector_index(uri, column="vector", index_type="IVF_FLAT", ivf_centroids=centroids) + uri, vectors = _dataset(tmp_path) + create_vector_index(uri, column="vector", index_type="IVF_FLAT", num_partitions=4, sample_rate=8) original_open = DatasetOpenContext.open_latest competing = {} @@ -90,7 +91,6 @@ def open_after_rebuild(self): name="vector_idx", metric="cosine", replace=True, - ivf_centroids=centroids, num_partitions=4, ) competing["version"] = dataset.version @@ -99,7 +99,13 @@ def open_after_rebuild(self): monkeypatch.setattr(DatasetOpenContext, "open_latest", open_after_rebuild) create_vector_index( - uri, column="vector", index_type="IVF_FLAT", ivf_centroids=centroids, replace=True, fragment_group_size=2 + uri, + column="vector", + index_type="IVF_FLAT", + num_partitions=4, + sample_rate=8, + replace=True, + fragment_group_size=2, ) dataset = lance.dataset(uri) @@ -111,7 +117,7 @@ def open_after_rebuild(self): def test_unrelated_append_after_live_check_allows_index_commit(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: - uri, vectors, centroids = _dataset(tmp_path) + uri, vectors = _dataset(tmp_path) original_validate = lance_vector_index._validate_segments_against_manifest def validate_after_append(dataset, index_metas, expected_fragment_ids): @@ -121,7 +127,9 @@ def validate_after_append(dataset, index_metas, expected_fragment_ids): return original_validate(dataset, index_metas, expected_fragment_ids) monkeypatch.setattr(lance_vector_index, "_validate_segments_against_manifest", validate_after_append) - create_vector_index(uri, column="vector", index_type="IVF_FLAT", ivf_centroids=centroids, fragment_group_size=2) + create_vector_index( + uri, column="vector", index_type="IVF_FLAT", num_partitions=4, sample_rate=8, fragment_group_size=2 + ) dataset = lance.dataset(uri) assert dataset.version == 3 @@ -135,8 +143,8 @@ def test_replace_rejects_concurrently_indexed_appended_fragments( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: """A replacement must not leave an old-metric segment outside its pinned build plan.""" - uri, vectors, centroids = _dataset(tmp_path) - create_vector_index(uri, column="vector", index_type="IVF_FLAT", ivf_centroids=centroids) + uri, vectors = _dataset(tmp_path) + create_vector_index(uri, column="vector", index_type="IVF_FLAT", num_partitions=4, sample_rate=8) original_open = DatasetOpenContext.open_latest more = np.random.default_rng(31).standard_normal((40, 8)).astype(np.float32) + 10 competing = {} @@ -165,7 +173,8 @@ def open_after_incremental_maintenance(self): uri, column="vector", index_type="IVF_FLAT", - ivf_centroids=centroids, + num_partitions=4, + sample_rate=8, metric="cosine", replace=True, fragment_group_size=2,