Fix methods using IterableDataset.map that lead to features=None

[alvarobartt]

As currently IterableDataset.map is setting the info.features to None every time as we don't know the output of the dataset in advance, IterableDataset methods such as rename_column, rename_columns, and remove_columns. that internally use map lead to the features being None.

This PR is related to #3888, #5245, and #5284

✅ Current solution

The code in this PR is basically making sure that if the features were there since the beginning and a rename_column/rename_columns happens, those are kept and the rename is applied to the Features too. Also, if the features were not there before applying rename_column, rename_columns or remove_columns, a batch is prefetched and the features are being inferred (that could potentially be part of IterableDataset.__init__ in case the info.features value is None).

💡 Ideas

Some ideas were proposed in #3888, but probably the most consistent solution even though it may take some time is to actually do the type inferencing during the IterableDataset.__init__ in case the provided info.features is None, otherwise, we can just use the provided features.

Additionally, as mentioned at #3888, we could also include a features parameter to the map function, but that's probably more tedious.

Also thanks to @lhoestq for sharing some ideas in both #3888 and #5245 🤗

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[alvarobartt]

Maybe other options are:

• Keep the info.features to None if those were initially None
• Infer the features with pre-fetching just if the info.features is None
• If the info.features are there, make sure that after map features is not None

[alvarobartt]

Hi @lhoestq something that's still not clear to me is: should we infer the features always when applying a map if those are initially None, or just assume that if the features are initially None those should be left that way unless the user specifically sets those (or during iter)?

In this PR I'm using from datasets.iterable_dataset import _infer_features_from_batch to infer the features when those are None using pre-fetch of self._head(), but I'm not sure if that's the expected behavior.

Thanks in advance for your help!

[alvarobartt]

Also, the PR still has some more work to do, but probably the most relevant thing to fix right now is that the features are being set to None in the functions IterableDataset.rename_column, IterableDataset.rename_columns, and IterableDataset.remove_columns when the features originally had a value. So once that's fixed maybe we can focus on improving the current map's behavior, so as to avoid this from happening also when the user uses map directly and not through the functions mentioned above.

[lhoestq]

Cool thank you ! Resolving the features can be expensive sometimes, so maybe we don't resolve the features and we can just rename/remove columns if the features are known (i.e. if they're not None). What do you think ?

[alvarobartt]

Cool thank you ! Resolving the features can be expensive sometimes, so maybe we don't resolve the features and we can just rename/remove columns if the features are known (i.e. if they're not None). What do you think ?

Thanks for the feedback! Makes sense to me 👍🏻 I'll commit the comments now!

[alvarobartt]

Already done @lhoestq, feel free to merge whenever you want! Also before merging, can you please link the following issues #3888, #5245, and #5284, so that those are closed upon merge? Thanks!

[lhoestq]

Awesome thank you !

I'll keep #3888 open since it's not fully completed