feat(python): add Bitmap binding for RoaringBitmap#7837
Open
wjones127 wants to merge 4 commits into
Open
Conversation
Two sites in python/src/transaction.rs (Index.fragment_ids, and DataOverlayFile.offsets added in lance-format#7540) copy a RoaringBitmap through a Python set/list on every Rust<->Python crossing. This adds a Bitmap type wrapping Arc<RoaringBitmap>, constructible from a Python iterable, range, or pyarrow integer Array/ChunkedArray, with copy-on-write mutation. Both sites now emit/accept Bitmap directly (falling back to plain sets/lists on input for backward compatibility), and the same fix is applied to the two IndexSegment(Description) getters in indices.rs that had the identical set-expansion pattern. Fixes lance-format#7695
github-actions
Bot
added
A-python
📝 WalkthroughWalkthroughAdds a Python Roaring Bitmap type with Arrow and iterable constructors, container and mutation APIs, and pickle support. Index fragment IDs and data overlay offsets now use Bitmap representations, with updated serialization, type stubs, bindings, and tests. ChangesPython Roaring Bitmap Integration
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant PythonCaller
participant Bitmap
participant IndexMetadata
participant DataOverlayFile
participant FragmentMetadata
PythonCaller->>Bitmap: construct or mutate bitmap
PythonCaller->>IndexMetadata: pass fragment IDs
IndexMetadata-->>PythonCaller: return Bitmap fragment IDs
PythonCaller->>DataOverlayFile: pass dense or sparse offsets
DataOverlayFile->>FragmentMetadata: serialize normalized offsets
FragmentMetadata-->>PythonCaller: return Bitmap-backed metadata
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
wjones127
commented
Jul 20, 2026
- DataOverlayFile.offsets is now typed Iterable[int] | List[Iterable[int]] instead of spelling out Bitmap/List[Bitmap]/List[int]/List[List[int]] explicitly, since Bitmap already satisfies Iterable[int]. - bitmap_from_pyarrow collects straight from the arrow buffer into the RoaringBitmap per dtype arm, instead of through an intermediate Vec<u32>. - Bitmap.__iter__ now returns a dedicated streaming iterator (BitmapIterator) instead of eagerly building a PyList of every value up front. - DataOverlayFile's dense/sparse extraction now goes through the same extract_bitmap helper for both flat and per-field offsets, dropping the separate raw-Vec<u32> path and its ascending-order validation. Offsets are always resolved in ascending rank order regardless of input order, so there's no longer a distinct ordering contract to enforce.
- bitmap_from_pyarrow's UInt64 arm cast to i64 before range-checking, so a u64 value above i64::MAX wrapped to a negative number and the resulting ValueError reported that bogus negative instead of the real value. Added a dedicated u64_to_u32 that checks the range directly. - Bitmap's __eq__/__ne__ raised TypeError when compared against a value that isn't a Bitmap or an iterable of ints (e.g. `bitmap([1]) == 5`), instead of returning False/True like normal Python equality and like __contains__ already does for the same case. - Minor clarity fixes from self-review: extracted extract_sparse_bitmaps out of a dense combinator chain, documented that duplicate offsets silently collapse (a RoaringBitmap is a set), documented the ChunkedArray-vs-Array duck-typing check, and documented why bitmap.py exports both `Bitmap` and `bitmap`.
Contributor
Author
|
@coderabbitai review |
✅ Action performedReview finished.
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
Note
Quiet mode is enabled, so only the most important comments were posted inline. Other review comments are grouped below.
🟡 Other comments (1)
python/python/lance/fragment.py-126-131 (1)
126-131: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winNormalize every sparse offset iterable before JSON serialization.
DataOverlayFile.offsetspermitsList[Iterable[int]], but this only converts nestedBitmapvalues. A valid sparse input such as[range(1, 3)]remains arange, causingjson.dumps(metadata.to_json())to fail. Materialize every nested sparse iterable as a list while preserving dense integer iterables.Proposed fix
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. + # Normalize dense and sparse offset iterables to JSON-compatible lists. if isinstance(offsets, Bitmap): return list(offsets) - return [list(o) if isinstance(o, Bitmap) else o for o in offsets] + offsets = list(offsets) + if not offsets or isinstance(offsets[0], int): + return offsets + return [list(offsets_for_field) for offsets_for_field in offsets]Add a round-trip case using sparse
rangeor tuple offsets. As per coding guidelines, “Every bugfix and feature must have corresponding tests.”🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/python/lance/fragment.py` around lines 126 - 131, Update _offsets_to_json to materialize every nested sparse offset iterable as a list of integers, not only Bitmap instances, while preserving the existing dense-offset handling. Add a round-trip test using sparse range or tuple offsets and verify JSON serialization/deserialization succeeds.Source: Coding guidelines
🧹 Nitpick comments (1)
python/src/bitmap.rs (1)
197-219: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winUse
py.NotImplemented()instead of raisingNotImplementedErrorfor unsupported ordering ops.For
Lt/Le/Gt/Ge, returningErr(PyNotImplementedError::new_err(...))raises Python'sNotImplementedErrorexception directly, rather than returning theNotImplementedsentinel. Per PyO3's own guide, the idiomatic way to "leave some operations unimplemented" in__richcmp__is to return the sentinel so Python raises the conventionalTypeError: '<' not supported between instances of .... Code paths that catchTypeErrorfor unsupported comparisons (e.g. viasorted()/total_ordering) would be surprised byNotImplementedErrorhere.♻️ Proposed fix using the `NotImplemented` sentinel
- fn __richcmp__(&self, other: &Bound<'_, PyAny>, op: CompareOp) -> PyResult<bool> { + fn __richcmp__<'py>( + &self, + other: &Bound<'py, PyAny>, + op: CompareOp, + py: Python<'py>, + ) -> PyResult<Bound<'py, PyAny>> { let other_bitmap = if let Ok(other) = other.extract::<Self>() { Some((*other.0).clone()) } else { other.try_iter().ok().and_then(|it| { it.map(|item| item?.extract::<u32>()) .collect::<PyResult<RoaringBitmap>>() .ok() }) }; match (op, other_bitmap) { - (CompareOp::Eq, Some(other)) => Ok(*self.0 == other), - (CompareOp::Eq, None) => Ok(false), - (CompareOp::Ne, Some(other)) => Ok(*self.0 != other), - (CompareOp::Ne, None) => Ok(true), - _ => Err(PyNotImplementedError::new_err( - "Only == and != are supported", - )), + (CompareOp::Eq, Some(other)) => Ok((*self.0 == other).into_pyobject(py)?.into_any()), + (CompareOp::Eq, None) => Ok(false.into_pyobject(py)?.into_any()), + (CompareOp::Ne, Some(other)) => Ok((*self.0 != other).into_pyobject(py)?.into_any()), + (CompareOp::Ne, None) => Ok(true.into_pyobject(py)?.into_any()), + _ => Ok(py.NotImplemented().into_bound(py)), } }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/src/bitmap.rs` around lines 197 - 219, Update __richcmp__ so unsupported ordering operations (Lt, Le, Gt, and Ge) return PyNotImplemented::new(py) as the Python NotImplemented sentinel rather than raising PyNotImplementedError; preserve the existing Eq and Ne behavior.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@python/src/transaction.rs`:
- Around line 195-219: Update extract_bitmap to detect repeated offsets before
collecting them into RoaringBitmap, and return a descriptive PyErr when a
duplicate is encountered. Ensure this validation applies to both direct bitmap
extraction and sparse coverage through extract_sparse_bitmaps, while preserving
normal deduplicated ordering for unique offsets.
---
Other comments:
In `@python/python/lance/fragment.py`:
- Around line 126-131: Update _offsets_to_json to materialize every nested
sparse offset iterable as a list of integers, not only Bitmap instances, while
preserving the existing dense-offset handling. Add a round-trip test using
sparse range or tuple offsets and verify JSON serialization/deserialization
succeeds.
---
Nitpick comments:
In `@python/src/bitmap.rs`:
- Around line 197-219: Update __richcmp__ so unsupported ordering operations
(Lt, Le, Gt, and Ge) return PyNotImplemented::new(py) as the Python
NotImplemented sentinel rather than raising PyNotImplementedError; preserve the
existing Eq and Ne behavior.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: QUIET
Plan: Pro Plus
Run ID: bd1e3e04-2ce8-42f3-b05b-f8fedf8b2b47
📒 Files selected for processing (14)
python/python/lance/bitmap.pypython/python/lance/dataset.pypython/python/lance/fragment.pypython/python/lance/lance/__init__.pyipython/python/lance/lance/bitmap.pyipython/python/lance/lance/indices/__init__.pyipython/python/tests/test_bitmap.pypython/python/tests/test_dataset.pypython/python/tests/test_indices.pypython/python/tests/test_scalar_index.pypython/src/bitmap.rspython/src/indices.rspython/src/lib.rspython/src/transaction.rs
DataOverlayFile.offsets is a low-level API and duplicate offsets are intentionally left unvalidated (a RoaringBitmap is a set, so a repeat silently collapses to one entry). Make that an explicit, documented trade-off rather than an easy-to-miss implementation detail: spell it out in the offsets docstring and cross-reference it from the Rust-side extract_bitmap comment so the two don't drift out of sync.
wjones127
marked this pull request as ready for review
July 20, 2026 18:57
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
python/src/transaction.rs (2)
232-240: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winDon't consume one-shot sparse iterables during shape detection.
The dense probe walksoffsetsfirst; if the first element is itself iterable, the sparse fallback reuses the same Python iterator after that item has already been consumed. A generator of per-field coverages can lose its first field and misalign the overlay. Buffer the outer iterable once before trying both interpretations, and add a generator-backed regression test.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/src/transaction.rs` around lines 232 - 240, Update the coverage-shape detection in the transaction overlay construction to materialize the outer offsets iterable once before calling extract_bitmap or extract_sparse_bitmaps, then reuse the buffered values for both interpretations. Preserve dense and sparse behavior while ensuring one-shot generators retain their first per-field coverage, and add a regression test using a generator-backed sparse iterable.
208-215: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winPreserve the inner conversion error for
DataOverlayFile.offsets. For inputs likeoffsets=[-1], the dense probe has the bad value/type, but that error gets dropped and the sparse fallback ends up with the generic shape message. Wrap or propagate the original conversion error, and add cases for negative, out-of-range, and non-integer offsets.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/src/transaction.rs` around lines 208 - 215, Update extract_bitmap so the failed dense PyBitmap extraction error is preserved when the iterable fallback also fails, including for DataOverlayFile.offsets inputs such as negative, out-of-range, or non-integer values. Ensure the resulting error reports the inner element conversion failure rather than the generic shape error, and add coverage for all three invalid offset categories.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@python/src/transaction.rs`:
- Around line 232-240: Update the coverage-shape detection in the transaction
overlay construction to materialize the outer offsets iterable once before
calling extract_bitmap or extract_sparse_bitmaps, then reuse the buffered values
for both interpretations. Preserve dense and sparse behavior while ensuring
one-shot generators retain their first per-field coverage, and add a regression
test using a generator-backed sparse iterable.
- Around line 208-215: Update extract_bitmap so the failed dense PyBitmap
extraction error is preserved when the iterable fallback also fails, including
for DataOverlayFile.offsets inputs such as negative, out-of-range, or
non-integer values. Ensure the resulting error reports the inner element
conversion failure rather than the generic shape error, and add coverage for all
three invalid offset categories.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: QUIET
Plan: Pro Plus
Run ID: b3bbfcc8-686b-4e0d-8b1e-73bccef75a20
📒 Files selected for processing (2)
python/python/lance/dataset.pypython/src/transaction.rs
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Two sites in
python/src/transaction.rscopy aRoaringBitmapthrough a Pythonset/liston every Rust↔Python crossing:IndexMetadata.fragment_bitmap↔Index.fragment_idsDataOverlayFile'sOverlayCoverage↔offsets(added in feat(python): expose DataOverlay commit operation #7540), whosebitmap_from_sorted_offsetshelper has a comment pointing straight at this issueThis adds a
Bitmaptype (lance.bitmap.bitmap) wrappingArc<RoaringBitmap>:range, or an integer pyarrowArray/ChunkedArray(read via a zero-copy buffer, not per-value Python objects).__len__/__contains__/__iter__/__eq__,add/discard/update(copy-on-write viaArc::make_mut), and pickling.Both
transaction.rssites now emitBitmapon output and accept it directly on input (anArcclone instead of a Python-set round trip), while still accepting a plainset/liston input for backward compatibility. The same set-expansion pattern existed in two more places —IndexSegment.fragment_idsandIndexSegmentDescription.fragment_idsinindices.rs(used bylist_indices()/describe_indices()) — so those were switched toBitmaptoo.FragmentMetadata.to_json()normalizesBitmap/List[Bitmap]offsets back to plain (nested) lists of ints for JSON serialization.Fixes #7695
Test plan
python/tests/test_bitmap.py: construction (list/range/pyarrow array of various int widths/chunked array),len/iter/in, equality, pickle round-trip, copy-on-write mutation isolation, error cases (negative values, out-of-range values, nulls).test_dataset.pyoverlay tests to coverBitmap-typedoffsets(dense + sparse) and assert the output type; updated the invalid-offsets error-message assertion.test_indices.py/test_scalar_index.pyto assertIndex.fragment_ids/IndexInformation.fragment_idsareBitmap.cargo fmt --all,cargo clippy -p pylance --lib --tests -- -D warnings,uv run ruff format/ruff check,uv run pyrightall clean.