This is an automated email from the ASF dual-hosted git repository.
jrmccluskey pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new 24d48c984ae [Python] Fix stateful ParDo taking the timer key coder
from a side input (#40380)
24d48c984ae is described below
commit 24d48c984aec65b2be14617a9d453616c49f4667
Author: Divyansh Yadav <[email protected]>
AuthorDate: Fri Oct 2 18:16:32 2026 +0530
[Python] Fix stateful ParDo taking the timer key coder from a side input
(#40380)
---
CHANGES.md | 1 +
sdks/python/apache_beam/transforms/core.py | 4 ++-
.../apache_beam/transforms/userstate_test.py | 42 ++++++++++++++++++++++
3 files changed, 46 insertions(+), 1 deletion(-)
diff --git a/CHANGES.md b/CHANGES.md
index 48f4a022538..788ffd08154 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -88,6 +88,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
diff --git a/sdks/python/apache_beam/transforms/core.py
b/sdks/python/apache_beam/transforms/core.py
index cd01253cade..07c6955af7e 100644
--- a/sdks/python/apache_beam/transforms/core.py
+++ b/sdks/python/apache_beam/transforms/core.py
@@ -1856,7 +1856,9 @@ class ParDo(PTransformWithSideInputs):
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:
diff --git a/sdks/python/apache_beam/transforms/userstate_test.py
b/sdks/python/apache_beam/transforms/userstate_test.py
index 45dba5c9e9e..80612ec32f3 100644
--- a/sdks/python/apache_beam/transforms/userstate_test.py
+++ b/sdks/python/apache_beam/transforms/userstate_test.py
@@ -59,6 +59,7 @@ from apache_beam.transforms.userstate import get_dofn_specs
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 @@ class InterfaceTest(unittest.TestCase):
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))
