SO question 60080589

pd-concat.py

import argparse, logging
import pandas as pd
from random import choice

import apache_beam as beam
from apache_beam.options.pipeline_options import SetupOptions

import apache_beam.transforms.combiners as combine
import apache_beam.pvalue as pvalue


def create_dataframes(element):
    project_list = ['Beam', 'Spark', 'Flink', 'Nemo', 'Aiflow']
        
    for _ in range(10):
        data = {'projects' : [choice(project_list)]}
        yield pd.DataFrame(data)


class merge_dataframes(beam.DoFn):
    def process(self, element):
        logging.info(element)
        logging.info(type(element))
        yield pd.concat(element).reset_index(drop=True)


def run_pipeline(argv=None):
    parser = argparse.ArgumentParser()
    known_args, pipeline_args = parser.parse_known_args(argv)

    pipeline_options = beam.pipeline.PipelineOptions(pipeline_args)
    pipeline_options.view_as(SetupOptions).save_main_session = True
    p = beam.Pipeline(options=pipeline_options)

    (p
        | 'Start' >> beam.Create(['Start'])
        | 'Create DataFrames' >> beam.FlatMap(create_dataframes)
        | 'Combine To List' >> beam.combiners.ToList()
        | 'Merge DataFrames' >> beam.ParDo(merge_dataframes())
        | 'Print results' >> beam.Map(print))

    result = p.run()
    result.wait_until_finish()

if __name__ == '__main__':
    logging.getLogger().setLevel(logging.INFO)
    run_pipeline()