Add ability for coders to set version tags for update compat checks. (#36316)

* Add ability for coders to set version tags for update compat checks.

* Fix lint.

* Update docstring.

* Update sdks/python/apache_beam/coders/coders.py

* Update sdks/python/apache_beam/coders/coders.py

---------

Co-authored-by: tvalentyn <tvalentyn@users.noreply.github.com>

Author: claudevdm

sdks/python/apache_beam/coders/coders.py

@@ -395,6 +395,15 @@ def from_runner_api_parameter(unused_payload, components, unused_context):
       else:
         return cls()
 
+  def version_tag(self) -> str:
+    """For internal use. Appends a version tag to the coder key in the pipeline
+    proto. Some runners (e.g. DataflowRunner) use coder key/id to verify if a
+    pipeline is update compatible. If the implementation of a coder changed
+    in an update incompatible way a version tag can be added to fail compat
+    compatibility checks.
+    """
+    return ""
+
 
 @Coder.register_urn(
     python_urns.PICKLED_CODER, google.protobuf.wrappers_pb2.BytesValue)

sdks/python/apache_beam/pipeline.py

@@ -1712,5 +1712,7 @@ def _unique_ref(self, obj=None, obj_type=None, label=None):
     prefix = self._normalize(
         '%s_%s_%s' %
         (self.namespace, obj_type.__name__, label or type(obj).__name__))[0:100]
+    if isinstance(obj, typecoders.coders.Coder) and obj.version_tag():
+      prefix = "%s_%s" % (prefix, obj.version_tag())
     self._counters[obj_type] += 1
     return '%s_%d' % (prefix, self._counters[obj_type])

sdks/python/apache_beam/pipeline_test.py

@@ -28,6 +28,7 @@
 import pytest
 
 import apache_beam as beam
+from apache_beam import coders
 from apache_beam import typehints
 from apache_beam.coders import BytesCoder
 from apache_beam.io import Read
@@ -1075,6 +1076,38 @@ def test_requirements(self):
         common_urns.requirements.REQUIRES_BUNDLE_FINALIZATION.urn,
         proto.requirements)
 
+  def test_coder_version_tag_included_in_runner_api_key(self):
+    class MyClass:
+      def __init__(self, value: int):
+        self.value = value
+
+    class VersionedCoder(coders.Coder):
+      def encode(self, value):
+        return str(value.value).encode()
+
+      def decode(self, encoded):
+        return MyClass(int(encoded.decode()))
+
+      def version_tag(self):
+        return "v269"
+
+      def to_type_hint(self):
+        return MyClass
+
+    coders.registry.register_coder(MyClass, VersionedCoder)
+    p = beam.Pipeline()
+    _ = (p | beam.Impulse() | beam.Map(lambda _: MyClass(1)))
+    pipeline_proto = p.to_runner_api()
+    coder_keys = sorted(list(pipeline_proto.components.coders.keys()))
+
+    self.assertListEqual(
+        coder_keys,
+        [
+            'ref_Coder_BytesCoder_1',
+            'ref_Coder_GlobalWindowCoder_2',
+            'ref_Coder_VersionedCoder_v269_3'
+        ])
+
   def test_annotations(self):
     some_proto = BytesCoder().to_runner_api(None)