From 233abedbc6b6981cc3383a6804167132700ea2df Mon Sep 17 00:00:00 2001 From: uttam12331 Date: Wed, 1 Jul 2026 11:10:19 +0530 Subject: [PATCH] Fix: preserve features after IterableDataset.add_column (GH#5752) add_column delegates to map() but does not pass features, causing map() to set info.features = None and losing the original schema. Fix by inferring the new column's feature type from the provided data, merging it with the existing features, and passing the result to map() so the returned dataset's .features are correctly set. --- src/datasets/iterable_dataset.py | 10 +++++++++- tests/test_iterable_dataset.py | 14 ++++++++++++++ 2 files changed, 23 insertions(+), 1 deletion(-) diff --git a/src/datasets/iterable_dataset.py b/src/datasets/iterable_dataset.py index 92fdea2ad4d..c235af0744e 100644 --- a/src/datasets/iterable_dataset.py +++ b/src/datasets/iterable_dataset.py @@ -4032,7 +4032,15 @@ 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) + # GH#5752: Infer feature type for the new column and merge with existing + # features so that the returned dataset's .features are preserved. + column_features = _infer_features_from_batch({name: list(column) if not isinstance(column, list) else column}) + if self.features is not None: + features = self.features.copy() + features.update(column_features) + else: + features = column_features + return self.map(partial(add_column_fn, name=name, column=column), with_indices=True, features=features) def rename_column(self, original_column_name: str, new_column_name: str) -> "IterableDataset": """ diff --git a/tests/test_iterable_dataset.py b/tests/test_iterable_dataset.py index 3fd22d63809..527848516e6 100644 --- a/tests/test_iterable_dataset.py +++ b/tests/test_iterable_dataset.py @@ -2327,6 +2327,20 @@ def test_iterable_dataset_add_column(dataset_with_several_columns: IterableDatas assert "new_column" in new_dataset.column_names +def test_iterable_dataset_add_column_preserves_features(dataset_with_several_columns: IterableDataset): + # GH#5752 - .features was lost (became None) after .add_column + ds_with_features = dataset_with_several_columns._resolve_features() + assert ds_with_features.features is not None + original_feature_names = set(ds_with_features.features) + + new_column = ["x"] * (3 * DEFAULT_N_EXAMPLES) + new_dataset = ds_with_features.add_column("new_column", new_column) + + assert new_dataset.features is not None, ".features should be preserved after add_column" + assert "new_column" in new_dataset.features, "new_column should appear in features" + assert original_feature_names.issubset(new_dataset.features), "original features should be preserved" + + def test_iterable_dataset_rename_column(dataset_with_several_columns: IterableDataset): new_dataset = dataset_with_several_columns.rename_column("id", "new_id") assert list(new_dataset) == [