Reject features with different columns in IterableDataset.cast

[LeSingh1]

IterableDataset.cast() assigns the given features directly:

features = _fix_for_backward_compatible_features(features)
info = self._info.copy()
info.features = features

There is no check that the columns match the dataset's, which Dataset.cast() does perform. The worst case is an extra column: it is accepted silently and then fabricated in the data, because the declared features are applied to the examples on the fly.

ds = Dataset.from_dict({"x": [1, 2], "y": ["a", "b"]}).to_iterable_dataset()

new_features = Features({"x": Value("int64"), "y": Value("string"), "z": Value("int64")})
list(ds.cast(new_features))
# [{'x': 1, 'y': 'a', 'z': None}, {'x': 2, 'y': 'b', 'z': None}]

No error at any point — the dataset just grows a column that was never in it. The map-style call raises:

ValueError: The columns in features (['x', 'y', 'z']) must be identical as the columns in the dataset: ['x', 'y']

Missing and renamed columns are accepted too, and fail later with a CastError: Couldn't cast from deep in the Arrow cast, which doesn't say the column set is wrong.

Fix

Apply the same check and the same message as Dataset.cast() when the features are known:

if info.features is not None and sorted(features) != sorted(info.features):
    raise ValueError(...)

sorted() is used on both sides, exactly like the map-style version, so reordering the columns keeps working. Datasets whose features aren't resolved yet (info.features is None) are untouched.

Test

test_iterable_dataset_cast_with_different_columns, parametrized over a missing column, an extra column and a renamed column. All three fail on main.

[shashvat-singham]

Ran the four interesting shapes against both Dataset.cast and IterableDataset.cast, on main and on this branch. The divergence is exactly as described, and the fix closes it without touching the case that should keep working:

main

features passed	Dataset.cast	IterableDataset.cast
missing a column	ValueError	OK → ['id']
extra column	ValueError	OK → ['id', 'label', 'extra']
renamed column	ValueError	OK → ['id', 'renamed']
reordered only	OK	OK

this branch

features passed	Dataset.cast	IterableDataset.cast
missing a column	ValueError	ValueError
extra column	ValueError	ValueError
renamed column	ValueError	ValueError
reordered only	OK	OK

Two things worth calling out in favour of the change:

The error text now matches the eager one exactly. Both paths produce The columns in features ([...]) must be identical as the columns in the dataset: [...], so someone who moved a pipeline from Dataset to streaming gets the same message rather than a new one. Reusing the wording was the right call even though it means duplicating the string.

Reordering still works, and that matters. sorted(features) != sorted(info.features) compares names, so Features({"label": ..., "id": ...}) against {"id": ..., "label": ...} is still accepted on both paths, which is what Dataset.cast does today. Worth stating explicitly in the PR description, because "reject features with different columns" reads like it might have caught reordering too, and a reviewer who assumes that would flag it as a regression. Your parametrized test covers missing/extra/renamed — adding the reorder case as an explicit positive assertion would pin that this is deliberate rather than incidental.

One small wording nit on the message you inherited: "must be identical as the columns" should be "identical to". It is pre-existing in arrow_dataset.py, so fixing it here would mean touching both to keep them in sync — probably not worth it in this PR, but worth knowing it is now duplicated in two places rather than one.

Verified on Windows 11 / Python 3.11.9, main vs pr-8453.