diff --git a/src/datasets/arrow_dataset.py b/src/datasets/arrow_dataset.py index 59451a640e6..247e12b30a2 100644 --- a/src/datasets/arrow_dataset.py +++ b/src/datasets/arrow_dataset.py @@ -2655,6 +2655,16 @@ def rename_columns(self, column_mapping: dict[str, str], new_fingerprint: Option if empty_new_columns: raise ValueError(f"New column names {empty_new_columns} are empty.") + colliding_new_columns = set(column_mapping.values()) & ( + set(dataset.column_names) - set(column_mapping.keys()) + ) + if colliding_new_columns: + raise ValueError( + f"New column names {colliding_new_columns} already in the dataset. " + f"Please choose column names which are not already in the dataset. " + f"Current columns in the dataset: {dataset._data.column_names}" + ) + def rename(columns): return [column_mapping[col] if col in column_mapping else col for col in columns] diff --git a/tests/test_arrow_dataset.py b/tests/test_arrow_dataset.py index 18a8a6038fe..6979390f76f 100644 --- a/tests/test_arrow_dataset.py +++ b/tests/test_arrow_dataset.py @@ -688,6 +688,17 @@ def test_rename_columns(self, in_memory): with self.assertRaises(ValueError): dset.rename_columns({"col_1": "new_name", "col_2": "new_name"}) + # New name collides with an existing column that isn't itself being renamed: + # this must raise instead of silently producing two columns named "col_2" + # while `features`/`column_names` disagree with the underlying arrow table. + with self.assertRaises(ValueError): + dset.rename_columns({"col_1": "col_2"}) + + # A swap (each renamed column's new name is itself being renamed away) must still work. + with dset.rename_columns({"col_1": "col_2", "col_2": "col_1"}) as new_dset: + self.assertEqual(new_dset.num_columns, 3) + self.assertListEqual(list(new_dset.column_names), ["col_2", "col_1", "col_3"]) + def test_select_columns(self, in_memory): with tempfile.TemporaryDirectory() as tmp_dir: with self._create_dummy_dataset(in_memory, tmp_dir, multiple_columns=True) as dset: