Pass memoryviews on pyarrow.Buffers instead of copying to bytes.

PiperOrigin-RevId: 627396776

Author: marcenacp

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 [bytes(element.as_buffer()) for element in elements]
+    return [memoryview(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 bytes(element[0].as_buffer())
+      return memoryview(element[0].as_buffer())
     raise IndexError(f'Could not find element at index {key}')
 
   def __len__(self) -> int:

tensorflow_datasets/core/example_parser.py

@@ -128,7 +128,7 @@ def __post_init__(self):
     self._flat_example_specs = utils.flatten_nest_dict(self.example_specs)
 
   def parse_example(
-      self, serialized_example: bytes
+      self, serialized_example: bytes | memoryview
   ) -> Mapping[str, Union[np.ndarray, list[Any]]]:
     example = tf_example_pb2.Example.FromString(serialized_example)
     np_example = _features_to_numpy(example.features, self._flat_example_specs)

tensorflow_datasets/core/features/top_level_feature.py

@@ -19,7 +19,7 @@
 
 import enum
 import functools
-from typing import Any, List, Union
+from typing import Any
 
 from tensorflow_datasets.core import example_parser
 from tensorflow_datasets.core import example_serializer
@@ -44,15 +44,15 @@ class TopLevelFeature(feature_lib.FeatureConnector):
   """
 
   @functools.cached_property
-  def flat_features(self) -> List[Any]:
+  def flat_features(self) -> list[Any]:
     return self._flatten(self)
 
   @functools.cached_property
-  def flat_serialized_info(self) -> List[Any]:
+  def flat_serialized_info(self) -> list[Any]:
     return self._flatten(self.get_serialized_info())
 
   @functools.cached_property
-  def flat_sequence_ranks(self) -> List[int]:
+  def flat_sequence_ranks(self) -> list[int]:
     return [_get_sequence_rank(s) for s in self.flat_serialized_info]  # pylint: disable=not-an-iterable
 
   def _decode_example_generic(
@@ -129,7 +129,7 @@ def serialize_example(self, example_data) -> bytes:
 
   def deserialize_example(
       self,
-      serialized_example: Union[tf.Tensor, bytes],
+      serialized_example: tf.Tensor | bytes,
       *,
       decoders=None,
   ) -> utils.TensorDict:
@@ -151,7 +151,7 @@ def deserialize_example(
 
   def deserialize_example_np(
       self,
-      serialized_example: Union[tf.Tensor, bytes],
+      serialized_example: tf.Tensor | bytes | memoryview,
       *,
       decoders=None,
   ) -> utils.NpArrayOrScalarDict: