Repository: huggingface/datasets
Pull request: Add `remove_columns` to `IterableDataset`
Fixes #2944

WIP
* Not tested yet.
* We might want to allow batched remove for efficiency.

@lhoestq Do you think it should have `batched=` and `batch_size=`?

Problem statement:
Add  `remove_columns` to `IterableDataset ` 
**Is your feature request related to a problem? Please describe.**
A clear and concise description of what the problem is.

```python
from datasets import load_dataset
dataset = load_dataset("c4", 'realnewslike', streaming =True, split='train')
dataset = dataset.remove_columns('url')
```
```
AttributeError: 'IterableDataset' object has no attribute 'remove_columns'
```

**Describe the solution you'd like**

It would be nice to have `.remove_columns()` to match the `Datasets` api. 


**Describe alternatives you've considered**

This can be done with a single call to `.map()`, 

I can try to help add this. 🤗

Hints from issue discussion:
Hi ! Good idea :)
If you are interested in contributing, feel free to give it a try and open a Pull Request. Also let me know if I can help you with this or if you have questions

Commit message: Add missing import

Patch to review:
diff --git a/src/datasets/iterable_dataset.py b/src/datasets/iterable_dataset.py
index b1df83dfcf1..4594af9bd56 100644
--- a/src/datasets/iterable_dataset.py
+++ b/src/datasets/iterable_dataset.py
@@ -1,7 +1,7 @@
 import copy
 from dataclasses import dataclass
 from itertools import cycle, islice, repeat
-from typing import Any, Callable, Dict, Iterable, Iterator, List, Optional
+from typing import Any, Callable, Dict, Iterable, Iterator, List, Optional, Union
 
 import numpy as np
 import pyarrow as pa
@@ -471,6 +471,29 @@ def take(self, n) -> "IterableDataset":
             shuffling=copy.deepcopy(self._shuffling),
         )
 
+    def remove_columns(self, column_names: Union[str, List[str]]) -> "IterableDataset":
+        """
+        Remove one or several column(s) in the dataset and the features associated to them.
+        The removal is done on-the-fly on the examples when iterating over the dataset.
+
+        You can also remove a column using :func:`Dataset.map` with `remove_columns` but the present method
+        is in-place (doesn't copy the data to a new dataset) and is thus faster.
+
+        Args:
+            column_names (:obj:`Union[str, List[str]]`): Name of the column(s) to remove.
+            new_fingerprint
+
+        Returns:
+            :class:`IterableDataset`: A copy of the dataset object without the columns to remove.
+        """
+        if isinstance(column_names, str):
+            column_names = [column_names]
+
+        def remove_fn(example):
+            return {k: v for k, v in example.items() if k not in column_names}
+
+        return self.map(remove_fn)
+
 
 def iterable_dataset(
     ex_iterable: Iterable,


Code review comments:
diff_hunk: @@ -471,6 +471,29 @@ def take(self, n) -> "IterableDataset":
             shuffling=copy.deepcopy(self._shuffling),
         )
 
+    def remove_columns(self, column_names: Union[str, List[str]]) -> "IterableDataset":
+        """
+        Remove one or several column(s) in the dataset and the features associated to them.
+        The removal is done on-the-fly on the examples when iterating over the dataset.
+
+        You can also remove a column using :func:`Dataset.map` with `remove_columns` but the present method
+        is in-place (doesn't copy the data to a new dataset) and is thus faster.
+
+        Args:
+            column_names (:obj:`Union[str, List[str]]`): Name of the column(s) to remove.
+            new_fingerprint
path: src/datasets/iterable_dataset.py
text: @user1:
```suggestion
```