Treat creation of side input views as a PTransform

This change introduces PCollectionView as a specialization of PValue.

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

Author: charlesccychen

google/cloud/dataflow/dataflow_test.py

@@ -250,19 +250,91 @@ def match(actual):
     assert_that(results, matcher(1, a_list, some_pairs))
     pipeline.run()
 
-  def test_as_list_without_unique_labels(self):
-    a_list = [1, 2, 3]
+  def test_as_singleton_without_unique_labels(self):
+    # This should succeed as calling AsSingleton on the same PCollection twice
+    # with the same defaults will return the same PCollectionView.
+    a_list = [2]
+    pipeline = Pipeline('DirectPipelineRunner')
+    main_input = pipeline | Create('main input', [1])
+    side_list = pipeline | Create('side list', a_list)
+    results = main_input | FlatMap(
+        'test',
+        lambda x, s1, s2: [[x, s1, s2]],
+        AsSingleton(side_list), AsSingleton(side_list))
+
+    def  matcher(expected_elem, expected_singleton):
+      def match(actual):
+        [[actual_elem, actual_singleton1, actual_singleton2]] = actual
+        equal_to([expected_elem])([actual_elem])
+        equal_to([expected_singleton])([actual_singleton1])
+        equal_to([expected_singleton])([actual_singleton2])
+      return match
+
+    assert_that(results, matcher(1, 2))
+    pipeline.run()
+
+  def test_as_singleton_with_different_defaults_without_unique_labels(self):
+    # This should fail as AsSingleton with distinct default values should create
+    # distinct PCollectionViews with the same full_label.
+    a_list = [2]
     pipeline = Pipeline('DirectPipelineRunner')
     main_input = pipeline | Create('main input', [1])
     side_list = pipeline | Create('side list', a_list)
+
     with self.assertRaises(RuntimeError) as e:
       _ = main_input | FlatMap(
           'test',
-          lambda x, ls1, ls2: [[x, ls1, ls2]],
-          AsList(side_list), AsList(side_list))
+          lambda x, s1, s2: [[x, s1, s2]],
+          AsSingleton(side_list), AsSingleton(side_list, default_value=3))
     self.assertTrue(
         e.exception.message.startswith(
-            'Transform "AsList" does not have a stable unique label.'))
+            'Transform "ViewAsSingleton(side list.None)" does not have a '
+            'stable unique label.'))
+
+  def test_as_singleton_with_different_defaults_with_unique_labels(self):
+    a_list = []
+    pipeline = Pipeline('DirectPipelineRunner')
+    main_input = pipeline | Create('main input', [1])
+    side_list = pipeline | Create('side list', a_list)
+    results = main_input | FlatMap(
+        'test',
+        lambda x, s1, s2: [[x, s1, s2]],
+        AsSingleton('si1', side_list, default_value=2),
+        AsSingleton('si2', side_list, default_value=3))
+
+    def  matcher(expected_elem, expected_singleton1, expected_singleton2):
+      def match(actual):
+        [[actual_elem, actual_singleton1, actual_singleton2]] = actual
+        equal_to([expected_elem])([actual_elem])
+        equal_to([expected_singleton1])([actual_singleton1])
+        equal_to([expected_singleton2])([actual_singleton2])
+      return match
+
+    assert_that(results, matcher(1, 2, 3))
+    pipeline.run()
+
+  def test_as_list_without_unique_labels(self):
+    # This should succeed as calling AsList on the same PCollection twice will
+    # return the same PCollectionView.
+    a_list = [1, 2, 3]
+    pipeline = Pipeline('DirectPipelineRunner')
+    main_input = pipeline | Create('main input', [1])
+    side_list = pipeline | Create('side list', a_list)
+    results = main_input | FlatMap(
+        'test',
+        lambda x, ls1, ls2: [[x, ls1, ls2]],
+        AsList(side_list), AsList(side_list))
+
+    def  matcher(expected_elem, expected_list):
+      def match(actual):
+        [[actual_elem, actual_list1, actual_list2]] = actual
+        equal_to([expected_elem])([actual_elem])
+        equal_to(expected_list)(actual_list1)
+        equal_to(expected_list)(actual_list2)
+      return match
+
+    assert_that(results, matcher(1, [1, 2, 3]))
+    pipeline.run()
 
   def test_as_list_with_unique_labels(self):
     a_list = [1, 2, 3]
@@ -282,6 +354,9 @@ def match(actual):
         equal_to(expected_list)(actual_list2)
       return match
 
+    assert_that(results, matcher(1, [1, 2, 3]))
+    pipeline.run()
+
   def test_as_dict_with_unique_labels(self):
     some_kvs = [('a', 1), ('b', 2)]
     pipeline = Pipeline('DirectPipelineRunner')

google/cloud/dataflow/internal/apiclient.py

@@ -124,6 +124,8 @@ class Step(object):
   """Wrapper for a dataflow Step protobuf."""
 
   def __init__(self, step_kind, step_name):
+    self.step_kind = step_kind
+    self.step_name = step_name
     self.proto = dataflow.Step(kind=step_kind, name=step_name)
     self.proto.properties = {}
 

google/cloud/dataflow/pipeline.py

@@ -137,6 +137,9 @@ def __init__(self, runner=None, options=None, argv=None):
     # If a transform is applied and the full label is already in the set
     # then the transform will have to be cloned with a new label.
     self.applied_labels = set()
+    # Store cache of views created from PCollections.  For reference, see
+    # pvalue._cache_view().
+    self._view_cache = {}
 
   def _add_pvalue(self, pval):
     """Adds a PValue to the pipeline's node list."""
@@ -263,7 +266,8 @@ def apply(self, transform, pvalueish=None):
       # TODO(robertwb): Multi-input, multi-output inference.
       # TODO(robertwb): Ideally we'd do intersection here.
       if (type_options is not None and type_options.pipeline_type_check and
-          isinstance(result, pvalue.PCollection) and not result.element_type):
+          isinstance(result, (pvalue.PCollection, pvalue.PCollectionView))
+          and not result.element_type):
         input_element_type = (
             inputs[0].element_type
             if len(inputs) == 1
@@ -377,7 +381,7 @@ def real_producer(pv):
         if not isinstance(main_input, pvalue.PBegin):
           real_producer(main_input).refcounts[main_input.tag] += 1
       for side_input in self.side_inputs:
-        real_producer(side_input.pvalue).refcounts[side_input.pvalue.tag] += 1
+        real_producer(side_input).refcounts[side_input.tag] += 1
 
   def add_output(self, output):
     assert (isinstance(output, pvalue.PValue) or
@@ -409,8 +413,7 @@ def visit(self, visitor, pipeline, visited):
 
     # Visit side inputs.
     for pval in self.side_inputs:
-      if isinstance(pval, pvalue.AsSideInput) and pval.pvalue not in visited:
-        pval = pval.pvalue  # Unpack marker-object-wrapped pvalue.
+      if isinstance(pval, pvalue.PCollectionView) and pval not in visited:
         assert pval.producer is not None
         pval.producer.visit(visitor, pipeline, visited)
         # The value should be visited now since we visit outputs too.