[FR] `map` should reuse unchanged columns from the previous dataset to avoid disk usage

[NightMachinery]

Feature request

Currently adding a new column with map will cause all the data in the dataset to be duplicated and stored/cached on the disk again. It should reuse unchanged columns.

Motivation

This allows having datasets with different columns but sharing some basic columns. Currently, these datasets would become too expensive to store and one would need some kind of on-the-fly join; which also doesn't seem implemented.

Your contribution

_

[mariosasko]

You can use the remove_columns parameter in map to avoid duplicating the columns (and save disk space) and then concatenate the original dataset with the map result:

from datasets import concatenate_datasets
# dummy example
ds_new = ds.map(lambda x: {"new_col": x["col"] + 2}, remove_columns=ds.column_names)
ds_combined = concatenate_datasets([ds, ds_new], axis=1)

Doing this automatically is hard to implement efficiently unless we know ahead of time which existing columns will be modified by a map transform. We have this info when input_columns are specified, so I think this is the only case we can optimize.

[ArjunJagdale]

Hi @mariosasko 👋 I’d like to start working on this issue.