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 89fb1dcfedf [Python] Add UnboundedSource ValidatesRunner test (#38892)
89fb1dcfedf is described below
commit 89fb1dcfedfdb9728af3b803d2d812f3b5cd543c
Author: Elia Liu <[email protected]>
AuthorDate: Thu Jul 9 01:18:40 2026 +1000
[Python] Add UnboundedSource ValidatesRunner test (#38892)
---
...PostCommit_Python_ValidatesRunner_Dataflow.json | 3 +-
...am_PostCommit_Python_ValidatesRunner_Flink.json | 3 +-
...am_PostCommit_Python_ValidatesRunner_Spark.json | 3 +-
sdks/python/apache_beam/io/iobase_test.py | 8 +-
sdks/python/apache_beam/io/unbounded_source.py | 140 +++++++++++++++++++++
.../python/apache_beam/io/unbounded_source_test.py | 136 +-------------------
.../portability/fn_api_runner/fn_runner_test.py | 29 +++++
.../runners/portability/spark_runner_test.py | 4 +
8 files changed, 187 insertions(+), 139 deletions(-)
diff --git
a/.github/trigger_files/beam_PostCommit_Python_ValidatesRunner_Dataflow.json
b/.github/trigger_files/beam_PostCommit_Python_ValidatesRunner_Dataflow.json
index f1ba03a243e..462412f7bd3 100644
--- a/.github/trigger_files/beam_PostCommit_Python_ValidatesRunner_Dataflow.json
+++ b/.github/trigger_files/beam_PostCommit_Python_ValidatesRunner_Dataflow.json
@@ -1,4 +1,5 @@
{
"comment": "Modify this file in a trivial way to cause this test suite to
run",
- "modification": 5
+ "modification": 5,
+ "https://github.com/apache/beam/pull/38892": "UnboundedSource portable VR
test"
}
diff --git
a/.github/trigger_files/beam_PostCommit_Python_ValidatesRunner_Flink.json
b/.github/trigger_files/beam_PostCommit_Python_ValidatesRunner_Flink.json
index b08e2bb8115..962795fcb0a 100644
--- a/.github/trigger_files/beam_PostCommit_Python_ValidatesRunner_Flink.json
+++ b/.github/trigger_files/beam_PostCommit_Python_ValidatesRunner_Flink.json
@@ -1,5 +1,6 @@
{
"https://github.com/apache/beam/pull/32648": "testing addition of Flink 1.19
support",
"https://github.com/apache/beam/pull/34830": "testing",
- "trigger-2026-04-04": "portable_runner expand_sdf opt-in"
+ "trigger-2026-04-04": "portable_runner expand_sdf opt-in",
+ "https://github.com/apache/beam/pull/38892": "UnboundedSource portable VR
test"
}
diff --git
a/.github/trigger_files/beam_PostCommit_Python_ValidatesRunner_Spark.json
b/.github/trigger_files/beam_PostCommit_Python_ValidatesRunner_Spark.json
index bd7b284625b..f4288649d8a 100644
--- a/.github/trigger_files/beam_PostCommit_Python_ValidatesRunner_Spark.json
+++ b/.github/trigger_files/beam_PostCommit_Python_ValidatesRunner_Spark.json
@@ -1,5 +1,6 @@
{
"https://github.com/apache/beam/pull/34830": "testing",
"https://github.com/apache/beam/issues/35429": "testing",
- "trigger-2026-04-04": "portable_runner expand_sdf opt-in"
+ "trigger-2026-04-04": "portable_runner expand_sdf opt-in",
+ "https://github.com/apache/beam/pull/38892": "UnboundedSource portable VR
test"
}
diff --git a/sdks/python/apache_beam/io/iobase_test.py
b/sdks/python/apache_beam/io/iobase_test.py
index dbedf4681f4..60c26156315 100644
--- a/sdks/python/apache_beam/io/iobase_test.py
+++ b/sdks/python/apache_beam/io/iobase_test.py
@@ -224,22 +224,22 @@ class UseSdfBoundedSourcesTests(unittest.TestCase):
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).
+ ``unbounded_source`` as a finite fake source (no network).
"""
def test_read_end_to_end_unbounded(self):
- from apache_beam.io.unbounded_source_test import UnboundedCountingSource
+ from apache_beam.io.unbounded_source 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
+ from apache_beam.io.unbounded_source 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
+ from apache_beam.io.unbounded_source import UnboundedCountingSource
p = beam.Pipeline()
p | 'ReadIt' >> beam.io.Read(UnboundedCountingSource(3))
diff --git a/sdks/python/apache_beam/io/unbounded_source.py
b/sdks/python/apache_beam/io/unbounded_source.py
index 3246602252e..ac5f17ce45a 100644
--- a/sdks/python/apache_beam/io/unbounded_source.py
+++ b/sdks/python/apache_beam/io/unbounded_source.py
@@ -93,6 +93,8 @@ from typing import Callable
from typing import Iterable
from typing import Optional
+from typing_extensions import override
+
from apache_beam import coders
from apache_beam.coders.coders import BooleanCoder
from apache_beam.coders.coders import Coder
@@ -994,3 +996,141 @@ class ReadFromUnboundedSource(PTransform):
def _infer_output_coder(self, input_type=None, input_coder=None):
return self._source.default_output_coder()
+
+
+#
------------------------------------------------------------------------------
+# Internal in-memory counting source, shared by Beam's own tests and examples.
+#
------------------------------------------------------------------------------
+
+# Realistic event-time base away from the Unix epoch.
+_EVENT_TIME_BASE = Timestamp(1729987200) # 2024-10-27T00:00:00Z
+
+
+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):
+ """In-memory counting source for Beam's own tests and examples.
+
+ Internal: not part of the public ``UnboundedSource`` API. Emits integers
+ ``0..count-1`` with event time ``_EVENT_TIME_BASE + index``, self-terminates
+ at EOF, resumes from ``last_index + 1``, and optionally splits into
+ independent even/odd sub-sources.
+ """
+ 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()
diff --git a/sdks/python/apache_beam/io/unbounded_source_test.py
b/sdks/python/apache_beam/io/unbounded_source_test.py
index 3b63d0d2de9..6c6651cd0db 100644
--- a/sdks/python/apache_beam/io/unbounded_source_test.py
+++ b/sdks/python/apache_beam/io/unbounded_source_test.py
@@ -29,11 +29,15 @@ 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 _EVENT_TIME_BASE
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 UnboundedCountingSource
from apache_beam.io.unbounded_source import UnboundedReader
from apache_beam.io.unbounded_source import UnboundedSource
+from apache_beam.io.unbounded_source import _CountingCheckpointMark
+from apache_beam.io.unbounded_source import _CountingReader
from apache_beam.io.unbounded_source import _FinalizeCheckpointOnce
from apache_beam.io.unbounded_source import _ReaderCache
from apache_beam.io.unbounded_source import _ReadFromUnboundedSourceDoFn
@@ -55,138 +59,6 @@ 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
diff --git
a/sdks/python/apache_beam/runners/portability/fn_api_runner/fn_runner_test.py
b/sdks/python/apache_beam/runners/portability/fn_api_runner/fn_runner_test.py
index 72a5ed4412c..884dd342310 100644
---
a/sdks/python/apache_beam/runners/portability/fn_api_runner/fn_runner_test.py
+++
b/sdks/python/apache_beam/runners/portability/fn_api_runner/fn_runner_test.py
@@ -46,6 +46,7 @@ import apache_beam as beam
from apache_beam.coders import coders
from apache_beam.coders.coders import StrUtf8Coder
from apache_beam.io import restriction_trackers
+from apache_beam.io.unbounded_source import UnboundedCountingSource
from apache_beam.io.watermark_estimators import ManualWatermarkEstimator
from apache_beam.metrics import monitoring_infos
from apache_beam.metrics.execution import MetricKey
@@ -1170,6 +1171,25 @@ class FnApiRunnerTest(unittest.TestCase):
actual = (p | beam.Create([10]) | beam.ParDo(SimpleSDF()))
assert_that(actual, equal_to(range(5)))
+ def test_unbounded_source_read(self):
+ """Reads a self-terminating UnboundedSource end to end.
+
+ Exercises the runner-API round trip and the EOF ``MAX_TIMESTAMP`` watermark
+ propagation that lets the downstream window fire.
+ """
+ with self.create_pipeline() as p:
+ out = p | beam.io.Read(UnboundedCountingSource(5))
+ self.assertFalse(out.is_bounded)
+ assert_that(out, equal_to([0, 1, 2, 3, 4]), label='assert_elements')
+ windowed = (
+ out
+ | beam.WindowInto(window.FixedWindows(100))
+ | beam.Map(lambda v: ('all', v))
+ | beam.GroupByKey()
+ | beam.MapTuple(lambda _key, values: sorted(values)))
+ assert_that(
+ windowed, equal_to([[0, 1, 2, 3, 4]]), label='assert_windowed')
+
def test_group_by_key(self):
with self.create_pipeline() as p:
res = (
@@ -2038,6 +2058,9 @@ class FnApiRunnerTestWithMultiWorkers(FnApiRunnerTest):
def test_sdf_with_watermark_tracking(self):
raise unittest.SkipTest("This test is for a single worker only.")
+ def test_unbounded_source_read(self):
+ raise unittest.SkipTest("This test is for a single worker only.")
+
def test_sdf_with_dofn_as_watermark_estimator(self):
raise unittest.SkipTest("This test is for a single worker only.")
@@ -2071,6 +2094,9 @@ class
FnApiRunnerTestWithGrpcAndMultiWorkers(FnApiRunnerTest):
def test_sdf_with_watermark_tracking(self):
raise unittest.SkipTest("This test is for a single worker only.")
+ def test_unbounded_source_read(self):
+ raise unittest.SkipTest("This test is for a single worker only.")
+
def test_sdf_with_dofn_as_watermark_estimator(self):
raise unittest.SkipTest("This test is for a single worker only.")
@@ -2115,6 +2141,9 @@ class
FnApiRunnerTestWithBundleRepeatAndMultiWorkers(FnApiRunnerTest):
def test_sdf_with_watermark_tracking(self):
raise unittest.SkipTest("This test is for a single worker only.")
+ def test_unbounded_source_read(self):
+ raise unittest.SkipTest("This test is for a single worker only.")
+
def test_sdf_with_dofn_as_watermark_estimator(self):
raise unittest.SkipTest("This test is for a single worker only.")
diff --git a/sdks/python/apache_beam/runners/portability/spark_runner_test.py
b/sdks/python/apache_beam/runners/portability/spark_runner_test.py
index fe7160d4e18..54238bbf192 100644
--- a/sdks/python/apache_beam/runners/portability/spark_runner_test.py
+++ b/sdks/python/apache_beam/runners/portability/spark_runner_test.py
@@ -146,6 +146,10 @@ class
SparkRunnerTest(portable_runner_test.PortableRunnerTest):
# Skip until Spark runner supports SDF.
raise unittest.SkipTest("https://github.com/apache/beam/issues/19468";)
+ def test_unbounded_source_read(self):
+ # Skip until Spark runner supports SDF.
+ raise unittest.SkipTest("https://github.com/apache/beam/issues/19468";)
+
def test_sdf_with_watermark_tracking(self):
# Skip until Spark runner supports SDF.
raise unittest.SkipTest("https://github.com/apache/beam/issues/19468";)
