[IterableDatasets] Add column followed by interleave datasets gives bogus outputs
Opened by sanchit-gandhi (closed)

### Describe the bug

If we add a new column to our iterable dataset using the hack described in #5752, when we then interleave datasets the new column is pinned to one value.

### Steps to reproduce the bug

What we're going to do here is:
1. Load an iterable dataset in streaming mode (`original_dataset`)
2. Add a new column to this dataset using the hack in #5752 (`modified_dataset_1`)
3. Create another new dataset by adding a column with the same key but different values (`modified_dataset_2`)
4. Interleave our new datasets (`modified_dataset_1` + `modified_dataset_2`)
5. Check the value of our newly added column (`new_column`)

```python
from datasets import load_dataset

# load an iterable dataset
original_dataset = load_dataset("librispeech_asr", "clean", split="validation", streaming=True)

# now add a new column to our streaming dataset using our hack from 5752
name = "new_column"
column = [f"new dataset 1, row {i}" for i in range(50)]

new_features = original_dataset.features.copy()
new_features[name] = new_features["file"]  # I know that "file" has the right column type to match our new feature

def add_column_fn(example, idx):
    if name in example:
        raise ValueError(f"Error when adding {name}: column {name} is already in the dataset.")
    return {name: column[idx]}

modified_dataset_1 = original_dataset.map(add_column_fn, with_indices=True, features=new_features)

# now create a second modified dataset using the same trick
column = [f"new dataset 2, row {i}" for i in range(50)]

def add_column_fn(example, idx):
    if name in example:
        raise ValueError(f"Error when adding {name}: column {name} is already in the dataset.")
    return {name: column[idx]}

modified_dataset_2 = original_dataset.map(add_column_fn, with_indices=True, features=new_features)

# interleave these datasets
interleaved_dataset = interleave_datasets([modified_dataset_1, modified_dataset_2])

# now check what the value of the added column is
for i, sample in enumerate(interleaved_dataset):
    print(sample["new_column"])
    if i == 10:
        break
```
**Print Output:**
```
new dataset 2, row 0
new dataset 2, row 0
new dataset 2, row 1
new dataset 2, row 1
new dataset 2, row 2
new dataset 2, row 2
new dataset 2, row 3
new dataset 2, row 3
new dataset 2, row 4
new dataset 2, row 4
new dataset 2, row 5
```
We see that we only get outputs from our second dataset.

### Expected behavior

We should interleave between dataset 1 and 2 and increase in row value:
```
new dataset 1, row 0
new dataset 2, row 0
new dataset 1, row 1
new dataset 2, row 1
new dataset 1, row 2
new dataset 2, row 2
...
```

### Environment info

- datasets version: 2.10.2.dev0
- Platform: Linux-4.19.0-23-cloud-amd64-x86_64-with-glibc2.28
- Python version: 3.9.16
- Huggingface_hub version: 0.13.3
- PyArrow version: 10.0.1
- Pandas version: 1.5.2

Comments:
Problem with the code snippet! Using global vars and functions was not a good idea with iterable datasets!

If we update to:
```python
from datasets import load_dataset

original_dataset = load_dataset("librispeech_asr", "clean", split="validation", streaming=True)

# now add a new column to our streaming dataset using our hack
name = "new_column"
column_1 = [f"new dataset 1, row {i}" for i in range(50)]

new_features = original_dataset.features.copy()
new_features[name] = new_features["file"]  # I know that "file" has the right column type to match our new feature

def add_column_fn_1(example, idx):
    if name in example:
        raise ValueError(f"Error when adding {name}: column {name} is already in the dataset.")
    return {name: column_1[idx]}

modified_dataset_1 = original_dataset.map(add_column_fn_1, with_indices=True, features=new_features)

# now create a second modified dataset using the same trick
column_2 = [f"new dataset 2, row {i}" for i in range(50)]

def add_column_fn_2(example, idx):
    if name in example:
        raise ValueError(f"Error when adding {name}: column {name} is already in the dataset.")
    return {name: column_2[idx]}

modified_dataset_2 = original_dataset.map(add_column_fn_2, with_indices=True, features=new_features)

interleaved_dataset = interleave_datasets([modified_dataset_1, modified_dataset_2])

for i, sample in enumerate(interleaved_dataset):
    print(sample["new_column"])
    if i == 10:
        break
```
we get the correct outputs:
```python
new dataset 1, row 0
new dataset 2, row 0
new dataset 1, row 1
new dataset 2, row 1
new dataset 1, row 2
new dataset 2, row 2
new dataset 1, row 3
new dataset 2, row 3
new dataset 1, row 4
new dataset 2, row 4
new dataset 1, row 5
```


Thanks @sanchit-gandhi, this solo performance is very helpful! :)