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..b19b0925e7db 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. @@ -730,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 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