From 9b7a27abacbdcca290cc7673e2f2cea144681922 Mon Sep 17 00:00:00 2001 From: Dharshika-11 Date: Mon, 17 Aug 2026 11:10:14 +0530 Subject: [PATCH] Bound PGBKCVOperation precombine table by in-memory size (#39754) --- .../apache_beam/runners/worker/operations.pxd | 7 +- .../apache_beam/runners/worker/operations.py | 74 ++++++++-- .../runners/worker/operations_test.py | 136 ++++++++++++++++++ 3 files changed, 205 insertions(+), 12 deletions(-) create mode 100644 sdks/python/apache_beam/runners/worker/operations_test.py diff --git a/sdks/python/apache_beam/runners/worker/operations.pxd b/sdks/python/apache_beam/runners/worker/operations.pxd index 52211e4d8ce8..c0f28cf52f3a 100644 --- a/sdks/python/apache_beam/runners/worker/operations.pxd +++ b/sdks/python/apache_beam/runners/worker/operations.pxd @@ -138,8 +138,11 @@ cdef class PGBKCVOperation(Operation): cdef public bint is_default_windowing cdef public object timestamp_combiner cdef dict table - cdef long max_keys - cdef long key_count + cdef public long max_keys + cdef public long key_count + cdef public long max_bytes + cdef public long estimated_bytes + cdef bint _is_tiny_accumulator cpdef add_key_value(self, wkey, value, timestamp) cpdef output_key(self, wkey, value, timestamp) diff --git a/sdks/python/apache_beam/runners/worker/operations.py b/sdks/python/apache_beam/runners/worker/operations.py index bfa1c0f3d47c..6df85b890a9b 100644 --- a/sdks/python/apache_beam/runners/worker/operations.py +++ b/sdks/python/apache_beam/runners/worker/operations.py @@ -23,6 +23,7 @@ # ruff: noqa: UP006 import collections import logging +import sys import threading import warnings from typing import TYPE_CHECKING @@ -51,6 +52,7 @@ from apache_beam.runners.worker import operation_specs from apache_beam.runners.worker import sideinputs from apache_beam.runners.worker.data_sampler import DataSampler +from apache_beam.runners.worker.statecache import get_deep_size from apache_beam.transforms import combiners from apache_beam.transforms import core from apache_beam.transforms import sideinputs as apache_sideinputs @@ -1234,6 +1236,16 @@ def flush(self, target): self.output(windowed_value) +def _safe_get_deep_size(*objs): + try: + return get_deep_size(*objs) + except Exception: + try: + return sum(sys.getsizeof(o) for o in objs) + except Exception: + return 128 + + class PGBKCVOperation(Operation): """Partial group-by-key operation. @@ -1241,7 +1253,14 @@ class PGBKCVOperation(Operation): a combine function applied. """ def __init__( - self, name_context, spec, counter_factory, state_sampler, windowing=None): + self, + name_context, + spec, + counter_factory, + state_sampler, + windowing=None, + max_bytes=None, + max_keys=None): super(PGBKCVOperation, self).__init__(name_context, spec, counter_factory, state_sampler) # Combiners do not accept deferred side-inputs (the ignored fourth @@ -1265,16 +1284,27 @@ def __init__( self.timestamp_combiner = None # Optimization for the (known tiny accumulator, often wide keyspace) # combine functions. - # TODO(b/36567833): Bound by in-memory size rather than key count. - self.max_keys = ( - 1000 * 1000 if + self._is_tiny_accumulator = ( isinstance(fn, (combiners.CountCombineFn, combiners.MeanCombineFn)) or # TODO(b/36597732): Replace this 'or' part by adding the 'cy' optimized # combiners to the short list above. ( isinstance(fn, core.CallableWrapperCombineFn) and - fn._fn in (min, max, sum)) else 100 * 1000) # pylint: disable=protected-access + fn._fn in (min, max, sum))) # pylint: disable=protected-access + + if max_keys is not None: + self.max_keys = max_keys + else: + self.max_keys = 1000 * 1000 if self._is_tiny_accumulator else 100 * 1000 + + if max_bytes is not None: + self.max_bytes = max_bytes + else: + self.max_bytes = ( + 100 * 1024 * 1024 if self._is_tiny_accumulator else 10 * 1024 * 1024) + self.key_count = 0 + self.estimated_bytes = 0 self.table = {} def setup(self, data_sampler=None): @@ -1301,34 +1331,58 @@ def process(self, wkv): def add_key_value(self, wkey, value, timestamp): entry = self.table.get(wkey, None) if entry is None: - if self.key_count >= self.max_keys: - target = self.key_count * 9 // 10 + if (self.key_count >= self.max_keys + or self.estimated_bytes >= self.max_bytes): + target_keys = self.key_count * 9 // 10 + target_bytes = self.max_bytes * 9 // 10 old_wkeys = [] # TODO(robertwb): Use an LRU cache? for old_wkey, old_wvalue in self.table.items(): old_wkeys.append(old_wkey) # Can't mutate while iterating. self.output_key(old_wkey, old_wvalue[0], old_wvalue[1]) self.key_count -= 1 - if self.key_count <= target: + if len(old_wvalue) > 3: + self.estimated_bytes -= old_wvalue[3] + if (self.key_count <= target_keys + and self.estimated_bytes <= target_bytes): break + if self.estimated_bytes < 0: + self.estimated_bytes = 0 for old_wkey in reversed(old_wkeys): del self.table[old_wkey] self.key_count += 1 - # We save the accumulator as a one element list so we can efficiently + # We save the accumulator as a list so we can efficiently # mutate when new values are added without searching the cache again. + # Format: [accumulator, timestamp, input_count, entry_bytes] entry = self.table[wkey] = [ - self.combine_fn.create_accumulator(), timestamp + self.combine_fn.create_accumulator(), timestamp, 0, 0 ] entry[0] = self.combine_fn_add_input(entry[0], value) if not self.is_default_windowing and self.timestamp_combiner: entry[1] = self.timestamp_combiner.combine(entry[1], timestamp) + entry[2] += 1 + count = entry[2] + if count == 1: + entry_bytes = _safe_get_deep_size(wkey, entry[0]) + entry[3] = entry_bytes + self.estimated_bytes += entry_bytes + elif not self._is_tiny_accumulator: + val_size = sys.getsizeof(value) + if (count & (count - 1)) == 0 or count % 64 == 0 or val_size > 10000: + new_size = _safe_get_deep_size(wkey, entry[0]) + delta = new_size - entry[3] + self.estimated_bytes += delta + entry[3] = new_size + def finish(self): # type: () -> None for wkey, value in self.table.items(): self.output_key(wkey, value[0], value[1]) self.table = {} self.key_count = 0 + self.estimated_bytes = 0 + super(PGBKCVOperation, self).finish() def teardown(self): # type: () -> None diff --git a/sdks/python/apache_beam/runners/worker/operations_test.py b/sdks/python/apache_beam/runners/worker/operations_test.py new file mode 100644 index 000000000000..629e6ace46f4 --- /dev/null +++ b/sdks/python/apache_beam/runners/worker/operations_test.py @@ -0,0 +1,136 @@ +# +# 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. +# + +"""Tests for operations.py, specifically PGBKCVOperation memory bounding.""" + +# pytype: skip-file + +import unittest + +from apache_beam.internal import pickler +from apache_beam.runners import common +from apache_beam.runners.worker import operation_specs +from apache_beam.runners.worker import operations +from apache_beam.transforms import combiners +from apache_beam.transforms import core +from apache_beam.utils.windowed_value import WindowedValue + + +class ListCombineFn(core.CombineFn): + def create_accumulator(self): + return [] + + def add_input(self, accumulator, input_val): + accumulator.append(input_val) + return accumulator + + def merge_accumulators(self, accumulators): + res = [] + for a in accumulators: + res.extend(a) + return res + + def extract_output(self, accumulator): + return accumulator + + +class MockOutputReceiver(common.Receiver): + def __init__(self): + self.output_values = [] + + def receive(self, windowed_value): + self.output_values.append(windowed_value) + + +class PGBKCVOperationTest(unittest.TestCase): + def _create_operation(self, combine_fn, max_bytes=None, max_keys=None): + spec = operation_specs.WorkerPartialGroupByKey( + combine_fn=pickler.dumps((combine_fn, [], {})), + input=None, + output_coders=[None]) + op = operations.PGBKCVOperation( + name_context=common.NameContext('test_step'), + spec=spec, + counter_factory=None, + state_sampler=None, + max_bytes=max_bytes, + max_keys=max_keys) + receiver = MockOutputReceiver() + op.add_receiver(receiver, 0) + op.setup() + op.is_default_windowing = True + return op, receiver + + def test_pgbkcv_max_bytes_bounding(self): + max_bytes = 2000 + max_keys = 100000 + op, receiver = self._create_operation( + ListCombineFn(), max_bytes=max_bytes, max_keys=max_keys) + + large_payload = 'x' * 500 + for i in range(100): + key = f'key_{i}' + wv = WindowedValue((key, large_payload), 0, ()) + op.process(wv) + if len(receiver.output_values) > 0: + break + + self.assertTrue( + len(receiver.output_values) > 0, + 'Operation should flush due to max_bytes memory limit being exceeded') + self.assertLess( + op.key_count, max_keys, + 'Eviction should occur well before max_keys limit') + + def test_pgbkcv_max_keys_bounding(self): + max_keys = 10 + max_bytes = 10 * 1024 * 1024 + op, receiver = self._create_operation( + combiners.CountCombineFn(), max_bytes=max_bytes, max_keys=max_keys) + + for i in range(15): + wv = WindowedValue((f'key_{i}', 1), 0, ()) + op.process(wv) + + self.assertTrue( + len(receiver.output_values) > 0, + 'Operation should flush when key_count reaches max_keys') + + def test_pgbkcv_finish_outputs_all_remaining(self): + op, receiver = self._create_operation( + combiners.CountCombineFn(), max_bytes=100000, max_keys=1000) + + keys = ['a', 'b', 'c', 'a', 'b'] + for k in keys: + op.process(WindowedValue((k, 1), 0, ())) + + op.finish() + + output_map = {wv.value[0]: wv.value[1] for wv in receiver.output_values} + self.assertEqual(output_map['a'], 2) + self.assertEqual(output_map['b'], 2) + self.assertEqual(output_map['c'], 1) + + def test_custom_parameters(self): + op, _ = self._create_operation( + ListCombineFn(), max_bytes=12345, max_keys=678) + self.assertEqual(op.max_bytes, 12345) + self.assertEqual(op.max_keys, 678) + + +if __name__ == '__main__': + unittest.main()