gemini-code-assist[bot] commented on code in PR #39090:
URL: https://github.com/apache/beam/pull/39090#discussion_r3524080603


##########
sdks/python/apache_beam/io/watch.py:
##########
@@ -0,0 +1,838 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Experimental ``Watch`` transform for the Python SDK.
+
+``Watch`` continuously watches a growing set of outputs for each input element,
+calling a user poll function on an interval until a per-input termination
+condition fires. It is the engine behind periodic file-discovery and any
+periodic polling source.
+
+For every input element the transform runs an independent loop::
+
+    poll -> keep never-seen-before outputs -> emit them (timestamped) ->
+    update watermark -> check termination -> wait(poll_interval) -> poll -> ...
+
+The output is an unbounded ``PCollection`` of ``(input, output)`` pairs. Each
+output carries the event time the poll function first reported it. By default
+dedup is by value identity, using a stable 128-bit hash of the encoded output,
+so the output coder must be deterministic for dedup to hold across workers and
+restarts. :meth:`Watch.with_timestamp_cursor` switches to bounded-state dedup 
by
+event time (see Scalability below).
+
+Example::
+
+    from apache_beam.io.watch import Watch, PollResult, after_total_of
+    from apache_beam.transforms.window import TimestampedValue
+    from apache_beam.utils.timestamp import Duration, Timestamp
+
+    def poll(prefix):
+      now = Timestamp.now()
+      outputs = [TimestampedValue(prefix + str(i), now) for i in range(3)]
+      return PollResult.complete(outputs)
+
+    watched = (inputs
+               | Watch.growth_of(poll)
+                      .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. Default 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. By default dedup is exact and
+by value identity: the per-input state retains one hash per distinct output and
+is not garbage-collected, so it grows with the number of distinct outputs an
+input ever produces. For a long-lived, high-cardinality source whose outputs
+carry strictly increasing event-time timestamps, call
+:meth:`Watch.with_timestamp_cursor` to dedup by a high-water-mark timestamp
+instead: the per-input state is a single timestamp that never grows and the 
poll
+result is not hashed, so state and per-checkpoint encoding stay O(1) regardless
+of how many outputs the input produces. See its docstring for the ordering
+precondition. 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.
+"""
+
+import collections
+import dataclasses
+import hashlib
+import logging
+import time
+from typing import Any
+from typing import Callable
+from typing import Iterable
+from typing import List
+from typing import Optional
+from typing import Tuple
+
+from apache_beam import coders
+from apache_beam.coders.coders import Coder
+from apache_beam.coders.coders import NullableCoder
+from apache_beam.coders.coders import TimestampCoder
+from apache_beam.coders.coders import TupleCoder
+from apache_beam.io import iobase
+from apache_beam.io.watermark_estimators import ManualWatermarkEstimator
+from apache_beam.runners import sdf_utils
+from apache_beam.transforms import PTransform
+from apache_beam.transforms import core
+from apache_beam.transforms.window import TimestampedValue
+from apache_beam.utils.timestamp import MAX_TIMESTAMP
+from apache_beam.utils.timestamp import Duration
+from apache_beam.utils.timestamp import Timestamp
+
+__all__ = [
+    'Watch',
+    'PollResult',
+    'PollFn',
+    'TerminationCondition',
+    'never',
+    'after_total_of',
+]
+
+_LOGGER = logging.getLogger(__name__)
+
+_HASH_DIGEST_SIZE = 16  # 128-bit digest width.
+
+# 
------------------------------------------------------------------------------
+# Public API.
+# 
------------------------------------------------------------------------------
+
+
[email protected](frozen=True)
+class PollResult:
+  """Outputs produced by one poll, plus an optional explicit watermark.
+
+  ``watermark`` of ``None`` lets the transform infer the watermark from the
+  earliest new output. A watermark of ``MAX_TIMESTAMP`` (set by
+  :meth:`complete`) marks the input finished, so polling stops.
+  """
+  outputs: Tuple[TimestampedValue, ...]
+  watermark: Optional[Timestamp] = None
+
+  @property
+  def is_complete(self) -> bool:
+    return self.watermark == MAX_TIMESTAMP
+
+  @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:
+      default_ts = Timestamp.of(timestamp)
+    normalized = []
+    for output in outputs:
+      if isinstance(output, TimestampedValue):
+        normalized.append(output)
+      else:
+        normalized.append(TimestampedValue(output, default_ts))

Review Comment:
   ![medium](https://www.gstatic.com/codereviewagent/medium-priority.svg)
   
   When iterating over `outputs`, if the user passes a single-use generator or 
iterator, iterating over it to normalize elements will consume it. Since 
`outputs` is iterated again when constructing the `PollResult` or during 
normalization, this can lead to empty outputs or unexpected behavior. Consider 
converting `outputs` to a list or tuple first before normalization.
   
   ```suggestion
       outputs = list(outputs)
       normalized = []
       for output in outputs:
         if isinstance(output, TimestampedValue):
           normalized.append(output)
         else:
           normalized.append(TimestampedValue(output, default_ts))
   ```



##########
sdks/python/apache_beam/io/watch.py:
##########
@@ -0,0 +1,838 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Experimental ``Watch`` transform for the Python SDK.
+
+``Watch`` continuously watches a growing set of outputs for each input element,
+calling a user poll function on an interval until a per-input termination
+condition fires. It is the engine behind periodic file-discovery and any
+periodic polling source.
+
+For every input element the transform runs an independent loop::
+
+    poll -> keep never-seen-before outputs -> emit them (timestamped) ->
+    update watermark -> check termination -> wait(poll_interval) -> poll -> ...
+
+The output is an unbounded ``PCollection`` of ``(input, output)`` pairs. Each
+output carries the event time the poll function first reported it. By default
+dedup is by value identity, using a stable 128-bit hash of the encoded output,
+so the output coder must be deterministic for dedup to hold across workers and
+restarts. :meth:`Watch.with_timestamp_cursor` switches to bounded-state dedup 
by
+event time (see Scalability below).
+
+Example::
+
+    from apache_beam.io.watch import Watch, PollResult, after_total_of
+    from apache_beam.transforms.window import TimestampedValue
+    from apache_beam.utils.timestamp import Duration, Timestamp
+
+    def poll(prefix):
+      now = Timestamp.now()
+      outputs = [TimestampedValue(prefix + str(i), now) for i in range(3)]
+      return PollResult.complete(outputs)
+
+    watched = (inputs
+               | Watch.growth_of(poll)
+                      .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. Default 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. By default dedup is exact and
+by value identity: the per-input state retains one hash per distinct output and
+is not garbage-collected, so it grows with the number of distinct outputs an
+input ever produces. For a long-lived, high-cardinality source whose outputs
+carry strictly increasing event-time timestamps, call
+:meth:`Watch.with_timestamp_cursor` to dedup by a high-water-mark timestamp
+instead: the per-input state is a single timestamp that never grows and the 
poll
+result is not hashed, so state and per-checkpoint encoding stay O(1) regardless
+of how many outputs the input produces. See its docstring for the ordering
+precondition. 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.
+"""
+
+import collections
+import dataclasses
+import hashlib
+import logging
+import time
+from typing import Any
+from typing import Callable
+from typing import Iterable
+from typing import List
+from typing import Optional
+from typing import Tuple
+
+from apache_beam import coders
+from apache_beam.coders.coders import Coder
+from apache_beam.coders.coders import NullableCoder
+from apache_beam.coders.coders import TimestampCoder
+from apache_beam.coders.coders import TupleCoder
+from apache_beam.io import iobase
+from apache_beam.io.watermark_estimators import ManualWatermarkEstimator
+from apache_beam.runners import sdf_utils
+from apache_beam.transforms import PTransform
+from apache_beam.transforms import core
+from apache_beam.transforms.window import TimestampedValue
+from apache_beam.utils.timestamp import MAX_TIMESTAMP
+from apache_beam.utils.timestamp import Duration
+from apache_beam.utils.timestamp import Timestamp
+
+__all__ = [
+    'Watch',
+    'PollResult',
+    'PollFn',
+    'TerminationCondition',
+    'never',
+    'after_total_of',
+]
+
+_LOGGER = logging.getLogger(__name__)
+
+_HASH_DIGEST_SIZE = 16  # 128-bit digest width.
+
+# 
------------------------------------------------------------------------------
+# Public API.
+# 
------------------------------------------------------------------------------
+
+
[email protected](frozen=True)
+class PollResult:
+  """Outputs produced by one poll, plus an optional explicit watermark.
+
+  ``watermark`` of ``None`` lets the transform infer the watermark from the
+  earliest new output. A watermark of ``MAX_TIMESTAMP`` (set by
+  :meth:`complete`) marks the input finished, so polling stops.
+  """
+  outputs: Tuple[TimestampedValue, ...]
+  watermark: Optional[Timestamp] = None
+
+  @property
+  def is_complete(self) -> bool:
+    return self.watermark == MAX_TIMESTAMP
+
+  @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:
+      default_ts = Timestamp.of(timestamp)
+    normalized = []
+    for output in outputs:
+      if isinstance(output, TimestampedValue):
+        normalized.append(output)
+      else:
+        normalized.append(TimestampedValue(output, default_ts))
+    return tuple(normalized)
+
+  @staticmethod
+  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. 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)
+
+  @staticmethod
+  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. 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))
+
+
+class PollFn(object):
+  """Optional base for a poll function ``input -> PollResult``.
+
+  Any callable with that signature works; subclass only to attach an output
+  coder hint via :meth:`default_output_coder`.
+  """
+  def __call__(self, element: Any) -> PollResult:
+    raise NotImplementedError
+
+  def default_output_coder(self) -> Optional[Coder]:
+    return None
+
+
+class TerminationCondition(object):
+  """Per-input stop policy with immutable, encodable state.
+
+  Hooks follow the lifecycle of one input's polling loop. ``state`` flows from
+  :meth:`for_new_input` through the per-round hooks and is serialized with
+  :meth:`state_coder`.
+  """
+  def for_new_input(self, now: Timestamp, element: Any) -> Any:
+    raise NotImplementedError
+
+  def on_seen_new_output(self, now: Timestamp, state: Any) -> Any:
+    return state
+
+  def on_poll_complete(self, state: Any) -> Any:
+    return state
+
+  def can_stop_polling(self, now: Timestamp, state: Any) -> bool:
+    raise NotImplementedError
+
+  def state_coder(self) -> Coder:
+    raise NotImplementedError
+
+
+class _Never(TerminationCondition):
+  """Polls until the poll function returns :meth:`PollResult.complete`."""
+  def for_new_input(self, now, element):
+    return 0
+
+  def can_stop_polling(self, now, state):
+    return False
+
+  def state_coder(self):
+    return coders.VarIntCoder()
+
+
+class _AfterTotalOf(TerminationCondition):
+  """Stops once the wall-clock time since the input was first seen exceeds a
+  fixed duration."""
+  def __init__(self, duration: Duration):
+    self._duration_micros = duration.micros
+
+  def for_new_input(self, now, element):
+    return (now, self._duration_micros)
+
+  def can_stop_polling(self, now, state):
+    start, duration_micros = state
+    return (now - start).micros > duration_micros
+
+  def state_coder(self):
+    return TupleCoder([TimestampCoder(), coders.VarIntCoder()])
+
+
+def never() -> TerminationCondition:
+  """Polls until :meth:`PollResult.complete`."""
+  return _Never()
+
+
+def after_total_of(duration) -> TerminationCondition:
+  """Stops polling an input after ``duration`` (a :class:`Duration` or seconds)
+  has elapsed since it was first seen."""
+  return _AfterTotalOf(_as_duration(duration))
+
+
+# 
------------------------------------------------------------------------------
+# Restriction state.
+# 
------------------------------------------------------------------------------
+
+
[email protected](frozen=True)
+class _PollingGrowthState:
+  """Keep-polling state: dedup state, watermark, termination state.
+
+  In the default (hash) mode ``completed`` maps a 16-byte output hash to the
+  event time it was first seen; it is insertion-ordered and treated as
+  immutable. In timestamp-cursor mode ``completed`` is empty and ``cursor``
+  holds the greatest event time emitted so far, so the state is O(1).
+  """
+  completed: 'collections.OrderedDict[bytes, Timestamp]'
+  poll_watermark: Optional[Timestamp]
+  termination_state: Any
+  cursor: Optional[Timestamp] = None
+
+
[email protected](frozen=True)
+class _NonPollingGrowthState:
+  """Replay-then-stop state: the outputs already emitted this round.
+
+  Produced as the checkpoint primary so a bundle retry re-emits exactly those
+  outputs.
+  """
+  pending: PollResult
+
+
+_GrowthState = Any  # Union[_PollingGrowthState, _NonPollingGrowthState]
+
+# 
------------------------------------------------------------------------------
+# Coders.
+# 
------------------------------------------------------------------------------
+
+
+class _HashCode128Coder(Coder):
+  """Fixed-width coder for a 16-byte output hash.
+
+  Encodes and decodes exactly 16 bytes and raises on any other length, so a
+  corrupt restriction surfaces at decode time.
+  """
+  def encode(self, value: bytes) -> bytes:
+    if len(value) != _HASH_DIGEST_SIZE:
+      raise ValueError(
+          'hash must be %d bytes, got %d' % (_HASH_DIGEST_SIZE, len(value)))
+    return value
+
+  def decode(self, encoded: bytes) -> bytes:
+    if len(encoded) != _HASH_DIGEST_SIZE:
+      raise ValueError(
+          'hash must be %d bytes, got %d' % (_HASH_DIGEST_SIZE, len(encoded)))
+    return encoded
+
+  def is_deterministic(self) -> bool:
+    return True
+
+
+class _TimestampedValueCoder(Coder):
+  """Coder for :class:`TimestampedValue`.
+
+  The Python SDK ships no coder for ``TimestampedValue``, so this encodes the
+  ``(value, timestamp)`` pair with a :class:`TupleCoder` and rebuilds the
+  ``TimestampedValue`` on decode.
+  """
+  def __init__(self, value_coder: Coder):
+    self._tuple_coder = TupleCoder([value_coder, TimestampCoder()])
+
+  def encode(self, value: TimestampedValue) -> bytes:
+    return self._tuple_coder.encode((value.value, value.timestamp))
+
+  def decode(self, encoded: bytes) -> TimestampedValue:
+    value, timestamp = self._tuple_coder.decode(encoded)
+    return TimestampedValue(value, timestamp)
+
+  def is_deterministic(self) -> bool:
+    return self._tuple_coder.is_deterministic()
+
+
+class _GrowthStateCoder(Coder):
+  """Encodes a :class:`_PollingGrowthState` or :class:`_NonPollingGrowthState`.
+
+  A ``(tag, payload)`` envelope selects the variant; the payload is a
+  variant-specific :class:`TupleCoder`. ``completed`` is encoded as an ordered
+  list of ``(hash, timestamp)`` pairs so insertion order survives a round trip,
+  and the cursor is a nullable timestamp. This format is internal to the Python
+  SDK.
+  """
+  def __init__(self, output_coder: Coder, termination: TerminationCondition):
+    nullable_ts = NullableCoder(TimestampCoder())
+    self._envelope_coder = TupleCoder(
+        [coders.VarIntCoder(), coders.BytesCoder()])
+    self._polling_coder = TupleCoder([
+        termination.state_coder(),
+        nullable_ts,
+        coders.ListCoder(TupleCoder([_HashCode128Coder(), TimestampCoder()])),
+        nullable_ts,
+    ])
+    self._non_polling_coder = TupleCoder([
+        nullable_ts,
+        coders.ListCoder(_TimestampedValueCoder(output_coder)),
+    ])
+
+  def encode(self, state: _GrowthState) -> bytes:
+    if isinstance(state, _PollingGrowthState):
+      payload = self._polling_coder.encode((
+          state.termination_state,
+          state.poll_watermark,
+          list(state.completed.items()),
+          state.cursor))
+      return self._envelope_coder.encode((0, payload))
+    payload = self._non_polling_coder.encode(
+        (state.pending.watermark, list(state.pending.outputs)))
+    return self._envelope_coder.encode((1, payload))
+
+  def decode(self, encoded: bytes) -> _GrowthState:
+    tag, payload = self._envelope_coder.decode(encoded)
+    if tag == 0:
+      termination_state, poll_watermark, items, cursor = (
+          self._polling_coder.decode(payload))
+      return _PollingGrowthState(
+          collections.OrderedDict(items),
+          poll_watermark,
+          termination_state,
+          cursor)
+    if tag == 1:
+      watermark, outputs = self._non_polling_coder.decode(payload)
+      return _NonPollingGrowthState(PollResult(tuple(outputs), watermark))
+    raise ValueError('unknown Watch growth state tag: %r' % (tag, ))
+
+  def is_deterministic(self) -> bool:
+    return False
+
+
+# 
------------------------------------------------------------------------------
+# Restriction tracker.
+# 
------------------------------------------------------------------------------
+
+
+class _GrowthRestrictionTracker(iobase.RestrictionTracker):
+  """Drives one input's polling loop.
+
+  ``process()`` only sees a ``RestrictionTrackerView`` whose ``try_claim``
+  returns a bool, so the poll happens inside ``try_claim`` and its result is
+  returned through a two-slot holder list passed as the claim position:
+  ``holder[0]`` carries the input element in, ``holder[1]`` carries the work
+  out. At most one claim succeeds per ``process()``.
+
+  The poll runs while the tracker lock is held, so a ``PollFn`` must be bounded
+  or timeout-safe; a blocking poll delays runner-initiated checkpoints.
+  """
+  def __init__(
+      self,
+      restriction: _GrowthState,
+      poll_fn: Callable[[Any], PollResult],
+      key_coder: Coder,
+      termination: TerminationCondition,
+      now_fn: Callable[[], float],
+      cursor_mode: bool = False):
+    self._restriction = restriction
+    self._poll_fn = poll_fn
+    self._key_coder = key_coder
+    self._termination = termination
+    self._now = now_fn
+    self._cursor_mode = cursor_mode
+    self._should_stop = False
+    self._primary = None  # type: Optional[_GrowthState]
+    self._residual = None  # type: Optional[_GrowthState]
+
+  def current_restriction(self) -> _GrowthState:
+    return self._restriction
+
+  def _hash_output(self, value: Any) -> bytes:
+    return hashlib.blake2b(
+        self._key_coder.encode(value), digest_size=_HASH_DIGEST_SIZE).digest()
+
+  def try_claim(self, holder: list) -> bool:
+    """Performs one poll round (or one replay) and reports it via ``holder``.
+
+    Returns ``False`` only when a checkpoint already stopped this invocation,
+    in which case ``process()`` must emit nothing.
+    """
+    if self._should_stop:
+      return False
+    restriction = self._restriction
+    if isinstance(restriction, _NonPollingGrowthState):
+      holder[1] = ('replay', restriction.pending)
+      self._should_stop = True
+      return True
+
+    element = holder[0]
+    now = Timestamp.of(self._now())
+    result = self._poll_fn(element)
+
+    claimed = []  # type: List[Tuple[bytes, Timestamp]]
+    if self._cursor_mode:
+      # Dedup by a high-water-mark timestamp: keep only outputs strictly past
+      # the cursor, so the state is one timestamp and the poll is not hashed.
+      cursor = restriction.cursor
+      new_outputs = [
+          output for output in result.outputs
+          if cursor is None or output.timestamp > cursor
+      ]

Review Comment:
   ![high](https://www.gstatic.com/codereviewagent/high-priority.svg)
   
   In `_GrowthRestrictionTracker.try_claim`, when `self._cursor_mode` is 
enabled, the code compares `output.timestamp > cursor`. However, 
`output.timestamp` can be a raw float or integer if the user directly 
instantiated `TimestampedValue` without using `Timestamp` objects (as 
`TimestampedValue` accepts `int` or `float` for its timestamp). Since `cursor` 
is a `Timestamp` object (or `None`), comparing a `Timestamp` with a float or 
integer will raise a `TypeError` in Python 3. To prevent comparison errors, 
normalize `output.timestamp` to a `Timestamp` object using 
`Timestamp.of(output.timestamp)` before comparison.
   
   ```suggestion
         cursor = restriction.cursor
         new_outputs = [
             output for output in result.outputs
             if cursor is None or Timestamp.of(output.timestamp) > cursor
         ]
   ```



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]

Reply via email to