diff --git a/sdks/python/apache_beam/pipeline.py b/sdks/python/apache_beam/pipeline.py index 0ed5a435e788..791adfbfec17 100644 --- a/sdks/python/apache_beam/pipeline.py +++ b/sdks/python/apache_beam/pipeline.py @@ -837,10 +837,10 @@ def apply( self._infer_result_type(transform, tuple(inputs.values()), result) assert isinstance(result.producer.inputs, tuple) - # The DoOutputsTuple adds the PCollection to the outputs when accessed - # except for the main tag. Add the main tag here. if isinstance(result, pvalue.DoOutputsTuple): - current.add_output(result, result._main_tag) + for tag, pc in list(result._pcolls.items()): + if tag not in current.outputs: + current.add_output(pc, tag) continue # If there is already a tag with the same name, increase a counter for diff --git a/sdks/python/apache_beam/pipeline_test.py b/sdks/python/apache_beam/pipeline_test.py index dc0d9a7cc58f..ab9da013bcf9 100644 --- a/sdks/python/apache_beam/pipeline_test.py +++ b/sdks/python/apache_beam/pipeline_test.py @@ -1562,6 +1562,59 @@ def file_artifact(path, hash, staged_name): self.assertEqual(len(proto.components.environments), 6) + def test_multiple_outputs_composite_ptransform(self): + """ + Test that a composite PTransform with multiple outputs is represented + correctly in the pipeline proto. + """ + class SalesSplitter(beam.DoFn): + def process(self, element): + price = element['price'] + if price > 100: + yield beam.pvalue.TaggedOutput('premium_sales', element) + else: + yield beam.pvalue.TaggedOutput('standard_sales', element) + + class ParentSalesSplitter(beam.PTransform): + def expand(self, pcoll): + return pcoll | beam.ParDo(SalesSplitter()).with_outputs( + 'premium_sales', 'standard_sales') + + sales_data = [ + { + 'item': 'Laptop', 'price': 1200 + }, + { + 'item': 'Mouse', 'price': 25 + }, + { + 'item': 'Keyboard', 'price': 75 + }, + { + 'item': 'Monitor', 'price': 350 + }, + { + 'item': 'Headphones', 'price': 90 + }, + ] + + with beam.Pipeline() as pipeline: + sales_records = pipeline | 'Create Sales' >> beam.Create(sales_data) + _ = sales_records | 'Split Sales' >> ParentSalesSplitter() + current_transforms = list(pipeline.transforms_stack) + all_applied_transforms = { + xform.full_label: xform + for xform in current_transforms + } + while current_transforms: + xform = current_transforms.pop() + all_applied_transforms[xform.full_label] = xform + current_transforms.extend(xform.parts) + xform = all_applied_transforms['Split Sales'] + # Confirm that Split Sales correctly has two outputs as specified by + # ParDo.with_outputs in ParentSalesSplitter. + assert len(xform.outputs) == 2 + if __name__ == '__main__': unittest.main()