Add features parameter to IterableDatasetDict.map

IterableDataset.map accepts a features parameter to declare the
output schema, but IterableDatasetDict.map did not expose it.
This meant users of IterableDatasetDict had no way to preserve
feature metadata through map operations.

Author: HaukurPall

src/datasets/dataset_dict.py

@@ -2105,6 +2105,7 @@ def map(
         batch_size: int = 1000,
         drop_last_batch: bool = False,
         remove_columns: Optional[Union[str, list[str]]] = None,
+        features: Optional[Features] = None,
         fn_kwargs: Optional[dict] = None,
     ) -> "IterableDatasetDict":
         """
@@ -2156,6 +2157,8 @@ def map(
                 Remove a selection of columns while doing the mapping.
                 Columns will be removed before updating the examples with the output of `function`, i.e. if `function` is adding
                 columns with names in `remove_columns`, these columns will be kept.
+            features (`[Features]`, *optional*, defaults to `None`):
+                Feature types of the resulting dataset.
             fn_kwargs (`Dict`, *optional*, defaults to `None`):
                 Keyword arguments to be passed to `function`
 
@@ -2187,6 +2190,7 @@ def map(
                 batch_size=batch_size,
                 drop_last_batch=drop_last_batch,
                 remove_columns=remove_columns,
+                features=features,
                 fn_kwargs=fn_kwargs,
             )