[Python] Fix stateful ParDo taking the timer key coder from a side input

CHANGES.md

@@ -87,6 +87,7 @@
 * (Go) Fixed pubsubio importing a `google.golang.org/genproto` package removed in recent releases, which broke builds of Go modules depending on a current `genproto` version ([#40018](https://github.com/apache/beam/issues/40018)).
 * (Java) BigQueryIO now treats a 404 when deleting a temporary table or dataset as success, so a replayed work item whose earlier attempt already deleted it no longer retries forever ([#24997](https://github.com/apache/beam/issues/24997)).
 * (Java) IcebergIO now writes rows containing `EnumerationType` (proto enum) fields as strings, instead of throwing `Unsupported Beam logical type Enum` ([#40299](https://github.com/apache/beam/issues/40299)).
+* (Python) Fixed stateful DoFns with side inputs sometimes taking the timer key coder from a side input instead of the main input, which could make the worker fail to decode timer keys with `Unknown type tag` ([#40374](https://github.com/apache/beam/issues/40374)).
 * Fixed X (Java/Python) ([#X](https://github.com/apache/beam/issues/X)).
 
 ## Security Fixes

sdks/python/apache_beam/transforms/core.py

@@ -1856,7 +1856,9 @@ def _do_fn_info(self):
   def _get_key_and_window_coder(self, named_inputs):
     if named_inputs is None or not self._signature.is_stateful_dofn():
       return None, None
-    main_input = list(set(named_inputs.keys()) - set(self.side_inputs))[0]
+    main_input = [
+        tag for tag in named_inputs if not tag.startswith(SIDE_INPUT_PREFIX)
+    ][0]
     input_pcoll = named_inputs[main_input]
     kv_type_hint = input_pcoll.element_type
     if kv_type_hint and kv_type_hint != typehints.Any:

sdks/python/apache_beam/transforms/userstate_test.py

@@ -59,6 +59,7 @@
 from apache_beam.transforms.userstate import is_stateful_dofn
 from apache_beam.transforms.userstate import on_timer
 from apache_beam.transforms.userstate import validate_stateful_dofn
+from apache_beam.utils import proto_utils
 
 
 class TestStatefulDoFn(DoFn):
@@ -195,6 +196,47 @@ def test_state_spec_proto_conversion(self):
         beam_runner_api_pb2.FunctionSpec(urn=common_urns.user_state.BAG.urn),
         state_proto.protocol)
 
+  def test_timer_key_coder_ignores_side_inputs(self):
+    class StatefulDoFnWithSideInputs(DoFn):
+      EXPIRY_TIMER = TimerSpec('expiry', TimeDomain.WATERMARK)
+
+      def process(
+          self,
+          element,
+          side_1,
+          side_2,
+          side_3,
+          timer=DoFn.TimerParam(EXPIRY_TIMER)):
+        yield element
+
+      @on_timer(EXPIRY_TIMER)
+      def expiry_callback(self):
+        yield 'expired'
+
+    p = beam.Pipeline()
+    main_input = p | 'Main' >> beam.Create([('key', 1)])
+    # The side inputs have a different key type than the main input, so the
+    # timer key coder shows which input it was taken from.
+    side_inputs = [
+        beam.pvalue.AsDict(p | 'Side%d' % i >> beam.Create([(i, 'side')]))
+        for i in range(3)
+    ]
+    _ = main_input | 'Stateful' >> beam.ParDo(
+        StatefulDoFnWithSideInputs(), *side_inputs)
+
+    proto = p.to_runner_api()
+    pardo_proto, = [
+        t for t in proto.components.transforms.values()
+        if t.unique_name == 'Stateful'
+    ]
+    pardo_payload = proto_utils.parse_Bytes(
+        pardo_proto.spec.payload, beam_runner_api_pb2.ParDoPayload)
+    timer_spec = pardo_payload.timer_family_specs[
+        StatefulDoFnWithSideInputs.EXPIRY_TIMER.name]
+    timer_coder = proto.components.coders[timer_spec.timer_family_coder_id]
+    key_coder = proto.components.coders[timer_coder.component_coder_ids[0]]
+    self.assertEqual(common_urns.coders.STRING_UTF8.urn, key_coder.spec.urn)
+
   def test_param_construction(self):
     with self.assertRaises(ValueError):
       DoFn.StateParam(TimerSpec('timer', TimeDomain.WATERMARK))