muassif
62 lines
2.0 KiB
Python
62 lines
2.0 KiB
Python
# wordcount.py: count words in a text file
|
|
|
|
import argparse
|
|
import os
|
|
import re
|
|
|
|
import apache_beam as beam
|
|
from apache_beam.io import ReadFromText
|
|
from apache_beam.io import WriteToText
|
|
from apache_beam.options.pipeline_options import PipelineOptions
|
|
from apache_beam.options.pipeline_options import SetupOptions
|
|
|
|
def run(argv=None, save_main_session=True):
|
|
os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = "/Users/muasif/gcd-projs/gcp-key/word-count-316612-f22f7ffcc2dd.json"
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument(
|
|
'--input',
|
|
dest='input',
|
|
default='gs://muasif/input/sample.txt',
|
|
help='Input file to process.')
|
|
parser.add_argument(
|
|
'--output',
|
|
dest='output',
|
|
default='gs://muasif/output/result',
|
|
help='Output file to write results to.')
|
|
known_args, pipeline_args = parser.parse_known_args(argv)
|
|
pipeline_args.extend([
|
|
'--runner=DataflowRunner',
|
|
'--project=word-count-316612',
|
|
'--region=us-central1',
|
|
'--staging_location=gs://muasif/staging',
|
|
'--temp_location=gs://muasif/temp',
|
|
'--job_name=my-wordcount-job',
|
|
])
|
|
|
|
pipeline_options = PipelineOptions(pipeline_args)
|
|
pipeline_options.view_as(SetupOptions).\
|
|
save_main_session = save_main_session
|
|
with beam.Pipeline(options=pipeline_options) as p:
|
|
|
|
lines = p | ReadFromText(known_args.input)
|
|
# Count the occurrences of each word.
|
|
counts = (
|
|
lines
|
|
| 'Split words' >> (
|
|
beam.FlatMap(
|
|
lambda x: re.findall(r'[A-Za-z\']+', x)).
|
|
with_output_types(str))
|
|
| 'Pair with 1' >> beam.Map(lambda x: (x, 1))
|
|
| 'Group & Sum' >> beam.CombinePerKey(sum))
|
|
|
|
# Format the word counts into a PCollection of strings.
|
|
def format_result(word_count):
|
|
(word, count) = word_count
|
|
return '%s: %s' % (word, count)
|
|
|
|
output = counts | 'Format' >> beam.Map(format_result)
|
|
output | WriteToText(known_args.output)
|
|
|
|
if __name__ == '__main__':
|
|
run()
|