Add distributed full vector index builds and atomic replacement

.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

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

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",

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,