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 8a4b690..45b573c 100644 --- a/README.md +++ b/README.md @@ -104,6 +104,87 @@ 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` — 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` — 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). + +`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`, `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 +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. +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 +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 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 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 ```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..f383c65 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,179 @@ 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, + num_bits: int = 8, + sample_rate: int = 256, + 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, + **kwargs: Any, +) -> None: + """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 + ``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 + ``commit_existing_index_segments``. Segments built against the same + 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``, ``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://`) + 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". 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 + 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. + 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 + ``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. + 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. + 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. + 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. + **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; 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. + + Returns: + None + + 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``, 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) + + 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) + + 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 + + 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, + num_bits=num_bits, + sample_rate=sample_rate, + fragment_group_size=fragment_group_size, + max_concurrency=max_concurrency, + **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..fa5ec2d --- /dev/null +++ b/daft_lance/lance_vector_index.py @@ -0,0 +1,359 @@ +from __future__ import annotations + +import inspect +import logging +import pickle +from functools import lru_cache +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_lance.lance_scalar_index import ( + _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"}) + +# 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", + "index_type", + "name", + "replace", + "train", + "fragment_ids", + "metric", + "num_partitions", + "num_sub_vectors", + "num_bits", + "ivf_centroids", + "pq_codebook", + "storage_options", + "index_uuid", + "ivf_centroids_file", + } +) + + +@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 explicit parameters and the HNSW options parsed from ``**kwargs``. + """ + params = inspect.signature(lance.LanceDataset.create_index_uncommitted).parameters + return (frozenset(params) | {"m", "max_level", "ef_construction"}) - _HANDLER_MANAGED_KWARGS + + +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." + ) + 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 " + "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.") + 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())}." + ) + + +def _index_exists(dataset: lance.LanceDataset, name: str) -> bool: + return any(index.name == name for index in dataset.describe_indices()) + + +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 + + 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. + # 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, + 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, + num_bits: int = 8, + sample_rate: int = 256, + fragment_group_size: int | None = None, + max_concurrency: 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``. 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. + 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. 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") + + 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}") + 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) + + # 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" + + 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.") + + # 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") + + # 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. + builder = lance.indices.IndicesBuilder(lance_ds, column) + 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) + + 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=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) + + 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, + } + + # 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) + + 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).into_partitions(len(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() + 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 + # 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 new file mode 100644 index 0000000..46753f9 --- /dev/null +++ b/tests/io/lance/test_lance_vector_index.py @@ -0,0 +1,598 @@ +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, lance_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, + 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 + + 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) +def test_every_index_type_builds_and_answers_queries(tmp_path: Path, index_type: str) -> None: + """All six advertised vector index types build multiple segments and return results.""" + 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=8, 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 + 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_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.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, + max_rows_per_file=256, + ) + 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) + 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: + """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", 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) + 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: + """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_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_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( + 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 + + 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 + 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_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, sample_rate=32, 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) + + # 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) + + +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 + + +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 + 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: + """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() + + +@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=2048, rows_per_file=512) + 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() + + +@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() + + +@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() + + +@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 new file mode 100644 index 0000000..eeeea45 --- /dev/null +++ b/tests/io/lance/test_lance_vector_index_concurrency.py @@ -0,0 +1,191 @@ +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, optimize_indices +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, + ) + return uri, vectors + + +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 = _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", 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", num_partitions=4, sample_rate=8, 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 = _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", 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", num_partitions=4, sample_rate=8, 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_replace_true_replaces_concurrently_rebuilt_index(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + 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 = {} + + def open_after_rebuild(self): + dataset = lance.dataset(uri) + dataset.create_index( + "vector", + "IVF_FLAT", + name="vector_idx", + metric="cosine", + replace=True, + num_partitions=4, + ) + competing["version"] = dataset.version + competing["segments"] = _segments(dataset) + return original_open(self) + + monkeypatch.setattr(DatasetOpenContext, "open_latest", open_after_rebuild) + create_vector_index( + uri, + column="vector", + index_type="IVF_FLAT", + num_partitions=4, + sample_rate=8, + replace=True, + fragment_group_size=2, + ) + + dataset = lance.dataset(uri) + 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() + + +def test_unrelated_append_after_live_check_allows_index_commit(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + uri, vectors = _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", num_partitions=4, sample_rate=8, 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() + + +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 = _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 = {} + 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", + num_partitions=4, + sample_rate=8, + 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"]