datasets 4.8.2__tar.gz → 4.8.4__tar.gz

{datasets-4.8.2/src/datasets.egg-info → datasets-4.8.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: datasets
-Version: 4.8.2
+Version: 4.8.4
 Summary: HuggingFace community-driven open-source library of datasets
 Home-page: https://github.com/huggingface/datasets
 Download-URL: https://github.com/huggingface/datasets/tags

{datasets-4.8.2 → datasets-4.8.4}/setup.py

@@ -233,7 +233,7 @@ EXTRAS_REQUIRE = {
 
 setup(
     name="datasets",
-    version="4.8.2",  # expected format is one of x.y.z.dev0, or x.y.z.rc1 or x.y.z (no to dashes, yes to dots)
+    version="4.8.4",  # expected format is one of x.y.z.dev0, or x.y.z.rc1 or x.y.z (no to dashes, yes to dots)
     description="HuggingFace community-driven open-source library of datasets",
     long_description=open("README.md", encoding="utf-8").read(),
     long_description_content_type="text/markdown",

{datasets-4.8.2 → datasets-4.8.4}/src/datasets/__init__.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-__version__ = "4.8.2"
+__version__ = "4.8.4"
 
 from .arrow_dataset import Column, Dataset
 from .arrow_reader import ReadInstruction

{datasets-4.8.2 → datasets-4.8.4}/src/datasets/features/features.py

@@ -1239,7 +1239,7 @@ class Json:
             raise RuntimeError("Decoding is disabled for this feature. Please use Json(decode=True) instead.")
         return ujson_loads(example_data)
 
-    def cast_storage(self, storage: Union[pa.Array]) -> pa.Int64Array:
+    def cast_storage(self, storage: Union[pa.Array]) -> pa.JsonArray:
         """Cast an Arrow array to the `Json` arrow storage type.
 
         Args:
@@ -1247,7 +1247,7 @@ class Json:
                 PyArrow array to cast.
 
         Returns:
-            `pa.Int64Array`: Array in the `ClassLabel` arrow storage type.
+            `pa.JsonArray`: Array in the `Json` arrow storage type.
         """
         if isinstance(storage, pa.JsonArray):
             return storage

{datasets-4.8.2 → datasets-4.8.4}/src/datasets/formatting/jax_formatter.py

@@ -107,10 +107,13 @@ class JaxFormatter(TensorFormatter[Mapping, "jax.Array", Mapping]):
             if isinstance(value, PIL.Image.Image):
                 value = np.asarray(value)
         if config.TORCHVISION_AVAILABLE and "torchvision" in sys.modules:
-            from torchvision.io import VideoReader
+            try:
+                from torchvision.io import VideoReader
 
-            if isinstance(value, VideoReader):
-                return value  # TODO(QL): set output to jax arrays ?
+                if isinstance(value, VideoReader):
+                    return value  # TODO(QL): set output to jax arrays ?
+            except ImportError:
+                pass
         if config.TORCHCODEC_AVAILABLE and "torchcodec" in sys.modules:
             from torchcodec.decoders import AudioDecoder, VideoDecoder
 

{datasets-4.8.2 → datasets-4.8.4}/src/datasets/formatting/np_formatter.py

@@ -64,10 +64,13 @@ class NumpyFormatter(TensorFormatter[Mapping, np.ndarray, Mapping]):
             if isinstance(value, PIL.Image.Image):
                 return np.asarray(value, **self.np_array_kwargs)
         if config.TORCHVISION_AVAILABLE and "torchvision" in sys.modules:
-            from torchvision.io import VideoReader
+            try:
+                from torchvision.io import VideoReader
 
-            if isinstance(value, VideoReader):
-                return value  # TODO(QL): set output to np arrays ?
+                if isinstance(value, VideoReader):
+                    return value  # TODO(QL): set output to np arrays ?
+            except ImportError:
+                pass
         if config.TORCHCODEC_AVAILABLE and "torchcodec" in sys.modules:
             from torchcodec.decoders import AudioDecoder, VideoDecoder
 

{datasets-4.8.2 → datasets-4.8.4}/src/datasets/formatting/tf_formatter.py

@@ -71,10 +71,13 @@ class TFFormatter(TensorFormatter[Mapping, "tf.Tensor", Mapping]):
             if isinstance(value, PIL.Image.Image):
                 value = np.asarray(value)
         if config.TORCHVISION_AVAILABLE and "torchvision" in sys.modules:
-            from torchvision.io import VideoReader
+            try:
+                from torchvision.io import VideoReader
 
-            if isinstance(value, VideoReader):
-                return value  # TODO(QL): set output to tf tensors ?
+                if isinstance(value, VideoReader):
+                    return value  # TODO(QL): set output to tf tensors ?
+            except ImportError:
+                pass
         if config.TORCHCODEC_AVAILABLE and "torchcodec" in sys.modules:
             from torchcodec.decoders import AudioDecoder, VideoDecoder
 

{datasets-4.8.2 → datasets-4.8.4}/src/datasets/formatting/torch_formatter.py

@@ -77,10 +77,13 @@ class TorchFormatter(TensorFormatter[Mapping, "torch.Tensor", Mapping]):
 
                 value = value.transpose((2, 0, 1))
         if config.TORCHVISION_AVAILABLE and "torchvision" in sys.modules:
-            from torchvision.io import VideoReader
+            try:
+                from torchvision.io import VideoReader
 
-            if isinstance(value, VideoReader):
-                return value  # TODO(QL): set output to torch tensors ?
+                if isinstance(value, VideoReader):
+                    return value  # TODO(QL): set output to torch tensors ?
+            except ImportError:
+                pass
         if config.TORCHCODEC_AVAILABLE and "torchcodec" in sys.modules:
             from torchcodec.decoders import AudioDecoder, VideoDecoder
 

{datasets-4.8.2 → datasets-4.8.4}/src/datasets/iterable_dataset.py

@@ -424,15 +424,22 @@ class ArrowExamplesIterable(_BaseExamplesIterable):
 
 
 class RebatchedArrowExamplesIterable(_BaseExamplesIterable):
-    def __init__(self, ex_iterable: _BaseExamplesIterable, batch_size: Optional[int], drop_last_batch: bool = False):
+    def __init__(
+        self,
+        ex_iterable: _BaseExamplesIterable,
+        batch_size: Optional[int],
+        drop_last_batch: bool = False,
+        force_convert_to_arrow: bool = False,
+    ):
         super().__init__()
         self.ex_iterable = ex_iterable
         self.batch_size = batch_size
         self.drop_last_batch = drop_last_batch
+        self.force_convert_to_arrow = force_convert_to_arrow
 
     @property
     def iter_arrow(self):
-        return self._iter_arrow
+        return self._iter_arrow if self.ex_iterable.iter_arrow or self.force_convert_to_arrow else None
 
     @property
     def is_typed(self):
@@ -462,8 +469,12 @@ class RebatchedArrowExamplesIterable(_BaseExamplesIterable):
             self.ex_iterable.load_state_dict(self._state_dict["previous_state"])
         if self.ex_iterable.iter_arrow:
             iterator = self.ex_iterable.iter_arrow()
-        else:
+        elif self.force_convert_to_arrow:
             iterator = _convert_to_arrow(self.ex_iterable, batch_size=1)
+        else:
+            raise RuntimeError(
+                "_iter_arrow is not available in RebatchedArrowExamplesIterable, use an examples iterable that implements _iter_arrow() or pass force_convert_to_arrow=True"
+            )
         if self.batch_size is None or self.batch_size <= 0:
             if self._state_dict and self._state_dict["batch_idx"] > 0:
                 return
@@ -544,7 +555,10 @@ class RebatchedArrowExamplesIterable(_BaseExamplesIterable):
 
     def shuffle_data_sources(self, generator: np.random.Generator) -> "RebatchedArrowExamplesIterable":
         return RebatchedArrowExamplesIterable(
-            self.ex_iterable.shuffle_data_sources(generator), self.batch_size, self.drop_last_batch
+            self.ex_iterable.shuffle_data_sources(generator),
+            self.batch_size,
+            self.drop_last_batch,
+            self.force_convert_to_arrow,
         )
 
     def shard_data_sources(self, num_shards: int, index: int, contiguous=True) -> "RebatchedArrowExamplesIterable":
@@ -552,13 +566,12 @@ class RebatchedArrowExamplesIterable(_BaseExamplesIterable):
             self.ex_iterable.shard_data_sources(num_shards, index, contiguous=contiguous),
             self.batch_size,
             self.drop_last_batch,
+            self.force_convert_to_arrow,
         )
 
     def reshard_data_sources(self) -> "RebatchedArrowExamplesIterable":
         return RebatchedArrowExamplesIterable(
-            self.ex_iterable.reshard_data_sources(),
-            self.batch_size,
-            self.drop_last_batch,
+            self.ex_iterable.reshard_data_sources(), self.batch_size, self.drop_last_batch, self.force_convert_to_arrow
         )
 
     @property
@@ -634,7 +647,11 @@ class StepExamplesIterable(_BaseExamplesIterable):
         return self.ex_iterable.features
 
     def _init_state_dict(self) -> dict:
-        self._state_dict = self.ex_iterable._init_state_dict()
+        self._state_dict = {
+            "examples_iterable": self.ex_iterable._init_state_dict(),
+            "stepped": 0,
+            "type": self.__class__.__name__,
+        }
         return self._state_dict
 
     def __iter__(self):
@@ -647,9 +664,15 @@ class StepExamplesIterable(_BaseExamplesIterable):
                 break
 
     def _iter_arrow(self):
+        stepped = self._state_dict["stepped"] if self._state_dict else 0
         for key, pa_table in self.ex_iterable.iter_arrow():
-            pa_table = pa_table.take(pa.array(range(self.offset, len(pa_table), self.step), type=pa.int64()))
-            yield key, pa_table
+            stepped_pa_table = pa_table.take(
+                pa.array(range((self.offset - stepped) % self.step, len(pa_table), self.step), type=pa.int64())
+            )
+            stepped = (stepped + len(pa_table)) % self.step
+            if self._state_dict:
+                self._state_dict["stepped"] = stepped
+            yield key, stepped_pa_table
 
     def shuffle_data_sources(self, generator: np.random.Generator) -> "StepExamplesIterable":
         return StepExamplesIterable(
@@ -1249,12 +1272,18 @@ class MappedExamplesIterable(_BaseExamplesIterable):
             # batch_size should match for iter_arrow
             if not isinstance(ex_iterable, RebatchedArrowExamplesIterable):
                 raise ValueError(
-                    f"The {formatting.format_type.capitalize()}-formatted {type(self).__name__} has underlying iterable"
+                    f"The {formatting.format_type.capitalize()}-formatted {type(self).__name__} has underlying iterable "
                     f"that is a {type(ex_iterable).__name__} instead of a RebatchedArrowExamplesIterable."
                 )
+            elif not ex_iterable.iter_arrow:
+                raise ValueError(
+                    f"The {formatting.format_type.capitalize()}-formatted {type(self).__name__} has underlying iterable "
+                    f"that is a {type(ex_iterable).__name__} but doesnt' implement iter_arrow(), a possible fix could be "
+                    "to use RebatchedArrowExamplesIterable(..., force_convert_to_arrow=True)."
+                )
             elif ex_iterable.batch_size != (batch_size if batched else 1):
                 raise ValueError(
-                    f"The {formatting.format_type.capitalize()}-formatted {type(self).__name__} has batch_size={batch_size if batched else 1} which is"
+                    f"The {formatting.format_type.capitalize()}-formatted {type(self).__name__} has batch_size={batch_size if batched else 1} which is "
                     f"different from {ex_iterable.batch_size=} from its underlying iterable."
                 )
         # to enable graceful ends
@@ -2131,22 +2160,6 @@ def _apply_feature_types_on_example(
     return decoded_example
 
 
-def _apply_feature_types_on_batch(
-    batch: dict, features: Features, token_per_repo_id: dict[str, Union[str, bool, None]]
-) -> dict:
-    batch = dict(batch)
-    # add missing columns
-    n_examples = len(batch[next(iter(batch))])
-    for column_name in features:
-        if column_name not in batch:
-            batch[column_name] = [None] * n_examples
-    # we encode the batch for ClassLabel feature types for example
-    encoded_batch = features.encode_batch(batch)
-    # Decode batch for Audio feature, e.g.
-    decoded_batch = features.decode_batch(encoded_batch, token_per_repo_id=token_per_repo_id)
-    return decoded_batch
-
-
 @dataclass
 class FormattingConfig:
     format_type: Optional[str]
@@ -2167,16 +2180,18 @@ class FormattedExamplesIterable(_BaseExamplesIterable):
         formatting: Optional[FormattingConfig],
         features: Optional[Features],
         token_per_repo_id: dict[str, Union[str, bool, None]],
+        force_convert_to_python: bool = False,
     ):
         super().__init__()
         self.ex_iterable = ex_iterable
         self._features = features
         self.formatting = formatting
         self.token_per_repo_id = token_per_repo_id
+        self.force_convert_to_python = force_convert_to_python
 
     @property
     def iter_arrow(self):
-        if self.ex_iterable.iter_arrow and (not self.formatting or self.formatting.is_table):
+        if self.ex_iterable.iter_arrow and not self.force_convert_to_python:
             return self._iter_arrow
 
     @property
@@ -2203,6 +2218,10 @@ class FormattedExamplesIterable(_BaseExamplesIterable):
                 features=self._features if not self.ex_iterable.is_typed else None,
                 token_per_repo_id=self.token_per_repo_id,
             )
+
+        # It's ok to use _iter_arrow here without fancy state_dict logic since it's
+        # used with RebatchedArrowExamplesIterable with the right batch_size to
+        # never lose examples
         if self.ex_iterable.iter_arrow:
             # feature casting (inc column addition) handled within self._iter_arrow()
             for key, pa_table in self._iter_arrow():
@@ -2248,6 +2267,7 @@ class FormattedExamplesIterable(_BaseExamplesIterable):
             features=self.features,
             token_per_repo_id=self.token_per_repo_id,
             formatting=self.formatting,
+            force_convert_to_python=self.force_convert_to_python,
         )
 
     def shard_data_sources(self, num_shards: int, index: int, contiguous=True) -> "FormattedExamplesIterable":
@@ -2257,6 +2277,7 @@ class FormattedExamplesIterable(_BaseExamplesIterable):
             features=self.features,
             token_per_repo_id=self.token_per_repo_id,
             formatting=self.formatting,
+            force_convert_to_python=self.force_convert_to_python,
         )
 
     def reshard_data_sources(self) -> "FormattedExamplesIterable":
@@ -2265,6 +2286,7 @@ class FormattedExamplesIterable(_BaseExamplesIterable):
             features=self.features,
             token_per_repo_id=self.token_per_repo_id,
             formatting=self.formatting,
+            force_convert_to_python=self.force_convert_to_python,
         )
 
     @property
@@ -2565,11 +2587,7 @@ class IterableDataset(DatasetInfoMixin):
 
             if self._formatting and (ex_iterable.iter_arrow or self._formatting.is_table):
                 formatter = get_formatter(self._formatting.format_type, features=self.features)
-                if ex_iterable.iter_arrow:
-                    iterator = ex_iterable.iter_arrow()
-                else:
-                    iterator = _convert_to_arrow(ex_iterable, batch_size=1)
-                for key, pa_table in iterator:
+                for key, pa_table in ex_iterable.iter_arrow():
                     yield formatter.format_row(pa_table)
                 return
             else:
@@ -2599,14 +2617,7 @@ class IterableDataset(DatasetInfoMixin):
         self, batch_size: int = 1, drop_last_batch: bool = False
     ) -> _BaseExamplesIterable:
         ex_iterable = self._ex_iterable
-        if (
-            self._formatting
-            and (ex_iterable.iter_arrow or self._formatting.is_table)
-            or (self.features and ex_iterable.features != self.features)
-        ):
-            ex_iterable = RebatchedArrowExamplesIterable(
-                ex_iterable, batch_size=batch_size, drop_last_batch=drop_last_batch
-            )
+
         if self.epoch:
             ex_iterable = ex_iterable.shuffle_data_sources(np.random.default_rng(self.epoch))
             ex_iterable = shift_ex_examples_rngs(ex_iterable, self.epoch)
@@ -2634,6 +2645,15 @@ class IterableDataset(DatasetInfoMixin):
                     )
                 ex_iterable = StepExamplesIterable(ex_iterable, step=world_size, offset=rank)
 
+        if ex_iterable.iter_arrow:
+            ex_iterable = RebatchedArrowExamplesIterable(
+                ex_iterable, batch_size=batch_size, drop_last_batch=drop_last_batch
+            )
+        elif self._formatting and self._formatting.is_table:
+            ex_iterable = RebatchedArrowExamplesIterable(
+                ex_iterable, batch_size=batch_size, drop_last_batch=drop_last_batch, force_convert_to_arrow=True
+            )
+
         if self._formatting or (self.features and ex_iterable.features != self.features):
             ex_iterable = FormattedExamplesIterable(
                 ex_iterable,
@@ -2663,11 +2683,7 @@ class IterableDataset(DatasetInfoMixin):
         ex_iterable = self._prepare_ex_iterable_for_iteration()
         if self._formatting and (ex_iterable.iter_arrow or self._formatting.is_table):
             formatter = get_formatter(self._formatting.format_type, features=self.features)
-            if ex_iterable.iter_arrow:
-                iterator = ex_iterable.iter_arrow()
-            else:
-                iterator = _convert_to_arrow(ex_iterable, batch_size=1)
-            for key, pa_table in iterator:
+            for key, pa_table in ex_iterable.iter_arrow():
                 yield formatter.format_row(pa_table)
             return
 
@@ -2692,11 +2708,7 @@ class IterableDataset(DatasetInfoMixin):
 
         ex_iterable = self._prepare_ex_iterable_for_iteration(batch_size=batch_size, drop_last_batch=drop_last_batch)
         if self._formatting and (ex_iterable.iter_arrow or self._formatting.is_table):
-            if ex_iterable.iter_arrow:
-                iterator = ex_iterable.iter_arrow()
-            else:
-                iterator = _convert_to_arrow(ex_iterable, batch_size=batch_size, drop_last_batch=drop_last_batch)
-            for key, pa_table in iterator:
+            for key, pa_table in ex_iterable.iter_arrow():
                 yield formatter.format_batch(pa_table)
             return
 
@@ -3410,7 +3422,10 @@ class IterableDataset(DatasetInfoMixin):
                 token_per_repo_id=self._token_per_repo_id,
             )
             ex_iterable = RebatchedArrowExamplesIterable(
-                ex_iterable, batch_size=batch_size if batched else 1, drop_last_batch=drop_last_batch
+                ex_iterable,
+                batch_size=batch_size if batched else 1,
+                drop_last_batch=drop_last_batch,
+                force_convert_to_arrow=True,
             )
         else:
             if self._formatting and self._ex_iterable.iter_arrow:
@@ -3424,6 +3439,7 @@ class IterableDataset(DatasetInfoMixin):
                     formatting=copy.deepcopy(self._formatting),
                     features=input_features,
                     token_per_repo_id=self._token_per_repo_id,
+                    force_convert_to_python=True,
                 )
 
         ex_iterable = MappedExamplesIterable(
@@ -4074,7 +4090,7 @@ class IterableDataset(DatasetInfoMixin):
 
         * audio -> dict of "array" and "sampling_rate" and "path"
         * image -> PIL.Image
-        * video -> torchvision.io.VideoReader
+        * video -> torchcodec.decoders.VideoDecoder
 
         You can enable multithreading using `num_threads`. This is especially useful to speed up remote
         data streaming. However it can be slower than `num_threads=0` for local data on fast disks.

{datasets-4.8.2 → datasets-4.8.4}/src/datasets/packaged_modules/json/json.py

@@ -190,8 +190,13 @@ class Json(datasets.ArrowBasedBuilder):
                                 and self.info.features is None
                                 and self.config.on_mixed_types == "use_json"
                             ):
-                                examples = [ujson_loads(line) for line in batch.splitlines()]
-                                json_field_paths += find_mixed_struct_types_field_paths(examples)
+                                try:
+                                    examples = [ujson_loads(line) for line in batch.splitlines()]
+                                except ValueError:
+                                    # the file is likely not JSON Lines and may contain one single multi-line JSON object
+                                    pass
+                                else:
+                                    json_field_paths += find_mixed_struct_types_field_paths(examples)
                             # Re-encode JSON fields
                             original_batch = batch
                             if json_field_paths:

{datasets-4.8.2 → datasets-4.8.4/src/datasets.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: datasets
-Version: 4.8.2
+Version: 4.8.4
 Summary: HuggingFace community-driven open-source library of datasets
 Home-page: https://github.com/huggingface/datasets
 Download-URL: https://github.com/huggingface/datasets/tags

{datasets-4.8.2 → datasets-4.8.4}/tests/test_iterable_dataset.py

@@ -752,7 +752,9 @@ def test_mapped_examples_iterable_input_columns(n, func, batched, batch_size, in
 )
 def test_mapped_examples_iterable_arrow_format(n, func, batched, batch_size):
     base_ex_iterable = ExamplesIterable(generate_examples_fn, {"n": n})
-    base_ex_iterable = RebatchedArrowExamplesIterable(base_ex_iterable, batch_size=batch_size if batched else 1)
+    base_ex_iterable = RebatchedArrowExamplesIterable(
+        base_ex_iterable, batch_size=batch_size if batched else 1, force_convert_to_arrow=True
+    )
     ex_iterable = MappedExamplesIterable(
         base_ex_iterable,
         func,
@@ -832,7 +834,9 @@ def test_mapped_examples_iterable_arrow_format_from_arrow_examples_iterable(n, f
 )
 def test_mapped_examples_iterable_drop_last_batch_and_arrow_format(n, func, batched, batch_size):
     base_ex_iterable = ExamplesIterable(generate_examples_fn, {"n": n})
-    base_ex_iterable = RebatchedArrowExamplesIterable(base_ex_iterable, batch_size=batch_size if batched else 1)
+    base_ex_iterable = RebatchedArrowExamplesIterable(
+        base_ex_iterable, batch_size=batch_size if batched else 1, force_convert_to_arrow=True
+    )
     ex_iterable = MappedExamplesIterable(
         base_ex_iterable,
         func,
@@ -895,7 +899,9 @@ def test_mapped_examples_iterable_drop_last_batch_and_arrow_format(n, func, batc
 )
 def test_mapped_examples_iterable_with_indices_and_arrow_format(n, func, batched, batch_size):
     base_ex_iterable = ExamplesIterable(generate_examples_fn, {"n": n})
-    base_ex_iterable = RebatchedArrowExamplesIterable(base_ex_iterable, batch_size=batch_size if batched else 1)
+    base_ex_iterable = RebatchedArrowExamplesIterable(
+        base_ex_iterable, batch_size=batch_size if batched else 1, force_convert_to_arrow=True
+    )
     ex_iterable = MappedExamplesIterable(
         base_ex_iterable,
         func,
@@ -946,7 +952,9 @@ def test_mapped_examples_iterable_with_indices_and_arrow_format(n, func, batched
 )
 def test_mapped_examples_iterable_remove_columns_arrow_format(n, func, batched, batch_size, remove_columns):
     base_ex_iterable = ExamplesIterable(generate_examples_fn, {"n": n, "extra_column": "foo"})
-    base_ex_iterable = RebatchedArrowExamplesIterable(base_ex_iterable, batch_size=batch_size if batched else 1)
+    base_ex_iterable = RebatchedArrowExamplesIterable(
+        base_ex_iterable, batch_size=batch_size if batched else 1, force_convert_to_arrow=True
+    )
     ex_iterable = MappedExamplesIterable(
         base_ex_iterable,
         func,
@@ -991,7 +999,9 @@ def test_mapped_examples_iterable_remove_columns_arrow_format(n, func, batched,
 )
 def test_mapped_examples_iterable_fn_kwargs_and_arrow_format(n, func, batched, batch_size, fn_kwargs):
     base_ex_iterable = ExamplesIterable(generate_examples_fn, {"n": n})
-    base_ex_iterable = RebatchedArrowExamplesIterable(base_ex_iterable, batch_size=batch_size if batched else 1)
+    base_ex_iterable = RebatchedArrowExamplesIterable(
+        base_ex_iterable, batch_size=batch_size if batched else 1, force_convert_to_arrow=True
+    )
     ex_iterable = MappedExamplesIterable(
         base_ex_iterable,
         func,
@@ -1031,7 +1041,9 @@ def test_mapped_examples_iterable_fn_kwargs_and_arrow_format(n, func, batched, b
 )
 def test_mapped_examples_iterable_input_columns_and_arrow_format(n, func, batched, batch_size, input_columns):
     base_ex_iterable = ExamplesIterable(generate_examples_fn, {"n": n})
-    base_ex_iterable = RebatchedArrowExamplesIterable(base_ex_iterable, batch_size=batch_size if batched else 1)
+    base_ex_iterable = RebatchedArrowExamplesIterable(
+        base_ex_iterable, batch_size=batch_size if batched else 1, force_convert_to_arrow=True
+    )
     ex_iterable = MappedExamplesIterable(
         base_ex_iterable,
         func,
@@ -1232,6 +1244,48 @@ def test_take_examples_iterable():
     assert_load_state_dict_resumes_iteration(take_ex_iterable)
 
 
+def test_step_examples_iterable():
+    total, step, offset = 10, 2, 1
+    base_ex_iterable = ExamplesIterable(generate_examples_fn, {"n": total})
+    step_ex_iterable = StepExamplesIterable(base_ex_iterable, step=step, offset=offset)
+    expected = list(generate_examples_fn(n=total))[offset::step]
+    assert list(step_ex_iterable) == expected
+    assert_load_state_dict_resumes_iteration(step_ex_iterable)
+
+
+def test_skip_arrow_examples_iterable():
+    total, count = 10, 2
+    base_ex_iterable = ArrowExamplesIterable(generate_tables_fn, {"n": total})
+    skip_ex_iterable = SkipExamplesIterable(base_ex_iterable, n=count)
+    expected = [x for _, pa_table in generate_tables_fn(n=total) for x in pa_table.to_pylist()][count:]
+    assert [example for _, example in skip_ex_iterable] == expected
+    assert skip_ex_iterable.shuffle_data_sources(np.random.default_rng(42)) is skip_ex_iterable, (
+        "skip examples makes the shards order fixed"
+    )
+    assert_load_state_dict_resumes_iteration(skip_ex_iterable)
+
+
+def test_take_arrow_examples_iterable():
+    total, count = 10, 2
+    base_ex_iterable = ArrowExamplesIterable(generate_tables_fn, {"n": total})
+    take_ex_iterable = TakeExamplesIterable(base_ex_iterable, n=count)
+    expected = [x for _, pa_table in generate_tables_fn(n=total) for x in pa_table.to_pylist()][:count]
+    assert [example for _, example in take_ex_iterable] == expected
+    assert take_ex_iterable.shuffle_data_sources(np.random.default_rng(42)) is take_ex_iterable, (
+        "skip examples makes the shards order fixed"
+    )
+    assert_load_state_dict_resumes_iteration(take_ex_iterable)
+
+
+def test_step_arrow_examples_iterable():
+    total, step, offset = 10, 2, 1
+    base_ex_iterable = ArrowExamplesIterable(generate_tables_fn, {"n": total})
+    step_ex_iterable = StepExamplesIterable(base_ex_iterable, step=step, offset=offset)
+    expected = [x for _, pa_table in generate_tables_fn(n=total) for x in pa_table.to_pylist()][offset::step]
+    assert [example for _, example in step_ex_iterable] == expected
+    assert_load_state_dict_resumes_iteration(step_ex_iterable)
+
+
 @pytest.mark.parametrize(
     "n, num_times",
     [
@@ -1345,7 +1399,9 @@ def test_no_iter_arrow(ex_iterable: _BaseExamplesIterable):
         HorizontallyConcatenatedMultiSourcesExamplesIterable([ArrowExamplesIterable(generate_tables_fn, {})]),
         # RandomlyCyclingMultiSourcesExamplesIterable([ArrowExamplesIterable(generate_tables_fn, {})], np.random.default_rng(42)),  # not implemented
         MappedExamplesIterable(
-            RebatchedArrowExamplesIterable(ExamplesIterable(generate_examples_fn, {}), batch_size=1),
+            RebatchedArrowExamplesIterable(
+                ExamplesIterable(generate_examples_fn, {}), batch_size=1, force_convert_to_arrow=True
+            ),
             lambda t: t,
             formatting=FormattingConfig(format_type="arrow"),
         ),
@@ -1355,12 +1411,16 @@ def test_no_iter_arrow(ex_iterable: _BaseExamplesIterable):
             formatting=FormattingConfig(format_type="arrow"),
         ),
         FilteredExamplesIterable(
-            RebatchedArrowExamplesIterable(ExamplesIterable(generate_examples_fn, {}), batch_size=1),
+            RebatchedArrowExamplesIterable(
+                ExamplesIterable(generate_examples_fn, {}), batch_size=1, force_convert_to_arrow=True
+            ),
             lambda t: True,
             formatting=FormattingConfig(format_type="arrow"),
         ),
         FilteredExamplesIterable(
-            RebatchedArrowExamplesIterable(ArrowExamplesIterable(generate_tables_fn, {}), batch_size=1),
+            RebatchedArrowExamplesIterable(
+                ArrowExamplesIterable(generate_tables_fn, {}), batch_size=1, force_convert_to_arrow=True
+            ),
             lambda t: True,
             formatting=FormattingConfig(format_type="arrow"),
         ),