Move tensorflow-io import to transform_fn_io from __init__.

PiperOrigin-RevId: 439928415

Author: varshaan

tensorflow_transform/beam/tft_beam_io/__init__.py

@@ -16,10 +16,3 @@
 from tensorflow_transform.beam.tft_beam_io.beam_metadata_io import WriteMetadata
 from tensorflow_transform.beam.tft_beam_io.transform_fn_io import ReadTransformFn
 from tensorflow_transform.beam.tft_beam_io.transform_fn_io import WriteTransformFn
-# TF 2.6 split support for filesystems such as Amazon S3 out to the
-# `tensorflow_io` package. Hence, this import is needed wherever we touch the
-# filesystem.
-try:
-  import tensorflow_io as _  # pytype: disable=import-error # pylint: disable=g-import-not-at-top
-except ModuleNotFoundError:
-  pass

tensorflow_transform/beam/tft_beam_io/transform_fn_io.py

@@ -15,6 +15,7 @@
 
 import os
 
+from absl import logging
 import apache_beam as beam
 import tensorflow_transform as tft
 from tensorflow_transform import impl_helper
@@ -99,6 +100,13 @@ def expand(self, transform_fn):
     def publish_outputs(unused_element, metadata_source_path,
                         transform_fn_source_path):
       import tensorflow as tf  # pylint: disable=g-import-not-at-top
+      # TF 2.6 split support for filesystems such as Amazon S3 out to the
+      # `tensorflow_io` package. Hence, this import is needed wherever we touch
+      # the filesystem.
+      try:
+        import tensorflow_io as _  # pytype: disable=import-error # pylint: disable=g-import-not-at-top
+      except ModuleNotFoundError:
+        logging.info('tensorflow_io is not available.')
       if not tf.io.gfile.exists(self._path):
         tf.io.gfile.makedirs(self._path)