[BEAM-11862] Updates ExternalTransformFinder to consider external leaf transforms

sdks/python/apache_beam/pipeline.py

@@ -995,12 +995,26 @@ def contains_external_transforms(pipeline):
     pipeline.visit(visitor)
     return visitor._contains_external_transforms
 
-  def visit_transform(self, transform_node):
-    # type: (AppliedPTransform) -> None
+  def _perform_exernal_transform_test(self, transform):
+    if not transform:
+      return
     from apache_beam.transforms import ExternalTransform
-    if isinstance(transform_node.transform, ExternalTransform):
+    if isinstance(transform, ExternalTransform):
       self._contains_external_transforms = True
 
+  def visit_transform(self, transform_node):
+    # type: (AppliedPTransform) -> None
+    self._perform_exernal_transform_test(transform_node.transform)
+
+  def enter_composite_transform(self, transform_node):
+    # type: (AppliedPTransform) -> None
+    # Python SDK object graph may represent an external transform that is a leaf
+    # of the pipeline graph as a composite without sub-transforms.
+    # Note that this visitor is just used to identify pipelines with external
+    # transforms. A Runner API pipeline proto generated from the Pipeline object
+    # will include external sub-transform.
+    self._perform_exernal_transform_test(transform_node.transform)
+
 
 class AppliedPTransform(object):
   """For internal use only; no backwards-compatibility guarantees.

sdks/python/apache_beam/runners/portability/expansion_service_test.py

@@ -56,6 +56,9 @@
 TEST_PARTITION_URN = "beam:transforms:xlang:test:partition"
 TEST_PYTHON_BS4_URN = "beam:transforms:xlang:test:python_bs4"
 
+# A transform that does not produce an output.
+TEST_NO_OUTPUT_URN = "beam:transforms:xlang:test:nooutput"
+
 
 @ptransform.PTransform.register_urn('beam:transforms:xlang:count', None)
 class CountPerElementTransform(ptransform.PTransform):
@@ -296,6 +299,23 @@ def from_runner_api_parameter(unused_ptransform, level, unused_context):
     return FibTransform(int(level.decode('ascii')))
 
 
+@ptransform.PTransform.register_urn(TEST_NO_OUTPUT_URN, None)
+class NoOutputTransform(ptransform.PTransform):
+  def expand(self, pcoll):
+    def log_val(val):
+      logging.debug('Got value: %r', val)
+
+    # Logging without returning anything
+    _ = (pcoll | 'TestLabel' >> beam.ParDo(log_val))
+
+  def to_runner_api_parameter(self, unused_context):
+    return TEST_NO_OUTPUT_URN, None
+
+  @staticmethod
+  def from_runner_api_parameter(unused_ptransform, payload, unused_context):
+    return NoOutputTransform(parse_string_payload(payload)['data'])
+
+
 def parse_string_payload(input_byte):
   payload = ExternalConfigurationPayload()
   payload.ParseFromString(input_byte)

sdks/python/apache_beam/transforms/external_test.py

@@ -259,6 +259,33 @@ def test_unique_name(self):
     self.assertEqual(
         len(set(pcolls)), len(pcolls), msg='PCollection names are not unique.')
 
+  def test_external_transform_finder_non_leaf(self):
+    pipeline = beam.Pipeline()
+    _ = (
+        pipeline
+        | beam.Create(['a', 'b'])
+        | beam.ExternalTransform(
+            'beam:transforms:xlang:test:prefix',
+            ImplicitSchemaPayloadBuilder({'data': u'0'}),
+            expansion_service.ExpansionServiceServicer())
+        | beam.Map(lambda x: x))
+    pipeline.run().wait_until_finish()
+
+    self.assertTrue(pipeline.contains_external_transforms)
+
+  def test_external_transform_finder_leaf(self):
+    pipeline = beam.Pipeline()
+    _ = (
+        pipeline
+        | beam.Create(['a', 'b'])
+        | beam.ExternalTransform(
+            'beam:transforms:xlang:test:nooutput',
+            ImplicitSchemaPayloadBuilder({'data': u'0'}),
+            expansion_service.ExpansionServiceServicer()))
+    pipeline.run().wait_until_finish()
+
+    self.assertTrue(pipeline.contains_external_transforms)
+
 
 if __name__ == '__main__':
   logging.getLogger().setLevel(logging.INFO)