diff --git a/CHANGES.md b/CHANGES.md index 7ea2cccde299..6bf160021516 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -88,6 +88,7 @@ * Fixed unbounded checkpoint state growth for splittable DoFns that self-checkpoint on the portable Flink runner (Java) ([#27648](https://github.com/apache/beam/issues/27648)). * Improved Java pipeline performance by avoiding repeated `DoFn` type descriptor resolution when creating cached invokers ([#39309](https://github.com/apache/beam/issues/39309)). * (Python) Fixed a memory leak in Python SDK caused by storing exceptions with potentially large stack frames in a cache ([#39406](https://github.com/apache/beam/issues/39406)). +* (Python) Fixed `PeriodicImpulse`/`PeriodicSequence` watermark regression where the reported watermark stalled at the last emitted element's timestamp, causing a saw-tooth watermark age lagging by up to one `fire_interval` ([#39026](https://github.com/apache/beam/issues/39026)). ## Security Fixes diff --git a/sdks/python/apache_beam/transforms/periodicsequence.py b/sdks/python/apache_beam/transforms/periodicsequence.py index e2bdc3c6c0f8..6a221c618b69 100644 --- a/sdks/python/apache_beam/transforms/periodicsequence.py +++ b/sdks/python/apache_beam/transforms/periodicsequence.py @@ -166,9 +166,26 @@ def process( current_output_timestamp = start + interval * current_output_index if current_output_timestamp > time.time(): - # we are too ahead of time, let's wait. - restriction_tracker.defer_remainder( - timestamp.Timestamp(current_output_timestamp)) + # We are ahead of time, let's wait. No element will be produced before + # the next fire time, so advance the watermark up to that timestamp. + # Without this the reported watermark stalls at the last emitted + # element's timestamp and lags by up to one fire_interval, producing a + # saw-tooth watermark age (regression introduced in 2.67.0 by the + # process() rewrite in #35412, see + # https://github.com/apache/beam/issues/39026). + # + # For pre-timestamped data the provided event times may be out of order + # (a later element can carry an earlier timestamp, treated as a late + # event), so we must not advance the watermark past what has already + # been emitted. + next_output_timestamp = timestamp.Timestamp(current_output_timestamp) + if not self._is_pre_timestamped: + current_watermark = watermark_estimator.current_watermark() + if current_watermark is None or \ + next_output_timestamp > current_watermark: + # ensure watermark is monotonic + watermark_estimator.set_watermark(next_output_timestamp) + restriction_tracker.defer_remainder(next_output_timestamp) break if not restriction_tracker.try_claim(current_output_index): diff --git a/sdks/python/apache_beam/transforms/periodicsequence_test.py b/sdks/python/apache_beam/transforms/periodicsequence_test.py index 6fbc68daed8b..367d36d8a4c9 100644 --- a/sdks/python/apache_beam/transforms/periodicsequence_test.py +++ b/sdks/python/apache_beam/transforms/periodicsequence_test.py @@ -29,12 +29,17 @@ 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 @@ -368,5 +373,61 @@ def test_rebase_timestamp_with_wrong_setting(self): 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()