diff --git a/daft_lance/_lance.py b/daft_lance/_lance.py index 1832771..1bb7430 100644 --- a/daft_lance/_lance.py +++ b/daft_lance/_lance.py @@ -388,16 +388,19 @@ 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 segmented execution this supports "INVERTED", "FTS", and "BTREE". - Other scalar index types supported by Lance (for example "BITMAP", "NGRAM", "ZONEMAP", + For distributed segmented 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(...)``. Segmented BTREE/INVERTED/FTS + ``LanceDataset.create_scalar_index(...)``. Segmented 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 BITMAP indexing with the default ``segmented=False`` behavior, + ``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. diff --git a/daft_lance/lance_scalar_index.py b/daft_lance/lance_scalar_index.py index 35b0875..bd8e60c 100644 --- a/daft_lance/lance_scalar_index.py +++ b/daft_lance/lance_scalar_index.py @@ -19,7 +19,7 @@ logger = logging.getLogger(__name__) # 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: @@ -86,7 +86,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)", @@ -94,6 +94,11 @@ 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: + # Lance's BITMAP segment builder needs a stable shard number to + # distinguish independently-built bitmap segments before merge/commit. + 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 @@ -106,7 +111,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) @@ -142,12 +147,12 @@ def create_scalar_index_internal( ) -> None: """Internal implementation of distributed scalar index creation. - 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 + When ``segmented=True``, ``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. + (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") @@ -182,9 +187,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( @@ -196,16 +206,28 @@ def create_scalar_index_internal( ) return + if index_type == "BITMAP" and not segmented: + logger.info( + "Falling back to Lance scalar index creation for non-segmented BITMAP index %s.", + name, + ) + lance_ds.create_scalar_index( + column=column, + index_type=index_type, + name=name, + replace=replace, + **kwargs, + ) + return + # Generate index name if not provided if name is None: name = f"{column}_{index_type.lower()}_idx" - use_segmented_workflow = segmented - # Handle replace parameter - check for existing index with same name - if not replace or use_segmented_workflow: + if not replace or segmented: existing_names = _existing_index_names(lance_ds) - if name in existing_names and use_segmented_workflow: + if name in existing_names and segmented: 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." @@ -213,9 +235,10 @@ 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: + if index_type == "BTREE" and not segmented: logger.info( - "Falling back to Lance scalar index creation for non-segmented BTREE index %s.", + "Falling back to Lance scalar index creation for non-segmented %s index %s.", + index_type, name, ) lance_ds.create_scalar_index( @@ -265,7 +288,7 @@ def create_scalar_index_internal( # 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: + if segmented: _create_segmented_index( lance_ds=lance_ds, uri=uri, @@ -329,13 +352,31 @@ def _create_segmented_index( **kwargs, ) + segment_data: list[dict[str, Any]] + if index_type == "BITMAP": + # Give each worker-built BITMAP segment a stable Lance shard id. This + # is separate from Lance fragment ids: shard_id identifies the bitmap + # segment, while fragment_ids identify the rows covered by that segment. + 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 93a5bf8..1b325e6 100644 --- a/tests/io/lancedb/test_lancedb_scalar_index.py +++ b/tests/io/lancedb/test_lancedb_scalar_index.py @@ -687,6 +687,164 @@ 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_segmented_bitmap_respects_fragment_group_size(self, monkeypatch): + """Test that segmented BITMAP can group multiple fragments per segment.""" + + class FakeFragment: + def __init__(self, fragment_id): + self.fragment_id = fragment_id + + def count_rows(self): + return 1 + + class FakeLanceDataset: + schema = pa.schema([("flag", pa.int64())]) + + def describe_indices(self): + return [] + + def get_fragments(self): + return [FakeFragment(i) for i in range(4)] + + calls = [] + + def fake_create_segmented_index(**kwargs): + calls.append(kwargs) + + monkeypatch.setattr(lance_scalar_index, "_create_segmented_index", fake_create_segmented_index) + + create_scalar_index_internal( + lance_ds=FakeLanceDataset(), + uri="memory://bitmap", + column="flag", + index_type="BITMAP", + name="flag_bitmap_idx", + segmented=True, + fragment_group_size=2, + ) + + assert [len(group["fragment_ids"]) for group in calls[0]["fragment_data"]] == [2, 2] + + 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_bitmap_replace_true_default_name_preserves_lance_replacement(self, temp_dir): + """Test that default BITMAP replacement preserves Lance's default index name.""" + path = Path(temp_dir) / "bitmap_default_replace.lance" + table = pa.table( + { + "flag": pa.array([1, 2, 1, 3], type=pa.int64()), + } + ) + lance.write_dataset(table, str(path), max_rows_per_file=2) + + initial_dataset = lance.dataset(str(path)) + initial_dataset.create_scalar_index(column="flag", index_type="BITMAP") + initial_names = [idx["name"] for idx in lance.dataset(str(path)).list_indices()] + assert len(initial_names) == 1 + + create_scalar_index( + uri=path, + column="flag", + index_type="BITMAP", + replace=True, + ) + + final_names = [idx["name"] for idx in lance.dataset(str(path)).list_indices()] + assert final_names == initial_names + def test_segmented_btree_basic(self, temp_dir): """Test basic segmented BTree index creation and query.""" data = {