`IterableDatasetDict.map()` call removes `column_names` (in fact info.features)
Opened by mombip (open)

When calling `IterableDatasetDict.map()`, each split’s `IterableDataset.map()` is invoked without a `features` argument. While omitting the argument isn’t itself incorrect, the implementation then sets `info.features = features`, which destroys the original `features` content. Since `IterableDataset.column_names` relies on `info.features`, it ends up broken (`None`).

**Reproduction**

1.  Define an IterableDatasetDict with a non-None features schema.
2.  my_iterable_dataset_dict contains "text" column.
3. Call:
```Python
new_dict = my_iterable_dataset_dict.map(
    function=my_fn,
    with_indices=False,
    batched=True,
    batch_size=16,
)
```
4. Observe
```Python
new_dict["train"].info.features  # {'text': Value(dtype='string', id=None)}
new_dict["train"].column_names   # ['text']
```
5. Call:
```Python
new_dict = my_iterable_dataset_dict.map(
    function=my_fn,
    with_indices=False,
    batched=True,
    batch_size=16,
    remove_columns=["foo"]
)
```
6. Observe:
```Python
new_dict["train"].info.features  # → None
new_dict["train"].column_names   # → None
```
5. Internally, in dataset_dict.py this loop omits features ([code](https://github.com/huggingface/datasets/blob/b9efdc64c3bfb8f21f8a4a22b21bddd31ecd5a31/src/datasets/dataset_dict.py#L2047C5-L2056C14)):
```Python
for split, dataset in self.items():
    dataset_dict[split] = dataset.map(
        function=function,
        with_indices=with_indices,
        input_columns=input_columns,
        batched=batched,
        batch_size=batch_size,
        drop_last_batch=drop_last_batch,
        remove_columns=remove_columns,
        fn_kwargs=fn_kwargs,
        # features omitted → defaults to None
    )
```
7. Then inside IterableDataset.map() ([code](https://github.com/huggingface/datasets/blob/b9efdc64c3bfb8f21f8a4a22b21bddd31ecd5a31/src/datasets/iterable_dataset.py#L2619C1-L2622C37)) correct `info.features` is replaced by features which is None:
```Python
info = self.info.copy()
info.features = features  # features is None here
return IterableDataset(..., info=info, ...)
```

**Suggestion**
It looks like this replacement was added intentionally but maybe should be done only if `features` is `not None`.

**Workarround:**
`SFTTrainer` calls `dataset.map()` several times and then fails on `NoneType` when iterating `dataset.column_names`.
I decided to write this patch - works form me.

```python
def patch_iterable_dataset_map():
    _orig_map = IterableDataset.map

    def _patched_map(self, *args, **kwargs):
        if "features" not in kwargs or kwargs["features"] is None:
            kwargs["features"] = self.info.features
        return _orig_map(self, *args, **kwargs)

    IterableDataset.map = _patched_map
```






Comments:
Hi ! IterableDataset doesn't know what's the output of the function you pass to map(), so it's not possible to know in advance the features of the output dataset.

There is a workaround though: either do `ds = ds.map(..., features=features)`, or you can do `ds = ds._resolve_features()` which iterates on the first rows to infer the dataset features.

Thank you. I understand that “IterableDataset doesn't know what's the output of the function”—that’s true, but:

Unfortunately, the workaround you proposed **doesn’t solve** the problem. `ds.map()` is called multiple times by third-party code (i.e. `SFTTrainer`). To apply your approach, I would have to modify external library code. That’s why I decided to patch the _class_ rather than update `dataset` _objects_ (in fact, updating the object after `map()` was my initial approach, but then I realized I’m not the only one mapping an already-mapped dataset.)

As a user, I expected that after mapping I would get a new dataset with the correct column names. If, for some reason, that can’t be the default behavior, I would expect an argument—i.e. `auto_resolve_features: bool = False` — to control how my dataset is mapped if following mapping operation are called.

It’s also problematic that `column_names` are tied to `features`, which is even more confusing and forces you to inspect the source code to understand what’s going on.

**New version of workaround:**
```python
def patch_iterable_dataset_map():
    _orig_map = IterableDataset.map

    def _patched_map(self, *args, **kwargs):
        ds = _orig_map(self, *args, **kwargs)
        return ds._resolve_features()

    IterableDataset.map = _patched_map
```

I see, maybe `.resolve_features()` should be called by default in this case in the SFTTrainer ? (or pass `features=` if the data processing always output the same features)

We can even support a new parameter `features="infer"` if it would be comfortable to not use internal methods in SFTTrainer

I think most straightforward solution would be to reinitialize `features` from data after mapping if `feature` argument is not passed. I hink it is more intuitive behavior than just cleaning features. There is also problem in usage `.resolve_features()` in this context. I observed that it leads to `_head()` method execution and it then causes that 5 batches from dataset are iterated (`_head()` defaults to 5 batches). 
I'm not sure how it influences whole process. Are those 5 batches (in my case it's 5000 rows) used only to find `features`. Does final training/eval process "see" this items? How it affects IterableDataset state (current position)?

I checked the source code and while it indeed iterates on the first 5 rows. As a normal iteration, it does record the state in case you call `.state_dict()`, but it doesn't change the starting state. The starting state is always the beginning of the dataset, unless it is explicitly set with `.load_state_dict()`. To be clear, if you iterate on the dataset after `._resolve_features()`, it will start from the beginning of the dataset (or from a state you manually pass using `.load_state_dict()`)

Hi!
I’ve opened a PR #7658 to address this issue.

The fix ensures that info.features is only updated if features is not None, preventing accidental loss of schema and column_names.
Please let me know if you see any edge cases or have additional concerns!
Also, if a test is needed for this case, happy to discuss—the fix is small, but I can add one if the maintainers prefer.

Thanks everyone for the clear diagnosis and suggestions in this thread!