Lazy load metadata to save disk access when user doesn't use metadata

PiperOrigin-RevId: 694450576

Author: tomvdw

tensorflow_datasets/core/dataset_info.py

@@ -747,12 +747,16 @@ def read_from_directory(self, dataset_info_dir: epath.PathLike) -> None:
           dataset_info_dir
       )
 
-    # Restore the MetaDataDict from metadata.json if there is any
-    if _metadata_filepath(dataset_info_dir).exists():
-      # If the dataset was loaded from file, self.metadata will be `None`, so
-      # we create a MetadataDict first.
-      if self._metadata is None:
-        self._metadata = MetadataDict()
+    # If the dataset was loaded from file, self.metadata will be `None`, so
+    # we create a MetadataDict first.
+    if self._metadata is None:
+      self._metadata = LazyMetadataDict(dataset_info_dir)
+    elif isinstance(self._metadata, MetadataDict):
+      lazy_metadata = LazyMetadataDict(dataset_info_dir)
+      lazy_metadata.update(self._metadata)
+      self._metadata = lazy_metadata
+    elif _metadata_filepath(dataset_info_dir).exists():
+      # Restore the MetaDataDict from metadata.json if there is any
       self._metadata.load_metadata(dataset_info_dir)
 
     # Update fields which are not defined in the code. This means that
@@ -1375,6 +1379,12 @@ def add_tfds_data_source_access(
   )
 
 
+def _load_metadata_from_file(data_dir: epath.PathLike) -> dict[str, Any]:
+  """Loads metadata from file."""
+  with _metadata_filepath(data_dir).open(mode="r") as f:
+    return json.load(f)
+
+
 class MetadataDict(Metadata, dict):
   """A `tfds.core.Metadata` object that acts as a `dict`.
 
@@ -1389,8 +1399,41 @@ def save_metadata(self, data_dir):
   def load_metadata(self, data_dir):
     """Restore the metadata."""
     self.clear()
-    with _metadata_filepath(data_dir).open(mode="r") as f:
-      self.update(json.load(f))
+    self.update(_load_metadata_from_file(data_dir))
+
+
+class LazyMetadataDict(MetadataDict):
+  """A `tfds.core.Metadata` object that acts as a `dict`.
+
+  Content is lazily loaded from the given data directory.
+  """
+
+  def __init__(self, data_dir: epath.PathLike) -> None:
+    self._data_dir = epath.Path(data_dir)
+    self._data_is_loaded = False
+    super().__init__()
+
+  def _load_metadata(self):
+    if not self._data_is_loaded:
+      if _metadata_filepath(self._data_dir).exists():
+        self.load_metadata(self._data_dir)
+      self._data_is_loaded = True
+
+  def __getitem__(self, key, /):
+    self._load_metadata()
+    return super().__getitem__(key)
+
+  def __eq__(self, value, /):
+    self._load_metadata()
+    return super().__eq__(value)
+
+  def keys(self):
+    self._load_metadata()
+    return super().keys()
+
+  def items(self):
+    self._load_metadata()
+    return super().items()
 
 
 class BeamMetadataDict(MetadataDict):

tensorflow_datasets/core/dataset_info_test.py

@@ -13,16 +13,16 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Tests for tensorflow_datasets.core.dataset_info."""
-
 import json
 import os
 import pathlib
 import re
 import tempfile
 import time
 from typing import Union
+from unittest import mock
 
+from etils import epath
 import numpy as np
 import pytest
 import tensorflow as tf
@@ -836,5 +836,59 @@ def test_get_split_info_from_proto_unavailable_format(self):
 # pylint: enable=g-inconsistent-quotes
 
 
+class LazyMetadataDictTest(testing.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    self.data_dir = epath.Path(self.tmp_dir)
+    metadata_file = self.data_dir / "metadata.json"
+    metadata_file.write_text(json.dumps({"test": "test123"}))
+
+  def test_load_metadata(self):
+    metadata_dict = dataset_info.LazyMetadataDict(self.data_dir)
+    self.assertEqual(metadata_dict["test"], "test123")
+
+  @mock.patch.object(dataset_info, "_load_metadata_from_file")
+  def test_get_item(self, mock_load_metadata_from_file):
+    mock_load_metadata_from_file.return_value = {"test": "test123"}
+
+    metadata_dict = dataset_info.LazyMetadataDict(self.data_dir)
+    mock_load_metadata_from_file.assert_not_called()
+
+    self.assertEqual(metadata_dict["test"], "test123")
+    mock_load_metadata_from_file.assert_called_with(self.data_dir)
+
+    self.assertEqual(sorted(metadata_dict.keys()), ["test"])
+
+    # Update the metadata.
+    metadata_dict["abc"] = "def"
+    self.assertEqual(metadata_dict["abc"], "def")
+
+  @mock.patch.object(dataset_info, "_load_metadata_from_file")
+  def test_keys(self, mock_load_metadata_from_file):
+    mock_load_metadata_from_file.return_value = {"test": "test123"}
+    metadata_dict = dataset_info.LazyMetadataDict(self.data_dir)
+    mock_load_metadata_from_file.assert_not_called()
+    actual_keys = sorted(metadata_dict.keys())
+    mock_load_metadata_from_file.assert_called_with(self.data_dir)
+    self.assertEqual(actual_keys, ["test"])
+
+  @mock.patch.object(dataset_info, "_load_metadata_from_file")
+  def test_items(self, mock_load_metadata_from_file):
+    mock_load_metadata_from_file.return_value = {"test": "test123"}
+    metadata_dict = dataset_info.LazyMetadataDict(self.data_dir)
+    mock_load_metadata_from_file.assert_not_called()
+    self.assertNotEmpty(metadata_dict.items())
+    mock_load_metadata_from_file.assert_called_with(self.data_dir)
+
+  @mock.patch.object(dataset_info, "_load_metadata_from_file")
+  def test_updating_before_loading(self, mock_load_metadata_from_file):
+    metadata_dict = dataset_info.LazyMetadataDict(self.data_dir)
+    mock_load_metadata_from_file.assert_not_called()
+
+    metadata_dict.update({"abc": "def", "ghi": "jkl"})
+    mock_load_metadata_from_file.assert_not_called()
+
+
 if __name__ == "__main__":
   testing.test_main()