[Python] Polish UnboundedSource SDF wrapper API

Refine the experimental Python UnboundedSource wrapper for a user-facing API:

- document source, reader, and checkpoint contracts with a beam.io.Read example

- keep no-data polling as a private implementation detail

- simplify ReadFromUnboundedSource expansion with Create([source])

- make cleanup diagnostics visible without user-facing escalation wording

- tighten tests with a splittable UnboundedCountingSource, realistic timestamps, @OverRide annotations, and real Read expansion coverage

- defer the CHANGES entry until ValidatesRunner coverage is in place

Author: Eliaaazzz

sdks/python/apache_beam/io/iobase.py

@@ -918,11 +918,10 @@ def __init__(self, source: SourceBase) -> None:
     """Initializes a Read transform.
 
     Args:
-      source: Data source to read from. A ``BoundedSource`` is wrapped in the
-        bounded SDF reader; an ``UnboundedSource`` is dispatched through
-        :class:`apache_beam.io.unbounded_source.ReadFromUnboundedSource` with
-        the default poll interval (users wanting a custom poll cadence must
-        instantiate ``ReadFromUnboundedSource`` directly).
+      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
@@ -948,15 +947,10 @@ def expand(self, pbegin):
           | 'EmitSource' >>
           core.Map(lambda _: self.source).with_output_types(BoundedSource)
           | SDFBoundedSourceReader(display_data))
-    # Lazy import to break the iobase <-> unbounded_source cycle: the
-    # unbounded_source module imports iobase (UnboundedSource extends
-    # SourceBase). Pattern matches the _PubSubSource lazy import below.
+    # 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):
-      # Delegate to the dedicated SDF PTransform; identical to the user
-      # writing `p | ReadFromUnboundedSource(self.source)` directly. Custom
-      # poll_interval_seconds requires using ReadFromUnboundedSource directly.
       return pbegin | ReadFromUnboundedSource(self.source)
     elif isinstance(self.source, ptransform.PTransform):
       # The Read transform can also admit a full PTransform as an input
@@ -999,15 +993,9 @@ def to_runner_api_parameter(
               timestamp_attribute=self.source.timestamp_attribute,
               with_attributes=self.source.with_attributes,
               id_attribute=self.source.id_label))
-    # Lazy import to avoid the iobase <-> unbounded_source cycle.
+    # Local import to avoid a circular dependency.
     from apache_beam.io.unbounded_source import UnboundedSource
     if isinstance(self.source, (BoundedSource, UnboundedSource)):
-      # READ.urn covers both source flavours; the IsBounded enum distinguishes
-      # them. NB: today the bundle_processor.py IMPULSE_READ_TRANSFORM handler
-      # only consumes BOUNDED - the UNBOUNDED branch round-trips correctly
-      # through the protobuf graph but execution still flows through this
-      # composite's expanded sub-transforms (Impulse | Map | SDF-ParDo), not
-      # through the URN-handler. Runner-side UNBOUNDED dispatch is W2 work.
       return (
           common_urns.deprecated_primitives.READ.urn,
           beam_runner_api_pb2.ReadPayload(

sdks/python/apache_beam/io/iobase_test.py

@@ -220,46 +220,33 @@ def test_sdf_wrap_range_source(self):
 
 
 class UseSdfUnboundedSourcesTests(unittest.TestCase):
-  """Mirrors UseSdfBoundedSourcesTests for the new UnboundedSource branch in
-  iobase.Read.expand(). Uses CountingSource from unbounded_source_test as the
-  fake finite UnboundedSource (avoids dragging the network in).
+  """Covers the UnboundedSource branch in
+  ``iobase.Read.expand()``. Uses ``UnboundedCountingSource`` from
+  ``unbounded_source_test`` as a finite fake source (no network).
   """
-  def test_read_dispatches_to_read_from_unbounded_source(self):
-    from apache_beam.io.unbounded_source_test import CountingSource
-    with mock.patch(
-        'apache_beam.io.unbounded_source.ReadFromUnboundedSource.expand'
-    ) as mock_expand:
-      mock_expand.side_effect = (
-          lambda pbegin: pbegin | beam.Impulse() | beam.Map(lambda _: 'fake'))
-      with beam.Pipeline() as p:
-        out = p | beam.io.Read(CountingSource(3))
-        assert_that(out, equal_to(['fake']))
-      mock_expand.assert_called_once()
-
   def test_read_end_to_end_unbounded(self):
-    from apache_beam.io.unbounded_source_test import CountingSource
+    from apache_beam.io.unbounded_source_test import UnboundedCountingSource
     with beam.Pipeline() as p:
-      out = p | beam.io.Read(CountingSource(5))
+      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 CountingSource
+    from apache_beam.io.unbounded_source_test import UnboundedCountingSource
     with beam.Pipeline() as p:
-      out = p | beam.io.Read(CountingSource(3))
+      out = p | beam.io.Read(UnboundedCountingSource(3))
       self.assertFalse(out.is_bounded)
 
   def test_to_runner_api_emits_unbounded_read_payload(self):
     """``Read.to_runner_api_parameter`` must serialize an UnboundedSource as
-    ``READ.urn`` with ``IsBounded.UNBOUNDED``. The runner-side handler is W2
-    and ignores this enum today, but the wire format must round-trip
-    consistently for pipeline persistence / cross-runner submission.
+    ``READ.urn`` with ``IsBounded.UNBOUNDED`` so the wire format round-trips
+    consistently for pipeline persistence and cross-runner submission.
     """
-    from apache_beam.io.unbounded_source_test import CountingSource
+    from apache_beam.io.unbounded_source_test import UnboundedCountingSource
     from apache_beam.portability import common_urns
     from apache_beam.portability.api import beam_runner_api_pb2
     from apache_beam.runners.pipeline_context import PipelineContext
 
-    read = beam.io.Read(CountingSource(5))
+    read = beam.io.Read(UnboundedCountingSource(5))
     urn, payload = read.to_runner_api_parameter(PipelineContext())
 
     self.assertEqual(urn, common_urns.deprecated_primitives.READ.urn)
@@ -274,12 +261,11 @@ def test_read_unbounded_round_trips_through_runner_api(self):
     transform must be a ``Read`` wrapping an equivalent UnboundedSource.
     """
     from apache_beam.io.unbounded_source import UnboundedSource
-    from apache_beam.io.unbounded_source_test import CountingSource
-    from apache_beam.portability import common_urns
+    from apache_beam.io.unbounded_source_test import UnboundedCountingSource
     from apache_beam.portability.api import beam_runner_api_pb2
     from apache_beam.runners.pipeline_context import PipelineContext
 
-    original = beam.io.Read(CountingSource(7))
+    original = beam.io.Read(UnboundedCountingSource(7))
     context = PipelineContext()
     urn, payload = original.to_runner_api_parameter(context)
 
@@ -290,7 +276,7 @@ def test_read_unbounded_round_trips_through_runner_api(self):
 
     self.assertIsInstance(restored, iobase.Read)
     self.assertIsInstance(restored.source, UnboundedSource)
-    self.assertIsInstance(restored.source, CountingSource)
+    self.assertIsInstance(restored.source, UnboundedCountingSource)
     self.assertFalse(restored.source.is_bounded())
     # Verify the source's internal state survived pickle round-trip.
     self.assertEqual(restored.source._count, 7)