Skip to content

Fix IterableDataset.remove_columns when a single column name is passed - #8384

Open
LeSingh1 wants to merge 1 commit into
huggingface:mainfrom
LeSingh1:fix/iterable-remove-columns-str
Open

Fix IterableDataset.remove_columns when a single column name is passed#8384
LeSingh1 wants to merge 1 commit into
huggingface:mainfrom
LeSingh1:fix/iterable-remove-columns-str

Conversation

@LeSingh1

@LeSingh1 LeSingh1 commented Aug 4, 2026

Copy link
Copy Markdown

IterableDataset.remove_columns accepts either a single column name or a list, but it only ever tests col in column_names. When a string is passed, that becomes a substring test, so any existing column whose name is contained in the one being removed is dropped from info.features:

ds = Dataset.from_dict({"label": [0, 1], "label_text": ["a", "b"]})
it = ds.to_iterable_dataset().remove_columns("label_text")

it.column_names  # []  -> should be ['label']
it.features      # {}
next(iter(it))   # {'label': 0}  -> the column is still being yielded

features and column_names then disagree with the actual examples, which breaks anything that reads the schema afterwards — select_columns, cast, or writing to parquet. The map-style Dataset.remove_columns("label_text") correctly returns ['label'], so the two paths disagree.

The existing test_iterable_dataset_remove_columns passes only incidentally: it removes "id" and no sibling column name contains "id".

This normalizes a string argument to a list first, which is what select_columns in the same class already does.

Added coverage for the single-string case with a sibling whose name is a substring; it fails on main with assert [] == ['label'] and passes with the change. tests/test_iterable_dataset.py and tests/test_dataset_dict.py pass (531 passed, 30 skipped; the 2 remaining failures are pre-existing network-dependent tests unrelated to this change).

…atures

remove_columns accepts either a single column name or a list, but it only ever
tested `col in column_names`. With a string argument that is a substring test,
so every existing column whose name is contained in the one being removed was
deleted from info.features while its values were still yielded:

    ds = Dataset.from_dict({"label": [0, 1], "label_text": ["a", "b"]})
    it = ds.to_iterable_dataset().remove_columns("label_text")
    it.column_names  # [] instead of ["label"]
    next(iter(it))   # {"label": 0}

features and column_names then disagree with the actual examples, which breaks
anything reading the schema afterwards (select_columns, cast, writing).

Normalize a string argument to a list first, like select_columns already does.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant