fix: respect segmented scalar index flag

daft_lance/_lance.py

@@ -388,14 +388,14 @@ def create_scalar_index(
         io_config: A custom IOConfig to use when accessing LanceDB data. Defaults to None.
         column: Column name to index
         index_type: Type of index to build.
-            For distributed execution this supports "INVERTED", "FTS", and "BTREE".
+            For distributed segmented execution this supports "INVERTED", "FTS", and "BTREE".
             Other scalar index types supported by Lance (for example "BITMAP", "NGRAM", "ZONEMAP",
             "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 False.
             This is only supported by scalar index types that are passed directly to
-            ``LanceDataset.create_scalar_index(...)``. Distributed BTREE/INVERTED/FTS
+            ``LanceDataset.create_scalar_index(...)``. Segmented 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.
@@ -414,11 +414,10 @@ def create_scalar_index(
             If None, Daft will use its default concurrency setting. Must be a positive integer.
         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``.
+            ``commit_existing_index_segments``. ``"FTS"`` is normalized to Lance's
+            inverted full-text index. If False, scalar index creation uses the legacy
+            partitioned workflow or Lance's direct ``create_scalar_index`` path.
+        **kwargs: Additional keyword arguments forwarded to the selected Lance index creation API.
 
     Returns:
         None

daft_lance/lance_scalar_index.py

@@ -18,9 +18,6 @@
 
 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"}
 
@@ -145,12 +142,12 @@ def create_scalar_index_internal(
 ) -> None:
     """Internal implementation of distributed scalar index creation.
 
-    ``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.
+    When ``segmented=True``, ``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")
@@ -203,7 +200,7 @@ def create_scalar_index_internal(
     if name is None:
         name = f"{column}_{index_type.lower()}_idx"
 
-    use_segmented_workflow = segmented or index_type in SEGMENTED_INDEX_TYPES
+    use_segmented_workflow = segmented
 
     # Handle replace parameter - check for existing index with same name
     if not replace or use_segmented_workflow:
@@ -216,6 +213,20 @@ def create_scalar_index_internal(
         if name in existing_names:
             raise ValueError(f"Index with name '{name}' already exists. Set replace=True to replace it.")
 
+    if index_type == "BTREE" and not use_segmented_workflow:
+        logger.info(
+            "Falling back to Lance scalar index creation for non-segmented BTREE index %s.",
+            name,
+        )
+        lance_ds.create_scalar_index(
+            column=column,
+            index_type=index_type,
+            name=name,
+            replace=replace,
+            **kwargs,
+        )
+        return
+
     # Get available fragment IDs to use
     fragments = lance_ds.get_fragments()
     fragment_ids_to_use = [fragment.fragment_id for fragment in fragments]

tests/io/lancedb/test_lancedb_scalar_index.py

@@ -9,8 +9,8 @@
 import pytest
 
 import daft
-from daft.dependencies import pd
-from daft_lance import create_scalar_index
+from daft.dependencies import pa, pd
+from daft_lance import create_scalar_index, lance_scalar_index
 from daft_lance.lance_scalar_index import (
     SegmentedFragmentIndexHandler,
     _existing_index_names,
@@ -101,12 +101,11 @@ def test_build_distributed_index_search_functionality(self, multi_fragment_lance
         )
         updated_dataset = lance.dataset(dataset_uri)
 
-        # Verify the index was created with usable details.
-        described = updated_dataset.describe_indices()
-        text_index = next((idx for idx in described if "text" in idx.name), None)
-
-        assert text_index is not None, "Text index not found"
-        assert "Inverted" in text_index.type_url, f"Expected Inverted index, got {text_index.type_url}"
+        # Verify the index was created. The default non-segmented path does not
+        # populate Lance index details, so use list_indices() here.
+        indices = updated_dataset.list_indices()
+        index_names = [idx["name"] for idx in indices]
+        assert "text_inverted_idx" in index_names, f"Text index not found in {index_names}"
 
         # Test full-text search functionality
         search_term = "Python"
@@ -289,8 +288,8 @@ def test_build_distributed_index_replace_false_existing_index(self, multi_fragme
         error_msg = str(exc_info.value)
         assert "already exists" in error_msg and index_name in error_msg
 
-    def test_build_distributed_index_replace_true_existing_index_is_rejected(self, multi_fragment_lance_dataset):
-        """Test that segmented replace=True rejects existing indexes without corrupting them."""
+    def test_build_distributed_index_replace_true_overwrites_existing(self, multi_fragment_lance_dataset):
+        """Test that default non-segmented replace=True overwrites existing indexes."""
         dataset_uri = multi_fragment_lance_dataset
         index_name = "test_replace_true_index"
 
@@ -303,37 +302,35 @@ def test_build_distributed_index_replace_true_existing_index_is_rejected(self, m
         )
 
         updated_dataset = lance.dataset(dataset_uri)
-        initial_indices = updated_dataset.describe_indices()
+        initial_indices = updated_dataset.list_indices()
         assert len(initial_indices) > 0, "Initial index creation failed"
 
         # Find our initial index
-        initial_index = next((idx for idx in initial_indices if idx.name == index_name), None)
+        initial_index = next((idx for idx in initial_indices if idx["name"] == index_name), None)
         assert initial_index is not None, "Initial index not found"
 
-        with pytest.raises(ValueError, match="cannot atomically replace existing index"):
-            create_scalar_index(
-                uri=dataset_uri,
-                column="text",
-                index_type="INVERTED",
-                name=index_name,
-                replace=True,
-            )
+        create_scalar_index(
+            uri=dataset_uri,
+            column="text",
+            index_type="INVERTED",
+            name=index_name,
+            replace=True,
+        )
 
         updated_dataset = lance.dataset(dataset_uri)
-        final_indices = updated_dataset.describe_indices()
-        final_index = next((idx for idx in final_indices if idx.name == index_name), None)
+        final_indices = updated_dataset.list_indices()
+        final_index = next((idx for idx in final_indices if idx["name"] == index_name), None)
 
-        assert final_index is not None, "Existing index should still exist after rejected replacement"
-        assert "Inverted" in final_index.type_url, "Index type should remain Inverted"
+        assert final_index is not None, "Index should still exist after replacement"
 
-        # Test that the original index still works for searching
+        # Test that the replaced index still works for searching
         search_term = "Python"
         results = updated_dataset.scanner(
             full_text_query=search_term,
             columns=["id", "text"],
         ).to_table()
 
-        assert results.num_rows > 0, f"No results found for search term '{search_term}' after rejected replacement"
+        assert results.num_rows > 0, f"No results found for search term '{search_term}' after index replacement"
 
     def test_build_distributed_index_auto_adjust_workers(self, temp_dir):
         """Test that concurrency is automatically adjusted if it exceeds fragment count."""
@@ -375,11 +372,9 @@ def test_build_distributed_index_fragment_group_size(self, multi_fragment_lance_
         )
 
         updated_dataset = lance.dataset(dataset_uri)
-        described = updated_dataset.describe_indices()
-        assert len(described) == 1
-        assert described[0].name == index_name
-        assert described[0].type_url == "/lance.table.InvertedIndexDetails"
-        assert described[0].num_rows_indexed == 8
+        indices = updated_dataset.list_indices()
+        index_names = [idx["name"] for idx in indices]
+        assert index_name in index_names, f"Index {index_name!r} not found in {index_names}"
 
         results = updated_dataset.scanner(
             full_text_query="Python",
@@ -419,10 +414,9 @@ def test_build_distributed_index_fts_type(self, multi_fragment_lance_dataset):
         )
 
         updated_dataset = lance.dataset(dataset_uri)
-        described = updated_dataset.describe_indices()
-        assert len(described) == 1
-        assert described[0].name == index_name
-        assert described[0].type_url == "/lance.table.InvertedIndexDetails"
+        indices = updated_dataset.list_indices()
+        index_names = [idx["name"] for idx in indices]
+        assert index_name in index_names, f"FTS index not found in {index_names}"
 
         # Test search functionality
         search_term = "Python"
@@ -890,31 +884,39 @@ def test_segmented_btree_integer_column(self, temp_dir):
         results = updated_dataset.scanner(filter="count > 500", columns=["id", "count"]).to_table()
         assert results.num_rows == 3  # 600, 700, 800
 
-    def test_segmented_false_uses_segmented_flow_for_btree(self, temp_dir):
-        """Test that BTREE uses the segmented flow even when segmented=False."""
-        data = {
-            "id": [1, 2, 3, 4, 5, 6, 7, 8],
-            "price": [10.5, 20.75, 30.0, 40.25, 50.5, 60.75, 70.0, 80.25],
-        }
-        dataset = daft.from_pydict(data)
-        path = Path(temp_dir) / "segmented_false.lance"
-        dataset.write_lance(uri=path, max_rows_per_file=2)
+    def test_segmented_false_uses_lance_scalar_flow_for_btree(self, monkeypatch):
+        """Test that segmented=False opts out of distributed BTREE workflows."""
 
-        # segmented=False is retained for API compatibility; BTREE still uses
-        # Lance's public segmented-index workflow by default.
-        create_scalar_index(
-            uri=path,
+        class FakeLanceDataset:
+            schema = pa.schema([("price", pa.float64())])
+
+            def describe_indices(self):
+                return []
+
+            def create_scalar_index(self, **kwargs):
+                calls.append(("lance_scalar", kwargs))
+
+        calls = []
+
+        def fake_create_segmented_index(**kwargs):
+            calls.append(("segmented", kwargs))
+
+        def fake_create_partitioned_index(**kwargs):
+            calls.append(("partitioned", kwargs))
+
+        monkeypatch.setattr(lance_scalar_index, "_create_segmented_index", fake_create_segmented_index)
+        monkeypatch.setattr(lance_scalar_index, "_create_partitioned_index", fake_create_partitioned_index)
+
+        create_scalar_index_internal(
+            lance_ds=FakeLanceDataset(),
+            uri="memory://btree",
             column="price",
             index_type="BTREE",
             name="price_btree_idx",
             segmented=False,
         )
 
-        updated_dataset = lance.dataset(path)
-        described = updated_dataset.describe_indices()
-        assert len(described) == 1
-        assert described[0].name == "price_btree_idx"
-        assert "BTree" in described[0].type_url
+        assert [call[0] for call in calls] == ["lance_scalar"]
 
     def test_segmented_inverted_creates_index(self, multi_fragment_lance_dataset):
         """Test that segmented=True with INVERTED creates an index."""