Don't use KV coder for ungrouped shuffle reads/writes

This fixes an issue where ungrouped reshuffles were incorrectly requiring KV
coders.

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

Author: charlesccychen

google/cloud/dataflow/worker/executor.py

@@ -22,6 +22,7 @@
 import random
 
 
+from google.cloud.dataflow.coders import BytesCoder
 from google.cloud.dataflow.internal import pickler
 from google.cloud.dataflow.pvalue import EmptySideInput
 from google.cloud.dataflow.runners import common
@@ -278,8 +279,9 @@ def __init__(self, spec, shuffle_source=None):
   def start(self):
     super(GroupedShuffleReadOperation, self).start()
     if self.shuffle_source is None:
+      coders = (self.spec.coder.key_coder(), self.spec.coder.value_coder())
       self.shuffle_source = shuffle.GroupedShuffleSource(
-          self.spec.shuffle_reader_config, coder=self.spec.coders,
+          self.spec.shuffle_reader_config, coder=coders,
           start_position=self.spec.start_shuffle_position,
           end_position=self.spec.end_shuffle_position)
     with self.shuffle_source.reader() as reader:
@@ -308,8 +310,9 @@ def __init__(self, spec, shuffle_source=None):
   def start(self):
     super(UngroupedShuffleReadOperation, self).start()
     if self.shuffle_source is None:
+      coders = (BytesCoder(), self.spec.coder)
       self.shuffle_source = shuffle.UngroupedShuffleSource(
-          self.spec.shuffle_reader_config, coder=self.spec.coders,
+          self.spec.shuffle_reader_config, coder=coders,
           start_position=self.spec.start_shuffle_position,
           end_position=self.spec.end_shuffle_position)
     with self.shuffle_source.reader() as reader:
@@ -337,13 +340,17 @@ def __init__(self, spec, shuffle_sink=None):
 
   def start(self):
     super(ShuffleWriteOperation, self).start()
-    # TODO(silviuc): Shuffle 'kind' is ignored!
+    self.is_ungrouped = self.spec.shuffle_kind == 'ungrouped'
+    coder = self.spec.coder
+    if self.is_ungrouped:
+      coders = (BytesCoder(), coder)
+    else:
+      coders = (coder.key_coder(), coder.value_coder())
     if self.shuffle_sink is None:
       self.shuffle_sink = shuffle.ShuffleSink(
-          self.spec.shuffle_writer_config, coder=self.spec.coders)
+          self.spec.shuffle_writer_config, coder=coders)
     self.writer = self.shuffle_sink.writer()
     self.writer.__enter__()
-    self.is_ungrouped = self.spec.shuffle_kind == 'ungrouped'
 
   def finish(self):
     logging.debug('Finishing %s', self)

google/cloud/dataflow/worker/executor_test.py

@@ -98,7 +98,7 @@ def reader(self):
 
 class ExecutorTest(unittest.TestCase):
 
-  SHUFFLE_CODERS = (coders.PickleCoder(), coders.PickleCoder())
+  SHUFFLE_CODER = coders.PickleCoder()
 
   def create_temp_file(self, content_text):
     """Creates a temporary file with content and returns the path to it."""
@@ -178,7 +178,7 @@ def test_read_do_shuffle_write(self):
         maptask.WorkerShuffleWrite(shuffle_kind='group_keys',
                                    shuffle_writer_config='none',
                                    input=(1, 0),
-                                   coders=self.SHUFFLE_CODERS)
+                                   coder=self.SHUFFLE_CODER)
     ]
     shuffle_sink_mock = mock.MagicMock()
     executor.MapTaskExecutor().execute(
@@ -195,7 +195,7 @@ def test_shuffle_read_do_write(self):
         maptask.WorkerGroupingShuffleRead(shuffle_reader_config='none',
                                           start_shuffle_position='aaa',
                                           end_shuffle_position='zzz',
-                                          coders=self.SHUFFLE_CODERS),
+                                          coder=self.SHUFFLE_CODER),
         maptask.WorkerDoFn(serialized_fn=pickle_with_side_inputs(
             ptransform.CallableWrapperDoFn(
                 lambda (k, vs): [str((k, v)) for v in vs])),
@@ -223,7 +223,7 @@ def test_ungrouped_shuffle_read_and_write(self):
         maptask.WorkerUngroupedShuffleRead(shuffle_reader_config='none',
                                            start_shuffle_position='aaa',
                                            end_shuffle_position='zzz',
-                                           coders=self.SHUFFLE_CODERS),
+                                           coder=self.SHUFFLE_CODER),
         maptask.WorkerWrite(
             fileio.TextFileSink(file_path_prefix=output_path,
                                 append_trailing_newlines=True,

google/cloud/dataflow/worker/maptask.py

@@ -97,7 +97,7 @@ def worker_object_to_string(worker_object):
 WorkerGroupingShuffleRead = build_worker_instruction(
     'WorkerGroupingShuffleRead',
     ['start_shuffle_position', 'end_shuffle_position',
-     'shuffle_reader_config', 'coders'])
+     'shuffle_reader_config', 'coder'])
 """Worker details needed to read from a grouping shuffle source.
 
 Attributes:
@@ -108,14 +108,14 @@ def worker_object_to_string(worker_object):
   shuffle_reader_config: An opaque string used to initialize the shuffle
     reader. Contains things like connection endpoints for the shuffle
     server appliance and various options.
-  coders: A 2-tuple of coders (key, value) to decode shuffle entries.
+  coder: The KV coder used to decode shuffle entries.
 """
 
 
 WorkerUngroupedShuffleRead = build_worker_instruction(
     'WorkerUngroupedShuffleRead',
     ['start_shuffle_position', 'end_shuffle_position',
-     'shuffle_reader_config', 'coders'])
+     'shuffle_reader_config', 'coder'])
 """Worker details needed to read from an ungrouped shuffle source.
 
 Attributes:
@@ -126,7 +126,7 @@ def worker_object_to_string(worker_object):
   shuffle_reader_config: An opaque string used to initialize the shuffle
     reader. Contains things like connection endpoints for the shuffle
     server appliance and various options.
-  coders: A 2-tuple of coders (key, value) to decode shuffle entries.
+  coder: The value coder used to decode shuffle entries.
 """
 
 
@@ -160,7 +160,7 @@ def worker_object_to_string(worker_object):
 
 WorkerShuffleWrite = build_worker_instruction(
     'WorkerShuffleWrite',
-    ['shuffle_kind', 'shuffle_writer_config', 'input', 'coders'])
+    ['shuffle_kind', 'shuffle_writer_config', 'input', 'coder'])
 """Worker details needed to write to a shuffle sink.
 
 Attributes:
@@ -173,7 +173,8 @@ def worker_object_to_string(worker_object):
   input: A (producer index, output index) tuple representing the
     ParallelInstruction operation whose output feeds into this operation.
     The output index is 0 except for multi-output operations (like ParDo).
-  coders: A 2-tuple of coders (key, value) to encode shuffle entries.
+  coder: The coder for input elements. If the shuffle_kind is grouping, this is
+    expected to be a KV coder.
 """
 
 
@@ -370,20 +371,23 @@ def get_read_work_item(work, env, context):
   if source:
     return WorkerRead(source, tag=None)
 
-  # TODO(mairbek) create Shuffler Source/Reader
-  kv_coders = get_coder_from_spec(codec_specs, kv_pair=True)
+  coder = get_coder_from_spec(codec_specs)
+  # TODO(ccy): Reconcile WindowedValueCoder wrappings for sources with custom
+  # coders so this special case won't be necessary.
+  if isinstance(coder, coders.WindowedValueCoder):
+    coder = coder.wrapped_value_coder
   if specs['@type'] == 'GroupingShuffleSource':
     return WorkerGroupingShuffleRead(
         start_shuffle_position=specs['start_shuffle_position']['value'],
         end_shuffle_position=specs['end_shuffle_position']['value'],
         shuffle_reader_config=specs['shuffle_reader_config']['value'],
-        coders=kv_coders)
+        coder=coder)
   elif specs['@type'] == 'UngroupedShuffleSource':
     return WorkerUngroupedShuffleRead(
         start_shuffle_position=specs['start_shuffle_position']['value'],
         end_shuffle_position=specs['end_shuffle_position']['value'],
         shuffle_reader_config=specs['shuffle_reader_config']['value'],
-        coders=kv_coders)
+        coder=coder)
   else:
     raise NotImplementedError('Unknown source type: %r' % specs)
 
@@ -452,14 +456,17 @@ def get_write_work_item(work, env, context):
   sink = env.parse_sink(specs, codec_specs, context)
   if sink:
     return WorkerWrite(sink, input=get_input_spec(work.write.input))
-  # TODO(mairbek) create Shuffler Sink/Writer
   if specs['@type'] == 'ShuffleSink':
-    kv_coders = get_coder_from_spec(codec_specs, kv_pair=True)
+    coder = get_coder_from_spec(codec_specs)
+    # TODO(ccy): Reconcile WindowedValueCoder wrappings for sources with custom
+    # coders so this special case won't be necessary.
+    if isinstance(coder, coders.WindowedValueCoder):
+      coder = coder.wrapped_value_coder
     return WorkerShuffleWrite(
         shuffle_kind=specs['shuffle_kind']['value'],
         shuffle_writer_config=specs['shuffle_writer_config']['value'],
         input=get_input_spec(work.write.input),
-        coders=kv_coders)
+        coder=coder)
   else:
     raise NotImplementedError('Unknown sink type: %r' % specs)
 

google/cloud/dataflow/worker/workitem_test.py

@@ -413,7 +413,7 @@ def test_concat_source_to_shuffle_sink(self):
                 shuffle_kind='group_keys',
                 shuffle_writer_config='opaque',
                 input=(1, 0),
-                coders=(CODER.key_coder(), CODER.value_coder()))]))
+                coder=CODER)]))
 
   def test_text_source_to_shuffle_sink(self):
     work = workitem.get_work_items(get_text_source_to_shuffle_sink_message())
@@ -433,7 +433,7 @@ def test_text_source_to_shuffle_sink(self):
                 shuffle_kind='group_keys',
                 shuffle_writer_config='opaque',
                 input=(1, 0),
-                coders=(CODER.key_coder(), CODER.value_coder()))]))
+                coder=CODER)]))
 
   def test_shuffle_source_to_text_sink(self):
     work = workitem.get_work_items(
@@ -445,7 +445,7 @@ def test_shuffle_source_to_text_sink(self):
                 start_shuffle_position='opaque',
                 end_shuffle_position='opaque',
                 shuffle_reader_config='opaque',
-                coders=(CODER.key_coder(), CODER.value_coder())),
+                coder=CODER),
             maptask.WorkerWrite(io.TextFileSink(
                 file_path_prefix='gs://somefile',
                 append_trailing_newlines=True,
@@ -461,7 +461,7 @@ def test_ungrouped_shuffle_source_to_text_sink(self):
                 start_shuffle_position='opaque',
                 end_shuffle_position='opaque',
                 shuffle_reader_config='opaque',
-                coders=(CODER.key_coder(), CODER.value_coder())),
+                coder=CODER),
             maptask.WorkerWrite(io.TextFileSink(
                 file_path_prefix='gs://somefile',
                 append_trailing_newlines=True,