Iterable Dataset: rename column clashes with remove column
Opened by sanchit-gandhi (closed)

### Describe the bug

Suppose I have a two iterable datasets, one with the features: 
* `{"audio", "text", "column_a"}`

And the other with the features:
* `{"audio", "sentence", "column_b"}`

I want to combine both datasets using `interleave_datasets`, which requires me to unify the column names. I would typically do this by:
1. Renaming the common columns to the same name (e.g. `"text"` -> `"sentence"`)
2. Removing the unwanted columns (e.g. `"column_a"`, `"column_b"`)

However, the process of renaming and removing columns in an iterable dataset doesn't work, since we need to preserve the original text column, meaning we can't combine the datasets.

### Steps to reproduce the bug

```python
from datasets import load_dataset

# load LS in streaming mode
dataset = load_dataset("librispeech_asr", "clean", split="validation", streaming=True)

# check original features
dataset_features = dataset.features.keys()
print("Original features: ", dataset_features)

# rename "text" -> "sentence"
dataset = dataset.rename_column("text", "sentence")

# remove unwanted columns
COLUMNS_TO_KEEP = {"audio", "sentence"}
dataset = dataset.remove_columns(set(dataset_features - COLUMNS_TO_KEEP))

# stream first sample, should return "audio" and "sentence" columns
print(next(iter(dataset)))
```
Traceback:
```python
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
Cell In[5], line 17
     14 COLUMNS_TO_KEEP = {"audio", "sentence"}
     15 dataset = dataset.remove_columns(set(dataset_features - COLUMNS_TO_KEEP))
---> 17 print(next(iter(dataset)))

File ~/datasets/src/datasets/iterable_dataset.py:1353, in IterableDataset.__iter__(self)
   1350         yield formatter.format_row(pa_table)
   1351     return
-> 1353 for key, example in ex_iterable:
   1354     if self.features:
   1355         # `IterableDataset` automatically fills missing columns with None.
   1356         # This is done with `_apply_feature_types_on_example`.
   1357         example = _apply_feature_types_on_example(
   1358             example, self.features, token_per_repo_id=self._token_per_repo_id
   1359         )

File ~/datasets/src/datasets/iterable_dataset.py:652, in MappedExamplesIterable.__iter__(self)
    650     yield from ArrowExamplesIterable(self._iter_arrow, {})
    651 else:
--> 652     yield from self._iter()

File ~/datasets/src/datasets/iterable_dataset.py:729, in MappedExamplesIterable._iter(self)
    727 if self.remove_columns:
    728     for c in self.remove_columns:
--> 729         del transformed_example[c]
    730 yield key, transformed_example
    731 current_idx += 1

KeyError: 'text'
```
=> we see that `datasets` is looking for the column "text", even though we've renamed this to "sentence" and then removed the un-wanted "text" column from our dataset.

### Expected behavior

Should be able to rename and remove columns from iterable dataset.

### Environment info

- `datasets` version: 2.15.1.dev0
- Platform: macOS-13.5.1-arm64-arm-64bit
- Python version: 3.11.6
- `huggingface_hub` version: 0.19.4
- PyArrow version: 14.0.1
- Pandas version: 2.1.2
- `fsspec` version: 2023.9.2

Comments:
Column "text" doesn't exist anymore so you can't remove it

You can get the expected result by fixing typos in the snippet :)
```python
from datasets import load_dataset

# load LS in streaming mode
dataset = load_dataset("librispeech_asr", "clean", split="validation", streaming=True)

# check original features
dataset_features = dataset.features.keys()
print("Original features: ", dataset_features)

# rename "text" -> "sentence"
dataset = dataset.rename_column("text", "sentence")

# remove unwanted columns
COLUMNS_TO_KEEP = {"audio", "sentence"}
dataset = dataset.remove_columns(set(dataset.features) - COLUMNS_TO_KEEP)

# stream first sample, should return "audio" and "sentence" columns
print(next(iter(dataset)))
```

Fixed code:

```python
from datasets import load_dataset

# load LS in streaming mode
dataset = load_dataset("librispeech_asr", "clean", split="validation", streaming=True)

# check original features
dataset_features = dataset.features.keys()
print("Original features: ", dataset_features)

# rename "text" -> "sentence"
dataset = dataset.rename_column("text", "sentence")
dataset_features = dataset.features.keys()

# remove unwanted columns
COLUMNS_TO_KEEP = {"audio", "sentence"}
dataset = dataset.remove_columns(set(dataset_features - COLUMNS_TO_KEEP))

# stream first sample, should return "audio" and "sentence" columns
print(next(iter(dataset)))
```

Whoops 😅 Thanks for the swift reply both! Works like a charm!