Several changes to support cross-language transforms on DataflowRunner

chamikaramj · chamikaramj · commit 14e5632ea053 · 2019-09-18T16:44:42.000-07:00
diff --git a/sdks/python/apache_beam/runners/dataflow/dataflow_runner.py b/sdks/python/apache_beam/runners/dataflow/dataflow_runner.py
@@ -858,13 +858,31 @@ def run_ParDo(self, transform_node, options):
     outputs = []
     step.encoding = self._get_encoded_output_coder(transform_node)
 
+    all_output_tags = transform_proto.outputs.keys()
+
+    from apache_beam.transforms.core import RunnerAPIPTransformHolder
+    external_transform = isinstance(transform, RunnerAPIPTransformHolder)
+
+    # Some external transforms require output tags to not be modified.
+    # So we randomly select one of the output tags as the main output and
+    # leave others as side outputs. Transform execution should not change
+    # dependending on which output tag we choose as the main output here.
+    # Also, some SDKs do not work correctly if output tags are modified. So for
+    # external transforms, we leave tags unmodified.
+    main_output_tag = (
+        all_output_tags[0] if external_transform else PropertyNames.OUT)
+
+    # Python SDK uses 'None' as the tag of the main output.
+    tag_to_ignore = main_output_tag if external_transform else 'None'
+    side_output_tags = set(all_output_tags).difference({tag_to_ignore})
+
     # Add the main output to the description.
     outputs.append(
         {PropertyNames.USER_NAME: (
             '%s.%s' % (transform_node.full_label, PropertyNames.OUT)),
          PropertyNames.ENCODING: step.encoding,
-         PropertyNames.OUTPUT_NAME: PropertyNames.OUT})
-    for side_tag in transform.output_tags:
+         PropertyNames.OUTPUT_NAME: main_output_tag})
+    for side_tag in side_output_tags:
       # The assumption here is that all outputs will have the same typehint
       # and coder as the main output. This is certainly the case right now
       # but conceivably it could change in the future.
@@ -873,7 +891,8 @@ def run_ParDo(self, transform_node, options):
               '%s.%s' % (transform_node.full_label, side_tag)),
            PropertyNames.ENCODING: step.encoding,
            PropertyNames.OUTPUT_NAME: (
-               '%s_%s' % (PropertyNames.OUT, side_tag))})
+               side_tag if external_transform
+               else '%s_%s' % (PropertyNames.OUT, side_tag))})
 
     step.add_property(PropertyNames.OUTPUT_INFO, outputs)
 
diff --git a/sdks/python/apache_beam/runners/pipeline_context.py b/sdks/python/apache_beam/runners/pipeline_context.py
@@ -40,10 +40,10 @@ class Environment(object):
   Provides consistency with how the other componentes are accessed.
   """
   def __init__(self, proto):
-    self._proto = proto
+    self.proto = proto
 
   def to_runner_api(self, context):
-    return self._proto
+    return self.proto
 
   @staticmethod
   def from_runner_api(proto, context):
diff --git a/sdks/python/apache_beam/testing/test_pipeline.py b/sdks/python/apache_beam/testing/test_pipeline.py
@@ -67,7 +67,8 @@ def __init__(self,
                options=None,
                argv=None,
                is_integration_test=False,
-               blocking=True):
+               blocking=True,
+               additional_pipeline_args=None):
     """Initialize a pipeline object for test.
 
     Args:
@@ -88,14 +89,18 @@ def __init__(self,
         test, :data:`False` otherwise.
       blocking (bool): Run method will wait until pipeline execution is
         completed.
+      additional_pipeline_args (List[str]): additional pipeline arguments to be
+        included when construction the pipeline options object.
 
     Raises:
       ~exceptions.ValueError: if either the runner or options argument is not
         of the expected type.
     """
     self.is_integration_test = is_integration_test
     self.not_use_test_runner_api = False
-    self.options_list = self._parse_test_option_args(argv)
+    additional_pipeline_args = additional_pipeline_args or []
+    self.options_list = (
+        self._parse_test_option_args(argv) + additional_pipeline_args)
     self.blocking = blocking
     if options is None:
       options = PipelineOptions(self.options_list)
diff --git a/sdks/python/apache_beam/transforms/core.py b/sdks/python/apache_beam/transforms/core.py
@@ -362,10 +362,16 @@ def to_runner_api(self, context, has_parts=False):
       else:
         env1 = id_to_proto_map[env_id]
         env2 = context.environments[env_id]
-        assert env1.SerializeToString() == env2.SerializeToString(), (
+        assert env1.urn == env2.proto.urn, (
             'Expected environments with the same ID to be equal but received '
+            'environments with different URNs '
             '%r and %r',
-            env1, env2)
+            env1.urn, env2.proto.urn)
+        assert env1.payload == env2.proto.payload, (
+            'Expected environments with the same ID to be equal but received '
+            'environments with different payloads '
+            '%r and %r',
+            env1.payload, env2.proto.payload)
     return self._proto
 
   def get_restriction_coder(self):
