Windowing Support for the DaskRunner

.github/workflows/dask_runner_tests.yml

@@ -65,7 +65,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"},
@@ -85,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

@@ -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!')
@@ -178,5 +184,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

@@ -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):

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/cloudml/'
  'apache_beam/testing/benchmarks/inference/'

sdks/python/setup.py

@@ -352,8 +352,8 @@ def get_portability_package_data():
  ],
  'dataframe': dataframe_dependency,
  'dask': [
- 'dask >= 2022.6',
- 'distributed >= 2022.6',
+ 'dask >= 2022.6.0',
+ 'distributed >= 2022.6.0',
  ],
  },
  zip_safe=False,