periodicsequence_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.
#

"""Unit tests for the PTransform and descendants."""

# pytype: skip-file

import inspect
import logging
import random
import time
import unittest

from parameterized import parameterized

import apache_beam as beam
from apache_beam.io.restriction_trackers import OffsetRange
from apache_beam.io.restriction_trackers import OffsetRestrictionTracker
from apache_beam.io.watermark_estimators import ManualWatermarkEstimator
from apache_beam.runners.sdf_utils import RestrictionTrackerView
from apache_beam.runners.sdf_utils import ThreadsafeRestrictionTracker
from apache_beam.testing.test_pipeline import TestPipeline
from apache_beam.testing.util import assert_that
from apache_beam.testing.util import equal_to
from apache_beam.testing.util import is_empty
from apache_beam.transforms import trigger
from apache_beam.transforms import window
from apache_beam.transforms.periodicsequence import ImpulseSeqGenDoFn
from apache_beam.transforms.periodicsequence import PeriodicImpulse
from apache_beam.transforms.periodicsequence import PeriodicSequence
from apache_beam.transforms.periodicsequence import RebaseMode
from apache_beam.transforms.periodicsequence import _sequence_backlog_bytes
from apache_beam.transforms.window import FixedWindows
from apache_beam.utils.timestamp import Timestamp

# Disable frequent lint warning due to pipe operator for chaining transforms.
# pylint: disable=expression-not-assigned


class PeriodicSequenceTest(unittest.TestCase):
  def test_periodicsequence_outputs_valid_sequence(self):
    start_offset = 1
    start_time = time.time() + start_offset
    duration = 1
    end_time = start_time + duration
    interval = 0.25

    with TestPipeline() as p:
      result = (
          p
          | 'ImpulseElement' >> beam.Create([(start_time, end_time, interval)])
          | 'ImpulseSeqGen' >> PeriodicSequence())

      k = [
          start_time + x * interval
          for x in range(0, int(duration / interval), 1)
      ]
      self.assertEqual(result.is_bounded, False)
      assert_that(result, equal_to(k))

  def test_periodicsequence_outputs_valid_sequence_in_past(self):
    start_offset = -10000
    it = time.time() + start_offset
    duration = 5
    et = it + duration
    interval = 1

    with TestPipeline() as p:
      result = (
          p
          | 'ImpulseElement' >> beam.Create([(it, et, interval)])
          | 'ImpulseSeqGen' >> PeriodicSequence())

      k = [it + x * interval for x in range(0, int(duration / interval), 1)]
      self.assertEqual(result.is_bounded, False)
      assert_that(result, equal_to(k))

  def test_periodicsequence_output_size(self):
    element = [0, 1000000000, 10]
    self.assertEqual(
        _sequence_backlog_bytes(element, 100, OffsetRange(10, 100000000)), 0)
    self.assertEqual(
        _sequence_backlog_bytes(element, 100, OffsetRange(9, 100000000)), 8)
    self.assertEqual(
        _sequence_backlog_bytes(element, 100, OffsetRange(8, 100000000)), 16)
    self.assertEqual(
        _sequence_backlog_bytes(element, 101, OffsetRange(9, 100000000)), 8)
    self.assertEqual(
        _sequence_backlog_bytes(element, 10000, OffsetRange(0, 100000000)),
        8 * 10000 / 10)
    self.assertEqual(
        _sequence_backlog_bytes(element, 10000, OffsetRange(1002, 1003)), 0)
    self.assertEqual(
        _sequence_backlog_bytes(element, 10100, OffsetRange(1002, 1003)), 8)


class PeriodicImpulseTest(unittest.TestCase):
  def test_windowing_on_si(self):
    start_offset = -15
    it = time.time() + start_offset
    duration = 15
    et = it + duration
    interval = 5

    with TestPipeline() as p:
      si = (
          p
          | 'PeriodicImpulse' >> PeriodicImpulse(it, et, interval, True)
          | 'AddKey' >> beam.Map(lambda v: ('key', v))
          | 'GBK' >> beam.GroupByKey()
          | 'SortGBK' >> beam.MapTuple(lambda k, vs: (k, sorted(vs))))

      actual = si
      k = [('key', [it + x * interval])
           for x in range(0, int(duration / interval), 1)]
      assert_that(actual, equal_to(k))

  def test_default_start(self):
    default_parameters = inspect.signature(PeriodicImpulse.__init__).parameters
    it = default_parameters["start_timestamp"].default
    duration = 1
    et = it + duration
    interval = 0.5

    # Check default `stop_timestamp` is the same type `start_timestamp`
    is_same_type = isinstance(
        it, type(default_parameters["stop_timestamp"].default))
    error = "'start_timestamp' and 'stop_timestamp' have different type"
    assert is_same_type, error

    with TestPipeline() as p:
      result = p | 'PeriodicImpulse' >> PeriodicImpulse(it, et, interval)

      k = [it + x * interval for x in range(0, int(duration / interval))]
      self.assertEqual(result.is_bounded, False)
      assert_that(result, equal_to(k))

  @unittest.skip("hard to determine warm-up time and threshold for runners.")
  def test_processing_time(self):
    warmup_time = 3
    threshold = 0.5
    with TestPipeline() as p:
      ret = (
          p | PeriodicImpulse(
              start_timestamp=Timestamp.now() + warmup_time,
              data=[10, 20, 30],
              fire_interval=2)
          | beam.Map(lambda _: time.time())
          | beam.WindowInto(
              window.GlobalWindows(),
              trigger=trigger.Repeatedly(trigger.AfterCount(3)),
              accumulation_mode=trigger.AccumulationMode.DISCARDING,
          )
          | beam.GroupBy()
          | beam.FlatMap(lambda x: [v - min(x[1]) for v in x[1]]))
      expected = [0, 2, 4]
      assert_that(ret, equal_to(expected, lambda x, y: abs(x - y) < threshold))

  @parameterized.expand([0.5, 1, 2, 10])
  def test_stop_over_by_epsilon(self, interval):
    with TestPipeline() as p:
      ret = (
          p | PeriodicImpulse(
              start_timestamp=Timestamp(seconds=1),
              stop_timestamp=Timestamp(seconds=1, micros=1),
              data=[1, 2],
              fire_interval=interval)
          | beam.WindowInto(FixedWindows(interval))
          | beam.WithKeys(0)
          | beam.GroupByKey())
      expected = [
          (0, [1]),
      ]
      assert_that(ret, equal_to(expected))

  @parameterized.expand([1, 2])
  def test_stop_over_by_interval(self, interval):
    with TestPipeline() as p:
      ret = (
          p | PeriodicImpulse(
              start_timestamp=Timestamp(seconds=1),
              stop_timestamp=Timestamp(seconds=1 + interval),
              data=[1, 2],
              fire_interval=interval)
          | beam.WindowInto(FixedWindows(interval))
          | beam.WithKeys(0)
          | beam.GroupByKey())
      expected = [(0, [1])]
      assert_that(ret, equal_to(expected))

  @parameterized.expand([1, 2])
  def test_stop_over_by_interval_and_epsilon(self, interval):
    with TestPipeline() as p:
      ret = (
          p | PeriodicImpulse(
              start_timestamp=Timestamp(seconds=1),
              stop_timestamp=Timestamp(seconds=1 + interval, micros=1),
              data=[1, 2],
              fire_interval=interval)
          | beam.WindowInto(FixedWindows(interval))
          | beam.WithKeys(0)
          | beam.GroupByKey())
      expected = [(0, [1]), (0, [2])]
      assert_that(ret, equal_to(expected))

  def test_interval(self):
    with TestPipeline() as p:
      ret = (
          p | PeriodicImpulse(data=[1, 2, 3, 4], fire_interval=0.5)
          | beam.WindowInto(FixedWindows(0.5))
          | beam.WithKeys(0)
          | beam.GroupByKey())
      expected = [(0, [1]), (0, [2]), (0, [3]), (0, [4])]
      assert_that(ret, equal_to(expected))

  def test_repeat(self):
    now = Timestamp.now()
    with self.assertWarnsRegex(UserWarning, "not enough to span"):
      with TestPipeline() as p:
        ret = (
            p | PeriodicImpulse(
                start_timestamp=now,
                stop_timestamp=now + 2.6,
                data=[1, 2, 3, 4],
                fire_interval=0.5)
            | beam.WindowInto(FixedWindows(0.5))
            | beam.WithKeys(0)
            | beam.GroupByKey())
        expected = [(0, [1]), (0, [2]), (0, [3]), (0, [4]), (0, [1]), (0, [2])]
        assert_that(ret, equal_to(expected))

  def test_timestamped_value(self):
    with TestPipeline() as p:
      ret = (
          p | PeriodicImpulse(
              data=[(Timestamp(1), 1), (Timestamp(3), 2), (Timestamp(2), 3),
                    (Timestamp(1), 4)],
              fire_interval=0.5)
          | beam.WindowInto(FixedWindows(0.5))
          | beam.WithKeys(0)
          | beam.GroupByKey())
      expected = [(0, [1, 4]), (0, [2]), (0, [3])]
      assert_that(ret, equal_to(expected))

  def test_not_enough_timestamped_value(self):
    now = Timestamp.now()
    data = [(Timestamp(1), 1), (Timestamp(2), 2), (Timestamp(3), 3)]
    with self.assertRaisesRegex(ValueError, "not enough to span"):
      with TestPipeline() as p:
        _ = (
            p | PeriodicImpulse(
                start_timestamp=now,
                stop_timestamp=now + 2.6,
                data=data,
                fire_interval=0.5))

  def test_fuzzy_length_and_interval(self):
    times = 30
    for _ in range(times):
      seed = int(time.time() * 1000)
      random.seed(seed)
      n = int(random.randint(1, 100))
      data = list(range(n))
      m = random.randint(1, 1000)
      interval = m / 1e6
      now = Timestamp.now()
      try:
        with TestPipeline() as p:
          ret = (
              p | PeriodicImpulse(
                  start_timestamp=now, data=data, fire_interval=interval))
          assert_that(ret, equal_to(data))
      except Exception as e:  # pylint: disable=broad-except
        logging.error("Error occurred at random seed=%d", seed)
        raise e

  def test_fuzzy_length_at_minimal_interval(self):
    times = 30
    for _ in range(times):
      seed = int(time.time() * 1000)
      random.seed(seed)
      n = int(random.randint(1, 100))
      data = list(range(n))
      interval = 1e-6
      now = Timestamp.now()
      try:
        with TestPipeline() as p:
          ret = (
              p | PeriodicImpulse(
                  start_timestamp=now, data=data, fire_interval=interval))
          assert_that(ret, equal_to(data))
      except Exception as e:  # pylint: disable=broad-except
        logging.error("Error occurred at random seed=%d", seed)
        raise e

  def test_int_type_input(self):
    # This test is to verify that if input timestamps and interval are integers,
    # the generated timestamped values are also integers.
    # This is necessary for the following test to pass:
    # apache_beam.examples.snippets.snippets_test.SlowlyChangingSideInputsTest
    with TestPipeline() as p:
      ret = (
          p | PeriodicImpulse(
              start_timestamp=1, stop_timestamp=5, fire_interval=1))
      expected = [1, 2, 3, 4]
      assert_that(
          ret, equal_to(expected, lambda x, y: type(x) is type(y) and x == y))

  def test_float_type_input(self):
    with TestPipeline() as p:
      ret = (
          p | PeriodicImpulse(
              start_timestamp=1.0, stop_timestamp=5.0, fire_interval=1))
      expected = [1.0, 2.0, 3.0, 4.0]
      assert_that(
          ret, equal_to(expected, lambda x, y: type(x) is type(y) and x == y))

  def test_timestamp_type_input(self):
    with TestPipeline() as p:
      ret = (
          p | PeriodicImpulse(
              start_timestamp=Timestamp.of(1),
              stop_timestamp=Timestamp.of(5),
              fire_interval=1))
      expected = [1.0, 2.0, 3.0, 4.0]
      assert_that(
          ret, equal_to(expected, lambda x, y: type(x) is type(y) and x == y))

  def test_rebase_timestamp(self):
    class CheckTimeStamp(beam.DoFn):
      def process(self, elem):
        ts = Timestamp.of(elem)
        now = Timestamp.now()
        # When rebase is enabled, the timestamp should be closer to now than the
        # original start.
        if (ts - Timestamp.of(1)) < (now - ts):
          yield "wrong"

    with TestPipeline() as p:
      ret = (
          p | PeriodicImpulse(
              start_timestamp=Timestamp.of(1),
              stop_timestamp=Timestamp.of(5),
              fire_interval=1,
              rebase=RebaseMode.REBASE_ALL)
          | beam.ParDo(CheckTimeStamp()))
      assert_that(ret, is_empty())

  def test_rebase_timestamp_with_wrong_setting(self):
    with self.assertRaises(Exception):
      # exception is raised because start_timestamp is rebased to the pipeline
      # execution time, but the stop_timestamp is 5 seconds after unix epoch.
      with TestPipeline() as p:
        _ = (
            p | PeriodicImpulse(
                start_timestamp=Timestamp.of(1),
                stop_timestamp=Timestamp.of(5),
                fire_interval=1,
                rebase=RebaseMode.REBASE_START))


class ImpulseSeqGenDoFnWatermarkTest(unittest.TestCase):
  """Drives ``ImpulseSeqGenDoFn.process`` directly to assert the reported
  watermark when the DoFn defers because the next fire time is in the future.
  """
  @staticmethod
  def _run_process(dofn, element, restriction, initial_watermark=None):
    tracker = ThreadsafeRestrictionTracker(
        OffsetRestrictionTracker(restriction))
    view = RestrictionTrackerView(tracker)
    estimator = ManualWatermarkEstimator(initial_watermark)
    outputs = list(
        dofn.process(
            element, restriction_tracker=view, watermark_estimator=estimator))
    return outputs, estimator

  def test_watermark_advances_to_next_fire_on_defer(self):
    # Regression test for https://github.com/apache/beam/issues/39026.
    # With a long fire_interval, index 0 fires now and index 1 is scheduled far
    # in the future, so the DoFn defers. The reported watermark must advance to
    # the next fire time rather than stalling at the last emitted element's
    # timestamp (which caused a saw-tooth watermark age from 2.74.0).
    interval = 100
    start = time.time() - 5
    element = (start, start + 10000, interval)

    outputs, estimator = self._run_process(
        ImpulseSeqGenDoFn(), element, OffsetRange(0, 5))

    self.assertEqual(len(outputs), 1)
    last_emitted = Timestamp(start)
    next_fire = Timestamp(start + interval)
    self.assertEqual(outputs[0].timestamp, last_emitted)
    # Watermark advanced past the last emitted element, up to the next fire.
    self.assertEqual(estimator.current_watermark(), next_fire)
    self.assertGreater(estimator.current_watermark(), last_emitted)

  def test_watermark_not_advanced_past_emitted_for_pre_timestamped(self):
    # For pre-timestamped data the provided event times may be out of order, so
    # the watermark must not be advanced to the schedule time on defer; it stays
    # at the latest emitted event timestamp.
    interval = 100
    start = time.time() - 5
    element = (start, start + 10000, interval)
    data = [(Timestamp(1), 'a'), (Timestamp(2), 'b'), (Timestamp(3), 'c')]

    outputs, estimator = self._run_process(
        ImpulseSeqGenDoFn(data), element, OffsetRange(0, 5))

    self.assertEqual(len(outputs), 1)
    self.assertEqual(outputs[0].value, 'a')
    self.assertEqual(outputs[0].timestamp, Timestamp(1))
    # Watermark stays at the emitted event time, not the future schedule time.
    self.assertEqual(estimator.current_watermark(), Timestamp(1))
    self.assertLess(estimator.current_watermark(), Timestamp(start + interval))


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