Make size_in_bytes reflect final dataset size instead of downloaded size.

PiperOrigin-RevId: 292190740

Author: adarob

tensorflow_datasets/core/dataset_builder.py

@@ -309,9 +309,15 @@ def download_and_prepare(self, download_dir=None, download_config=None):
 
     logging.info("Generating dataset %s (%s)", self.name, self._data_dir)
     if not utils.has_sufficient_disk_space(
-        self.info.size_in_bytes, directory=self._data_dir_root):
-      raise IOError("Not enough disk space. Needed: %s" %
-                    units.size_str(self.info.size_in_bytes))
+        self.info.size_in_bytes + self.info.download_size,
+        directory=self._data_dir_root):
+      raise IOError(
+          "Not enough disk space. Needed: {} (download: {}, generated: {})"
+          .format(
+              units.size_str(self.info.size_in_bytes + self.info.download_size),
+              units.size_str(self.info.download_size),
+              units.size_str(self.info.size_in_bytes),
+          ))
     self._log_download_bytes()
 
     dl_manager = self._make_download_manager(
@@ -352,7 +358,7 @@ def download_and_prepare(self, download_dir=None, download_config=None):
           else:  # Mode is forced or stats do not exists yet
             logging.info("Computing statistics.")
             self.info.compute_dynamic_properties()
-          self.info.size_in_bytes = dl_manager.downloaded_size
+          self.info.downloaded_size = dl_manager.downloaded_size
           # Write DatasetInfo to disk, even if we haven't computed statistics.
           self.info.write_to_directory(self._data_dir)
     self._log_download_done()
@@ -1055,8 +1061,9 @@ def _prepare_split(self, split_generator, max_examples_per_split):
                                   total=split_info.num_examples, leave=False):
       example = self.info.features.encode_example(record)
       writer.write(key, example)
-    shard_lengths = writer.finalize()
+    shard_lengths, total_size = writer.finalize()
     split_generator.split_info.shard_lengths.extend(shard_lengths)
+    split_generator.split_info.num_bytes = total_size
 
 
 class BeamBasedBuilder(FileAdapterBuilder):
@@ -1148,10 +1155,11 @@ def _download_and_prepare(self, dl_manager, download_config):
     split_dict = self.info.splits
     for split_name, beam_writer in self._beam_writers.items():
       logging.info("Retrieving shard lengths for %s...", split_name)
-      shard_lengths = beam_writer.finalize()
+      shard_lengths, total_size = beam_writer.finalize()
       split_info = split_dict[split_name]
       split_info.shard_lengths.extend(shard_lengths)
       split_info.num_shards = len(shard_lengths)
+      split_info.num_bytes = total_size
     logging.info("Updating split info...")
     self.info.update_splits_if_different(split_dict)
 

tensorflow_datasets/core/dataset_info.py

@@ -205,11 +205,17 @@ def citation(self):
 
   @property
   def size_in_bytes(self):
-    return self.as_proto.size_in_bytes
+    size_in_bytes = sum(split.num_bytes for split in self.splits.values())
+    # Fall back to deprecated proto field if `num_bytes` fields are empty.
+    return size_in_bytes or self.as_proto.size_in_bytes
 
-  @size_in_bytes.setter
-  def size_in_bytes(self, size):
-    self.as_proto.size_in_bytes = size
+  @property
+  def download_size(self):
+    return self.as_proto.download_size
+
+  @download_size.setter
+  def download_size(self, size):
+    self.as_proto.download_size = size
 
   @property
   def features(self):

tensorflow_datasets/core/dataset_info_test.py

@@ -113,6 +113,7 @@ def test_reading(self):
 
     # Assert that this is computed correctly.
     self.assertEqual(40, info.splits.total_num_examples)
+    self.assertEqual(11594722, info.size_in_bytes)
 
     self.assertEqual("image", info.supervised_keys[0])
     self.assertEqual("label", info.supervised_keys[1])
@@ -169,7 +170,9 @@ def test_restore_after_modification(self):
         citation="some citation",
         redistribution_info={"license": "some license"}
     )
-    info.size_in_bytes = 456
+    info.download_size = 456
+    info.as_proto.splits.add(name="train", num_bytes=512)
+    info.as_proto.splits.add(name="validation", num_bytes=64)
     info.as_proto.schema.feature.add()
     info.as_proto.schema.feature.add()  # Add dynamic statistics
     info.download_checksums = {
@@ -198,7 +201,8 @@ def test_restore_after_modification(self):
           citation="some citation (new)",
           redistribution_info={"license": "some license (new)"}
       )
-      restored_info.size_in_bytes = 789
+      restored_info.download_size = 789
+      restored_info.as_proto.splits.add(name="validation", num_bytes=288)
       restored_info.as_proto.schema.feature.add()
       restored_info.as_proto.schema.feature.add()
       restored_info.as_proto.schema.feature.add()
@@ -219,7 +223,8 @@ def test_restore_after_modification(self):
       self.assertEqual(restored_info.citation, "some citation (new)")
       self.assertEqual(restored_info.redistribution_info.license,
                        "some license (new)")
-      self.assertEqual(restored_info.size_in_bytes, 789)
+      self.assertEqual(restored_info.download_size, 789)
+      self.assertEqual(restored_info.size_in_bytes, 576)
       self.assertEqual(len(restored_info.as_proto.schema.feature), 4)
       self.assertEqual(restored_info.download_checksums, {
           "url2": "some other checksum (new)",

tensorflow_datasets/core/proto/dataset_info.proto

@@ -28,11 +28,13 @@ message SplitInfo {
   int64 num_shards = 2;
   // The number of examples in each shard.
   repeated int64 shard_lengths = 4;
+  // The number of bytes in the split.
+  int64 num_bytes = 5;
 
   // The concrete statistics about this split.
   tensorflow.metadata.v0.DatasetFeatureStatistics statistics = 3;
 
-  // Next available: 5.
+  // Next available: 6.
 }
 
 // This message indicates which feature in the dataset schema is the input and
@@ -59,8 +61,12 @@ message DatasetInfo {
   // A citation string if one exists for this dataset.
   string citation = 3;
 
+  // DEPRECATED
   // *Approximate* size in bytes of this dataset on disk.
-  int64 size_in_bytes = 4;
+  int64 size_in_bytes = 4 [deprecated=true];
+
+  // Size in bytes of downloaded files.
+  int64 download_size = 12;
 
   // Canonical location of the dataset.
   DatasetLocation location = 5;
@@ -81,5 +87,5 @@ message DatasetInfo {
 
   RedistributionInfo redistribution_info = 11;
 
-  // Next available: 12
+  // Next available: 13
 }