[Python] Bound and harden the Watch transform

sdks/python/apache_beam/io/watch.py

@@ -48,6 +48,45 @@ def poll(prefix):
                       .with_poll_interval(Duration(seconds=5))
                       .with_termination_per_input(after_total_of(60)))
 
+Watermark and event-time contract
+---------------------------------
+Each emitted output carries the event time the poll function reported for it.
+Raw (non-``TimestampedValue``) outputs default to processing time
+(``Timestamp.now()`` at poll), so wrap outputs in ``TimestampedValue`` or pass
+``timestamp=`` to :meth:`PollResult.incomplete`/:meth:`PollResult.complete` when
+the data has a real event time. The watermark for an input is derived per round
+as: (a) the explicit ``with_watermark`` if the poll supplies one; else (b) the
+earliest event time of this round's new outputs; else (c) held unchanged when a
+round yields nothing; and it is released to ``MAX_TIMESTAMP`` on
+:meth:`PollResult.complete`. The watermark only ever advances.
+
+Policy (b) is safe only for poll functions that enumerate outputs in
+non-decreasing event-time order. If a later round returns a brand-new output
+whose event time is *below* the already-advanced watermark, that output is
+emitted at its true (earlier) time and is therefore late: downstream event-time
+windowing may drop it. Watch logs a throttled warning when this happens. For
+out-of-order sources, have the poll supply an explicit watermark via
+:meth:`PollResult.with_watermark` that bounds the earliest event time any future
+output may have. Dedup is by output value identity (a re-seen value keeps its
+first-seen timestamp), and the output coder must be deterministic for dedup to
+hold across workers and restarts.
+
+Scalability
+-----------
+Parallelism is per input element: each input is one restriction watched on one
+worker, with no intra-key splitting, so throughput scales with the number of
+input elements, not with a single input's growth. The
+per-input dedup set retains one hash per distinct output and is not
+garbage-collected by default, so it grows with the number of distinct outputs an
+input ever produces. For high-cardinality, long-lived inputs, bound it with
+:meth:`Watch.with_dedup_horizon`, which drops outputs more than the horizon
+behind the watermark and collects their dedup state, so memory tracks a sliding
+event-time window. A re-surfaced aged output is suppressed rather than
+re-emitted, so this is safe for full-relisting polls; the trade-off is that a
+genuinely new output arriving more than the horizon behind the watermark is
+dropped as late (see its docstring). The poll function runs synchronously on the
+watch path, so keep it bounded and timeout-safe.
+
 This API is experimental and may change in backwards-incompatible ways.
 """
 
@@ -113,6 +152,11 @@ def is_complete(self) -> bool:
 
   @staticmethod
   def _normalize(outputs, timestamp) -> Tuple[TimestampedValue, ...]:
+    # The default timestamp is computed once per call (not per output) so all
+    # raw outputs in one poll share an event time and the inferred watermark is
+    # not jittered by wall-clock reads. ``timestamp=None`` means "use processing
+    # time"; pass ``timestamp=`` or wrap outputs in ``TimestampedValue`` to set
+    # a real event time.
     if timestamp is None:
       default_ts = Timestamp.now()
     else:
@@ -130,7 +174,10 @@ def incomplete(outputs: Iterable, timestamp=None) -> 'PollResult':
     """Reports outputs and expects more; the transform infers the watermark.
 
     A raw (non-:class:`TimestampedValue`) output is stamped with ``timestamp``
-    when given, else with the current processing time.
+    when given, else with the current processing time. With no explicit
+    watermark, the transform holds the watermark at the earliest event time of
+    this poll's new outputs, which is only safe for non-decreasing event-time
+    enumerations; out-of-order sources should call :meth:`with_watermark`.
     """
     return PollResult(PollResult._normalize(outputs, timestamp), watermark=None)
 
@@ -139,12 +186,16 @@ def complete(outputs: Iterable, timestamp=None) -> 'PollResult':
     """Reports the final outputs for an input, after which polling stops.
 
     A raw (non-:class:`TimestampedValue`) output is stamped with ``timestamp``
-    when given, else with the current processing time.
+    when given, else with the current processing time. The watermark is released
+    to ``MAX_TIMESTAMP`` so downstream event-time windows for this input close.
     """
     return PollResult(
         PollResult._normalize(outputs, timestamp), watermark=MAX_TIMESTAMP)
 
   def with_watermark(self, watermark) -> 'PollResult':
+    """Sets an explicit watermark: a promise that no future output for this
+    input will have an event time below ``watermark``. Use this for sources that
+    can surface outputs out of event-time order across poll rounds."""
     return dataclasses.replace(self, watermark=Timestamp.of(watermark))
 
 
@@ -374,12 +425,14 @@ def __init__(
       poll_fn: Callable[[Any], PollResult],
       key_coder: Coder,
       termination: TerminationCondition,
-      now_fn: Callable[[], float]):
+      now_fn: Callable[[], float],
+      dedup_horizon: Optional[Duration] = None):
     self._restriction = restriction
     self._poll_fn = poll_fn
     self._key_coder = key_coder
     self._termination = termination
     self._now = now_fn
+    self._dedup_horizon = dedup_horizon
     self._should_stop = False
     self._primary = None  # type: Optional[_GrowthState]
     self._residual = None  # type: Optional[_GrowthState]
@@ -409,10 +462,23 @@ def try_claim(self, holder: list) -> bool:
     now = Timestamp.of(self._now())
     result = self._poll_fn(element)
 
+    # An output more than the dedup horizon behind the committed watermark is
+    # past its dedup lifetime: its hash has been (or is about to be) evicted, so
+    # it is dropped as late instead of re-emitted. Without this, a re-listing
+    # poll that re-surfaces an aged output every round would resurrect the
+    # evicted hash, re-emit a duplicate, and pin the inferred watermark to the
+    # output's old event time. The suppression cutoff equals the prior round's
+    # eviction cutoff, so an evicted output is suppressed on its next appearance
+    # with no resurrection gap. None when no horizon or no committed watermark.
+    suppress_cutoff = self._suppress_cutoff(restriction.poll_watermark)
     new_outputs = []  # type: List[TimestampedValue]
     claimed = []  # type: List[Tuple[bytes, Timestamp]]
     seen_this_round = set()  # type: set
+    suppressed = 0
     for output in result.outputs:
+      if suppress_cutoff is not None and output.timestamp < suppress_cutoff:
+        suppressed += 1
+        continue
       key_hash = self._hash_output(output.value)
       if key_hash in restriction.completed or key_hash in seen_this_round:
         continue
@@ -447,17 +513,26 @@ def try_claim(self, holder: list) -> bool:
       # initiated or via defer_remainder) resumes a state that emits nothing.
       self._residual = _NonPollingGrowthState(PollResult((), watermark))
     else:
-      merged = collections.OrderedDict(restriction.completed)
-      for key_hash, first_seen in claimed:
-        merged[key_hash] = first_seen
       residual_watermark = self._max_watermark(
           restriction.poll_watermark, watermark)
+      completed = self._build_completed(
+          restriction.completed, claimed, residual_watermark)
       self._residual = _PollingGrowthState(
-          merged, residual_watermark, termination_state)
-    holder[1] = ('poll', new_outputs, watermark, stop)
+          completed, residual_watermark, termination_state)
+    holder[1] = ('poll', new_outputs, watermark, stop, suppressed)
     self._should_stop = True
     return True
 
+  def _suppress_cutoff(self,
+                       watermark: Optional[Timestamp]) -> Optional[Timestamp]:
+    # Event-time floor below which an output is past its dedup horizon. Derived
+    # from the committed (prior-round) watermark, the one value known before
+    # this round's outputs are inspected, so it matches the prior round's
+    # eviction cutoff.
+    if self._dedup_horizon is None or watermark is None:
+      return None
+    return watermark - self._dedup_horizon
+
   @staticmethod
   def _max_watermark(left: Optional[Timestamp],
                      right: Optional[Timestamp]) -> Optional[Timestamp]:
@@ -467,6 +542,43 @@ def _max_watermark(left: Optional[Timestamp],
       return left
     return max(left, right)
 
+  def _build_completed(
+      self,
+      base: 'collections.OrderedDict[bytes, Timestamp]',
+      claimed: List[Tuple[bytes, Timestamp]],
+      watermark: Optional[Timestamp],
+  ) -> 'collections.OrderedDict[bytes, Timestamp]':
+    """Builds the residual dedup map.
+
+    Reuses the parent map untouched when a round adds nothing and no dedup
+    horizon is configured, so a steady-state idle round avoids the O(N) copy.
+    Otherwise copies, appends this round's new hashes, and (when a horizon is
+    set) evicts entries whose event time has fallen behind the watermark.
+    """
+    horizon = self._dedup_horizon
+    if not claimed and horizon is None:
+      return base
+    merged = collections.OrderedDict(base)
+    for key_hash, first_seen in claimed:
+      merged[key_hash] = first_seen
+    if horizon is not None and watermark is not None:
+      self._evict_stale(merged, watermark - horizon)
+    return merged
+
+  @staticmethod
+  def _evict_stale(
+      completed: 'collections.OrderedDict[bytes, Timestamp]',
+      cutoff: Timestamp) -> None:
+    # Drop every entry first seen at an event time below the cutoff, regardless
+    # of insertion order. Insertion order is not event-time order in general:
+    # with an explicit held-back watermark a later round can add an
+    # earlier-timestamped output after a future one, so a prefix scan would miss
+    # genuinely-stale entries and fail to bound the state. This full scan is
+    # O(N) in the (horizon-bounded) map size, paid only when a horizon is set.
+    stale = [h for h, first_seen in completed.items() if first_seen < cutoff]
+    for h in stale:
+      del completed[h]
+
   def try_split(self, fraction_of_remainder):
     # Only self-checkpoint (fraction 0) is supported; decline dynamic splits.
     if fraction_of_remainder != 0:
@@ -521,14 +633,20 @@ def __init__(
       termination: TerminationCondition,
       poll_interval: Duration,
       output_coder: Coder,
-      now_fn: Optional[Callable[[], float]] = None):
+      now_fn: Optional[Callable[[], float]] = None,
+      dedup_horizon: Optional[Duration] = None):
     self._poll_fn = poll_fn
     self._termination = termination
     self._poll_interval = poll_interval
     self._output_coder = output_coder
     self._key_coder = output_coder
     self._now = now_fn or time.time
+    self._dedup_horizon = dedup_horizon
     self._restriction_coder = _GrowthStateCoder(output_coder, termination)
+    # Counts of late emissions and horizon-suppressed outputs seen on this
+    # worker, for throttled warnings.
+    self._late_count = 0
+    self._suppressed_count = 0
 
   def initial_restriction(self, element) -> _PollingGrowthState:
     now = Timestamp.of(self._now())
@@ -543,7 +661,8 @@ def create_tracker(self, restriction) -> _GrowthRestrictionTracker:
         self._poll_fn,
         self._key_coder,
         self._termination,
-        self._now)
+        self._now,
+        self._dedup_horizon)
 
   def split(self, element, restriction):
     # Watch fans out by input element, so each restriction stays whole.
@@ -581,8 +700,18 @@ def process(
       for output in work[1].outputs:
         yield TimestampedValue((element, output.value), output.timestamp)
       return
-    new_outputs, watermark, stop = work[1], work[2], work[3]
+    new_outputs, watermark, stop, suppressed = (
+        work[1], work[2], work[3], work[4])
+    if suppressed:
+      self._warn_suppressed(element, suppressed)
+    # Emit outputs first, then advance the watermark (emit-then-advance), so a
+    # round that reports both outputs and a higher watermark cannot push the
+    # watermark past those outputs' event times. An output behind the current
+    # watermark is already late; warn rather than silently dropping it.
+    current_watermark = watermark_estimator.current_watermark()
     for output in new_outputs:
+      if current_watermark is not None and output.timestamp < current_watermark:
+        self._warn_late(element, output.timestamp, current_watermark)
       yield TimestampedValue((element, output.value), output.timestamp)
     if stop:
       # The input is finished, so release the watermark hold to MAX.
@@ -592,6 +721,39 @@ def process(
       _set_watermark_if_greater(watermark_estimator, watermark)
     tracker.defer_remainder(self._poll_interval)
 
+  def _warn_late(self, element, output_timestamp, watermark) -> None:
+    # Throttle to powers of two so an ongoing problem stays visible without
+    # flooding worker logs.
+    self._late_count += 1
+    if self._late_count & (self._late_count - 1) == 0:
+      _LOGGER.warning(
+          'Watch emitted an output for input %r at event time %s, behind the '
+          'current watermark %s, so downstream event-time windowing may drop '
+          'it as late. This happens when the poll returns outputs out of '
+          'event-time order across rounds without an explicit watermark; call '
+          'PollResult.with_watermark(...) for out-of-order sources. '
+          '(%d late emissions so far on this worker)',
+          element,
+          output_timestamp,
+          watermark,
+          self._late_count)
+
+  def _warn_suppressed(self, element, count) -> None:
+    # Throttle to powers of two so an ongoing drop stays visible without
+    # flooding worker logs.
+    self._suppressed_count += count
+    if self._suppressed_count & (self._suppressed_count - 1) == 0:
+      _LOGGER.warning(
+          'Watch suppressed %d output(s) for input %r more than the dedup '
+          'horizon behind the watermark, dropping them as late. For a '
+          're-listing poll this is expected for already-emitted outputs; for a '
+          'source that surfaces genuinely new late data, widen '
+          'with_dedup_horizon(...) or supply PollResult.with_watermark(...). '
+          '(%d suppressed so far on this worker)',
+          count,
+          element,
+          self._suppressed_count)
+
 
 def _set_watermark_if_greater(watermark_estimator, new_watermark) -> None:
   # set_watermark raises on regression, so only ever advance the watermark.
@@ -617,13 +779,15 @@ def __init__(
       termination: Optional[TerminationCondition] = None,
       poll_interval: Optional[Duration] = None,
       output_coder: Optional[Coder] = None,
-      now_fn: Optional[Callable[[], float]] = None):
+      now_fn: Optional[Callable[[], float]] = None,
+      dedup_horizon: Optional[Duration] = None):
     super().__init__()
     self._poll_fn = poll_fn
     self._termination = termination or never()
     self._poll_interval = poll_interval
     self._output_coder = output_coder
     self._now = now_fn
+    self._dedup_horizon = dedup_horizon
 
   @classmethod
   def growth_of(cls, poll_fn: Callable[[Any], PollResult]) -> 'Watch':
@@ -635,7 +799,8 @@ def _replace(self, **changes) -> 'Watch':
         termination=self._termination,
         poll_interval=self._poll_interval,
         output_coder=self._output_coder,
-        now_fn=self._now)
+        now_fn=self._now,
+        dedup_horizon=self._dedup_horizon)
     spec.update(changes)
     return Watch(**spec)
 
@@ -649,6 +814,40 @@ def with_termination_per_input(
   def with_output_coder(self, output_coder: Coder) -> 'Watch':
     return self._replace(output_coder=output_coder)
 
+  def with_dedup_horizon(self, horizon) -> 'Watch':
+    """Bounds the per-input dedup state by event time.
+
+    By default the dedup set is exact and unbounded: it retains one hash per
+    distinct output forever, so its size grows with the number of distinct
+    outputs an input produces. With a horizon set, Watch drops every output
+    whose event time is more than ``horizon`` (a :class:`Duration` or seconds)
+    behind the input's watermark and collects its dedup entry, so the set tracks
+    a sliding event-time window instead of the full history. This bounds memory,
+    the per-round copy, and the per-checkpoint encode to the horizon window.
+
+    An aged-out output that the poll re-surfaces at its original (now
+    below-horizon) event time is suppressed, not re-emitted: the suppression
+    cutoff equals the prior round's eviction cutoff, so a re-listing poll that
+    returns the whole collection every round stays exactly-once and its
+    watermark keeps advancing. The trade-off is that a genuinely new output
+    arriving more than ``horizon`` behind the watermark is also dropped as late;
+    Watch logs a throttled warning whenever it suppresses an output. For
+    out-of-order sources that must keep such outputs, widen the horizon or
+    supply an explicit :meth:`PollResult.with_watermark`.
+
+    Residual hazard: an aged-out value that the poll re-returns at a fresh,
+    at-or-after-watermark timestamp is on time, so it is not suppressed and is
+    re-emitted as a duplicate. Cursor or time-windowed polls that never re-emit
+    an aged-out value at a fresh timestamp avoid this.
+
+    ``horizon`` must be non-negative.
+    """
+    horizon = _as_duration(horizon)
+    if horizon < Duration(0):
+      raise ValueError(
+          'dedup horizon must be non-negative, got %s' % (horizon, ))
+    return self._replace(dedup_horizon=horizon)
+
   def expand(self, pcoll):
     if self._poll_interval is None:
       raise ValueError('Watch requires with_poll_interval(...)')
@@ -669,7 +868,8 @@ def expand(self, pcoll):
             self._termination,
             self._poll_interval,
             output_coder,
-            self._now))
+            self._now,
+            self._dedup_horizon))
 
 
 def _as_duration(value) -> Duration: