From 70ff65dc6b20390a64b8518ff99153168f77e506 Mon Sep 17 00:00:00 2001 From: Muhammad Abdullah Rasheed Date: Mon, 27 Jul 2026 01:57:36 +0100 Subject: [PATCH 1/2] Reject duplicate column names in Dataset.select_columns select_columns validated that every requested column exists, but not that the request was free of duplicates. Passing the same name twice built a dataset whose Arrow table carried the column twice while its features carried it once, since the features are rebuilt with a dict comprehension that collapses the repeat. The result is a dataset that is broken for ordinary use: num_columns disagrees with len(features), and both indexing and map raise KeyError('Field "x" exists 2 times in schema'). Raise a ValueError instead, matching how rename_columns already rejects a mapping that would produce duplicate names. --- src/datasets/arrow_dataset.py | 7 +++++++ tests/test_arrow_dataset.py | 17 +++++++++++++++++ 2 files changed, 24 insertions(+) diff --git a/src/datasets/arrow_dataset.py b/src/datasets/arrow_dataset.py index ba2dbe13d61..14d067bb911 100644 --- a/src/datasets/arrow_dataset.py +++ b/src/datasets/arrow_dataset.py @@ -2716,6 +2716,13 @@ def select_columns(self, column_names: Union[str, list[str]], new_fingerprint: O f"{self._data.column_names}." ) + number_of_duplicates = len(column_names) - len(set(column_names)) + if number_of_duplicates != 0: + raise ValueError( + "Selected column names must all be different, but this selection " + f"has {number_of_duplicates} duplicates." + ) + dataset = copy.deepcopy(self) dataset._data = dataset._data.select(column_names) dataset._info.features = Features({col: self._info.features[col] for col in dataset._data.column_names}) diff --git a/tests/test_arrow_dataset.py b/tests/test_arrow_dataset.py index fcf05cece1b..e7639509da4 100644 --- a/tests/test_arrow_dataset.py +++ b/tests/test_arrow_dataset.py @@ -729,6 +729,23 @@ def test_select_columns(self, in_memory): self.assertNotEqual(new_dset._fingerprint, fingerprint) assert_arrow_metadata_are_synced_with_dataset_features(new_dset) + def test_select_columns_duplicates(self, in_memory): + # Selecting the same column twice used to build a dataset whose Arrow table had + # the column twice while its features had it once. That object then raised + # KeyError('Field "col_1" exists 2 times in schema') on ordinary access. + with tempfile.TemporaryDirectory() as tmp_dir: + with self._create_dummy_dataset(in_memory, tmp_dir, multiple_columns=True) as dset: + with self.assertRaises(ValueError): + dset.select_columns(column_names=["col_1", "col_1"]) + + with self.assertRaises(ValueError): + dset.select_columns(column_names=["col_1", "col_2", "col_1"]) + + # Selecting distinct columns, including reordered, is unaffected. + with dset.select_columns(column_names=["col_2", "col_1"]) as new_dset: + self.assertListEqual(list(new_dset.column_names), ["col_2", "col_1"]) + self.assertEqual(new_dset.num_columns, len(new_dset.features)) + def test_concatenate(self, in_memory): data1, data2, data3 = {"id": [0, 1, 2]}, {"id": [3, 4, 5]}, {"id": [6, 7]} info1 = DatasetInfo(description="Dataset1") From a39964312866ff91a0394c854d20919d4a834c25 Mon Sep 17 00:00:00 2001 From: Muhammad Abdullah Rasheed Date: Mon, 27 Jul 2026 13:40:25 +0100 Subject: [PATCH 2/2] Extend duplicate select_columns check to IterableDataset Dataset.select_columns already rejects duplicate column names. Apply the same guard to IterableDataset.select_columns so both classes behave consistently. Add a matching test. --- src/datasets/iterable_dataset.py | 7 +++++++ tests/test_iterable_dataset.py | 14 ++++++++++++++ 2 files changed, 21 insertions(+) diff --git a/src/datasets/iterable_dataset.py b/src/datasets/iterable_dataset.py index 0f917fde272..691056e5859 100644 --- a/src/datasets/iterable_dataset.py +++ b/src/datasets/iterable_dataset.py @@ -4181,6 +4181,13 @@ def select_columns(self, column_names: Union[str, list[str]]) -> "IterableDatase if isinstance(column_names, str): column_names = [column_names] + number_of_duplicates = len(column_names) - len(set(column_names)) + if number_of_duplicates != 0: + raise ValueError( + "Selected column names must all be different, but this selection " + f"has {number_of_duplicates} duplicates." + ) + if self._info: info = deepcopy(self._info) if self._info.features is not None: diff --git a/tests/test_iterable_dataset.py b/tests/test_iterable_dataset.py index f578e3bccdc..6a80e433ab1 100644 --- a/tests/test_iterable_dataset.py +++ b/tests/test_iterable_dataset.py @@ -2484,6 +2484,20 @@ def test_iterable_dataset_select_columns(dataset_with_several_columns: IterableD assert all(c in new_dataset.column_names for c in ["id", "filepath"]) +def test_iterable_dataset_select_columns_duplicates(dataset_with_several_columns: IterableDataset): + with pytest.raises(ValueError, match="duplicates"): + dataset_with_several_columns.select_columns(["id", "id"]) + + with pytest.raises(ValueError, match="duplicates"): + dataset_with_several_columns.select_columns(["id", "filepath", "id"]) + + # Selecting distinct columns, including reordered, is unaffected. + new_dataset = dataset_with_several_columns.select_columns(["filepath", "id"]) + assert list(new_dataset) == [ + {k: v for k, v in example.items() if k in ("id", "filepath")} for example in dataset_with_several_columns + ] + + def test_iterable_dataset_cast_column(): ex_iterable = ExamplesIterable(generate_examples_fn, {"label": 10}) features = Features({"id": Value("int64"), "label": Value("int64")})