Preserve non-input_colums in Dataset.map if input_columns are specified (#4971)

* Preserve non-`input_colums` in `Dataset.map`

* Remove unnecessary assert

* Address review comments

Author: mariosasko

src/datasets/arrow_dataset.py

@@ -2680,7 +2680,7 @@ def apply_function_on_filtered_inputs(inputs, indices, check_same_num_examples=F
                 validate_function_output(processed_inputs, indices)
             if not update_data:
                 return None  # Nothing to update, let's move on
-            if self._format_type is not None:
+            if self._format_type is not None or input_columns:
                 inputs = self._getitem(
                     key=(indices if isinstance(indices, int) else slice(indices[0], indices[-1] + 1)),
                     format_type=None,

tests/test_arrow_dataset.py

@@ -1341,6 +1341,23 @@ 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, multiple_columns=True) as dset:
+                with dset.map(lambda col_1: {"label": col_1 % 2}, input_columns="col_1") as mapped_dset:
+                    self.assertEqual(mapped_dset[0].keys(), {"col_1", "col_2", "col_3", "label"})
+                    self.assertEqual(
+                        mapped_dset.features,
+                        Features(
+                            {
+                                "col_1": Value("int64"),
+                                "col_2": Value("string"),
+                                "col_3": Value("bool"),
+                                "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: