[Python] Refactor MatchContinuously onto the Watch transform

sdks/python/apache_beam/io/fileio.py

@@ -93,28 +93,34 @@
 import random
 import uuid
 from collections import namedtuple
-from functools import partial
 from typing import Any
 from typing import BinaryIO  # pylint: disable=unused-import
 from typing import Callable
 from typing import Iterable
+from typing import Optional
 from typing import Union
 
 import apache_beam as beam
+from apache_beam.coders.coders import VarIntCoder
 from apache_beam.io import filesystem
 from apache_beam.io import filesystems
 from apache_beam.io.filesystem import BeamIOError
 from apache_beam.io.filesystem import CompressionTypes
+from apache_beam.io.watch import PollFn
+from apache_beam.io.watch import PollResult
+from apache_beam.io.watch import TerminationCondition
+from apache_beam.io.watch import Watch
+from apache_beam.io.watch import never
 from apache_beam.options.pipeline_options import GoogleCloudOptions
 from apache_beam.options.value_provider import StaticValueProvider
 from apache_beam.options.value_provider import ValueProvider
 from apache_beam.transforms.periodicsequence import PeriodicImpulse
-from apache_beam.transforms.userstate import CombiningValueStateSpec
 from apache_beam.transforms.window import BoundedWindow
 from apache_beam.transforms.window import FixedWindows
 from apache_beam.transforms.window import GlobalWindow
 from apache_beam.transforms.window import IntervalWindow
 from apache_beam.utils.timestamp import MAX_TIMESTAMP
+from apache_beam.utils.timestamp import Duration
 from apache_beam.utils.timestamp import Timestamp
 
 __all__ = [
@@ -251,6 +257,83 @@ def process(
     yield ReadableFile(metadata, self._compression)
 
 
+class _PollClock(object):
+  """Shares one clock reading per poll round, so the start gate and the poll
+  budget judge the ``start_timestamp`` boundary consistently."""
+  def __init__(self):
+    self.last_poll_micros: Optional[int] = None
+
+
+class _WatchWindowTermination(TerminationCondition):
+  """Stops after the polls that fall in the ``[start, stop)`` window.
+
+  ``max_polls`` is the ``PeriodicImpulse`` tick count
+  ``ceil((stop - start) / interval)``; polls before ``start`` are waiting
+  rounds and do not consume the budget.
+  """
+  def __init__(self, clock: _PollClock, start_micros: int, max_polls: int):
+    self._clock = clock
+    self._start_micros = start_micros
+    self._max_polls = max_polls
+
+  def for_new_input(self, now, element):
+    return 0
+
+  def on_poll_complete(self, state):
+    poll_micros = self._clock.last_poll_micros
+    if poll_micros is not None and poll_micros >= self._start_micros:
+      return state + 1
+    return state
+
+  def can_stop_polling(self, now, state):
+    return state >= self._max_polls
+
+  def state_coder(self):
+    return VarIntCoder()
+
+
+def _file_path_key(metadata: filesystem.FileMetadata) -> str:
+  return metadata.path
+
+
+def _file_path_and_mtime_key(
+    metadata: filesystem.FileMetadata) -> tuple[str, float]:
+  # Keying on the last-modified time makes a changed file look new again. A
+  # missing (zero) timestamp is rejected because updates could never be seen.
+  if not metadata.last_updated_in_seconds:
+    raise BeamIOError(
+        'MatchContinuously(match_updated_files=True) requires file '
+        'last-modified times, but %s reports none.' % metadata.path)
+  return metadata.path, metadata.last_updated_in_seconds
+
+
+class _MatchContinuouslyPollFn(PollFn):
+  """Polls a file pattern, honoring empty-match rules.
+
+  A poll before ``start_timestamp`` emits nothing. Matches carry the poll time
+  as their event time, and the watermark advances to the poll time so
+  event-time windows progress even when nothing new matches.
+  """
+  def __init__(self, empty_match_treatment, start_timestamp, clock=None):
+    self._empty_match_treatment = empty_match_treatment
+    self._start_micros = Timestamp.of(start_timestamp).micros
+    self._clock = clock if clock is not None else _PollClock()
+
+  def __call__(self, file_pattern: str) -> PollResult[filesystem.FileMetadata]:
+    now = Timestamp.now()
+    self._clock.last_poll_micros = now.micros
+    if now.micros < self._start_micros:
+      return PollResult.incomplete(())
+    match_result = filesystems.FileSystems.match([file_pattern])[0]
+    if (not match_result.metadata_list and
+        not EmptyMatchTreatment.allow_empty_match(file_pattern,
+                                                  self._empty_match_treatment)):
+      raise BeamIOError(
+          'Empty match for pattern %s. Disallowed.' % file_pattern)
+    return PollResult.incomplete(
+        match_result.metadata_list, timestamp=now).with_watermark(now)
+
+
 class MatchContinuously(beam.PTransform):
   """Checks for new files for a given pattern every interval.
 
@@ -261,10 +344,11 @@ class MatchContinuously(beam.PTransform):
   guarantees.
 
   Matching continuously scales poorly, as it is stateful, and requires storing
-  file ids in memory. In addition, because it is memory-only, if a pipeline is
-  restarted, already processed files will be reprocessed. Consider an alternate
-  technique, such as Pub/Sub Notifications
-  (https://cloud.google.com/storage/docs/pubsub-notifications)
+  file ids for every file the pattern has matched. With ``has_deduplication``
+  enabled those ids are kept in the splittable DoFn restriction, so a runner
+  with checkpointing enabled restores them after a restart and does not
+  reprocess files. Consider an alternate technique, such as Pub/Sub
+  Notifications (https://cloud.google.com/storage/docs/pubsub-notifications)
   when using GCS if possible.
   """
   def __init__(
@@ -306,37 +390,76 @@ def __init__(
         'if possible')
 
   def expand(self, pbegin) -> beam.PCollection[filesystem.FileMetadata]:
-    # invoke periodic impulse
-    impulse = pbegin | PeriodicImpulse(
-        start_timestamp=self.start_ts,
-        stop_timestamp=self.stop_ts,
-        fire_interval=self.interval)
-
-    # match file pattern periodically
-    file_pattern = self.file_pattern
-    match_files = (
-        impulse
-        | 'GetFilePattern' >> beam.Map(lambda x: file_pattern)
-        | MatchAll(self.empty_match_treatment))
-
-    # apply deduplication strategy if required
+    if Duration.of(self.interval).micros <= 0:
+      raise ValueError('MatchContinuously interval must be positive.')
     if self.has_deduplication:
-      # Making a Key Value so each file has its own state.
-      match_files = match_files | 'ToKV' >> beam.Map(lambda x: (x.path, x))
-      if self.match_upd:
-        match_files = match_files | 'RemoveOldAlreadyRead' >> beam.ParDo(
-            _RemoveOldDuplicates())
-      else:
-        match_files = match_files | 'RemoveAlreadyRead' >> beam.ParDo(
-            _RemoveDuplicates())
-
-    # apply windowing if required. Apply at last because deduplication relies on
-    # the global window.
+      match_files = self._match_deduplicated(pbegin)
+    else:
+      match_files = self._match_all_each_poll(pbegin)
+
+    # Apply windowing last because dedup relies on the global window.
     if self.apply_windowing:
       match_files = match_files | beam.WindowInto(FixedWindows(self.interval))
 
     return match_files
 
+  def _match_deduplicated(self,
+                          pbegin) -> beam.PCollection[filesystem.FileMetadata]:
+    # Watch emits each file once per key: the path, joined by the mtime when
+    # matching updated files; stop_timestamp bounds the polls to [start, stop).
+    clock = _PollClock()
+    if self.stop_ts == MAX_TIMESTAMP:
+      termination = never()
+    else:
+      start_ts = Timestamp.of(self.start_ts)
+      stop_ts = Timestamp.of(self.stop_ts)
+      if stop_ts < start_ts:
+        raise ValueError(
+            'MatchContinuously stop_timestamp %s precedes start_timestamp %s' %
+            (stop_ts, start_ts))
+      interval_micros = Duration.of(self.interval).micros
+      span_micros = (stop_ts - start_ts).micros
+      # Ceiling division reproduces PeriodicImpulse's tick count; the window
+      # upper bound is exclusive.
+      max_polls = -(-span_micros // interval_micros)
+      if max_polls == 0:
+        # An empty [start, stop) window never ticks; the impulse path keeps
+        # the output empty without Watch's unconditional first poll.
+        return self._match_all_each_poll(pbegin)
+      termination = _WatchWindowTermination(clock, start_ts.micros, max_polls)
+    if self.match_upd:
+      output_key_fn = _file_path_and_mtime_key
+    else:
+      output_key_fn = _file_path_key
+    poll_fn = _MatchContinuouslyPollFn(
+        self.empty_match_treatment, self.start_ts, clock)
+    # The key coder is inferred from the key function's return annotation.
+    watch = Watch(
+        poll_fn,
+        poll_interval=self.interval,
+        termination=termination,
+        output_key_fn=output_key_fn)
+    # Watch emits (pattern, file) pairs; keep the FileMetadata output type so
+    # downstream transforms stay typed instead of falling back to Any.
+    return (
+        pbegin
+        | 'Impulse' >> beam.Create([self.file_pattern])
+        | 'Watch' >> watch
+        | 'DropPattern' >> beam.Map(lambda kv: kv[1]).with_output_types(
+            filesystem.FileMetadata))
+
+  def _match_all_each_poll(self,
+                           pbegin) -> beam.PCollection[filesystem.FileMetadata]:
+    # No deduplication: re-emit every match on each poll.
+    return (
+        pbegin
+        | PeriodicImpulse(
+            start_timestamp=self.start_ts,
+            stop_timestamp=self.stop_ts,
+            fire_interval=self.interval)
+        | 'GetFilePattern' >> beam.Map(lambda x: self.file_pattern)
+        | MatchAll(self.empty_match_treatment))
+
 
 class ReadMatches(beam.PTransform):
   """Converts each result of MatchFiles() or MatchAll() to a ReadableFile.
@@ -892,50 +1015,3 @@ def finish_bundle(self):
               timestamp=key[1].start,
               windows=[key[1]]  # TODO(pabloem) HOW DO WE GET THE PANE
           ))
-
-
-class _RemoveDuplicates(beam.DoFn):
-  """Internal DoFn that filters out filenames already seen (even though the file
-  has updated)."""
-  COUNT_STATE = CombiningValueStateSpec('count', combine_fn=sum)
-
-  def process(
-      self,
-      element: tuple[str, filesystem.FileMetadata],
-      count_state=beam.DoFn.StateParam(COUNT_STATE)
-  ) -> Iterable[filesystem.FileMetadata]:
-
-    path = element[0]
-    file_metadata = element[1]
-    counter = count_state.read()
-
-    if counter == 0:
-      count_state.add(1)
-      _LOGGER.debug('Generated entry for file %s', path)
-      yield file_metadata
-    else:
-      _LOGGER.debug('File %s was already read, seen %d times', path, counter)
-
-
-class _RemoveOldDuplicates(beam.DoFn):
-  """Internal DoFn that filters out filenames already seen and timestamp
-  unchanged."""
-  TIME_STATE = CombiningValueStateSpec(
-      'count', combine_fn=partial(max, default=0.0))
-
-  def process(
-      self,
-      element: tuple[str, filesystem.FileMetadata],
-      time_state=beam.DoFn.StateParam(TIME_STATE)
-  ) -> Iterable[filesystem.FileMetadata]:
-    path = element[0]
-    file_metadata = element[1]
-    new_ts = file_metadata.last_updated_in_seconds
-    old_ts = time_state.read()
-
-    if old_ts < new_ts:
-      time_state.add(new_ts)
-      _LOGGER.debug('Generated entry for file %s', path)
-      yield file_metadata
-    else:
-      _LOGGER.debug('File %s was already read', path)