Optimize shuffle writing

* Cache the key and value coder impls on the shuffle
  readers and writers themselves.
* Don't re-encode the key in the secondary key slot.
  Instead, let it be empty, as we're not using it now.

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

Author: robertwb

google/cloud/dataflow/worker/executor.py

@@ -328,10 +328,9 @@ def process(self, o):
       k, v = str(random.getrandbits(64)), o.value
     else:
       k, v = o.value
-    # TODO(silviuc): Figure out what is the proper value for the secondary key.
-    # For now the secondary key is a duplicate of the primary key just because
-    # they both use the same coder.
-    self.writer.Write(k, k, v)
+    # TODO(silviuc): Use timestamps for the secondary key to get values in
+    # times-sorted order.
+    self.writer.Write(k, '', v)
 
 
 class DoOperation(Operation):

google/cloud/dataflow/worker/executor_test.py

@@ -186,8 +186,8 @@ def test_read_do_shuffle_write(self):
         test_shuffle_sink=shuffle_sink_mock)
     # Make sure we have seen all the (k, v) writes.
     shuffle_sink_mock.writer().Write.assert_has_calls(
-        [mock.call('a', 'a', 1), mock.call('b', 'b', 1),
-         mock.call('c', 'c', 1), mock.call('d', 'd', 1)])
+        [mock.call('a', '', 1), mock.call('b', '', 1),
+         mock.call('c', '', 1), mock.call('d', '', 1)])
 
   def test_shuffle_read_do_write(self):
     output_path = self.create_temp_file('n/a')

google/cloud/dataflow/worker/shuffle.py

@@ -309,6 +309,8 @@ def __init__(self, shuffle_source, reader=None):
     self.source = shuffle_source
     self.reader = reader
     self.entries_iterable = None
+    self.key_coder = self.source.key_coder.get_impl()
+    self.value_coder = self.source.value_coder.get_impl()
 
   def __enter__(self):
     if self.reader is None:
@@ -340,7 +342,7 @@ def __iter__(self):
       entries_iterator.push_back(entry)
       key_values = ShuffleKeyValuesIterable(
           entries_iterator,
-          entry.key, self.source.value_coder, entry.position)
+          entry.key, self.value_coder, entry.position)
       group_start = entry.position
 
       last_group_start = self._range_tracker.last_group_start
@@ -354,7 +356,7 @@ def __iter__(self):
         # source.
         return
 
-      yield (self.source.key_coder.decode(entry.key), key_values)
+      yield (self.key_coder.decode(entry.key), key_values)
       # We need to drain the iterator returned just in case this
       # was not done by the caller. Otherwise we will not properly advance
       # to the next key but rather return the next entry for the current
@@ -406,7 +408,7 @@ def __init__(self, shuffle_source, reader=None):
 
   def __iter__(self):
     for entry in self.entries_iterable:
-      yield self.source.value_coder.decode(entry.value)
+      yield self.value_coder.decode(entry.value)
 
 
 class ShuffleSourceBase(iobase.Source):
@@ -451,6 +453,8 @@ def __init__(self, shuffle_sink, writer=None):
     self.writer = writer
     self.stream = StringIO.StringIO()
     self.bytes_buffered = 0
+    self.key_coder = self.sink.key_coder.get_impl()
+    self.value_coder = self.sink.value_coder.get_impl()
 
   def __enter__(self):
     if self.writer is None:
@@ -468,9 +472,9 @@ def __exit__(self, exception_type, exception_value, traceback):
 
   def Write(self, key, secondary_key, value):
     entry = ShuffleEntry(
-        self.sink.key_coder.encode(key),
-        self.sink.key_coder.encode(secondary_key),
-        self.sink.value_coder.encode(value),
+        self.key_coder.encode(key),
+        secondary_key,
+        self.value_coder.encode(value),
         position=None)
     entry.to_bytes(self.stream, with_position=False)
     self.bytes_buffered += entry.size

google/cloud/dataflow/worker/shuffle_test.py

@@ -19,14 +19,15 @@
 import logging
 import unittest
 
+from google.cloud.dataflow import coders
 from google.cloud.dataflow.io import iobase
 from google.cloud.dataflow.worker.shuffle import GroupedShuffleSource
 from google.cloud.dataflow.worker.shuffle import ShuffleEntry
 from google.cloud.dataflow.worker.shuffle import ShuffleSink
 from google.cloud.dataflow.worker.shuffle import UngroupedShuffleSource
 
 
-class Base64Coder(object):
+class Base64Coder(coders.Coder):
   """Simple base64 coder used throughout the tests."""
 
   def decode(self, o):
@@ -67,7 +68,7 @@ def _make_chunk(self, descriptor, start_index):
     for key, value in descriptor:
       ShuffleEntry(
           coder.encode(key),
-          coder.encode('2nd-%s' % key),
+          '',
           coder.encode(value),
           position=str(position)).to_bytes(stream)
       position += 1
@@ -348,7 +349,7 @@ class TestShuffleSink(unittest.TestCase):
 
   def test_basics(self):
     source = ShuffleSink(config_bytes='not used', coder=Base64Coder())
-    entries = [('a', '2nd-a', '1'), ('b', '2nd-b', '0'), ('b', '2nd-b', '1')]
+    entries = [('a', '', '1'), ('b', '', '0'), ('b', '', '1')]
     fake_writer = FakeShuffleWriter()
     with source.writer(test_writer=fake_writer) as writer:
       for entry in entries: