Allow operations to override the coder passed to update_counters

Some operations want to override at read or write time the coder
that was declared when the work item was initialized, so they
can provide a coder specific to this element.
----Release Notes----
[]
-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=120007941

Author: gildea

google/cloud/dataflow/worker/executor.pxd

@@ -28,7 +28,7 @@ cdef class Operation(object):
   cpdef finish(self)
 
   @cython.locals(receiver=Operation)
-  cpdef output(self, windowed_value, int output_index=*)
+  cpdef output(self, windowed_value, object coder=*, int output_index=*)
 
 cdef class ReadOperation(Operation):
   cdef object _current_progress
@@ -45,6 +45,7 @@ cdef class CombineOperation(Operation):
 cdef class ShuffleWriteOperation(Operation):
   cdef object shuffle_sink
   cdef object writer
+  cdef object _write_coder
   cdef bint is_ungrouped
 
 cdef class GroupedShuffleReadOperation(Operation):

google/cloud/dataflow/worker/executor.py

@@ -23,6 +23,8 @@
 
 
 from google.cloud.dataflow.coders import BytesCoder
+from google.cloud.dataflow.coders import TupleCoder
+from google.cloud.dataflow.coders import WindowedValueCoder
 from google.cloud.dataflow.internal import pickler
 from google.cloud.dataflow.pvalue import EmptySideInput
 from google.cloud.dataflow.runners import common
@@ -65,15 +67,15 @@ def start(self, step_name):
     self.opcounter = opcounters.OperationCounters(
         self.counter_factory, step_name, self.coder, self.output_index)
 
-  def output(self, windowed_value):
-    self.update_counters_start(windowed_value)
+  def output(self, windowed_value, coder=None):
+    self.update_counters_start(windowed_value, coder)
     for receiver in self.receivers:
       receiver.process(windowed_value)
     self.update_counters_finish()
 
-  def update_counters_start(self, windowed_value):
+  def update_counters_start(self, windowed_value, coder=None):
     if self.opcounter:
-      self.opcounter.update_from(windowed_value)
+      self.opcounter.update_from(windowed_value, coder)
 
   def update_counters_finish(self):
     if self.opcounter:
@@ -130,8 +132,8 @@ def process(self, o):
     """Process element in operation."""
     pass
 
-  def output(self, windowed_value, output_index=0):
-    self.receivers[output_index].output(windowed_value)
+  def output(self, windowed_value, coder=None, output_index=0):
+    self.receivers[output_index].output(windowed_value, coder)
 
   def add_receiver(self, operation, output_index=0):
     """Adds a receiver operation for the specified output."""
@@ -282,8 +284,10 @@ def __init__(self, spec, counter_factory, shuffle_source=None):
 
   def start(self):
     super(GroupedShuffleReadOperation, self).start()
+    write_coder = None
     if self.shuffle_source is None:
       coders = (self.spec.coder.key_coder(), self.spec.coder.value_coder())
+      write_coder = WindowedValueCoder(TupleCoder(coders))
       self.shuffle_source = shuffle.GroupedShuffleSource(
           self.spec.shuffle_reader_config, coder=coders,
           start_position=self.spec.start_shuffle_position,
@@ -292,7 +296,7 @@ def start(self):
       for key, key_values in reader:
         self._reader = reader
         windowed_value = GlobalWindows.WindowedValue((key, key_values))
-        self.output(windowed_value)
+        self.output(windowed_value, coder=write_coder)
 
   def get_progress(self):
     if self._reader is not None:
@@ -313,8 +317,10 @@ def __init__(self, spec, counter_factory, shuffle_source=None):
 
   def start(self):
     super(UngroupedShuffleReadOperation, self).start()
+    write_coder = None
     if self.shuffle_source is None:
       coders = (BytesCoder(), self.spec.coder)
+      write_coder = WindowedValueCoder(TupleCoder(coders))
       self.shuffle_source = shuffle.UngroupedShuffleSource(
           self.spec.shuffle_reader_config, coder=coders,
           start_position=self.spec.start_shuffle_position,
@@ -323,7 +329,7 @@ def start(self):
       for value in reader:
         self._reader = reader
         windowed_value = GlobalWindows.WindowedValue(value)
-        self.output(windowed_value)
+        self.output(windowed_value, coder=write_coder)
 
   def get_progress(self):
     # 'UngroupedShuffleReader' does not support progress reporting.
@@ -350,6 +356,7 @@ def start(self):
       coders = (BytesCoder(), coder)
     else:
       coders = (coder.key_coder(), coder.value_coder())
+    self._write_coder = WindowedValueCoder(TupleCoder(coders))
     if self.shuffle_sink is None:
       self.shuffle_sink = shuffle.ShuffleSink(
           self.spec.shuffle_writer_config, coder=coders)
@@ -364,7 +371,7 @@ def process(self, o):
     if self.debug_logging_enabled:
       logging.debug('Processing [%s] in %s', o, self)
     assert isinstance(o, WindowedValue)
-    self.receivers[0].update_counters_start(o)
+    self.receivers[0].update_counters_start(o, coder=self._write_coder)
     # 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

google/cloud/dataflow/worker/opcounters.py

@@ -31,15 +31,17 @@ def __init__(self, counter_factory, step_name, coder, output_index):
     self.coder = coder
     self._active_accumulators = []
 
-  def update_from(self, windowed_value):
+  def update_from(self, windowed_value, coder=None):
     """Add one value to this counter."""
     self.element_counter.update(1)
     byte_size_accumulator = Accumulator(self.mean_byte_counter.name)
     self._active_accumulators.append(byte_size_accumulator)
     # TODO(gildea):
     # Actually compute the encoded size of this value.
     # In spirit, something like this:
-    #     self.coder.store_estimated_size(windowed_value, byte_size_accumulator)
+    #     if coder is None:
+    #       coder = self.coder
+    #     coder.store_estimated_size(windowed_value, byte_size_accumulator)
     # but will need to do sampling.
 
   def update_collect(self):