diff --git a/daft_lance/_lance.py b/daft_lance/_lance.py index c4d9d58..6e36321 100644 --- a/daft_lance/_lance.py +++ b/daft_lance/_lance.py @@ -388,16 +388,18 @@ 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". - Other scalar index types supported by Lance (for example "BITMAP", "NGRAM", "ZONEMAP", + For distributed execution this supports "BITMAP", "BTREE", "INVERTED", and "FTS". + Other scalar index types supported by Lance (for example "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(...)``. Distributed BITMAP/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. + For default BITMAP indexing, ``replace=True`` on an existing index falls back + to Lance's scalar replacement path to preserve existing API behavior. storage_options: Storage options for the dataset. version: Version of the dataset to use. asof: Timestamp to use for time travel queries. @@ -414,11 +416,11 @@ 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``. Distributed ``"BITMAP"``, ``"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 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..77f89da 100644 --- a/daft_lance/lance_scalar_index.py +++ b/daft_lance/lance_scalar_index.py @@ -19,10 +19,10 @@ 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 = {"BITMAP", "BTREE", "INVERTED"} # Segmented index types whose worker-built segments must be merged before commit. -MERGED_SEGMENTED_INDEX_TYPES = {"INVERTED"} +MERGED_SEGMENTED_INDEX_TYPES = {"BITMAP", "INVERTED"} class FragmentIndexHandler: @@ -89,7 +89,7 @@ def __init__( self.name = name self.kwargs = kwargs - def __call__(self, fragment_ids: list[int]) -> bytes: + def __call__(self, fragment_ids: list[int], shard_id: int | None = None) -> bytes: """Build an independent index segment and return its pickled metadata.""" logger.info( "Building segmented index segment for fragments %s (column=%s, type=%s)", @@ -97,6 +97,9 @@ def __call__(self, fragment_ids: list[int]) -> bytes: self.column, self.index_type, ) + segment_kwargs = self.kwargs.copy() + if self.index_type == "BITMAP" and shard_id is not None: + segment_kwargs["shard_id"] = shard_id # Create one uncommitted index segment. ``pylance 8.0.0`` supports # scalar index segments through this public API. Segment creation always @@ -109,7 +112,7 @@ def __call__(self, fragment_ids: list[int]) -> bytes: replace=False, train=True, fragment_ids=fragment_ids, - **self.kwargs, + **segment_kwargs, ) return pickle.dumps(index_meta) @@ -145,12 +148,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. + ``BITMAP``, ``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") @@ -185,9 +188,14 @@ def create_scalar_index_internal( and not pa.types.is_string(value_type) ): raise TypeError(f"Column {column} must be numeric or string type for BTREE index, got {value_type}") + case "BITMAP": + # BITMAP supports multiple physical Arrow types depending on the + # Lance release. Leave final type validation to Lance rather than + # duplicating a narrower Python-side allowlist here. + pass case _: logger.warning( - "Distributed indexing currently only supports 'INVERTED' and 'BTREE' index types, not '%s'. So we are falling back to single-threaded index creation.", + "Distributed indexing currently only supports 'BITMAP', 'INVERTED', and 'BTREE' index types, not '%s'. So we are falling back to single-threaded index creation.", index_type, ) lance_ds.create_scalar_index( @@ -208,6 +216,18 @@ def create_scalar_index_internal( # Handle replace parameter - check for existing index with same name if not replace or use_segmented_workflow: existing_names = _existing_index_names(lance_ds) + if name in existing_names and index_type == "BITMAP" and replace and use_segmented_workflow and not segmented: + logger.info( + "Falling back to Lance's scalar BITMAP replacement path because segmented index replacement is not atomic.", + ) + lance_ds.create_scalar_index( + column=column, + index_type=index_type, # type: ignore[arg-type] + name=name, + replace=replace, + **kwargs, + ) + return if name in existing_names and use_segmented_workflow: raise ValueError( f"Index with name '{name}' already exists and cannot atomically replace existing index " @@ -233,6 +253,12 @@ def create_scalar_index_internal( fragment_group_size, ) + if index_type == "BITMAP" and fragment_group_size != 1: + logger.info( + "Adjusted BITMAP segmented fragment_group_size to 1 because released Lance Python APIs require one fragment per bitmap segment.", + ) + fragment_group_size = 1 + logger.info("Starting fragment-parallel processing and creating DataFrame with fragment batches") fragment_data = distribute_fragments_balanced(fragments, fragment_group_size) @@ -318,13 +344,28 @@ def _create_segmented_index( **kwargs, ) + segment_data: list[dict[str, Any]] + if index_type == "BITMAP": + segment_data = [ + { + **group, + "shard_id": shard_id, + } + for shard_id, group in enumerate(fragment_data) + ] + else: + segment_data = fragment_data + with execution_config_ctx(maintain_order=False): if num_partitions is not None and num_partitions > 1: - df = from_pylist(fragment_data).repartition(num_partitions) + df = from_pylist(segment_data).repartition(num_partitions) else: - df = from_pylist(fragment_data) + df = from_pylist(segment_data) - df = df.select(handler(df["fragment_ids"]).alias("index_meta")) + if index_type == "BITMAP": + df = df.select(handler(df["fragment_ids"], df["shard_id"]).alias("index_meta")) + else: + df = df.select(handler(df["fragment_ids"]).alias("index_meta")) collected = df.collect() # Deserialise the Index metadata returned by each worker. diff --git a/tests/io/lancedb/test_lancedb_scalar_index.py b/tests/io/lancedb/test_lancedb_scalar_index.py index 5952be7..8fa1589 100644 --- a/tests/io/lancedb/test_lancedb_scalar_index.py +++ b/tests/io/lancedb/test_lancedb_scalar_index.py @@ -9,7 +9,7 @@ import pytest import daft -from daft.dependencies import pd +from daft.dependencies import pa, pd from daft_lance import create_scalar_index from daft_lance.lance_scalar_index import ( SegmentedFragmentIndexHandler, @@ -693,6 +693,101 @@ def list_indices(self): assert _existing_index_names(FakeLanceDataset()) == {"legacy_idx"} + def test_segmented_bitmap_handler_forwards_shard_id(self): + """Test that BITMAP segment creation forwards the required shard id.""" + + class FakeLanceDataset: + def __init__(self): + self.calls = [] + + @property + def _ds(self): + raise AssertionError("public BITMAP segment creation should not use fallback") + + def create_index_uncommitted(self, **kwargs): + self.calls.append(kwargs) + return {"segment": "bitmap-metadata"} + + fake_ds = FakeLanceDataset() + handler = SegmentedFragmentIndexHandler( + lance_ds=fake_ds, + column="flag", + index_type="BITMAP", + name="flag_idx", + ) + + raw_segment = handler([1, 2], shard_id=7) + + assert pickle.loads(raw_segment) == {"segment": "bitmap-metadata"} + assert fake_ds.calls == [ + { + "column": "flag", + "index_type": "BITMAP", + "name": "flag_idx", + "replace": False, + "train": True, + "fragment_ids": [1, 2], + "shard_id": 7, + } + ] + + def test_segmented_bitmap_segments_are_merged_before_commit(self): + """Test that BITMAP segments are merged to avoid one physical segment per fragment.""" + + class FakeLanceDataset: + def __init__(self): + self.calls = [] + + def merge_existing_index_segments(self, segments): + self.calls.append(segments) + return {"segment": "merged-bitmap"} + + fake_ds = FakeLanceDataset() + segments = [{"segment": "a"}, {"segment": "b"}] + + prepared = _prepare_index_segments_for_commit(fake_ds, "BITMAP", segments) + + assert prepared == [{"segment": "merged-bitmap"}] + assert fake_ds.calls == [segments] + + def test_bitmap_replace_true_existing_index_preserves_lance_replacement(self): + """Test that default BITMAP indexing keeps replace=True behavior for existing indexes.""" + + class ExistingIndex: + name = "flag_bitmap_idx" + + class FakeLanceDataset: + schema = pa.schema([("flag", pa.int64())]) + + def __init__(self): + self.calls = [] + + def describe_indices(self): + return [ExistingIndex()] + + def create_scalar_index(self, **kwargs): + self.calls.append(kwargs) + + fake_ds = FakeLanceDataset() + + create_scalar_index_internal( + lance_ds=fake_ds, + uri="memory://bitmap", + column="flag", + index_type="BITMAP", + name="flag_bitmap_idx", + replace=True, + ) + + assert fake_ds.calls == [ + { + "column": "flag", + "index_type": "BITMAP", + "name": "flag_bitmap_idx", + "replace": True, + } + ] + def test_segmented_btree_basic(self, temp_dir): """Test basic segmented BTree index creation and query.""" data = {