Add segmented index workflow for INVERTED and FTS

daft_lance/_lance.py

@@ -362,7 +362,7 @@ def create_scalar_index(
     column: str,
     index_type: str = "INVERTED",
     name: str | None = None,
-    replace: bool = True,
+    replace: bool = False,
     storage_options: dict[str, Any] | None = None,
     version: int | str | None = None,
     asof: str | None = None,
@@ -390,9 +390,14 @@ def create_scalar_index(
         index_type: Type of index to build.
             For distributed execution this supports "INVERTED", "FTS", and "BTREE".
             Other scalar index types supported by Lance (for example "BITMAP", "NGRAM", "ZONEMAP",
-            "LABEL_LIST", "BLOOMFILTER") are passed through to Lance's scalar index implementation.
+            "LABEL_LIST", "BLOOMFILTER") are passed directly to
+            ``LanceDataset.create_scalar_index(...)``.
         name: Name of the index (generated if None).
-        replace: Whether to replace an existing index with the same name. Defaults to True.
+        replace: Whether to replace an existing index with the same name. Defaults to False.
+            This is only supported by scalar index types that are passed directly to
+            ``LanceDataset.create_scalar_index(...)``. Distributed BTREE/INVERTED/FTS
+            indexes use Lance's public segmented-index commit API, which does not
+            currently expose atomic replacement, so existing index names are rejected.
         storage_options: Storage options for the dataset.
         version: Version of the dataset to use.
         asof: Timestamp to use for time travel queries.
@@ -407,11 +412,12 @@ def create_scalar_index(
             greater than 1 enable additional parallelism on distributed runners; values <= 1 or None will use the default partitioning.
         max_concurrency: Maximum number of concurrent tasks to use for processing fragment batches.
             If None, Daft will use its default concurrency setting. Must be a positive integer.
-        segmented: If True and ``index_type`` is ``"BTREE"``, use the segmented index
-            workflow where each worker builds a fully independent index segment and the
-            coordinator commits them atomically via ``commit_existing_index_segments``.
-            This produces proper ``index_details`` metadata so ``describe_indices()``
-            works correctly.  Defaults to False (legacy partitioned-and-merged flow).
+        segmented: If True, force the segmented index workflow where each worker builds
+            a fully independent index segment and the coordinator commits them via
+            ``commit_existing_index_segments``. Distributed ``"BTREE"``, ``"INVERTED"``,
+            and ``"FTS"`` indexes use this workflow by default; ``"FTS"`` is normalized
+            to Lance's inverted full-text index. Other scalar index types are passed
+            directly to ``LanceDataset.create_scalar_index(...)``.
         **kwargs: Additional keyword arguments forwarded to ``lance.LanceDataset.create_scalar_index``.
 
     Returns:

daft_lance/lance_scalar_index.py

@@ -18,6 +18,12 @@
 
 logger = logging.getLogger(__name__)
 
+# Scalar index types that use Lance's public segmented-index workflow by default.
+SEGMENTED_INDEX_TYPES = {"BTREE", "INVERTED"}
+
+# Segmented index types whose worker-built segments must be merged before commit.
+MERGED_SEGMENTED_INDEX_TYPES = {"INVERTED"}
+
 
 class FragmentIndexHandler:
     """Handler for distributed scalar index creation on fragment batches."""
@@ -62,13 +68,11 @@ def __call__(self, fragment_ids: list[int]) -> bool:
 class SegmentedFragmentIndexHandler:
     """Handler for segmented scalar index creation on fragment batches.
 
-    Unlike ``FragmentIndexHandler``, which writes partial index files sharing
-    a single UUID, this handler builds a fully independent index segment per
-    worker via the low-level ``_ds.create_index`` binding.  The returned
-    ``lance.Index`` metadata (including ``index_details``) is serialised
-    (pickled) so it can cross Daft process/serialisation boundaries.  The
-    coordinator then commits all segments atomically with
-    ``commit_existing_index_segments``.
+    Each Daft worker receives a subset of Lance fragment IDs and builds an
+    uncommitted index segment for just those fragments using Lance's public
+    ``create_index_uncommitted`` API.  The worker returns the segment metadata
+    to the coordinator, which commits all segments into the dataset manifest
+    with ``commit_existing_index_segments``.
     """
 
     def __init__(
@@ -77,14 +81,12 @@ def __init__(
         column: str,
         index_type: str,
         name: str,
-        replace: bool,
         **kwargs: Any,
     ) -> None:
         self.lance_ds = lance_ds
         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:
@@ -96,29 +98,44 @@ def __call__(self, fragment_ids: list[int]) -> bytes:
             self.index_type,
         )
 
-        # _ds.create_index returns a lance.Index dataclass when fragment_ids
-        # is provided (uncommitted segment mode).
-        index_meta: lance.Index = self.lance_ds._ds.create_index(  # type: ignore[call-arg]
-            [self.column],
-            self.index_type,
+        # Create one uncommitted index segment. ``pylance 8.0.0`` supports
+        # scalar index segments through this public API. Segment creation always
+        # uses ``replace=False`` because replacement, if supported, must happen
+        # in the final manifest commit rather than independently in each worker.
+        index_meta = self.lance_ds.create_index_uncommitted(
+            column=self.column,
+            index_type=self.index_type,
             name=self.name,
-            replace=self.replace,
+            replace=False,
             train=True,
-            storage_options=None,
-            kwargs={"fragment_ids": fragment_ids, **self.kwargs},
+            fragment_ids=fragment_ids,
+            **self.kwargs,
         )
 
         return pickle.dumps(index_meta)
 
 
+def _existing_index_names(lance_ds: lance.LanceDataset) -> set[str]:
+    """Return existing index names, falling back for legacy indexes with bad details."""
+    try:
+        return {idx.name for idx in lance_ds.describe_indices()}
+    except Exception:
+        pass
+
+    try:
+        return {cast(dict[str, Any], idx)["name"] for idx in lance_ds.list_indices()}
+    except Exception:
+        return set()
+
+
 def create_scalar_index_internal(
     lance_ds: lance.LanceDataset,
     uri: str | pathlib.Path,
     *,
     column: str,
     index_type: str = "INVERTED",
     name: str | None = None,
-    replace: bool = True,
+    replace: bool = False,
     storage_options: dict[str, Any] | None = None,
     fragment_group_size: int | None = None,
     num_partitions: int | None = None,
@@ -128,17 +145,12 @@ def create_scalar_index_internal(
 ) -> None:
     """Internal implementation of distributed scalar index creation.
 
-    INVERTED and BTREE use a 3-phase distributed workflow (fragment-parallel build,
-    merge_index_metadata, then commit). ``FTS`` is normalized to ``INVERTED`` (same Lance
-    index); see Lance Rust/Python bindings: ``INVERTED`` and ``FTS`` map to the same
-    inverted full-text index type.
-
-    When ``segmented=True`` and ``index_type`` is ``BTREE``, a cleaner segmented workflow
-    is used instead: each worker builds a fully independent index segment via the low-level
-    ``_ds.create_index`` binding (which returns ``lance.Index`` metadata including
-    ``index_details``), and the coordinator commits them atomically with
-    ``commit_existing_index_segments``.  This resolves a known issue where ``index_details``
-    was left empty in the legacy path, preventing ``describe_indices()`` from working.
+    ``BTREE`` and ``INVERTED`` use Lance's public segment-index workflow: each
+    worker builds a fully independent index segment, and the coordinator commits
+    them atomically with ``commit_existing_index_segments``. ``FTS`` is
+    normalized to ``INVERTED`` (same Lance index); see Lance Rust/Python
+    bindings: ``INVERTED`` and ``FTS`` map to the same inverted full-text index
+    type.
     """
     if not column:
         raise ValueError("Column name cannot be empty")
@@ -190,15 +202,17 @@ def create_scalar_index_internal(
     # Generate index name if not provided
     if name is None:
         name = f"{column}_{index_type.lower()}_idx"
+
+    use_segmented_workflow = segmented or index_type in SEGMENTED_INDEX_TYPES
+
     # Handle replace parameter - check for existing index with same name
-    if not replace:
-        existing_indices = []
-        try:
-            existing_indices = lance_ds.describe_indices()
-        except Exception:
-            # If we can't check existing indices, continue
-            pass
-        existing_names = {idx.name for idx in existing_indices}
+    if not replace or use_segmented_workflow:
+        existing_names = _existing_index_names(lance_ds)
+        if name in existing_names and use_segmented_workflow:
+            raise ValueError(
+                f"Index with name '{name}' already exists and cannot atomically replace existing index "
+                "with Lance's public segmented index API. Drop the existing index first or use a different name."
+            )
         if name in existing_names:
             raise ValueError(f"Index with name '{name}' already exists. Set replace=True to replace it.")
 
@@ -237,17 +251,16 @@ def create_scalar_index_internal(
         segmented,
     )
 
-    # Choose between the segmented workflow and the legacy partitioned-and-merged
-    # workflow.  Segmented mode produces proper ``index_details`` so
-    # ``describe_indices()`` works correctly.
-    if segmented and index_type == "BTREE":
+    # Use segment-index creation for Lance scalar index types that expose the
+    # public uncommitted segment API.  The legacy path is kept as a fallback for
+    # older/unsupported distributed scalar index types.
+    if use_segmented_workflow:
         _create_segmented_index(
             lance_ds=lance_ds,
             uri=uri,
             column=column,
             index_type=index_type,
             name=name,
-            replace=replace,
             storage_options=storage_options,
             fragment_data=fragment_data,
             fragment_ids_to_use=fragment_ids_to_use,
@@ -279,7 +292,6 @@ def _create_segmented_index(
     column: str,
     index_type: str,
     name: str,
-    replace: bool,
     storage_options: dict[str, Any] | None,
     fragment_data: list[dict[str, list[int]]],
     fragment_ids_to_use: list[int],
@@ -289,11 +301,10 @@ def _create_segmented_index(
 ) -> None:
     """Segmented index workflow: each worker builds an independent segment.
 
-    Workers call the low-level ``_ds.create_index`` binding (which returns
-    ``lance.Index`` metadata with ``index_details`` populated), pickle the
-    result so it can traverse Daft serialisation boundaries, and return it.
-    The coordinator unpickles all segments and commits them atomically via
-    ``commit_existing_index_segments``.
+    Workers call Lance's uncommitted index segment API, pickle the returned
+    ``lance.Index`` metadata so it can traverse Daft serialisation boundaries,
+    and return it.  The coordinator unpickles all segments and commits them
+    atomically via ``commit_existing_index_segments``.
     """
     handler_cls = daft.cls(
         SegmentedFragmentIndexHandler,
@@ -304,7 +315,6 @@ def _create_segmented_index(
         column=column,
         index_type=index_type,
         name=name,
-        replace=replace,
         **kwargs,
     )
 
@@ -322,20 +332,34 @@ def _create_segmented_index(
         pickle.loads(raw) for raw in collected.to_pydict()["index_meta"]
     ]
 
+    # Reload dataset to pick up the latest version (segment files were written
+    # by workers against the version that was current at their invocation time).
+    lance_ds = lance.LanceDataset(uri, storage_options=storage_options)
+    index_metas = _prepare_index_segments_for_commit(lance_ds, index_type, index_metas)
+
     logger.info(
         "Collected %d index segments; committing as segmented index %s",
         len(index_metas),
         name,
     )
-
-    # Reload dataset to pick up the latest version (segment files were written
-    # by workers against the version that was current at their invocation time).
-    lance_ds = lance.LanceDataset(uri, storage_options=storage_options)
     lance_ds.commit_existing_index_segments(name, column, index_metas)
 
     logger.info("Segmented index %s committed successfully", name)
 
 
+def _prepare_index_segments_for_commit(
+    lance_ds: lance.LanceDataset,
+    index_type: str,
+    index_metas: list[lance.Index | lance.indices.IndexSegment],
+) -> list[lance.Index | lance.indices.IndexSegment]:
+    """Prepare worker-built segments for the final manifest commit."""
+    if index_type not in MERGED_SEGMENTED_INDEX_TYPES or len(index_metas) <= 1:
+        return index_metas
+
+    merged = lance_ds.merge_existing_index_segments([cast(lance.Index, segment) for segment in index_metas])
+    return [merged]
+
+
 def _create_partitioned_index(
     lance_ds: lance.LanceDataset,
     uri: str | pathlib.Path,

pyproject.toml

@@ -2,13 +2,13 @@
 name = "daft-lance"
 version = "0.4.0"
 description = "Lance integration for Daft - compaction, indexing, merge columns, and REST operations"
-requires-python = ">=3.10"
+requires-python = ">=3.10,<3.14"
 license = "Apache-2.0"
 readme = "README.md"
 dependencies = [
   "lance-namespace>=0.6.0",
   "lance-namespace-urllib3-client>=0.6.0",
-  "pylance>=7.0.0"
+  "pylance>=8.0.0"
 ]
 
 [dependency-groups]