From 52cc6f239d9050a99006dc4deeee45f4d3e1a445 Mon Sep 17 00:00:00 2001 From: xgbtxy <115146404+xgbtxy@users.noreply.github.com> Date: Mon, 26 Oct 2026 12:21:22 +0800 Subject: [PATCH 1/2] [Python] Preserve argument binding for defaulted DoFn elements Ignore the ordinary default on the implicitly supplied first input when all positional parameters have defaults. This keeps injected parameters and side inputs in their intended positions for process and process_batch. Cover implicit and explicit element parameters, scalar and deferred side inputs, and the existing FnApiRunner batch execution configurations. Fixes #19631 Assisted-by: OpenAI Codex --- CHANGES.md | 1 + sdks/python/apache_beam/runners/common.py | 7 +++- .../python/apache_beam/runners/common_test.py | 37 +++++++++++++++++++ .../fn_api_runner/fn_runner_test.py | 23 ++++++++++++ 4 files changed, 66 insertions(+), 2 deletions(-) diff --git a/CHANGES.md b/CHANGES.md index 8fa59e7d5ea0..c43f9ef2ec44 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -108,6 +108,7 @@ ## Bugfixes +* (Python) Fixed timestamp and side-input argument binding when a `DoFn.process` or `DoFn.process_batch` method gives its element or batch parameter a default value ([#19631](https://github.com/apache/beam/issues/19631)). * (Python) Fixed BigQuery `FILE_LOADS` silently omitting files when dynamic destinations use different spellings of the same table ([#40001](https://github.com/apache/beam/pull/40001)). * (Python) DaskRunner now preserves outputs emitted by `DoFn.finish_bundle()` instead of silently dropping them ([#40519](https://github.com/apache/beam/issues/40519)). * (Python) BigQuery streaming inserts now reset the retry budget and backoff for each batch instead of sharing them across a bundle ([#39995](https://github.com/apache/beam/pull/39995)). diff --git a/sdks/python/apache_beam/runners/common.py b/sdks/python/apache_beam/runners/common.py index b1ae4d4dd34e..b8f47b7f7247 100644 --- a/sdks/python/apache_beam/runners/common.py +++ b/sdks/python/apache_beam/runners/common.py @@ -717,8 +717,11 @@ def __init__(self, placeholder): self.placeholder = placeholder if all(core.DoFn.ElementParam != arg for arg in default_arg_values): - # TODO(https://github.com/apache/beam/issues/19631): Handle cases in which - # len(arg_names) == len(default_arg_values). + if (default_arg_values and len(arg_names) == len(default_arg_values) and + not isinstance(default_arg_values[0], core._DoFnParam) and + default_arg_values[0] is not core.DoFn.BundleFinalizerParam): + # The implicit element occupies the first slot even if it has a default. + default_arg_values = default_arg_values[1:] args_to_pick = len(arg_names) - len(default_arg_values) - 1 # Positional argument values for process(), with placeholders for special # values such as the element, timestamp, etc. diff --git a/sdks/python/apache_beam/runners/common_test.py b/sdks/python/apache_beam/runners/common_test.py index c1162f653cad..5af15d24114a 100644 --- a/sdks/python/apache_beam/runners/common_test.py +++ b/sdks/python/apache_beam/runners/common_test.py @@ -20,6 +20,7 @@ import unittest import hamcrest as hc +from parameterized import parameterized import apache_beam as beam from apache_beam.io.restriction_trackers import OffsetRange @@ -32,6 +33,8 @@ from apache_beam.runners.sdf_utils import SplitResultResidual from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.test_stream import TestStream +from apache_beam.testing.util import assert_that +from apache_beam.testing.util import equal_to from apache_beam.transforms import trigger from apache_beam.transforms import window from apache_beam.transforms.core import DoFn @@ -121,6 +124,40 @@ def process(self, element): return RecordDoFn() + @parameterized.expand([ + ('implicit_default', 1, 'default'), + ('implicit_positional', 1, 'positional'), + ('implicit_keyword', 1, 'keyword'), + ('implicit_side_input', 1, 'side_input'), + ('explicit_default', DoFn.ElementParam, 'default'), + ('explicit_positional', DoFn.ElementParam, 'positional'), + ('explicit_keyword', DoFn.ElementParam, 'keyword'), + ('explicit_side_input', DoFn.ElementParam, 'side_input'), + ]) + def test_dofn_process_default_element(self, name, element_default, side_mode): + class ProcessDefaults(DoFn): + def process( + self, + element=element_default, + timestamp=DoFn.TimestampParam, + side_input=3): + yield element, timestamp, side_input + + with TestPipeline() as p: + args, kwargs = [], {} + if side_mode == 'positional': + args = [2] + elif side_mode == 'keyword': + kwargs = {'side_input': 2} + elif side_mode == 'side_input': + args = [beam.pvalue.AsSingleton(p | 'SideInput' >> beam.Create([2]))] + result = ( + p + | 'Input' >> beam.Create([window.TimestampedValue(5, 123)]) + | beam.ParDo(ProcessDefaults(), *args, **kwargs)) + expected_side = 3 if side_mode == 'default' else 2 + assert_that(result, equal_to([(5, Timestamp(123), expected_side)])) + def test_dofn_process_keyparam(self): class DoFnProcessWithKeyparam(DoFn): def process(self, element, mykey=DoFn.KeyParam): 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 884dd3423109..255008ab13f8 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 @@ -269,6 +269,29 @@ def infer_output_type(self, input_type): assert_that(res, equal_to([0, 1, 0, 3, 0, 5, 0, 7, 0, 9])) + def test_batch_pardo_default_element(self): + class ProcessDefaults(beam.DoFn): + @no_type_check + def process_batch( + self, + batch: np.ndarray = None, + ts=beam.DoFn.TimestampParam, + side_input=3) -> Iterator[np.ndarray]: + yield batch * (ts.seconds() + side_input) + + def infer_output_type(self, input_type): + return input_type + + with self.create_pipeline() as p: + result = ( + p + | beam.Create(np.array([1, 2, 3], dtype=np.int64)).with_output_types( + np.int64) + | beam.Map(lambda value: window.TimestampedValue(value, 123)). + with_output_types(np.int64) + | beam.ParDo(ProcessDefaults(), 2)) + assert_that(result, equal_to([125, 250, 375])) + def test_batch_pardo_window_param(self): class PerWindowDoFn(beam.DoFn): @no_type_check From 0dbcf731ed20388b3343328e7d6e1c27e3748a06 Mon Sep 17 00:00:00 2001 From: xgbtxy <115146404+xgbtxy@users.noreply.github.com> Date: Wed, 28 Oct 2026 13:26:16 +0800 Subject: [PATCH 2/2] fixup! [Python] Preserve argument binding for defaulted DoFn elements Skip the default-argument loop when no defaults remain, avoiding the negative-zero slice noted in review. Assisted-by: OpenAI Codex --- sdks/python/apache_beam/runners/common.py | 73 ++++++++++++----------- 1 file changed, 37 insertions(+), 36 deletions(-) diff --git a/sdks/python/apache_beam/runners/common.py b/sdks/python/apache_beam/runners/common.py index b8f47b7f7247..b19b0925e7db 100644 --- a/sdks/python/apache_beam/runners/common.py +++ b/sdks/python/apache_beam/runners/common.py @@ -733,42 +733,43 @@ def __init__(self, placeholder): # Fill the OtherPlaceholders for context, key, window or timestamp remaining_args_iter = iter(input_args[args_to_pick:]) - for a, d in zip(arg_names[-len(default_arg_values):], default_arg_values): - if core.DoFn.ElementParam == d: - args_with_placeholders.append(ArgPlaceholder(d)) - elif core.DoFn.KeyParam == d: - args_with_placeholders.append(ArgPlaceholder(d)) - elif core.DoFn.WindowParam == d: - args_with_placeholders.append(ArgPlaceholder(d)) - elif core.DoFn.WindowedValueParam == d: - args_with_placeholders.append(ArgPlaceholder(d)) - elif core.DoFn.TimestampParam == d: - args_with_placeholders.append(ArgPlaceholder(d)) - elif core.DoFn.PaneInfoParam == d: - args_with_placeholders.append(ArgPlaceholder(d)) - elif core.DoFn.SideInputParam == d: - # If no more args are present then the value must be passed via kwarg - try: - args_with_placeholders.append(next(remaining_args_iter)) - except StopIteration: - if a not in input_kwargs: - raise ValueError("Value for sideinput %s not provided" % a) - elif isinstance(d, core.DoFn.StateParam): - args_with_placeholders.append(ArgPlaceholder(d)) - elif isinstance(d, core.DoFn.TimerParam): - args_with_placeholders.append(ArgPlaceholder(d)) - elif isinstance(d, type) and core.DoFn.BundleFinalizerParam == d: - args_with_placeholders.append(ArgPlaceholder(d)) - elif isinstance(d, core.DoFn.BundleContextParam): - args_with_placeholders.append(ArgPlaceholder(d)) - elif isinstance(d, core.DoFn.SetupContextParam): - args_with_placeholders.append(ArgPlaceholder(d)) - else: - # If no more args are present then the value must be passed via kwarg - try: - args_with_placeholders.append(next(remaining_args_iter)) - except StopIteration: - pass + if default_arg_values: + for a, d in zip(arg_names[-len(default_arg_values):], default_arg_values): + if core.DoFn.ElementParam == d: + args_with_placeholders.append(ArgPlaceholder(d)) + elif core.DoFn.KeyParam == d: + args_with_placeholders.append(ArgPlaceholder(d)) + elif core.DoFn.WindowParam == d: + args_with_placeholders.append(ArgPlaceholder(d)) + elif core.DoFn.WindowedValueParam == d: + args_with_placeholders.append(ArgPlaceholder(d)) + elif core.DoFn.TimestampParam == d: + args_with_placeholders.append(ArgPlaceholder(d)) + elif core.DoFn.PaneInfoParam == d: + args_with_placeholders.append(ArgPlaceholder(d)) + elif core.DoFn.SideInputParam == d: + # If no more args are present then the value must be passed via kwarg + try: + args_with_placeholders.append(next(remaining_args_iter)) + except StopIteration: + if a not in input_kwargs: + raise ValueError("Value for sideinput %s not provided" % a) + elif isinstance(d, core.DoFn.StateParam): + args_with_placeholders.append(ArgPlaceholder(d)) + elif isinstance(d, core.DoFn.TimerParam): + args_with_placeholders.append(ArgPlaceholder(d)) + elif isinstance(d, type) and core.DoFn.BundleFinalizerParam == d: + args_with_placeholders.append(ArgPlaceholder(d)) + elif isinstance(d, core.DoFn.BundleContextParam): + args_with_placeholders.append(ArgPlaceholder(d)) + elif isinstance(d, core.DoFn.SetupContextParam): + args_with_placeholders.append(ArgPlaceholder(d)) + else: + # If no more args are present then the value must be passed via kwarg + try: + args_with_placeholders.append(next(remaining_args_iter)) + except StopIteration: + pass args_with_placeholders.extend(list(remaining_args_iter)) # Stash the list of placeholder positions for performance