def create_groups(group_ids, corpus, word, ignore_corpus, ignore_word):
"""Generate groups given the input PCollections."""
def attach_corpus_fn(group, corpus, ignore):
selected = None
len_corpus = len(corpus)
while not selected:
c = corpus[randrange(0, len_corpus - 1)].values()[0]
if c != ignore:
selected = c
yield (group, selected)
def attach_word_fn(group, words, ignore):
selected = None
len_words = len(words)
while not selected:
c = words[randrange(0, len_words - 1)].values()[0]
if c != ignore:
selected = c
yield group + (selected, )
return (group_ids
| df.FlatMap('attach corpus', attach_corpus_fn, AsList(corpus),
AsSingleton(ignore_corpus))
| df.FlatMap('attach word', attach_word_fn, AsIter(word),
AsSingleton(ignore_word)))
def test_pardo_using_flatmap(self):
words = ['aa', 'bbb', 'c']
# [START model_pardo_using_flatmap]
word_lengths = words | df.FlatMap(lambda word: [len(word)])
# [END model_pardo_using_flatmap]
self.assertEqual({2, 3, 1}, set(word_lengths))
def apply(self, pcoll):
return (pcoll
| (df.FlatMap('split', lambda x: re.findall(r'[A-Za-z\']+', x)
).with_output_types(unicode))
| df.Map('pair_with_one', lambda x: (x, 1))
| df.GroupByKey('group')
| df.Map('count', lambda (word, ones): (word, sum(ones))))
def run(argv=None):
"""Build and run the pipeline."""
parser = argparse.ArgumentParser()
parser.add_argument(
'--input_topic',
required=True,
help='Input PubSub topic of the form "/topics/<PROJECT>/<TOPIC>".')
parser.add_argument(
'--output_topic',
required=True,
help='Output PubSub topic of the form "/topics/<PROJECT>/<TOPIC>".')
known_args, pipeline_args = parser.parse_known_args(argv)
p = df.Pipeline(argv=pipeline_args)
# Read the text file[pattern] into a PCollection.
lines = p | df.io.Read('read', df.io.PubSubSource(known_args.input_topic))
# Capitalize the characters in each line.
transformed = (lines
| (df.FlatMap('split', lambda x: re.findall(
r'[A-Za-z\']+', x)).with_output_types(unicode))
| df.Map('pair_with_one', lambda x: (x, 1))
| df.WindowInto(window.FixedWindows(15, 0))
| df.GroupByKey('group')
| df.Map('count', lambda (word, ones): (word, sum(ones)))
| df.Map('format', lambda tup: '%s: %d' % tup))
# Write to PubSub.
# pylint: disable=expression-not-assigned
transformed | df.io.Write('pubsub_write',
df.io.PubSubSink(known_args.output_topic))
p.run()
def run(argv=None):
"""Main entry point; defines and runs the wordcount pipeline."""
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',
# CHANGE 1/5: The Google Cloud Storage path is required
# for outputting the results.
default='gs://YOUR_OUTPUT_BUCKET/AND_OUTPUT_PREFIX',
help='Output file to write results to.')
known_args, pipeline_args = parser.parse_known_args(argv)
pipeline_args.extend([
# CHANGE 2/5: (OPTIONAL) Change this to BlockingDataflowPipelineRunner to
# run your pipeline on the Google Cloud Dataflow Service.
'--runner=DirectPipelineRunner',
# CHANGE 3/5: Your project ID is required in order to run your pipeline on
# the Google Cloud Dataflow Service.
'--project=SET_YOUR_PROJECT_ID_HERE',
# CHANGE 4/5: Your Google Cloud Storage path is required for staging local
# files.
'--staging_location=gs://YOUR_BUCKET_NAME/AND_STAGING_DIRECTORY',
# CHANGE 5/5: Your Google Cloud Storage path is required for temporary
# files.
'--temp_location=gs://YOUR_BUCKET_NAME/AND_TEMP_DIRECTORY',
'--job_name=your-wordcount-job',
])
p = df.Pipeline(argv=pipeline_args)
# Read the text file[pattern] into a PCollection.
lines = p | df.io.Read('read', df.io.TextFileSource(known_args.input))
# Count the occurrences of each word.
counts = (lines
| (df.FlatMap('split', lambda x: re.findall(r'[A-Za-z\']+', x)).
with_output_types(unicode))
| df.Map('pair_with_one', lambda x: (x, 1))
| df.GroupByKey('group')
| df.Map('count', lambda (word, ones): (word, sum(ones))))
# Format the counts into a PCollection of strings.
output = counts | df.Map('format', lambda (word, c): '%s: %s' % (word, c))
# Write the output using a "Write" transform that has side effects.
# pylint: disable=expression-not-assigned
output | df.io.Write('write', df.io.TextFileSink(known_args.output))
# Actually run the pipeline (all operations above are deferred).
p.run()
def assert_that(pcoll, matcher):
"""Asserts that the give PCollection satisfies the constraints of the matcher
in a way that is runnable locally or on a remote service.
"""
singleton = pcoll.pipeline | df.Create('create_singleton', [None])
def check_matcher(_, side_value):
assert matcher(side_value)
return []
singleton | df.FlatMap(check_matcher, AsIter(pcoll))
def test_pardo_using_flatmap_yield(self):
words = ['aA', 'bbb', 'C']
# [START model_pardo_using_flatmap_yield]
def capitals(word):
for letter in word:
if 'A' <= letter <= 'Z':
yield letter
all_capitals = words | df.FlatMap(capitals)
# [END model_pardo_using_flatmap_yield]
self.assertEqual({'A', 'C'}, set(all_capitals))
def test_pardo_side_input(self):
p = df.Pipeline('DirectPipelineRunner')
words = p | df.Create('start', ['a', 'bb', 'ccc', 'dddd'])
# [START model_pardo_side_input]
# Callable takes additional arguments.
def filter_using_length(word, lower_bound, upper_bound=float('inf')):
if lower_bound <= len(word) <= upper_bound:
yield word
# Construct a deferred side input.
avg_word_len = words | df.Map(len) | df.CombineGlobally(
df.combiners.MeanCombineFn())
# Call with explicit side inputs.
small_words = words | df.FlatMap('small', filter_using_length, 0, 3)
# A single deferred side input.
larger_than_average = words | df.FlatMap(
'large',
filter_using_length,
lower_bound=pvalue.AsSingleton(avg_word_len))
# Mix and match.
small_but_nontrivial = words | df.FlatMap(
filter_using_length,
lower_bound=2,
upper_bound=pvalue.AsSingleton(avg_word_len))
# [END model_pardo_side_input]
df.assert_that(small_words, df.equal_to(['a', 'bb', 'ccc']))
df.assert_that(larger_than_average,
df.equal_to(['ccc', 'dddd']),
label='larger_than_average')
df.assert_that(small_but_nontrivial,
df.equal_to(['bb']),
label='small_but_not_trivial')
p.run()
def apply(self, words):
"""Compute the most common words for each possible prefixes.
Args:
words: a PCollection of strings
Returns:
A PCollection of most common words with each prefix, in the form
(prefix, [(count, word), (count, word), ...])
"""
return (words
| df.combiners.Count.PerElement()
| df.FlatMap(extract_prefixes)
| df.combiners.Top.LargestPerKey(self._count))
def apply(self, uri_to_content):
# Compute the total number of documents, and prepare a singleton
# PCollection to use as side input.
total_documents = (uri_to_content
| df.Keys('get uris')
| df.RemoveDuplicates('get unique uris')
| df.combiners.Count.Globally(' count uris'))
# Create a collection of pairs mapping a URI to each of the words
# in the document associated with that that URI.
def split_into_words((uri, line)):
return [(uri, w.lower()) for w in re.findall(r'[A-Za-z\']+', line)]
uri_to_words = (uri_to_content
| df.FlatMap('split words', split_into_words))
# Compute a mapping from each word to the total number of documents
# in which it appears.
word_to_doc_count = (
uri_to_words
| df.RemoveDuplicates('get unique words per doc')
| df.Values('get words')
| df.combiners.Count.PerElement('count docs per word'))
# Compute a mapping from each URI to the total number of words in the
# document associated with that URI.
uri_to_word_total = (
uri_to_words
| df.Keys(' get uris')
| df.combiners.Count.PerElement('count words in doc'))
# Count, for each (URI, word) pair, the number of occurrences of that word
# in the document associated with the URI.
uri_and_word_to_count = (
uri_to_words
| df.combiners.Count.PerElement('count word-doc pairs'))
# Adjust the above collection to a mapping from (URI, word) pairs to counts
# into an isomorphic mapping from URI to (word, count) pairs, to prepare
# for a join by the URI key.
uri_to_word_and_count = (uri_and_word_to_count
| df.Map(
'shift keys', lambda ((uri, word), count):
(uri, (word, count))))
def test_pardo_with_undeclared_side_outputs(self):
numbers = [1, 2, 3, 4, 5, 10, 20]
# [START model_pardo_with_side_outputs_undeclared]
def even_odd(x):
yield pvalue.SideOutputValue('odd' if x % 2 else 'even', x)
if x % 10 == 0:
yield x
results = numbers | df.FlatMap(even_odd).with_outputs()
evens = results.even
odds = results.odd
tens = results[None] # the undeclared main output
# [END model_pardo_with_side_outputs_undeclared]
self.assertEqual({2, 4, 10, 20}, set(evens))
self.assertEqual({1, 3, 5}, set(odds))
self.assertEqual({10, 20}, set(tens))
def run(argv=None):
"""Runs the workflow computing total points from a collection of matches."""
parser = argparse.ArgumentParser()
parser.add_argument('--input',
required=True,
help='Input file to process.')
parser.add_argument('--output',
required=True,
help='Output file to write results to.')
known_args, pipeline_args = parser.parse_known_args(argv)
p = df.Pipeline(argv=pipeline_args)
(p # pylint: disable=expression-not-assigned
| df.io.Read('read',
df.io.TextFileSource(known_args.input, coder=JsonCoder()))
| df.FlatMap('points', compute_points) | df.CombinePerKey(sum)
| df.io.Write('write',
df.io.TextFileSink(known_args.output, coder=JsonCoder())))
p.run()
def test_after_count(self):
p = Pipeline('DirectPipelineRunner')
result = (p
| df.Create([1, 2, 3, 4, 5, 10, 11])
| df.FlatMap(lambda t: [('A', t), ('B', t + 5)])
| df.Map(lambda (k, t): TimestampedValue((k, t), t))
| df.WindowInto(
FixedWindows(10),
trigger=AfterCount(3),
accumulation_mode=AccumulationMode.DISCARDING)
| df.GroupByKey()
| df.Map(lambda (k, v): ('%s-%s' % (k, len(v)), set(v))))
assert_that(
result,
equal_to({
'A-5': {1, 2, 3, 4, 5},
# A-10, A-11 never emitted due to AfterCount(3) never firing.
'B-4': {6, 7, 8, 9},
'B-3': {10, 15, 16},
}.iteritems()))
def run(argv=None):
parser = argparse.ArgumentParser()
parser.add_argument('--input',
required=True,
help='Input file to process.')
parser.add_argument('--output',
required=True,
help='Output file to write results to.')
known_args, pipeline_args = parser.parse_known_args(argv)
p = df.Pipeline(argv=pipeline_args)
(p # pylint: disable=expression-not-assigned
| df.io.Read('read', df.io.TextFileSource(known_args.input))
| df.FlatMap('split', lambda x: re.findall(r'[A-Za-z\']+', x))
| TopPerPrefix('TopPerPrefix', 5)
| df.Map('format', lambda (prefix, candidates): '%s: %s' %
(prefix, candidates))
| df.io.Write('write', df.io.TextFileSink(known_args.output)))
p.run()
def count_tornadoes(input_data):
"""Workflow computing the number of tornadoes for each month that had one.
Args:
input_data: a PCollection of dictionaries representing table rows. Each
dictionary will have a 'month' and a 'tornado' key as described in the
module comment.
Returns:
A PCollection of dictionaries containing 'month' and 'tornado_count' keys.
Months without tornadoes are skipped.
"""
return (input_data
| df.FlatMap(
'months with tornadoes', lambda row: [(int(row['month']), 1)]
if row['tornado'] else [])
| df.CombinePerKey('monthly count', sum)
| df.Map('format', lambda (k, v): {
'month': k,
'tornado_count': v
}))
| df.CoGroupByKey('cogroup by uri'))
# Compute a mapping from each word to a (URI, term frequency) pair for each
# URI. A word's term frequency for a document is simply the number of times
# that word occurs in the document divided by the total number of words in
# the document.
def compute_term_frequency((uri, count_and_total)):
word_and_count = count_and_total['word counts']
# We have an iterable for one element that we want extracted.
[word_total] = count_and_total['word totals']
for word, count in word_and_count:
yield word, (uri, float(count) / word_total)
word_to_uri_and_tf = (uri_to_word_and_count_and_total
| df.FlatMap('compute term frequencies',
compute_term_frequency))
# Compute a mapping from each word to its document frequency.
# A word's document frequency in a corpus is the number of
# documents in which the word appears divided by the total
# number of documents in the corpus.
#
# This calculation uses a side input, a Dataflow-computed auxiliary value
# presented to each invocation of our MapFn lambda. The second argument to
# the lambda (called total---note that we are unpacking the first argument)
# receives the value we listed after the lambda in Map(). Additional side
# inputs (and ordinary Python values, too) can be provided to MapFns and
# DoFns in this way.
word_to_df = (word_to_doc_count
| df.Map(
'compute doc frequencies', lambda