feat: add optimize_indices for incremental index maintenance

README.md

@@ -74,12 +74,38 @@ no `fragment_ids` instead).
 
 
 > **Breaking change (from 0.5.0):** the `segmented` parameter was removed —
-> the distributed segment-index workflow is now the only code path, and the
-> single-node fallbacks are gone. `replace` now defaults to `True` (matching
+> the distributed segment-index workflow is now the only code path, and
+> the single-node fallbacks are gone. `replace` now defaults to `True` (matching
 > pylance). Indexes created by older versions of
 > `create_scalar_index` on `INVERTED` columns may carry empty index metadata
 > (see #69); rebuilding them with `replace=True` records full metadata.
 
+#### Index Maintenance
+
+Appended data is not indexed automatically — queries stay correct (uncovered
+fragments fall back to scans) but slow down as the unindexed share grows.
+`optimize_indices` restores index health on the dataset's latest version:
+it indexes newly appended fragments, merges small segments, and heals stale
+fragment IDs left inside mixed segments by deletes as part of a commit that
+indexes or merges new data. It commits no new version when there is no new
+data to index and no segments to merge.
+
+```python
+from daft_lance import optimize_indices
+
+stats = optimize_indices("s3://bucket/my_dataset")
+stats = optimize_indices("s3://bucket/my_dataset", indices=["name_idx"], num_indices_to_merge=4)
+```
+
+Like lance-ray, `optimize_indices` delegates to pylance's
+`DatasetOptimizer.optimize_indices` and runs in the coordinator process; for
+a distributed rebuild use `create_scalar_index(..., replace=True)`. It
+returns `OptimizeIndicesStats` with the dataset versions immediately before
+and after the call, the duration, and per-index segment/coverage counts
+(coverage counts only fragments still live in the manifest); unknown or
+empty `indices` raise `ValueError`, and duplicates are ignored.
+
+
 ### Column Merging
 
 ```python

daft_lance/__init__.py

@@ -9,15 +9,19 @@
     create_scalar_index,
     merge_columns,
     merge_columns_df,
+    optimize_indices,
     read_lance,
     write_lance,
 )
+from .lance_scalar_index import OptimizeIndicesStats
 
 __all__ = [
+    "OptimizeIndicesStats",
     "compact_files",
     "create_scalar_index",
     "merge_columns",
     "merge_columns_df",
+    "optimize_indices",
     "read_lance",
     "take_blobs",
     "write_lance",

daft_lance/_lance.py

@@ -16,7 +16,7 @@
 from .lance_compaction import compact_files_internal
 from .lance_data_sink import LanceDataSink
 from .lance_merge_column import merge_columns_from_df, merge_columns_internal
-from .lance_scalar_index import create_scalar_index_internal
+from .lance_scalar_index import OptimizeIndicesStats, create_scalar_index_internal, optimize_indices_internal
 from .lance_scan import LanceScanOperator
 from .namespace import validate_uri_or_namespace
 from .utils import construct_lance_dataset_handle
@@ -550,6 +550,110 @@ def create_scalar_index(
     )
 
 
+@PublicAPI
+def optimize_indices(
+    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, str] | None = None,
+    indices: list[str] | None = None,
+    num_indices_to_merge: int | None = None,
+    storage_options: dict[str, Any] | 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,
+) -> OptimizeIndicesStats:
+    """Incrementally optimize existing indexes.
+
+    As data is appended it is not added to existing indexes automatically:
+    queries keep working (uncovered fragments fall back to scans) but they
+    get slower as the unindexed share grows. This function restores index
+    health on the dataset's latest version: newly appended fragments are
+    indexed, small segments are merged, and stale fragment IDs left inside
+    mixed segments by deletes are healed as part of a commit that indexes
+    or merges new data. It commits no new version when there is no new
+    data to index and no segments to merge.
+
+    Delegates to pylance's ``DatasetOptimizer.optimize_indices`` — the same
+    choice lance-ray makes — so it runs in the coordinator process (Lance
+    core parallelizes the underlying scans). A full distributed rebuild is
+    ``create_scalar_index(..., replace=True)``.
+
+    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 Lance 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".
+        indices: Names of the indexes to optimize. ``None`` (the default)
+            optimizes every index on the dataset. Unknown names and an empty
+            list raise ``ValueError``; duplicates are ignored.
+        num_indices_to_merge: How many segments to merge when compacting an
+            index (passed to pylance). ``0`` indexes the new data into a new
+            segment instead of merging; ``None`` uses pylance's default.
+        storage_options: Storage options for the dataset.
+        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 for the dataset.
+        metadata_cache_size_bytes: Size of the metadata cache in bytes.
+
+    Returns:
+        OptimizeIndicesStats: versions of the dataset's latest snapshot
+        immediately before and after the call, wall-clock duration, and
+        per-index segment/coverage counts (coverage counts only fragments
+        still live in the manifest). ``changed`` is ``True`` when a new
+        version became visible during the call — with no concurrent
+        writers, exactly when this run committed one.
+
+    Raises:
+        ValueError: If ``indices`` is empty or names indexes that do not
+            exist on the dataset.
+
+    Examples:
+        >>> import daft_lance
+        >>> stats = daft_lance.optimize_indices("s3://my-bucket/dataset/")  # doctest: +SKIP
+        >>> stats.changed  # doctest: +SKIP
+        True
+        >>> [i.name for i in stats.indices]  # doctest: +SKIP
+        ['name_idx']
+
+        Optimize one index and merge its small segments:
+
+        >>> daft_lance.optimize_indices(  # doctest: +SKIP
+        ...     "s3://my-bucket/dataset/", indices=["name_idx"], num_indices_to_merge=4
+        ... )
+    """
+    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,
+        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,
+    )
+
+    return optimize_indices_internal(
+        dataset_handle.dataset,
+        dataset_handle.worker_open_context(),
+        indices=indices,
+        num_indices_to_merge=num_indices_to_merge,
+    )
+
+
 @PublicAPI
 def compact_files(
     uri: str | pathlib.Path | None = None,

daft_lance/lance_scalar_index.py

@@ -1,7 +1,9 @@
 from __future__ import annotations
 
+import dataclasses
 import logging
 import pickle
+import time
 from typing import TYPE_CHECKING, Any, cast
 
 import daft
@@ -464,3 +466,163 @@ def _validate_segments_against_manifest(
                 "not covered by any built segment (a worker result was likely lost). Re-run the "
                 "index build."
             )
+
+
+@dataclasses.dataclass(frozen=True)
+class OptimizedIndexStats:
+    """Per-index outcome of an ``optimize_indices`` run.
+
+    ``fragments_covered_*`` count only fragments that still exist in the
+    manifest: stale IDs left inside a segment by deletes do not count as
+    coverage. ``*_after`` fields are zero when the index is absent from the
+    after-snapshot (e.g. a concurrent writer dropped it).
+    """
+
+    name: str
+    segments_before: int
+    segments_after: int
+    fragments_covered_before: int
+    fragments_covered_after: int
+
+
+@dataclasses.dataclass(frozen=True)
+class OptimizeIndicesStats:
+    """Outcome of an ``optimize_indices`` run over one dataset.
+
+    The versions are sampled from the dataset's latest version immediately
+    before and after the optimize call, so they describe the same lineage a
+    single writer sees.
+    """
+
+    version_before: int
+    version_after: int
+    duration_seconds: float
+    indices: list[OptimizedIndexStats]
+
+    @property
+    def changed(self) -> bool:
+        """Whether a new dataset version became visible during the run.
+
+        With no concurrent writers this is exactly whether this run
+        committed a version; a concurrent commit during the call also makes
+        it ``True``.
+        """
+        return self.version_after != self.version_before
+
+
+def _index_snapshot(lance_ds: lance.LanceDataset) -> dict[str, tuple[int, int]]:
+    """Map index name to ``(segment count, live covered-fragment count)``.
+
+    Coverage counts only fragment IDs that still exist in the manifest, so
+    stale IDs left by deletes do not masquerade as coverage. Datasets with
+    legacy manifests that ``describe_indices`` cannot parse fall back to
+    ``list_indices`` names with unknown counts, the same degradation
+    ``create_scalar_index`` uses (``_existing_index_names``).
+    """
+    live_fragments = {fragment.fragment_id for fragment in lance_ds.get_fragments()}
+    try:
+        descriptions = lance_ds.describe_indices()
+    except Exception:
+        logger.warning("describe_indices() failed; reporting index names only", exc_info=True)
+        try:
+            return {cast(dict[str, Any], idx)["name"]: (0, 0) for idx in lance_ds.list_indices()}
+        except Exception:
+            return {}
+
+    snapshot: dict[str, tuple[int, int]] = {}
+    for desc in descriptions:
+        segments = desc.segments or []
+        covered = {fid for segment in segments for fid in (segment.fragment_ids or ())}
+        snapshot[desc.name] = (len(segments), len(covered & live_fragments))
+    return snapshot
+
+
+def optimize_indices_internal(
+    lance_ds: lance.LanceDataset,
+    open_context: DatasetOpenContext,
+    *,
+    indices: list[str] | None = None,
+    num_indices_to_merge: int | None = None,
+) -> OptimizeIndicesStats:
+    """Incrementally maintain existing indexes.
+
+    Delegates to pylance's ``DatasetOptimizer.optimize_indices`` — the same
+    choice lance-ray makes — because Lance core owns the delta-index
+    semantics. One run indexes newly appended fragments, merges small
+    segments (``num_indices_to_merge``) when there are deltas to merge, and
+    heals stale fragment IDs left inside mixed segments **as part of a
+    commit that indexes or merges new data**; with no new data to index it
+    commits nothing and stale-only coverage is left as-is. Heavier changes
+    are a distributed rebuild: ``create_scalar_index(..., replace=True)``.
+
+    ``indices`` is our parameter and gets deterministic semantics here
+    because pylance silently ignores unknown names: an empty list raises,
+    unknown names raise listing the available indexes, and duplicates are
+    ignored (each index is optimized and reported once). Merge-count
+    validation and everything about index internals belong to Lance.
+
+    Both stat snapshots are sampled from the dataset's latest version, so
+    they are never mixed across snapshots even if the caller pinned the
+    handle's version.
+    """
+    if indices is not None:
+        if len(indices) == 0:
+            raise ValueError("indices must be a non-empty list of index names; pass None to optimize all indexes.")
+        unique_indices = list(dict.fromkeys(indices))
+        duplicates = sorted({name for name in unique_indices if indices.count(name) > 1})
+        if duplicates:
+            logger.warning("Duplicate index names %s were given; each index is optimized once.", duplicates)
+        indices = unique_indices
+
+    base = open_context.open_latest()
+    before = _index_snapshot(base)
+    if indices is not None:
+        unknown = sorted(set(indices) - before.keys())
+        if unknown:
+            raise ValueError(f"indices {unknown} do not exist on the dataset. Available index names: {sorted(before)}")
+
+    call_kwargs: dict[str, Any] = {}
+    if indices is not None:
+        call_kwargs["index_names"] = list(indices)
+    if num_indices_to_merge is not None:
+        call_kwargs["num_indices_to_merge"] = num_indices_to_merge
+
+    logger.info(
+        "Optimizing indices: uri=%s, indices=%s, num_indices_to_merge=%s",
+        open_context.uri,
+        indices if indices is not None else "(all)",
+        num_indices_to_merge,
+    )
+    version_before = base.version
+    start = time.monotonic()
+    lance_ds.optimize.optimize_indices(**call_kwargs)  # type: ignore[no-untyped-call]
+    duration = time.monotonic() - start
+
+    latest = open_context.open_latest()
+    after = _index_snapshot(latest)
+
+    selected = indices if indices is not None else sorted(before)
+    per_index = [
+        OptimizedIndexStats(
+            name=name,
+            segments_before=before.get(name, (0, 0))[0],
+            segments_after=after.get(name, (0, 0))[0],
+            fragments_covered_before=before.get(name, (0, 0))[1],
+            fragments_covered_after=after.get(name, (0, 0))[1],
+        )
+        for name in selected
+    ]
+
+    logger.info(
+        "Optimized %d indices in %.2fs: version %d -> %d",
+        len(per_index),
+        duration,
+        version_before,
+        latest.version,
+    )
+    return OptimizeIndicesStats(
+        version_before=version_before,
+        version_after=latest.version,
+        duration_seconds=duration,
+        indices=per_index,
+    )