diff --git a/paimon-python/pypaimon/read/merge_engine_support.py b/paimon-python/pypaimon/read/merge_engine_support.py index 0624bf1e6504..f1c477ed8882 100644 --- a/paimon-python/pypaimon/read/merge_engine_support.py +++ b/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" diff --git a/paimon-python/pypaimon/read/reader/aggregate/aggregators.py b/paimon-python/pypaimon/read/reader/aggregate/aggregators.py index 0f7c4d4c3123..4d1a960db960 100644 --- a/paimon-python/pypaimon/read/reader/aggregate/aggregators.py +++ b/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) +) diff --git a/paimon-python/pypaimon/tests/test_field_aggregators.py b/paimon-python/pypaimon/tests/test_field_aggregators.py index 6d548499ca5e..7ce1e013bb5b 100644 --- a/paimon-python/pypaimon/tests/test_field_aggregators.py +++ b/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