rebatch arrow iterables before formatted iterable

Author: lhoestq

src/datasets/iterable_dataset.py

@@ -2194,7 +2194,9 @@ def _prepare_ex_iterable_for_iteration(
         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):
+        if (self._formatting or (self.features and ex_iterable.features != self.features)) and (
+            ex_iterable.iter_arrow or self._formatting.is_table
+        ):
             ex_iterable = RebatchedArrowExamplesIterable(
                 ex_iterable, batch_size=batch_size, drop_last_batch=drop_last_batch
             )