Windowing Support for the Dask Runner

.github/workflows/dask_runner_tests.yml

@@ -65,9 +65,8 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        os: [ubuntu-latest, macos-latest, windows-latest]
+        os: [ubuntu-latest, macos-latest]
         params: [
-          {"py_ver": "3.7", "tox_env": "py37"},
           {"py_ver": "3.8", "tox_env": "py38"},
           {"py_ver": "3.9", "tox_env": "py39"},
           {"py_ver": "3.10", "tox_env": "py310" },
@@ -86,15 +85,11 @@ jobs:
         run: pip install tox
       - name: Install SDK with dask
         working-directory: ./sdks/python
-        run: pip install setuptools --upgrade && pip install -e .[gcp,dask,test]
+        run: pip install setuptools --upgrade && pip install -e .[dask,test,dataframes]
       - name: Run tests basic unix
         if: startsWith(matrix.os, 'ubuntu') || startsWith(matrix.os, 'macos')
         working-directory: ./sdks/python
         run: tox -c tox.ini -e ${{ matrix.params.tox_env }}-dask
-      - name: Run tests basic windows
-        if: startsWith(matrix.os, 'windows')
-        working-directory: ./sdks/python
-        run: tox -c tox.ini -e ${{ matrix.params.tox_env }}-win-dask
       - name: Upload test logs
         uses: actions/upload-artifact@v3
         if: always()

sdks/python/apache_beam/runners/dask/dask_runner.py

@@ -159,6 +159,8 @@ def is_fnapi_compatible():
     return False
 
   def run_pipeline(self, pipeline, options):
+    import dask
+
     # TODO(alxr): Create interactive notebook support.
     if is_in_notebook():
       raise NotImplementedError('interactive support will come later!')
@@ -178,5 +180,6 @@ def run_pipeline(self, pipeline, options):
     dask_visitor = self.to_dask_bag_visitor()
     pipeline.visit(dask_visitor)
 
-    futures = client.compute(list(dask_visitor.bags.values()))
+    opt_graph = dask.optimize(dask_visitor.bags.values())
+    futures = client.compute(opt_graph)
     return DaskRunnerResult(client, futures)

sdks/python/apache_beam/runners/dask/dask_runner_test.py

@@ -24,10 +24,11 @@
 from apache_beam.testing.util import equal_to
 
 try:
-  from apache_beam.runners.dask.dask_runner import DaskOptions
-  from apache_beam.runners.dask.dask_runner import DaskRunner
   import dask
   import dask.distributed as ddist
+
+  from apache_beam.runners.dask.dask_runner import DaskOptions  # pylint: disable=ungrouped-imports
+  from apache_beam.runners.dask.dask_runner import DaskRunner  # pylint: disable=ungrouped-imports
 except (ImportError, ModuleNotFoundError):
   raise unittest.SkipTest('Dask must be installed to run tests.')
 
@@ -73,6 +74,11 @@ def test_create(self):
       pcoll = p | beam.Create([1])
       assert_that(pcoll, equal_to([1]))
 
+  def test_create_multiple(self):
+    with self.pipeline as p:
+      pcoll = p | beam.Create([1, 2, 3])
+      assert_that(pcoll, equal_to([1, 2, 3]))
+
   def test_create_and_map(self):
     def double(x):
       return x * 2
@@ -81,6 +87,22 @@ def double(x):
       pcoll = p | beam.Create([1]) | beam.Map(double)
       assert_that(pcoll, equal_to([2]))
 
+  def test_create_and_map_multiple(self):
+    def double(x):
+      return x * 2
+
+    with self.pipeline as p:
+      pcoll = p | beam.Create([1, 2]) | beam.Map(double)
+      assert_that(pcoll, equal_to([2, 4]))
+
+  def test_create_and_map_many(self):
+    def double(x):
+      return x * 2
+
+    with self.pipeline as p:
+      pcoll = p | beam.Create(list(range(1, 11))) | beam.Map(double)
+      assert_that(pcoll, equal_to(list(range(2, 21, 2))))
+
   def test_create_map_and_groupby(self):
     def double(x):
       return x * 2, x
@@ -89,6 +111,31 @@ def double(x):
       pcoll = p | beam.Create([1]) | beam.Map(double) | beam.GroupByKey()
       assert_that(pcoll, equal_to([(2, [1])]))
 
+  def test_create_map_and_groupby_multiple(self):
+    def double(x):
+      return x * 2, x
+
+    with self.pipeline as p:
+      pcoll = p | beam.Create([1, 2, 1, 2, 3
+                               ]) | beam.Map(double) | beam.GroupByKey()
+      assert_that(pcoll, equal_to([(2, [1, 1]), (4, [2, 2]), (6, [3])]))
+
+  def test_map_with_side_inputs(self):
+    def mult_by(x, y):
+      return x * y
+
+    with self.pipeline as p:
+      pcoll = p | beam.Create([1]) | beam.Map(mult_by, 3)
+      assert_that(pcoll, equal_to([3]))
+
+  def test_map_with_named_side_inputs(self):
+    def mult_by(x, y):
+      return x * y
+
+    with self.pipeline as p:
+      pcoll = p | beam.Create([1]) | beam.Map(mult_by, y=3)
+      assert_that(pcoll, equal_to([3]))
+
 
 if __name__ == '__main__':
   unittest.main()

sdks/python/apache_beam/runners/dask/overrides.py

@@ -73,7 +73,6 @@ def infer_output_type(self, input_type):
 @typehints.with_input_types(t.Tuple[K, t.Iterable[V]])
 @typehints.with_output_types(t.Tuple[K, t.Iterable[V]])
 class _GroupAlsoByWindow(beam.ParDo):
-  """Not used yet..."""
   def __init__(self, windowing):
     super().__init__(_GroupAlsoByWindowDoFn(windowing))
     self.windowing = windowing
@@ -86,7 +85,11 @@ def expand(self, input_or_inputs):
 @typehints.with_output_types(t.Tuple[K, t.Iterable[V]])
 class _GroupByKey(beam.PTransform):
   def expand(self, input_or_inputs):
-    return input_or_inputs | "GroupByKey" >> _GroupByKeyOnly()
+    return (
+        input_or_inputs
+        | "ReifyWindows" >> beam.ParDo(beam.GroupByKey.ReifyWindows())
+        | "GroupByKey" >> _GroupByKeyOnly()
+        | "GroupByWindow" >> _GroupAlsoByWindow(input_or_inputs.windowing))
 
 
 class _Flatten(beam.PTransform):

sdks/python/apache_beam/runners/dask/transform_evaluator.py

@@ -25,15 +25,73 @@
 import abc
 import dataclasses
 import typing as t
+from dataclasses import field
 
 import apache_beam
 import dask.bag as db
+from apache_beam import DoFn
+from apache_beam import TaggedOutput
 from apache_beam.pipeline import AppliedPTransform
+from apache_beam.runners.common import DoFnContext
+from apache_beam.runners.common import DoFnInvoker
+from apache_beam.runners.common import DoFnSignature
+from apache_beam.runners.common import Receiver
+from apache_beam.runners.common import _OutputHandler
 from apache_beam.runners.dask.overrides import _Create
 from apache_beam.runners.dask.overrides import _Flatten
 from apache_beam.runners.dask.overrides import _GroupByKeyOnly
+from apache_beam.transforms.window import GlobalWindow
+from apache_beam.transforms.window import TimestampedValue
+from apache_beam.transforms.window import WindowFn
+from apache_beam.utils.windowed_value import WindowedValue
 
 OpInput = t.Union[db.Bag, t.Sequence[db.Bag], None]
+PCollVal = t.Union[WindowedValue, t.Any]
+
+
+def get_windowed_value(item: t.Any, window_fn: WindowFn) -> WindowedValue:
+  if isinstance(item, TaggedOutput):
+    item = item.value
+
+  if isinstance(item, WindowedValue):
+    windowed_value = item
+  elif isinstance(item, TimestampedValue):
+    assign_context = WindowFn.AssignContext(item.timestamp, item.value)
+    windowed_value = WindowedValue(
+        item.value, item.timestamp, tuple(window_fn.assign(assign_context)))
+  else:
+    windowed_value = WindowedValue(item, 0, (GlobalWindow(), ))
+
+  return windowed_value
+
+
+def defenestrate(x):
+  if isinstance(x, WindowedValue):
+    return x.value
+  return x
+
+
+@dataclasses.dataclass
+class TaggingReceiver(Receiver):
+  tag: str
+  values: t.List[PCollVal]
+
+  def receive(self, windowed_value: WindowedValue):
+    if self.tag:
+      output = TaggedOutput(self.tag, windowed_value)
+    else:
+      output = windowed_value
+    self.values.append(output)
+
+
+@dataclasses.dataclass
+class OneReceiver(dict):
+  values: t.List[PCollVal] = field(default_factory=list)
+
+  def __missing__(self, key):
+    if key not in self:
+      self[key] = TaggingReceiver(key, self.values)
+    return self[key]
 
 
 @dataclasses.dataclass
@@ -65,15 +123,54 @@ def apply(self, input_bag: OpInput) -> db.Bag:
 class ParDo(DaskBagOp):
   def apply(self, input_bag: db.Bag) -> db.Bag:
     transform = t.cast(apache_beam.ParDo, self.transform)
-    return input_bag.map(
-        transform.fn.process, *transform.args, **transform.kwargs).flatten()
 
+    args, kwargs = transform.raw_side_inputs
+    args = list(args)
+    main_input = next(iter(self.applied.main_inputs.values()))
+    window_fn = main_input.windowing.windowfn if hasattr(
+        main_input, "windowing") else None
 
-class Map(DaskBagOp):
-  def apply(self, input_bag: db.Bag) -> db.Bag:
-    transform = t.cast(apache_beam.Map, self.transform)
-    return input_bag.map(
-        transform.fn.process, *transform.args, **transform.kwargs)
+    context = DoFnContext(transform.label, state=None)
+    bundle_finalizer_param = DoFn.BundleFinalizerParam()
+    do_fn_signature = DoFnSignature(transform.fn)
+
+    tagged_receivers = OneReceiver()
+
+    output_processor = _OutputHandler(
+        window_fn=window_fn,
+        main_receivers=tagged_receivers[None],
+        tagged_receivers=tagged_receivers,
+        per_element_output_counter=None,
+        output_batch_converter=None,
+        process_yields_batches=False,
+        process_batch_yields_elements=False)
+
+    do_fn_invoker = DoFnInvoker.create_invoker(
+        do_fn_signature,
+        output_processor,
+        context,
+        None,
+        args,
+        kwargs,
+        user_state_context=None,
+        bundle_finalizer_param=bundle_finalizer_param)
+
+    def apply_dofn_to_bundle(items):
+      do_fn_invoker.invoke_setup()
+      do_fn_invoker.invoke_start_bundle()
+
+      for it in items:
+        do_fn_invoker.invoke_process(it)
+
+      results = [v.value for v in tagged_receivers.values]
+
+      do_fn_invoker.invoke_finish_bundle()
+      do_fn_invoker.invoke_teardown()
+
+      return results
+
+    return input_bag.map(get_windowed_value,
+                         window_fn).map_partitions(apply_dofn_to_bundle)
 
 
 class GroupByKey(DaskBagOp):
@@ -83,21 +180,20 @@ def key(item):
 
     def value(item):
       k, v = item
-      return k, [elm[1] for elm in v]
+      return k, [defenestrate(elm[1]) for elm in v]
 
     return input_bag.groupby(key).map(value)
 
 
 class Flatten(DaskBagOp):
-  def apply(self, input_bag: OpInput) -> db.Bag:
+  def apply(self, input_bag: t.List[db.Bag]) -> db.Bag:
     assert type(input_bag) is list, 'Must take a sequence of bags!'
     return db.concat(input_bag)
 
 
 TRANSLATIONS = {
     _Create: Create,
     apache_beam.ParDo: ParDo,
-    apache_beam.Map: Map,
     _GroupByKeyOnly: GroupByKey,
     _Flatten: Flatten,
 }

sdks/python/scripts/generate_pydoc.sh

@@ -64,6 +64,7 @@ excluded_patterns=(
     'apache_beam/runners/portability/'
     'apache_beam/runners/test/'
     'apache_beam/runners/worker/'
+    'apache_beam/runners/dask/transform_evaluator.*'
     'apache_beam/testing/benchmarks/chicago_taxi/'
     'apache_beam/testing/benchmarks/inference/'
     'apache_beam/testing/benchmarks/data/'

sdks/python/setup.py

@@ -352,8 +352,8 @@ def get_portability_package_data():
           'dataframe': ['pandas>=1.0,<1.5;python_version<"3.10"',
                         'pandas>=1.4.3,<1.5;python_version>="3.10"'],
           'dask': [
-            'dask >= 2022.6',
-            'distributed >= 2022.6',
+            'dask >= 2022.6.0',
+            'distributed >= 2022.6.0',
           ],
       },
       zip_safe=False,

sdks/python/tox.ini

@@ -17,7 +17,7 @@
 
 [tox]
 # new environments will be excluded by default unless explicitly added to envlist.
-envlist = py37,py38,py39,py310,py37-{cloud,cython,lint,mypy,dask},py38-{cloud,cython,docs,cloudcoverage,dask},py39-{cloud,cython},py310-{cloud,cython,dask},whitespacelint
+envlist = py37,py38,py39,py310,py37-{cloud,cython,lint,mypy,dask},py38-{cloud,cython,docs,cloudcoverage,dask},py39-{cloud,cython,dask},py310-{cloud,cython,dask},whitespacelint
 toxworkdir = {toxinidir}/target/{env:ENV_NAME:.tox}
 
 [pycodestyle]
@@ -93,9 +93,10 @@ commands =
   {toxinidir}/scripts/run_pytest.sh {envname} "{posargs}"
 
 [testenv:py{37,38,39}-dask]
-extras = test,dask
+extras = test,dask,dataframes
 commands =
   {toxinidir}/scripts/run_pytest.sh {envname} "{posargs}"
+
 [testenv:py38-cloudcoverage]
 deps =
   codecov