tfrecord writers: clearer error message when no examples are yielded or empty pcollection as input.

PiperOrigin-RevId: 289162141

Author: pierrot0

tensorflow_datasets/core/tfrecords_writer.py

@@ -103,6 +103,8 @@ def _get_shard_specs(num_examples, total_size, bucket_lengths, path):
 
 
 def _get_shard_boundaries(num_examples, number_of_shards):
+  if num_examples == 0:
+    raise AssertionError("No examples were yielded.")
   if num_examples < number_of_shards:
     raise AssertionError("num_examples ({}) < number_of_shards ({})".format(
         num_examples, number_of_shards))
@@ -306,6 +308,8 @@ def _get_boundaries_per_bucket_shard(self, shard_len_sizes):
         examples in the bucket and size is the total size in bytes of the
         elements in that bucket. Buckets with no elements are not mentioned.
     """
+    if not shard_len_sizes:
+      raise AssertionError("Not a single example present in the PCollection!")
     total_num_examples = 0
     total_size = 0
     bucket2length = {}

tensorflow_datasets/core/tfrecords_writer_test.py

@@ -131,6 +131,9 @@ def _read_records(path):
 
 class WriterTest(testing.TestCase):
 
+  EMPTY_SPLIT_ERROR = 'No examples were yielded.'
+  TOO_SMALL_SPLIT_ERROR = 'num_examples (1) < number_of_shards (2)'
+
   @absltest.mock.patch.object(
       example_serializer, 'ExampleSerializer', testing.DummySerializer)
   def _write(self, to_write, path, salt=''):
@@ -174,9 +177,29 @@ def test_write_duplicated_keys(self):
           AssertionError, 'Two records share the same hashed key!'):
         self._write(to_write, path)
 
+  def test_empty_split(self):
+    path = os.path.join(self.tmp_dir, 'foo.tfrecord')
+    to_write = []
+    with absltest.mock.patch.object(tfrecords_writer, '_get_number_shards',
+                                    return_value=1):
+      with self.assertRaisesWithPredicateMatch(
+          AssertionError, self.EMPTY_SPLIT_ERROR):
+        self._write(to_write, path)
+
+  def test_too_small_split(self):
+    path = os.path.join(self.tmp_dir, 'foo.tfrecord')
+    to_write = [(1, b'a')]
+    with absltest.mock.patch.object(tfrecords_writer, '_get_number_shards',
+                                    return_value=2):
+      with self.assertRaisesWithPredicateMatch(
+          AssertionError, self.TOO_SMALL_SPLIT_ERROR):
+        self._write(to_write, path)
+
 
 class TfrecordsWriterBeamTest(WriterTest):
 
+  EMPTY_SPLIT_ERROR = 'Not a single example present in the PCollection!'
+
   @absltest.mock.patch.object(
       example_serializer, 'ExampleSerializer', testing.DummySerializer)
   def _write(self, to_write, path, salt=''):