Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions CHANGES.md
Original file line number Diff line number Diff line change
Expand Up @@ -74,6 +74,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)).

## Deprecations

Expand Down Expand Up @@ -126,6 +127,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

Expand Down
1 change: 1 addition & 0 deletions sdks/python/apache_beam/runners/common.pxd
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
35 changes: 35 additions & 0 deletions sdks/python/apache_beam/runners/common.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
# pytype: skip-file

# ruff: noqa: UP006
import inspect
import logging
import sys
import threading
Expand Down Expand Up @@ -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,
Expand All @@ -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:
Expand Down Expand Up @@ -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``.

Expand All @@ -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
Expand All @@ -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):
Expand All @@ -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 = []

Expand Down
58 changes: 58 additions & 0 deletions sdks/python/apache_beam/runners/common_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
Loading