huggingface / huggingface/datasets

`IterableDataset` lost but not keep columns when map function adding columns with names in `remove_columns`

Open
#6,408 0 comments 0 reactions 0 assignees View on GitHub
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.