diff --git a/python/python/lance/bitmap.py b/python/python/lance/bitmap.py new file mode 100644 index 00000000000..4008d1e3b61 --- /dev/null +++ b/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"] diff --git a/python/python/lance/dataset.py b/python/python/lance/dataset.py index def635d58a6..2a83f43ba5c 100644 --- a/python/python/lance/dataset.py +++ b/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 diff --git a/python/python/lance/fragment.py b/python/python/lance/fragment.py index bca62e2b280..d3313c190ba 100644 --- a/python/python/lance/fragment.py +++ b/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 diff --git a/python/python/lance/lance/__init__.pyi b/python/python/lance/lance/__init__.pyi index 03ec9f46db7..8cf4326f4d9 100644 --- a/python/python/lance/lance/__init__.pyi +++ b/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 diff --git a/python/python/lance/lance/bitmap.pyi b/python/python/lance/lance/bitmap.pyi new file mode 100644 index 00000000000..e18a1faf3e5 --- /dev/null +++ b/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: ... diff --git a/python/python/lance/lance/indices/__init__.pyi b/python/python/lance/lance/indices/__init__.pyi index 181ec93b2d1..a97f6e1066f 100644 --- a/python/python/lance/lance/indices/__init__.pyi +++ b/python/python/lance/lance/indices/__init__.pyi @@ -18,6 +18,7 @@ from typing import Optional import pyarrow as pa from .. import _Fragment +from ..bitmap import Bitmap class IndexConfig: index_type: str @@ -25,7 +26,7 @@ class IndexConfig: 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] diff --git a/python/python/tests/test_bitmap.py b/python/python/tests/test_bitmap.py new file mode 100644 index 00000000000..21413f0e6a9 --- /dev/null +++ b/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]) diff --git a/python/python/tests/test_dataset.py b/python/python/tests/test_dataset.py index 898bdf93069..3d6ca5c1fd6 100644 --- a/python/python/tests/test_dataset.py +++ b/python/python/tests/test_dataset.py @@ -5109,9 +5109,84 @@ def test_data_overlay_sparse_per_field(tmp_path: Path): assert result.column("val").to_pylist()[2] == 20 +def test_data_overlay_offsets_accept_bitmap(tmp_path: Path): + from lance.bitmap import bitmap + + base_dir = tmp_path / "test" + table = pa.table( + { + "id": pa.array(range(10), pa.int32()), + "val": pa.array([i * 10 for i in range(10)], pa.int32()), + } + ) + dataset = lance.write_dataset(table, base_dir) + + # Dense coverage as a single Bitmap. + dense_file = _write_overlay_file( + dataset, + base_dir, + "dense.lance", + pa.table({"val": pa.array([111, 444], pa.int32())}), + fields=[1], + ) + dataset = lance.LanceDataset.commit( + dataset, + lance.LanceOperation.DataOverlay( + [ + lance.LanceOperation.DataOverlayGroup( + 0, + [ + lance.LanceOperation.DataOverlayFile( + dense_file, offsets=bitmap([1, 4]) + ) + ], + ) + ] + ), + read_version=dataset.version, + ) + result = dataset.to_table() + assert result.column("val").to_pylist() == [0, 111, 20, 30, 444, 50, 60, 70, 80, 90] + + # Sparse coverage as a list of Bitmaps. + sparse_file = _write_overlay_file( + dataset, + base_dir, + "sparse.lance", + pa.table( + { + "id": pa.array([777], pa.int32()), + "val": pa.array([330], pa.int32()), + } + ), + fields=[0, 1], + ) + dataset = lance.LanceDataset.commit( + dataset, + lance.LanceOperation.DataOverlay( + [ + lance.LanceOperation.DataOverlayGroup( + 0, + [ + lance.LanceOperation.DataOverlayFile( + sparse_file, offsets=[bitmap([2]), bitmap([3])] + ) + ], + ) + ] + ), + read_version=dataset.version, + ) + result = dataset.to_table() + assert result.column("id").to_pylist()[2] == 777 + assert result.column("val").to_pylist()[3] == 330 + + def test_data_overlay_round_trips_through_fragment_metadata(tmp_path: Path): import json + from lance.bitmap import bitmap as Bitmap + base_dir = tmp_path / "test" table = pa.table( { @@ -5141,6 +5216,7 @@ def test_data_overlay_round_trips_through_fragment_metadata(tmp_path: Path): # Reading the fragment surfaces its overlays, stamped with the commit version. metadata = dataset.get_fragments()[0].metadata assert len(metadata.overlays) == 1 + assert isinstance(metadata.overlays[0].offsets, Bitmap) assert metadata.overlays[0].offsets == [1, 4] assert metadata.overlays[0].committed_version == overlay_version @@ -5174,9 +5250,9 @@ def test_data_overlay_rejects_invalid_offsets(tmp_path: Path): fields=[0], ) - # offsets is neither a flat list of ints (dense) nor a list of per-field int - # lists (sparse), so the coverage shape can't be resolved. - with pytest.raises(ValueError, match="offsets must be a list"): + # offsets is neither an iterable of ints (dense) nor an iterable of int + # iterables (sparse), so the coverage shape can't be resolved. + with pytest.raises(ValueError, match="offsets must be an iterable"): lance.LanceDataset.commit( dataset, lance.LanceOperation.DataOverlay( @@ -5199,15 +5275,14 @@ def test_data_overlay_rejects_invalid_offsets(tmp_path: Path): "offsets", [ [2, 1], # dense, descending - [1, 1], # dense, duplicate [[2, 1]], # sparse, descending - [[1, 1]], # sparse, duplicate ], ) -def test_data_overlay_rejects_unsorted_offsets(tmp_path: Path, offsets): - # Offsets map positionally to value rows in data_file. A RoaringBitmap would - # silently reorder/dedup them, so a non-ascending list must be rejected up - # front rather than corrupting the row mapping. +def test_data_overlay_accepts_any_offset_order(tmp_path: Path, offsets): + """Offsets are always resolved in ascending order (the smallest covered + offset maps to value-file row 0, the next-smallest to row 1, ...) + regardless of what order the caller lists them in — an out-of-order list + is accepted and resolves identically to a pre-sorted one.""" base_dir = tmp_path / "test" table = pa.table({"val": pa.array([0, 1, 2], pa.int32())}) dataset = lance.write_dataset(table, base_dir) @@ -5215,27 +5290,23 @@ def test_data_overlay_rejects_unsorted_offsets(tmp_path: Path, offsets): dataset, base_dir, "ov.lance", - pa.table({"val": pa.array([9, 9], pa.int32())}), + pa.table({"val": pa.array([100, 200], pa.int32())}), fields=[0], ) - with pytest.raises(ValueError, match="strictly ascending"): - lance.LanceDataset.commit( - dataset, - lance.LanceOperation.DataOverlay( - [ - lance.LanceOperation.DataOverlayGroup( - 0, - [ - lance.LanceOperation.DataOverlayFile( - data_file, offsets=offsets - ) - ], - ) - ] - ), - read_version=dataset.version, - ) + overlay = lance.LanceOperation.DataOverlayFile(data_file, offsets=offsets) + dataset = lance.LanceDataset.commit( + dataset, + lance.LanceOperation.DataOverlay( + [lance.LanceOperation.DataOverlayGroup(0, [overlay])] + ), + read_version=dataset.version, + ) + + result = dataset.to_table().column("val").to_pylist() + # rank 0 (offset 1) -> value row 0 (100); rank 1 (offset 2) -> value row 1 (200) + assert result[1] == 100 + assert result[2] == 200 def test_schema_project_drop_column(tmp_path: Path): diff --git a/python/python/tests/test_indices.py b/python/python/tests/test_indices.py index ae51e6ddb0d..48a68583efc 100644 --- a/python/python/tests/test_indices.py +++ b/python/python/tests/test_indices.py @@ -300,7 +300,10 @@ def test_ivf_centroids_multivector_fragment_ids(tmpdir): ivf_centroids=centroids, ) + from lance.bitmap import bitmap as Bitmap + assert index.uuid == "00000000-0000-4000-8000-000000000001" + assert isinstance(index.fragment_ids, Bitmap) assert index.fragment_ids == set(fragment_ids) assert index.name == "embeddings_idx" diff --git a/python/python/tests/test_scalar_index.py b/python/python/tests/test_scalar_index.py index 985072fe5de..253520ab80f 100644 --- a/python/python/tests/test_scalar_index.py +++ b/python/python/tests/test_scalar_index.py @@ -179,6 +179,8 @@ def test_list_indices_characterization(indexed_dataset: lance.LanceDataset): Index dataclasses. This characterization test guards the dict keys and values so the deprecated method stays backwards compatible. """ + from lance.bitmap import bitmap as Bitmap + with pytest.warns(DeprecationWarning): indices = indexed_dataset.list_indices() @@ -199,7 +201,7 @@ def test_list_indices_characterization(indexed_dataset: lance.LanceDataset): assert set(idx) == expected_keys assert isinstance(idx["uuid"], str) and len(idx["uuid"]) > 0 assert isinstance(idx["fields"], list) - assert isinstance(idx["fragment_ids"], set) + assert isinstance(idx["fragment_ids"], Bitmap) assert isinstance(idx["version"], int) assert idx["type"] != "Unknown" assert idx["base_id"] is None diff --git a/python/src/bitmap.rs b/python/src/bitmap.rs new file mode 100644 index 00000000000..29b6c0d78c4 --- /dev/null +++ b/python/src/bitmap.rs @@ -0,0 +1,262 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::sync::Arc; + +use arrow::pyarrow::FromPyArrow; +use arrow_array::{cast::AsArray, make_array}; +use arrow_data::ArrayData; +use arrow_schema::DataType; +use pyo3::basic::CompareOp; +use pyo3::exceptions::{PyNotImplementedError, PyValueError}; +use pyo3::prelude::*; +use pyo3::types::PyBytes; +use roaring::RoaringBitmap; + +/// A lazy, streaming iterator over a `Bitmap`'s values — yields one Python +/// `int` per `__next__` call rather than materializing them all up front. +#[pyclass(name = "BitmapIterator", module = "lance.bitmap")] +pub struct PyBitmapIter(roaring::bitmap::IntoIter); + +#[pymethods] +impl PyBitmapIter { + fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> { + slf + } + + fn __next__(mut slf: PyRefMut<'_, Self>) -> Option { + slf.0.next() + } +} + +/// A set of non-negative integers backed by a `RoaringBitmap`. +/// +/// Cheap to clone (an `Arc` bump) 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. +#[pyclass(name = "Bitmap", module = "lance.bitmap", from_py_object)] +#[derive(Clone, Debug, Default)] +pub struct PyBitmap(pub Arc); + +impl PyBitmap { + pub fn new(bitmap: RoaringBitmap) -> Self { + Self(Arc::new(bitmap)) + } +} + +fn i64_to_u32(value: i64) -> PyResult { + u32::try_from(value).map_err(|_| { + PyValueError::new_err(format!( + "Bitmap values must fit in an unsigned 32-bit integer, got {value}" + )) + }) +} + +fn u64_to_u32(value: u64) -> PyResult { + u32::try_from(value).map_err(|_| { + PyValueError::new_err(format!( + "Bitmap values must fit in an unsigned 32-bit integer, got {value}" + )) + }) +} + +/// Read an integer pyarrow array's values into a `RoaringBitmap`, without +/// going through per-value Python objects or an intermediate `Vec` — each +/// arm collects straight from the array's native buffer into the bitmap. +fn bitmap_from_pyarrow(ob: &Bound<'_, PyAny>) -> PyResult { + let data = ArrayData::from_pyarrow_bound(ob)?; + if data.null_count() > 0 { + return Err(PyValueError::new_err( + "Bitmap cannot be constructed from an array containing nulls", + )); + } + let array = make_array(data); + match array.data_type() { + DataType::UInt8 => Ok(array + .as_primitive::() + .values() + .iter() + .map(|v| *v as u32) + .collect()), + DataType::UInt16 => Ok(array + .as_primitive::() + .values() + .iter() + .map(|v| *v as u32) + .collect()), + DataType::UInt32 => Ok(array + .as_primitive::() + .values() + .iter() + .copied() + .collect()), + DataType::UInt64 => array + .as_primitive::() + .values() + .iter() + .map(|v| u64_to_u32(*v)) + .collect(), + DataType::Int8 => array + .as_primitive::() + .values() + .iter() + .map(|v| i64_to_u32(*v as i64)) + .collect(), + DataType::Int16 => array + .as_primitive::() + .values() + .iter() + .map(|v| i64_to_u32(*v as i64)) + .collect(), + DataType::Int32 => array + .as_primitive::() + .values() + .iter() + .map(|v| i64_to_u32(*v as i64)) + .collect(), + DataType::Int64 => array + .as_primitive::() + .values() + .iter() + .map(|v| i64_to_u32(*v)) + .collect(), + other => Err(PyValueError::new_err(format!( + "Bitmap can only be constructed from an integer pyarrow array, got {other:?}" + ))), + } +} + +#[pymethods] +impl PyBitmap { + /// Construct a Bitmap from an iterable of non-negative ints (list, set, + /// range, generator, ...) or an integer pyarrow Array/ChunkedArray. + #[new] + #[pyo3(signature = (values=None))] + fn new_py(values: Option<&Bound<'_, PyAny>>) -> PyResult { + let Some(values) = values else { + return Ok(Self::default()); + }; + if let Ok(existing) = values.extract::() { + return Ok(existing); + } + // A pyarrow `Array` supports the `__arrow_c_array__` Arrow C Data + // Interface export directly; a `ChunkedArray` doesn't (it only + // supports the streaming `__arrow_c_stream__` form), so detect it by + // its `combine_chunks()` method and flatten it to a single `Array` + // first. + if values.hasattr("combine_chunks")? { + let combined = values.call_method0("combine_chunks")?; + return Ok(Self::new(bitmap_from_pyarrow(&combined)?)); + } + if values.hasattr("__arrow_c_array__")? { + return Ok(Self::new(bitmap_from_pyarrow(values)?)); + } + let bitmap = values + .try_iter()? + .map(|item| item?.extract::()) + .collect::>()?; + Ok(Self::new(bitmap)) + } + + fn __len__(&self) -> usize { + self.0.len() as usize + } + + fn __contains__(&self, value: &Bound<'_, PyAny>) -> PyResult { + match value.extract::() { + Ok(v) => Ok(self.0.contains(v)), + Err(_) => Ok(false), + } + } + + fn __iter__(&self, py: Python<'_>) -> PyResult> { + // Cloning the bitmap here is a native Rust-side copy of its compressed + // containers, not a per-value Python allocation — the point is that + // `PyBitmapIter` then streams values lazily instead of eagerly + // building a Python list of boxed ints up front. + Py::new(py, PyBitmapIter((*self.0).clone().into_iter())) + } + + fn __repr__(&self) -> String { + const MAX_VALUES_SHOWN: usize = 20; + let len = self.0.len(); + let values: Vec = self + .0 + .iter() + .take(MAX_VALUES_SHOWN) + .map(|v| v.to_string()) + .collect(); + if (len as usize) > MAX_VALUES_SHOWN { + format!("Bitmap({{{}, ...}}, len={})", values.join(", "), len) + } else { + format!("Bitmap({{{}}})", values.join(", ")) + } + } + + fn __richcmp__(&self, other: &Bound<'_, PyAny>, op: CompareOp) -> PyResult { + // A value that isn't a Bitmap or an iterable of ints (e.g. `5`) is + // simply unequal, matching `__contains__` and normal Python `==` + // semantics — it shouldn't raise just because the type differs. + let other_bitmap = if let Ok(other) = other.extract::() { + Some((*other.0).clone()) + } else { + other.try_iter().ok().and_then(|it| { + it.map(|item| item?.extract::()) + .collect::>() + .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", + )), + } + } + + /// Add a value, cloning the underlying bitmap first if it is shared with + /// another `Bitmap`. + fn add(&mut self, value: u32) { + Arc::make_mut(&mut self.0).insert(value); + } + + /// Remove a value if present, cloning the underlying bitmap first if it + /// is shared with another `Bitmap`. + fn discard(&mut self, value: u32) { + Arc::make_mut(&mut self.0).remove(value); + } + + /// Add all values from an iterable, cloning the underlying bitmap first + /// if it is shared with another `Bitmap`. + fn update(&mut self, values: &Bound<'_, PyAny>) -> PyResult<()> { + let bitmap = Arc::make_mut(&mut self.0); + for item in values.try_iter()? { + bitmap.insert(item?.extract::()?); + } + Ok(()) + } + + fn __reduce__(&self, py: Python<'_>) -> PyResult<(Py, (Py,))> { + let mut buf = Vec::new(); + self.0 + .serialize_into(&mut buf) + .map_err(|e| PyValueError::new_err(format!("Failed to serialize Bitmap: {e}")))?; + let ctor = py + .import("lance.bitmap")? + .getattr("Bitmap")? + .getattr("_from_bytes")?; + let bytes = PyBytes::new(py, &buf); + Ok((ctor.unbind(), (bytes.into(),))) + } + + #[staticmethod] + fn _from_bytes(data: &[u8]) -> PyResult { + let bitmap = RoaringBitmap::deserialize_from(data) + .map_err(|e| PyValueError::new_err(format!("Failed to deserialize Bitmap: {e}")))?; + Ok(Self::new(bitmap)) + } +} diff --git a/python/src/indices.rs b/python/src/indices.rs index 1d5c476d795..261a21d06bb 100644 --- a/python/src/indices.rs +++ b/python/src/indices.rs @@ -1,7 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The Lance Authors -use std::collections::HashSet; use std::fmt::Write; use std::sync::Arc; @@ -81,8 +80,8 @@ impl PyIndexSegment { } #[getter] - fn fragment_ids(&self) -> HashSet { - self.inner.fragment_bitmap().iter().collect() + fn fragment_ids(&self) -> crate::bitmap::PyBitmap { + crate::bitmap::PyBitmap::new(self.inner.fragment_bitmap().clone()) } #[getter] @@ -91,10 +90,11 @@ impl PyIndexSegment { } fn __repr__(&self) -> String { + let fragment_ids: Vec = self.inner.fragment_bitmap().iter().collect(); format!( "IndexSegment(uuid={}, fragment_ids={:?}, index_version={})", self.uuid(), - self.fragment_ids(), + fragment_ids, self.index_version() ) } @@ -627,7 +627,7 @@ pub struct PyIndexSegmentDescription { /// The dataset version at which the index segment was last updated pub dataset_version_at_last_update: u64, /// The fragment ids that are covered by the index segment - pub fragment_ids: HashSet, + pub fragment_ids: crate::bitmap::PyBitmap, /// The version of the index pub index_version: i32, /// The timestamp when the index segment was created @@ -642,11 +642,8 @@ pub struct PyIndexSegmentDescription { impl PyIndexSegmentDescription { pub fn from_metadata(segment: &lance_table::format::IndexMetadata) -> Self { - let fragment_ids = segment - .fragment_bitmap - .as_ref() - .map(|bitmap| bitmap.iter().collect::>()) - .unwrap_or_default(); + let fragment_ids = + crate::bitmap::PyBitmap::new(segment.fragment_bitmap.clone().unwrap_or_default()); let size_bytes = segment.total_size_bytes(); Self { @@ -661,11 +658,12 @@ impl PyIndexSegmentDescription { } pub fn __repr__(&self) -> String { + let fragment_ids: Vec = self.fragment_ids.0.iter().collect(); format!( "IndexSegmentDescription(uuid={}, dataset_version_at_last_update={}, fragment_ids={:?}, index_version={}, created_at={:?}, size_bytes={:?}, base_id={:?})", self.uuid, self.dataset_version_at_last_update, - self.fragment_ids, + fragment_ids, self.index_version, self.created_at, self.size_bytes, diff --git a/python/src/lib.rs b/python/src/lib.rs index c5631d26f10..b82a13a6920 100644 --- a/python/src/lib.rs +++ b/python/src/lib.rs @@ -62,6 +62,7 @@ use std::ffi::CString; use std::ptr::NonNull; pub(crate) mod arrow; +pub(crate) mod bitmap; pub(crate) mod blob; #[cfg(feature = "datagen")] pub(crate) mod datagen; @@ -292,6 +293,8 @@ fn lance(py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; + m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; diff --git a/python/src/transaction.rs b/python/src/transaction.rs index 37085d60de7..16acfa19c0e 100644 --- a/python/src/transaction.rs +++ b/python/src/transaction.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The Lance Authors +use crate::bitmap::PyBitmap; use crate::dataset::DatasetBasePath; use crate::schema::LanceSchema; use crate::utils::{PyLance, class_name, export_vec, extract_vec}; @@ -14,7 +15,6 @@ use lance::datatypes::Schema; use lance_table::format::overlay::{DataOverlayFile, OverlayCoverage}; use lance_table::format::{BasePath, DataFile, Fragment, IndexFile, IndexMetadata}; use pyo3::exceptions::PyValueError; -use pyo3::types::PySet; use pyo3::{Bound, FromPyObject, PyAny, PyResult, Python}; use pyo3::{intern, prelude::*}; use roaring::RoaringBitmap; @@ -71,13 +71,7 @@ impl FromPyObject<'_, '_> for PyLance { let fragment_ids = ob.getattr("fragment_ids")?; let created_at = ob.getattr("created_at")?.extract()?; - let fragment_ids_ref: &Bound<'_, PySet> = fragment_ids.cast()?; - let fragment_bitmap = Some( - fragment_ids_ref - .into_iter() - .map(|id| id.extract::()) - .collect::>()?, - ); + let fragment_bitmap = Some(extract_bitmap(&fragment_ids)?); let base_id: Option = ob .getattr("base_id")? .extract::>()? @@ -124,16 +118,7 @@ impl<'py> IntoPyObject<'py> for PyLance<&IndexMetadata> { let fields = &self.0.fields; let dataset_version = self.0.dataset_version; let index_version = self.0.index_version; - let fragment_ids = self.0.fragment_bitmap.as_ref().map_or_else( - || PySet::empty(py).unwrap(), - |bitmap| { - let set = PySet::empty(py).unwrap(); - for id in bitmap.iter() { - set.add(id).unwrap(); - } - set - }, - ); + let fragment_ids = PyBitmap::new(self.0.fragment_bitmap.clone().unwrap_or_default()); let created_at = self.0.created_at; let base_id = self.0.base_id.map(|id| id as i64); let files = self @@ -207,19 +192,35 @@ impl<'py> IntoPyObject<'py> for PyLance<&DataReplacementGroup> { } } -// The Nth offset in an overlay list positionally maps to the Nth value row in -// `data_file`, but `RoaringBitmap` stores offsets in ascending order and drops -// duplicates. A caller-supplied list that isn't strictly ascending would be -// silently reordered, breaking that mapping, so reject it here instead. This can -// go away once we expose RoaringBitmap directly to Python (issue #7695). -fn bitmap_from_sorted_offsets(offsets: Vec) -> PyResult { - if offsets.windows(2).any(|w| w[0] >= w[1]) { - return Err(PyValueError::new_err( - "DataOverlayFile.offsets must be strictly ascending with no duplicates; \ - each offset positionally maps to a value row in data_file", - )); +// Accept either a `Bitmap` (cheap `Arc` clone) or any other iterable of ints +// (a `set`, `list`, etc.), collected in whatever order the iterable yields — +// `RoaringBitmap` itself defines the canonical (ascending, deduplicated) +// order, so there's no separate ordering contract to validate here. +// +// Deliberate footgun, not an oversight: a `RoaringBitmap` is a set, so a +// duplicate value in the input silently collapses to one entry. That's +// harmless for `IndexMetadata.fragment_bitmap` (just a set of fragment +// ids), but for `DataOverlayFile.offsets` (see its docstring in +// dataset.py) it would silently shift every later offset onto the wrong +// row of the value file. This is intentionally left unvalidated — it's a +// low-level API and callers are expected to pass distinct offsets — but +// don't remove this comment without also updating that docstring. +fn extract_bitmap(ob: &Bound<'_, PyAny>) -> PyResult { + if let Ok(bitmap) = ob.extract::() { + return Ok((*bitmap.0).clone()); } - Ok(RoaringBitmap::from_sorted_iter(offsets).expect("offsets verified strictly ascending")) + ob.try_iter()? + .map(|item| item?.extract::()) + .collect::>() +} + +// Extract `offsets` as a sparse (per-field) coverage: an iterable of +// `Bitmap`s/int iterables, one per field. +fn extract_sparse_bitmaps(offsets: &Bound<'_, PyAny>) -> PyResult> { + offsets + .try_iter()? + .map(|item| extract_bitmap(&item?)) + .collect() } impl FromPyObject<'_, '_> for PyLance { @@ -228,22 +229,19 @@ impl FromPyObject<'_, '_> for PyLance { let data_file = ob.getattr("data_file")?.extract::>()?.0; let offsets = ob.getattr("offsets")?; - // A flat list of offsets is a dense overlay (one coverage shared by every - // field); a list of per-field lists is a sparse overlay. Differentiate by - // shape, trying the dense form first. - let coverage = if let Ok(shared) = offsets.extract::>() { - OverlayCoverage::dense(bitmap_from_sorted_offsets(shared)?) - } else if let Ok(per_field) = offsets.extract::>>() { - OverlayCoverage::sparse( - per_field - .into_iter() - .map(bitmap_from_sorted_offsets) - .collect::>>()?, - ) + // A `Bitmap`/flat iterable of ints is a dense overlay (one coverage + // shared by every field); an iterable of `Bitmap`/int iterables is a + // sparse overlay (one per field). Differentiate by shape, trying the + // dense form first: it fails to extract if `offsets`' elements aren't + // ints (e.g. they're themselves iterables), falling through to sparse. + let coverage = if let Ok(bitmap) = extract_bitmap(&offsets) { + OverlayCoverage::dense(bitmap) + } else if let Ok(per_field) = extract_sparse_bitmaps(&offsets) { + OverlayCoverage::sparse(per_field) } else { return Err(PyValueError::new_err( - "DataOverlayFile.offsets must be a list of ints (dense coverage shared by \ - every field) or a list of per-field int lists (sparse coverage)", + "DataOverlayFile.offsets must be an iterable of ints (dense coverage shared \ + by every field) or an iterable of per-field int iterables (sparse coverage)", )); }; @@ -281,15 +279,17 @@ impl<'py> IntoPyObject<'py> for PyLance<&DataOverlayFile> { let committed_version = self.0.committed_version; - // Mirror the read side: a dense overlay becomes a flat list of offsets, a - // sparse overlay a list of per-field lists. + // Mirror the read side: a dense overlay becomes a single Bitmap, a + // sparse overlay a list of per-field Bitmaps. match &self.0.coverage { OverlayCoverage::Shared(bitmap) => { - let offsets: Vec = bitmap.iter().collect(); + // `bitmap` is already an `Arc` — cloning it is a + // cheap refcount bump, not a deep copy. + let offsets = PyBitmap(bitmap.clone()); cls.call1((data_file, offsets, committed_version)) } OverlayCoverage::PerField(bitmaps) => { - let offsets: Vec> = bitmaps.iter().map(|b| b.iter().collect()).collect(); + let offsets: Vec = bitmaps.iter().cloned().map(PyBitmap).collect(); cls.call1((data_file, offsets, committed_version)) } }