[Python] Implement combiner deferred side inputs

sdks/python/apache_beam/runners/direct/direct_runner.py

@@ -389,7 +389,7 @@ def _get_transform_overrides(pipeline_options):
 
  # Importing following locally to avoid a circular dependency.
  from apache_beam.pipeline import PTransformOverride
- from apache_beam.runners.direct.helper_transforms import LiftedCombinePerKey
+ from apache_beam.transforms.combiners import LiftedCombinePerKey
  from apache_beam.runners.direct.sdf_direct_runner import ProcessKeyedElementsViaKeyedWorkItemsOverride
  from apache_beam.runners.direct.sdf_direct_runner import SplittableParDoOverride
 

sdks/python/apache_beam/transforms/combiners.py

@@ -41,6 +41,7 @@
 from apache_beam.typehints import with_output_types
 from apache_beam.utils.timestamp import Duration
 from apache_beam.utils.timestamp import Timestamp
+from apache_beam.utils.windowed_value import WindowedValue
 
 __all__ = [
  'Count',
@@ -985,3 +986,126 @@ def merge_accumulators(self, accumulators):
 
  def extract_output(self, accumulator):
  return accumulator[0]
+
+
+class LiftedCombinePerKey(core.PTransform):
+ """An implementation of CombinePerKey that does mapper-side pre-combining.
+
+ This shouldn't generally be used directly except for use-cases where a
+ runner doesn't support CombinePerKey. This implementation manually implements
+ a CombinePerKey using ParDos, as opposed to runner implementations which may
+ use a more efficient implementation.
+ """
+ def __init__(self, combine_fn, args, kwargs):
+ side_inputs = _pack_side_inputs(args, kwargs)
+ self._side_inputs: dict = side_inputs
+ if not isinstance(combine_fn, core.CombineFn):
+ combine_fn = core.CombineFn.from_callable(combine_fn)
+ self._combine_fn = combine_fn
+
+ def expand(self, pcoll):
+ return (
+ pcoll
+ | core.ParDo(
+ _PartialGroupByKeyCombiningValues(self._combine_fn),
+ **self._side_inputs)
+ | core.GroupByKey()
+ | core.ParDo(_FinishCombine(self._combine_fn), **self._side_inputs))
+
+
+def _pack_side_inputs(side_input_args, side_input_kwargs):
+ if len(side_input_args) >= 10:
+ # If we have more than 10 side inputs, we can't use the
+ # _side_input_arg_{i} as our keys since they won't sort
+ # correctly. Just punt for now, more than 10 args probably
+ # doesn't happen often.
+ raise NotImplementedError
+ side_inputs = {}
+ for i, si in enumerate(side_input_args):
+ side_inputs[f'_side_input_arg_{i}'] = si
+ for k, v in side_input_kwargs.items():
+ side_inputs[k] = v
+ return side_inputs
+
+
+def _unpack_side_inputs(side_inputs):
+ side_input_args = []
+ side_input_kwargs = {}
+ for k, v in sorted(side_inputs.items(), key=lambda x: x[0]):
+ if k.startswith('_side_input_arg_'):
+ side_input_args.append(v)
+ else:
+ side_input_kwargs[k] = v
+ return side_input_args, side_input_kwargs
+
+
+class _PartialGroupByKeyCombiningValues(core.DoFn):
+ """Aggregates values into a per-key-window cache.
+
+ As bundles are in-memory-sized, we don't bother flushing until the very end.
+ """
+ def __init__(self, combine_fn):
+ self._combine_fn = combine_fn
+ self.side_input_args = []
+ self.side_input_kwargs = {}
+
+ def setup(self):
+ self._combine_fn.setup()
+
+ def start_bundle(self):
+ self._cache = dict()
+ self._cached_windowed_side_inputs = {}
+
+ def process(self, element, window=core.DoFn.WindowParam, **side_inputs):
+ k, vi = element
+ side_input_args, side_input_kwargs = _unpack_side_inputs(side_inputs)
+ if (k, window) not in self._cache:
+ self._cache[(k, window)] = self._combine_fn.create_accumulator(
+ *side_input_args, **side_input_kwargs)
+
+ self._cache[k, window] = self._combine_fn.add_input(
+ self._cache[k, window], vi, *side_input_args, **side_input_kwargs)
+ self._cached_windowed_side_inputs[window] = (
+ side_input_args, side_input_kwargs)
+
+ def finish_bundle(self):
+ for (k, w), va in self._cache.items():
+ # We compact the accumulator since a GBK (which necessitates encoding)
+ # will follow.
+ side_input_args, side_input_kwargs = (
+ self._cached_windowed_side_inputs[w])
+ yield WindowedValue((
+ k,
+ self._combine_fn.compact(va, *side_input_args, **side_input_kwargs)),
+ w.end, (w, ))
+
+ def teardown(self):
+ self._combine_fn.teardown()
+
+
+class _FinishCombine(core.DoFn):
+ """Merges partially combined results.
+ """
+ def __init__(self, combine_fn):
+ self._combine_fn = combine_fn
+
+ def setup(self):
+ self._combine_fn.setup()
+
+ def process(self, element, window=core.DoFn.WindowParam, **side_inputs):
+
+ k, vs = element
+ side_input_args, side_input_kwargs = _unpack_side_inputs(side_inputs)
+ return [(
+ k,
+ self._combine_fn.extract_output(
+ self._combine_fn.merge_accumulators(
+ vs, *side_input_args, **side_input_kwargs),
+ *side_input_args,
+ **side_input_kwargs))]
+
+ def teardown(self):
+ try:
+ self._combine_fn.teardown()
+ except AttributeError:
+ pass