Reject duplicate column names in Dataset.select_columns - #8368
Reject duplicate column names in Dataset.select_columns#8368AbdullahRasheed45 wants to merge 2 commits into
Conversation
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.
ebarkhordar
left a comment
There was a problem hiding this comment.
The duplicate check is the right call, and it holds up on a run. On main, select_columns(["col_1", "col_1"]) returns a dataset whose Arrow table carries the column twice while features has it once, and d["col_1"], select, sort, flatten_indices and save_to_disk then all fail with KeyError: 'Field "col_1" exists 2 times in schema'. The guard removes that, and the new test actually detects it: deleting just the seven added lines in arrow_dataset.py takes test_select_columns_duplicates from 2 passed to 2 failed with "ValueError not raised". tests/test_arrow_dataset.py -k select_columns stays green (4 passed).
One scope question. IterableDataset.select_columns (iterable_dataset.py:4156) takes the same argument and is untouched here, so after this PR the two classes answer the same call differently. At this branch:
IterableDataset.select_columns(['col_1','col_1'])
features : ['col_1']
rows : [{'col_1': 0}, {'col_1': 1}, {'col_1': 2}]
Dataset.select_columns(['col_1','col_1'])
ValueError: Selected column names must all be different, but this selection has 1 duplicates.
The iterable path dedupes silently instead of corrupting anything, so it is not the same bug, but it does mean the validation is now class-dependent. DatasetDict and IterableDatasetDict each inherit whichever their element class does (dataset_dict.py:507 and :2261), so the split follows through to them. Is keeping this to Dataset deliberate, or would the same check be welcome in the iterable path?
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.
The problem
Dataset.select_columnschecks that every requested column exists, but never checks that the request is free of duplicates:Passing the same name twice then produces a dataset that is internally inconsistent. PyArrow's
selecthonours the repeat and returns two columns, while the features are rebuilt with a dict comprehension that collapses it back to one:On current
main:So the call succeeds and hands back an object that raises on its own basic operations. The failure surfaces later, at an unrelated line, with a message that points at the Arrow schema rather than at the
select_columnscall that caused it.The change
Raise a
ValueErrorwhen the selection contains duplicates.This mirrors
rename_columns, which already rejects the analogous case a few methods up:Selecting distinct columns — including reordering, a single string, or an empty list — is unaffected.
A consistency question for you
IterableDataset.select_columns(["x", "x"])does not fail today: it silently de-duplicates and yields{"x": ...}. So the two classes disagree about what a duplicated selection means.I chose to raise on
Datasetrather than de-duplicate, because silently dropping part of what the caller asked for hides a mistake, and becauserename_columnssets the precedent for rejecting duplicates outright. But that leaves the divergence in place, and aligning the two is a call for you rather than me — I'm happy to follow up either by makingIterableDatasetraise too, or by switching this to de-duplicate instead, whichever you prefer.Adjacent, not touched here:
remove_columns(["x", "x"])raises a bareKeyError: 'x', which at least fails rather than corrupting, but is not a clear message either.Tests
Added
test_select_columns_duplicatesnext to the existingtest_select_columns, following the same_create_dummy_datasetpattern. It covers the plain duplicate and a duplicate mixed among distinct columns, and asserts that a reordered distinct selection still works withnum_columns == len(features)— the invariant the bug broke.It fails on
mainin both the in-memory and on-disk parametrizations, and passes with the fix.pytest tests/test_arrow_dataset.py -k "select_columns or remove_columns or rename_column or concatenate or flatten"→ 45 passed.ruff checkandruff formatclean.