Skip to content

Raise on Dataset.rename_columns name collision with an existing column #8323

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
roli-lpci wants to merge 1 commit into
base: main
Choose a base branch
from
+21 −0
Open

Raise on Dataset.rename_columns name collision with an existing column #8323

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 10 additions & 0 deletions src/datasets/arrow_dataset.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()) & (

@Sanjays2402 Sanjays2402 Jul 31, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

IterableDataset.rename_columns has no equivalent guard, so the two APIs diverge after this: streaming still only catches the collision inside _rename_columns_fn at iteration time, and that check is any(col in example for col in column_mapping.values()), which also rejects the swap case you added a test for here ({"col_1": "col_2", "col_2": "col_1"} — both new names are present in the example). Worth mirroring the eager check in iterable_dataset.py so the swap works and the collision fails at call time rather than on first next()?

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. "

@Sanjays2402 Sanjays2402 Jul 30, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

the streaming sibling has the same hole but worse: IterableDataset.rename_columns does no check at all, and _rename_columns_fn tests any(col in example for col in column_mapping.values()) before the old columns are removed, so on 5.0.1 even the swap case your test blesses here ({"a": "b", "b": "a"}) blows up at iteration with columns set() are already in the dataset. — and {"a": "b"} silently gives features == {"b": ...} while erroring only once you iterate. worth mirroring this eager check there so both paths fail the same way?

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]

Expand Down
11 changes: 11 additions & 0 deletions tests/test_arrow_dataset.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down