apache-beam

import re import argparse import apache_beam as beam from apache_beam.io import ReadFromText, WriteToText from apache_beam.options.pipeline_options import PipelineOptions def main(argv=None): parser = argparse.ArgumentParser() parser.add_argument( '--input', dest='input', default='gs://dataflow-samples/shakespeare/kinglear.txt', help='Input file to process.') parser.add_argument( '--output', dest='output', default='output.txt', help='Output file to write results to.') known_args, pipeline_args = parser.parse_known_args(argv) pipeline_options = PipelineOptions(pipeline_args) with beam.Pipeline(options=pipeline_options) as p: # Read the text file into a PCollection lines = p | ReadFromText(known_args.input) # Count the occurrences of each word counts = ( lines | 'Split' >> beam.FlatMap(lambda x: re.findall(r'[A-Za-z\']+', x)) | 'PairWithOne' >> beam.Map(lambda x: (x, 1)) | 'GroupAndSum' >> beam.CombinePerKey(sum) ) # Format the counts into strings output = counts | 'Format' >> beam.Map( lambda word_count: '%s: %s' % (word_count[0], word_count[1])) # Write the output output | WriteToText(known_args.output) if __name__ == '__main__': print("Running Beam WordCount pipeline locally...") main() print("Pipeline finished. Check 'output.txt' for results.")