[Python] Refactor MatchContinuously to use Watch

Author: Eliaaazzz

sdks/python/apache_beam/io/fileio.py

@@ -93,28 +93,35 @@
 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 Union
 
 import apache_beam as beam
+from apache_beam.coders.coders import StrUtf8Coder
+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.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__ = [
@@ -251,6 +258,79 @@ def process(
     yield ReadableFile(metadata, self._compression)
 
 
+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)``, so the number of polls is independent of
+  how fast the runner reschedules deferred work. Only polls at or after
+  ``start`` count toward the budget; earlier polls are deferred waits that must
+  not consume it, matching ``PeriodicImpulse``, which never advances a tick
+  while waiting for the start time.
+  """
+  def __init__(self, start_micros: int, max_polls: int):
+    self._start_micros = start_micros
+    self._max_polls = max_polls
+
+  def for_new_input(self, now, element):
+    return 0
+
+  def on_poll_complete(self, now, state):
+    if now.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
+
+
+class _MatchContinuouslyPollFn(PollFn):
+  """Polls a file pattern for ``MatchContinuously``, honoring empty-match rules.
+
+  Emits no outputs before ``start_timestamp`` so polling can start ahead of the
+  first intended match. Normally each match is stamped with the poll time as
+  its event time. When matching updated files, each match is stamped with the
+  file's last-modified time so ``Watch(use_timestamp=True)`` can emit the same
+  path again when that timestamp changes. The watermark advances to the poll
+  time so downstream event-time windows keep progressing even when a poll finds
+  no new files.
+  """
+  def __init__(
+      self,
+      empty_match_treatment,
+      start_timestamp,
+      use_metadata_timestamp: bool = False):
+    self._empty_match_treatment = empty_match_treatment
+    self._start_micros = Timestamp.of(start_timestamp).micros
+    self._use_metadata_timestamp = use_metadata_timestamp
+
+  def __call__(self, file_pattern: str) -> PollResult:
+    now = Timestamp.now()
+    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)
+    if self._use_metadata_timestamp:
+      outputs = [
+          TimestampedValue(metadata, metadata.last_updated_in_seconds)
+          for metadata in match_result.metadata_list
+      ]
+      return PollResult.incomplete(outputs).with_watermark(now)
+    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.
 
@@ -306,37 +386,72 @@ 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 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]:
+    # The Watch transform polls the pattern and emits each file once per path.
+    # When matching updated files, the poll function timestamps each output with
+    # the file's last-modified time and Watch includes that timestamp in dedup.
+    # stop_timestamp bounds the watch to the polls that fall in [start, stop).
+    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
+      if interval_micros <= 0:
+        raise ValueError('MatchContinuously interval must be positive.')
+      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)
+      termination = _WatchWindowTermination(start_ts.micros, max_polls)
+    file_pattern = self.file_pattern
+    poll_fn = _MatchContinuouslyPollFn(
+        self.empty_match_treatment, self.start_ts, self.match_upd)
+    watch = Watch(
+        poll_fn,
+        poll_interval=self.interval,
+        termination=termination,
+        output_key_fn=_file_path_key,
+        output_key_coder=StrUtf8Coder(),
+        use_timestamp=self.match_upd)
+    # 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([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.
+    file_pattern = self.file_pattern
+    return (
+        pbegin
+        | PeriodicImpulse(
+            start_timestamp=self.start_ts,
+            stop_timestamp=self.stop_ts,
+            fire_interval=self.interval)
+        | 'GetFilePattern' >> beam.Map(lambda x: file_pattern)
+        | MatchAll(self.empty_match_treatment))
+
 
 class ReadMatches(beam.PTransform):
   """Converts each result of MatchFiles() or MatchAll() to a ReadableFile.
@@ -892,50 +1007,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)

sdks/python/apache_beam/io/fileio_test.py

@@ -28,12 +28,11 @@
 import uuid
 import warnings
 
-import pytest
-from hamcrest.library.text import stringmatches
-
 import apache_beam as beam
+import pytest
 from apache_beam.io import fileio
 from apache_beam.io.filebasedsink_test import _TestCaseWithTempDirCleanUp
+from apache_beam.io.filesystem import BeamIOError
 from apache_beam.io.filesystem import CompressionTypes
 from apache_beam.io.filesystems import FileSystems
 from apache_beam.options.pipeline_options import PipelineOptions
@@ -49,6 +48,7 @@
 from apache_beam.transforms.window import GlobalWindow
 from apache_beam.transforms.window import IntervalWindow
 from apache_beam.utils.timestamp import Timestamp
+from hamcrest.library.text import stringmatches
 
 warnings.filterwarnings(
     'ignore', category=FutureWarning, module='apache_beam.io.fileio_test')
@@ -420,6 +420,65 @@ def _create_extra_file(element):
 
       assert_that(match_continiously, equal_to(files))
 
+  def test_poll_fn_gates_on_start_timestamp(self):
+    tempdir = '%s%s' % (self._new_tempdir(), os.sep)
+    self._create_temp_file(dir=tempdir)
+    pattern = FileSystems.join(tempdir, '*')
+
+    future_start = fileio._MatchContinuouslyPollFn(
+        fileio.EmptyMatchTreatment.ALLOW, Timestamp.now() + 3600)
+    self.assertEqual((), future_start(pattern).outputs)
+
+    past_start = fileio._MatchContinuouslyPollFn(
+        fileio.EmptyMatchTreatment.ALLOW, Timestamp.now() - 3600)
+    self.assertEqual(1, len(past_start(pattern).outputs))
+
+  def test_poll_fn_disallows_empty_match(self):
+    tempdir = '%s%s' % (self._new_tempdir(), os.sep)
+    poll_fn = fileio._MatchContinuouslyPollFn(
+        fileio.EmptyMatchTreatment.DISALLOW, Timestamp.now() - 3600)
+    with self.assertRaises(BeamIOError):
+      poll_fn(FileSystems.join(tempdir, 'no-such-file'))
+
+  def test_poll_fn_can_timestamp_outputs_with_file_mtime(self):
+    tempdir = '%s%s' % (self._new_tempdir(), os.sep)
+    self._create_temp_file(dir=tempdir)
+    poll_fn = fileio._MatchContinuouslyPollFn(
+        fileio.EmptyMatchTreatment.ALLOW,
+        Timestamp.now() - 3600,
+        use_metadata_timestamp=True)
+    result = poll_fn(FileSystems.join(tempdir, '*'))
+    self.assertEqual(1, len(result.outputs))
+    output = result.outputs[0]
+    self.assertEqual(
+        Timestamp.of(output.value.last_updated_in_seconds), output.timestamp)
+
+  def test_watch_window_termination_ignores_pre_start_polls(self):
+    # Polls before start_timestamp are deferred waits and must not consume the
+    # budget, otherwise a future start_timestamp silently drops all output.
+    start_micros = Timestamp.of(1000).micros
+    term = fileio._WatchWindowTermination(start_micros, max_polls=2)
+    before = Timestamp.of(999)
+    after = Timestamp.of(1000)
+    state = term.for_new_input(before, 'pattern')
+    state = term.on_poll_complete(before, state)
+    state = term.on_poll_complete(before, state)
+    self.assertFalse(term.can_stop_polling(before, state))
+    state = term.on_poll_complete(after, state)
+    self.assertFalse(term.can_stop_polling(after, state))
+    state = term.on_poll_complete(after, state)
+    self.assertTrue(term.can_stop_polling(after, state))
+
+  def test_poll_fn_advances_watermark_on_empty_match(self):
+    # An empty (but allowed) match still carries a watermark so downstream
+    # event-time windows keep progressing when no new files appear.
+    tempdir = '%s%s' % (self._new_tempdir(), os.sep)
+    poll_fn = fileio._MatchContinuouslyPollFn(
+        fileio.EmptyMatchTreatment.ALLOW, Timestamp.now() - 3600)
+    result = poll_fn(FileSystems.join(tempdir, '*'))
+    self.assertEqual((), result.outputs)
+    self.assertIsNotNone(result.watermark)
+
 
 class WriteFilesTest(_TestCaseWithTempDirCleanUp):