Simple Beam CAN Pipeline

simple_beam_can_pipeline.py

import apache_beam as beam

from apache_beam.options.pipeline_options import PipelineOptions

from libs.datasets import combined_datasets
from pyseir.run import OneRegionPipeline


class ProcessOneRegion(beam.DoFn):
    def process(self, one_region):
        yield OneRegionPipeline.run(one_region).infer_df.to_csv()


beam_options = PipelineOptions()
with beam.Pipeline(options=beam_options) as pipeline:
    regions = combined_datasets.load_us_timeseries_dataset().get_subset(
        states=['NY', 'CA'])
    regions = [one_region for _, one_region in regions.iter_one_regions()]
    (pipeline
     | 'Read Regions' >> beam.Create(regions)
     | 'Process One Region' >> beam.ParDo(ProcessOneRegion())
     | 'Write CSV' >> beam.io.WriteToText('out.csv'))