From d365379b03ccea17e860dc4ce89c0c7c86d7b773 Mon Sep 17 00:00:00 2001 From: AuroraVoyage Date: Sun, 19 Jul 2026 15:29:21 +0800 Subject: [PATCH] [python] Introduce merge_map aggregator function. --- .../pypaimon/read/merge_engine_support.py | 1 + .../read/reader/aggregate/aggregators.py | 98 ++++++++++++++++++- .../pypaimon/tests/test_field_aggregators.py | 53 +++++++++- 3 files changed, 149 insertions(+), 3 deletions(-) diff --git a/paimon-python/pypaimon/read/merge_engine_support.py b/paimon-python/pypaimon/read/merge_engine_support.py index cefcee2e6701..4519c17af0cd 100644 --- a/paimon-python/pypaimon/read/merge_engine_support.py +++ b/paimon-python/pypaimon/read/merge_engine_support.py @@ -64,6 +64,7 @@ "bool_or", "bool_and", "listagg", "nested_update", + "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 11b182edaf44..3ad2a1eebddd 100644 --- a/paimon-python/pypaimon/read/reader/aggregate/aggregators.py +++ b/paimon-python/pypaimon/read/reader/aggregate/aggregators.py @@ -33,13 +33,13 @@ 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 from pypaimon.read.reader.aggregate import register_aggregator from pypaimon.read.reader.aggregate.field_aggregator import FieldAggregator -from pypaimon.schema.data_types import AtomicType, DataType, ArrayType, RowType +from pypaimon.schema.data_types import AtomicType, DataType, ArrayType, RowType, MapType from pypaimon.table.row.internal_row import InternalRow # aggregator input type hints variables @@ -60,6 +60,7 @@ NAME_BOOL_AND = "bool_and" NAME_LISTAGG = "listagg" NAME_NESTED_UPDATE = "nested_update" +NAME_MERGE_MAP = "merge_map" # Base SQL type names treated as numeric for sum/product-style @@ -664,6 +665,96 @@ def _apply_nested_key_null_strategy(self, key: Tuple[Any, ...]) -> bool: ) +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 @@ -740,3 +831,6 @@ def _factory(field_type, field_name, options): register_aggregator( NAME_NESTED_UPDATE, _build_field_options(FieldNestedUpdateAgg, NAME_NESTED_UPDATE) ) +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 9c9d5876fe9f..41fd381e4b63 100644 --- a/paimon-python/pypaimon/tests/test_field_aggregators.py +++ b/paimon-python/pypaimon/tests/test_field_aggregators.py @@ -46,8 +46,9 @@ FieldSumAgg, FieldListaggAgg, FieldNestedUpdateAgg, + FieldMergeMapAgg, ) -from pypaimon.schema.data_types import AtomicType, DataField, RowType, ArrayType +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 @@ -1558,6 +1559,56 @@ def test_field_nested_update_with_non_sequential_field_ids(self): self.assertCountEqual(accumulator, [self.row(0, 0, "A", 1), ]) +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