From 361daa3ac3954c2c90eb633778a702639f44181e Mon Sep 17 00:00:00 2001 From: Sai Sridhar Date: Fri, 10 Jul 2026 09:45:26 +0530 Subject: [PATCH] Preserve features after IterableDataset.add_column `add_column` delegates to `map` but never updates the returned dataset's features schema. Calling `add_column` on a typed IterableDataset silently drops the `features` attribute (returns `None`), breaking any code that inspects schema after the operation. Follow the same pattern as `rename_columns`: save the existing features before the map call, infer the new column's type from the column data via PyArrow, then write the merged Features back to `_info.features`. Fixes #5752. --- src/datasets/iterable_dataset.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) 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": """