From 37606982664f56d4892768a12e69354ea39bd60c Mon Sep 17 00:00:00 2001 From: Eliaazzz Date: Fri, 24 Jul 2026 00:36:50 +1000 Subject: [PATCH] [Python] Refactor MatchContinuously onto the Watch transform Route MatchContinuously through Watch when deduplication is enabled. The polling loop and the set of already-matched file ids now live in the splittable DoFn restriction, replacing the per-key state DoFns. Because the matched ids are part of the restriction, a runner with checkpointing enabled restores them after a restart and does not reprocess files. The docstring is updated accordingly. Verified on Flink 1.20: after killing the TaskManager mid stream the job restored from a checkpoint and every file was still emitted exactly once. has_deduplication=False keeps the previous PeriodicImpulse behaviour. --- sdks/python/apache_beam/io/fileio.py | 255 +++++++++++++++------- sdks/python/apache_beam/io/fileio_test.py | 120 +++++++++- 2 files changed, 294 insertions(+), 81 deletions(-) diff --git a/sdks/python/apache_beam/io/fileio.py b/sdks/python/apache_beam/io/fileio.py index a333b7c89775..32f1b05edfd3 100644 --- a/sdks/python/apache_beam/io/fileio.py +++ b/sdks/python/apache_beam/io/fileio.py @@ -93,28 +93,37 @@ 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 FloatCoder +from apache_beam.coders.coders import StrUtf8Coder +from apache_beam.coders.coders import TupleCoder +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 +260,98 @@ def process( yield ReadableFile(metadata, self._compression) +class _PollClock(object): + """The poll-time clock reading one ``MatchContinuously`` round shares. + + The start gate (a poll before ``start_timestamp`` emits nothing) and the + poll budget (only polls at or after ``start_timestamp`` consume it) must + agree on a single reading. With independent clock reads, a round straddling + the boundary could consume the budget without having matched anything. The + poll function writes the reading and the termination condition consumes it + within the same round; instances are not shared across bundle threads. + """ + 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)``, 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, judged by the poll's own clock reading in + ``clock``; 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, 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 file whose timestamp changed look + # new again, mirroring the Java SDK's ExtractFilenameAndLastUpdateFn — which + # also rejects a missing (zero) timestamp, since 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 for ``MatchContinuously``, honoring empty-match rules. + + Emits no outputs before ``start_timestamp`` so polling can start ahead of the + first intended match. Each match is stamped with the poll time as its event + time, and the watermark advances to the poll time so downstream event-time + windows keep progressing even when a poll finds no new files. Each round's + clock reading is recorded in ``clock`` so ``_WatchWindowTermination`` judges + the start boundary by the same reading as the gate below. + """ + 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: + 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 +362,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 +408,81 @@ 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]: + # The Watch transform polls the pattern and emits each file once per key: + # the path, joined by the last-modified time when matching updated files. + # stop_timestamp bounds the watch to the polls that fall in [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 in PeriodicImpulse. Fall + # back to the impulse path — with zero ticks dedup is moot — so the + # output stays empty and unbounded, rather than let Watch run its + # 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 + output_key_coder = TupleCoder([StrUtf8Coder(), FloatCoder()]) + else: + output_key_fn = _file_path_key + output_key_coder = StrUtf8Coder() + poll_fn = _MatchContinuouslyPollFn( + self.empty_match_treatment, self.start_ts, clock) + watch = Watch( + poll_fn, + poll_interval=self.interval, + termination=termination, + output_key_fn=output_key_fn, + output_key_coder=output_key_coder) + # 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 +1038,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) diff --git a/sdks/python/apache_beam/io/fileio_test.py b/sdks/python/apache_beam/io/fileio_test.py index ce535265ef2f..42ce976c673e 100644 --- a/sdks/python/apache_beam/io/fileio_test.py +++ b/sdks/python/apache_beam/io/fileio_test.py @@ -28,13 +28,13 @@ 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.filesystem import FileMetadata from apache_beam.io.filesystems import FileSystems from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import StandardOptions @@ -49,6 +49,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 +421,119 @@ 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_stamps_outputs_with_poll_time(self): + # Matches always carry the poll time as event time, matching the Java + # SDK's MatchPollFn; matching updated files must not change that. + tempdir = '%s%s' % (self._new_tempdir(), os.sep) + self._create_temp_file(dir=tempdir) + poll_fn = fileio._MatchContinuouslyPollFn( + fileio.EmptyMatchTreatment.ALLOW, Timestamp.now() - 3600) + before = Timestamp.now() + result = poll_fn(FileSystems.join(tempdir, '*')) + after = Timestamp.now() + self.assertEqual(1, len(result.outputs)) + output = result.outputs[0] + self.assertLessEqual(before, output.timestamp) + self.assertLessEqual(output.timestamp, after) + self.assertEqual(result.watermark, output.timestamp) + + def test_match_updated_files_keys_on_path_and_mtime(self): + # An updated file dedups as new because its key changes, mirroring the + # Java SDK's ExtractFilenameAndLastUpdateFn. + metadata = FileMetadata('/tmp/a', 1, 1234.5) + self.assertEqual(('/tmp/a', 1234.5), + fileio._file_path_and_mtime_key(metadata)) + + def test_match_updated_files_rejects_missing_mtime(self): + # Java's ExtractFilenameAndLastUpdateFn rejects a zero last-modified time: + # without mtimes, updates could never be detected. + with self.assertRaises(BeamIOError): + fileio._file_path_and_mtime_key(FileMetadata('/tmp/a', 1)) + + def test_start_equals_stop_matches_nothing(self): + # PeriodicImpulse's [start, stop) tick window is empty when start == stop; + # the deduplicated path must skip Watch's unconditional first poll by + # falling back to the impulse path, which also keeps the output unbounded. + tempdir = '%s%s' % (self._new_tempdir(), os.sep) + self._create_temp_file(dir=tempdir) + start = Timestamp.now() + with TestPipeline() as p: + match_continiously = ( + p + | fileio.MatchContinuously( + file_pattern=FileSystems.join(tempdir, '*'), + interval=0.2, + start_timestamp=start, + stop_timestamp=start)) + assert_that(match_continiously, equal_to([])) + + def test_rejects_nonpositive_interval(self): + tempdir = '%s%s' % (self._new_tempdir(), os.sep) + with self.assertRaisesRegex(ValueError, 'interval must be positive'): + with TestPipeline() as p: + _ = p | fileio.MatchContinuously( + file_pattern=FileSystems.join(tempdir, '*'), interval=0) + + 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. The + # boundary is judged by the poll's own clock reading, so a round straddling + # the start cannot consume the budget without having matched. + start_micros = Timestamp.of(1000).micros + clock = fileio._PollClock() + term = fileio._WatchWindowTermination(clock, start_micros, max_polls=2) + now = Timestamp.of(999) + state = term.for_new_input(now, 'pattern') + clock.last_poll_micros = Timestamp.of(999).micros + state = term.on_poll_complete(state) + state = term.on_poll_complete(state) + self.assertFalse(term.can_stop_polling(now, state)) + clock.last_poll_micros = Timestamp.of(1000).micros + state = term.on_poll_complete(state) + self.assertFalse(term.can_stop_polling(now, state)) + state = term.on_poll_complete(state) + self.assertTrue(term.can_stop_polling(now, state)) + + def test_poll_fn_records_its_clock_reading_for_the_termination(self): + # The gate and the poll budget share one reading per round; see _PollClock. + tempdir = '%s%s' % (self._new_tempdir(), os.sep) + clock = fileio._PollClock() + poll_fn = fileio._MatchContinuouslyPollFn( + fileio.EmptyMatchTreatment.ALLOW, Timestamp.now() + 3600, clock) + self.assertIsNone(clock.last_poll_micros) + poll_fn(FileSystems.join(tempdir, '*')) + self.assertIsNotNone(clock.last_poll_micros) + + 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):