Revert input_columns change (#5006)

revert input_columns change
huggingface · Sep 21, 2022 · 4889d5d · 4889d5d · github-actions · Sep 21, 2022
1 parent cec23d5
commit 4889d5d
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 4 deletions.
diff --git a/src/datasets/arrow_dataset.py b/src/datasets/arrow_dataset.py
@@ -2733,7 +2733,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 or input_columns:
+            if self._format_type:
                 inputs = self._getitem(
                     key=(indices if isinstance(indices, int) else slice(indices[0], indices[-1] + 1)),
                     format_type=None,

diff --git a/tests/test_arrow_dataset.py b/tests/test_arrow_dataset.py
@@ -1362,14 +1362,12 @@ 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[0].keys(), {"col_1", "label"})
                     self.assertEqual(
                         mapped_dset.features,
                         Features(
                             {
                                 "col_1": Value("int64"),
-                                "col_2": Value("string"),
-                                "col_3": Value("bool"),
                                 "label": Value("int64"),
                             }
                         ),