Use side inputs to implement non-compressing to singleton global combines.

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())

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]))