Preserve features after IterableDataset.add_column - #8379
Conversation
| return self.map(partial(add_column_fn, name=name, column=column), with_indices=True) | ||
| features = self.features.copy() if self.features is not None else None | ||
| if features is not None: | ||
| features.update(_infer_features_from_batch({name: column})) |
There was a problem hiding this comment.
inferring off {name: column} builds an arrow table from the whole column, which both copies the full column up front and narrows what add_column accepts: the signature allows np.array, and pa.Table.from_pydict raises ArrowInvalid: only handle 1-dimensional arrays for a 2D ndarray, so a call that used to just end up untyped now blows up before map. inferring from column[:1] (and falling back to None on ArrowInvalid) would keep that working. also worth mirroring Dataset.add_column's feature= arg so the caller can skip inference entirely.
| return self.map(partial(add_column_fn, name=name, column=column), with_indices=True) | ||
| features = self.features.copy() if self.features is not None else None | ||
| if features is not None: | ||
| features.update(_infer_features_from_batch({name: column})) |
There was a problem hiding this comment.
the docstring types column as Union[list, np.array], and a 2-D ndarray (adding an embedding column) goes straight into pa.Table.from_pydict here and raises ArrowInvalid: only handle 1-dimensional arrays. that used to work since features stayed None. worth wrapping the infer in a try and falling back to features=None, or converting to a list of rows first.
2 participants
Summary
IterableDatasetfeatures whenadd_columnappends a streaming columnadd_columnafter resolving iterable dataset featuresFixes #5752.
Testing
IterableDataset.from_generator(...).add_column(...)python -m pytest tests/test_iterable_dataset.py::test_iterable_dataset_add_column tests/test_iterable_dataset.py::test_iterable_dataset_add_column_preserves_featurespython -m pytest tests/test_iterable_dataset.py::test_pickle_after_many_transforms