From f7b5352f2809cbc6d83da21062188f557839a15d Mon Sep 17 00:00:00 2001 From: Joey Tran Date: Sat, 20 Sep 2025 06:37:32 -0400 Subject: [PATCH 1/3] add test and potential fix --- sdks/python/apache_beam/pipeline.py | 4 +- sdks/python/apache_beam/pipeline_test.py | 68 ++++++++++++++++++++++++ 2 files changed, 71 insertions(+), 1 deletion(-) diff --git a/sdks/python/apache_beam/pipeline.py b/sdks/python/apache_beam/pipeline.py index 0ed5a435e788..94dc9307b8ae 100644 --- a/sdks/python/apache_beam/pipeline.py +++ b/sdks/python/apache_beam/pipeline.py @@ -840,7 +840,9 @@ def apply( # 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..503bb2b0c690 100644 --- a/sdks/python/apache_beam/pipeline_test.py +++ b/sdks/python/apache_beam/pipeline_test.py @@ -1562,6 +1562,74 @@ def file_artifact(path, hash, staged_name): self.assertEqual(len(proto.components.environments), 6) + def test_multiple_outputs_composite_ptransform(self): + from apache_beam.runners.render import RenderRunner + + class NoOpTransform(beam.PTransform): + def expand(self, pcoll): + return pcoll | 'No-Op' >> beam.Map(lambda x: x) + + 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 | 'Split Sales' >> 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 + }, + ] + + optsions = beam.options.pipeline_options.PipelineOptions( + "--render_output=./twopt.dot".split()) + with beam.Pipeline(runner=RenderRunner(), options=optsions) as pipeline: + sales_records = pipeline | 'Create Sales' >> beam.Create(sales_data) + + # Split the PCollection into two tagged outputs + split_collections = sales_records | 'Split Sales' >> ParentSalesSplitter() + + premium_sales = split_collections.premium_sales + standard_sales = split_collections.standard_sales + + # Apply the same composite transform to both PCollections + premium_sales_noop = premium_sales | 'Premium No-Op' >> NoOpTransform() + standard_sales_noop = standard_sales | 'Standard No-Op' >> NoOpTransform() + + # Print the results to verify they are still there + _ = premium_sales_noop | 'Print Premium After No-Op' >> beam.Map(print) + _ = standard_sales_noop | 'Print Standard After No-Op' >> beam.Map(print) + 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'] + assert len(xform.outputs) == 2 + if __name__ == '__main__': unittest.main() From a337dfc47a4f06fad6fda78d26e2cfbc0627abe6 Mon Sep 17 00:00:00 2001 From: Joey Tran Date: Mon, 22 Sep 2025 08:19:56 -0400 Subject: [PATCH 2/3] Simplify test --- sdks/python/apache_beam/pipeline.py | 2 -- sdks/python/apache_beam/pipeline_test.py | 29 ++++-------------------- 2 files changed, 4 insertions(+), 27 deletions(-) diff --git a/sdks/python/apache_beam/pipeline.py b/sdks/python/apache_beam/pipeline.py index 94dc9307b8ae..791adfbfec17 100644 --- a/sdks/python/apache_beam/pipeline.py +++ b/sdks/python/apache_beam/pipeline.py @@ -837,8 +837,6 @@ 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): for tag, pc in list(result._pcolls.items()): if tag not in current.outputs: diff --git a/sdks/python/apache_beam/pipeline_test.py b/sdks/python/apache_beam/pipeline_test.py index 503bb2b0c690..0aeb30dfbb59 100644 --- a/sdks/python/apache_beam/pipeline_test.py +++ b/sdks/python/apache_beam/pipeline_test.py @@ -1563,12 +1563,6 @@ def file_artifact(path, hash, staged_name): self.assertEqual(len(proto.components.environments), 6) def test_multiple_outputs_composite_ptransform(self): - from apache_beam.runners.render import RenderRunner - - class NoOpTransform(beam.PTransform): - def expand(self, pcoll): - return pcoll | 'No-Op' >> beam.Map(lambda x: x) - class SalesSplitter(beam.DoFn): def process(self, element): price = element['price'] @@ -1579,8 +1573,8 @@ def process(self, element): class ParentSalesSplitter(beam.PTransform): def expand(self, pcoll): - return pcoll | 'Split Sales' >> beam.ParDo( - SalesSplitter()).with_outputs('premium_sales', 'standard_sales') + return pcoll | beam.ParDo(SalesSplitter()).with_outputs( + 'premium_sales', 'standard_sales') sales_data = [ { @@ -1600,24 +1594,9 @@ def expand(self, pcoll): }, ] - optsions = beam.options.pipeline_options.PipelineOptions( - "--render_output=./twopt.dot".split()) - with beam.Pipeline(runner=RenderRunner(), options=optsions) as pipeline: + with beam.Pipeline() as pipeline: sales_records = pipeline | 'Create Sales' >> beam.Create(sales_data) - - # Split the PCollection into two tagged outputs - split_collections = sales_records | 'Split Sales' >> ParentSalesSplitter() - - premium_sales = split_collections.premium_sales - standard_sales = split_collections.standard_sales - - # Apply the same composite transform to both PCollections - premium_sales_noop = premium_sales | 'Premium No-Op' >> NoOpTransform() - standard_sales_noop = standard_sales | 'Standard No-Op' >> NoOpTransform() - - # Print the results to verify they are still there - _ = premium_sales_noop | 'Print Premium After No-Op' >> beam.Map(print) - _ = standard_sales_noop | 'Print Standard After No-Op' >> beam.Map(print) + _ = sales_records | 'Split Sales' >> ParentSalesSplitter() current_transforms = list(pipeline.transforms_stack) all_applied_transforms = { xform.full_label: xform From bf04449e8bb2e86e6f94ba00b26b259f983d9ecc Mon Sep 17 00:00:00 2001 From: Joey Tran Date: Mon, 22 Sep 2025 08:24:24 -0400 Subject: [PATCH 3/3] add some comments --- sdks/python/apache_beam/pipeline_test.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/sdks/python/apache_beam/pipeline_test.py b/sdks/python/apache_beam/pipeline_test.py index 0aeb30dfbb59..ab9da013bcf9 100644 --- a/sdks/python/apache_beam/pipeline_test.py +++ b/sdks/python/apache_beam/pipeline_test.py @@ -1563,6 +1563,10 @@ 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'] @@ -1607,6 +1611,8 @@ def expand(self, pcoll): 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