Support large iterable side inputs

This change provides support in the DataflowPipelineRunner for large iterable
side inputs that do not fit in memory.

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

Author: charlesccychen

google/cloud/dataflow/pvalue.py

@@ -263,6 +263,11 @@ class ListPCollectionView(PCollectionView):
   pass
 
 
+class DictPCollectionView(PCollectionView):
+  """A PCollectionView that can be treated as a dict."""
+  pass
+
+
 def _get_cached_view(pipeline, key):
   return pipeline._view_cache.get(key, None)  # pylint: disable=protected-access
 
@@ -411,24 +416,20 @@ def AsList(pcoll, label=None):  # pylint: disable=invalid-name
 
 @can_take_label_as_first_argument
 def AsDict(pcoll, label=None):  # pylint: disable=invalid-name
-  """Convenience function packaging an entire PCollection as a side input dict.
+  """Create a DictPCollectionView from the elements of input PCollection.
 
-  Intended for use in side-argument specification---the same places where
-  AsSingleton and AsIter are used. Unlike those wrapper classes, AsDict (as
-  implemented) is a function that schedules a Combiner to condense pcoll into a
-  single dict, then wraps the resulting one-element PCollection in AsSingleton.
+  The contents of the given PCollection whose elements are 2-tuples of key and
+  value will be available as a dict-like object in PTransforms that use the
+  returned PCollectionView as a side input.
 
   Args:
-    pcoll: Input pcollection. All elements should be key-value pairs (i.e.
-       2-tuples) with unique keys.
+    pcoll: Input pcollection containing 2-tuples of key and value.
     label: Label to be specified if several AsDict's for the same PCollection.
 
   Returns:
-    A singleton PCollectionView containing the dict as above.
+    A dict PCollectionView containing the dict as above.
   """
   label = label or _format_view_label(pcoll)
-  singleton_label = 'ToDict(%s)' % label
-  combine_label = 'CombineToDict(%s)' % label
 
   # Don't recreate the view if it was already created.
   cache_key = (pcoll, AsDict)
@@ -439,8 +440,8 @@ def AsDict(pcoll, label=None):  # pylint: disable=invalid-name
   # Local import is required due to dependency loop; even though the
   # implementation of this function requires concepts defined in modules that
   # depend on pvalue, it lives in this module to reduce user workload.
-  from google.cloud.dataflow.transforms import combiners  # pylint: disable=g-import-not-at-top
-  view = AsSingleton(singleton_label, pcoll | combiners.ToDict(combine_label))
+  from google.cloud.dataflow.transforms import sideinputs  # pylint: disable=g-import-not-at-top
+  view = (pcoll | sideinputs.ViewAsDict(label=label))
   _cache_view(pcoll.pipeline, cache_key, view)
   return view
 

google/cloud/dataflow/runners/direct_runner.py

@@ -28,6 +28,7 @@
 from google.cloud.dataflow import coders
 from google.cloud.dataflow import error
 from google.cloud.dataflow.io import fileio
+from google.cloud.dataflow.pvalue import DictPCollectionView
 from google.cloud.dataflow.pvalue import EmptySideInput
 from google.cloud.dataflow.pvalue import IterablePCollectionView
 from google.cloud.dataflow.pvalue import ListPCollectionView
@@ -119,6 +120,8 @@ def run_CreatePCollectionView(self, transform_node):
       result = [v.value for v in values]
     elif isinstance(view, ListPCollectionView):
       result = [v.value for v in values]
+    elif isinstance(view, DictPCollectionView):
+      result = dict(v.value for v in values)
     else:
       raise NotImplementedError
 

google/cloud/dataflow/transforms/sideinputs.py

@@ -115,3 +115,31 @@ def apply(self, pcoll):
             | CreatePCollectionView(pvalue.ListPCollectionView(pcoll.pipeline))
             .with_input_types(input_type)
             .with_output_types(output_type))
+
+K = typehints.TypeVariable('K')
+V = typehints.TypeVariable('V')
+@typehints.with_input_types(typehints.Tuple[K, V])
+@typehints.with_output_types(typehints.Dict[K, V])
+class ViewAsDict(PTransform):  # pylint: disable=g-wrong-blank-lines
+  """Transform to view PCollection as a dict PCollectionView.
+
+  Important: this transform is an implementation detail and should not be used
+  directly by pipeline writers. Use pvalue.AsDict(...) instead.
+  """
+
+  def __init__(self, label=None):
+    if label:
+      label = 'ViewAsDict(%s)' % label
+    super(ViewAsDict, self).__init__(label=label)
+
+  def apply(self, pcoll):
+    self._check_pcollection(pcoll)
+    input_type = pcoll.element_type
+    key_type, value_type = (
+        typehints.trivial_inference.key_value_types(input_type))
+    output_type = typehints.Dict[key_type, value_type]
+    return (pcoll
+            | CreatePCollectionView(
+                pvalue.DictPCollectionView(pcoll.pipeline))
+            .with_input_types(input_type)
+            .with_output_types(output_type))

google/cloud/dataflow/worker/executor.py

@@ -43,6 +43,7 @@
 from google.cloud.dataflow.worker import maptask
 from google.cloud.dataflow.worker import opcounters
 from google.cloud.dataflow.worker import shuffle
+from google.cloud.dataflow.worker import sideinputs
 
 
 class ReceiverSet(object):
@@ -201,14 +202,6 @@ def start(self):
           windowed_value = GlobalWindows.WindowedValue(value)
         self.output(windowed_value)
 
-  def side_read_all(self, singleton=False):
-    # TODO(mairbek): Should we return WindowedValue here?
-    with self.spec.source.reader() as reader:
-      for value in reader:
-        yield value
-        if singleton:
-          return
-
   def request_dynamic_split(self, dynamic_split_request):
     if self._reader is not None:
       return self._reader.request_dynamic_split(dynamic_split_request)
@@ -423,36 +416,46 @@ def _read_side_inputs(self, tags_and_types):
     # specification. This can happen for instance if the source has been
     # sharded into several files.
     for side_tag, view_class, view_options in tags_and_types:
-      # Note that currently, the implementation of Iterable and List views
-      # are identical. This may change in the future once we allow very large
-      # side input collections.
-      is_singleton = view_class == pvalue.SingletonPCollectionView
+      sources = []
       # Using the side_tag in the lambda below will trigger a pylint warning.
       # However in this case it is fine because the lambda is used right away
       # while the variable has the value assigned by the current iteration of
       # the for loop.
       # pylint: disable=cell-var-from-loop
-      results = []
       for si in itertools.ifilter(
           lambda o: o.tag == side_tag, self.spec.side_inputs):
-        if isinstance(si, maptask.WorkerSideInputSource):
-          op = ReadOperation(si, self.counter_factory)
-        else:
+        if not isinstance(si, maptask.WorkerSideInputSource):
           raise NotImplementedError('Unknown side input type: %r' % si)
-        for v in op.side_read_all(singleton=is_singleton):
-          results.append(v)
-          if is_singleton:
-            break
-      if is_singleton:
+        sources.append(si.source)
+      iterator_fn = sideinputs.get_iterator_fn_for_sources(sources)
+
+      if view_class == pvalue.SingletonPCollectionView:
         has_default, default = view_options
-        if results:
-          yield results[0]
+        has_result = False
+        result = None
+        for v in iterator_fn():
+          has_result = True
+          result = v
+          break
+        if has_result:
+          yield result
         elif has_default:
           yield default
         else:
           yield EmptySideInput()
+      elif view_class == pvalue.IterablePCollectionView:
+        yield sideinputs.EmulatedIterable(iterator_fn)
+      elif view_class == pvalue.ListPCollectionView:
+        # TODO(ccy): this is not yet suitable for lists that do not fit in
+        # memory on a single machine.
+        yield list(iterator_fn())
+      elif view_class == pvalue.DictPCollectionView:
+        # TODO(ccy): this is not yet suitable for dictionaries that do not fit
+        # in memory on a single machine.
+        yield dict(iterator_fn())
       else:
-        yield results
+        raise NotImplementedError('Unknown PCollectionView type: %s' %
+                                  view_class)
 
   def start(self):
     super(DoOperation, self).start()

google/cloud/dataflow/worker/executor_test.py

@@ -400,7 +400,9 @@ def test_create_do_with_singleton_side_bigquery_write(self):
       # Setup the reader so it will yield the values in 'side_elements'.
       reader_mock = mock_class.return_value
       reader_mock.__enter__.return_value = reader_mock
-      reader_mock.__iter__.return_value = (x for x in side_elements)
+      # Use a lambda so that multiple readers can be created, each reading the
+      # entirety of the side elements.
+      reader_mock.__iter__.side_effect = lambda: (x for x in side_elements)
 
       pickled_elements = [pickler.dumps(e) for e in elements]
       executor.MapTaskExecutor().execute(make_map_task([
@@ -442,7 +444,9 @@ def test_create_do_with_collection_side_bigquery_write(self):
       # Setup the reader so it will yield the values in 'side_elements'.
       reader_mock = mock_class.return_value
       reader_mock.__enter__.return_value = reader_mock
-      reader_mock.__iter__.return_value = (x for x in side_elements)
+      # Use a lambda so that multiple readers can be created, each reading the
+      # entirety of the side elements.
+      reader_mock.__iter__.side_effect = lambda: (x for x in side_elements)
 
       executor.MapTaskExecutor().execute(make_map_task([
           maptask.WorkerRead(

google/cloud/dataflow/worker/sideinputs.py

@@ -0,0 +1,38 @@
+# Copyright 2016 Google Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the 'License');
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an 'AS IS' BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Utilities for handling side inputs."""
+
+import collections
+
+
+
+def get_iterator_fn_for_sources(sources):
+  """Returns callable that returns iterator over elements for given sources."""
+  def _inner():
+    for source in sources:
+      with source.reader() as reader:
+        for value in reader:
+          yield value
+  return _inner
+
+
+class EmulatedIterable(collections.Iterable):
+  """Emulates an iterable for a side input."""
+
+  def __init__(self, iterator_fn):
+    self.iterator_fn = iterator_fn
+
+  def __iter__(self):
+    return self.iterator_fn()

google/cloud/dataflow/worker/sideinputs_test.py

@@ -0,0 +1,49 @@
+# Copyright 2016 Google Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for side input utilities."""
+
+import logging
+import unittest
+
+
+from google.cloud.dataflow.worker import sideinputs
+
+
+class EmulatedCollectionsTest(unittest.TestCase):
+
+  def test_emulated_iterable(self):
+    def _iterable_fn():
+      for i in range(10):
+        yield i
+    iterable = sideinputs.EmulatedIterable(_iterable_fn)
+    # Check that multiple iterations are supported.
+    for _ in range(0, 5):
+      for i, j in enumerate(iterable):
+        self.assertEqual(i, j)
+
+  def test_large_iterable_values(self):
+    def _iterable_fn():
+      for i in range(10):
+        yield ('%d' % i) * (200 * 1024 * 1024)
+    iterable = sideinputs.EmulatedIterable(_iterable_fn)
+    # Check that multiple iterations are supported.
+    for _ in range(0, 3):
+      for i, j in enumerate(iterable):
+        self.assertEqual(('%d' % i) * (200 * 1024 * 1024), j)
+
+
+if __name__ == '__main__':
+  logging.getLogger().setLevel(logging.INFO)
+  unittest.main()