diff --git a/sdks/python/apache_beam/io/watch.py b/sdks/python/apache_beam/io/watch.py index b0ff2bd6a075..f2eadaf4ace8 100644 --- a/sdks/python/apache_beam/io/watch.py +++ b/sdks/python/apache_beam/io/watch.py @@ -77,6 +77,7 @@ def poll(prefix) -> PollResult[str]: 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 @@ def expand(self, pcoll): 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 8c1f6571da66..472177ceaee6 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 @@ def test_infers_output_coder_from_return_annotation(self): | 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 = (