Side effect when filtering data due to `does_function_return_dict` call in `Dataset.map()`

Hi there!

In my codebase I have a function to filter rows in a dataset, selecting only a certain number of examples per class. The function passes a extra argument to maintain a counter of the number of dataset rows/examples already selected per each class, which are the ones I want to keep in the end:

```python
      def fill_train_examples_per_class(example, per_class_limit: int, counter: collections.Counter):
          label = int(example['label'])
          current_counter = counter.get(label, 0)
          if current_counter < per_class_limit:
              counter[label] = current_counter + 1
              return True
          return False
```

At some point I invoke it through the `Dataset.filter()` method in the `arrow_dataset.py` module like this:

```python
...
kwargs = {"per_class_limit": train_examples_per_class_limit, "counter": Counter()}
datasets['train'] = datasets['train'].filter(fill_train_examples_per_class,  num_proc=1, fn_kwargs=kwargs)
...
```

The problem is that, passing a stateful container (the counter,) provokes a side effect in the new filtered dataset obtained. This is due to the fact that at some point in `filter()`, the `map()`'s function `does_function_return_dict` is invoked in  line [1290](https://github.com/huggingface/datasets/blob/96578adface7e4bc1f3e8bafbac920d72ca1ca60/src/datasets/arrow_dataset.py#L1290). 

When this occurs, the state of the counter is initially modified by the effects of the function call on the 1 or 2 rows selected in lines 1288 and 1289 of the same file (which are marked as `test_inputs` & `test_indices` respectively in lines 1288 and 1289. This happens out of the control of the user (which for example can't reset the state of the counter before continuing the execution,) provoking in the end an undesired side effect in the results obtained. 

In my case, the resulting dataset -despite of the counter results are ok- lacks an instance of the classes 0 and 1 (which happen to be the classes of the first two examples of my dataset.) The rest of the classes I have in my dataset, contain the right number of examples as they were not affected by the effects of `does_function_return_dict` call.

I've debugged my code extensively and made a workaround myself hardcoding the necessary stuff (basically putting `update_data=True` in line 1290,) and then I obtain the results I expected without the side effect.

Is there a way to avoid that call to `does_function_return_dict` in map()'s line 1290 ? (e.g. extracting the required information that `does_function_return_dict` returns without making the testing calls to the user function on dataset rows 0 & 1) 

Thanks in advance,

Francisco Perez-Sorrosal



Labels: name: enhancement

Thanks for the report !

Currently we don't have a way to let the user easily disable this behavior.
However I agree that we should support stateful processing functions, ideally by removing `does_function_return_dict`.

We needed this function in order to know whether the `map` functions needs to write data or not. if `does_function_return_dict` returns False then we don't write anything.

Instead of checking the output of the processing function outside of the for loop that iterates through the dataset to process it, we can check the output of the first processed example and at that point decide if we need to write data or not.

Therefore it's definitely possible to fix this unwanted behavior, any contribution going into this direction is welcome :)

Thanks @mariosasko for the PR!