Add nondeterministic_order to dataset_info.py

PiperOrigin-RevId: 691798490

Author: The TensorFlow Datasets Authors

tensorflow_datasets/core/dataset_info.py

@@ -186,6 +186,7 @@ def __init__(
       features: feature_lib.FeatureConnector | None = None,
       supervised_keys: SupervisedKeysType | None = None,
       disable_shuffling: bool = False,
+      nondeterministic_order: bool = False,
       homepage: str | None = None,
       citation: str | None = None,
       metadata: Metadata | None = None,
@@ -228,7 +229,11 @@ def __init__(
 
         Note that selecting features in nested `tfds.features.FeaturesDict`
         objects is not supported.
-      disable_shuffling: `bool`, specify whether to shuffle the examples.
+      disable_shuffling: `bool`, specifies whether to shuffle the examples.
+      nondeterministic_order: `bool`, if True and the dataset uses beam, it will
+        use `NoShuffleBeamWriter` which does not assure deterministic
+        shuffling when writing' examples to disk. This might result in quicker
+        dataset preparation.
       homepage: `str`, optional, the homepage for this dataset.
       citation: `str`, optional, the citation to use for this dataset.
       metadata: `tfds.core.Metadata`, additonal object which will be
@@ -268,6 +273,7 @@ def __init__(
         version=str(self._identity.version),
         release_notes=self._identity.release_notes,
         disable_shuffling=disable_shuffling,
+        nondeterministic_order=nondeterministic_order,
         config_name=self._identity.config_name,
         config_description=self._identity.config_description,
         config_tags=self._identity.config_tags,
@@ -342,6 +348,7 @@ def from_proto(
         features=features,
         supervised_keys=supervised_keys,
         disable_shuffling=proto.disable_shuffling,
+        nondeterministic_order=proto.nondeterministic_order,
         citation=proto.citation,
         license=proto.redistribution_info.license,
         split_dict=splits_lib.SplitDict.from_proto(
@@ -400,6 +407,13 @@ def release_notes(self) -> dict[str, str] | None:
   def disable_shuffling(self) -> bool:
     return self.as_proto.disable_shuffling
 
+  @property
+  def nondeterministic_order(self) -> bool:
+    return self._info_proto.nondeterministic_order
+
+  def set_nondeterministic_order(self, nondeterministic_order: bool) -> None:
+    self._info_proto.nondeterministic_order = nondeterministic_order
+
   @property
   def homepage(self) -> str:
     urls = self.as_proto.location.urls
@@ -923,6 +937,7 @@ def __repr__(self):
         ("features", _indent(repr(self.features))),
         ("supervised_keys", self.supervised_keys),
         ("disable_shuffling", self.disable_shuffling),
+        ("nondeterministic_order", self.nondeterministic_order),
         ("splits", splits),
         ("citation", _indent(f'"""{self.citation}"""')),
         # Proto add a \n that we strip.
@@ -940,6 +955,7 @@ def __getstate__(self):
         "features": self.features,
         "supervised_keys": self.supervised_keys,
         "disable_shuffling": self.disable_shuffling,
+        "nondeterministic_order": self.nondeterministic_order,
         "homepage": self.homepage,
         "citation": self.citation,
         "metadata": self.metadata,
@@ -956,6 +972,7 @@ def __setstate__(self, state):
         features=state["features"],
         supervised_keys=state["supervised_keys"],
         disable_shuffling=state["disable_shuffling"],
+        nondeterministic_order=state["nondeterministic_order"],
         homepage=state["homepage"],
         citation=state["citation"],
         metadata=state["metadata"],

tensorflow_datasets/core/dataset_info_test.py

@@ -818,6 +818,7 @@ def test_get_split_info_from_proto_unavailable_format(self):
     }),
     supervised_keys=('image', 'label'),
     disable_shuffling=False,
+    nondeterministic_order=False,
     splits={
         'test': <SplitInfo num_examples=20, num_shards=1>,
         'train': <SplitInfo num_examples=20, num_shards=1>,