Implement decode_example_np for tfds.core.features.Dataset.

PiperOrigin-RevId: 569395247

Author: marcenacp

tensorflow_datasets/core/data_sources/python.py

@@ -0,0 +1,43 @@
+# coding=utf-8
+# Copyright 2023 The TensorFlow Datasets Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Python DataSource base class."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+import dataclasses
+from typing import Any, Callable
+
+
+@dataclasses.dataclass
+class PythonDataSource(Sequence):
+  """Python data source backed by Python objects: length and __getitem__."""
+
+  length: int
+  # If you have pickling issues for this function, define it in the upper scope
+  # or get inspiration from _getitem in
+  # tensorflow_datasets/core/features/dataset_feature.py.
+  getitem: Callable[[int], Any]
+
+  def __len__(self) -> int:
+    return self.length
+
+  def __iter__(self):
+    for i in range(self.length):
+      yield self[i]
+
+  def __getitem__(self, i: int) -> Any:
+    return self.getitem(i)

tensorflow_datasets/core/data_sources/python_test.py

@@ -0,0 +1,49 @@
+# coding=utf-8
+# Copyright 2023 The TensorFlow Datasets Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for the PythonDataSource."""
+
+import pickle
+
+from tensorflow_datasets.core.data_sources import python
+
+
+def getitem(i):
+  return i
+
+
+def test_create_a_python_data_source():
+  source = python.PythonDataSource(length=2, getitem=getitem)
+  assert len(source) == 2
+  assert source[0] == 0
+  assert source[1] == 1
+  assert source[2] == 2
+
+
+def test_iterate_on_a_python_data_source():
+  source = python.PythonDataSource(length=42, getitem=getitem)
+  i = 0
+  for i, j in enumerate(iter(source)):
+    assert i == j
+  assert i == 41
+
+
+def test_python_data_source_is_pickable():
+  def _getitem(i):
+    return i
+
+  source = python.PythonDataSource(length=42, getitem=_getitem)
+  source = pickle.loads(pickle.dumps(source))
+  assert source[0] == 0

tensorflow_datasets/core/decode/base.py

@@ -84,6 +84,18 @@ def decode_example(self, serialized_example):
     """
     raise NotImplementedError('Abstract class')
 
+  def decode_example_np(self, serialized_example):
+    """Decode the example feature field for NumPy (eg: image).
+
+    Args:
+      serialized_example: `np.array` as decoded, the dtype/shape should be
+        identical to `feature.get_serialized_info()`.
+
+    Returns:
+      example: Decoded example. Defaults to `decode_example`.
+    """
+    return self.decode_example(serialized_example)
+
   def decode_batch_example(self, serialized_example):
     """See `FeatureConnector.decode_batch_example` for details."""
     return tf.map_fn(

tensorflow_datasets/core/features/dataset_feature.py

@@ -16,17 +16,33 @@
 """Dataset feature for nested datasets."""
 from __future__ import annotations
 
+import dataclasses
 import functools
-from typing import Any, Dict, Iterator, Union
+from typing import Any, Callable, Dict, Iterator, Union
 
+from tensorflow_datasets.core.data_sources import python
 from tensorflow_datasets.core.features import feature as feature_lib
 from tensorflow_datasets.core.features import sequence_feature
+from tensorflow_datasets.core.features import tensor_feature
+from tensorflow_datasets.core.features import top_level_feature
 from tensorflow_datasets.core.utils import py_utils
 from tensorflow_datasets.core.utils import tree_utils
 from tensorflow_datasets.core.utils import type_utils
 from tensorflow_datasets.core.utils.lazy_imports_utils import tensorflow as tf
 
 
+@dataclasses.dataclass(frozen=True)
+class _getitem:  # pylint: disable=invalid-name
+  """A pickable version of getitem that can be fed to Beam pipelines."""
+
+  decode_fn: Callable[[Any], Any]
+  nest: Callable[[Any], Any]
+  flat_example: list[Any]
+
+  def __call__(self, i):
+    return self.decode_fn(self.nest([v[i] for v in self.flat_example]))
+
+
 class Dataset(sequence_feature.Sequence):
   """A Dataset feature encodes a nested dataset.
 
@@ -35,6 +51,12 @@ class Dataset(sequence_feature.Sequence):
   top-level `tf.data.Dataset` returned by `tfds.load`. At generation time, an
   iterable over the dataset elements is given.
 
+  If you use tfds.data_source and the NumPy path, `Dataset` will return
+  a [data
+  source](https://www.tensorflow.org/datasets/api_docs/python/tfds/data_source).
+  The advantage of having a data source is that decoding will be lazily executed
+  when you access each example in the dataset.
+
   This is an experimental feature. Currently, only one level of nesting is
   supported and TF1 graph is not supported either.
 
@@ -146,6 +168,54 @@ def decode_example(self, serialized_example, decoders=None):
     )
     return ds
 
+  def decode_example_np(
+      self, serialized_example, decoders=None
+  ) -> python.PythonDataSource:
+    """See base class for details."""
+    flatten = self.feature._flatten  # pylint: disable=protected-access
+    nest = self.feature._nest  # pylint: disable=protected-access
+    flat_example = flatten(serialized_example)
+    flat_features = flatten(self.feature)
+    num_slices: int | None = None
+
+    # First discover the number of slices in the Dataset. Notably, it's possible
+    # that tensors have to be reshaped. We call slice a record in the Dataset.
+    # We don't use `example` to avoid confusion with the `serialized_example`.
+    for i, feature in enumerate(flat_features):
+      if isinstance(feature, tensor_feature.Tensor) and feature.shape:
+        try:
+          flat_example[i] = flat_example[i].reshape((-1,) + feature.shape)
+        except ValueError as e:
+          raise ValueError(
+              "The length of all elements of one slice should be the same."
+          ) from e
+        feature_num_slices = flat_example[i].shape[0]
+      else:
+        feature_num_slices = len(flat_example[i])
+      if num_slices is None:
+        num_slices = feature_num_slices
+      else:
+        if feature_num_slices != num_slices:
+          raise ValueError(
+              "The length of elements of all slices should be the same. Got"
+              f" {num_slices} and {feature_num_slices}"
+          )
+    if num_slices is None:
+      raise ValueError("no feature was found.")
+
+    # Then, we apply the decode function on each slice.
+    if isinstance(self.feature, top_level_feature.TopLevelFeature):
+      # Only top-level features accept decoders.
+      decode_fn = functools.partial(
+          self.feature.decode_example_np, decoders=decoders
+      )
+    else:
+      decode_fn = self.feature.decode_example_np
+
+    return python.PythonDataSource(
+        length=num_slices, getitem=_getitem(decode_fn, nest, flat_example)
+    )
+
   def _flatten(self, x):
     """See base class for details."""
     return [x]