def test_multi_valued_singleton_side_input(self):
pipeline = Pipeline('DirectPipelineRunner')
pcol = pipeline | Create('start', [1, 2])
side = pipeline | Create('side', [3, 4]) # 2 values in side input.
pcol | FlatMap('compute', lambda x, s: [x * s], AsSingleton(side))
with self.assertRaises(ValueError) as e:
pipeline.run()
def test_default_value_singleton_side_input(self):
pipeline = Pipeline('DirectPipelineRunner')
pcol = pipeline | Create('start', [1, 2])
side = pipeline | Create('side', []) # 0 values in side input.
result = (pcol | FlatMap('compute', lambda x, s: [x * s],
AsSingleton(side, 10)))
assert_that(result, equal_to([10, 20]))
pipeline.run()
def test_iterable_side_input(self):
pipeline = Pipeline('DirectPipelineRunner')
pcol = pipeline | Create('start', [1, 2])
side = pipeline | Create('side', [3, 4]) # 2 values in side input.
result = pcol | FlatMap('compute', lambda x, s: [x * y for y in s],
AllOf(side))
assert_that(result, equal_to([3, 4, 6, 8]))
pipeline.run()
def test_create(self):
pipeline = Pipeline('DirectPipelineRunner')
pcoll = pipeline | Create('label1', [1, 2, 3])
assert_that(pcoll, equal_to([1, 2, 3]))
# Test if initial value is an iterator object.
pcoll2 = pipeline | Create('label2', iter((4, 5, 6)))
pcoll3 = pcoll2 | FlatMap('do', lambda x: [x + 10])
assert_that(pcoll3, equal_to([14, 15, 16]), label='pcoll3')
pipeline.run()
def test_reuse_cloned_custom_transform_instance(self):
pipeline = Pipeline(DirectPipelineRunner())
pcoll1 = pipeline | Create('pcoll1', [1, 2, 3])
pcoll2 = pipeline | Create('pcoll2', [4, 5, 6])
transform = PipelineTest.CustomTransform()
result1 = pcoll1 | transform
result2 = pcoll2 | transform.clone('new label')
assert_that(result1, equal_to([2, 3, 4]), label='r1')
assert_that(result2, equal_to([5, 6, 7]), label='r2')
pipeline.run()
def test_par_do_with_side_input_as_arg(self):
pipeline = Pipeline('DirectPipelineRunner')
words_list = ['aa', 'bb', 'cc']
words = pipeline | Create('SomeWords', words_list)
prefix = pipeline | Create('SomeString', ['xyz']) # side in
suffix = 'zyx'
result = words | FlatMap(
'DecorateWords', lambda x, pfx, sfx: ['%s-%s-%s' % (pfx, x, sfx)],
AsSingleton(prefix), suffix)
assert_that(result, equal_to(['xyz-%s-zyx' % x for x in words_list]))
pipeline.run()
def test_pcollectionview_not_recreated(self):
pipeline = Pipeline('DirectPipelineRunner')
value = pipeline | Create('create1', [1, 2, 3])
value2 = pipeline | Create('create2', [(1, 1), (2, 2), (3, 3)])
self.assertEqual(AsSingleton(value), AsSingleton(value))
self.assertEqual(AsSingleton('new', value, default_value=1),
AsSingleton('new', value, default_value=1))
self.assertNotEqual(AsSingleton(value),
AsSingleton('new', value, default_value=1))
self.assertEqual(AsIter(value), AsIter(value))
self.assertEqual(AsList(value), AsList(value))
self.assertEqual(AsDict(value2), AsDict(value2))
def test_empty_singleton_side_input(self):
pipeline = Pipeline('DirectPipelineRunner')
pcol = pipeline | Create('start', [1, 2])
side = pipeline | Create('side', []) # Empty side input.
def my_fn(k, s):
v = ('empty' if isinstance(s, EmptySideInput) else 'full')
return [(k, v)]
result = pcol | FlatMap('compute', my_fn, AsSingleton(side))
assert_that(result, equal_to([(1, 'empty'), (2, 'empty')]))
pipeline.run()
def test_reuse_custom_transform_instance(self):
pipeline = Pipeline(DirectPipelineRunner())
pcoll1 = pipeline | Create('pcoll1', [1, 2, 3])
pcoll2 = pipeline | Create('pcoll2', [4, 5, 6])
transform = PipelineTest.CustomTransform()
pcoll1 | transform
with self.assertRaises(RuntimeError) as cm:
pipeline.apply(transform, pcoll2)
self.assertEqual(
cm.exception.message,
'Transform "CustomTransform" does not have a stable unique label. '
'This will prevent updating of pipelines. '
'To clone a transform with a new label use: '
'transform.clone("NEW LABEL").')
def test_as_singleton_with_different_defaults_without_unique_labels(self):
# This should fail as AsSingleton with distinct default values should create
# distinct PCollectionViews with the same full_label.
a_list = [2]
pipeline = Pipeline('DirectPipelineRunner')
main_input = pipeline | Create('main input', [1])
side_list = pipeline | Create('side list', a_list)
with self.assertRaises(RuntimeError) as e:
_ = main_input | FlatMap('test', lambda x, s1, s2: [[x, s1, s2]],
AsSingleton(side_list),
AsSingleton(side_list, default_value=3))
self.assertTrue(
e.exception.message.startswith(
'Transform "ViewAsSingleton(side list.None)" does not have a '
'stable unique label.'))
def test_timestamped_with_combiners(self):
p = Pipeline('DirectPipelineRunner')
result = (p
# Create some initial test values.
| Create('start', [(k, k) for k in range(10)])
# The purpose of the WindowInto transform is to establish a
# FixedWindows windowing function for the PCollection.
# It does not bucket elements into windows since the timestamps
# from Create are not spaced 5 ms apart and very likely they all
# fall into the same window.
| WindowInto('w', FixedWindows(5))
# Generate timestamped values using the values as timestamps.
# Now there are values 5 ms apart and since Map propagates the
# windowing function from input to output the output PCollection
# will have elements falling into different 5ms windows.
| Map(lambda (x, t): TimestampedValue(x, t))
# We add a 'key' to each value representing the index of the
# window. This is important since there is no guarantee of
# order for the elements of a PCollection.
| Map(lambda v: (v / 5, v)))
# Sum all elements associated with a key and window. Although it
# is called CombinePerKey it is really CombinePerKeyAndWindow the
# same way GroupByKey is really GroupByKeyAndWindow.
sum_per_window = result | CombinePerKey(sum)
# Compute mean per key and window.
mean_per_window = result | combiners.Mean.PerKey()
assert_that(sum_per_window, equal_to([(0, 10), (1, 35)]),
label='assert:sum')
assert_that(mean_per_window, equal_to([(0, 2.0), (1, 7.0)]),
label='assert:mean')
p.run()
def test_map(self):
pipeline = Pipeline('DirectPipelineRunner')
lines = pipeline | Create('input', ['a', 'b', 'c'])
result = (lines
| Map('upper', str.upper)
| Map('prefix', lambda x, prefix: prefix + x, 'foo-'))
assert_that(result, equal_to(['foo-A', 'foo-B', 'foo-C']))
pipeline.run()
def test_pipeline_as_context(self):
def raise_exception(exn):
raise exn
with self.assertRaises(ValueError):
with Pipeline(self.runner_name) as p:
# pylint: disable=expression-not-assigned
p | Create([ValueError]) | Map(raise_exception)
def test_word_count_using_get(self):
pipeline = Pipeline('DirectPipelineRunner')
lines = pipeline | Create('SomeWords', [DataflowTest.SAMPLE_DATA])
result = ((lines | FlatMap('GetWords',
lambda x: re.findall(r'\w+', x))).apply(
'CountWords', DataflowTest.Count))
assert_that(result, equal_to(DataflowTest.SAMPLE_RESULT))
pipeline.run()
def test_par_do_with_do_fn_object(self):
class SomeDoFn(DoFn):
"""A custom DoFn for a FlatMap transform."""
def process(self, context, prefix, suffix):
return ['%s-%s-%s' % (prefix, context.element, suffix)]
pipeline = Pipeline('DirectPipelineRunner')
words_list = ['aa', 'bb', 'cc']
words = pipeline | Create('SomeWords', words_list)
prefix = 'zyx'
suffix = pipeline | Create('SomeString', ['xyz']) # side in
result = words | ParDo('DecorateWordsDoFn',
SomeDoFn(),
prefix,
suffix=AsSingleton(suffix))
assert_that(result, equal_to(['zyx-%s-xyz' % x for x in words_list]))
pipeline.run()
def test_cached_pvalues_are_refcounted(self):
"""Test that cached PValues are refcounted and deleted.
The intermediary PValues computed by the workflow below contain
one million elements so if the refcounting does not work the number of
objects tracked by the garbage collector will increase by a few millions
by the time we execute the final Map checking the objects tracked.
Anything that is much larger than what we started with will fail the test.
"""
def check_memory(value, count_threshold):
gc.collect()
objects_count = len(gc.get_objects())
if objects_count > count_threshold:
raise RuntimeError(
'PValues are not refcounted: %s, %s' % (
objects_count, count_threshold))
return value
def create_dupes(o, _):
yield o
yield SideOutputValue('side', o)
pipeline = Pipeline('DirectPipelineRunner')
gc.collect()
count_threshold = len(gc.get_objects()) + 10000
biglist = pipeline | Create('oom:create', ['x'] * 1000000)
dupes = (
biglist
| Map('oom:addone', lambda x: (x, 1))
| FlatMap('oom:dupes', create_dupes,
AsIter(biglist)).with_outputs('side', main='main'))
result = (
(dupes.side, dupes.main, dupes.side)
| Flatten('oom:flatten')
| CombinePerKey('oom:combine', sum)
| Map('oom:check', check_memory, count_threshold))
assert_that(result, equal_to([('x', 3000000)]))
pipeline.run()
self.assertEqual(
pipeline.runner.debug_counters['element_counts'],
{
'oom:flatten': 3000000,
('oom:combine/GroupByKey/reify_windows', None): 3000000,
('oom:dupes/oom:dupes', 'side'): 1000000,
('oom:dupes/oom:dupes', None): 1000000,
'oom:create': 1000000,
('oom:addone', None): 1000000,
'oom:combine/GroupByKey/group_by_key': 1,
('oom:check', None): 1,
'assert_that/singleton': 1,
('assert_that/Map(match)', None): 1,
('oom:combine/GroupByKey/group_by_window', None): 1,
('oom:combine/Combine/ParDo(CombineValuesDoFn)', None): 1})
def test_timestamped_value(self):
p = Pipeline('DirectPipelineRunner')
result = (p
| Create('start', [(k, k) for k in range(10)])
| Map(lambda (x, t): TimestampedValue(x, t))
| WindowInto('w', FixedWindows(5))
| Map(lambda v: ('key', v))
| GroupByKey())
assert_that(result, equal_to([('key', [0, 1, 2, 3, 4]),
('key', [5, 6, 7, 8, 9])]))
p.run()
def test_as_dict_with_unique_labels(self):
some_kvs = [('a', 1), ('b', 2)]
pipeline = Pipeline('DirectPipelineRunner')
main_input = pipeline | Create('main input', [1])
side_kvs = pipeline | Create('side kvs', some_kvs)
results = main_input | FlatMap(
'test', lambda x, dct1, dct2: [[x, dct1, dct2]], AsDict(side_kvs),
AsDict(side_kvs, label='label'))
def matcher(expected_elem, expected_kvs):
def match(actual):
[[actual_elem, actual_dict1, actual_dict2]] = actual
equal_to([expected_elem])([actual_elem])
equal_to(expected_kvs)(actual_dict1.iteritems())
equal_to(expected_kvs)(actual_dict2.iteritems())
return match
assert_that(results, matcher(1, some_kvs))
pipeline.run()
def test_as_list_with_unique_labels(self):
a_list = [1, 2, 3]
pipeline = Pipeline('DirectPipelineRunner')
main_input = pipeline | Create('main input', [1])
side_list = pipeline | Create('side list', a_list)
results = main_input | FlatMap(
'test', lambda x, ls1, ls2: [[x, ls1, ls2]], AsList(side_list),
AsList(side_list, label='label'))
def matcher(expected_elem, expected_list):
def match(actual):
[[actual_elem, actual_list1, actual_list2]] = actual
equal_to([expected_elem])([actual_elem])
equal_to(expected_list)(actual_list1)
equal_to(expected_list)(actual_list2)
return match
assert_that(results, matcher(1, [1, 2, 3]))
pipeline.run()
def test_visit_node_sub_graph(self):
pipeline = Pipeline('DirectPipelineRunner')
pcoll1 = pipeline | Create('pcoll', [1, 2, 3])
pcoll2 = pcoll1 | FlatMap('do1', lambda x: [x + 1])
pcoll3 = pcoll2 | FlatMap('do2', lambda x: [x + 1])
pcoll4 = pcoll2 | FlatMap('do3', lambda x: [x + 1])
visitor = PipelineTest.Visitor(visited=[])
pipeline.visit(visitor, node=pcoll3)
self.assertFalse(pcoll4 in visitor.visited)
self.assertEqual(set([pcoll1, pcoll2, pcoll3]), set(visitor.visited))
def test_empty_side_outputs(self):
pipeline = Pipeline('DirectPipelineRunner')
nums = pipeline | Create('Some Numbers', [1, 3, 5])
results = nums | FlatMap(
'ClassifyNumbers', lambda x:
[x, SideOutputValue('even'
if x % 2 == 0 else 'odd', x)]).with_outputs()
assert_that(results[None], equal_to([1, 3, 5]))
assert_that(results.odd, equal_to([1, 3, 5]), label='assert:odd')
assert_that(results.even, equal_to([]), label='assert:even')
pipeline.run()
def test_as_singleton_with_different_defaults_with_unique_labels(self):
a_list = []
pipeline = Pipeline('DirectPipelineRunner')
main_input = pipeline | Create('main input', [1])
side_list = pipeline | Create('side list', a_list)
results = main_input | FlatMap(
'test', lambda x, s1, s2: [[x, s1, s2]],
AsSingleton('si1', side_list, default_value=2),
AsSingleton('si2', side_list, default_value=3))
def matcher(expected_elem, expected_singleton1, expected_singleton2):
def match(actual):
[[actual_elem, actual_singleton1, actual_singleton2]] = actual
equal_to([expected_elem])([actual_elem])
equal_to([expected_singleton1])([actual_singleton1])
equal_to([expected_singleton2])([actual_singleton2])
return match
assert_that(results, matcher(1, 2, 3))
pipeline.run()
def test_as_list_without_unique_labels(self):
# This should succeed as calling AsList on the same PCollection twice will
# return the same PCollectionView.
a_list = [1, 2, 3]
pipeline = Pipeline('DirectPipelineRunner')
main_input = pipeline | Create('main input', [1])
side_list = pipeline | Create('side list', a_list)
results = main_input | FlatMap(
'test', lambda x, ls1, ls2: [[x, ls1, ls2]], AsList(side_list),
AsList(side_list))
def matcher(expected_elem, expected_list):
def match(actual):
[[actual_elem, actual_list1, actual_list2]] = actual
equal_to([expected_elem])([actual_elem])
equal_to(expected_list)(actual_list1)
equal_to(expected_list)(actual_list2)
return match
assert_that(results, matcher(1, [1, 2, 3]))
pipeline.run()
def test_as_singleton_without_unique_labels(self):
# This should succeed as calling AsSingleton on the same PCollection twice
# with the same defaults will return the same PCollectionView.
a_list = [2]
pipeline = Pipeline('DirectPipelineRunner')
main_input = pipeline | Create('main input', [1])
side_list = pipeline | Create('side list', a_list)
results = main_input | FlatMap('test', lambda x, s1, s2: [[x, s1, s2]],
AsSingleton(side_list),
AsSingleton(side_list))
def matcher(expected_elem, expected_singleton):
def match(actual):
[[actual_elem, actual_singleton1, actual_singleton2]] = actual
equal_to([expected_elem])([actual_elem])
equal_to([expected_singleton])([actual_singleton1])
equal_to([expected_singleton])([actual_singleton2])
return match
assert_that(results, matcher(1, 2))
pipeline.run()
def test_undeclared_side_outputs(self):
pipeline = Pipeline('DirectPipelineRunner')
nums = pipeline | Create('Some Numbers', [1, 2, 3, 4])
results = nums | FlatMap(
'ClassifyNumbers', lambda x:
[x, SideOutputValue('even'
if x % 2 == 0 else 'odd', x)]).with_outputs()
# TODO(silviuc): Revisit this test to check for undeclared side outputs.
# This should work with .with_outputs() without any tags declared and
# the results[None] should work also.
assert_that(results[None], equal_to([1, 2, 3, 4]))
assert_that(results.odd, equal_to([1, 3]), label='assert:odd')
assert_that(results.even, equal_to([2, 4]), label='assert:even')
pipeline.run()
def test_as_list_and_as_dict_side_inputs(self):
a_list = [5, 1, 3, 2, 9]
some_pairs = [('crouton', 17), ('supreme', None)]
pipeline = Pipeline('DirectPipelineRunner')
main_input = pipeline | Create('main input', [1])
side_list = pipeline | Create('side list', a_list)
side_pairs = pipeline | Create('side pairs', some_pairs)
results = main_input | FlatMap(
'concatenate',
lambda x, the_list, the_dict: [[x, the_list, the_dict]],
AsList(side_list), AsDict(side_pairs))
def matcher(expected_elem, expected_list, expected_pairs):
def match(actual):
[[actual_elem, actual_list, actual_dict]] = actual
equal_to([expected_elem])([actual_elem])
equal_to(expected_list)(actual_list)
equal_to(expected_pairs)(actual_dict.iteritems())
return match
assert_that(results, matcher(1, a_list, some_pairs))
pipeline.run()
def test_par_do_with_multiple_outputs_and_using_return(self):
def some_fn(v):
if v % 2 == 0:
return [v, SideOutputValue('even', v)]
else:
return [v, SideOutputValue('odd', v)]
pipeline = Pipeline('DirectPipelineRunner')
nums = pipeline | Create('Some Numbers', [1, 2, 3, 4])
results = nums | FlatMap('ClassifyNumbers', some_fn).with_outputs(
'odd', 'even', main='main')
assert_that(results.main, equal_to([1, 2, 3, 4]))
assert_that(results.odd, equal_to([1, 3]), label='assert:odd')
assert_that(results.even, equal_to([2, 4]), label='assert:even')
pipeline.run()
def test_visit_entire_graph(self):
pipeline = Pipeline(self.runner_name)
pcoll1 = pipeline | Create('pcoll', [1, 2, 3])
pcoll2 = pcoll1 | FlatMap('do1', lambda x: [x + 1])
pcoll3 = pcoll2 | FlatMap('do2', lambda x: [x + 1])
pcoll4 = pcoll2 | FlatMap('do3', lambda x: [x + 1])
transform = PipelineTest.CustomTransform()
pcoll5 = pcoll4 | transform
visitor = PipelineTest.Visitor(visited=[])
pipeline.visit(visitor)
self.assertEqual(set([pcoll1, pcoll2, pcoll3, pcoll4, pcoll5]),
set(visitor.visited))
self.assertEqual(set(visitor.enter_composite),
set(visitor.leave_composite))
self.assertEqual(2, len(visitor.enter_composite))
self.assertEqual(visitor.enter_composite[1].transform, transform)
self.assertEqual(visitor.leave_composite[0].transform, transform)
def test_par_do_with_multiple_outputs_and_using_yield(self):
class SomeDoFn(DoFn):
"""A custom DoFn using yield."""
def process(self, context):
yield context.element
if context.element % 2 == 0:
yield SideOutputValue('even', context.element)
else:
yield SideOutputValue('odd', context.element)
pipeline = Pipeline('DirectPipelineRunner')
nums = pipeline | Create('Some Numbers', [1, 2, 3, 4])
results = nums | ParDo('ClassifyNumbers', SomeDoFn()).with_outputs(
'odd', 'even', main='main')
assert_that(results.main, equal_to([1, 2, 3, 4]))
assert_that(results.odd, equal_to([1, 3]), label='assert:odd')
assert_that(results.even, equal_to([2, 4]), label='assert:even')
pipeline.run()
def test_window_transform(self):
class TestWindowFn(WindowFn):
"""Windowing function adding two disjoint windows to each element."""
def assign(self, assign_context):
_ = assign_context
return [IntervalWindow(10, 20), IntervalWindow(20, 30)]
def merge(self, existing_windows):
return existing_windows
pipeline = Pipeline('DirectPipelineRunner')
numbers = pipeline | Create('KVs', [(1, 10), (2, 20), (3, 30)])
result = (numbers
| WindowInto('W', windowfn=TestWindowFn())
| GroupByKey('G'))
assert_that(
result,
equal_to([(1, [10]), (1, [10]), (2, [20]), (2, [20]), (3, [30]),
(3, [30])]))
pipeline.run()