This is an automated email from the ASF dual-hosted git repository.
Abacn pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new 9d3694165e1 [Python] Add UnboundedSource SDF wrapper (#19137) (#38724)
9d3694165e1 is described below
commit 9d3694165e1106fa7bfeee1febf0d9e229400f9b
Author: Elia Liu <[email protected]>
AuthorDate: Fri Jun 26 04:47:06 2026 +1000
[Python] Add UnboundedSource SDF wrapper (#19137) (#38724)
* [Python] Add UnboundedSource SDF wrapper (#19137)
Brings Java's ``UnboundedSource`` / ``UnboundedReader`` / ``CheckpointMark``
abstractions to the Python SDK as a Splittable-DoFn wrapper runnable on the
portable Fn API (DirectRunner / FnApiRunner). Wires the new source type into
``iobase.Read.expand()`` so ``p | beam.io.Read(my_unbounded_source)``
dispatches alongside the existing ``BoundedSource`` branch. Loosely inspired
by Java's ``Read.UnboundedSourceAsSDFWrapperFn``; the streaming-SDF template
followed for the process loop / watermark / defer plumbing is
``apache_beam.transforms.periodicsequence``.
---
sdks/python/apache_beam/io/iobase.py | 14 +-
sdks/python/apache_beam/io/iobase_test.py | 49 +-
sdks/python/apache_beam/io/unbounded_source.py | 996 ++++++++++++++
.../python/apache_beam/io/unbounded_source_test.py | 1380 ++++++++++++++++++++
4 files changed, 2434 insertions(+), 5 deletions(-)
diff --git a/sdks/python/apache_beam/io/iobase.py
b/sdks/python/apache_beam/io/iobase.py
index afc977406af..aa03280050f 100644
--- a/sdks/python/apache_beam/io/iobase.py
+++ b/sdks/python/apache_beam/io/iobase.py
@@ -918,7 +918,10 @@ class Read(ptransform.PTransform):
"""Initializes a Read transform.
Args:
- source: Data source to read from.
+ source: the data source to read from. May be a ``BoundedSource``, an
+ ``UnboundedSource``, or a ``PTransform`` (which is applied directly).
+ For any other source ``Read`` is treated as a primitive and relayed to
+ the runner implementation.
"""
super().__init__()
self.source = source
@@ -944,6 +947,11 @@ class Read(ptransform.PTransform):
| 'EmitSource' >>
core.Map(lambda _: self.source).with_output_types(BoundedSource)
| SDFBoundedSourceReader(display_data))
+ # Local import to avoid a circular dependency.
+ from apache_beam.io.unbounded_source import ReadFromUnboundedSource
+ from apache_beam.io.unbounded_source import UnboundedSource
+ if isinstance(self.source, UnboundedSource):
+ return pbegin | ReadFromUnboundedSource(self.source)
elif isinstance(self.source, ptransform.PTransform):
# The Read transform can also admit a full PTransform as an input
# rather than an anctual source. If the input is a PTransform, then
@@ -993,6 +1001,10 @@ class Read(ptransform.PTransform):
is_bounded=beam_runner_api_pb2.IsBounded.BOUNDED
if self.source.is_bounded() else
beam_runner_api_pb2.IsBounded.UNBOUNDED))
+ # Local import to avoid a circular dependency.
+ from apache_beam.io.unbounded_source import UnboundedSource
+ if isinstance(self.source, UnboundedSource):
+ return super().to_runner_api_parameter(context)
elif isinstance(self.source, ptransform.PTransform):
return self.source.to_runner_api_parameter(context)
raise NotImplementedError(
diff --git a/sdks/python/apache_beam/io/iobase_test.py
b/sdks/python/apache_beam/io/iobase_test.py
index eb9617cfae3..dbedf4681f4 100644
--- a/sdks/python/apache_beam/io/iobase_test.py
+++ b/sdks/python/apache_beam/io/iobase_test.py
@@ -21,15 +21,16 @@
import unittest
-import mock
-
import apache_beam as beam
-from apache_beam.io.concat_source import ConcatSource
-from apache_beam.io.concat_source_test import RangeSource
+import mock
from apache_beam.io import iobase
from apache_beam.io import range_trackers
+from apache_beam.io.concat_source import ConcatSource
+from apache_beam.io.concat_source_test import RangeSource
from apache_beam.io.iobase import SourceBundle
from apache_beam.options.pipeline_options import DebugOptions
+from apache_beam.portability import common_urns
+from apache_beam.portability import python_urns
from apache_beam.testing.util import assert_that
from apache_beam.testing.util import equal_to
@@ -220,5 +221,45 @@ class UseSdfBoundedSourcesTests(unittest.TestCase):
self._run_sdf_wrapper_pipeline(RangeSource(0, 4), [0, 1, 2, 3])
+class UseSdfUnboundedSourcesTests(unittest.TestCase):
+ """Covers the UnboundedSource branch in
+ ``iobase.Read.expand()``. Uses ``UnboundedCountingSource`` from
+ ``unbounded_source_test`` as a finite fake source (no network).
+ """
+ def test_read_end_to_end_unbounded(self):
+ from apache_beam.io.unbounded_source_test import UnboundedCountingSource
+ with beam.Pipeline() as p:
+ out = p | beam.io.Read(UnboundedCountingSource(5))
+ assert_that(out, equal_to([0, 1, 2, 3, 4]))
+
+ def test_read_unbounded_pcollection_is_unbounded(self):
+ from apache_beam.io.unbounded_source_test import UnboundedCountingSource
+ p = beam.Pipeline()
+ out = p | beam.io.Read(UnboundedCountingSource(3))
+ self.assertFalse(out.is_bounded)
+
+ def test_read_unbounded_serializes_as_expanded_composite(self):
+ from apache_beam.io.unbounded_source_test import UnboundedCountingSource
+ p = beam.Pipeline()
+ p | 'ReadIt' >> beam.io.Read(UnboundedCountingSource(3))
+
+ proto = p.to_runner_api(use_fake_coders=True)
+ transforms = proto.components.transforms.values()
+ deprecated_reads = [
+ transform.unique_name for transform in transforms
+ if transform.spec.urn == common_urns.deprecated_primitives.READ.urn
+ ]
+ read_transforms = [
+ transform for transform in proto.components.transforms.values()
+ if transform.unique_name == 'ReadIt'
+ ]
+
+ self.assertEqual([], deprecated_reads)
+ self.assertEqual(1, len(read_transforms))
+ self.assertEqual(
+ python_urns.GENERIC_COMPOSITE_TRANSFORM, read_transforms[0].spec.urn)
+ self.assertTrue(read_transforms[0].subtransforms)
+
+
if __name__ == '__main__':
unittest.main()
diff --git a/sdks/python/apache_beam/io/unbounded_source.py
b/sdks/python/apache_beam/io/unbounded_source.py
new file mode 100644
index 00000000000..3246602252e
--- /dev/null
+++ b/sdks/python/apache_beam/io/unbounded_source.py
@@ -0,0 +1,996 @@
+#
+# 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 ``UnboundedSource`` support for the Python SDK.
+
+``UnboundedSource`` support is currently experimental in the Python SDK; the
+API may change in backwards-incompatible ways.
+
+An unbounded source reads an effectively infinite stream of records (message
+queues, change-data-capture feeds, and similar) with checkpoint-based
+pause/resume, watermark reporting, and bundle finalization.
+
+To define a source, implement :class:`UnboundedSource`, an
+:class:`UnboundedReader`, and (when the reader has a resumable position) a
+:class:`CheckpointMark`::
+
+ import apache_beam as beam
+ from apache_beam.io.unbounded_source import (
+ CheckpointMark, UnboundedReader, UnboundedSource)
+ from apache_beam.utils.timestamp import MAX_TIMESTAMP
+
+ class MyCheckpointMark(CheckpointMark):
+ def __init__(self, position):
+ self.position = position
+
+ def finalize_checkpoint(self):
+ # Commit/acknowledge records up to ``position`` upstream, e.g. ack the
+ # consumed messages on a queue.
+ ...
+
+ class MyReader(UnboundedReader):
+ def start(self):
+ # Position at the first record; return whether one is available.
+ ...
+
+ def advance(self):
+ # Move to the next record; ``False`` means no data is available now.
+ ...
+
+ def get_current(self):
+ ...
+
+ def get_current_timestamp(self):
+ ... # event time of the current record
+
+ def get_watermark(self):
+ # Lower bound on the timestamps of future records. Return
+ # ``MAX_TIMESTAMP`` to signal the reader has permanently finished.
+ ...
+
+ def get_checkpoint_mark(self):
+ return MyCheckpointMark(...)
+
+ class MySource(UnboundedSource):
+ def split(self, desired_num_splits, options=None):
+ # Return independent sub-sources, or ``[self]`` when not splittable.
+ return [self]
+
+ def create_reader(self, options, checkpoint_mark):
+ # Build a reader; resume after ``checkpoint_mark`` when it is not None.
+ return MyReader(...)
+
+ def get_checkpoint_mark_coder(self):
+ return ... # a Coder for MyCheckpointMark
+
+Read the source in a pipeline with :class:`apache_beam.io.Read`::
+
+ with beam.Pipeline() as p:
+ p | beam.io.Read(MySource()) | beam.Map(print)
+"""
+
+import collections
+import dataclasses
+import logging
+import threading
+import time
+from typing import Any
+from typing import Callable
+from typing import Iterable
+from typing import Optional
+
+from apache_beam import coders
+from apache_beam.coders.coders import BooleanCoder
+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.coders.coders import _MemoizingPickleCoder
+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 MIN_TIMESTAMP
+from apache_beam.utils.timestamp import Duration
+from apache_beam.utils.timestamp import Timestamp
+
+__all__ = [
+ 'CheckpointMark',
+ 'UnboundedReader',
+ 'UnboundedSource',
+ 'ReadFromUnboundedSource',
+]
+
+_LOGGER = logging.getLogger(__name__)
+
+# Sentinel used when a reader has no data available right now.
+# This is distinct from end-of-stream.
+_NO_DATA = object()
+
+_DEFAULT_POLL_INTERVAL_SECONDS = 1.0
+_DEFAULT_DESIRED_NUM_SPLITS = 20
+_DEFAULT_MAX_RECORDS_PER_BUNDLE = 10000
+_DEFAULT_MAX_READ_TIME_SECONDS = 10.0
+# A reader parked by a residual that never resumes is closed once idle this
+# long; the cache also caps its entry count as a memory backstop.
+_DEFAULT_READER_CACHE_IDLE_SECONDS = 60.0
+_DEFAULT_READER_CACHE_MAX_SIZE = 100
+
+# Encodes a source to a structural cache key. Internally consistent across park
+# and acquire; need not match the restriction wire coder.
+_SOURCE_KEY_CODER = _MemoizingPickleCoder()
+
+#
------------------------------------------------------------------------------
+# Public abstract base classes.
+#
------------------------------------------------------------------------------
+
+
+class CheckpointMark(object):
+ """A durable, serializable position in an :class:`UnboundedSource`.
+
+ Produced by :meth:`UnboundedReader.get_checkpoint_mark`, encoded with
+ :meth:`UnboundedSource.get_checkpoint_mark_coder`, and used to resume a
reader
+ (see :meth:`UnboundedSource.create_reader`).
+ """
+ def finalize_checkpoint(self) -> None:
+ """Called once the runner has durably committed work up to this mark.
+
+ Override to acknowledge/commit upstream (for example, ack the consumed
+ messages on a queue). The default is a no-op.
+
+ The runner calls this at most once for a committed checkpoint mark.
+ Finalization is best effort; a mark may never be finalized. An exception
+ raised here is logged. On bundle retry an uncommitted mark may be re-cut
+ over an overlapping span, so this method must be idempotent (acknowledge by
+ absolute position).
+ """
+ pass
+
+
+class UnboundedReader(object):
+ """Reads records from an :class:`UnboundedSource`.
+
+ Lifecycle: exactly one :meth:`start`, then any number of :meth:`advance`
+ calls; whenever one returns ``True`` the current record is available via
+ :meth:`get_current` / :meth:`get_current_timestamp`. A ``False`` return means
+ "no data available right now", which is distinct from end-of-stream: a reader
+ signals a permanent end by reporting a watermark of ``MAX_TIMESTAMP``.
+ """
+ def start(self) -> bool:
+ """Positions at the first record; returns whether one is available."""
+ raise NotImplementedError
+
+ def advance(self) -> bool:
+ """Advances to the next record. ``False`` means no data is available now.
+
+ Should not block. The wrapper enforces the per-bundle record and time caps
+ only between records, so a blocking ``start``/``advance`` can overrun the
+ time cap and stall the bundle. Return ``False`` when no data is currently
+ available instead of waiting.
+ """
+ raise NotImplementedError
+
+ def get_current(self) -> Any:
+ """Returns the record claimed by the last successful start/advance."""
+ raise NotImplementedError
+
+ def get_current_timestamp(self) -> Timestamp:
+ """Returns the event-time timestamp of the current record."""
+ raise NotImplementedError
+
+ def get_watermark(self) -> Timestamp:
+ """An approximate lower bound on timestamps of future records.
+
+ Treated as monotonic by the wrapper. Return ``MAX_TIMESTAMP`` to signal
that
+ this reader has permanently finished.
+ """
+ raise NotImplementedError
+
+ def get_checkpoint_mark(self) -> CheckpointMark:
+ """Returns a durable mark to resume from. Call only at a bundle
boundary."""
+ raise NotImplementedError
+
+ def close(self) -> None:
+ """Releases reader resources. Default no-op."""
+ pass
+
+
+class UnboundedSource(iobase.SourceBase):
+ """A source producing an unbounded stream of records with checkpointing.
+
+ Read it in a pipeline with :class:`apache_beam.io.Read`::
+
+ p | beam.io.Read(MyUnboundedSource())
+ """
+ def split(self,
+ desired_num_splits: int,
+ options: Optional[Any] = None) -> Iterable['UnboundedSource']:
+ """Splits into at most ``desired_num_splits`` independent sub-sources.
+
+ Each returned sub-source must be independent and must not share mutable
+ state with siblings (the runner may execute them concurrently across
+ workers). Return ``[self]`` if the source cannot be split. Splitting is
+ performed once, before any checkpoint exists; once a reader has
+ checkpointed, the restriction is kept intact.
+ """
+ raise NotImplementedError
+
+ def create_reader(
+ self, options: Optional[Any],
+ checkpoint_mark: Optional[CheckpointMark]) -> UnboundedReader:
+ """Creates a reader, optionally resuming from ``checkpoint_mark``.
+
+ Contract:
+ * When ``checkpoint_mark`` is ``None``, the returned reader's ``start()``
+ produces the very first record of the source (or returns ``False`` if
+ none yet).
+ * When ``checkpoint_mark`` is not ``None``, the returned reader's
+ ``start()`` produces the first record strictly after the position
+ encoded by ``checkpoint_mark``. The reader must not re-deliver records
+ already covered by the prior bundle.
+ """
+ raise NotImplementedError
+
+ def get_checkpoint_mark_coder(self) -> Coder:
+ """Returns the coder for this source's :class:`CheckpointMark` instances.
+
+ The SDK may call this while encoding or decoding source restrictions.
+ Implementations should be deterministic, side-effect free, and should not
+ perform I/O.
+ """
+ raise NotImplementedError(
+ '%s must override get_checkpoint_mark_coder() to return a Coder for '
+ 'its CheckpointMark subclass.' % type(self).__name__)
+
+ def is_bounded(self) -> bool:
+ # SourceBase.is_bounded raises; an unbounded source is, by definition, not.
+ return False
+
+ def default_output_coder(self) -> Coder:
+ # Permissive default; override for a tighter coder.
+ return coders.registry.get_coder(object)
+
+
+#
------------------------------------------------------------------------------
+# SDF wrapper internals: a private implementation detail of
+# ReadFromUnboundedSource.
+#
------------------------------------------------------------------------------
+
+
[email protected](frozen=True)
+class _UnboundedSourceRestriction(object):
+ """Durable SDF restriction describing where a reader should (re)start.
+
+ Holds only serializable state -- never a live reader. ``is_done`` marks the
+ terminal (MAX-watermark) transition. ``finalization_checkpoint_mark`` is kept
+ separate from ``checkpoint_mark`` so a done primary can carry a commit hook
+ without polluting the resume state.
+
+ Field roles:
+ * ``checkpoint_mark`` -- RESUME state. A reader rebuilt from this mark
+ must produce the first record strictly after it.
+ * ``finalization_checkpoint_mark`` -- COMMIT hook. Only set on a done
+ primary that was just cut this bundle. Registered with the runner's
+ bundle finalizer to acknowledge upstream. Independent of
+ ``checkpoint_mark`` so a residual's resume state can be ``None`` while
+ still recording what should be committed.
+ """
+ source: UnboundedSource
+ checkpoint_mark: Optional[CheckpointMark] = None
+ watermark: Timestamp = MIN_TIMESTAMP
+ is_done: bool = False
+ finalization_checkpoint_mark: Optional[CheckpointMark] = None
+
+
+class _UnboundedSourceRestrictionCoder(Coder):
+ """Encodes :class:`_UnboundedSourceRestriction` as a fixed 5-tuple.
+
+ Stateless: at encode time the source's own
+ :meth:`UnboundedSource.get_checkpoint_mark_coder` is looked up from the
+ restriction; at decode time the source is decoded first and its coder
+ drives the checkpoint-mark decoding. This avoids passing source-specific
+ coder state into the coder's constructor, which in turn lets
+ :class:`_UnboundedSourceRestrictionProvider` and
+ :class:`_ReadFromUnboundedSourceDoFn` be module-level classes.
+
+ Wire shape: source_bytes / checkpoint_bytes / watermark / done /
+ finalization_checkpoint_bytes -- the checkpoint and finalization bytes
+ are independently encoded with the (source-declared) checkpoint coder
+ wrapped in :class:`NullableCoder`.
+ """
+ def __init__(self):
+ self._source_coder = _MemoizingPickleCoder()
+ self._bytes_coder = coders.BytesCoder()
+ self._tuple_coder = TupleCoder((
+ self._bytes_coder, # source (pickled bytes)
+ self._bytes_coder, # checkpoint_mark (nullable-encoded bytes)
+ TimestampCoder(), # watermark
+ BooleanCoder(), # is_done
+ self._bytes_coder)) # finalization_checkpoint_mark (nullable-encoded)
+
+ def _checkpoint_coder(self, source: UnboundedSource) -> Coder:
+ return NullableCoder(source.get_checkpoint_mark_coder())
+
+ def encode(self, restriction: '_UnboundedSourceRestriction') -> bytes:
+ source_bytes = self._source_coder.encode(restriction.source)
+ cp_coder = self._checkpoint_coder(restriction.source)
+ return self._tuple_coder.encode((
+ source_bytes,
+ cp_coder.encode(restriction.checkpoint_mark),
+ restriction.watermark,
+ restriction.is_done,
+ cp_coder.encode(restriction.finalization_checkpoint_mark)))
+
+ def decode(self, encoded: bytes) -> '_UnboundedSourceRestriction':
+ (source_bytes, checkpoint_bytes, watermark, is_done,
+ finalization_bytes) = self._tuple_coder.decode(encoded)
+ source = self._source_coder.decode(source_bytes)
+ cp_coder = self._checkpoint_coder(source)
+ return _UnboundedSourceRestriction(
+ source=source,
+ checkpoint_mark=cp_coder.decode(checkpoint_bytes),
+ watermark=watermark,
+ is_done=is_done,
+ finalization_checkpoint_mark=cp_coder.decode(finalization_bytes))
+
+ def is_deterministic(self) -> bool:
+ # Pickled source and checkpoint are not guaranteed deterministic.
+ return False
+
+
+class _ReaderCache(object):
+ """Holds live readers between an SDF self-checkpoint and its resume.
+
+ A fresh tracker is built for every bundle, so a reader parked at a
+ self-checkpoint would otherwise be closed and rebuilt from its checkpoint
+ mark on the next bundle. Parking it here lets the resuming bundle reclaim the
+ same started reader, keyed by the residual's structural ``(source, checkpoint
+ mark)``. ``acquire`` removes the entry, so two trackers never drive one
+ reader.
+
+ A residual may be reassigned to another worker and never resume here; such a
+ reader is released once it falls idle past ``idle_seconds`` or when the entry
+ count exceeds ``max_size``, and the owning DoFn's teardown releases the rest.
+ One DoFn instance drives several trackers across threads, so access is
locked.
+ """
+ def __init__(
+ self,
+ idle_seconds: float = _DEFAULT_READER_CACHE_IDLE_SECONDS,
+ max_size: int = _DEFAULT_READER_CACHE_MAX_SIZE,
+ now: Optional[Callable[[], float]] = None):
+ self._idle_seconds = idle_seconds
+ self._max_size = max_size
+ self._now = now or time.monotonic
+ self._lock = threading.Lock()
+ # key -> (reader, started, parked_at); ordered oldest-first for eviction.
+ self._entries = collections.OrderedDict(
+ ) # type: collections.OrderedDict[Any, tuple[UnboundedReader, bool,
float]]
+
+ def acquire(self, key: Any) -> Optional[tuple['UnboundedReader', bool]]:
+ """Removes and returns ``(reader, started)`` for ``key``, or None."""
+ with self._lock:
+ entry = self._entries.pop(key, None)
+ stale = self._evict_idle()
+ self._close_readers(stale)
+ if entry is None:
+ return None
+ return entry[0], entry[1]
+
+ def park(self, key: Any, reader: 'UnboundedReader', started: bool) -> None:
+ """Stores ``reader`` under ``key`` for a later bundle to reclaim. A reader
+ already parked under ``key`` is closed."""
+ with self._lock:
+ replaced = self._entries.pop(key, None)
+ self._entries[key] = (reader, started, self._now())
+ stale = self._evict_idle()
+ while len(self._entries) > self._max_size:
+ _, oldest = self._entries.popitem(last=False)
+ stale.append(oldest[0])
+ if replaced is not None and replaced[0] is not reader:
+ stale.append(replaced[0])
+ self._close_readers(stale)
+
+ def close_all(self) -> None:
+ """Closes every parked reader. Called from the owning DoFn's teardown."""
+ with self._lock:
+ entries = list(self._entries.values())
+ self._entries.clear()
+ self._close_readers(entry[0] for entry in entries)
+
+ def _evict_idle(self) -> list:
+ # Caller holds the lock. Pops entries idle past the window (oldest first)
+ # and returns their readers for the caller to close after unlocking.
+ deadline = self._now() - self._idle_seconds
+ stale = []
+ while self._entries:
+ key, entry = next(iter(self._entries.items()))
+ if entry[2] > deadline:
+ break
+ del self._entries[key]
+ stale.append(entry[0])
+ return stale
+
+ def _close_readers(self, readers) -> None:
+ for reader in readers:
+ try:
+ reader.close()
+ except Exception: # pylint: disable=broad-except
+ _LOGGER.warning('Error closing UnboundedReader', exc_info=True)
+
+
+class _UnboundedSourceRestrictionTracker(iobase.RestrictionTracker):
+ """Drives an :class:`UnboundedReader` for one SDF restriction.
+
+ Owns the live reader (lazily created, never serialized): both
runner-initiated
+ ``defer_remainder`` self-checkpoints with ``try_split(0)``, which must
+ checkpoint the live reader.
+
+ A self-checkpoint parks the reader in the DoFn's :class:`_ReaderCache` for
the
+ next bundle to reclaim, keeping one started reader alive across bundles. The
+ DoFn injects ``_reader_cache`` at the start of ``process()``; with no cache
+ the tracker builds a fresh reader each bundle.
+
+ ``process()`` only sees a ``RestrictionTrackerView``, which hides custom
+ methods and whose ``try_claim`` returns just a bool, so the freshly-read
+ record is handed back through a one-element holder list passed as the
+ ``try_claim`` *position* argument.
+ """
+ def __init__(
+ self,
+ restriction: _UnboundedSourceRestriction,
+ options: Optional[Any] = None):
+ self._restriction = restriction
+ self._options = options
+ self._reader = None # type: Optional[UnboundedReader]
+ self._started = False
+ # True once a checkpoint has been cut this bundle (EOF or self-checkpoint).
+ self._checkpoint_taken = False
+ # Cross-bundle reader cache, injected by the DoFn; None disables caching.
+ self._reader_cache = None # type: Optional[_ReaderCache]
+
+ def _ensure_reader(self) -> None:
+ if self._reader is not None:
+ return
+ cached = self._acquire_cached_reader()
+ if cached is not None:
+ # A parked reader is already started and positioned past its checkpoint.
+ self._reader, self._started = cached
+ return
+ self._reader = self._restriction.source.create_reader(
+ self._options, self._restriction.checkpoint_mark)
+
+ def _cache_key(self,
+ restriction: _UnboundedSourceRestriction) -> Optional[Any]:
+ """Structural ``(source, checkpoint)`` key, or None when uncacheable.
+
+ Built from the source pickle and the source's own checkpoint coder so a
+ parked reader and its resuming restriction map to the same entry. A None
+ key disables caching for that restriction; the resume then rebuilds from
+ the checkpoint mark under the source's ``create_reader`` contract.
+ """
+ try:
+ source_bytes = _SOURCE_KEY_CODER.encode(restriction.source)
+ cp_coder = NullableCoder(restriction.source.get_checkpoint_mark_coder())
+ return source_bytes, cp_coder.encode(restriction.checkpoint_mark)
+ except Exception: # pylint: disable=broad-except
+ return None
+
+ def _acquire_cached_reader(self) -> Optional[tuple['UnboundedReader', bool]]:
+ if self._reader_cache is None:
+ return None
+ key = self._cache_key(self._restriction)
+ if key is None:
+ return None
+ return self._reader_cache.acquire(key)
+
+ def _park_or_close_reader(
+ self, residual: _UnboundedSourceRestriction) -> None:
+ """Hands the live reader to the cache for ``residual`` to reclaim, or
+ closes it when no cache is available or the restriction is uncacheable."""
+ if self._reader is None:
+ return
+ key = (
+ self._cache_key(residual) if self._reader_cache is not None else None)
+ if key is None:
+ self._close_reader_if_open()
+ return
+ reader, self._reader = self._reader, None
+ self._reader_cache.park(key, reader, self._started)
+
+ def _clone_checkpoint(
+ self, checkpoint: Optional[CheckpointMark]) -> Optional[CheckpointMark]:
+ """Returns an independent copy of a mark via the source's checkpoint coder.
+
+ Used to keep the primary's finalize hook and the residual's resume state
+ from sharing one object, since a user ``finalize_checkpoint()`` may mutate
+ the mark.
+ """
+ if checkpoint is None:
+ return None
+ coder = self._restriction.source.get_checkpoint_mark_coder()
+ return coder.decode(coder.encode(checkpoint))
+
+ def _close_reader_if_open(self) -> None:
+ """Idempotent reader release. Called by the EOF and split paths, and by
+ the DoFn's ``finally`` so an exception inside ``process()`` does not leak
+ sockets / file descriptors held by the live :class:`UnboundedReader`.
+ """
+ if self._reader is None:
+ return
+ try:
+ self._reader.close()
+ except Exception: # pylint: disable=broad-except
+ _LOGGER.warning('Error closing UnboundedReader', exc_info=True)
+ finally:
+ self._reader = None
+
+ def current_restriction(self) -> _UnboundedSourceRestriction:
+ return self._restriction
+
+ def try_claim(self, out: list[Any]) -> bool:
+ """Advances the reader by one record.
+
+ ``out[0]`` receives ``(value, record_timestamp, source_watermark)`` on the
+ has-data path, or the :data:`_NO_DATA` sentinel otherwise. The watermark is
+ the source's reported watermark, not the record's event time: the DoFn
+ advances the output watermark with the former and timestamps the record
+ with the latter. The argument doubles as the output holder (the
+ ``RestrictionTracker`` ABC treats it as a claim position), which the
+ threadsafe-tracker chain forwards opaquely.
+ """
+ try:
+ return self._try_claim_inner(out)
+ except Exception:
+ # Reader state is now indeterminate; release it before re-raising.
+ self._close_reader_if_open()
+ raise
+
+ def _try_claim_inner(self, out: list[Any]) -> bool:
+ if self._restriction.is_done:
+ out[0] = _NO_DATA
+ return False
+ self._ensure_reader()
+ if not self._started:
+ has_data = self._reader.start()
+ else:
+ has_data = self._reader.advance()
+ self._started = True
+ if has_data:
+ # Emit an available record before checking the watermark: a reader may
+ # report its last record and a MAX_TIMESTAMP watermark on the same call,
+ # and EOF is realized on the next data-less claim.
+ out[0] = (
+ self._reader.get_current(),
+ self._reader.get_current_timestamp(),
+ self._reader.get_watermark())
+ return True
+ watermark = self._reader.get_watermark()
+ if watermark >= MAX_TIMESTAMP:
+ # No data and watermark at MAX: cut a final checkpoint, close, mark done.
+ checkpoint = self._reader.get_checkpoint_mark()
+ self._close_reader_if_open()
+ self._restriction = dataclasses.replace(
+ self._restriction,
+ checkpoint_mark=None, # nothing left to resume from
+ watermark=MAX_TIMESTAMP,
+ is_done=True,
+ finalization_checkpoint_mark=checkpoint)
+ self._checkpoint_taken = True
+ out[0] = _NO_DATA
+ return False
+ # No data is available now. Refresh the watermark before deferring.
+ self._restriction = dataclasses.replace(
+ self._restriction, watermark=watermark)
+ out[0] = _NO_DATA
+ return True
+
+ def try_split(
+ self, fraction_of_remainder
+ ) -> Optional[tuple[_UnboundedSourceRestriction,
+ _UnboundedSourceRestriction]]:
+ """Cuts a checkpoint, returning (primary, residual) or None.
+
+ The cut checkpoint goes into ``primary.finalization_checkpoint_mark`` so
+ the DoFn can register a bundle-finalize callback for it. The same
+ checkpoint object also goes into ``residual.checkpoint_mark`` so the
+ resumed reader rebuilds at the correct position. The two fields are
+ independent on purpose (see :class:`_UnboundedSourceRestriction`
+ docstring): a runner that re-processes the primary alone must not see
+ a stale resume state, and a residual must not register finalize again
+ until ITS checkpoint is cut in a future bundle.
+ """
+ try:
+ return self._try_split_inner(fraction_of_remainder)
+ except Exception:
+ self._close_reader_if_open()
+ raise
+
+ def _try_split_inner(self, fraction_of_remainder):
+ # Only self-checkpoint (fraction 0) is supported; decline runner-initiated
+ # dynamic splits.
+ if fraction_of_remainder != 0:
+ return None
+ if self._reader is None or not self._started or self._restriction.is_done:
+ return None
+ checkpoint = self._reader.get_checkpoint_mark()
+ # The residual watermark is advisory; the SDF watermark estimator state is
+ # the authoritative cross-bundle watermark.
+ watermark = self._reader.get_watermark()
+ # Keep the two channels independent: the primary carries only the finalize
+ # hook, the residual only the resume state. The residual gets its own clone
+ # so a finalize_checkpoint() that mutates the primary's mark cannot corrupt
+ # the residual's resume position before the runner encodes it.
+ primary = dataclasses.replace(
+ self._restriction,
+ checkpoint_mark=None,
+ is_done=True,
+ finalization_checkpoint_mark=checkpoint)
+ residual = _UnboundedSourceRestriction(
+ source=self._restriction.source,
+ checkpoint_mark=self._clone_checkpoint(checkpoint),
+ watermark=watermark,
+ is_done=False,
+ finalization_checkpoint_mark=None)
+ self._restriction = primary
+ self._checkpoint_taken = True
+ # Park the reader so the resuming bundle reclaims it; on a cache miss the
+ # residual rebuilds one from its checkpoint mark.
+ self._park_or_close_reader(residual)
+ return primary, residual
+
+ def check_done(self) -> bool:
+ # Called after every process(); must raise if work is left unaccounted for.
+ if self._restriction.is_done or self._checkpoint_taken:
+ return True
+ raise ValueError(
+ 'UnboundedSource restriction was neither finished nor checkpointed; '
+ 'process() must self-checkpoint via defer_remainder() or run to EOF: '
+ '%r' % (self._restriction, ))
+
+ def current_progress(self) -> 'iobase.RestrictionProgress':
+ # Backlog-based progress is not implemented; report a coarse done/not-done
+ # signal via ``completed`` / ``remaining``.
+ if self._restriction.is_done:
+ return iobase.RestrictionProgress(completed=1.0, remaining=0.0)
+ return iobase.RestrictionProgress(completed=0.0, remaining=1.0)
+
+ def is_bounded(self) -> bool:
+ return False
+
+
+class _UnboundedSourceRestrictionProvider(core.RestrictionProvider):
+ """Wraps an :class:`UnboundedSource` element as an SDF restriction.
+
+ Stateless module-level singleton (see :data:`_PROVIDER`): all
+ source-specific state (e.g. the source's checkpoint coder) is derived
+ per-call from the restriction's ``source`` field, which lets
+ :class:`_ReadFromUnboundedSourceDoFn` live at module level too. The provider
+ currently passes ``None`` for the ``options`` forwarded to
+ :meth:`UnboundedSource.split`.
+ """
+ def __init__(self):
+ self._restriction_coder = _UnboundedSourceRestrictionCoder()
+
+ def initial_restriction(
+ self, element: UnboundedSource) -> _UnboundedSourceRestriction:
+ if not isinstance(element, UnboundedSource):
+ raise TypeError(
+ 'ReadFromUnboundedSource expected an UnboundedSource element, got %r'
+ % (element, ))
+ return _UnboundedSourceRestriction(source=element)
+
+ def create_tracker(
+ self, restriction: _UnboundedSourceRestriction
+ ) -> _UnboundedSourceRestrictionTracker:
+ return _UnboundedSourceRestrictionTracker(restriction)
+
+ def split(self, element,
+ restriction) -> Iterable[_UnboundedSourceRestriction]:
+ if restriction.is_done or restriction.checkpoint_mark is not None:
+ yield restriction
+ return
+
+ # ``source.split`` is user code and may refuse to split; fall back to a
+ # single restriction on error.
+ try:
+ split_sources = list(
+ restriction.source.split(_DEFAULT_DESIRED_NUM_SPLITS, None))
+ except Exception: # pylint: disable=broad-except
+ _LOGGER.warning(
+ 'Exception while splitting UnboundedSource. Source not split.',
+ exc_info=True)
+ yield restriction
+ return
+
+ if not split_sources:
+ yield restriction
+ return
+
+ # A non-UnboundedSource split result is a contract violation, not a
+ # refusal, so fail loudly (outside the try/except above).
+ for split_source in split_sources:
+ if not isinstance(split_source, UnboundedSource):
+ raise TypeError(
+ 'UnboundedSource.split() produced %r, expected UnboundedSource' %
+ (split_source, ))
+
+ for split_source in split_sources:
+ yield dataclasses.replace(
+ restriction,
+ source=split_source,
+ checkpoint_mark=None,
+ is_done=False,
+ finalization_checkpoint_mark=None)
+
+ def restriction_size(self, element, restriction) -> int:
+ # TODO(https://github.com/apache/beam/issues/19137): implement backlog
+ # estimation.
+ return 1
+
+ def restriction_coder(self) -> Coder:
+ return self._restriction_coder
+
+ def truncate(self, element, restriction):
+ # On drain, stop emitting new records.
+ return None
+
+
+# Module-level stateless singleton, captured via ``RestrictionParam`` at the
+# DoFn's class-definition time.
+_PROVIDER = _UnboundedSourceRestrictionProvider()
+
+
+class _FinalizeCheckpointOnce(object):
+ def __init__(self, checkpoint_mark: CheckpointMark):
+ self._checkpoint_mark = checkpoint_mark
+ # The lock keeps finalization idempotent if a runner ever invokes the
+ # callback more than once.
+ self._lock = threading.Lock()
+ self._finalized = False
+
+ def __call__(self) -> None:
+ with self._lock:
+ if self._finalized:
+ return
+ self._finalized = True
+ # Finalization is best effort: log and swallow so a failing user override
+ # does not fail the bundle (matches CheckpointMark.finalize_checkpoint).
+ try:
+ self._checkpoint_mark.finalize_checkpoint()
+ except Exception: # pylint: disable=broad-except
+ _LOGGER.warning(
+ 'Error finalizing UnboundedSource checkpoint mark.', exc_info=True)
+
+
+class _ReadFromUnboundedSourceDoFn(core.DoFn):
+ """SDF wrapper driving an :class:`UnboundedReader` for one restriction.
+
+ Module-level so stdlib pickle and cloudpickle can serialise the DoFn. The
+ restriction provider is the module-level :data:`_PROVIDER` singleton.
+ """
+ def __init__(
+ self,
+ poll_interval: float = _DEFAULT_POLL_INTERVAL_SECONDS,
+ max_records_per_bundle: int = _DEFAULT_MAX_RECORDS_PER_BUNDLE,
+ max_read_time_seconds: float = _DEFAULT_MAX_READ_TIME_SECONDS,
+ _now: Optional[Callable[[], float]] = None):
+ self._poll_interval = poll_interval
+ self._max_records_per_bundle = max_records_per_bundle
+ self._max_read_time_seconds = max_read_time_seconds
+ # Monotonic clock seam; tests inject a deterministic clock.
+ self._now = _now
+ # Per-worker reader cache; created in setup(), never serialized.
+ self._reader_cache = None # type: Optional[_ReaderCache]
+
+ def setup(self):
+ self._reader_cache = _ReaderCache()
+
+ def teardown(self):
+ if self._reader_cache is not None:
+ self._reader_cache.close_all()
+ self._reader_cache = None
+
+ @core.DoFn.unbounded_per_element()
+ def process(
+ self,
+ unused_element,
+ bundle_finalizer=core.DoFn.BundleFinalizerParam,
+ tracker=core.DoFn.RestrictionParam(_PROVIDER),
+ watermark_estimator=core.DoFn.WatermarkEstimatorParam(
+ ManualWatermarkEstimator.default_provider())):
+ # Positional params (element, bundle finalizer) must precede the
+ # kwarg-injected ones (tracker, watermark estimator).
+ assert isinstance(tracker, sdf_utils.RestrictionTrackerView)
+ inner_tracker = _unwrap_tracker(tracker)
+ if inner_tracker is not None and self._reader_cache is not None:
+ # Let this bundle reclaim a reader parked by the prior bundle and re-park
+ # it on self-checkpoint. No cache means setup() was skipped.
+ inner_tracker._reader_cache = self._reader_cache
+ initial = tracker.current_restriction()
+ now = self._now or time.monotonic
+ records_emitted = 0
+ # Armed on the first emitted record so reader startup is excluded.
+ read_deadline = None # type: Optional[float]
+ try:
+ while True:
+ holder = [None]
+ if not tracker.try_claim(holder):
+ # EOF: advance the estimator to the tracker's MAX watermark so
+ # downstream event-time windows can close.
+ _set_watermark_if_greater(
+ watermark_estimator, tracker.current_restriction().watermark)
+ break
+ record = holder[0]
+ if record is _NO_DATA:
+ # No data now: advance the watermark and self-checkpoint with the
+ # poll delay so an idle source backs off before resuming.
+ _set_watermark_if_greater(
+ watermark_estimator, tracker.current_restriction().watermark)
+ tracker.defer_remainder(Duration(seconds=self._poll_interval))
+ break
+ # The third tuple field is the source watermark. The record timestamp
+ # remains the output event time. Emit the element before advancing the
+ # estimator so a reader that reports MAX on the same claim as its final
+ # record cannot push the output watermark past that record first.
+ value, record_timestamp, source_watermark = record
+ yield TimestampedValue(value, record_timestamp)
+ _set_watermark_if_greater(watermark_estimator, source_watermark)
+ records_emitted += 1
+ if read_deadline is None:
+ read_deadline = now() + self._max_read_time_seconds
+ # A busy reader never hits the EOF or no-data branch. Bound the bundle
+ # by record count and elapsed time so the runner commits the checkpoint
+ # and runs finalization, then resume with no delay. The deadline is
+ # checked between records; a reader that blocks inside advance() can
+ # overrun it, so the record cap is the hard backstop.
+ reached_record_cap = records_emitted >= self._max_records_per_bundle
+ if reached_record_cap or now() >= read_deadline:
+ tracker.defer_remainder()
+ break
+ finally:
+ current = tracker.current_restriction()
+ try:
+ # Register finalization only when a checkpoint was cut this bundle.
+ # The SDK bundle finalizer applies no deadline, so finalization is
+ # unbounded best effort.
+ finalize_mark = current.finalization_checkpoint_mark
+ if current is not initial and finalize_mark is not None:
+ bundle_finalizer.register(_FinalizeCheckpointOnce(finalize_mark))
+ finally:
+ # The EOF and self-checkpoint paths already closed or parked the
+ # reader, so this is a no-op there. It closes a reader still held when
+ # process() exits early, e.g. a downstream yield raised before any
+ # checkpoint.
+ if inner_tracker is not None:
+ inner_tracker._close_reader_if_open()
+ else:
+ _LOGGER.warning(
+ 'UnboundedSource DoFn could not close a reader because the SDF '
+ 'tracker wrapper did not expose '
+ '_UnboundedSourceRestrictionTracker (got %s). Reader resources '
+ 'may remain open until garbage collection.',
+ type(tracker).__name__)
+
+
+def _unwrap_tracker(
+ tracker: Any) -> Optional['_UnboundedSourceRestrictionTracker']:
+ """Returns the :class:`_UnboundedSourceRestrictionTracker` behind the SDF
+ view and threadsafe wrappers, or None when the chain is unexpected."""
+ inner = tracker
+ if hasattr(inner, '_threadsafe_restriction_tracker'):
+ inner = inner._threadsafe_restriction_tracker
+ if hasattr(inner, '_restriction_tracker'):
+ inner = inner._restriction_tracker
+ if isinstance(inner, _UnboundedSourceRestrictionTracker):
+ return inner
+ return None
+
+
+def _set_watermark_if_greater(
+ watermark_estimator, new_watermark: Timestamp) -> None:
+ # ManualWatermarkEstimator.set_watermark raises on regression, so only ever
+ # advance it (a regressing reader watermark is absorbed here).
+ current = watermark_estimator.current_watermark()
+ if current is None or new_watermark > current:
+ watermark_estimator.set_watermark(new_watermark)
+
+
+class ReadFromUnboundedSource(PTransform):
+ """Reads an :class:`UnboundedSource`.
+
+ Most users should prefer :class:`apache_beam.io.Read`, which dispatches an
+ ``UnboundedSource`` here automatically::
+
+ p | beam.io.Read(MyUnboundedSource())
+
+ Args:
+ source: the :class:`UnboundedSource` to read.
+ poll_interval: resume delay in seconds applied when the reader has no data,
+ which bounds how often an idle source is polled. Must be >= 0.
+ max_records_per_bundle: a busy reader self-checkpoints after emitting this
+ many records in one bundle. Must be >= 1. Defaults to 10000.
+ max_read_time_seconds: a busy reader self-checkpoints after this many
+ seconds in one bundle. Must be > 0. Defaults to 10.0. The deadline is
+ checked between records, so a reader that blocks inside ``advance()`` may
+ overrun it; ``max_records_per_bundle`` is the hard backstop.
+
+ The bundle self-checkpoints as soon as either cap is reached.
+ """
+ def __init__(
+ self,
+ source: UnboundedSource,
+ poll_interval: float = _DEFAULT_POLL_INTERVAL_SECONDS,
+ max_records_per_bundle: int = _DEFAULT_MAX_RECORDS_PER_BUNDLE,
+ max_read_time_seconds: float = _DEFAULT_MAX_READ_TIME_SECONDS):
+ if not isinstance(source, UnboundedSource):
+ raise TypeError('source must be an UnboundedSource, got %r' % (source, ))
+ if max_records_per_bundle < 1:
+ raise ValueError(
+ 'max_records_per_bundle must be >= 1, got %r' %
+ (max_records_per_bundle, ))
+ if max_read_time_seconds <= 0:
+ raise ValueError(
+ 'max_read_time_seconds must be > 0, got %r' %
+ (max_read_time_seconds, ))
+ if poll_interval < 0:
+ raise ValueError(
+ 'poll_interval must be >= 0, got %r' % (poll_interval, ))
+ super().__init__()
+ self._source = source
+ self._poll_interval = poll_interval
+ self._max_records_per_bundle = max_records_per_bundle
+ self._max_read_time_seconds = max_read_time_seconds
+
+ def expand(self, pbegin):
+ source = self._source
+ output_coder = source.default_output_coder()
+ # The source is the SDF element used to derive the initial restriction.
+ # process() reads from the restriction, so it does not use the element
+ # directly.
+ output = (
+ pbegin
+ | 'Create' >> core.Create([source])
+ | 'ReadUnbounded' >> core.ParDo(
+ _ReadFromUnboundedSourceDoFn(
+ self._poll_interval,
+ self._max_records_per_bundle,
+ self._max_read_time_seconds)))
+ # Surface an element type only when the global registry already maps it to
+ # an equivalent coder. Avoid mutating ``coders.registry`` for a
+ # parameterized coder whose instance state would be lost.
+ try:
+ type_hint = output_coder.to_type_hint()
+ except NotImplementedError:
+ type_hint = None
+ if type_hint is not None:
+ try:
+ registered_coder = coders.registry.get_coder(type_hint)
+ except Exception: # pylint: disable=broad-except
+ _LOGGER.warning(
+ 'Could not look up the registered coder for element type %s.',
+ type_hint,
+ exc_info=True)
+ else:
+ if registered_coder == output_coder:
+ output.element_type = type_hint
+ return output
+
+ def _infer_output_coder(self, input_type=None, input_coder=None):
+ return self._source.default_output_coder()
diff --git a/sdks/python/apache_beam/io/unbounded_source_test.py
b/sdks/python/apache_beam/io/unbounded_source_test.py
new file mode 100644
index 00000000000..3b63d0d2de9
--- /dev/null
+++ b/sdks/python/apache_beam/io/unbounded_source_test.py
@@ -0,0 +1,1380 @@
+#
+# 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.
+#
+
+"""Tests for apache_beam.io.unbounded_source.
+
+Semantics are covered by deterministic unit tests; the end-to-end DirectRunner
+tests assert ordering and termination only (no flaky defer-timing assertions).
+"""
+
+import logging
+import unittest
+
+from typing_extensions import override
+
+import apache_beam as beam
+from apache_beam import coders
+from apache_beam.io import unbounded_source as _unbounded_source_module
+from apache_beam.io.unbounded_source import _NO_DATA
+from apache_beam.io.unbounded_source import CheckpointMark
+from apache_beam.io.unbounded_source import ReadFromUnboundedSource
+from apache_beam.io.unbounded_source import UnboundedReader
+from apache_beam.io.unbounded_source import UnboundedSource
+from apache_beam.io.unbounded_source import _FinalizeCheckpointOnce
+from apache_beam.io.unbounded_source import _ReaderCache
+from apache_beam.io.unbounded_source import _ReadFromUnboundedSourceDoFn
+from apache_beam.io.unbounded_source import _set_watermark_if_greater
+from apache_beam.io.unbounded_source import _UnboundedSourceRestriction
+from apache_beam.io.unbounded_source import _UnboundedSourceRestrictionCoder
+from apache_beam.io.unbounded_source import _UnboundedSourceRestrictionProvider
+from apache_beam.io.unbounded_source import _UnboundedSourceRestrictionTracker
+from apache_beam.io.watermark_estimators import ManualWatermarkEstimator
+from apache_beam.runners import sdf_utils
+from apache_beam.testing.test_pipeline import TestPipeline
+from apache_beam.testing.util import assert_that
+from apache_beam.testing.util import equal_to
+from apache_beam.transforms import core
+from apache_beam.transforms.window import FixedWindows
+from apache_beam.utils.timestamp import MAX_TIMESTAMP
+from apache_beam.utils.timestamp import MIN_TIMESTAMP
+from apache_beam.utils.timestamp import Timestamp
+
+# pylint: disable=expression-not-assigned
+
+# Realistic event-time base away from the Unix epoch.
+_EVENT_TIME_BASE = Timestamp(1729987200) # 2024-10-27T00:00:00Z
+
+#
------------------------------------------------------------------------------
+# In-memory demo source emitting integers 0..count-1 with event time
+# ``_EVENT_TIME_BASE + index``. It self-terminates at EOF, resumes from
+# ``last_index + 1``, and splits into even/odd sub-sources when configured.
+#
------------------------------------------------------------------------------
+
+
+class _CountingCheckpointMark(CheckpointMark):
+ def __init__(self, last_index, finalize_log=None):
+ self.last_index = last_index
+ self._finalize_log = finalize_log
+
+ @override
+ def finalize_checkpoint(self):
+ if self._finalize_log is not None:
+ self._finalize_log.append(self.last_index)
+
+ def __eq__(self, other):
+ return (
+ isinstance(other, _CountingCheckpointMark) and
+ other.last_index == self.last_index)
+
+ def __hash__(self):
+ return hash(self.last_index)
+
+ def __repr__(self):
+ return '_CountingCheckpointMark(last_index=%r)' % (self.last_index, )
+
+
+class _CountingReader(UnboundedReader):
+ def __init__(
+ self, count, start_index, finalize_log=None, modulus=1, residue=0):
+ self._count = count
+ self._next = start_index
+ self._modulus = modulus
+ self._residue = residue
+ self._current = None
+ self._exhausted = False
+ self._finalize_log = finalize_log
+ self.closed = False
+
+ def _read_next(self):
+ while self._next < self._count:
+ index = self._next
+ self._next += 1
+ if index % self._modulus == self._residue:
+ self._current = index
+ return True
+ self._exhausted = True
+ return False
+
+ @override
+ def start(self):
+ return self._read_next()
+
+ @override
+ def advance(self):
+ return self._read_next()
+
+ @override
+ def get_current(self):
+ return self._current
+
+ @override
+ def get_current_timestamp(self):
+ return _EVENT_TIME_BASE + self._current
+
+ @override
+ def get_watermark(self):
+ if self._exhausted:
+ return MAX_TIMESTAMP
+ if self._current is None:
+ return MIN_TIMESTAMP
+ return _EVENT_TIME_BASE + self._current
+
+ @override
+ def get_checkpoint_mark(self):
+ last = self._current if self._current is not None else self._next - 1
+ return _CountingCheckpointMark(last, finalize_log=self._finalize_log)
+
+ @override
+ def close(self):
+ self.closed = True
+
+
+class UnboundedCountingSource(UnboundedSource):
+ def __init__(
+ self,
+ count,
+ finalize_log=None,
+ is_splittable=False,
+ modulus=1,
+ residue=0):
+ self._count = count
+ self._finalize_log = finalize_log
+ self._is_splittable = is_splittable
+ self._modulus = modulus
+ self._residue = residue
+ self.last_reader = None
+
+ @override
+ def split(self, desired_num_splits, options=None):
+ if not self._is_splittable or desired_num_splits < 2:
+ return [self]
+ # Split into independent even/odd sub-sources (each non-splittable).
+ return [
+ UnboundedCountingSource(
+ self._count,
+ finalize_log=self._finalize_log,
+ modulus=2,
+ residue=residue) for residue in (0, 1)
+ ]
+
+ @override
+ def create_reader(self, options, checkpoint_mark):
+ start_index = (
+ 0 if checkpoint_mark is None else checkpoint_mark.last_index + 1)
+ self.last_reader = _CountingReader(
+ self._count,
+ start_index,
+ finalize_log=self._finalize_log,
+ modulus=self._modulus,
+ residue=self._residue)
+ return self.last_reader
+
+ @override
+ def get_checkpoint_mark_coder(self):
+ return coders.PickleCoder()
+
+
+class _StringCountingReader(_CountingReader):
+ @override
+ def get_current(self):
+ return 'v%s' % self._current
+
+
+class _StringCountingSource(UnboundedCountingSource):
+ @override
+ def create_reader(self, options, checkpoint_mark):
+ start_index = (
+ 0 if checkpoint_mark is None else checkpoint_mark.last_index + 1)
+ self.last_reader = _StringCountingReader(
+ self._count, start_index, finalize_log=self._finalize_log)
+ return self.last_reader
+
+ @override
+ def default_output_coder(self):
+ return coders.StrUtf8Coder()
+
+
+class _PrefixStrCoder(coders.Coder):
+ def __init__(self, prefix):
+ self._prefix = prefix
+
+ @override
+ def encode(self, value):
+ if not value.startswith(self._prefix):
+ raise ValueError('expected %r prefix' % self._prefix)
+ return value[len(self._prefix):].encode('utf-8')
+
+ @override
+ def decode(self, value):
+ return self._prefix + value.decode('utf-8')
+
+ @override
+ def is_deterministic(self):
+ return True
+
+ @override
+ def to_type_hint(self):
+ return str
+
+
+class _PrefixStringReader(_StringCountingReader):
+ @override
+ def get_current(self):
+ return 'prefix:%s' % super().get_current()
+
+
+class _PrefixStringSource(_StringCountingSource):
+ @override
+ def create_reader(self, options, checkpoint_mark):
+ start_index = (
+ 0 if checkpoint_mark is None else checkpoint_mark.last_index + 1)
+ self.last_reader = _PrefixStringReader(
+ self._count, start_index, finalize_log=self._finalize_log)
+ return self.last_reader
+
+ @override
+ def default_output_coder(self):
+ return _PrefixStrCoder('prefix:')
+
+
+class _NoDataReader(UnboundedReader):
+ """Always reports temporary absence of data with watermark below MAX."""
+ @override
+ def start(self):
+ return False
+
+ @override
+ def advance(self):
+ return False
+
+ @override
+ def get_current(self):
+ raise AssertionError('no data available')
+
+ @override
+ def get_current_timestamp(self):
+ raise AssertionError('no data available')
+
+ @override
+ def get_watermark(self):
+ return _EVENT_TIME_BASE
+
+ @override
+ def get_checkpoint_mark(self):
+ return _CountingCheckpointMark(-1)
+
+
+class _NoDataSource(UnboundedSource):
+ @override
+ def split(self, desired_num_splits, options=None):
+ return [self]
+
+ @override
+ def create_reader(self, options, checkpoint_mark):
+ return _NoDataReader()
+
+ @override
+ def get_checkpoint_mark_coder(self):
+ return coders.PickleCoder()
+
+
+class _MutatingCheckpointMark(CheckpointMark):
+ """A mark that mutates itself on finalize, to test primary/residual mark
+ isolation across a checkpoint cut."""
+ def __init__(self, last_index):
+ self.last_index = last_index
+
+ @override
+ def finalize_checkpoint(self):
+ self.last_index = -999
+
+
+class _MutatingReader(UnboundedReader):
+ def __init__(self):
+ self._index = -1
+
+ @override
+ def start(self):
+ self._index = 0
+ return True
+
+ @override
+ def advance(self):
+ self._index += 1
+ return True
+
+ @override
+ def get_current(self):
+ return self._index
+
+ @override
+ def get_current_timestamp(self):
+ return _EVENT_TIME_BASE + self._index
+
+ @override
+ def get_watermark(self):
+ return _EVENT_TIME_BASE + self._index
+
+ @override
+ def get_checkpoint_mark(self):
+ return _MutatingCheckpointMark(self._index)
+
+
+class _MutatingSource(UnboundedSource):
+ @override
+ def split(self, desired_num_splits, options=None):
+ return [self]
+
+ @override
+ def create_reader(self, options, checkpoint_mark):
+ return _MutatingReader()
+
+ @override
+ def get_checkpoint_mark_coder(self):
+ return coders.PickleCoder()
+
+
+class _MaxOnLastReader(UnboundedReader):
+ """Returns its only record with a MAX_TIMESTAMP watermark on the same claim,
+ then EOF on the next claim."""
+ @override
+ def start(self):
+ return True
+
+ @override
+ def advance(self):
+ return False
+
+ @override
+ def get_current(self):
+ return 7
+
+ @override
+ def get_current_timestamp(self):
+ return _EVENT_TIME_BASE
+
+ @override
+ def get_watermark(self):
+ return MAX_TIMESTAMP
+
+ @override
+ def get_checkpoint_mark(self):
+ return _CountingCheckpointMark(0)
+
+
+class _MaxOnLastSource(UnboundedSource):
+ @override
+ def split(self, desired_num_splits, options=None):
+ return [self]
+
+ @override
+ def create_reader(self, options, checkpoint_mark):
+ return _MaxOnLastReader()
+
+ @override
+ def get_checkpoint_mark_coder(self):
+ return coders.PickleCoder()
+
+
+def _new_tracker(source, checkpoint=None):
+ restriction = _UnboundedSourceRestriction(
+ source=source, checkpoint_mark=checkpoint)
+ return _UnboundedSourceRestrictionTracker(restriction)
+
+
+def _claim(tracker):
+ """Claims once; returns (claimed_bool, holder_value)."""
+ holder = [None]
+ claimed = tracker.try_claim(holder)
+ return claimed, holder[0]
+
+
+#
------------------------------------------------------------------------------
+# Tests
+#
------------------------------------------------------------------------------
+
+
+class AbcContractTest(unittest.TestCase):
+ def test_checkpointmark_default_finalize_is_noop(self):
+ self.assertIsNone(CheckpointMark().finalize_checkpoint())
+
+ def test_unboundedsource_is_bounded_false(self):
+ self.assertFalse(UnboundedCountingSource(3).is_bounded())
+
+ def test_reader_lifecycle_start_advance_eof(self):
+ reader = UnboundedCountingSource(3).create_reader(None, None)
+ self.assertTrue(reader.start())
+ self.assertEqual(reader.get_current(), 0)
+ self.assertEqual(reader.get_current_timestamp(), _EVENT_TIME_BASE)
+ self.assertTrue(reader.advance())
+ self.assertEqual(reader.get_current(), 1)
+ self.assertTrue(reader.advance())
+ self.assertEqual(reader.get_current(), 2)
+ self.assertFalse(reader.advance())
+ self.assertEqual(reader.get_watermark(), MAX_TIMESTAMP)
+
+
+class RestrictionCoderTest(unittest.TestCase):
+ def test_roundtrip_no_checkpoint(self):
+ source = UnboundedCountingSource(3)
+ coder = _UnboundedSourceRestrictionCoder()
+ decoded = coder.decode(
+ coder.encode(_UnboundedSourceRestriction(source=source)))
+ self.assertIsNone(decoded.checkpoint_mark)
+ self.assertEqual(decoded.watermark, MIN_TIMESTAMP)
+ self.assertFalse(decoded.is_done)
+ reader = decoded.source.create_reader(None, None)
+ self.assertTrue(reader.start())
+ self.assertEqual(reader.get_current(), 0)
+
+ def test_roundtrip_with_checkpoint_resumes(self):
+ source = UnboundedCountingSource(5)
+ coder = _UnboundedSourceRestrictionCoder()
+ restriction = _UnboundedSourceRestriction(
+ source=source,
+ checkpoint_mark=_CountingCheckpointMark(1),
+ watermark=Timestamp(1),
+ is_done=False)
+ decoded = coder.decode(coder.encode(restriction))
+ self.assertEqual(decoded.checkpoint_mark.last_index, 1)
+ self.assertEqual(decoded.watermark, Timestamp(1))
+ self.assertFalse(decoded.is_done)
+ # A reader built from the decoded checkpoint resumes at the next index.
+ reader = decoded.source.create_reader(None, decoded.checkpoint_mark)
+ self.assertTrue(reader.start())
+ self.assertEqual(reader.get_current(), 2)
+
+
+class RestrictionProviderTest(unittest.TestCase):
+ def test_initial_split_calls_source_split(self):
+ split_log = []
+
+ class _NamedSource(UnboundedCountingSource):
+ def __init__(self, name):
+ super().__init__(0)
+ self.name = name
+
+ @override
+ def split(self, desired_num_splits, options=None):
+ split_log.append((desired_num_splits, options))
+ return [_NamedSource('a'), _NamedSource('b')]
+
+ source = _NamedSource('root')
+ provider = _UnboundedSourceRestrictionProvider()
+ restriction = _UnboundedSourceRestriction(
+ source=source, watermark=Timestamp(7))
+
+ splits = list(provider.split(source, restriction))
+
+ # The provider is a stateless module-level singleton, so it always
+ # passes ``None`` as the ``options`` argument to ``UnboundedSource.split``.
+ self.assertEqual(split_log, [(20, None)])
+ self.assertEqual([split.source.name for split in splits], ['a', 'b'])
+ self.assertEqual([split.watermark for split in splits], [Timestamp(7)] * 2)
+ self.assertTrue(all(split.checkpoint_mark is None for split in splits))
+ self.assertTrue(
+ all(split.finalization_checkpoint_mark is None for split in splits))
+
+ def test_initial_split_does_not_split_checkpointed_restriction(self):
+ split_log = []
+
+ class _SplitSource(UnboundedCountingSource):
+ @override
+ def split(self, desired_num_splits, options=None):
+ split_log.append((desired_num_splits, options))
+ return [self]
+
+ source = _SplitSource(5)
+ provider = _UnboundedSourceRestrictionProvider()
+ restriction = _UnboundedSourceRestriction(
+ source=source, checkpoint_mark=_CountingCheckpointMark(2))
+
+ self.assertEqual(list(provider.split(source, restriction)), [restriction])
+ self.assertEqual(split_log, [])
+
+ def test_initial_split_falls_back_to_original_on_split_error(self):
+ class _BoomSource(UnboundedCountingSource):
+ @override
+ def split(self, desired_num_splits, options=None):
+ raise RuntimeError('split boom')
+
+ source = _BoomSource(5)
+ provider = _UnboundedSourceRestrictionProvider()
+ restriction = _UnboundedSourceRestriction(source=source)
+
+ self.assertEqual(list(provider.split(source, restriction)), [restriction])
+
+ def test_truncate_returns_none_for_drain(self):
+ # On drain the SDF stops emitting; truncate yields no residual.
+ provider = _UnboundedSourceRestrictionProvider()
+ source = UnboundedCountingSource(5)
+ restriction = _UnboundedSourceRestriction(source=source)
+ self.assertIsNone(provider.truncate(source, restriction))
+
+ def test_splittable_source_partitions_into_independent_subsources(self):
+ # A splittable source fans out into two sub-sources; reading each in
+ # isolation yields the even and the odd integers, and their union is the
+ # full sequence with no overlap.
+ source = UnboundedCountingSource(6, is_splittable=True)
+ provider = _UnboundedSourceRestrictionProvider()
+ restriction = _UnboundedSourceRestriction(source=source)
+
+ splits = list(provider.split(source, restriction))
+ self.assertEqual(len(splits), 2)
+
+ shards = []
+ for split in splits:
+ tracker = _UnboundedSourceRestrictionTracker(split)
+ shard = []
+ while True:
+ claimed, record = _claim(tracker)
+ if not claimed:
+ break
+ if record is not _NO_DATA:
+ shard.append(record[0])
+ shards.append(shard)
+ self.assertEqual(sorted(shards), [[0, 2, 4], [1, 3, 5]])
+
+
+class RestrictionTrackerTest(unittest.TestCase):
+ def test_claim_emits_in_order(self):
+ tracker = _new_tracker(UnboundedCountingSource(3))
+ values = []
+ while True:
+ claimed, record = _claim(tracker)
+ if not claimed:
+ break
+ self.assertIsNot(record, _NO_DATA)
+ values.append(record[0])
+ self.assertEqual(values, [0, 1, 2])
+ self.assertTrue(tracker.check_done())
+
+ def test_claim_emits_final_record_when_watermark_is_max(self):
+ # A reader may return its last record with a MAX_TIMESTAMP watermark on the
+ # same call; the record must still be emitted (EOF comes on the next
claim).
+ class _FinalRecordReader(UnboundedReader):
+ @override
+ def start(self):
+ return True
+
+ @override
+ def advance(self):
+ return False
+
+ @override
+ def get_current(self):
+ return 'last'
+
+ @override
+ def get_current_timestamp(self):
+ return _EVENT_TIME_BASE
+
+ @override
+ def get_watermark(self):
+ return MAX_TIMESTAMP
+
+ @override
+ def get_checkpoint_mark(self):
+ return _CountingCheckpointMark(0)
+
+ class _FinalSource(UnboundedSource):
+ @override
+ def split(self, desired_num_splits, options=None):
+ return [self]
+
+ @override
+ def create_reader(self, options, checkpoint_mark):
+ return _FinalRecordReader()
+
+ @override
+ def get_checkpoint_mark_coder(self):
+ return coders.PickleCoder()
+
+ tracker = _new_tracker(_FinalSource())
+ claimed, record = _claim(tracker)
+ self.assertTrue(claimed)
+ self.assertIsNot(record, _NO_DATA)
+ self.assertEqual(record[0], 'last')
+ # The next claim observes EOF and finishes (no second, phantom record).
+ claimed_again, _ = _claim(tracker)
+ self.assertFalse(claimed_again)
+ self.assertTrue(tracker.check_done())
+
+ def test_try_split_zero_produces_resumable_residual(self):
+ source = UnboundedCountingSource(5)
+ tracker = _new_tracker(source)
+ # Claim 0 and 1.
+ self.assertEqual(_claim(tracker)[1][0], 0)
+ self.assertEqual(_claim(tracker)[1][0], 1)
+
+ split = tracker.try_split(0)
+ self.assertIsNotNone(split)
+ primary, residual = split
+ self.assertTrue(primary.is_done)
+ self.assertFalse(residual.is_done)
+ # Resume / finalize channel separation: primary carries only the
+ # finalize hook, residual carries only the resume state.
+ self.assertIsNone(primary.checkpoint_mark)
+ self.assertIsNotNone(primary.finalization_checkpoint_mark)
+ self.assertEqual(primary.finalization_checkpoint_mark.last_index, 1)
+ self.assertEqual(residual.checkpoint_mark.last_index, 1)
+ self.assertIsNone(residual.finalization_checkpoint_mark)
+ # check_done passes on the (now done) primary.
+ self.assertTrue(tracker.check_done())
+
+ # Resuming from the residual continues at index 2.
+ resumed = _new_tracker(source, checkpoint=residual.checkpoint_mark)
+ self.assertEqual(_claim(resumed)[1][0], 2)
+
+ def test_try_split_isolates_residual_from_finalize_mutation(self):
+ # The primary's finalize hook and the residual's resume state must not
+ # share one object, so a mutating finalize_checkpoint() cannot corrupt the
+ # residual's resume position.
+ tracker = _new_tracker(_MutatingSource())
+ _claim(tracker) # start -> index 0
+ split = tracker.try_split(0)
+ self.assertIsNotNone(split)
+ primary, residual = split
+ self.assertIsNot(
+ primary.finalization_checkpoint_mark, residual.checkpoint_mark)
+ self.assertEqual(residual.checkpoint_mark.last_index, 0)
+ primary.finalization_checkpoint_mark.finalize_checkpoint()
+ self.assertEqual(residual.checkpoint_mark.last_index, 0)
+
+ def test_try_split_nonzero_declined(self):
+ source = UnboundedCountingSource(5)
+ tracker = _new_tracker(source)
+ self.assertEqual(_claim(tracker)[1][0], 0)
+
+ self.assertIsNone(tracker.try_split(0.5))
+ self.assertFalse(tracker.current_restriction().is_done)
+ self.assertIsNotNone(tracker._reader)
+ self.assertEqual(_claim(tracker)[1][0], 1)
+
+ def test_no_data_returns_sentinel_without_finishing(self):
+ tracker = _new_tracker(_NoDataSource())
+ claimed, record = _claim(tracker)
+ self.assertTrue(claimed)
+ self.assertIs(record, _NO_DATA)
+ # A self-checkpoint is still possible (poll/resume path).
+ self.assertIsNotNone(tracker.try_split(0))
+
+ def test_check_done_raises_when_not_done(self):
+ tracker = _new_tracker(UnboundedCountingSource(3))
+ with self.assertRaises(ValueError):
+ tracker.check_done()
+
+ def test_is_bounded_false(self):
+ self.assertFalse(_new_tracker(UnboundedCountingSource(3)).is_bounded())
+
+
+class _RecordingBundleFinalizer:
+ def __init__(self):
+ self.registered = []
+
+ def register(self, callback):
+ self.registered.append(callback)
+
+
+class _ManualClock:
+ """A deterministic monotonic clock for the time-cap tests."""
+ def __init__(self, now=0.0):
+ self.now = now
+
+ def __call__(self):
+ return self.now
+
+
+class BundleCapTest(unittest.TestCase):
+ """A busy reader self-checkpoints once the per-bundle record or time cap is
+ reached, so the runner can commit progress and run finalization."""
+ def _bundle(self, dofn, source, checkpoint=None, estimator=None):
+ """Builds the SDF tracker chain and returns the process() generator plus
the
+ tracker, threadsafe tracker, finalizer, and watermark estimator."""
+ tracker = _UnboundedSourceRestrictionTracker(
+ _UnboundedSourceRestriction(source=source, checkpoint_mark=checkpoint))
+ threadsafe = sdf_utils.ThreadsafeRestrictionTracker(tracker)
+ view = sdf_utils.RestrictionTrackerView(threadsafe)
+ finalizer = _RecordingBundleFinalizer()
+ estimator = estimator or ManualWatermarkEstimator(None)
+ gen = dofn.process(
+ None,
+ bundle_finalizer=finalizer,
+ tracker=view,
+ watermark_estimator=estimator)
+ return gen, tracker, threadsafe, finalizer, estimator
+
+ def test_record_cap_checkpoints_busy_source(self):
+ finalize_log = []
+ dofn = _ReadFromUnboundedSourceDoFn(
+ poll_interval=0, max_records_per_bundle=5, max_read_time_seconds=1e9)
+ # 1000 records is effectively unbounded against a cap of 5.
+ gen, tracker, threadsafe, finalizer, estimator = self._bundle(
+ dofn, UnboundedCountingSource(1000, finalize_log=finalize_log))
+ outputs = list(gen)
+
+ self.assertEqual([tv.value for tv in outputs], [0, 1, 2, 3, 4])
+ self.assertTrue(tracker.current_restriction().is_done)
+ self.assertTrue(tracker.check_done())
+ # The estimator holds the last emitted record's source watermark.
+ self.assertEqual(estimator.current_watermark(), _EVENT_TIME_BASE + 4)
+ # Residual resumes after the cut and carries no finalize hook.
+ residual, _ = threadsafe.deferred_status()
+ self.assertEqual(residual.checkpoint_mark.last_index, 4)
+ self.assertIsNone(residual.finalization_checkpoint_mark)
+ # Exactly one finalizer is registered; firing it commits the cut index
once.
+ self.assertEqual(len(finalizer.registered), 1)
+ finalizer.registered[0]()
+ finalizer.registered[0]()
+ self.assertEqual(finalize_log, [4])
+
+ def test_time_cap_checkpoints_busy_source(self):
+ clock = _ManualClock(1000.0)
+ dofn = _ReadFromUnboundedSourceDoFn(
+ poll_interval=0,
+ max_records_per_bundle=10**9,
+ max_read_time_seconds=5.0,
+ _now=clock)
+ gen, tracker, threadsafe, _, _ = self._bundle(
+ dofn, UnboundedCountingSource(1000))
+
+ # The deadline arms at 1000 + 5 after the first record and is checked
+ # between records, so records keep flowing until the clock passes it.
+ self.assertEqual(next(gen).value, 0)
+ self.assertEqual(next(gen).value, 1)
+ self.assertEqual(next(gen).value, 2)
+ clock.now = 1006.0
+ with self.assertRaises(StopIteration):
+ next(gen)
+
+ self.assertTrue(tracker.current_restriction().is_done)
+ residual, _ = threadsafe.deferred_status()
+ self.assertEqual(residual.checkpoint_mark.last_index, 2)
+
+ def test_cap_residual_resumes_in_next_bundle(self):
+ dofn = _ReadFromUnboundedSourceDoFn(
+ poll_interval=0, max_records_per_bundle=5, max_read_time_seconds=1e9)
+ source = UnboundedCountingSource(1000)
+ # Bundle 1 emits 0-4 and cuts a residual at index 4.
+ gen1, _, threadsafe1, _, _ = self._bundle(dofn, source)
+ self.assertEqual([tv.value for tv in gen1], [0, 1, 2, 3, 4])
+ residual1, _ = threadsafe1.deferred_status()
+
+ # Bundle 2 rebuilds the reader from the residual and emits 5-9.
+ gen2, _, threadsafe2, _, _ = self._bundle(
+ dofn, source, checkpoint=residual1.checkpoint_mark)
+ self.assertEqual([tv.value for tv in gen2], [5, 6, 7, 8, 9])
+ residual2, _ = threadsafe2.deferred_status()
+ self.assertEqual(residual2.checkpoint_mark.last_index, 9)
+
+ def test_busy_reader_is_reused_across_bundles(self):
+ # The self-checkpoint parks the reader; the resuming bundle reclaims the
+ # same started reader.
+ dofn = _ReadFromUnboundedSourceDoFn(
+ poll_interval=0, max_records_per_bundle=5, max_read_time_seconds=1e9)
+ dofn.setup() # creates the cross-bundle reader cache
+ source = UnboundedCountingSource(1000)
+ gen1, _, threadsafe1, _, _ = self._bundle(dofn, source)
+ self.assertEqual([tv.value for tv in gen1], [0, 1, 2, 3, 4])
+ reader1 = source.last_reader
+ self.assertFalse(reader1.closed) # parked, not closed
+ residual1, _ = threadsafe1.deferred_status()
+
+ gen2, _, _, _, _ = self._bundle(
+ dofn, source, checkpoint=residual1.checkpoint_mark)
+ self.assertEqual([tv.value for tv in gen2], [5, 6, 7, 8, 9])
+ # No new reader was created; source.last_reader is unchanged.
+ self.assertIs(source.last_reader, reader1)
+
+ def test_teardown_closes_parked_readers(self):
+ dofn = _ReadFromUnboundedSourceDoFn(
+ poll_interval=0, max_records_per_bundle=5, max_read_time_seconds=1e9)
+ dofn.setup()
+ source = UnboundedCountingSource(1000)
+ gen, _, _, _, _ = self._bundle(dofn, source)
+ list(gen) # the bundle parks its reader
+ reader = source.last_reader
+ self.assertFalse(reader.closed)
+ dofn.teardown()
+ self.assertTrue(reader.closed)
+
+ def test_eof_exactly_at_cap_resumes_then_finishes(self):
+ dofn = _ReadFromUnboundedSourceDoFn(
+ poll_interval=0, max_records_per_bundle=5, max_read_time_seconds=1e9)
+ source = UnboundedCountingSource(5) # exactly cap records
+ # Bundle 1 hits the cap on the last record before observing EOF.
+ gen1, t1, threadsafe1, _, _ = self._bundle(dofn, source)
+ self.assertEqual([tv.value for tv in gen1], [0, 1, 2, 3, 4])
+ self.assertTrue(t1.current_restriction().is_done)
+ residual1, _ = threadsafe1.deferred_status()
+ self.assertEqual(residual1.checkpoint_mark.last_index, 4)
+
+ # Bundle 2 resumes at index 5, finds EOF, and finishes with no output.
+ gen2, t2, threadsafe2, _, _ = self._bundle(
+ dofn, source, checkpoint=residual1.checkpoint_mark)
+ self.assertEqual(list(gen2), [])
+ self.assertTrue(t2.current_restriction().is_done)
+ self.assertIsNone(threadsafe2.deferred_status())
+
+ def test_eof_before_cap_finishes_without_residual(self):
+ dofn = _ReadFromUnboundedSourceDoFn(
+ poll_interval=0, max_records_per_bundle=100, max_read_time_seconds=1e9)
+ gen, tracker, threadsafe, _, _ = self._bundle(
+ dofn, UnboundedCountingSource(3))
+
+ self.assertEqual([tv.value for tv in gen], [0, 1, 2])
+ self.assertTrue(tracker.current_restriction().is_done)
+ self.assertIsNone(threadsafe.deferred_status())
+
+ def test_max_watermark_on_final_record_emitted_before_estimator_advances(
+ self):
+ # A reader that returns its final record with a MAX_TIMESTAMP watermark on
+ # the same claim must have that record emitted before the estimator reaches
+ # MAX, so the element is not stranded behind the output watermark.
+ dofn = _ReadFromUnboundedSourceDoFn(poll_interval=0)
+ gen, _, _, _, estimator = self._bundle(dofn, _MaxOnLastSource())
+
+ first = next(gen)
+ self.assertEqual(first.value, 7)
+ # The estimator has not yet been advanced to MAX when the record is
yielded.
+ self.assertNotEqual(estimator.current_watermark(), MAX_TIMESTAMP)
+ # Draining the bundle then advances the estimator to MAX on EOF.
+ self.assertEqual(list(gen), [])
+ self.assertEqual(estimator.current_watermark(), MAX_TIMESTAMP)
+
+
+class WatermarkTest(unittest.TestCase):
+ def test_set_watermark_is_monotonic(self):
+ estimator = ManualWatermarkEstimator(None)
+ _set_watermark_if_greater(estimator, Timestamp(5))
+ self.assertEqual(estimator.current_watermark(), Timestamp(5))
+ # A regression is ignored (would otherwise raise inside set_watermark).
+ _set_watermark_if_greater(estimator, Timestamp(3))
+ self.assertEqual(estimator.current_watermark(), Timestamp(5))
+ _set_watermark_if_greater(estimator, Timestamp(7))
+ self.assertEqual(estimator.current_watermark(), Timestamp(7))
+
+
+class FinalizationTest(unittest.TestCase):
+ def test_finalize_checkpoint_callback_is_at_most_once(self):
+ finalize_log = []
+ finalize_once = _FinalizeCheckpointOnce(
+ _CountingCheckpointMark(1, finalize_log=finalize_log))
+
+ finalize_once()
+ finalize_once()
+
+ self.assertEqual(finalize_log, [1])
+
+ def test_finalize_checkpoint_invoked(self):
+ # Unit-level finalize test (the e2e finalize may run in a worker process);
+ # the hook lives on the primary, independent of the residual's resume
state.
+ finalize_log = []
+ source = UnboundedCountingSource(5, finalize_log=finalize_log)
+ tracker = _new_tracker(source)
+ _claim(tracker) # 0
+ _claim(tracker) # 1
+ primary, _ = tracker.try_split(0)
+ primary.finalization_checkpoint_mark.finalize_checkpoint()
+ self.assertEqual(finalize_log, [1])
+
+
+class EndToEndTest(unittest.TestCase):
+ def test_direct_runner_emits_all_in_order(self):
+ with TestPipeline() as p:
+ out = p | ReadFromUnboundedSource(UnboundedCountingSource(5))
+ self.assertFalse(out.is_bounded)
+ assert_that(out, equal_to([0, 1, 2, 3, 4]))
+
+ def test_eof_lets_event_time_window_fire(self):
+ # On EOF the DoFn advances the watermark estimator to MAX_TIMESTAMP so the
+ # downstream FixedWindow closes and the GroupByKey fires; otherwise the
+ # output would be empty.
+ with TestPipeline() as p:
+ out = (
+ p
+ | ReadFromUnboundedSource(UnboundedCountingSource(5))
+ | beam.WindowInto(FixedWindows(100))
+ | beam.Map(lambda v: ('all', v))
+ | beam.GroupByKey()
+ | beam.MapTuple(lambda _key, values: sorted(values)))
+ assert_that(out, equal_to([[0, 1, 2, 3, 4]]))
+
+ def test_read_dispatches_through_iobase_read(self):
+ # ``beam.io.Read(source)`` must produce the same records as
+ # ``ReadFromUnboundedSource(source)``.
+ with TestPipeline() as p:
+ out = p | beam.io.Read(UnboundedCountingSource(5))
+ self.assertFalse(out.is_bounded)
+ assert_that(out, equal_to([0, 1, 2, 3, 4]))
+
+ def test_splittable_source_reads_all_records_across_splits(self):
+ # A splittable source fans out into even/odd sub-sources during initial
+ # SDF splitting; the union of all sub-source reads is the full sequence.
+ with TestPipeline() as p:
+ out = p | beam.io.Read(UnboundedCountingSource(6, is_splittable=True))
+ assert_that(out, equal_to([0, 1, 2, 3, 4, 5]))
+
+ def test_source_default_output_coder_sets_output_type(self):
+ with TestPipeline() as p:
+ out = p | ReadFromUnboundedSource(_StringCountingSource(2))
+ self.assertEqual(out.element_type, str)
+ assert_that(out, equal_to(['v0', 'v1']))
+
+ def test_small_cap_self_checkpoints_and_resumes_to_eof(self):
+ # A small per-bundle cap forces several self-checkpoint/resume cycles
+ # through the runner before EOF, exercising the cross-bundle reader cache
+ # over real residual encode/decode. All records still arrive in order.
+ with TestPipeline() as p:
+ out = p | ReadFromUnboundedSource(
+ UnboundedCountingSource(20), max_records_per_bundle=5)
+ assert_that(out, equal_to(list(range(20))))
+
+
+class ReadFromUnboundedSourceCoderTest(unittest.TestCase):
+ def test_parameterized_output_coder_does_not_mutate_global_registry(self):
+ try:
+ p = beam.Pipeline()
+ out = p | ReadFromUnboundedSource(_PrefixStringSource(1))
+
+ self.assertNotEqual(out.element_type, str)
+ self.assertEqual(coders.registry.get_coder(str), coders.StrUtf8Coder())
+ self.assertEqual(
+
ReadFromUnboundedSource(_PrefixStringSource(1))._infer_output_coder(),
+ _PrefixStrCoder('prefix:'))
+ finally:
+ coders.registry.register_coder(str, coders.StrUtf8Coder)
+
+
+#
------------------------------------------------------------------------------
+# Reader lifecycle, watermark, and contract regression tests (reader close on
+# every exit path, the NotImplementedError message, finalize idempotency).
+#
------------------------------------------------------------------------------
+
+
+class ReaderCloseTest(unittest.TestCase):
+ """Reader lifecycle: close() must run on every tracker-driven exit path."""
+ def test_tracker_closes_reader_on_eof(self):
+ source = UnboundedCountingSource(0) # immediately exhausted
+ tracker = _new_tracker(source)
+ holder = [None]
+ self.assertFalse(tracker.try_claim(holder))
+ self.assertIsNone(tracker._reader)
+ self.assertTrue(source.last_reader.closed)
+
+ def test_tracker_closes_reader_on_split_without_cache(self):
+ # With no cache injected, a self-checkpoint closes the reader.
+ source = UnboundedCountingSource(5)
+ tracker = _new_tracker(source)
+ _claim(tracker) # creates reader, claims 0
+ reader = source.last_reader
+ self.assertFalse(reader.closed)
+ split = tracker.try_split(0)
+ self.assertIsNotNone(split)
+ self.assertIsNone(tracker._reader)
+ self.assertTrue(reader.closed)
+
+ def test_tracker_parks_reader_on_split_with_cache(self):
+ # With a cache, a self-checkpoint parks the reader for the residual to
+ # reclaim.
+ source = UnboundedCountingSource(5)
+ cache = _ReaderCache()
+ tracker = _new_tracker(source)
+ tracker._reader_cache = cache
+ _claim(tracker)
+ reader = source.last_reader
+ _, residual = tracker.try_split(0)
+ self.assertIsNone(tracker._reader)
+ self.assertFalse(reader.closed)
+ # The residual's key reclaims the same started reader.
+ self.assertEqual(
+ cache.acquire(tracker._cache_key(residual)), (reader, True))
+
+ def test_close_helper_is_idempotent_and_safe_on_empty_tracker(self):
+ tracker = _new_tracker(UnboundedCountingSource(3))
+ # No reader yet -- helper must be a no-op.
+ tracker._close_reader_if_open()
+ _claim(tracker)
+ reader = tracker._reader
+ tracker._close_reader_if_open()
+ self.assertTrue(reader.closed)
+ self.assertIsNone(tracker._reader)
+ # Second call is a no-op (no reader to close).
+ tracker._close_reader_if_open()
+
+ def test_close_helper_swallows_reader_close_errors(self):
+ class _BoomReader(UnboundedReader):
+ @override
+ def start(self):
+ return True
+
+ @override
+ def advance(self):
+ return False
+
+ @override
+ def get_current(self):
+ return 'x'
+
+ @override
+ def get_current_timestamp(self):
+ return _EVENT_TIME_BASE
+
+ @override
+ def get_watermark(self):
+ return _EVENT_TIME_BASE
+
+ @override
+ def get_checkpoint_mark(self):
+ return CheckpointMark()
+
+ @override
+ def close(self):
+ raise RuntimeError('close blew up')
+
+ class _BoomSource(UnboundedSource):
+ @override
+ def split(self, desired_num_splits, options=None):
+ return [self]
+
+ @override
+ def create_reader(self, options, checkpoint_mark):
+ return _BoomReader()
+
+ @override
+ def get_checkpoint_mark_coder(self):
+ return coders.PickleCoder()
+
+ tracker = _new_tracker(_BoomSource())
+ _claim(tracker)
+ with self.assertLogs(_unbounded_source_module._LOGGER, 'WARNING') as logs:
+ tracker._close_reader_if_open()
+ self.assertTrue(
+ any('Error closing UnboundedReader' in line for line in logs.output))
+ self.assertIsNone(tracker._reader)
+
+
+class ReaderCacheTest(unittest.TestCase):
+ """The cache parks a reader under one key, hands it to the next acquirer,
+ bounds itself by idle time and entry count, and closes on teardown."""
+ def _reader(self):
+ class _FakeReader(UnboundedReader):
+ def __init__(self):
+ self.closed = False
+
+ @override
+ def close(self):
+ self.closed = True
+
+ return _FakeReader()
+
+ def test_park_then_acquire_returns_same_reader_and_started_flag(self):
+ cache = _ReaderCache()
+ reader = self._reader()
+ cache.park('k', reader, True)
+ self.assertEqual(cache.acquire('k'), (reader, True))
+ # Acquire removes the entry so two trackers cannot share one reader.
+ self.assertIsNone(cache.acquire('k'))
+ self.assertFalse(reader.closed)
+
+ def test_acquire_miss_returns_none(self):
+ self.assertIsNone(_ReaderCache().acquire('absent'))
+
+ def test_park_replacing_same_key_closes_displaced_reader(self):
+ # Two parks under one key without an intervening acquire (e.g. a bundle
+ # retry) must not leak the first reader.
+ cache = _ReaderCache()
+ old, new = self._reader(), self._reader()
+ cache.park('k', old, True)
+ cache.park('k', new, True)
+ self.assertTrue(old.closed)
+ self.assertFalse(new.closed)
+ self.assertEqual(cache.acquire('k'), (new, True))
+
+ def test_idle_reader_is_closed_on_next_touch(self):
+ clock = _ManualClock(1000.0)
+ cache = _ReaderCache(idle_seconds=30.0, now=clock)
+ reader = self._reader()
+ cache.park('k', reader, True)
+ clock.now = 1031.0 # past the idle window
+ cache.park('other', self._reader(), False) # triggers idle eviction
+ self.assertTrue(reader.closed)
+ self.assertIsNone(cache.acquire('k'))
+
+ def test_max_size_evicts_and_closes_oldest(self):
+ cache = _ReaderCache(max_size=2)
+ readers = [self._reader() for _ in range(3)]
+ cache.park('a', readers[0], True)
+ cache.park('b', readers[1], True)
+ cache.park('c', readers[2], True) # exceeds the cap, evicts 'a'
+ self.assertTrue(readers[0].closed)
+ self.assertIsNone(cache.acquire('a'))
+ self.assertIsNotNone(cache.acquire('b'))
+ self.assertIsNotNone(cache.acquire('c'))
+
+ def test_close_all_closes_every_parked_reader(self):
+ cache = _ReaderCache()
+ readers = [self._reader() for _ in range(3)]
+ for i, reader in enumerate(readers):
+ cache.park(str(i), reader, True)
+ cache.close_all()
+ self.assertTrue(all(reader.closed for reader in readers))
+ self.assertIsNone(cache.acquire('0'))
+
+ def test_close_all_swallows_reader_close_errors(self):
+ class _BoomReader(UnboundedReader):
+ @override
+ def close(self):
+ raise RuntimeError('close blew up')
+
+ cache = _ReaderCache()
+ cache.park('k', _BoomReader(), True)
+ with self.assertLogs(_unbounded_source_module._LOGGER, 'WARNING') as logs:
+ cache.close_all()
+ self.assertTrue(
+ any('Error closing UnboundedReader' in line for line in logs.output))
+
+
+class TrackerContractRegressionTest(unittest.TestCase):
+ """Tracker contract: source-watermark on the data path, finalize/resume
+ channel separation, and reader close on a reader-method failure."""
+ def test_data_path_holder_carries_source_watermark(self):
+ class _LaggingReader(UnboundedReader):
+ @override
+ def start(self):
+ return True
+
+ @override
+ def advance(self):
+ return False
+
+ @override
+ def get_current(self):
+ return 'rec'
+
+ @override
+ def get_current_timestamp(self):
+ return Timestamp(1000) # record event time
+
+ @override
+ def get_watermark(self):
+ return Timestamp(990) # source watermark lags 10us behind
+
+ @override
+ def get_checkpoint_mark(self):
+ return _CountingCheckpointMark(0)
+
+ class _LaggingSource(UnboundedSource):
+ @override
+ def split(self, desired_num_splits, options=None):
+ return [self]
+
+ @override
+ def create_reader(self, options, checkpoint_mark):
+ return _LaggingReader()
+
+ @override
+ def get_checkpoint_mark_coder(self):
+ return coders.PickleCoder()
+
+ tracker = _new_tracker(_LaggingSource())
+ claimed, record = _claim(tracker)
+ self.assertTrue(claimed)
+ self.assertIsNot(record, _NO_DATA)
+ value, record_timestamp, source_watermark = record
+ self.assertEqual(value, 'rec')
+ self.assertEqual(record_timestamp, Timestamp(1000))
+ # Critical: watermark slot is the SOURCE watermark, NOT record timestamp.
+ self.assertEqual(source_watermark, Timestamp(990))
+ self.assertNotEqual(source_watermark, record_timestamp)
+
+ def test_split_separates_finalize_and_resume_channels(self):
+ source = UnboundedCountingSource(5)
+ tracker = _new_tracker(source)
+ _claim(tracker) # claim 0 so reader has progress
+ primary, residual = tracker.try_split(0)
+ # Primary carries ONLY the finalize hook -- no resume state.
+ self.assertIsNone(primary.checkpoint_mark)
+ self.assertIsNotNone(primary.finalization_checkpoint_mark)
+ self.assertTrue(primary.is_done)
+ # Residual carries ONLY the resume state -- no finalize hook (a future
+ # bundle that splits THIS residual will produce ITS own finalize mark).
+ self.assertIsNotNone(residual.checkpoint_mark)
+ self.assertIsNone(residual.finalization_checkpoint_mark)
+ self.assertFalse(residual.is_done)
+ # The two marks reference the same underlying checkpoint object.
+ self.assertEqual(
+ primary.finalization_checkpoint_mark.last_index,
+ residual.checkpoint_mark.last_index)
+
+ def test_eof_populates_finalize_and_clears_resume(self):
+ # EOF transition: restriction.checkpoint_mark goes to None (no more
+ # records to resume from), finalization_checkpoint_mark carries the
+ # final commit hook.
+ source = UnboundedCountingSource(0) # immediately exhausted
+ tracker = _new_tracker(source)
+ holder = [None]
+ self.assertFalse(tracker.try_claim(holder))
+ r = tracker.current_restriction()
+ self.assertTrue(r.is_done)
+ self.assertEqual(r.watermark, MAX_TIMESTAMP)
+ self.assertIsNone(r.checkpoint_mark)
+ self.assertIsNotNone(r.finalization_checkpoint_mark)
+
+ def test_tracker_closes_reader_when_advance_raises(self):
+ # try_claim closes the reader before re-raising a reader-method failure, so
+ # the DoFn's finally need not traverse the SDF chain for these.
+ class _BoomReader(UnboundedReader):
+ def __init__(self):
+ self.closed = False
+
+ @override
+ def start(self):
+ return True
+
+ @override
+ def advance(self):
+ raise RuntimeError('advance boom')
+
+ @override
+ def get_current(self):
+ return 'first'
+
+ @override
+ def get_current_timestamp(self):
+ return _EVENT_TIME_BASE
+
+ @override
+ def get_watermark(self):
+ return _EVENT_TIME_BASE
+
+ @override
+ def get_checkpoint_mark(self):
+ return _CountingCheckpointMark(0)
+
+ @override
+ def close(self):
+ self.closed = True
+
+ class _BoomSource(UnboundedSource):
+ @override
+ def split(self, desired_num_splits, options=None):
+ return [self]
+
+ @override
+ def create_reader(self, options, checkpoint_mark):
+ return _BoomReader()
+
+ @override
+ def get_checkpoint_mark_coder(self):
+ return coders.PickleCoder()
+
+ src = _BoomSource()
+ tracker = _new_tracker(src)
+ # First claim succeeds (start returns True).
+ self.assertTrue(tracker.try_claim([None]))
+ reader_after_first = tracker._reader
+ self.assertIsNotNone(reader_after_first)
+ # The second claim's advance() raises; the tracker must close the reader
+ # before propagating.
+ with self.assertRaises(RuntimeError):
+ tracker.try_claim([None])
+ self.assertTrue(reader_after_first.closed)
+ self.assertIsNone(tracker._reader)
+
+ def test_tracker_closes_reader_when_get_watermark_raises(self):
+ # Reader method failures other than advance() also trigger close.
+ class _WatermarkBoomReader(UnboundedReader):
+ def __init__(self):
+ self.closed = False
+
+ @override
+ def start(self):
+ return False # no data -> drops into get_watermark path
+
+ @override
+ def advance(self):
+ return False
+
+ @override
+ def get_current(self):
+ raise AssertionError
+
+ @override
+ def get_current_timestamp(self):
+ raise AssertionError
+
+ @override
+ def get_watermark(self):
+ raise RuntimeError('watermark boom')
+
+ @override
+ def get_checkpoint_mark(self):
+ return _CountingCheckpointMark(0)
+
+ @override
+ def close(self):
+ self.closed = True
+
+ class _WatermarkBoomSource(UnboundedSource):
+ @override
+ def split(self, desired_num_splits, options=None):
+ return [self]
+
+ @override
+ def create_reader(self, options, checkpoint_mark):
+ return _WatermarkBoomReader()
+
+ @override
+ def get_checkpoint_mark_coder(self):
+ return coders.PickleCoder()
+
+ src = _WatermarkBoomSource()
+ tracker = _new_tracker(src)
+ with self.assertRaises(RuntimeError):
+ tracker.try_claim([None])
+ self.assertIsNone(tracker._reader)
+
+
+class UnboundedSourceContractTest(unittest.TestCase):
+ def test_get_checkpoint_mark_coder_default_names_subclass(self):
+ class MySource(UnboundedSource):
+ pass
+
+ with self.assertRaises(NotImplementedError) as cm:
+ MySource().get_checkpoint_mark_coder()
+ self.assertIn('MySource', str(cm.exception))
+
+
+class ReadFromUnboundedSourceValidationTest(unittest.TestCase):
+ def test_non_source_argument_raises(self):
+ with self.assertRaises(TypeError):
+ ReadFromUnboundedSource('not-a-source') # type: ignore[arg-type]
+
+ def test_invalid_caps_raise(self):
+ source = UnboundedCountingSource(1)
+ with self.assertRaises(ValueError):
+ ReadFromUnboundedSource(source, max_records_per_bundle=0)
+ with self.assertRaises(ValueError):
+ ReadFromUnboundedSource(source, max_read_time_seconds=0)
+ with self.assertRaises(ValueError):
+ ReadFromUnboundedSource(source, poll_interval=-1)
+
+
+if __name__ == '__main__':
+ logging.getLogger().setLevel(logging.INFO)
+ unittest.main()
