Windowing Support for the DaskRunner

.github/workflows/dask_runner_tests.yml

@@ -62,7 +62,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        os: [ubuntu-latest, macos-latest, windows-latest]
+        os: [ubuntu-latest, macos-latest]
         params: [
           {"py_ver": "3.8", "tox_env": "py38"},
           {"py_ver": "3.9", "tox_env": "py39"},
@@ -79,15 +79,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

@@ -37,6 +37,13 @@
 from apache_beam.runners.runner import PipelineState
 from apache_beam.utils.interactive_utils import is_in_notebook
 
+try:
+  # Added to try to prevent threading related issues, see
+  # https://github.com/pytest-dev/pytest/issues/3216#issuecomment-1502451456
+  from dask import distributed
+except ImportError:
+  distributed = {}
+
 
 class DaskOptions(PipelineOptions):
   @staticmethod
@@ -86,8 +93,6 @@ def _add_argparse_args(cls, parser: argparse.ArgumentParser) -> None:
 
 @dataclasses.dataclass
 class DaskRunnerResult(PipelineResult):
-  from dask import distributed
-
   client: distributed.Client
   futures: t.Sequence[distributed.Future]
 
@@ -99,8 +104,7 @@ def wait_until_finish(self, duration=None) -> str:
       if duration is not None:
         # Convert milliseconds to seconds
         duration /= 1000
-      self.client.wait_for_workers(timeout=duration)
-      self.client.gather(self.futures, errors='raise')
+      distributed.wait(self.futures, timeout=duration)
       self._state = PipelineState.DONE
     except:  # pylint: disable=broad-except
       self._state = PipelineState.FAILED
@@ -159,6 +163,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!')
@@ -177,6 +183,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(list(dask_visitor.bags.values()))
+    futures = client.compute(opt_graph)
     return DaskRunnerResult(client, futures)

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

@@ -14,6 +14,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
+import datetime
 import inspect
 import unittest
 
@@ -22,12 +23,14 @@
 from apache_beam.testing import test_pipeline
 from apache_beam.testing.util import assert_that
 from apache_beam.testing.util import equal_to
+from apache_beam.transforms import window
 
 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 +76,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 +89,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 +113,50 @@ 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]))
+
+  def test_groupby_with_fixed_windows(self):
+    def double(x):
+      return x * 2, x
+
+    def add_timestamp(pair):
+      delta = datetime.timedelta(seconds=pair[1] * 60)
+      now = (datetime.datetime.now() + delta).timestamp()
+      return window.TimestampedValue(pair, now)
+
+    with self.pipeline as p:
+      pcoll = (
+          p
+          | beam.Create([1, 2, 1, 2, 3])
+          | beam.Map(double)
+          | beam.WindowInto(window.FixedWindows(60))
+          | beam.Map(add_timestamp)
+          | beam.GroupByKey())
+      assert_that(pcoll, equal_to([(2, [1, 1]), (4, [2, 2]), (6, [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):