feat(python): add Bitmap binding for RoaringBitmap

python/python/lance/bitmap.py

@@ -0,0 +1,10 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright The Lance Authors
+
+# Both names refer to the same class: `Bitmap` for use in type hints and
+# isinstance checks, and the lowercase `bitmap` constructor-style alias
+# (like `list`/`set`) for `from lance.bitmap import bitmap; bitmap([1, 2])`.
+from .lance import Bitmap as Bitmap
+from .lance import Bitmap as bitmap
+
+__all__ = ["Bitmap", "bitmap"]

python/python/lance/dataset.py

@@ -27,7 +27,6 @@
     Literal,
     Optional,
     Sequence,
-    Set,
     Tuple,
     TypedDict,
     Union,
@@ -51,6 +50,7 @@
 from .fragment import DataFile, FragmentMetadata, LanceFragment
 from .indices import IndexConfig, IndexSegment, SupportedDistributedIndices
 from .lance import (
+    Bitmap,
     CleanupExplanation,
     CleanupStats,
     Compaction,
@@ -5586,7 +5586,9 @@ class Index:
     name: str
     fields: List[int]
     dataset_version: int
-    fragment_ids: Set[int]
+    fragment_ids: Bitmap
+    """The fragments covered by this index. A ``Set[int]``/``List[int]`` is
+    also accepted when constructing an ``Index``."""
     index_version: int
     created_at: Optional[datetime] = None
     base_id: Optional[int] = None
@@ -5603,7 +5605,7 @@ class IndexInformation(TypedDict):
     uuid: str
     fields: List[str]
     version: int
-    fragment_ids: Set[int]
+    fragment_ids: Bitmap
     base_id: Optional[int]
 
 
@@ -5999,11 +6001,12 @@ class DataOverlayFile:
         layered over the base data without rewriting the base files.
 
         The overlay is dense or sparse depending on the shape of ``offsets``:
-        pass a flat ``List[int]`` for a dense overlay (one offset list shared by
-        every field in ``data_file``) or a ``List[List[int]]`` for a sparse
-        overlay (one offset list per field, in the order of the file's fields).
-        Offsets are **physical** row offsets (positions in the base files,
-        counting deleted rows), like deletion vectors.
+        pass a single iterable of ints (e.g. a :class:`~lance.bitmap.Bitmap`
+        or a ``List[int]``) for a dense overlay (one offset set shared by
+        every field in ``data_file``), or a list of int iterables for a
+        sparse overlay (one offset set per field, in the order of the file's
+        fields). Offsets are **physical** row offsets (positions in the base
+        files, counting deleted rows), like deletion vectors.
 
         Attributes
         ----------
@@ -6012,12 +6015,20 @@ class DataOverlayFile:
             value column per covered field. The value at each covered offset is
             stored at the rank (0-based count of covered offsets below it) of
             that offset in the field's coverage.
-        offsets : Union[List[int], List[List[int]]]
-            The covered physical row offsets. A flat list is dense coverage
-            (shared by every field); a list of per-field lists is sparse
-            coverage (in field order). Each list must be strictly ascending
-            with no duplicates, since the Nth offset maps to the Nth value row
-            in ``data_file``; a non-ascending list raises ``ValueError``.
+        offsets : Iterable[int] | List[Iterable[int]]
+            The covered physical row offsets. A single int iterable is dense
+            coverage (shared by every field); a list of int iterables is
+            sparse coverage (in field order). Offsets are always resolved in
+            ascending order — the smallest covered offset maps to row 0 of
+            ``data_file``, the next-smallest to row 1, and so on — regardless
+            of the order values are given in, so a plain ``List[int]`` need
+            not be pre-sorted.
+
+            This is a low-level API with no built-in protection against a
+            duplicate offset: since the coverage is stored as a set, a
+            repeated offset silently collapses to one entry, shifting every
+            later offset onto the wrong row of ``data_file`` with no error.
+            Callers are responsible for passing distinct offsets.
         committed_version : Optional[int]
             The dataset version at which this overlay became effective. Leave as
             ``None`` when creating an overlay to commit — the commit stamps it.
@@ -6026,7 +6037,7 @@ class DataOverlayFile:
         """
 
         data_file: DataFile
-        offsets: Union[List[int], List[List[int]]]
+        offsets: Union[Iterable[int], List[Iterable[int]]]
         committed_version: Optional[int] = None
 
     @dataclass

python/python/lance/fragment.py

@@ -26,6 +26,7 @@
 
 import pyarrow as pa
 
+from .lance import Bitmap, _Fragment, _write_fragments, _write_fragments_transaction
 from .lance import (
     DeletionFile as DeletionFile,
 )
@@ -35,7 +36,6 @@
 from .lance import (
     RowIdMeta as RowIdMeta,
 )
-from .lance import _Fragment, _write_fragments, _write_fragments_transaction
 from .progress import FragmentWriteProgress, NoopFragmentWriteProgress
 from .types import _coerce_reader
 from .udf import BatchUDF, normalize_transform
@@ -123,11 +123,18 @@ def _data_file_to_json(f: DataFile) -> dict:
             d["path"] = d.pop("_path")
             return d
 
+        def _offsets_to_json(offsets):
+            # `offsets` is a Bitmap (dense) or a list of Bitmap/int-list (sparse,
+            # per field); normalize to plain (nested) lists of ints for JSON.
+            if isinstance(offsets, Bitmap):
+                return list(offsets)
+            return [list(o) if isinstance(o, Bitmap) else o for o in offsets]
+
         files = [_data_file_to_json(f) for f in self.files]
         overlays = [
             dict(
                 data_file=_data_file_to_json(o.data_file),
-                offsets=o.offsets,
+                offsets=_offsets_to_json(o.offsets),
                 committed_version=o.committed_version,
             )
             for o in self.overlays

python/python/lance/lance/__init__.pyi

@@ -53,6 +53,7 @@ from ..progress import FragmentWriteProgress as FragmentWriteProgress
 from ..progress import IndexProgress as IndexProgress
 from ..types import ReaderLike as ReaderLike
 from ..udf import BatchUDF as BatchUDF
+from .bitmap import Bitmap as Bitmap
 from .debug import format_fragment as format_fragment
 from .debug import format_manifest as format_manifest
 from .debug import format_schema as format_schema

python/python/lance/lance/bitmap.pyi

@@ -0,0 +1,31 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright The Lance Authors
+
+from typing import Iterable, Iterator, Optional, Union
+
+import pyarrow as pa
+
+class Bitmap:
+    """
+    A set of non-negative integers backed by a Rust ``RoaringBitmap``.
+
+    Cheap to clone and to pass into Lance APIs that accept a bitmap, since no
+    per-value Python object is created. Mutating methods (``add``,
+    ``discard``, ``update``) copy-on-write: cloning a ``Bitmap`` and mutating
+    one copy never affects the other.
+    """
+
+    def __init__(
+        self,
+        values: Optional[
+            Union[Iterable[int], pa.Array, pa.ChunkedArray, "Bitmap"]
+        ] = None,
+    ) -> None: ...
+    def __len__(self) -> int: ...
+    def __contains__(self, value: int) -> bool: ...
+    def __iter__(self) -> Iterator[int]: ...
+    def __eq__(self, other: object) -> bool: ...
+    def __ne__(self, other: object) -> bool: ...
+    def add(self, value: int) -> None: ...
+    def discard(self, value: int) -> None: ...
+    def update(self, values: Iterable[int]) -> None: ...

python/python/lance/lance/indices/__init__.pyi

@@ -18,14 +18,15 @@ from typing import Optional
 import pyarrow as pa
 
 from .. import _Fragment
+from ..bitmap import Bitmap
 
 class IndexConfig:
     index_type: str
     config: str
 
 class IndexSegment:
     uuid: str
-    fragment_ids: set[int]
+    fragment_ids: Bitmap
     index_version: int
 
     def __repr__(self) -> str: ...
@@ -74,7 +75,7 @@ def build_rq_model(
 class IndexSegmentDescription:
     uuid: str
     dataset_version_at_last_update: int
-    fragment_ids: set[int]
+    fragment_ids: Bitmap
     index_version: int
     created_at: Optional[datetime]
     size_bytes: Optional[int]

python/python/tests/test_bitmap.py

@@ -0,0 +1,157 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright The Lance Authors
+
+import pickle
+
+import pyarrow as pa
+import pytest
+from lance.bitmap import bitmap
+
+
+def test_construct_from_list():
+    b = bitmap([4, 1, 2, 1])
+    assert len(b) == 3
+    assert set(b) == {1, 2, 4}
+
+
+def test_construct_from_range():
+    b = bitmap(range(1000))
+    assert len(b) == 1000
+    assert 0 in b
+    assert 999 in b
+    assert 1000 not in b
+
+
+def test_construct_empty():
+    assert len(bitmap()) == 0
+    assert len(bitmap([])) == 0
+
+
+def test_construct_from_bitmap():
+    original = bitmap([1, 2, 3])
+    copy = bitmap(original)
+    assert copy == original
+
+
+@pytest.mark.parametrize(
+    "arrow_type",
+    [
+        pa.int8(),
+        pa.int16(),
+        pa.int32(),
+        pa.int64(),
+        pa.uint8(),
+        pa.uint16(),
+        pa.uint32(),
+        pa.uint64(),
+    ],
+)
+def test_construct_from_pyarrow_array(arrow_type):
+    arr = pa.array(range(100), type=arrow_type)
+    b = bitmap(arr)
+    assert len(b) == 100
+    assert set(b) == set(range(100))
+
+
+def test_construct_from_chunked_array():
+    chunked = pa.chunked_array(
+        [pa.array([1, 2, 3], type=pa.int32()), pa.array([4, 5], type=pa.int32())]
+    )
+    b = bitmap(chunked)
+    assert set(b) == {1, 2, 3, 4, 5}
+
+
+def test_construct_from_pyarrow_rejects_nulls():
+    arr = pa.array([1, 2, None], type=pa.int32())
+    with pytest.raises(ValueError):
+        bitmap(arr)
+
+
+def test_construct_rejects_negative_values():
+    with pytest.raises(OverflowError):
+        bitmap([-1])
+
+
+def test_construct_from_pyarrow_rejects_out_of_range():
+    arr = pa.array([2**33], type=pa.int64())
+    with pytest.raises(ValueError):
+        bitmap(arr)
+
+
+def test_construct_from_pyarrow_rejects_out_of_range_uint64():
+    # A u64 value that doesn't fit in u32 must be reported as its actual
+    # value, not silently wrapped (e.g. through a signed cast) into a
+    # misleading negative number.
+    huge = 2**64 - 1
+    arr = pa.array([huge], type=pa.uint64())
+    with pytest.raises(ValueError, match=str(huge)):
+        bitmap(arr)
+
+
+def test_len_iter_contains():
+    b = bitmap([1, 2, 4])
+    assert len(b) == 3
+    assert 1 in b
+    assert 3 not in b
+    assert "not an int" not in b
+    assert sorted(b) == [1, 2, 4]
+
+
+def test_iter_is_lazy():
+    """`iter()` returns a dedicated streaming iterator (not a `list_iterator`
+    over a pre-built list), yielding values one at a time on demand."""
+    b = bitmap(range(1000))
+    it = iter(b)
+    assert type(it).__name__ == "BitmapIterator"
+    assert next(it) == 0
+    assert next(it) == 1
+    assert list(it) == list(range(2, 1000))
+
+
+def test_equality():
+    assert bitmap([1, 2, 3]) == bitmap([3, 2, 1])
+    assert bitmap([1, 2, 3]) == {1, 2, 3}
+    assert bitmap([1, 2, 3]) != bitmap([1, 2])
+    assert bitmap([1, 2, 3]) != {1, 2}
+
+
+def test_equality_against_incompatible_type_is_false_not_error():
+    # Comparing to a value that isn't a Bitmap or an iterable of ints (e.g. a
+    # bare int) is simply unequal, matching normal Python `==` semantics —
+    # it must not raise just because the type differs.
+    b = bitmap([1, 2, 3])
+    assert (b == 5) is False
+    assert (b != 5) is True
+    assert (b == "not iterable") is False
+    assert (b != "not iterable") is True
+
+
+def test_repr():
+    assert repr(bitmap([1, 2, 3])) == "Bitmap({1, 2, 3})"
+
+
+def test_pickle_round_trip():
+    b = bitmap(range(10_000))
+    loaded = pickle.loads(pickle.dumps(b))
+    assert loaded == b
+
+
+def test_add_discard_update():
+    b = bitmap([1, 2, 3])
+    b.add(4)
+    assert 4 in b
+    b.discard(2)
+    assert 2 not in b
+    b.update([10, 11])
+    assert {10, 11}.issubset(set(b))
+
+
+def test_mutation_is_copy_on_write():
+    original = bitmap([1, 2, 3])
+    other = bitmap(original)
+
+    other.add(4)
+
+    assert 4 not in original
+    assert 4 in other
+    assert original == bitmap([1, 2, 3])