Cache deterministic types (#36032)

* Cache deterministically encoded types.

* Cache pickled deterministic type encodings.

* Fix lint.

---------

Co-authored-by: Claude <cvandermerwe@google.com>

Author: claudevdm

sdks/python/apache_beam/coders/coder_impl.pxd

@@ -93,6 +93,8 @@ cdef class FastPrimitivesCoderImpl(StreamCoderImpl):
   cdef encode_type(self, t, OutputStream stream)
   cdef decode_type(self, InputStream stream)
 
+cdef dict _pickled_types
+
 cdef dict _unpickled_types
 
 

sdks/python/apache_beam/coders/coder_impl.py

@@ -550,15 +550,19 @@ def encode_type_2_67_0(self, t, stream):
     """
     Encode special type with <=2.67.0 compatibility.
     """
-    _verify_dill_compat()
-    stream.write(dill.dumps(t), True)
+    if t not in _pickled_types:
+      _verify_dill_compat()
+      _pickled_types[t] = dill.dumps(t)
+    stream.write(_pickled_types[t], True)
 
   def encode_type(self, t, stream):
     if self.force_use_dill:
       return self.encode_type_2_67_0(t, stream)
-    bs = cloudpickle_pickler.dumps(
-        t, config=cloudpickle_pickler.NO_DYNAMIC_CLASS_TRACKING_CONFIG)
-    stream.write(bs, True)
+
+    if t not in _pickled_types:
+      _pickled_types[t] = cloudpickle_pickler.dumps(
+          t, config=cloudpickle_pickler.NO_DYNAMIC_CLASS_TRACKING_CONFIG)
+    stream.write(_pickled_types[t], True)
 
   def decode_type(self, stream):
     if self.force_use_dill:
@@ -620,6 +624,7 @@ def decode_from_stream(self, stream, nested):
       raise ValueError('Unknown type tag %x' % t)
 
 
+_pickled_types = {}  # type: Dict[type, bytes]
 _unpickled_types = {}  # type: Dict[bytes, type]
 
 

sdks/python/apache_beam/tools/coders_microbenchmark.py

@@ -42,6 +42,7 @@
 from apache_beam.coders import proto2_coder_test_messages_pb2 as test_message
 from apache_beam.coders import coder_impl
 from apache_beam.coders import coders
+from apache_beam.coders import coders_test_common
 from apache_beam.coders import row_coder
 from apache_beam.coders import typecoders
 from apache_beam.tools import utils
@@ -249,6 +250,10 @@ def row_coder_benchmark_factory(generate_fn):
   return coder_benchmark_factory(get_row_coder(generate_fn()), generate_fn)
 
 
+def importable_named_tuple():
+  return [coders_test_common.MyTypedNamedTuple('a', i) for i in range(1000)]
+
+
 def run_coder_benchmarks(
     num_runs, input_size, seed, verbose, filter_regex='.*'):
   random.seed(seed)
@@ -310,6 +315,11 @@ def run_coder_benchmarks(
       batch_row_coder_benchmark_factory(nullable_row, True),
       batch_row_coder_benchmark_factory(diverse_row, False),
       batch_row_coder_benchmark_factory(diverse_row, True),
+      coder_benchmark_factory(
+          coders.IterableCoder(
+              coders.FastPrimitivesCoder().as_deterministic_coder(
+                  step_label="step")),
+          importable_named_tuple),
   ]
 
   suite = [