Encode examples before writing them in the ShardBasedBuilder

PiperOrigin-RevId: 807804529

Author: tomvdw

tensorflow_datasets/core/__init__.py

@@ -23,6 +23,7 @@
 from tensorflow_datasets.core.dataset_builder import BuilderConfig
 from tensorflow_datasets.core.dataset_builder import DatasetBuilder
 from tensorflow_datasets.core.dataset_builder import GeneratorBasedBuilder
+from tensorflow_datasets.core.dataset_builder import ShardBasedBuilder
 from tensorflow_datasets.core.dataset_info import BeamMetadataDict
 from tensorflow_datasets.core.dataset_info import DatasetIdentity
 from tensorflow_datasets.core.dataset_info import DatasetInfo

tensorflow_datasets/core/split_builder.py

@@ -184,6 +184,7 @@ def submit_shard_based_generation(
     serializer = example_serializer.ExampleSerializer(serialized_info)
 
     shard_writer = writer_lib.ShardWriter(
+        features=self._features,
         serializer=serializer,
         example_writer=self._example_writer,
     )

tensorflow_datasets/core/writer.py

@@ -36,6 +36,7 @@
   from etils import epath
   from tensorflow_datasets.core import example_parser
   from tensorflow_datasets.core import example_serializer
+  from tensorflow_datasets.core import features as features_lib
   from tensorflow_datasets.core import file_adapters
   from tensorflow_datasets.core import hashing
   from tensorflow_datasets.core import naming
@@ -264,27 +265,34 @@ class ShardWriter:
 
   def __init__(
       self,
+      features: features_lib.FeatureConnector,
       serializer: example_serializer.Serializer,
       example_writer: ExampleWriter,
   ):
     """Initializes Writer.
 
     Args:
+      features: the features of the dataset.
       serializer: class that can serialize examples.
       example_writer: class that writes examples to disk or elsewhere.
     """
+    self._features = features
     self._serializer = serializer
     self._example_writer = example_writer
 
+  def _serialize_example(self, example: Example) -> Any:
+    """Encodes and serializes an example."""
+    return self._serializer.serialize_example(
+        self._features.encode_example(example)
+    )
+
   def write(
       self,
       examples: Iterable[KeyExample],
       path: epath.Path,
   ) -> int:
     """Returns the number of examples written to the given path."""
-    serialized_examples = [
-        (k, self._serializer.serialize_example(v)) for k, v in examples
-    ]
+    serialized_examples = [(k, self._serialize_example(v)) for k, v in examples]
     self._example_writer.write(path=path, examples=serialized_examples)
 
     return len(serialized_examples)