This is an automated email from the ASF dual-hosted git repository.
Abacn 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 8a1c19bc9d0 [Python] Convert typing and native generic hints in Watch
coder inference (#39547)
8a1c19bc9d0 is described below
commit 8a1c19bc9d0e93a020b2286d814419099d022dbf
Author: Elia Liu <[email protected]>
AuthorDate: Thu Jul 30 03:33:26 2026 +1000
[Python] Convert typing and native generic hints in Watch coder inference
(#39547)
registry.get_coder receives typing and native generic annotations such
as tuple[str, float] unconverted and falls back to pickling. Watch now
converts hints with convert_to_beam_type before the registry lookup, so
annotated poll functions and key functions infer real structured coders
and fully typed elements.
---
sdks/python/apache_beam/io/watch.py | 12 ++++++++++--
sdks/python/apache_beam/io/watch_test.py | 19 +++++++++++++++++++
2 files changed, 29 insertions(+), 2 deletions(-)
diff --git a/sdks/python/apache_beam/io/watch.py
b/sdks/python/apache_beam/io/watch.py
index b0ff2bd6a07..f2eadaf4ace 100644
--- a/sdks/python/apache_beam/io/watch.py
+++ b/sdks/python/apache_beam/io/watch.py
@@ -77,6 +77,7 @@ from apache_beam.runners import sdf_utils
from apache_beam.transforms import PTransform
from apache_beam.transforms import core
from apache_beam.transforms.window import TimestampedValue
+from apache_beam.typehints import native_type_compatibility
from apache_beam.utils.timestamp import MAX_TIMESTAMP
from apache_beam.utils.timestamp import Duration
from apache_beam.utils.timestamp import Timestamp
@@ -642,6 +643,13 @@ def _poll_output_type(poll_fn) -> Any:
return Any
+def _coder_for_hint(hint) -> Coder:
+ # typing and native generic hints such as tuple[str, float] must be
+ # converted to Beam typehints, or the registry falls back to pickling.
+ return coders.registry.get_coder(
+ native_type_compatibility.convert_to_beam_type(hint))
+
+
class Watch(PTransform):
"""Watches a growing set of outputs per input via a periodic poll function.
@@ -689,14 +697,14 @@ class Watch(PTransform):
if output_coder is None and isinstance(self._poll_fn, PollFn):
output_coder = self._poll_fn.default_output_coder()
if output_coder is None:
- output_coder =
coders.registry.get_coder(_poll_output_type(self._poll_fn))
+ output_coder = _coder_for_hint(_poll_output_type(self._poll_fn))
if self._output_key_fn is None:
# The output is its own dedup key, so the key coder is the output coder.
key_fn = _identity
key_coder = self._output_key_coder or output_coder
else:
key_fn = self._output_key_fn
- key_coder = self._output_key_coder or coders.registry.get_coder(
+ key_coder = self._output_key_coder or _coder_for_hint(
_return_type(self._output_key_fn))
# Dedup hashes the encoded key, so equal keys must encode equally; use the
# coder's deterministic form and reject coders that have none.
diff --git a/sdks/python/apache_beam/io/watch_test.py
b/sdks/python/apache_beam/io/watch_test.py
index 8c1f6571da6..472177ceaee 100644
--- a/sdks/python/apache_beam/io/watch_test.py
+++ b/sdks/python/apache_beam/io/watch_test.py
@@ -18,6 +18,7 @@
"""Tests for the Watch transform."""
import collections
+import typing
import unittest
import apache_beam as beam
@@ -444,6 +445,24 @@ class WatchEndToEndTest(unittest.TestCase):
| Watch(_complete_poll, poll_interval=Duration(1)))
self.assertEqual(typehints.Tuple[str, str], output.element_type)
+ def test_infers_coder_from_generic_annotations(self):
+ # tuple[str, float] and typing.Tuple[str, float] resolve to a tuple coder,
+ # not the pickling fallback.
+ def native_poll(element) -> PollResult[tuple[str, float]]:
+ return PollResult.complete([(element, 1.0)])
+
+ def typing_poll(
+ element) -> PollResult[typing.Tuple[str, float]]: # noqa: UP006
+ return PollResult.complete([(element, 1.0)])
+
+ for poll in (native_poll, typing_poll):
+ with self._in_memory_pipeline() as p:
+ output = (
+ p | beam.Create(['k:']) | Watch(poll, poll_interval=Duration(1)))
+ self.assertEqual(
+ typehints.Tuple[str, typehints.Tuple[str, float]],
+ output.element_type)
+
def test_uses_poll_fn_default_output_coder(self):
with self._in_memory_pipeline() as p:
output = (
