[python] Use pyroaring.BitMap64 for better performance

Author: JingsongLi

.github/workflows/paimon-python-checks.yml

@@ -135,7 +135,7 @@ jobs:
           else
             python -m pip install --upgrade pip
             pip install torch --index-url https://download.pytorch.org/whl/cpu
-            python -m pip install pyroaring readerwriterlock==1.0.9 fsspec==2024.3.1 cachetools==5.3.3 ossfs==2023.12.0 ray==2.48.0 fastavro==1.11.1 pyarrow==16.0.0 zstandard==0.24.0 polars==1.32.0 duckdb==1.3.2 numpy==1.24.3 pandas==2.0.3 pylance==0.39.0 cramjam flake8==4.0.1 pytest~=7.0 py4j==0.10.9.9 requests parameterized==0.9.0 faiss-cpu==1.7.4
+            python -m pip install pyroaring readerwriterlock==1.0.9 fsspec==2024.3.1 cachetools==5.3.3 ossfs==2023.12.0 ray==2.48.0 fastavro==1.11.1 pyarrow==16.0.0 pyroaring==1.0.3 zstandard==0.24.0 polars==1.32.0 duckdb==1.3.2 numpy==1.24.3 pandas==2.0.3 pylance==0.39.0 cramjam flake8==4.0.1 pytest~=7.0 py4j==0.10.9.9 requests parameterized==0.9.0 faiss-cpu==1.7.4
           fi
           df -h
       - name: Run lint-python.sh
@@ -170,7 +170,7 @@ jobs:
         run: |
             python -m pip install --upgrade pip
             pip install torch --index-url https://download.pytorch.org/whl/cpu
-            python -m pip install pyroaring readerwriterlock==1.0.9 fsspec==2024.3.1 cachetools==5.3.3 ossfs==2023.12.0 ray==2.48.0 fastavro==1.11.1 pyarrow==16.0.0 zstandard==0.24.0 polars==1.32.0 duckdb==1.3.2 numpy==1.24.3 pandas==2.0.3 pylance==0.39.0 flake8==4.0.1 pytest~=7.0 py4j==0.10.9.9 requests parameterized==0.9.0
+            python -m pip install pyroaring readerwriterlock==1.0.9 fsspec==2024.3.1 cachetools==5.3.3 ossfs==2023.12.0 ray==2.48.0 fastavro==1.11.1 pyarrow==16.0.0 pyroaring==1.0.3 zstandard==0.24.0 polars==1.32.0 duckdb==1.3.2 numpy==1.24.3 pandas==2.0.3 pylance==0.39.0 flake8==4.0.1 pytest~=7.0 py4j==0.10.9.9 requests parameterized==0.9.0
       - name: Run lint-python.sh
         shell: bash
         run: |
@@ -266,7 +266,7 @@ jobs:
           pip install torch --index-url https://download.pytorch.org/whl/cpu
           python -m pip install --no-cache-dir \
             pyroaring readerwriterlock==1.0.9 fsspec==2024.3.1 cachetools==5.3.3 ossfs==2023.12.0 \
-            fastavro==1.11.1 pyarrow==16.0.0 zstandard==0.24.0 polars==1.32.0 duckdb==1.3.2 \
+            fastavro==1.11.1 pyarrow==16.0.0 pyroaring==1.0.3 zstandard==0.24.0 polars==1.32.0 duckdb==1.3.2 \
             numpy==1.24.3 pandas==2.0.3 cramjam pytest~=7.0 py4j==0.10.9.9 requests \
             parameterized==0.9.0 packaging
       - name: Test Ray version compatibility

paimon-python/pypaimon/globalindex/roaring_bitmap.py

@@ -20,61 +20,98 @@
 Roaring Bitmap.
 """
 
-from typing import Iterator, Set
-import struct
+from typing import Iterator
+
+try:
+    from pyroaring import BitMap64
+    PYROARING_AVAILABLE = True
+except ImportError:
+    PYROARING_AVAILABLE = False
 
 
 class RoaringBitmap64:
     """
     A 64-bit roaring bitmap implementation.
+
     This class provides efficient storage and operations for sets of 64-bit integers.
-    It uses a set-based implementation for simplicity, which can be replaced with
-    a more efficient roaring bitmap library if needed.
+    If pyroaring is available, it uses BitMap64 for better performance and memory efficiency.
+    Otherwise, it falls back to a set-based implementation.
     """
 
     def __init__(self):
-        self._data: Set[int] = set()
+        if PYROARING_AVAILABLE:
+            self._data = BitMap64()
+        else:
+            self._data = set()
+            self._fallback = True
 
     def add(self, value: int) -> None:
         """Add a single value to the bitmap."""
-        self._data.add(value)
+        if PYROARING_AVAILABLE:
+            self._data.add(value)
+        else:
+            self._data.add(value)
 
     def add_range(self, from_: int, to: int) -> None:
         """Add a range of values [from_, to] to the bitmap."""
-        for i in range(from_, to + 1):
-            self._data.add(i)
+        if PYROARING_AVAILABLE:
+            self._data.add_range(from_, to + 1)
+        else:
+            for i in range(from_, to + 1):
+                self._data.add(i)
 
     def contains(self, value: int) -> bool:
         """Check if the bitmap contains the given value."""
-        return value in self._data
+        if PYROARING_AVAILABLE:
+            return value in self._data
+        else:
+            return value in self._data
 
     def is_empty(self) -> bool:
         """Check if the bitmap is empty."""
-        return len(self._data) == 0
+        if PYROARING_AVAILABLE:
+            return len(self._data) == 0
+        else:
+            return len(self._data) == 0
 
     def cardinality(self) -> int:
         """Return the number of elements in the bitmap."""
-        return len(self._data)
+        if PYROARING_AVAILABLE:
+            return len(self._data)
+        else:
+            return len(self._data)
 
     def __iter__(self) -> Iterator[int]:
         """Iterate over all values in the bitmap in sorted order."""
-        return iter(sorted(self._data))
+        if PYROARING_AVAILABLE:
+            return iter(self._data)
+        else:
+            return iter(sorted(self._data))
 
     def __len__(self) -> int:
         """Return the number of elements in the bitmap."""
-        return len(self._data)
+        if PYROARING_AVAILABLE:
+            return len(self._data)
+        else:
+            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()
+        if PYROARING_AVAILABLE:
+            self._data.clear()
+        else:
+            self._data.clear()
 
     def to_list(self) -> list:
         """Return a sorted list of all values in the bitmap."""
-        return sorted(self._data)
+        if PYROARING_AVAILABLE:
+            return list(self._data)
+        else:
+            return sorted(self._data)
 
     def to_range_list(self) -> list:
         """
@@ -85,76 +122,128 @@ def to_range_list(self) -> list:
         if self.is_empty():
             return []
 
-        sorted_values = sorted(self._data)
-        ranges = []
-        start = sorted_values[0]
-        end = start
-
-        for i in range(1, len(sorted_values)):
-            if sorted_values[i] == end + 1:
-                # Consecutive, extend the range
-                end = sorted_values[i]
-            else:
-                # Gap, close current range and start new one
-                ranges.append(Range(start, end))
-                start = sorted_values[i]
-                end = start
-
-        # Add the last range
-        ranges.append(Range(start, end))
-
-        return ranges
+        if PYROARING_AVAILABLE:
+            # Use pyroaring's efficient iteration
+            ranges = []
+            sorted_values = list(self._data)
+            start = sorted_values[0]
+            end = start
+
+            for i in range(1, len(sorted_values)):
+                if sorted_values[i] == end + 1:
+                    # Consecutive, extend the range
+                    end = sorted_values[i]
+                else:
+                    # Gap, close current range and start new one
+                    ranges.append(Range(start, end))
+                    start = sorted_values[i]
+                    end = start
+
+            # Add the last range
+            ranges.append(Range(start, end))
+
+            return ranges
+        else:
+            # Fallback implementation
+            sorted_values = sorted(self._data)
+            ranges = []
+            start = sorted_values[0]
+            end = start
+
+            for i in range(1, len(sorted_values)):
+                if sorted_values[i] == end + 1:
+                    # Consecutive, extend the range
+                    end = sorted_values[i]
+                else:
+                    # Gap, close current range and start new one
+                    ranges.append(Range(start, end))
+                    start = sorted_values[i]
+                    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
+        if PYROARING_AVAILABLE:
+            result._data = a._data & b._data
+        else:
+            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
+        if PYROARING_AVAILABLE:
+            result._data = a._data | b._data
+        else:
+            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
+        if PYROARING_AVAILABLE:
+            result._data = a._data - b._data
+        else:
+            result._data = a._data - b._data
         return result
 
     def serialize(self) -> bytes:
         """Serialize the bitmap to bytes."""
-        # Simple serialization format: count followed by sorted values
-        values = sorted(self._data)
-        data = struct.pack('>Q', len(values))  # 8-byte count
-        for v in values:
-            data += struct.pack('>q', v)  # 8-byte signed value
-        return data
+        if PYROARING_AVAILABLE:
+            return self._data.serialize()
+        else:
+            # Simple serialization format: count followed by sorted values
+            import struct
+            values = sorted(self._data)
+            data = struct.pack('>Q', len(values))  # 8-byte count
+            for v in values:
+                data += struct.pack('>q', v)  # 8-byte signed value
+            return data
 
     @staticmethod
     def deserialize(data: bytes) -> 'RoaringBitmap64':
         """Deserialize a bitmap from bytes."""
         result = RoaringBitmap64()
-        count = struct.unpack('>Q', data[:8])[0]
-        offset = 8
-        for _ in range(count):
-            value = struct.unpack('>q', data[offset:offset + 8])[0]
-            result.add(value)
-            offset += 8
+        if PYROARING_AVAILABLE:
+            result._data = BitMap64.deserialize(data)
+        else:
+            import struct
+            count = struct.unpack('>Q', data[:8])[0]
+            offset = 8
+            for _ in range(count):
+                value = struct.unpack('>q', data[offset:offset + 8])[0]
+                result.add(value)
+                offset += 8
         return result
 
     def __eq__(self, other: object) -> bool:
         if not isinstance(other, RoaringBitmap64):
             return False
-        return self._data == other._data
+        if PYROARING_AVAILABLE:
+            return self._data == other._data
+        else:
+            return self._data == other._data
 
     def __hash__(self) -> int:
-        return hash(frozenset(self._data))
+        if PYROARING_AVAILABLE:
+            return hash(tuple(sorted(self._data)))
+        else:
+            return hash(frozenset(self._data))
 
     def __repr__(self) -> str:
-        if len(self._data) <= 10:
-            return f"RoaringBitmap64({sorted(self._data)})"
-        return f"RoaringBitmap64({len(self._data)} elements)"
+        if PYROARING_AVAILABLE:
+            values = list(self._data)
+            if len(values) <= 10:
+                return f"RoaringBitmap64({values})"
+            return f"RoaringBitmap64({len(values)} elements)"
+        else:
+            if len(self._data) <= 10:
+                return f"RoaringBitmap64({sorted(self._data)})"
+            return f"RoaringBitmap64({len(self._data)} elements)"