Fix PickleCoder.as_deterministic_coder() raising TypeError

CHANGES.md

@@ -104,6 +104,7 @@
 * (Java) MongoDbIO read splitting now preserves non-ObjectId `_id` types (e.g. string ids) instead of failing to parse the generated range filters ([#39900](https://github.com/apache/beam/issues/39900)).
 * (Go) Fixed GCS glob matching silently dropping objects when the glob pattern contains multi-byte characters ([#39969](https://github.com/apache/beam/issues/39969)).
 * (Python) Fixed `TensorRTEngineHandlerNumPy` failing with `CUDA_ERROR_INVALID_VALUE` on models with a single-element input or output tensor ([#36306](https://github.com/apache/beam/issues/36306)).
+* (Python) Fixed `PickleCoder`/`_MemoizingPickleCoder.as_deterministic_coder()` raising `TypeError` instead of returning a working deterministic coder ([#28558](https://github.com/apache/beam/issues/28558)).
 
 ## Security Fixes
 

sdks/python/apache_beam/coders/coders.py

@@ -38,6 +38,7 @@
 
 import base64
 import decimal
+import logging
 import pickle
 from functools import lru_cache
 from typing import TYPE_CHECKING
@@ -114,6 +115,8 @@
     'PaneInfoCoder'
 ]
 
+_LOGGER = logging.getLogger(__name__)
+
 T = TypeVar('T')
 CoderT = TypeVar('CoderT', bound='Coder')
 ProtoCoderT = TypeVar('ProtoCoderT', bound='ProtoCoder')
@@ -899,7 +902,20 @@ def _nonhashable_dumps(x):
     return coder_impl.CallbackCoderImpl(_nonhashable_dumps, pickler.loads)
 
   def as_deterministic_coder(self, step_label, error_message=None):
-    return FastPrimitivesCoder(self, requires_deterministic=step_label)
+    _LOGGER.warning(
+        "PickleCoder was registered for a key type in '%s', but the runner "
+        "requires a deterministic key encoding and pickle is not "
+        "deterministic. Keys in this step will be encoded with the "
+        "deterministic fallback coder instead of pickle. That coder "
+        "supports primitives, containers, protobuf messages, frozen "
+        "dataclasses, NamedTuples, enums, and classes defining both "
+        "__getstate__ and __setstate__. Any other key type will fail at "
+        "encode time. If your key type is not one of these, register a "
+        "deterministic custom Coder for it or add a type hint so the "
+        "default coder is used.",
+        step_label)
+    return _update_compatible_deterministic_fast_primitives_coder(
+        self, step_label)
 
   def to_type_hint(self):
     return Any
@@ -914,7 +930,20 @@ def _create_impl(self):
         lambda x: dumps(x, protocol), pickle.loads)
 
   def as_deterministic_coder(self, step_label, error_message=None):
-    return FastPrimitivesCoder(self, requires_deterministic=step_label)
+    _LOGGER.warning(
+        "PickleCoder was registered for a key type in '%s', but the runner "
+        "requires a deterministic key encoding and pickle is not "
+        "deterministic. Keys in this step will be encoded with the "
+        "deterministic fallback coder instead of pickle. That coder "
+        "supports primitives, containers, protobuf messages, frozen "
+        "dataclasses, NamedTuples, enums, and classes defining both "
+        "__getstate__ and __setstate__. Any other key type will fail at "
+        "encode time. If your key type is not one of these, register a "
+        "deterministic custom Coder for it or add a type hint so the "
+        "default coder is used.",
+        step_label)
+    return _update_compatible_deterministic_fast_primitives_coder(
+        self, step_label)
 
   def to_type_hint(self):
     return Any

sdks/python/apache_beam/coders/coders_test.py

@@ -50,6 +50,21 @@ def test_equality(self):
     self.assertNotEqual(coders.Base64PickleCoder(), coders.PickleCoder())
     self.assertNotEqual(coders.Base64PickleCoder(), object())
 
+  def test_as_deterministic_coder(self):
+    # PickleCoder.as_deterministic_coder used to construct FastPrimitivesCoder
+    # with a requires_deterministic kwarg that constructor never accepted,
+    # raising TypeError on every call (see coders.FastPrimitivesCoder).
+    v = ('a' * 10, 'b' * 90)
+    deterministic = coders.PickleCoder().as_deterministic_coder('label')
+    self.assertTrue(deterministic.is_deterministic())
+    self.assertEqual(v, deterministic.decode(deterministic.encode(v)))
+
+    memoizing_deterministic = coders._MemoizingPickleCoder(
+    ).as_deterministic_coder('label')
+    self.assertTrue(memoizing_deterministic.is_deterministic())
+    self.assertEqual(
+        v, memoizing_deterministic.decode(memoizing_deterministic.encode(v)))
+
 
 class CodersTest(unittest.TestCase):
   def test_str_utf8_coder(self):