diff --git a/sdks/python/apache_beam/transforms/combiners.py b/sdks/python/apache_beam/transforms/combiners.py index c45ba4e89b9a..ddba28569d76 100644 --- a/sdks/python/apache_beam/transforms/combiners.py +++ b/sdks/python/apache_beam/transforms/combiners.py @@ -31,6 +31,7 @@ import numpy as np +from apache_beam import pvalue from apache_beam import typehints from apache_beam.transforms import core from apache_beam.transforms import cy_combiners @@ -49,6 +50,7 @@ 'Sample', 'Top', 'ToDict', + 'ToIterable', 'ToList', 'ToSet', 'Latest', @@ -787,16 +789,71 @@ def add_input(self, accumulator, element, *args, **kwargs): ] +class ToSingleElementCollection(CombinerWithoutDefaults): + """Base class for combining a PCollection into a single element. + + This is for combining a PCollection into a single element in the cases where + there is no value (and significant overhead) in using a combiner as the + output is no smaller than its input. + """ + def expand(self, pcoll): + raise NotImplementedError + + def _expand(self, pcoll, post_process): + if pcoll.windowing.is_default(): + if self.has_defaults: + return ( + pcoll.pipeline + | 'DoOnce' >> core.Create([None]) + | core.Map(lambda _, s: post_process(s), pvalue.AsIter(pcoll))) + else: + + def _if_non_empty(unused_elem, s): + for _ in s: + yield post_process(s) + break + + return ( + pcoll.pipeline + | 'DoOnce' >> core.Create([None]) + | core.FlatMap(_if_non_empty, pvalue.AsIter(pcoll))) + else: + if self.has_defaults: + raise ValueError( + 'Default values are not yet supported in %s() if the ' + 'output PCollection is not windowed by GlobalWindows. ' + 'Instead, use %s().without_defaults() to output ' + 'an empty PCollection if the input PCollection is empty.' % + (self.__class__.__name__, self.__class__.__name__)) + return ( + pcoll + | 'KeyWithVoid' >> core.ParDo(core._KeyWithNone()) + | core.GroupByKey() + | core.MapTuple(lambda _, vs: post_process(vs))) + + +@with_input_types(T) +@with_output_types(Iterable[T]) +class ToIterable(ToSingleElementCollection): + """A global Combine that condenses a PCollection into an iterable. + + Note that this forces the entire PCollection to be processed on a single + machine. + """ + def expand(self, pcoll): + return self._expand(pcoll, lambda x: x) + + @with_input_types(T) @with_output_types(list[T]) -class ToList(CombinerWithoutDefaults): - """A global CombineFn that condenses a PCollection into a single list.""" +class ToList(ToSingleElementCollection): + """A global Combine that condenses a PCollection into a single list. + + Note that this forces the entire PCollection to be processed on a single + machine. + """ def expand(self, pcoll): - if self.has_defaults: - return pcoll | self.label >> core.CombineGlobally(ToListCombineFn()) - else: - return pcoll | self.label >> core.CombineGlobally( - ToListCombineFn()).without_defaults() + return self._expand(pcoll, list) @with_input_types(T) @@ -836,19 +893,18 @@ def extract_output(self, accumulator): @with_input_types(tuple[K, V]) @with_output_types(dict[K, V]) -class ToDict(CombinerWithoutDefaults): - """A global CombineFn that condenses a PCollection into a single dict. +class ToDict(ToSingleElementCollection): + """A global Combine that condenses a PCollection into a single dict. PCollections should consist of 2-tuples, notionally (key, value) pairs. If multiple values are associated with the same key, only one of the values will be present in the resulting dict. + + Note that this forces the entire PCollection to be processed on a single + machine. """ def expand(self, pcoll): - if self.has_defaults: - return pcoll | self.label >> core.CombineGlobally(ToDictCombineFn()) - else: - return pcoll | self.label >> core.CombineGlobally( - ToDictCombineFn()).without_defaults() + return self._expand(pcoll, dict) @with_input_types(tuple[K, V]) @@ -876,7 +932,11 @@ def extract_output(self, accumulator): @with_input_types(T) @with_output_types(set[T]) class ToSet(CombinerWithoutDefaults): - """A global CombineFn that condenses a PCollection into a set.""" + """A global Combine that condenses a PCollection into a set. + + Note that this forces the entire PCollection to be processed on a single + machine. + """ def expand(self, pcoll): if self.has_defaults: return pcoll | self.label >> core.CombineGlobally(ToSetCombineFn()) diff --git a/sdks/python/apache_beam/transforms/combiners_test.py b/sdks/python/apache_beam/transforms/combiners_test.py index 14348bb8ce78..1dc927c49fa2 100644 --- a/sdks/python/apache_beam/transforms/combiners_test.py +++ b/sdks/python/apache_beam/transforms/combiners_test.py @@ -549,6 +549,144 @@ def match(actual): assert_that( result_windowed, matcher(set(the_list)), label='to-set-wo-defaults') + def test_to_iterable_to_list_to_dict_globally_windowed(self): + elements = [3, 1, 2] + pairs = [('a', 1), ('b', 2)] + + with TestPipeline() as p: + pcoll = p | 'CreateElems' >> Create(elements) + pcoll_pairs = p | 'CreatePairs' >> Create(pairs) + + # With defaults + res_iter = pcoll | 'ToIter' >> combine.ToIterable() + assert_that( + res_iter | 'SortIter' >> Map(sorted), + equal_to([[1, 2, 3]]), + label='IterDefault') + + res_list = pcoll | 'ToList' >> combine.ToList() + assert_that( + res_list | 'SortList' >> Map(sorted), + equal_to([[1, 2, 3]]), + label='ListDefault') + + res_dict = pcoll_pairs | 'ToDict' >> combine.ToDict() + assert_that(res_dict, equal_to([{'a': 1, 'b': 2}]), label='DictDefault') + + # Without defaults + res_iter_wo = ( + pcoll | 'ToIterWo' >> combine.ToIterable().without_defaults()) + assert_that( + res_iter_wo | 'SortIterWo' >> Map(sorted), + equal_to([[1, 2, 3]]), + label='IterWoDefault') + + res_list_wo = pcoll | 'ToListWo' >> combine.ToList().without_defaults() + assert_that( + res_list_wo | 'SortListWo' >> Map(sorted), + equal_to([[1, 2, 3]]), + label='ListWoDefault') + + res_dict_wo = ( + pcoll_pairs | 'ToDictWo' >> combine.ToDict().without_defaults()) + assert_that( + res_dict_wo, equal_to([{ + 'a': 1, 'b': 2 + }]), label='DictWoDefault') + + def test_to_iterable_to_list_to_dict_empty(self): + with TestPipeline() as p: + empty = p | 'CreateEmpty' >> Create([]) + empty_pairs = p | 'CreateEmptyPairs' >> Create([]) + + # Empty with defaults: [[]], [{}], [[]] + res_iter = empty | 'ToIter' >> combine.ToIterable() + assert_that( + res_iter | 'MapIter' >> Map(list), + equal_to([[]]), + label='EmptyIterDefault') + + res_list = empty | 'ToList' >> combine.ToList() + assert_that(res_list, equal_to([[]]), label='EmptyListDefault') + + res_dict = empty_pairs | 'ToDict' >> combine.ToDict() + assert_that(res_dict, equal_to([{}]), label='EmptyDictDefault') + + # Empty without defaults: [] + res_iter_wo = ( + empty | 'ToIterWo' >> combine.ToIterable().without_defaults()) + assert_that( + res_iter_wo | 'MapIterWo' >> Map(list), + equal_to([]), + label='EmptyIterWo') + + res_list_wo = empty | 'ToListWo' >> combine.ToList().without_defaults() + assert_that(res_list_wo, equal_to([]), label='EmptyListWo') + + res_dict_wo = ( + empty_pairs | 'ToDictWo' >> combine.ToDict().without_defaults()) + assert_that(res_dict_wo, equal_to([]), label='EmptyDictWo') + + def test_to_iterable_to_list_to_dict_windowed(self): + # Non-globally windowed with defaults raises ValueError + p = TestPipeline() + windowed = ( + p + | 'Create' >> Create([1, 2]) + | 'Window' >> WindowInto(FixedWindows(10))) + windowed_pairs = ( + p + | 'CreatePairs' >> Create([('a', 1), ('b', 2)]) + | 'WindowPairs' >> WindowInto(FixedWindows(10))) + with self.assertRaises(ValueError): + _ = windowed | combine.ToIterable() + with self.assertRaises(ValueError): + _ = windowed | combine.ToList() + with self.assertRaises(ValueError): + _ = windowed_pairs | combine.ToDict() + + # Non-globally windowed without defaults groups per window + with TestPipeline() as p: + timestamped_elems = ( + p + | 'CreateTimed' >> Create([(1, 0), (2, 5), (3, 10), (4, 15)]) + | 'AssignTimestamps' >> + Map(lambda x_ts: TimestampedValue(x_ts[0], x_ts[1]) + ).with_output_types(int) + | 'WindowTimed' >> WindowInto(FixedWindows(10))) + + timestamped_pairs = ( + p + | 'CreateTimedPairs' >> Create([(('a', 1), 0), (('b', 2), 5), + (('c', 3), 10), (('d', 4), 15)]) + | 'AssignTimestampsPairs' >> + Map(lambda p_ts: TimestampedValue(p_ts[0], p_ts[1]) + ).with_output_types(tuple[str, int]) + | 'WindowTimedPairs' >> WindowInto(FixedWindows(10))) + + res_iter = ( + timestamped_elems + | 'ToIter' >> combine.ToIterable().without_defaults() + | 'SortIter' >> Map(sorted)) + assert_that(res_iter, equal_to([[1, 2], [3, 4]]), label='WindowedIter') + + res_list = ( + timestamped_elems + | 'ToList' >> combine.ToList().without_defaults() + | 'SortList' >> Map(sorted)) + assert_that(res_list, equal_to([[1, 2], [3, 4]]), label='WindowedList') + + res_dict = ( + timestamped_pairs | 'ToDict' >> combine.ToDict().without_defaults()) + assert_that( + res_dict, + equal_to([{ + 'a': 1, 'b': 2 + }, { + 'c': 3, 'd': 4 + }]), + label='WindowedDict') + def test_combine_globally_with_default(self): with TestPipeline() as p: assert_that(p | Create([]) | CombineGlobally(sum), equal_to([0]))