Apache Beam - Python Inquiry [closed]

Question

I have created a python pipeline to be run in Dataflow that takes a column from a csv file and counts the number of listings and returns them to a dictionary key value pair.

I am very new to Apache Beam and I wanted to know if I am on the right track. Here is my code thus far. My collection of course is the .csv file, my transforms are in 2 functions, One to extract the proper column from the csv and the second to return the number of neighborhood listings from the csv as a dictionary value. The Neighborhood is the key whilst the number of listings is the value.

import apache_beam as beam 
import sys

PROJECT = 'sacred-ember-308717'
BUCKET = 'airbb-pythonpipelinedataflow'

class SplitFields(beam.DoFn):
    def getFields(self, element):
        id, name, host_id, host_name, neighbourhood_group, neighbourhood, latitude, longitude, room_type, price, minimum_nights, number_of_reviews, last_review, reviews_per_month, calculated_host_listings_count, availability_365 = element.split(",")
        return [{
            'neighbourhood' : int(neighbourhood)
        }]

class CollectListingsCount(beam.DoFn):
    def getFields(self, element):
        #Return a Dictionary containing number of listings per Neighborhood
        dict_result = {};
        for neighboorhood in element['neighbourhood']:
            if neighboorhood in dict_result:
                dict_result[neighboorhood] += 1
            else:
                dict_result[neighboorhood] = 1
        return dict_result

def run():
    argv = [
        '--project={0}'.format(PROJECT),
        '--job_name=examplejob2',
        '--save_main_session',
        '--staging_location=gs://{0}/staging/'.format(BUCKET),
        '--temp_location=gs://{0}/staging/'.format(BUCKET),
        '--region=us-central1',
        '--runner=DataflowRunner'
    ]

    myPipeline = beam.Pipeline(argv=argv)
    inputFile = BUCKET + '/*.csv'
    outputFile = BUCKET + '/tmp/output/'

    #Read from Input File
    (myPipeline 
        |'GetCSV' >> beam.io.ReadFromText(inputFile) 
        |'SplitFields' >> beam.ParDo(SplitFields())
        |'CollectListings' >> beam.ParDo(CollectListingsCount())
        |'FormulateDictionary' >> beam.GroupByKey()
        |'write' >> beam.io.WriteToText(outputFile)
    )

    myPipeline.run()

if __name__=='__main__':
    run()

Please advise! Thanks!

Answer 1

This is a good start, a few remarks / things to change:

• When subclassing beam.Don, the logic of the class needed for execution must be in a method process with the signature def process(self, element)

• I would advise you to use the argparse because writing the arguments as a string will make your life harder. By that I mean:

def run(argv=None):
   # your pipeline parameters 
   parser = argparse.ArgumentParser()
   parser.add_argument(
      '--bucket',
      dest='bucket',
      required=True,
      help='where to read and write files.')
   parser.add_argument(
      '--input',
      dest='input',
      required=True,
      help='input file.')
    parser.add_argument(
      '--output',
      dest='output',
      required=True,
      help='output file.')
    # this will tell apart all of the default pipeline parameters
    # from your parameters
    known_args, pipeline_args = parser.parse_known_args(argv)

    options = PipelineOptions(pipeline_args)
    myPipeline=beam.Pipeline(options=options)
    
    inputFile os.path.join('gs://', known_args.bucket, known_args.input)
    outputFile os.path.join('gs://', known_args.bucket, known_args.output)

    # Read from Input File
    (myPipeline 
       |'GetCSV' >> beam.io.ReadFromText(inputFile) 
       |'SplitFields' >> beam.ParDo(SplitFields())
       |'CollectListings' >> beam.ParDo(CollectListingsCount())
       |'FormulateDictionary' >> beam.GroupByKey()
       |'write' >> beam.io.WriteToText(outputFile)
    )

    myPipeline.run()

And now you can run your pipeline as a CLI as follows:

python -m yourFile.py \
  --runner DataflowRunner \
  --project PROJECT \
  --job_name examplejob2 \
  --region us-central1 \
  --temp_location gs://BUCKET/staging/ \
  --bucket BUCKET \
  --input "/*.csv" \
  --output "/tmp/output/"