[python] Support pyroaring without BitMap64

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

@@ -128,12 +128,12 @@ jobs:
           if [[ "${{ matrix.python-version }}" == "3.6.15" ]]; then
             python -m pip install --upgrade pip==21.3.1
             python --version
-            python -m pip install --no-cache-dir pyroaring readerwriterlock==1.0.9 'fsspec==2021.10.1' 'cachetools==4.2.4' 'ossfs==2021.8.0' pyarrow==6.0.1 pandas==1.1.5 'polars==0.9.12' 'fastavro==1.4.7' zstandard==0.19.0 dataclasses==0.8.0 flake8 pytest py4j==0.10.9.9 requests parameterized==0.8.1 2>&1 >/dev/null
+            python -m pip install --no-cache-dir 'pyroaring>=0.3.3,<0.3.4' readerwriterlock==1.0.9 'fsspec==2021.10.1' 'cachetools==4.2.4' 'ossfs==2021.8.0' pyarrow==6.0.1 pandas==1.1.5 'polars==0.9.12' 'fastavro==1.4.7' zstandard==0.19.0 dataclasses==0.8.0 flake8 pytest py4j==0.10.9.9 requests parameterized==0.8.1 2>&1 >/dev/null
             python -m pip install 'lumina-data>=${{ env.LUMINA_DATA_VERSION }}' -i https://pypi.org/simple/
           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.54.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 'daft>=0.7.6' pypaimon-rust==0.2.0 'datafusion>=52'
+            python -m pip install 'pyroaring>=1.0.0' readerwriterlock==1.0.9 fsspec==2024.3.1 cachetools==5.3.3 ossfs==2023.12.0 ray==2.54.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 'daft>=0.7.6' pypaimon-rust==0.2.0 'datafusion>=52'
             python -m pip install 'lumina-data>=${{ env.LUMINA_DATA_VERSION }}' -i https://pypi.org/simple/
             if python -c "import sys; sys.exit(0 if sys.version_info >= (3, 11) else 1)"; then
               python -m pip install vortex-data==0.70.0
@@ -185,7 +185,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.54.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>=1.0.0' readerwriterlock==1.0.9 fsspec==2024.3.1 cachetools==5.3.3 ossfs==2023.12.0 ray==2.54.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 'lumina-data>=${{ env.LUMINA_DATA_VERSION }}' -i https://pypi.org/simple/
       - name: Run lint-python.sh
         shell: bash
@@ -281,7 +281,7 @@ jobs:
           python -m pip install --upgrade pip
           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 \
+            'pyroaring>=1.0.0' 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 \
             numpy==1.24.3 pandas==2.0.3 cramjam pytest~=7.0 py4j==0.10.9.9 requests \
             parameterized==0.9.0 packaging

paimon-python/dev/requirements.txt

@@ -30,8 +30,8 @@ polars>=1,<2; python_version=="3.8"
 polars>=1.32,<2; python_version>="3.9"
 pyarrow>=6,<7; python_version < "3.8"
 pyarrow>=16,!=19.0.0,<20; python_version >= "3.8"
-pyroaring<=0.3.3; python_version < "3.7"
-pyroaring<=0.4.5; python_version == "3.7"
+pyroaring>=0.3.3,<0.3.4; python_version < "3.7"
+pyroaring>=0.4.5,<1.0.0; python_version >= "3.7" and python_version < "3.8"
 pyroaring>=1.0.0; python_version >= "3.8"
 readerwriterlock>=1,<2
 requests>=2.21.0,<3

paimon-python/pypaimon/tests/roaring_bitmap_test.py

@@ -0,0 +1,78 @@
+# 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.
+
+import unittest
+
+from pypaimon.utils import roaring_bitmap
+
+
+class RoaringBitmap64FallbackTest(unittest.TestCase):
+    def _bitmap(self):
+        return roaring_bitmap._ChunkedBitMap64()
+
+    def test_fallback_supports_values_across_32_bit_chunks(self):
+        bitmap = self._bitmap()
+        bitmap.add(1)
+        bitmap.add((1 << 32) + 2)
+        bitmap.add((2 << 32) + 3)
+
+        self.assertEqual([1, (1 << 32) + 2, (2 << 32) + 3], list(bitmap))
+        self.assertIn((1 << 32) + 2, bitmap)
+        self.assertEqual(3, len(bitmap))
+
+    def test_fallback_add_range_splits_chunks(self):
+        bitmap = self._bitmap()
+        bitmap.add_range((1 << 32) - 2, (1 << 32) + 2)
+
+        self.assertEqual(
+            [(1 << 32) - 2, (1 << 32) - 1, 1 << 32, (1 << 32) + 1],
+            list(bitmap),
+        )
+
+    def test_fallback_set_operations_do_not_alias_inputs(self):
+        left = self._bitmap()
+        right = self._bitmap()
+        for value in [1, 2, 1 << 32]:
+            left.add(value)
+        for value in [2, (1 << 32) + 3]:
+            right.add(value)
+
+        union = left | right
+        difference = left - right
+        union.add((1 << 32) + 4)
+        difference.add((1 << 32) + 1)
+
+        self.assertEqual([2], list(left & right))
+        self.assertEqual([1, 2, 1 << 32], list(left))
+        self.assertEqual([2, (1 << 32) + 3], list(right))
+        self.assertEqual([1, 2, 1 << 32, (1 << 32) + 3, (1 << 32) + 4],
+                         list(union))
+        self.assertEqual([1, 1 << 32, (1 << 32) + 1], list(difference))
+
+    def test_fallback_round_trip_serialization(self):
+        bitmap = self._bitmap()
+        for value in [0, 7, (1 << 32), (3 << 32) + 9]:
+            bitmap.add(value)
+
+        copy = roaring_bitmap._ChunkedBitMap64.deserialize(bitmap.serialize())
+
+        self.assertEqual(bitmap, copy)
+        self.assertEqual([0, 7, 1 << 32, (3 << 32) + 9], list(copy))
+
+
+if __name__ == "__main__":
+    unittest.main()

paimon-python/pypaimon/utils/roaring_bitmap.py

@@ -19,9 +19,166 @@
 Roaring Bitmap.
 """
 
+import struct
 from typing import Iterator
 
 
+_UINT32_MASK = (1 << 32) - 1
+_INT64_MAX = (1 << 63) - 1
+
+
+def _check_int64(value: int) -> None:
+    if value < 0 or value > _INT64_MAX:
+        raise ValueError("RoaringBitmap64 only supports values from 0 to 2^63 - 1.")
+
+
+class _ChunkedBitMap64:
+    """64-bit bitmap fallback built from pyroaring.BitMap chunks."""
+
+    def __init__(self):
+        from pyroaring import BitMap
+
+        self._bitmap_cls = BitMap
+        self._chunks = {}
+
+    def add(self, value: int) -> None:
+        _check_int64(value)
+        key = value >> 32
+        low = value & _UINT32_MASK
+        self._chunks.setdefault(key, self._bitmap_cls()).add(low)
+
+    def add_range(self, from_: int, to: int) -> None:
+        _check_int64(from_)
+        _check_int64(to - 1 if to > from_ else from_)
+        current = from_
+        while current < to:
+            key = current >> 32
+            next_key_start = (key + 1) << 32
+            chunk_end = min(to, next_key_start)
+            self._chunks.setdefault(key, self._bitmap_cls()).add_range(
+                current & _UINT32_MASK,
+                ((chunk_end - 1) & _UINT32_MASK) + 1
+            )
+            current = chunk_end
+
+    def __contains__(self, value: int) -> bool:
+        if value < 0 or value > _INT64_MAX:
+            return False
+        key = value >> 32
+        chunk = self._chunks.get(key)
+        return chunk is not None and (value & _UINT32_MASK) in chunk
+
+    def __iter__(self) -> Iterator[int]:
+        for key in sorted(self._chunks):
+            base = key << 32
+            for low in self._chunks[key]:
+                yield base + low
+
+    def __len__(self) -> int:
+        return sum(len(chunk) for chunk in self._chunks.values())
+
+    def clear(self) -> None:
+        self._chunks.clear()
+
+    def _copy_chunk(self, chunk):
+        return self._bitmap_cls.deserialize(chunk.serialize())
+
+    def serialize(self) -> bytes:
+        chunks = [
+            (key, chunk)
+            for key, chunk in sorted(self._chunks.items())
+            if len(chunk) > 0
+        ]
+        result = bytearray(struct.pack("<q", len(chunks)))
+        for key, chunk in chunks:
+            result.extend(struct.pack("<i", key))
+            result.extend(chunk.serialize())
+        return bytes(result)
+
+    @classmethod
+    def deserialize(cls, data: bytes) -> '_ChunkedBitMap64':
+        from pyroaring import BitMap
+
+        result = cls()
+        offset = 0
+        count = struct.unpack_from("<q", data, offset)[0]
+        offset += 8
+        if count < 0:
+            raise ValueError("Invalid 64-bit roaring bitmap chunk count: %s" % count)
+
+        last_key = -1
+        for _ in range(count):
+            key = struct.unpack_from("<i", data, offset)[0]
+            offset += 4
+            if key < 0:
+                raise ValueError("Invalid 64-bit roaring bitmap chunk key: %s" % key)
+            if key <= last_key:
+                raise ValueError("64-bit roaring bitmap chunk keys must be ascending.")
+
+            chunk = BitMap.deserialize(data[offset:])
+            consumed = len(chunk.serialize())
+            if consumed <= 0:
+                raise ValueError("Invalid empty 32-bit roaring bitmap payload.")
+            offset += consumed
+            if len(chunk) > 0:
+                result._chunks[key] = chunk
+            last_key = key
+
+        if offset != len(data):
+            raise ValueError("Trailing bytes in 64-bit roaring bitmap payload.")
+
+        return result
+
+    def __and__(self, other: '_ChunkedBitMap64') -> '_ChunkedBitMap64':
+        result = _ChunkedBitMap64()
+        for key in set(self._chunks).intersection(other._chunks):
+            chunk = self._chunks[key] & other._chunks[key]
+            if len(chunk) > 0:
+                result._chunks[key] = chunk
+        return result
+
+    def __or__(self, other: '_ChunkedBitMap64') -> '_ChunkedBitMap64':
+        result = _ChunkedBitMap64()
+        for key in set(self._chunks).union(other._chunks):
+            left = self._chunks.get(key)
+            right = other._chunks.get(key)
+            if left is None:
+                chunk = self._copy_chunk(right)
+            elif right is None:
+                chunk = self._copy_chunk(left)
+            else:
+                chunk = left | right
+            if len(chunk) > 0:
+                result._chunks[key] = chunk
+        return result
+
+    def __sub__(self, other: '_ChunkedBitMap64') -> '_ChunkedBitMap64':
+        result = _ChunkedBitMap64()
+        for key, chunk in self._chunks.items():
+            other_chunk = other._chunks.get(key)
+            if other_chunk is None:
+                diff = self._copy_chunk(chunk)
+            else:
+                diff = chunk - other_chunk
+            if len(diff) > 0:
+                result._chunks[key] = diff
+        return result
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, _ChunkedBitMap64):
+            return False
+        return self._chunks == other._chunks
+
+
+def _new_bitmap64():
+    try:
+        from pyroaring import BitMap64
+
+        return BitMap64()
+    except ImportError:
+        return _ChunkedBitMap64()
+
+
 class RoaringBitmap64:
     """
     A 64-bit roaring bitmap implementation.
@@ -31,15 +188,17 @@ class RoaringBitmap64:
     """
 
     def __init__(self):
-        from pyroaring import BitMap64
-        self._data = BitMap64()
+        self._data = _new_bitmap64()
 
     def add(self, value: int) -> None:
         """Add a single value to the bitmap."""
+        _check_int64(value)
         self._data.add(value)
 
     def add_range(self, from_: int, to: int) -> None:
         """Add a range of values [from_, to] to the bitmap."""
+        _check_int64(from_)
+        _check_int64(to)
         self._data.add_range(from_, to + 1)
 
     def contains(self, value: int) -> bool:
@@ -132,8 +291,12 @@ def serialize(self) -> bytes:
     def deserialize(data: bytes) -> 'RoaringBitmap64':
         """Deserialize a bitmap from bytes."""
         result = RoaringBitmap64()
-        from pyroaring import BitMap64
-        result._data = BitMap64.deserialize(data)
+        try:
+            from pyroaring import BitMap64
+
+            result._data = BitMap64.deserialize(data)
+        except ImportError:
+            result._data = _ChunkedBitMap64.deserialize(data)
         return result
 
     def __eq__(self, other: object) -> bool: