Add a flag in PeriodicImpulse to rebase timestamp. (#35569)

* Add a flag in PeriodicImpulse to rebase timestamp.

* Add a unit test.

* Change the boolean flag to an enum.

* Clarify the test that raises exception.

Author: shunping

sdks/python/apache_beam/transforms/periodicsequence.py

@@ -15,6 +15,7 @@
 # limitations under the License.
 #
 
+import enum
 import math
 import time
 import warnings
@@ -216,6 +217,21 @@ def expand(self, pcoll):
         | 'MapToTimestamped' >> beam.Map(lambda tt: TimestampedValue(tt, tt)))
 
 
+class RebaseMode(enum.Enum):
+  '''Controls how the start and stop timestamps are rebased to execution time.
+
+  Attributes:
+    REBASE_NONE: Timestamps are not changed.
+    REBASE_ALL: Both start and stop timestamps are rebased, preserving the
+      original duration.
+    REBASE_START: Only the start timestamp is rebased; the stop timestamp
+      is unchanged.
+  '''
+  REBASE_NONE = 0
+  REBASE_ALL = 1
+  REBASE_START = 2
+
+
 class PeriodicImpulse(PTransform):
   '''
   PeriodicImpulse transform generates an infinite sequence of elements with
@@ -270,7 +286,8 @@ def __init__(
       stop_timestamp: TimestampTypes = MAX_TIMESTAMP,
       fire_interval: float = 360.0,
       apply_windowing: bool = False,
-      data: Optional[Sequence[Any]] = None):
+      data: Optional[Sequence[Any]] = None,
+      rebase: RebaseMode = RebaseMode.REBASE_NONE):
     '''
     :param start_timestamp: Timestamp for first element.
     :param stop_timestamp: Timestamp at or after which no elements will be
@@ -301,21 +318,36 @@ def __init__(
         sufficient to cover the duration defined by `start_timestamp`,
         `stop_timestamp`, and `fire_interval`; otherwise, a `ValueError` is
         raised.
+
+    :param rebase: Controls how the start and stop timestamps are rebased to
+      execution time. See `RebaseMode` for more details. Defaults to
+      `REBASE_NONE`.
     '''
     self.start_ts = start_timestamp
     self.stop_ts = stop_timestamp
     self.interval = fire_interval
     self.apply_windowing = apply_windowing
     self.data = data
+    self.rebase = rebase
 
     if self.data:
       self._validate_and_adjust_duration()
 
   def expand(self, pbegin):
+    if self.rebase == RebaseMode.REBASE_ALL:
+      duration = Timestamp.of(self.stop_ts) - Timestamp.of(self.start_ts)
+      impulse_element = pbegin | beam.Impulse() | beam.Map(
+          lambda _:
+          [Timestamp.now(), Timestamp.now() + duration, self.interval])
+    elif self.rebase == RebaseMode.REBASE_START:
+      impulse_element = pbegin | beam.Impulse() | beam.Map(
+          lambda _: [Timestamp.now(), self.stop_ts, self.interval])
+    else:
+      impulse_element = pbegin | 'ImpulseElement' >> beam.Create(
+          [(self.start_ts, self.stop_ts, self.interval)])
+
     result = (
-        pbegin
-        | 'ImpulseElement' >> beam.Create(
-            [(self.start_ts, self.stop_ts, self.interval)])
+        impulse_element
         | 'GenSequence' >> beam.ParDo(ImpulseSeqGenDoFn(self.data)))
 
     if not self.data:

sdks/python/apache_beam/transforms/periodicsequence_test.py

@@ -32,10 +32,12 @@
 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 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
@@ -285,7 +287,6 @@ def test_fuzzy_length_at_minimal_interval(self):
     times = 30
     for _ in range(times):
       seed = int(time.time() * 1000)
-      seed = 1751135957975
       random.seed(seed)
       n = int(random.randint(1, 100))
       data = list(range(n))
@@ -334,6 +335,38 @@ def test_timestamp_type_input(self):
       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))
+
 
 if __name__ == '__main__':
   unittest.main()