huggingface / huggingface/datasets
`IterableDataset` lost but not keep columns when map function adding columns with names in `remove_columns`
- Dominant language
- Python
- Stars
- 22k
- Forks
- 3.4k
- Avg merge
- 5d 7h
- Merged PRs (30d)
- 17
Description
### Describe the bug
IterableDataset lost but not keep columns when map function adding columns with names in remove_columns,
Dataset not.
May be related to the code below:
https://github.com/huggingface/datasets/blob/06c3ffb8d068b6307b247164b10f7c7311cefed4/src/datasets/iterable_dataset.py#L750-L756
### Steps to reproduce the bug
```python
dataset: IterableDataset = load_dataset("Anthropic/hh-rlhf", streaming=True, split="train")
column_names = list(next(iter(dataset)).keys()) # ['chosen', 'rejected']
# map_fn will return dict {"chosen": xxx, "rejected": xxx, "prompt": xxx, "history": xxxx}
dataset = dataset.map(map_fn, batched=True, remove_columns=column_names)
next(iter(dataset))
# output
# {'prompt': 'xxx, 'history': xxx}
```
```python
# when load_dataset with streaming=False, the column_names are kept:
dataset: Dataset = load_dataset("Anthropic/hh-rlhf", streaming=False, split="train")
column_names = list(next(iter(dataset)).keys()) # ['chosen', 'rejected']
# map_fn will return dict {"chosen": xxx, "rejected": xxx, "prompt": xxx, "history": xxxx}
dataset = dataset.map(map_fn, batched=True, remove_columns=column_names)
next(iter(dataset))
# output
# {'prompt': 'xxx, 'history': xxx, "chosen": xxx, "rejected": xxx}
```
### Expected behavior
IterableDataset keep columns when map function adding columns with names in remove_columns
### Environment info
datasets==2.14.6
Contributor guide
Assessment
This issue has not been assessed yet.