New class ReceiverSet in the worker

A ReceiverSet holds all the per-receiver information an Operation has.

The motivation for introducing it was to have a place to hold the
coder for an output (to be introduced later), and it is useful even
before that addition as a way to combine receiving operations and
the OperationCounter.
----Release Notes----
[]
-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=117833057

Author: gildea

google/cloud/dataflow/runners/common.pxd

@@ -20,11 +20,9 @@ cdef class DoFnRunner(object):
   cdef object window_fn
   cdef object context
   cdef object tagged_receivers
-  cdef object tagged_counters
   cdef object logger
   cdef object step_name
 
-  cdef list main_receivers
-  cdef object main_counters
+  cdef object main_receivers
 
   cpdef _process_outputs(self, element, results)

google/cloud/dataflow/runners/common.py

@@ -48,7 +48,6 @@ def __init__(self,
                windowing,
                context,
                tagged_receivers,
-               tagged_counters,
                logger=None,
                step_name=None):
     if not args and not kwargs:
@@ -70,13 +69,11 @@ def finish_bundle(self, context):
     self.window_fn = windowing.windowfn
     self.context = context
     self.tagged_receivers = tagged_receivers
-    self.tagged_counters = tagged_counters
     self.logger = logger or FakeLogger()
     self.step_name = step_name
 
     # Optimize for the common case.
     self.main_receivers = tagged_receivers[None]
-    self.main_counters = tagged_counters[None]
 
   def start(self):
     self.context.set_element(None)
@@ -150,15 +147,10 @@ def _process_outputs(self, element, results):
       else:
         windowed_value = WindowedValue(
             result, element.timestamp, element.windows)
-      # TODO(robertwb): Should the counters be on the context?
       if tag is None:
-        self.main_counters.update(windowed_value)
-        for receiver in self.main_receivers:
-          receiver.process(windowed_value)
+        self.main_receivers.output(windowed_value)
       else:
-        self.tagged_counters[tag].update(windowed_value)
-        for receiver in self.tagged_receivers[tag]:
-          receiver.process(windowed_value)
+        self.tagged_receivers[tag].output(windowed_value)
 
 class NoContext(WindowFn.AssignContext):
   """An uninspectable WindowFn.AssignContext."""

google/cloud/dataflow/runners/direct_runner.py

@@ -114,19 +114,15 @@ def get_side_input_value(si):
     transform.dofn = OutputCheckWrapperDoFn(
         transform.dofn, transform_node.full_label)
 
-    class NoOpCounters(object):
-      def update(self, element):
-        pass
-
-    class RecordingReciever(object):
+    class RecordingReceiverSet(object):
       def __init__(self, tag):
         self.tag = tag
-      def process(self, element):
+      def output(self, element):
         results[self.tag].append(element)
 
-    class TaggedRecievers(dict):
+    class TaggedReceivers(dict):
       def __missing__(self, key):
-        return [RecordingReciever(key)]
+        return RecordingReceiverSet(key)
 
     results = collections.defaultdict(list)
     # Some tags may be empty.
@@ -135,8 +131,7 @@ def __missing__(self, key):
 
     runner = DoFnRunner(transform.dofn, transform.args, transform.kwargs,
                         side_inputs, transform_node.inputs[0].windowing,
-                        context, TaggedRecievers(),
-                        collections.defaultdict(NoOpCounters),
+                        context, TaggedReceivers(),
                         step_name=transform_node.full_label)
     runner.start()
     for v in self._cache.get_pvalue(transform_node.inputs[0]):

google/cloud/dataflow/worker/executor.py

@@ -42,6 +42,45 @@
 from google.cloud.dataflow.worker import shuffle
 
 
+class ReceiverSet(object):
+  """A ReceiverSet represents a graph edge between two Operation nodes.
+
+  The ReceiverSet object collects information from the output of the
+  Operation at one end of its edge and the input of the Operation at
+  the other edge.
+  ReceiverSets are attached to the outputting Operation.
+  """
+
+  def __init__(self, output_index=0):
+    self.receivers = []
+    self.opcounter = None
+    self.output_index = output_index
+
+  def start(self, step_name):
+    self.opcounter = opcounters.OperationCounters(step_name, self.output_index)
+
+  def add_receiver(self, receiving_operation):
+    self.receivers.append(receiving_operation)
+
+  def output(self, windowed_value):
+    self.update_counters(windowed_value)
+    for receiver in self.receivers:
+      receiver.process(windowed_value)
+
+  def update_counters(self, windowed_value):
+    if self.opcounter:
+      self.opcounter.update(windowed_value)
+
+  def itercounters(self):
+    if self.opcounter:
+      for counter in self.opcounter:
+        yield counter
+
+  def __str__(self):
+    return '[%s]' % ' '.join([r.str_internal(is_recursive=True)
+                              for r in self.receivers])
+
+
 class Operation(object):
   """An operation representing the live version of a work item specification.
 
@@ -58,21 +97,22 @@ def __init__(self, spec):
       spec: A maptask.Worker* instance.
     """
     self.spec = spec
-    self.receivers = [[]]
-    # Initially we have no counters.  Initializing this here makes it
-    # safe to call itercounters() at any time, even if start() has
-    # not been called yet.
-    self.counters = []
+    # Create the ReceiverSet for the default output.
+    # We need this in several cases:
+    # A. There may be no receiver explicitly created for an output:
+    #  1. ParDo without anything following it, executed for side effect.
+    #  2. Partition, which generates a default output that isn't used.
+    # B. Write operations want opcounters, even though they have no outputs.
+    self.receivers = [ReceiverSet()]
 
   def start(self):
     """Start operation."""
-    # If the operation has receivers, create one counter set per receiver.
-    self.counters = [opcounters.OperationCounters(self.step_name, output_index)
-                     for output_index in range(len(self.receivers))]
+    for receiver in self.receivers:
+      receiver.start(self.step_name)
 
   def itercounters(self):
-    for opcounter in self.counters:
-      for counter in opcounter:
+    for receiver in self.receivers:
+      for counter in receiver.itercounters():
         yield counter
 
   def finish(self):
@@ -84,15 +124,13 @@ def process(self, o):
     pass
 
   def output(self, windowed_value, output_index=0):
-    self.counters[output_index].update(windowed_value)
-    for receiver in self.receivers[output_index]:
-      receiver.process(windowed_value)
+    self.receivers[output_index].output(windowed_value)
 
   def add_receiver(self, operation, output_index=0):
     """Adds a receiver operation for the specified output."""
     while len(self.receivers) <= output_index:
-      self.receivers.append([])
-    self.receivers[output_index].append(operation)
+      self.receivers.append(ReceiverSet(len(self.receivers)))
+    self.receivers[output_index].add_receiver(operation)
 
   def __str__(self):
     """Generates a useful string for this object.
@@ -127,9 +165,7 @@ def str_internal(self, is_recursive=False):
 
     if not is_recursive and getattr(self, 'receivers', []):
       printable_fields.append('receivers=[%s]' % ', '.join([
-          rop.str_internal(is_recursive=True)
-          for oplist in self.receivers
-          for rop in oplist]))
+          str(receiver) for receiver in self.receivers]))
 
     return '<%s %s>' % (printable_name, ', '.join(printable_fields))
 
@@ -206,7 +242,7 @@ def finish(self):
   def process(self, o):
     logging.debug('Processing [%s] in %s', o, self)
     assert isinstance(o, WindowedValue)
-    self.counters[0].update(o)
+    self.receivers[0].update_counters(o)
     if self.use_windowed_value:
       self.writer.Write(o)
     else:
@@ -223,7 +259,7 @@ def __init__(self, spec):
   def process(self, o):
     logging.debug('Processing [%s] in %s', o, self)
     assert isinstance(o, WindowedValue)
-    self.counters[0].update(o)
+    self.receivers[0].update_counters(o)
     self.spec.output_buffer.append(o.value)
 
 
@@ -312,7 +348,7 @@ def finish(self):
   def process(self, o):
     logging.debug('Processing [%s] in %s', o, self)
     assert isinstance(o, WindowedValue)
-    self.counters[0].update(o)
+    self.receivers[0].update_counters(o)
     # We typically write into shuffle key/value pairs. This is the reason why
     # the else branch below expects the value attribute of the WindowedValue
     # argument to be a KV pair. However the service may write to shuffle in
@@ -409,7 +445,6 @@ def start(self):
     # by the DoFn function to the appropriate receivers. The main output is
     # tagged with None and is associated with its corresponding index.
     tagged_receivers = {}
-    tagged_counters = {}
     output_tag_prefix = PropertyNames.OUT + '_'
     for index, tag in enumerate(self.spec.output_tags):
       if tag == PropertyNames.OUT:
@@ -418,19 +453,11 @@ def start(self):
         original_tag = tag[len(output_tag_prefix):]
       else:
         raise ValueError('Unexpected output name for operation: %s' % tag)
-      # There may be no receiver for this output, in which case the
-      # lookup will create one, and this value will be processed
-      # for any side effect.  This is desirable.  There are two (known)
-      # cases where there is no receiver for an output:
-      #  1. ParDo without anything following it, executed for side effect.
-      #  2. Partition (shows up here in the worker as Flatten), which
-      #     generates a default output that isn't used.
       tagged_receivers[original_tag] = self.receivers[index]
-      tagged_counters[original_tag] = self.counters[index]
 
     self.dofn_runner = common.DoFnRunner(
         fn, args, kwargs, self._read_side_inputs(tags_and_types),
-        window_fn, self.context, tagged_receivers, tagged_counters,
+        window_fn, self.context, tagged_receivers,
         logger, self.step_name)
 
     self.dofn_runner.start()
@@ -773,12 +800,12 @@ def execute(self, map_task, test_shuffle_source=None, test_shuffle_sink=None):
 
       # Add receiver operations to the appropriate producers.
       if hasattr(op.spec, 'input'):
-        producer, index = op.spec.input
-        self._ops[producer].add_receiver(op, index)
+        producer, output_index = op.spec.input
+        self._ops[producer].add_receiver(op, output_index)
       # Flatten has 'inputs', not 'input'
       if hasattr(op.spec, 'inputs'):
-        for producer, index in op.spec.inputs:
-          self._ops[producer].add_receiver(op, index)
+        for producer, output_index in op.spec.inputs:
+          self._ops[producer].add_receiver(op, output_index)
 
     # Inject the step names into the operations.
     # This is used for logging and assigning names to counters.