Allow opposite of remove_columns on Dataset and DatasetDict

[hollance]

Feature request

In this blog post https://huggingface.co/blog/audio-datasets, I noticed the following code:

COLUMNS_TO_KEEP = ["text", "audio"]
all_columns = gigaspeech["train"].column_names
columns_to_remove = set(all_columns) - set(COLUMNS_TO_KEEP)

gigaspeech = gigaspeech.remove_columns(columns_to_remove)

This kind of thing happens a lot when you don't need to keep all columns from the dataset. It would be more convenient (and less error prone) if you could just write:

gigaspeech = gigaspeech.keep_columns(["text", "audio"])

Internally, keep_columns could still call remove_columns, but it expresses more clearly what the user's intent is.

Motivation

Less code to write for the user of the dataset.

Your contribution

[mariosasko]

Hi! I agree it would be nice to have a method like that. Instead of keep_columns, we can name it select_columns to be more aligned with PyArrow's naming convention (pa.Table.select).

[avisinghal6]

Hi, I am a newbie to open source and would like to contribute. @mariosasko can I take up this issue ?

[vijaybirju]

Hey, I also want to work on this issue I am a newbie to open source.

[lhoestq]

This sounds related to #5474

I'm fine with select_columns, or we could also override select to also accept a list of columns maybe ?

[avisinghal6]

@lhoestq, I am planning to add a member function to the dataset class to perform the selection operation. Do you think its the right way to proceed? or there is a better option ?

[lhoestq]

Unless @mariosasko thinks otherwise, I think it can go in Dataset.select() :)
Though some parameters like keep_in_memory, indices_cache_file_name or writer_batch_size wouldn't when selecting columns, so we would need to update the docstring as well

[lhoestq]

If someone wants to give it a shot, feel free to comment #self-assign and it will assign the issue to you.

Feel free to ping us here if you have questions or if we can help :)

[mariosasko]

I would rather have this functionality as a separate method. IMO it's always better to be explicit than to have an API where a single method can do different/uncorrelated things (somewhat reminds me of Pandas, and there is probably a good reason why PyArrow is more rigid in this aspect).

[lhoestq]

In the end I also think it would be nice to have it as a separate method, this way we can also have it for IterableDataset (which can't have select for indices)