diff --git a/src/datasets/iterable_dataset.py b/src/datasets/iterable_dataset.py index 17b9a2020fc..1c00a2cba54 100644 --- a/src/datasets/iterable_dataset.py +++ b/src/datasets/iterable_dataset.py @@ -4036,7 +4036,14 @@ def add_column(self, name: str, column: Union[list, np.array]) -> "IterableDatas Returns: `IterableDataset` """ - return self.map(partial(add_column_fn, name=name, column=column), with_indices=True) + original_features = self._info.features.copy() if self._info.features else None + ds_iterable = self.map(partial(add_column_fn, name=name, column=column), with_indices=True) + if original_features is not None: + col_table = pa.table({name: column}) + new_features = original_features.copy() + new_features[name] = Features.from_arrow_schema(col_table.schema)[name] + ds_iterable._info.features = new_features + return ds_iterable def rename_column(self, original_column_name: str, new_column_name: str) -> "IterableDataset": """