Segment index bitmap

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

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,14 +89,17 @@ 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)",
             fragment_ids,
             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.

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 = {