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

#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
import datetime
import inspect
import typing as t
import unittest

import apache_beam as beam
from apache_beam.options.pipeline_options import PipelineOptions
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:
  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.')


class DaskOptionsTest(unittest.TestCase):
  def test_parses_connection_timeout__defaults_to_none(self):
    default_options = PipelineOptions([])
    default_dask_options = default_options.view_as(DaskOptions)
    self.assertEqual(None, default_dask_options.timeout)

  def test_parses_connection_timeout__parses_int(self):
    conn_options = PipelineOptions('--dask_connection_timeout 12'.split())
    dask_conn_options = conn_options.view_as(DaskOptions)
    self.assertEqual(12, dask_conn_options.timeout)

  def test_parses_connection_timeout__handles_bad_input(self):
    err_options = PipelineOptions('--dask_connection_timeout foo'.split())
    dask_err_options = err_options.view_as(DaskOptions)
    self.assertEqual(dask.config.no_default, dask_err_options.timeout)

  def test_parser_destinations__agree_with_dask_client(self):
    options = PipelineOptions(
        '--dask_client_address localhost:8080 --dask_connection_timeout 600 '
        '--dask_scheduler_file foobar.cfg --dask_client_name charlie '
        '--dask_connection_limit 1024'.split())
    dask_options = options.view_as(DaskOptions)

    # Get the argument names for the constructor.
    client_args = list(inspect.signature(ddist.Client).parameters)

    for opt_name in dask_options.get_all_options(drop_default=True).keys():
      with self.subTest(f'{opt_name} in dask.distributed.Client constructor'):
        self.assertIn(opt_name, client_args)

  def test_parser_extract_bag_kwargs__deletes_dask_kwargs(self):
    options = PipelineOptions('--dask_npartitions 8'.split())
    dask_options = options.view_as(DaskOptions).get_all_options()

    self.assertIn('npartitions', dask_options)
    bag_kwargs = DaskOptions._extract_bag_kwargs(dask_options)
    self.assertNotIn('npartitions', dask_options)
    self.assertEqual(bag_kwargs, {'npartitions': 8})

  def test_parser_extract_bag_kwargs__unconfigured(self):
    options = PipelineOptions()
    dask_options = options.view_as(DaskOptions).get_all_options()

    # It's present as a default option.
    self.assertIn('npartitions', dask_options)
    bag_kwargs = DaskOptions._extract_bag_kwargs(dask_options)
    self.assertNotIn('npartitions', dask_options)
    self.assertEqual(bag_kwargs, {})


class DaskRunnerRunPipelineTest(unittest.TestCase):
  """Test class used to introspect the dask runner via a debugger."""
  def setUp(self) -> None:
    self.pipeline = test_pipeline.TestPipeline(runner=DaskRunner())

  def test_create(self):
    with self.pipeline as p:
      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, 4])
      assert_that(pcoll, equal_to([1, 2, 3, 4]))

  def test_create_and_map(self):
    def double(x):
      return x * 2

    with self.pipeline as p:
      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

    with self.pipeline as p:
      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_positional_side_input(self):
    def mult_by(x, y):
      return x * y

    with self.pipeline as p:
      side = p | "side" >> beam.Create([3])
      pcoll = (
          p
          | "main" >> beam.Create([1])
          | beam.Map(mult_by, beam.pvalue.AsSingleton(side)))
      assert_that(pcoll, equal_to([3]))

  def test_map_with_keyword_side_input(self):
    def mult_by(x, y):
      return x * y

    with self.pipeline as p:
      side = p | "side" >> beam.Create([3])
      pcoll = (
          p
          | "main" >> beam.Create([1])
          | beam.Map(mult_by, y=beam.pvalue.AsSingleton(side)))
      assert_that(pcoll, equal_to([3]))

  def test_pardo_side_inputs(self):
    def cross_product(elem, sides):
      for side in sides:
        yield elem, side

    with self.pipeline as p:
      main = p | "main" >> beam.Create(["a", "b", "c"])
      side = p | "side" >> beam.Create(["x", "y"])
      assert_that(
          main | beam.FlatMap(cross_product, beam.pvalue.AsList(side)),
          equal_to([
              ("a", "x"),
              ("b", "x"),
              ("c", "x"),
              ("a", "y"),
              ("b", "y"),
              ("c", "y"),
          ]),
      )

  def test_pardo_side_input_dependencies(self):
    with self.pipeline as p:
      inputs = [p | beam.Create([None])]
      for k in range(1, 10):
        inputs.append(
            inputs[0]
            | beam.ParDo(
                ExpectingSideInputsFn(f"Do{k}"),
                *[beam.pvalue.AsList(inputs[s]) for s in range(1, k)],
            ))

  def test_pardo_side_input_sparse_dependencies(self):
    with self.pipeline as p:
      inputs = []

      def choose_input(s):
        return inputs[(389 + s * 5077) % len(inputs)]

      for k in range(20):
        num_inputs = int((k * k % 16)**0.5)
        if num_inputs == 0:
          inputs.append(p | f"Create{k}" >> beam.Create([f"Create{k}"]))
        else:
          inputs.append(
              choose_input(0)
              | beam.ParDo(
                  ExpectingSideInputsFn(f"Do{k}"),
                  *[
                      beam.pvalue.AsList(choose_input(s))
                      for s in range(1, num_inputs)
                  ],
              ))

  @unittest.expectedFailure
  def test_pardo_windowed_side_inputs(self):
    with self.pipeline as p:
      # Now with some windowing.
      pcoll = (
          p
          | beam.Create(list(range(10)))
          | beam.Map(lambda t: window.TimestampedValue(t, t)))
      # Intentionally choosing non-aligned windows to highlight the transition.
      main = pcoll | "WindowMain" >> beam.WindowInto(window.FixedWindows(5))
      side = pcoll | "WindowSide" >> beam.WindowInto(window.FixedWindows(7))
      res = main | beam.Map(
          lambda x, s: (x, sorted(s)), beam.pvalue.AsList(side))
      assert_that(
          res,
          equal_to([
              # The window [0, 5) maps to the window [0, 7).
              (0, list(range(7))),
              (1, list(range(7))),
              (2, list(range(7))),
              (3, list(range(7))),
              (4, list(range(7))),
              # The window [5, 10) maps to the window [7, 14).
              (5, list(range(7, 10))),
              (6, list(range(7, 10))),
              (7, list(range(7, 10))),
              (8, list(range(7, 10))),
              (9, list(range(7, 10))),
          ]),
          label="windowed",
      )

  def test_flattened_side_input(self, with_transcoding=True):
    with self.pipeline as p:
      main = p | "main" >> beam.Create([None])
      side1 = p | "side1" >> beam.Create([("a", 1)])
      side2 = p | "side2" >> beam.Create([("b", 2)])
      if with_transcoding:
        # Also test non-matching coder types (transcoding required)
        third_element = [("another_type")]
      else:
        third_element = [("b", 3)]
      side3 = p | "side3" >> beam.Create(third_element)
      side = (side1, side2) | beam.Flatten()
      assert_that(
          main | beam.Map(lambda a, b: (a, b), beam.pvalue.AsDict(side)),
          equal_to([(None, {
              "a": 1, "b": 2
          })]),
          label="CheckFlattenAsSideInput",
      )
      assert_that(
          (side, side3) | "FlattenAfter" >> beam.Flatten(),
          equal_to([("a", 1), ("b", 2)] + third_element),
          label="CheckFlattenOfSideInput",
      )

  def test_gbk_side_input(self):
    with self.pipeline as p:
      main = p | "main" >> beam.Create([None])
      side = p | "side" >> beam.Create([("a", 1)]) | beam.GroupByKey()
      assert_that(
          main | beam.Map(lambda a, b: (a, b), beam.pvalue.AsDict(side)),
          equal_to([(None, {
              "a": [1]
          })]),
      )

  def test_multimap_side_input(self):
    with self.pipeline as p:
      main = p | "main" >> beam.Create(["a", "b"])
      side = p | "side" >> beam.Create([("a", 1), ("b", 2), ("a", 3)])
      assert_that(
          main
          | beam.Map(
              lambda k, d: (k, sorted(d[k])), beam.pvalue.AsMultiMap(side)),
          equal_to([("a", [1, 3]), ("b", [2])]),
      )

  def test_multimap_multiside_input(self):
    # A test where two transforms in the same stage consume the same PCollection
    # twice as side input.
    with self.pipeline as p:
      main = p | "main" >> beam.Create(["a", "b"])
      side = p | "side" >> beam.Create([("a", 1), ("b", 2), ("a", 3)])
      assert_that(
          main
          | "first map" >> beam.Map(
              lambda k, d, l: (k, sorted(d[k]), sorted([e[1] for e in l])),
              beam.pvalue.AsMultiMap(side),
              beam.pvalue.AsList(side),
          )
          | "second map" >> beam.Map(
              lambda k, d, l:
              (k[0], sorted(d[k[0]]), sorted([e[1] for e in l])),
              beam.pvalue.AsMultiMap(side),
              beam.pvalue.AsList(side),
          ),
          equal_to([("a", [1, 3], [1, 2, 3]), ("b", [2], [1, 2, 3])]),
      )

  def test_multimap_side_input_type_coercion(self):
    with self.pipeline as p:
      main = p | "main" >> beam.Create(["a", "b"])
      # The type of this side-input is forced to Any (overriding type
      # inference). Without type coercion to Tuple[Any, Any], the usage of this
      # side-input in AsMultiMap() below should fail.
      side = p | "side" >> beam.Create([("a", 1), ("b", 2),
                                        ("a", 3)]).with_output_types(t.Any)
      assert_that(
          main
          | beam.Map(
              lambda k, d: (k, sorted(d[k])), beam.pvalue.AsMultiMap(side)),
          equal_to([("a", [1, 3]), ("b", [2])]),
      )

  def test_pardo_unfusable_side_inputs__one(self):
    def cross_product(elem, sides):
      for side in sides:
        yield elem, side

    with self.pipeline as p:
      pcoll = p | "Create1" >> beam.Create(["a", "b"])
      assert_that(
          pcoll |
          "FlatMap1" >> beam.FlatMap(cross_product, beam.pvalue.AsList(pcoll)),
          equal_to([("a", "a"), ("a", "b"), ("b", "a"), ("b", "b")]),
          label="assert_that1",
      )

  def test_pardo_unfusable_side_inputs__two(self):
    def cross_product(elem, sides):
      for side in sides:
        yield elem, side

    with self.pipeline as p:
      pcoll = p | "Create2" >> beam.Create(["a", "b"])

      derived = ((pcoll, )
                 | beam.Flatten()
                 | beam.Map(lambda x: (x, x))
                 | beam.GroupByKey()
                 | "Unkey" >> beam.Map(lambda kv: kv[0]))
      assert_that(
          pcoll | "FlatMap2" >> beam.FlatMap(
              cross_product, beam.pvalue.AsList(derived)),
          equal_to([("a", "a"), ("a", "b"), ("b", "a"), ("b", "b")]),
          label="assert_that2",
      )

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

  def test_groupby_string_keys(self):
    with self.pipeline as p:
      pcoll = (
          p
          | beam.Create([('a', 1), ('a', 2), ('b', 3), ('b', 4)])
          | beam.GroupByKey())
      assert_that(pcoll, equal_to([('a', [1, 2]), ('b', [3, 4])]))


class ExpectingSideInputsFn(beam.DoFn):
  def __init__(self, name):
    self._name = name

  def default_label(self):
    return self._name

  def process(self, element, *side_inputs):
    if not all(list(s) for s in side_inputs):
      raise ValueError(f"Missing data in side input {side_inputs}")
    yield self._name


if __name__ == '__main__':
  unittest.main()