Segment index bitmap

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.

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,14 +86,19 @@ 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:
+            # 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,26 +206,39 @@ 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."
             )
         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.

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