[Python] Preserve argument binding for defaulted DoFn elements

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)).

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

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):

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