Fix silent per-char iteration when DoFn returns str/bytes/dict (#38429)

* Fix silent per-char iteration when DoFn returns str/bytes/dict

* Use %r in error message for correct repr of dict/bytes

* Address review: 'is not allowed' wording, move entry to Breaking Changes

* Address review: check only first output, skip generator DoFns

---------

Co-authored-by: Jack McCluskey <34928439+jrmccluskey@users.noreply.github.com>

Author: chrisqiqiu

CHANGES.md

@@ -75,6 +75,7 @@
 ## Breaking Changes
 
 * (Python) Removed `google-perftools` from the SDK container images. Users who wish to use `--profiler_agent=tcmalloc` should install google-perftools APT package in their custom container images separately ([#39323](https://github.com/apache/beam/issues/39323)).
+* `DoFn.process` returning a `str`, `bytes`, or `dict` (instead of an iterable wrapping one) now raises a `TypeError` rather than silently iterating per-character/byte/key (Python) ([#18712](https://github.com/apache/beam/issues/18712)).
 * (Java) Added `DRAINING` and `DRAINED` states to `PipelineResult`, including runner state mappings and Dataflow update handling ([#39020](https://github.com/apache/beam/issues/39020)).
 
 ## Deprecations
@@ -129,6 +130,7 @@
 * (Python) Typehints of dataclass fields are honored during type inferences. To restore the behavior of fallback-to-any,
   use pipeline option `--exclude_infer_dataclass_field_type` ([#38797](https://github.com/apache/beam/issues/38797)).
   However fixing forward is recommended.
+* X behavior was changed ([#X](https://github.com/apache/beam/issues/X)).
 
 ## Bugfixes
 

sdks/python/apache_beam/runners/common.pxd

@@ -150,6 +150,7 @@ cdef class _OutputHandler(OutputHandler):
   cdef object output_batch_converter
   cdef bint _process_batch_yields_elements
   cdef bint _process_yields_batches
+  cdef bint _check_user_dofn_output
 
   @cython.locals(windowed_value=WindowedValue,
                  windowed_batch=WindowedBatch,

sdks/python/apache_beam/runners/common.py

@@ -23,6 +23,7 @@
 # pytype: skip-file
 
 # ruff: noqa: UP006
+import inspect
 import logging
 import sys
 import threading
@@ -1461,6 +1462,17 @@ def __init__(
     else:
       per_element_output_counter = None
 
+    # Only validate the output of user class-based DoFns, and only when
+    # process() uses `return` (not `yield`). A generator process() returns a
+    # generator object, which can never be a str/bytes/dict, so the #18712 bug
+    # is impossible there and the check would be pure overhead. Callable-wrapped
+    # DoFns (Map/FlatMap) are also excluded, since flattening a returned
+    # str/bytes/dict is a legitimate use case for them.
+    check_user_dofn_output = (
+        not isinstance(fn, core.CallableWrapperDoFn) and
+        not inspect.isgeneratorfunction(
+            do_fn_signature.process_method.method_value))
+
     output_handler = _OutputHandler(
         windowing.windowfn,
         main_receivers,
@@ -1475,6 +1487,7 @@ def __init__(
             do_fn_signature.process_batch_method.method_value,
             '_beam_yields_elements',
             False),
+        check_user_dofn_output=check_user_dofn_output,
     )
 
     if do_fn_signature.is_stateful_dofn() and not user_state_context:
@@ -1633,6 +1646,7 @@ def __init__(
       output_batch_converter,  # type: Optional[BatchConverter]
       process_yields_batches,  # type: bool
       process_batch_yields_elements,  # type: bool
+      check_user_dofn_output=False,  # type: bool
   ):
     """Initializes ``_OutputHandler``.
 
@@ -1642,6 +1656,12 @@ def __init__(
       tagged_receivers: main receiver object.
       per_element_output_counter: per_element_output_counter of one work_item.
                                   could be none if experimental flag turn off
+      check_user_dofn_output: if True, validate that a user-class DoFn does not
+                              return a str/bytes/dict (a common bug — see
+                              https://github.com/apache/beam/issues/18712).
+                              Skipped for callable-wrapped DoFns (Map/FlatMap)
+                              where iterating a returned str/bytes/dict is a
+                              legitimate flatten use case.
     """
     self.window_fn = window_fn
     self.main_receivers = main_receivers
@@ -1654,6 +1674,7 @@ def __init__(
     self.output_batch_converter = output_batch_converter
     self._process_yields_batches = process_yields_batches
     self._process_batch_yields_elements = process_batch_yields_elements
+    self._check_user_dofn_output = check_user_dofn_output
 
   def handle_process_outputs(
       self, windowed_input_element, results, watermark_estimator=None):
@@ -1664,6 +1685,20 @@ def handle_process_outputs(
     A value wrapped in a TaggedOutput object will be unwrapped and
     then dispatched to the appropriate indexed output.
     """
+    if self._check_user_dofn_output:
+      # This bug is deterministic per DoFn: if process() returns a
+      # str/bytes/dict once, it does so for every element. So we only need to
+      # validate the first output and can then disable the check to avoid
+      # per-element overhead (see
+      # https://github.com/apache/beam/issues/18712).
+      self._check_user_dofn_output = False
+      if isinstance(results, (str, bytes, dict)):
+        object_type = type(results).__name__
+        raise TypeError(
+            'Returning a %s from a ParDo or FlatMap is not allowed. '
+            'Please use list(%r) if you really want this behavior.' %
+            (object_type, results))
+
     if results is None:
       results = []
 

sdks/python/apache_beam/runners/common_test.py

@@ -154,6 +154,64 @@ def process(self, element, mykey=DoFn.KeyParam):
       test_stream = (TestStream().advance_watermark_to(10).add_elements([1, 2]))
       (p | test_stream | beam.ParDo(DoFnProcessWithKeyparam()))
 
+  def test_dofn_returning_str_raises_clear_error(self):
+    """Regression test for https://github.com/apache/beam/issues/18712.
+
+    A DoFn returning a str instead of an iterable wrapping one used to
+    silently iterate per-character. It should now raise a clear TypeError.
+    """
+    class BadDoFn(DoFn):
+      def process(self, element):
+        return 'hello'
+
+    # Use base Exception (matching existing convention in
+    # typecheck_test.py::test_do_fn_returning_non_iterable_throws_error)
+    # because the runner's _reraise_augmented wraps the TypeError before
+    # it surfaces to the test framework.
+    with self.assertRaisesRegex(
+        Exception, 'Returning a str from a ParDo or FlatMap is not allowed'):
+      with TestPipeline() as p:
+        _ = p | beam.Create([0]) | beam.ParDo(BadDoFn())
+
+  def test_dofn_returning_bytes_raises_clear_error(self):
+    """Regression test for https://github.com/apache/beam/issues/18712."""
+    class BadDoFn(DoFn):
+      def process(self, element):
+        return b'hello'
+
+    with self.assertRaisesRegex(
+        Exception, 'Returning a bytes from a ParDo or FlatMap is not allowed'):
+      with TestPipeline() as p:
+        _ = p | beam.Create([0]) | beam.ParDo(BadDoFn())
+
+  def test_dofn_returning_dict_raises_clear_error(self):
+    """Regression test for https://github.com/apache/beam/issues/18712."""
+    class BadDoFn(DoFn):
+      def process(self, element):
+        return {'k': 'v'}
+
+    with self.assertRaisesRegex(
+        Exception, 'Returning a dict from a ParDo or FlatMap is not allowed'):
+      with TestPipeline() as p:
+        _ = p | beam.Create([0]) | beam.ParDo(BadDoFn())
+
+  def test_dofn_yielding_str_is_not_flagged(self):
+    """A generator (yield) process() can't hit the #18712 bug.
+
+    Yielding a str emits the whole str as one element and must not be
+    treated as the "returned a str" error case.
+    """
+    class YieldDoFn(DoFn):
+      def process(self, element):
+        yield 'hello'
+
+    with TestPipeline() as p:
+      (
+          p | beam.Create([0]) | beam.ParDo(YieldDoFn())
+          | beam.ParDo(self.record_dofn()))
+
+    self.assertEqual(['hello'], DoFnProcessTest.all_records)
+
   def test_pardo_with_unbounded_per_element_dofn(self):
     class UnboundedDoFn(beam.DoFn):
       @beam.DoFn.unbounded_per_element()