'with_format' is extremely slow when used together with 'interleave_datasets' or 'shuffle' on IterableDatasets
Opened by tobycrisford (closed)

### Describe the bug

If you:

1. Interleave two iterable datasets together with the interleave_datasets function, or shuffle an iterable dataset
2. Set the output format to torch tensors with .with_format('torch')

Then iterating through the dataset becomes over 100x slower than it is if you don't apply the torch formatting.

### Steps to reproduce the bug

```python
import datasets
import torch
from tqdm import tqdm

rand_a = torch.randn(3,224,224)
rand_b = torch.randn(3,224,224)

a = torch.stack([rand_a] * 1000)
b = torch.stack([rand_b] * 1000)

features = datasets.Features({"tensor": datasets.Array3D(shape=(3,224,224), dtype="float32")})
ds_a = datasets.Dataset.from_dict({"tensor": a}, features=features).to_iterable_dataset()
ds_b = datasets.Dataset.from_dict({"tensor": b}, features=features).to_iterable_dataset()

# Iterating through either dataset with torch formatting is really fast (2000it/s on my machine)
for example in tqdm(ds_a.with_format('torch')):
    pass

# Iterating through either dataset shuffled is also pretty fast (100it/s on my machine)
for example in tqdm(ds_a.shuffle()):
    pass

# Iterating through this interleaved dataset is pretty fast (200it/s on my machine)
ds_fast = datasets.interleave_datasets([ds_a, ds_b])
for example in tqdm(ds_fast):
    pass

# Iterating through either dataset with torch formatting *after shuffling* is really slow... (<2it/s on my machine)
for example in tqdm(ds_a.shuffle().with_format('torch')):
    pass

# Iterating through this torch formatted interleaved dataset is also really slow (<2it/s on my machine)...
ds_slow = datasets.interleave_datasets([ds_a, ds_b]).with_format('torch')
for example in tqdm(ds_slow):
    pass

# Even doing this is way faster!! (70it/s on my machine)
for example in tqdm(ds_fast):
    test = torch.tensor(example['tensor'])

```

### Expected behavior

Applying torch formatting to the interleaved dataset shouldn't increase the time taken to iterate through the dataset by very much, since even explicitly converting every example is over 70x faster than calling .with_format('torch').

### Environment info

- `datasets` version: 2.16.1
- Platform: Linux-6.5.0-15-generic-x86_64-with-glibc2.38
- Python version: 3.11.6
- `huggingface_hub` version: 0.20.3
- PyArrow version: 15.0.0
- Pandas version: 2.2.0
- `fsspec` version: 2023.10.0


Comments:
The "torch" formatting is usually fast because we do zero-copy conversion from the Arrow data on your disk to Torch tensors. However IterableDataset shuffling seems to do data copies that slow down the pipeline, and it shuffles python objects instead of Arrow data.

To fix this we need to implement `BufferShuffledExamplesIterable.iter_arrow()` (same as regular `BufferShuffledExamplesIterable.__iter__()` but yields Arrow tables)

https://github.com/huggingface/datasets/blob/b7d854b7fd3e9a330e21b76ee8421d4a7ebb4a7a/src/datasets/iterable_dataset.py#L968-L974


@lhoestq Any update on this?


This hasn't been implemented yet unfortunately, contributions are welcome though

@lhoestq So I see the following.

``` 
ds_interleave = datasets.interleave_datasets(
    [ds_a, ds_b],
    probabilities=[0.999, 0.001],
    stopping_strategy="first_exhausted",
)
ds_fast = ds_fast.with_format("torch")
``` 
With above code the ex_iterable becomes `RandomlyCyclingMultiSourcesExamplesIterable` which also doesn't have a `_iter_arrow` yet. I believe because of the shuffling under the hood the individual iterators will be of type `BufferShuffledExamplesIterable`.
So to get `interleave_datasets` to work we need to implement both.

I'm really not familiar with the code (first time looking at it) but I came up with this two methods.



```
class RandomlyCyclingMultiSourcesExamplesIterable
...
@property
    def iter_arrow(self):
        return self._iter_arrow

    def _iter_arrow(self, max_chunksize: Optional[int] = None):
        # we use this to buffer one example of each iterator to know if an iterator is exhausted
        nexts = [None] * len(self.ex_iterables)
        # because of that, we need to rewind 1 example when reloading the state dict
        if self._state_dict:
            for i in range(len(self.ex_iterables)):
                if self._state_dict["previous_states"][i] is not None:
                    self.ex_iterables[i].load_state_dict(
                        self._state_dict["previous_states"][i]
                    )
        iterators = [ex_iterable._iter_arrow() for ex_iterable in self.ex_iterables]

        indices_iterator = self._get_indices_iterator()

        is_exhausted = (
            np.array(self._state_dict["is_exhausted"])
            if self._state_dict
            else np.full(len(self.ex_iterables), False)
        )
        for i in indices_iterator:
            # if the stopping criteria is met, break the main for loop
            if self.bool_strategy_func(is_exhausted):
                break
            # let's pick one example from the iterator at index i
            if nexts[i] is None:
                nexts[i] = next(iterators[i], False)
            result = nexts[i]
            if self._state_dict:
                self._state_dict["previous_states"][i] = deepcopy(
                    self._state_dict["ex_iterables"][i]
                )
            nexts[i] = next(iterators[i], False)

            # the iterator is exhausted
            if nexts[i] is False:
                is_exhausted[i] = True
                if self._state_dict:
                    self._state_dict["is_exhausted"][i] = True
                # we reset it in case the stopping crtieria isn't met yet
                nexts[i] = None
                if self._state_dict:
                    self._state_dict["ex_iterables"][i] = self.ex_iterables[
                        i
                    ]._init_state_dict()
                    self._state_dict["previous_states"][i] = None
                iterators[i] = iter(self.ex_iterables[i])

            if result is not False:
                yield result
```

and the following for the other one.

```
class BufferShuffledExamplesIterable
...
@property
    def iter_arrow(self):
        return self._iter_arrow

def _iter_arrow(self, max_chunksize: Optional[int] = None):
        buffer_size = self.buffer_size
        rng = deepcopy(self.generator)
        indices_iterator = self._iter_random_indices(rng, buffer_size)
        # this is the shuffle buffer that we keep in memory
        mem_buffer = []
        for key, pa_table in self.ex_iterable._iter_arrow():
            if (
                len(mem_buffer) == buffer_size
            ):  # if the buffer is full, pick and example from it
                i = next(indices_iterator)
                yield mem_buffer[i]
                mem_buffer[i] = key, pa_table  # replace the picked example by a new one
            else:  # otherwise, keep filling the buffer
                mem_buffer.append((key, pa_table))
        # when we run out of examples, we shuffle the remaining examples in the buffer and yield them
        rng.shuffle(mem_buffer)
        yield from mem_buffer
```

Can you provide some feedback before I open a PR?
There seems to be a lot of duplication, there's room for improvement, but wanted to check correctness first.
Thanks.

It looks like a good start ! Btw you can move the function from RandomlyCyclingMultiSourcesExamplesIterable to CyclingMultiSourcesExamplesIterable

I believe CyclingMultiSourcesExamplesIterable can iterate on arrow data only if all the sources can iterate on arrow data:

```diff
@property
    def iter_arrow(self):
-         return self._iter_arrow
+         return self._iter_arrow if all(ex_iterable.iter_arrow for ex_iterable in self.ex_iterables) else None
```

same for BufferShuffledExamplesIterable:

```diff
@property
    def iter_arrow(self):
-         return self._iter_arrow
+         return self._iter_arrow if self.ex_iterable.iter_arrow else None
```

I also think in both cases we should yield row by row.

so _interleave_iterable_datasets() should set the batch_size to 1 like this:

```diff
ex_iterables = [copy.deepcopy(d._ex_iterable) for d in datasets]
+ if all(ex_iterable.iter_arrow for ex_iterable in ex_iterables):
+     ex_iterables = [RebatchedArrowExamplesIterable(ex_iterable, batch_size=1) for ex_iterable in ex_iterables]
```

and same for shuffle():
```diff
ex_iterable=BufferShuffledExamplesIterable(
-     self._ex_iterable, buffer_size=buffer_size, generator=generator
+     RebatchedArrowExamplesIterable(self._ex_iterable, batch_size=1) if self._ex_iterable.iter_arrow else self._ex_iterable,
+     buffer_size=buffer_size,
+     generator=generator
)
```

also feel free to use `.iter_arrow()` instead or `._iter_arrow()`.

Feel free to ping me for reviewing a PR :)

@lhoestq I assume we can close this. Thanks!


Yep ! :)