`IterableDatasetDict.map()` call removes `column_names` (in fact info.features)

[mombip]

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:

new_dict = my_iterable_dataset_dict.map(
    function=my_fn,
    with_indices=False,
    batched=True,
    batch_size=16,
)

4. Observe

new_dict["train"].info.features  # {'text': Value(dtype='string', id=None)}
new_dict["train"].column_names   # ['text']

5. Call:

new_dict = my_iterable_dataset_dict.map(
    function=my_fn,
    with_indices=False,
    batched=True,
    batch_size=16,
    remove_columns=["foo"]
)

6. Observe:

new_dict["train"].info.features  # → None
new_dict["train"].column_names   # → None

5. Internally, in dataset_dict.py this loop omits features (code):

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) correct info.features is replaced by features which is None:

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.

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

[lhoestq]

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.

[mombip]

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:

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

[lhoestq]

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

[mombip]

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)?

[lhoestq]

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())

[ArjunJagdale]

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!