Repository: huggingface/datasets
Pull request: Preserve non-`input_colums` in `Dataset.map` if `input_columns` are specified
Currently, if the `input_columns` list in `Dataset.map` is specified, the columns not in that list are dropped after the `map` transform.

This makes the behavior inconsistent with `IterableDataset.map`.
 
(It seems this issue was introduced by mistake in https://github.com/huggingface/datasets/pull/2246) 

Fix https://github.com/huggingface/datasets/issues/4858

Problem statement:
map() function removes columns when input_columns is not None
## Describe the bug
The map function, removes features from the dataset that are not present in the _input_columns_ list of columns, despite the columns being removed not mentioned in the _remove_columns_ argument.
## Steps to reproduce the bug
```python
from datasets import Dataset
ds = Dataset.from_dict({"a" : [1,2,3],"b" : [0,1,0], "c" : [2,4,5]})

def double(x,y):
  x = x*2
  y = y*2
  return {"d" : x, "e" : y}

ds.map(double, input_columns=["a","c"])
```

## Expected results
```
Dataset({
    features: ['a', 'b', 'c', 'd', 'e'],
    num_rows: 3
})
```
## Actual results
```
Dataset({
    features: ['a', 'c', 'd', 'e'],
    num_rows: 3
})
```

In this specific example feature **b** should not be removed.

## Environment info
<!-- You can run the command `datasets-cli env` and copy-and-paste its output below. -->
- `datasets` version: 2.4.0
- Platform: linux (colab)
- Python version: 3.7.13
- PyArrow version: 6.0.1

Hints from issue discussion:

Commit message: Remove unnecessary assert

Patch to review:
diff --git a/src/datasets/arrow_dataset.py b/src/datasets/arrow_dataset.py
index edca235535a..9f2f0403dd5 100644
--- a/src/datasets/arrow_dataset.py
+++ b/src/datasets/arrow_dataset.py
@@ -2731,7 +2731,7 @@ def init_buffer_and_writer():
                 # Only load the columns we actually need
                 if input_columns:
                     input_dataset = self.with_format(
-                        self._format_type, columns=input_columns, output_all_columns=False, **self._format_kwargs
+                        self._format_type, columns=input_columns, output_all_columns=True, **self._format_kwargs
                     )
                     if remove_columns:
                         remove_columns = list(set(remove_columns) & set(input_columns))
diff --git a/tests/test_arrow_dataset.py b/tests/test_arrow_dataset.py
index 7184fe25255..1e7ae94d5e3 100644
--- a/tests/test_arrow_dataset.py
+++ b/tests/test_arrow_dataset.py
@@ -1341,6 +1341,18 @@ def func(example):
                     )
                     self.assertListEqual(dset_test[0]["tensor"], [1, 2, 3])
 
+    def test_map_input_columns(self, in_memory):
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            with self._create_dummy_dataset(in_memory, tmp_dir) as dset:
+                print("dset", dset[0])
+                with dset.map(
+                    lambda filename: {"label": int(filename.rsplit("_", 1)[-1]) % 2}, input_columns="filename"
+                ) as mapped_dset:
+                    self.assertTrue("filename" in mapped_dset[0] and "label" in mapped_dset[0])
+                    self.assertEqual(
+                        mapped_dset.features, Features({"filename": Value("string"), "label": Value("int64")})
+                    )
+
     def test_map_remove_columns(self, in_memory):
         with tempfile.TemporaryDirectory() as tmp_dir:
             with self._create_dummy_dataset(in_memory, tmp_dir) as dset:


Code review comments:
diff_hunk: @@ -1341,6 +1341,18 @@ def func(example):
                     )
                     self.assertListEqual(dset_test[0]["tensor"], [1, 2, 3])
 
+    def test_map_input_columns(self, in_memory):
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            with self._create_dummy_dataset(in_memory, tmp_dir) as dset:
+                print("dset", dset[0])
path: tests/test_arrow_dataset.py
text: @user1:
We could remove this print:
```suggestion
```

diff_hunk: @@ -2731,7 +2731,7 @@ def init_buffer_and_writer():
                 # Only load the columns we actually need
                 if input_columns:
                     input_dataset = self.with_format(
-                        self._format_type, columns=input_columns, output_all_columns=False, **self._format_kwargs
+                        self._format_type, columns=input_columns, output_all_columns=True, **self._format_kwargs
path: src/datasets/arrow_dataset.py
text: @user1:
If we do this I think we can remove this part in `apply_function_on_filtered_inputs` no ? Otherwise the data would be loaded twice if a format type is set

```python
            if self._format_type is not None:
                inputs = self._getitem(
                    key=(indices if isinstance(indices, int) else slice(indices[0], indices[-1] + 1)),
                    format_type=None,
                    format_columns=None,
                    format_kwargs=None,
                    decoded=False,
                )
```

@author:
We still need this for the scenario where `input_columns` is `None` and self._format_type is not `None` (to ignore the formatting)

@author:
Okay, now `output_all_columns` is still `False` if `input_columns` is specified, but we also execute the above block to add the non-`input_columns` .