From 1654d3dbc6fec09d8235a4799a534151241215a3 Mon Sep 17 00:00:00 2001 From: wangzheyan Date: Tue, 7 Jul 2026 19:46:35 +0800 Subject: [PATCH] fix: respect segmented scalar index flag --- daft_lance/_lance.py | 13 +-- daft_lance/lance_scalar_index.py | 31 +++-- tests/io/lancedb/test_lancedb_scalar_index.py | 108 +++++++++--------- 3 files changed, 82 insertions(+), 70 deletions(-) diff --git a/daft_lance/_lance.py b/daft_lance/_lance.py index c4d9d58..1832771 100644 --- a/daft_lance/_lance.py +++ b/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 diff --git a/daft_lance/lance_scalar_index.py b/daft_lance/lance_scalar_index.py index 05b37c0..35b0875 100644 --- a/daft_lance/lance_scalar_index.py +++ b/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] diff --git a/tests/io/lancedb/test_lancedb_scalar_index.py b/tests/io/lancedb/test_lancedb_scalar_index.py index 5952be7..93a5bf8 100644 --- a/tests/io/lancedb/test_lancedb_scalar_index.py +++ b/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."""