Fix some edge cases and add a fuzzy interval test.

Author: shunping

sdks/python/apache_beam/transforms/periodicsequence.py

@@ -47,7 +47,7 @@ def initial_restriction(self, element):
 
     assert start <= end
     assert interval > 0
-    total_outputs = math.ceil((end - start) / interval)
+    total_outputs = round((end - start) / interval)
     return OffsetRange(0, total_outputs)
 
   def create_tracker(self, restriction):
@@ -232,10 +232,10 @@ def _validate_and_adjust_duration(self):
 
     if isinstance(self.stop_ts, Timestamp):
       if self.stop_ts == MAX_TIMESTAMP:
-        # adjust stop timestamp to match the data duration
-        end = start + data_duration
-        if self.interval > 1e-6:
-          end += 1e-6
+        # When the stop timestamp is unbounded (MAX_TIMESTAMP), set it to the
+        # data's actual end time plus an extra fire interval, because the
+        # impulse duration's upper bound is exclusive.
+        end = start + data_duration + self.interval
         self.stop_ts = Timestamp.of(end)
       else:
         end = self.stop_ts.micros / 1000000
@@ -244,7 +244,7 @@ def _validate_and_adjust_duration(self):
 
     # The total time for the impulse signal which occurs in [start, end).
     impulse_duration = end - start
-    if data_duration + self.interval < impulse_duration:
+    if round(data_duration + self.interval, 6) < round(impulse_duration, 6):
       # We don't have enough data for the impulse.
       # If we can fit at least one more data point in the impulse duration,
       # then we will be in the repeat mode.
@@ -264,8 +264,8 @@ def _validate_and_adjust_duration(self):
 
   def __init__(
       self,
-      start_timestamp: Union[Timestamp, float] = Timestamp.now(),
-      stop_timestamp: Union[Timestamp, float] = MAX_TIMESTAMP,
+      start_timestamp: Timestamp = Timestamp.now(),
+      stop_timestamp: Timestamp = MAX_TIMESTAMP,
       fire_interval: float = 360.0,
       apply_windowing: bool = False,
       data: Optional[Sequence[Any]] = None):

sdks/python/apache_beam/transforms/periodicsequence_test.py

@@ -19,7 +19,9 @@
 
 # pytype: skip-file
 
+import logging
 import inspect
+import random
 import time
 import unittest
 
@@ -174,7 +176,7 @@ def test_repeat(self):
         ret = (
             p | PeriodicImpulse(
                 start_timestamp=now,
-                stop_timestamp=now + 2.6,
+                stop_timestamp=now + 3.0,
                 data=[1, 2, 3, 4],
                 fire_interval=0.5)
             | beam.WindowInto(FixedWindows(0.5))
@@ -208,15 +210,22 @@ def test_not_enough_timestamped_value(self):
                 data=data,
                 fire_interval=0.5))
 
-  def test_small_interval(self):
-    data = [(Timestamp(1), 1), (Timestamp(2), 2), (Timestamp(3), 3),
-            (Timestamp(6), 6), (Timestamp(4), 4), (Timestamp(5), 5),
-            (Timestamp(7), 7), (Timestamp(8), 8), (Timestamp(9), 9),
-            (Timestamp(10), 10)]
-    expected = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
-    with TestPipeline() as p:
-      ret = (p | PeriodicImpulse(data=data, fire_interval=0.0001))
-      assert_that(ret, equal_to(expected))
+  def test_fuzzy_interval(self):
+    seed = int(time.time() * 1000)
+    times = 25
+    logging.warning("random seed=%d", seed)
+    random.seed(seed)
+    for _ in range(times):
+      n = int(random.randint(1, 100))
+      data = list(range(n))
+      m = random.randint(1, 1000)
+      interval = m / 1e6
+      now = Timestamp.now()
+      with TestPipeline() as p:
+        ret = (
+            p | PeriodicImpulse(
+                start_timestamp=now, data=data, fire_interval=interval))
+        assert_that(ret, equal_to(data))
 
 
 if __name__ == '__main__':