Skip to content

Preserve features after IterableDataset.add_column#8316

Open
sridhar-3009 wants to merge 1 commit into
huggingface:mainfrom
sridhar-3009:fix/5752-iterable-add-column-features
Open

Preserve features after IterableDataset.add_column#8316
sridhar-3009 wants to merge 1 commit into
huggingface:mainfrom
sridhar-3009:fix/5752-iterable-add-column-features

Conversation

@sridhar-3009

Copy link
Copy Markdown

What

IterableDataset.add_column delegates to map but never propagates the schema. After the call, dataset.features is None even when the original dataset had a typed schema, silently breaking all downstream code that inspects features.

Root cause

The implementation is a one-liner:

return self.map(partial(add_column_fn, name=name, column=column), with_indices=True)

map without an explicit features argument does not infer schema from the mapper output.

Fix

Follow the same pattern used by rename_columns: snapshot the existing features before the map, infer the new column's type from the column data via PyArrow (one line, pyarrow is already a required dependency), merge, and write back to _info.features.

original_features = self._info.features.copy() if self._info.features else None
ds_iterable = self.map(...)
if original_features is not None:
    col_table = pa.table({name: column})
    new_features = original_features.copy()
    new_features[name] = Features.from_arrow_schema(col_table.schema)[name]
    ds_iterable._info.features = new_features
return ds_iterable

If the original dataset had no features schema (features is None), behavior is unchanged.

Fixes #5752.

`add_column` delegates to `map` but never updates the returned dataset's
features schema. Calling `add_column` on a typed IterableDataset silently
drops the `features` attribute (returns `None`), breaking any code that
inspects schema after the operation.

Follow the same pattern as `rename_columns`: save the existing features
before the map call, infer the new column's type from the column data via
PyArrow, then write the merged Features back to `_info.features`.

Fixes huggingface#5752.
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.

Streaming dataset looses .feature method after .add_column

1 participant