[python] Introduce rbm32 aggregator function

paimon-python/pypaimon/read/merge_engine_support.py

@@ -69,6 +69,7 @@
     "merge_map_with_keytime",
     "merge_map",
     "theta_sketch",
+    "rbm32",
 ])
 _FIELDS_PREFIX = "fields."
 _FIELD_SEQUENCE_GROUP_SUFFIX = ".sequence-group"

paimon-python/pypaimon/read/reader/aggregate/aggregators.py

@@ -44,6 +44,7 @@
 from pypaimon.schema.data_types import AtomicType, DataType, ArrayType, RowType, MapType
 from pypaimon.table.row.generic_row import GenericRow
 from pypaimon.table.row.internal_row import InternalRow
+from pypaimon.utils.roaring_bitmap import RoaringBitmap
 
 # aggregator input type hints variables
 Record = Union[InternalRow, Dict[str, Any]]
@@ -69,6 +70,7 @@
 NAME_MERGE_MAP_WITH_KEYTIME = "merge_map_with_keytime"
 NAME_MERGE_MAP = "merge_map"
 NAME_THETA_SKETCH = "theta_sketch"
+NAME_RBM32 = "rbm32"
 
 
 # Integer range limits used for overflow checking.
@@ -149,6 +151,18 @@ def _check_array_row(name: str, field_type: DataType) -> ArrayType:
     return field_type
 
 
+def _check_roaring_bitmap(name: str, field_type: DataType):
+    """Check field_type is VarBinaryType and return the VarBinaryType."""
+
+    base = _atomic_base_name(field_type)
+    if base not in ("VARBINARY", "BYTES"):
+        raise ValueError(
+            "Data type for {} column must be 'VARBINARY' or 'BYTES' but was "
+            "'{}'.".format(name, field_type)
+        )
+    return field_type
+
+
 def is_blank(s: str) -> bool:
     if s is None:
         return True
@@ -1410,6 +1424,21 @@ def agg(self, accumulator: Any, input_field: Any) -> Any:
         return union.get_result().serialize()
 
 
+class FieldRoaringBitmap32Agg(FieldAggregator):
+    """roaring bitmap 32 aggregate a field of a row."""
+
+    def agg(self, accumulator: Any, input_field: Any) -> Any:
+        if accumulator is None or input_field is None:
+            return input_field if accumulator is None else accumulator
+
+        try:
+            acc = RoaringBitmap.deserialize(accumulator)
+            input_bitmap = RoaringBitmap.deserialize(input_field)
+            return RoaringBitmap.or_(acc, input_bitmap).serialize()
+        except Exception as ex:
+            raise RuntimeError("Unable to se/deserialize roaring bitmap.") from ex
+
+
 # ---------------------------------------------------------------------------
 # Registration. Each builder binds an identifier to a factory that
 # optionally validates the column DataType before constructing the
@@ -1451,6 +1480,13 @@ def _factory(field_type, field_name, options):
     return _factory
 
 
+def _build_roaring_bitmap(cls, identifier: str):
+    def _factory(field_type, field_name, options):
+        _check_roaring_bitmap(identifier, field_type)
+        return cls(identifier, field_type)
+    return _factory
+
+
 register_aggregator(
     NAME_PRIMARY_KEY,
     _build_no_type_check(FieldPrimaryKeyAgg, NAME_PRIMARY_KEY),
@@ -1502,3 +1538,6 @@ def _factory(field_type, field_name, options):
 register_aggregator(
     NAME_THETA_SKETCH, _build_no_type_check(FieldThetaSketchAgg, NAME_THETA_SKETCH)
 )
+register_aggregator(
+    NAME_RBM32, _build_roaring_bitmap(FieldRoaringBitmap32Agg, NAME_RBM32)
+)

paimon-python/pypaimon/tests/test_field_aggregators.py

@@ -55,10 +55,12 @@
     FieldMergeMapWithKeyTimeAgg,
     FieldMergeMapAgg,
     FieldThetaSketchAgg,
+    FieldRoaringBitmap32Agg,
 )
 from pypaimon.schema.data_types import AtomicType, DataField, RowType, ArrayType, MapType
 from pypaimon.table.row.generic_row import GenericRow
 from pypaimon.table.row.internal_row import InternalRow
+from pypaimon.utils.roaring_bitmap import RoaringBitmap
 
 
 def _make(identifier, sql_type, options: CoreOptions = None):
@@ -2472,6 +2474,45 @@ def test_field_theta_sketch_agg(self):
         self.assertEqual(agg.agg(acc2, input_val), acc2)
 
 
+class FieldRoaringBitmap32AggTest(unittest.TestCase):
+
+    def test_field_roaring_bitmap32_agg(self):
+        agg = _make("rbm32", "VARBINARY(20)")
+        self.assertIsInstance(agg, FieldRoaringBitmap32Agg)
+
+        input_rbm = RoaringBitmap()
+        acc1_rbm = RoaringBitmap()
+        acc2_rbm = RoaringBitmap()
+
+        input_rbm.add(1)
+        acc1_rbm.add_range(2, 3)
+        acc2_rbm.add_range(1, 3)
+
+        input_val = input_rbm.serialize()
+        acc1 = acc1_rbm.serialize()
+        acc2 = acc2_rbm.serialize()
+
+        self.assertIsNone(agg.agg(None, None))
+
+        result1 = agg.agg(None, input_val)
+        self.assertEqual(result1, input_val)
+
+        result2 = agg.agg(acc1, None)
+        self.assertEqual(result2, acc1)
+
+        result3 = agg.agg(acc1, input_val)
+        self.assertEqual(result3, acc2)
+
+        result4 = agg.agg(acc2, input_val)
+        self.assertEqual(result4, acc2)
+
+    def test_field_roaring_bitmap32_requires_varbinary(self):
+        with self.assertRaises(ValueError) as ctx:
+            _make("rbm32", "VARCHAR(20)")
+
+        self.assertIn("VARBINARY", str(ctx.exception))
+
+
 class RegistrationTest(unittest.TestCase):
     """Sanity check that all 10 expected aggregators (the primary-key
     placeholder plus 9 value aggregators) are registered when the