Use multiple file rename threads in finalize_write

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

Author: aaltay

google/cloud/dataflow/io/fileio.py

@@ -19,6 +19,7 @@
 import glob
 import gzip
 import logging
+from multiprocessing.pool import ThreadPool
 import os
 import re
 import shutil
@@ -285,6 +286,13 @@ class FileSink(iobase.Sink):
   The output of this write is a PCollection of all written shards.
   """
 
+  # Approximate number of write results be assigned for each rename thread.
+  _WRITE_RESULTS_PER_RENAME_THREAD = 100
+
+  # Max number of threads to be used for renaming even if it means each thread
+  # will process more write results.
+  _MAX_RENAME_THREADS = 64
+
   def __init__(self,
                file_path_prefix,
                coder,
@@ -346,22 +354,56 @@ def open_writer(self, init_result, uid):
   def finalize_write(self, init_result, writer_results):
     writer_results = sorted(writer_results)
     num_shards = len(writer_results)
-    # TODO(robertwb): Threadpool?
     channel_factory = ChannelFactory()
+    num_threads = max(1, min(
+        num_shards / FileSink._WRITE_RESULTS_PER_RENAME_THREAD,
+        FileSink._MAX_RENAME_THREADS))
+
+    rename_ops = []
     for shard_num, shard in enumerate(writer_results):
       final_name = ''.join([
           self.file_path_prefix,
           self.shard_name_format % dict(shard_num=shard_num,
                                         num_shards=num_shards),
           self.file_name_suffix])
+      rename_ops.append((shard, final_name))
+
+    logging.info(
+        'Starting finalize_write threads with num_shards: %d, num_threads: %d',
+        num_shards, num_threads)
+    start_time = time.time()
+
+    # Use a thread pool for renaming operations.
+    def _rename_file(rename_op):
+      """_rename_file executes single (old_name, new_name) rename operation."""
+      old_name, final_name = rename_op
       try:
-        channel_factory.rename(shard, final_name)
-      except IOError:
+        channel_factory.rename(old_name, final_name)
+      except IOError as e:
         # May have already been copied.
-        print shard, final_name, os.path.exists(final_name)
-        if not channel_factory.exists(final_name):
-          raise
-      yield final_name
+        exists = channel_factory.exists(final_name)
+        if not exists:
+          logging.warning(('IOError in _rename_file. old_name: %s, '
+                           'final_name: %s, err: %s'), old_name, final_name, e)
+          return(None, e)
+      except Exception as e:  # pylint: disable=broad-except
+        logging.warning(('Exception in _rename_file. old_name: %s, '
+                         'final_name: %s, err: %s'), old_name, final_name, e)
+        return(None, e)
+      return (final_name, None)
+
+    rename_results = ThreadPool(num_threads).map(_rename_file, rename_ops)
+
+    for final_name, err in rename_results:
+      if err:
+        logging.warning('Error when processing rename_results: %s', err)
+        raise err
+      else:
+        yield final_name
+
+    logging.info('Renamed %d shards in %.2f seconds.',
+                 num_shards, time.time() - start_time)
+
     try:
       channel_factory.rmdir(init_result)
     except IOError:

google/cloud/dataflow/io/fileio_test.py

@@ -18,6 +18,7 @@
 import glob
 import gzip
 import logging
+import os
 import tempfile
 import unittest
 
@@ -426,7 +427,7 @@ def test_file_sink_writing(self):
     self.assertEqual(open(shard2).read(), '[start][x][y][z][end]')
 
     # Check that any temp files are deleted.
-    self.assertEqual([shard1, shard2], sorted(glob.glob(temp_path + '*')))
+    self.assertItemsEqual([shard1, shard2], glob.glob(temp_path + '*'))
 
   def test_empty_write(self):
     temp_path = tempfile.NamedTemporaryFile().name
@@ -436,7 +437,6 @@ def test_empty_write(self):
     p = df.Pipeline('DirectPipelineRunner')
     p | df.Create([]) | df.io.Write(sink)  # pylint: disable=expression-not-assigned
     p.run()
-
     self.assertEqual(open(temp_path + '-00000-of-00001.foo').read(),
                      '[start][end]')
 
@@ -457,6 +457,66 @@ def test_fixed_shard_write(self):
     self.assertTrue('][a][' in concat, concat)
     self.assertTrue('][b][' in concat, concat)
 
+  def test_file_sink_multi_shards(self):
+    temp_path = tempfile.NamedTemporaryFile().name
+    sink = MyFileSink(temp_path,
+                      file_name_suffix='.foo',
+                      coder=coders.ToStringCoder())
+
+    # Manually invoke the generic Sink API.
+    init_token = sink.initialize_write()
+
+    num_shards = 1000
+    writer_results = []
+    for i in range(num_shards):
+      uuid = 'uuid-%05d' % i
+      writer = sink.open_writer(init_token, uuid)
+      writer.write('a')
+      writer.write('b')
+      writer.write(uuid)
+      writer_results.append(writer.close())
+
+    res_first = list(sink.finalize_write(init_token, writer_results))
+    # Retry the finalize operation (as if the first attempt was lost).
+    res_second = list(sink.finalize_write(init_token, writer_results))
+
+    self.assertItemsEqual(res_first, res_second)
+
+    res = sorted(res_second)
+    for i in range(num_shards):
+      shard_name = '%s-%05d-of-%05d.foo' % (temp_path, i, num_shards)
+      uuid = 'uuid-%05d' % i
+      self.assertEqual(res[i], shard_name)
+      self.assertEqual(
+          open(shard_name).read(), ('[start][a][b][%s][end]' % uuid))
+
+    # Check that any temp files are deleted.
+    self.assertItemsEqual(res, glob.glob(temp_path + '*'))
+
+  def test_file_sink_io_error(self):
+    temp_path = tempfile.NamedTemporaryFile().name
+    sink = MyFileSink(temp_path,
+                      file_name_suffix='.foo',
+                      coder=coders.ToStringCoder())
+
+    # Manually invoke the generic Sink API.
+    init_token = sink.initialize_write()
+
+    writer1 = sink.open_writer(init_token, '1')
+    writer1.write('a')
+    writer1.write('b')
+    res1 = writer1.close()
+
+    writer2 = sink.open_writer(init_token, '2')
+    writer2.write('x')
+    writer2.write('y')
+    writer2.write('z')
+    res2 = writer2.close()
+
+    os.remove(res2)
+    with self.assertRaises(IOError):
+      list(sink.finalize_write(init_token, [res1, res2]))
+
 if __name__ == '__main__':
   logging.getLogger().setLevel(logging.INFO)
   unittest.main()