Bound PGBKCVOperation precombine table by in-memory size (#39754)

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)

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,14 +1236,31 @@ 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.
 
   This operation handles grouped values with
 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

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()