| # Licensed to the Apache Software Foundation (ASF) under one | |
| # or more contributor license agreements. See the NOTICE file | |
| # distributed with this work for additional information | |
| # regarding copyright ownership. The ASF licenses this file | |
| # to you under the Apache License, Version 2.0 (the | |
| # "License"); you may not use this file except in compliance | |
| # with the License. You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, | |
| # software distributed under the License is distributed on an | |
| # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY | |
| # KIND, either express or implied. See the License for the | |
| # specific language governing permissions and limitations | |
| # under the License. | |
| """ | |
| Roaring Bitmap. | |
| """ | |
| from typing import Iterator | |
| class RoaringBitmap64: | |
| """ | |
| A 64-bit roaring bitmap implementation. | |
| This class provides efficient storage and operations for sets of 64-bit integers. | |
| It uses pyroaring.BitMap64 for better performance and memory efficiency. | |
| """ | |
| def __init__(self): | |
| from pyroaring import BitMap64 | |
| self._data = BitMap64() | |
| def add(self, value: int) -> None: | |
| """Add a single value to the bitmap.""" | |
| self._data.add(value) | |
| def add_range(self, from_: int, to: int) -> None: | |
| """Add a range of values [from_, to] to the bitmap.""" | |
| self._data.add_range(from_, to + 1) | |
| def contains(self, value: int) -> bool: | |
| """Check if the bitmap contains the given value.""" | |
| return value in self._data | |
| def is_empty(self) -> bool: | |
| """Check if the bitmap is empty.""" | |
| return len(self._data) == 0 | |
| def cardinality(self) -> int: | |
| """Return the number of elements in the bitmap.""" | |
| return len(self._data) | |
| def __iter__(self) -> Iterator[int]: | |
| """Iterate over all values in the bitmap in sorted order.""" | |
| return iter(self._data) | |
| def __len__(self) -> int: | |
| """Return the number of elements in the bitmap.""" | |
| return len(self._data) | |
| def __contains__(self, value: int) -> bool: | |
| """Check if the bitmap contains the given value.""" | |
| return self.contains(value) | |
| def clear(self) -> None: | |
| """Clear all values from the bitmap.""" | |
| self._data.clear() | |
| def to_list(self) -> list: | |
| """Return a sorted list of all values in the bitmap.""" | |
| return list(self._data) | |
| def to_range_list(self) -> list: | |
| """ | |
| Convert the bitmap to a list of Range objects. | |
| """ | |
| from pypaimon.utils.range import Range | |
| if self.is_empty(): | |
| return [] | |
| ranges = [] | |
| iterator = iter(self._data) | |
| start = next(iterator) | |
| end = start | |
| for value in iterator: | |
| if value == end + 1: | |
| # Consecutive, extend the range | |
| end = value | |
| else: | |
| # Gap, close current range and start new one | |
| ranges.append(Range(start, end)) | |
| start = value | |
| end = start | |
| # Add the last range | |
| ranges.append(Range(start, end)) | |
| return ranges | |
| @staticmethod | |
| def and_(a: 'RoaringBitmap64', b: 'RoaringBitmap64') -> 'RoaringBitmap64': | |
| """Return the intersection of two bitmaps.""" | |
| result = RoaringBitmap64() | |
| result._data = a._data & b._data | |
| return result | |
| @staticmethod | |
| def or_(a: 'RoaringBitmap64', b: 'RoaringBitmap64') -> 'RoaringBitmap64': | |
| """Return the union of two bitmaps.""" | |
| result = RoaringBitmap64() | |
| result._data = a._data | b._data | |
| return result | |
| @staticmethod | |
| def remove_all(a: 'RoaringBitmap64', b: 'RoaringBitmap64') -> 'RoaringBitmap64': | |
| result = RoaringBitmap64() | |
| result._data = a._data - b._data | |
| return result | |
| def remove_all_inplace(self, other: 'RoaringBitmap64') -> None: | |
| """Remove all values contained in ``other`` from this bitmap, in place.""" | |
| self._data -= other._data | |
| def serialize(self) -> bytes: | |
| """Serialize the bitmap to bytes.""" | |
| return self._data.serialize() | |
| @staticmethod | |
| def deserialize(data: bytes) -> 'RoaringBitmap64': | |
| """Deserialize a bitmap from bytes.""" | |
| result = RoaringBitmap64() | |
| from pyroaring import BitMap64 | |
| result._data = BitMap64.deserialize(data) | |
| return result | |
| def __eq__(self, other: object) -> bool: | |
| if not isinstance(other, RoaringBitmap64): | |
| return False | |
| return self._data == other._data | |
| def __hash__(self) -> int: | |
| return hash(tuple(sorted(self._data))) | |
| def __repr__(self) -> str: | |
| values = list(self._data) | |
| if len(values) <= 10: | |
| return f"RoaringBitmap64({values})" | |
| return f"RoaringBitmap64({len(values)} elements)" | |
| class RoaringBitmap: | |
| """ | |
| A 32-bit roaring bitmap implementation. | |
| This class provides efficient storage and operations for sets of 32-bit integers. | |
| It uses pyroaring.BitMap for better performance and memory efficiency. | |
| """ | |
| def __init__(self): | |
| from pyroaring import BitMap | |
| self._data = BitMap() | |
| def add(self, value: int) -> None: | |
| """Add a single value to the bitmap.""" | |
| self._data.add(value) | |
| def add_range(self, from_: int, to: int) -> None: | |
| """Add a range of values [from_, to] to the bitmap.""" | |
| self._data.add_range(from_, to + 1) | |
| def contains(self, value: int) -> bool: | |
| """Check if the bitmap contains the given value.""" | |
| return value in self._data | |
| def is_empty(self) -> bool: | |
| """Check if the bitmap is empty.""" | |
| return len(self._data) == 0 | |
| def cardinality(self) -> int: | |
| """Return the number of elements in the bitmap.""" | |
| return len(self._data) | |
| def __iter__(self) -> Iterator[int]: | |
| """Iterate over all values in the bitmap in sorted order.""" | |
| return iter(self._data) | |
| def __len__(self) -> int: | |
| """Return the number of elements in the bitmap.""" | |
| return len(self._data) | |
| def __contains__(self, value: int) -> bool: | |
| """Check if the bitmap contains the given value.""" | |
| return self.contains(value) | |
| def clear(self) -> None: | |
| """Clear all values from the bitmap.""" | |
| self._data.clear() | |
| def to_list(self) -> list: | |
| """Return a sorted list of all values in the bitmap.""" | |
| return list(self._data) | |
| @staticmethod | |
| def and_(a: 'RoaringBitmap', b: 'RoaringBitmap') -> 'RoaringBitmap': | |
| """Return the intersection of two bitmaps.""" | |
| result = RoaringBitmap() | |
| result._data = a._data & b._data | |
| return result | |
| @staticmethod | |
| def or_(a: 'RoaringBitmap', b: 'RoaringBitmap') -> 'RoaringBitmap': | |
| """Return the union of two bitmaps.""" | |
| result = RoaringBitmap() | |
| result._data = a._data | b._data | |
| return result | |
| @staticmethod | |
| def remove_all(a: 'RoaringBitmap', b: 'RoaringBitmap') -> 'RoaringBitmap': | |
| result = RoaringBitmap() | |
| result._data = a._data - b._data | |
| return result | |
| def serialize(self) -> bytes: | |
| """Serialize the bitmap to bytes.""" | |
| return self._data.serialize() | |
| @staticmethod | |
| def deserialize(data: bytes) -> 'RoaringBitmap': | |
| """Deserialize a bitmap from bytes.""" | |
| result = RoaringBitmap() | |
| from pyroaring import BitMap | |
| result._data = BitMap.deserialize(data) | |
| return result | |
| def __eq__(self, other: object) -> bool: | |
| if not isinstance(other, RoaringBitmap): | |
| return False | |
| return self._data == other._data | |
| def __hash__(self) -> int: | |
| return hash(tuple(sorted(self._data))) | |
| def __repr__(self) -> str: | |
| values = list(self._data) | |
| if len(values) <= 10: | |
| return f"RoaringBitmap({values})" | |
| return f"RoaringBitmap({len(values)} elements)" |