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)