Lazy import tree in TFDS.

PiperOrigin-RevId: 627352990

Author: tomvdw

tensorflow_datasets/core/data_sources/base.py

@@ -26,7 +26,7 @@
 from tensorflow_datasets.core import splits as splits_lib
 from tensorflow_datasets.core.utils import shard_utils
 from tensorflow_datasets.core.utils import type_utils
-import tree
+from tensorflow_datasets.core.utils.lazy_imports_utils import tree
 
 T = TypeVar('T')
 

tensorflow_datasets/core/data_sources/parquet.py

@@ -38,14 +38,14 @@ def __getitems__(self, keys: Iterable[int]) -> Sequence[Any]:
       return []
     # All elements are written in the first column (see core.ParquetFileAdapter)
     elements = self.table.take(keys).column(0)
-    return [element.as_buffer() for element in elements]
+    return [bytes(element.as_buffer()) for element in elements]
 
   def __getitem__(self, key: int) -> Any:
     """Retrieves the n-th element from the Parquet table."""
     # The element is written in the first column (see core.ParquetFileAdapter)
     element = self.table.slice(key, 1).column(0)
     if len(element) == 1:
-      return element[0].as_buffer()
+      return bytes(element[0].as_buffer())
     raise IndexError(f'Could not find element at index {key}')
 
   def __len__(self) -> int:

tensorflow_datasets/core/dataset_builder.py

@@ -57,8 +57,8 @@
 from tensorflow_datasets.core.utils import type_utils
 from tensorflow_datasets.core.utils.lazy_imports_utils import apache_beam as beam
 from tensorflow_datasets.core.utils.lazy_imports_utils import tensorflow as tf
+from tensorflow_datasets.core.utils.lazy_imports_utils import tree
 import termcolor
-import tree
 
 
 ListOrTreeOrElem = type_utils.ListOrTreeOrElem

tensorflow_datasets/core/dataset_utils.py

@@ -29,7 +29,7 @@
 from tensorflow_datasets.core import utils
 from tensorflow_datasets.core.utils import type_utils
 from tensorflow_datasets.core.utils.lazy_imports_utils import tensorflow as tf
-import tree
+from tensorflow_datasets.core.utils.lazy_imports_utils import tree
 
 Tree = type_utils.Tree
 Tensor = type_utils.Tensor
@@ -208,9 +208,7 @@ def dataset_shape_is_fully_defined(ds):
 
 
 def features_shape_is_fully_defined(features):
-  return all(
-      [
-          tf.TensorShape(info.shape).is_fully_defined()
-          for info in tf.nest.flatten(features.get_tensor_info())
-      ]
-  )
+  return all([
+      tf.TensorShape(info.shape).is_fully_defined()
+      for info in tf.nest.flatten(features.get_tensor_info())
+  ])

tensorflow_datasets/core/decode/base.py

@@ -19,7 +19,7 @@
 import functools
 
 from tensorflow_datasets.core.utils.lazy_imports_utils import tensorflow as tf
-import tree
+from tensorflow_datasets.core.utils.lazy_imports_utils import tree
 
 
 class Decoder(abc.ABC):

tensorflow_datasets/core/features/dataset_feature.py

@@ -29,7 +29,7 @@
 from tensorflow_datasets.core.utils import py_utils
 from tensorflow_datasets.core.utils import type_utils
 from tensorflow_datasets.core.utils.lazy_imports_utils import tensorflow as tf
-import tree
+from tensorflow_datasets.core.utils.lazy_imports_utils import tree
 
 
 @dataclasses.dataclass(frozen=True)

tensorflow_datasets/core/features/feature.py

@@ -41,7 +41,7 @@
 from tensorflow_datasets.core.utils import type_utils
 from tensorflow_datasets.core.utils.lazy_imports_utils import tensorflow as tf
 from tensorflow_datasets.core.utils.lazy_imports_utils import tf_agents
-import tree
+from tensorflow_datasets.core.utils.lazy_imports_utils import tree
 
 from google.protobuf import descriptor
 from google.protobuf import json_format

tensorflow_datasets/core/features/sequence_feature.py

@@ -28,7 +28,7 @@
 from tensorflow_datasets.core.proto import feature_pb2
 from tensorflow_datasets.core.utils import py_utils
 from tensorflow_datasets.core.utils import type_utils
-import tree
+from tensorflow_datasets.core.utils.lazy_imports_utils import tree
 
 Json = type_utils.Json
 

tensorflow_datasets/core/reader.py

@@ -33,7 +33,7 @@
 from tensorflow_datasets.core.utils import read_config as read_config_lib
 from tensorflow_datasets.core.utils import shard_utils
 from tensorflow_datasets.core.utils.lazy_imports_utils import tensorflow as tf
-import tree
+from tensorflow_datasets.core.utils.lazy_imports_utils import tree
 
 Tree = utils.Tree
 TreeDict = utils.TreeDict

tensorflow_datasets/core/utils/lazy_imports_utils.py

@@ -112,6 +112,7 @@ def datasets_error_callback(module_name: Exception):
   import pandas
   import pyarrow
   from pyarrow import parquet
+  import tree
 
 
 with epy.lazy_imports(