Lazy load array_record.

This is a short-term solution to the installation problems on macOS since 4.9.0.
This also makes sense in the long-term, because array_record is not used by all users.

The long-term solution to fix macOS installation is to build proper wheels for array_record on macos/windows platforms, but this will take more time.

PiperOrigin-RevId: 523373538

Author: marcenacp

tensorflow_datasets/core/data_sources/array_record.py

@@ -28,10 +28,9 @@
 from tensorflow_datasets.core import file_adapters
 from tensorflow_datasets.core import splits as splits_lib
 from tensorflow_datasets.core.utils import type_utils
+from tensorflow_datasets.core.utils.lazy_imports_utils import array_record_data_source
 import tree
 
-from array_record.python import array_record_data_source
-
 T = TypeVar('T')
 
 _DEFAULT_ITERATION_STEP = 1000
@@ -55,9 +54,9 @@ class ArrayRecordDataSource(AbcSequence):
       None
   )
   iteration_step: int = _DEFAULT_ITERATION_STEP
-  data_source: array_record_data_source.ArrayRecordDataSource = (
-      dataclasses.field(init=False)
-  )
+  # In order to lazy load array_record, we don't load
+  # `array_record_data_source.ArrayRecordDataSource` here.
+  data_source: Any = dataclasses.field(init=False)
   length: int = dataclasses.field(init=False)
 
   def __post_init__(self):

tensorflow_datasets/core/data_sources/array_record_test.py

@@ -25,8 +25,7 @@
 from tensorflow_datasets.core import splits as splits_lib
 from tensorflow_datasets.core.data_sources import array_record
 from tensorflow_datasets.core.utils import shard_utils
-
-from array_record.python import array_record_data_source
+from tensorflow_datasets.core.utils.lazy_imports_utils import array_record_data_source
 
 
 _FILE_INSTRUCTIONS = [

tensorflow_datasets/core/dataset_builder.py

@@ -59,7 +59,6 @@
 if typing.TYPE_CHECKING:
   import apache_beam as beam
   from apache_beam.runners import runner
-  from array_record.python import array_record_data_source
 
 ListOrTreeOrElem = type_utils.ListOrTreeOrElem
 Tree = type_utils.Tree

tensorflow_datasets/core/file_adapters.py

@@ -24,10 +24,9 @@
 
 from etils import epath
 from tensorflow_datasets.core.utils import type_utils
+from tensorflow_datasets.core.utils.lazy_imports_utils import array_record_module
 from tensorflow_datasets.core.utils.lazy_imports_utils import tensorflow as tf
 
-from array_record.python import array_record_module
-
 ExamplePositions = List[Any]
 
 

tensorflow_datasets/core/folder_dataset/compute_split_utils.py

@@ -30,6 +30,7 @@
 from tensorflow_datasets.core import splits as split_lib
 from tensorflow_datasets.core import utils
 from tensorflow_datasets.core.proto import dataset_info_pb2
+from tensorflow_datasets.core.utils.lazy_imports_utils import array_record_module
 
 from google.protobuf import json_format
 

tensorflow_datasets/core/utils/lazy_imports_utils.py

@@ -185,9 +185,26 @@ def tf_success_callback(**kwargs):
   ensure_tf_version(kwargs["module"])
 
 
+def array_record_error_callback(**kwargs):
+  del kwargs
+  print("\n\n***************************************************************")
+  print(
+      "Failed to import ArrayRecord. This probably means that you are running"
+      " on macOS or Windows. ArrayRecord currently does not work for your"
+      " infrastructure, because it uses Python bindings in C++. We are actively"
+      " working on this issue. Thanks for your understanding."
+  )
+  print("***************************************************************\n\n")
+
+
 with lazy_imports(
     error_callback=tf_error_callback, success_callback=tf_success_callback
 ):
   import tensorflow as tf  # pylint: disable=g-import-not-at-top,unused-import
 
+
+with lazy_imports(error_callback=array_record_error_callback):
+  from array_record.python import array_record_data_source  # pylint: disable=g-import-not-at-top,unused-import
+  from array_record.python import array_record_module  # pylint: disable=g-import-not-at-top,unused-import
+
 tensorflow = tf

tensorflow_datasets/import_without_tf_test.py

@@ -60,6 +60,7 @@ def _generate_examples(self):
 def test_import_tfds_without_loading_tf():
   with mock.patch.object(logging, 'log_first_n') as log_first_n:
     assert 'tensorflow' not in sys.modules
+    assert 'array_record' not in sys.modules
 
     data_dir = '/tmp/import_without_tf'
     builder = DummyDataset(data_dir=data_dir)
@@ -77,3 +78,4 @@ def test_import_tfds_without_loading_tf():
     # No warning concerning TensorFlow DTypes was dispatched while loading
     assert not log_first_n.called
     assert 'tensorflow' not in sys.modules
+    assert 'array_record' in sys.modules