Implement continuous combining in pre-shuffle combining table.

This is particularly useful for global combines.

Fixed counters to be updated exactly once per output,
rather than once for every consumer.

Also optimized the callable combine fn and fixed tracebacks for the DoFn runner.

----Release Notes----
[]
-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=117146623

Author: robertwb

google/cloud/dataflow/runners/common.py

@@ -93,25 +93,18 @@ def process(self, element):
         self.context.set_element(element)
         self._process_outputs(element, self.dofn.process(self.context))
     except BaseException as exn:
-      raise self.augment_exception(exn)
-
-  def augment_exception(self, exn):
-    try:
-      if getattr(exn, '_tagged_with_step', False) or not self.step_name:
-        return exn
-      args = exn.args
-      if args and isinstance(args[0], str):
-        args = (args[0] + " [while running '%s']" % self.step_name,) + args[1:]
-        # Poor man's exception chaining.
-        try:
-          raise type(exn), args, sys.exc_info()[2]
-        except BaseException as new_exn:
-          new_exn._tagged_with_step = True
-          return new_exn
-      else:
-        return exn
-    except:
-      return exn
+      self.reraise_augmented(exn)
+
+  def reraise_augmented(self, exn):
+    if getattr(exn, '_tagged_with_step', False) or not self.step_name:
+      raise
+    args = exn.args
+    if args and isinstance(args[0], str):
+      args = (args[0] + " [while running '%s']" % self.step_name,) + args[1:]
+      # Poor man's exception chaining.
+      raise type(exn), args, sys.exc_info()[2]
+    else:
+      raise
 
   def _process_outputs(self, element, results):
     """Dispatch the result of computation to the appropriate receivers.

google/cloud/dataflow/transforms/combiners.py

@@ -452,37 +452,42 @@ def extract_output(self, accumulator):
     return accumulator
 
 
-class PhasedCombineFnExecutor(object):
-  """Executor for phases of combine operations."""
+def curry_combine_fn(fn, args, kwargs):
+  if not args and not kwargs:
+    return fn
 
-  def __init__(self, phase, fn, args, kwargs):
+  else:
 
-    if not args and not kwargs:
-      self.combine_fn = fn
-    else:
+    class CurriedFn(core.CombineFn):
+      """CombineFn that applies extra arguments."""
+
+      def create_accumulator(self):
+        return fn.create_accumulator(*args, **kwargs)
 
-      class CurriedFn(core.CombineFn):
-        """CombineFn that applies extra arguments."""
+      def add_input(self, accumulator, element):
+        return fn.add_input(accumulator, element, *args, **kwargs)
 
-        def create_accumulator(self):
-          return fn.create_accumulator(*args, **kwargs)
+      def add_inputs(self, accumulator, elements):
+        return fn.add_inputs(accumulator, elements, *args, **kwargs)
 
-        def add_input(self, accumulator, element):
-          return fn.add_input(accumulator, element, *args, **kwargs)
+      def merge_accumulators(self, accumulators):
+        return fn.merge_accumulators(accumulators, *args, **kwargs)
 
-        def add_inputs(self, accumulator, elements):
-          return fn.add_inputs(accumulator, elements, *args, **kwargs)
+      def extract_output(self, accumulator):
+        return fn.extract_output(accumulator, *args, **kwargs)
 
-        def merge_accumulators(self, accumulators):
-          return fn.merge_accumulators(accumulators, *args, **kwargs)
+      def apply(self, elements):
+        return fn.apply(elements, *args, **kwargs)
 
-        def extract_output(self, accumulator):
-          return fn.extract_output(accumulator, *args, **kwargs)
+    return CurriedFn()
 
-        def apply(self, elements):
-          return fn.apply(elements, *args, **kwargs)
 
-      self.combine_fn = CurriedFn()
+class PhasedCombineFnExecutor(object):
+  """Executor for phases of combine operations."""
+
+  def __init__(self, phase, fn, args, kwargs):
+
+    self.combine_fn = curry_combine_fn(fn, args, kwargs)
 
     if phase == 'all':
       self.apply = self.full_combine

google/cloud/dataflow/transforms/core.py

@@ -384,17 +384,22 @@ def create_accumulator(self, *args, **kwargs):
     return self._EMPTY
 
   def add_input(self, accumulator, element, *args, **kwargs):
-    return self.add_inputs([element], *args, **kwargs)
+    if accumulator is self._EMPTY:
+      return element
+    else:
+      return self._fn([accumulator, element], *args, **kwargs)
 
   def add_inputs(self, accumulator, elements, *args, **kwargs):
-    if accumulator is not self._EMPTY:
+    if accumulator is self._EMPTY:
+      return self._fn(elements, *args, **kwargs)
+    elif isinstance(elements, (list, tuple)):
+      return self._fn([accumulator] + elements, *args, **kwargs)
+    else:
       def union():
         yield accumulator
         for e in elements:
           yield e
-    else:
-      union = lambda: elements
-    return self._fn(union(), *args, **kwargs)
+      return self._fn(union(), *args, **kwargs)
 
   def merge_accumulators(self, accumulators, *args, **kwargs):
     # It's (weakly) assumed that self._fn is associative.

google/cloud/dataflow/worker/executor.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# cython: profile=True
+
 """Worker operations executor."""
 
 import collections
@@ -24,8 +26,10 @@
 from google.cloud.dataflow.pvalue import EmptySideInput
 from google.cloud.dataflow.runners import common
 import google.cloud.dataflow.transforms as ptransform
+from google.cloud.dataflow.transforms import combiners
 from google.cloud.dataflow.transforms import trigger
 from google.cloud.dataflow.transforms import window
+from google.cloud.dataflow.transforms.combiners import curry_combine_fn
 from google.cloud.dataflow.transforms.combiners import PhasedCombineFnExecutor
 from google.cloud.dataflow.transforms.trigger import InMemoryUnmergedState
 from google.cloud.dataflow.transforms.window import GlobalWindows
@@ -147,8 +151,8 @@ def start(self):
           windowed_value = value
         else:
           windowed_value = GlobalWindows.WindowedValue(value)
+        self.counters[0].update(windowed_value)
         for receiver in self.receivers[0]:
-          self.counters[0].update(windowed_value)
           receiver.process(windowed_value)
 
   def side_read_all(self, singleton=False):
@@ -239,9 +243,9 @@ def start(self):
     with self.shuffle_source.reader() as reader:
       for key, key_values in reader:
         self._reader = reader
+        windowed_value = GlobalWindows.WindowedValue((key, key_values))
+        self.counters[0].update(windowed_value)
         for receiver in self.receivers[0]:
-          windowed_value = GlobalWindows.WindowedValue((key, key_values))
-          self.counters[0].update(windowed_value)
           receiver.process(windowed_value)
 
   def get_progress(self):
@@ -271,9 +275,9 @@ def start(self):
     with self.shuffle_source.reader() as reader:
       for value in reader:
         self._reader = reader
+        windowed_value = GlobalWindows.WindowedValue(value)
+        self.counters[0].update(windowed_value)
         for receiver in self.receivers[0]:
-          windowed_value = GlobalWindows.WindowedValue(value)
-          self.counters[0].update(windowed_value)
           receiver.process(windowed_value)
 
   def get_progress(self):
@@ -463,11 +467,18 @@ def process(self, o):
     key, values = o.value
     windowed_result = WindowedValue(
         (key, self.phased_combine_fn.apply(values)), o.timestamp, o.windows)
+    self.counters[0].update(windowed_result)
     for receiver in self.receivers[0]:
-      self.counters[0].update(windowed_result)
       receiver.process(windowed_result)
 
 
+def create_pgbk_op(spec):
+  if spec.combine_fn:
+    return PGBKCVOperation(spec)
+  else:
+    return PGBKOperation(spec)
+
+
 class PGBKOperation(Operation):
   """Partial group-by-key operation.
 
@@ -478,16 +489,7 @@ class PGBKOperation(Operation):
 
   def __init__(self, spec):
     super(PGBKOperation, self).__init__(spec)
-    self.phased_combine_fn = None
-    if self.spec.combine_fn:
-      # Combiners do not accept deferred side-inputs (the ignored fourth
-      # argument) and therefore the code to handle the extra args/kwargs is
-      # simpler than for the DoFn's of ParDo.
-      #
-      # TODO(ccy): Combine as we go for each key instead of storing up state
-      # for combination when flushing.
-      fn, args, kwargs = pickler.loads(self.spec.combine_fn)[:3]
-      self.phased_combine_fn = PhasedCombineFnExecutor('add', fn, args, kwargs)
+    assert not self.spec.combine_fn
     self.table = collections.defaultdict(list)
     self.size = 0
     # TODO(robertwb) Make this configurable.
@@ -512,16 +514,58 @@ def flush(self, target):
       del self.table[kw]
       key, windows = kw
       output_value = [v.value[1] for v in vs]
-      if self.phased_combine_fn:
-        output_value = self.phased_combine_fn.apply(output_value)
       windowed_value = WindowedValue(
           (key, output_value),
           vs[0].timestamp, windows)
+      self.counters[0].update(windowed_value)
       for receiver in self.receivers[0]:
-        self.counters[0].update(windowed_value)
         receiver.process(windowed_value)
 
 
+class PGBKCVOperation(Operation):
+
+  def __init__(self, spec):
+    super(PGBKCVOperation, self).__init__(spec)
+    # Combiners do not accept deferred side-inputs (the ignored fourth
+    # argument) and therefore the code to handle the extra args/kwargs is
+    # simpler than for the DoFn's of ParDo.
+    fn, args, kwargs = pickler.loads(self.spec.combine_fn)[:3]
+    self.combine_fn = curry_combine_fn(fn, args, kwargs)
+    # Optimization for the (known tiny accumulator, often wide keyspace)
+    # count function.
+    # TODO(robertwb): Bound by in-memory size rather than key count.
+    self.max_keys = (
+        1000000 if isinstance(fn, combiners.CountCombineFn) else 10000)
+    self.key_count = 0
+    self.table = {}
+
+  def process(self, wkv):
+    key, value = wkv.value
+    wkey = tuple(wkv.windows), key
+    entry = self.table.get(wkey, None)
+    if entry is None:
+      if self.key_count >= self.max_keys:
+        old_wkey = self.table.iterkeys().next()  # Any key, could use LRU
+        self.output(old_wkey, self.table.pop(old_wkey)[0])
+      else:
+        self.key_count += 1
+      entry = self.table[wkey] = [self.combine_fn.create_accumulator()]
+    entry[0] = self.combine_fn.add_inputs(entry[0], [value])
+
+  def finish(self):
+    for wkey, value in self.table.iteritems():
+      self.output(wkey, value[0])
+    self.entries = {}
+    self.key_count = 0
+
+  def output(self, wkey, value):
+    windows, key = wkey
+    windowed_value = WindowedValue((key, value), windows[0].end, windows)
+    self.counters[0].update(windowed_value)
+    for receiver in self.receivers[0]:
+      receiver.process(windowed_value)
+
+
 class FlattenOperation(Operation):
   """Flatten operation.
 
@@ -533,8 +577,8 @@ def process(self, o):
     logging.debug('Processing [%s] in %s', o, self)
     assert isinstance(o, WindowedValue)
     windowed_result = WindowedValue(o.value, o.timestamp, o.windows)
+    self.counters[0].update(windowed_result)
     for receiver in self.receivers[0]:
-      self.counters[0].update(windowed_result)
       receiver.process(windowed_result)
 
 
@@ -559,8 +603,8 @@ def process(self, o):
             o.timestamp, o.windows))
 
   def output(self, windowed_result):
+    self.counters[0].update(windowed_result)
     for receiver in self.receivers[0]:
-      self.counters[0].update(windowed_result)
       receiver.process(windowed_result)
 
 
@@ -608,8 +652,8 @@ def process(self, o):
               window.WindowedValue((k, values), timestamp, [out_window]))
 
   def output(self, windowed_result):
+    self.counters[0].update(windowed_result)
     for receiver in self.receivers[0]:
-      self.counters[0].update(windowed_result)
       receiver.process(windowed_result)
 
 
@@ -646,8 +690,8 @@ def process(self, o):
                                          [out_window]))
 
   def output(self, windowed_result):
+    self.counters[0].update(windowed_result)
     for receiver in self.receivers[0]:
-      self.counters[0].update(windowed_result)
       receiver.process(windowed_result)
 
 
@@ -706,7 +750,7 @@ def execute(self, map_task, test_shuffle_source=None, test_shuffle_sink=None):
       elif isinstance(spec, maptask.WorkerCombineFn):
         op = CombineOperation(spec)
       elif isinstance(spec, maptask.WorkerPartialGroupByKey):
-        op = PGBKOperation(spec)
+        op = create_pgbk_op(spec)
       elif isinstance(spec, maptask.WorkerDoFn):
         op = DoOperation(spec)
       elif isinstance(spec, maptask.WorkerGroupingShuffleRead):