Workaround that hopefully fixes pickling segfaults with Python 3.10.

Beam depends on dill version with seemingly outdated logic for CodeType pickling which causes occasional segfaults in Python 3.10. This change patches the pickling logic to a newer version. This is a temporary workaround until Beam depends on a newer version of dill (or this workaround is propagated to Beam) and we update the beam dependency.

PiperOrigin-RevId: 546023626

Author: iindyk

tensorflow_transform/beam/impl.py

@@ -80,6 +80,7 @@
 from tensorflow_transform.tf_metadata import dataset_metadata
 from tensorflow_transform.tf_metadata import metadata_io
 from tensorflow_transform.tf_metadata import schema_utils
+from tfx_bsl import beam as tfx_bsl_beam
 from tfx_bsl.coders import example_coder
 from tfx_bsl.telemetry import collection as telemetry
 from tfx_bsl.telemetry import util as telemetry_util
@@ -94,6 +95,8 @@
 
 from tensorflow_metadata.proto.v0 import schema_pb2
 
+tfx_bsl_beam.fix_code_type_pickling()
+
 _TransformFnPathType = str
 
 Context = context.Context