Reject duplicate column names in Dataset.select_columns

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})

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:

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

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")})