Resolve floating point precision loss problem. Add tests.

Author: shunping

sdks/python/apache_beam/transforms/periodicsequence.py

@@ -21,7 +21,6 @@
 from typing import Any
 from typing import Optional
 from typing import Sequence
-from typing import Union
 
 import apache_beam as beam
 from apache_beam.io.restriction_trackers import OffsetRange
@@ -41,13 +40,21 @@ class ImpulseSeqGenRestrictionProvider(core.RestrictionProvider):
   def initial_restriction(self, element):
     start, end, interval = element
     if isinstance(start, Timestamp):
-      start = start.micros / 1000000
+      start_micros = start.micros
+    else:
+      start_micros = round(start * 1000000)
+
     if isinstance(end, Timestamp):
-      end = end.micros / 1000000
+      end_micros = end.micros
+    else:
+      end_micros = round(end * 1000000)
+
+    interval_micros = round(interval * 1000000)
 
-    assert start <= end
+    assert start_micros <= end_micros
     assert interval > 0
-    total_outputs = round((end - start) / interval)
+    delta_micros: int = end_micros - start_micros
+    total_outputs = math.ceil(delta_micros / interval_micros)
     return OffsetRange(0, total_outputs)
 
   def create_tracker(self, restriction):
@@ -236,7 +243,7 @@ def _validate_and_adjust_duration(self):
         # 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)
+        self.stop_ts = Timestamp(micros=end * 1000000)
       else:
         end = self.stop_ts.micros / 1000000
     else:

sdks/python/apache_beam/transforms/periodicsequence_test.py

@@ -25,6 +25,8 @@
 import time
 import unittest
 
+from parameterized import parameterized
+
 import apache_beam as beam
 from apache_beam.io.restriction_trackers import OffsetRange
 from apache_beam.testing.test_pipeline import TestPipeline
@@ -159,6 +161,53 @@ def test_processing_time(self):
       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 = (
@@ -176,7 +225,7 @@ def test_repeat(self):
         ret = (
             p | PeriodicImpulse(
                 start_timestamp=now,
-                stop_timestamp=now + 3.0,
+                stop_timestamp=now + 2.6,
                 data=[1, 2, 3, 4],
                 fire_interval=0.5)
             | beam.WindowInto(FixedWindows(0.5))
@@ -212,7 +261,7 @@ def test_not_enough_timestamped_value(self):
 
   def test_fuzzy_interval(self):
     seed = int(time.time() * 1000)
-    times = 25
+    times = 30
     logging.warning("random seed=%d", seed)
     random.seed(seed)
     for _ in range(times):