Add plumbing to pass a coder to the byte-size counter updater

 - add output_coders and write_coder to maptask.Worker* objects
 - pass a coder to the ReceiverSet
 - pass that coder to the OperationCounter, where we compute size
----Release Notes----
[]
-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=118396896

Author: gildea

google/cloud/dataflow/utils/counters.py

@@ -64,7 +64,8 @@ def __init__(self, name, aggregation_kind):
     """
     self.name = name
     self.aggregation_kind = aggregation_kind
-    assert aggregation_kind == self.SUM  # update only handles sum
+    # optimized update doesn't handle all types
+    assert aggregation_kind == self.SUM or aggregation_kind == self.MEAN
     self.c_total = 0
     self.py_total = 0
     self.elements = 0

google/cloud/dataflow/worker/executor.py

@@ -52,17 +52,19 @@ class ReceiverSet(object):
   ReceiverSets are attached to the outputting Operation.
   """
 
-  def __init__(self, output_index=0):
+  def __init__(self, coder, 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)
+    self.coder = coder
 
   def add_receiver(self, receiving_operation):
     self.receivers.append(receiving_operation)
 
+  def start(self, step_name):
+    self.opcounter = opcounters.OperationCounters(
+        step_name, self.coder, self.output_index)
+
   def output(self, windowed_value):
     self.update_counters(windowed_value)
     for receiver in self.receivers:
@@ -98,20 +100,20 @@ def __init__(self, spec):
       spec: A maptask.Worker* instance.
     """
     self.spec = spec
-    # 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()]
+    self.receivers = []
+    # Everything except WorkerSideInputSource, which is not a
+    # top-level operation, should have output_coders
+    if getattr(self.spec, 'output_coders', None):
+      for i, coder in enumerate(self.spec.output_coders):
+        self.receivers.append(ReceiverSet(coder, i))
 
   def start(self):
     """Start operation."""
-    for receiver in self.receivers:
-      receiver.start(self.step_name)
     self.debug_logging_enabled = logging.getLogger().isEnabledFor(
         logging.DEBUG)
+    # Start our receivers, now that we know our step name.
+    for receiver in self.receivers:
+      receiver.start(self.step_name)
 
   def itercounters(self):
     for receiver in self.receivers:
@@ -131,8 +133,6 @@ def output(self, windowed_value, output_index=0):
 
   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(ReceiverSet(len(self.receivers)))
     self.receivers[output_index].add_receiver(operation)
 
   def __str__(self):
@@ -341,7 +341,7 @@ def __init__(self, spec, shuffle_sink=None):
   def start(self):
     super(ShuffleWriteOperation, self).start()
     self.is_ungrouped = self.spec.shuffle_kind == 'ungrouped'
-    coder = self.spec.coder
+    coder = self.spec.output_coders[0]
     if self.is_ungrouped:
       coders = (BytesCoder(), coder)
     else:
@@ -418,7 +418,7 @@ def _read_side_inputs(self, tags_and_types):
       results = []
       for si in itertools.ifilter(
           lambda o: o.tag == side_tag, self.spec.side_inputs):
-        if isinstance(si, maptask.WorkerRead):
+        if isinstance(si, maptask.WorkerSideInputSource):
           op = ReadOperation(si)
         else:
           raise NotImplementedError('Unknown side input type: %r' % si)