def expand(self, pcoll):
_ = pcoll | 'key-with-none-a' >> beam.ParDo(
core._KeyWithNone())
_ = pcoll | 'key-with-none-b' >> beam.ParDo(
core._KeyWithNone())
_ = pcoll | 'key-with-none-c' >> beam.ParDo(
core._KeyWithNone())
def test_eliminate_common_key_with_void(self):
pipeline = beam.Pipeline()
pcoll = pipeline | 'Start' >> beam.Create([1, 2, 3])
_ = pcoll | 'TestKeyWithNoneA' >> beam.ParDo(core._KeyWithNone())
_ = pcoll | 'TestKeyWithNoneB' >> beam.ParDo(core._KeyWithNone())
pipeline_proto = pipeline.to_runner_api()
_, stages = translations.create_and_optimize_stages(
pipeline_proto, [translations.eliminate_common_key_with_none],
known_runner_urns=frozenset())
key_with_none_stages = [
stage for stage in stages if 'TestKeyWithNone' in stage.name
]
self.assertEqual(len(key_with_none_stages), 1)
