diff --git a/paimon-python/pypaimon/read/merge_engine_support.py b/paimon-python/pypaimon/read/merge_engine_support.py index ff551f9ab51e..aab3cf3dd669 100644 --- a/paimon-python/pypaimon/read/merge_engine_support.py +++ b/paimon-python/pypaimon/read/merge_engine_support.py @@ -66,6 +66,7 @@ "nested_update", "collect", "merge_map_with_keytime", + "merge_map", ]) _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 321a70dacaee..d973acb8f430 100644 --- a/paimon-python/pypaimon/read/reader/aggregate/aggregators.py +++ b/paimon-python/pypaimon/read/reader/aggregate/aggregators.py @@ -33,7 +33,7 @@ error rather than a silent fallback. """ -from typing import Any, List, Dict, Optional, Tuple, Union +from typing import Any, List, Dict, Optional, Tuple, Union, Set from pypaimon.common.options import CoreOptions from pypaimon.common.options.core_options import NestedKeyNullStrategy @@ -62,6 +62,7 @@ NAME_NESTED_UPDATE = "nested_update" NAME_COLLECT = "collect" NAME_MERGE_MAP_WITH_KEYTIME = "merge_map_with_keytime" +NAME_MERGE_MAP = "merge_map" # Base SQL type names treated as numeric for sum/product-style @@ -851,6 +852,96 @@ def _get_ts_field(self, input_row: Any) -> Any: ) +class FieldMergeMapAgg(FieldAggregator): + """ + Merge map values by combining all key-value pairs. + + When the same key exists in both maps, the value from the input map + overwrites the value from the accumulator map. + """ + + def __init__(self, name: str, field_type: DataType): + super().__init__(name, field_type) + if not isinstance(field_type, MapType): + raise ValueError( + "Data type for merge map column must be 'MAP' but was '{}'".format(field_type) + ) + + 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 + + result = {} + + self._put_to_map(result, accumulator) + self._put_to_map(result, input_field) + + return result + + def retract(self, accumulator: Any, retract_field: Any) -> Any: + # it's hard to mark the input is retracted without accumulator + if accumulator is None: + return None + + # nothing to be retracted + if retract_field is None: + return accumulator + + if len(retract_field) == 0: + return accumulator + + retract_keys = self._get_keys(retract_field) + acc = {} + self._put_to_map(acc, accumulator) + + result = { + key: value + for key, value in acc.items() + if key not in retract_keys + } + + return result + + def _put_to_map(self, maps: Dict[Any, Any], input_field: Any): + if isinstance(input_field, dict): + maps.update(input_field) + elif isinstance(input_field, list): + tmp_map = {} + for item in input_field: + if not isinstance(item, dict): + raise TypeError( + "list element must be dict, got {}".format(type(item)) + ) + tmp_map[item['key']] = item['value'] + + maps.update(tmp_map) + else: + raise TypeError( + "input_field must be dict or list[dict], got {}".format(type(input_field)) + ) + + def _get_keys(self, retract_field: Any) -> Set[Any]: + keys = set() + + if isinstance(retract_field, dict): + keys.update(retract_field.keys()) + + elif isinstance(retract_field, list): + for item in retract_field: + if not isinstance(item, dict): + raise TypeError( + "list element must be dict, got {}".format(type(item)) + ) + + keys.add(item["key"]) + else: + raise TypeError( + "retract_field must be dict or list[dict], got {}".format(type(retract_field)) + ) + + return keys + + # --------------------------------------------------------------------------- # Registration. Each builder binds an identifier to a factory that # optionally validates the column DataType before constructing the @@ -933,3 +1024,6 @@ def _factory(field_type, field_name, options): register_aggregator( NAME_MERGE_MAP_WITH_KEYTIME, _build_field_options(FieldMergeMapWithKeyTimeAgg, NAME_MERGE_MAP_WITH_KEYTIME) ) +register_aggregator( + NAME_MERGE_MAP, _build_no_type_check(FieldMergeMapAgg, NAME_MERGE_MAP) +) diff --git a/paimon-python/pypaimon/tests/test_field_aggregators.py b/paimon-python/pypaimon/tests/test_field_aggregators.py index c6a188d90af4..72b8767ec865 100644 --- a/paimon-python/pypaimon/tests/test_field_aggregators.py +++ b/paimon-python/pypaimon/tests/test_field_aggregators.py @@ -48,6 +48,7 @@ FieldNestedUpdateAgg, FieldCollectAgg, FieldMergeMapWithKeyTimeAgg, + FieldMergeMapAgg, ) from pypaimon.schema.data_types import AtomicType, DataField, RowType, ArrayType, MapType from pypaimon.table.row.generic_row import GenericRow @@ -1920,6 +1921,56 @@ def test_field_merge_map_with_key_time_agg_retract(self): ) +class FieldMergeMapAggTest(unittest.TestCase): + + def _make(self, options: CoreOptions = None): + if options is None: + options = CoreOptions(Options.from_none()) + + return create_field_aggregator( + MapType(True, AtomicType("INT"), AtomicType("STRING")), + "field0", "merge_map", options=options + ) + + def test_field_merge_map_agg(self): + agg = self._make() + self.assertIsInstance(agg, FieldMergeMapAgg) + + self.assertIsNone(agg.agg(None, None)) + self.assertEqual(agg.agg({1: "A"}, None), {1: "A"}) + self.assertEqual(agg.agg(None, {1: "A"}), {1: "A"}) + + acc = agg.agg(None, {1: "A"}) + self.assertEqual(acc, {1: "A"}) + + acc = agg.agg(acc, {1: "A", 2: "B"}) + self.assertEqual(acc, {1: "A", 2: "B"}) + + acc = agg.agg(acc, {1: "a", 3: "c"}) + self.assertEqual(acc, {1: "a", 2: "B", 3: "c"}) + + def test_field_merge_map_agg_retract(self): + agg = self._make() + + result = agg.retract( + {1: "A", 2: "B", 3: "C"}, + {1: "A", 2: "A"}, + ) + self.assertEqual(result, {3: "C"}) + self.assertEqual(agg.retract(None, {1: "A"}), None) + self.assertEqual(agg.retract(result, None), {3: "C"}) + self.assertEqual(agg.retract(result, {}), {3: "C"}) + + def test_field_merge_map_agg_for_pyarrow(self): + agg = self._make() + acc = agg.agg(None, [{'key': 1, 'value': 'A'}]) + acc = agg.agg(acc, [{'key': 1, 'value': 'a'}, {'key': 2, 'value': 'B'}]) + self.assertEqual(acc, {1: "a", 2: "B"}) + + acc = agg.retract(acc, [{'key': 1, 'value': 'A'}, {'key': 3, 'value': 'C'}]) + self.assertEqual(acc, {2: "B"}) + + class RegistrationTest(unittest.TestCase): """Sanity check that all 10 expected aggregators (the primary-key placeholder plus 9 value aggregators) are registered when the