Title: Dataset/DatasetDict.filter() cannot have `batched=True` due to `mask` (numpy array?) being non-iterable.
Author: tianjianjiang (CONTRIBUTOR)
State: closed
Created: 2022-09-06T09:47:48Z
Labels: bug

## Describe the bug
`Dataset/DatasetDict.filter()` cannot have `batched=True` due to `mask` (numpy array?) being non-iterable.

## Steps to reproduce the bug
(In a python 3.7.12 env, I've tried 2.4.0 and 2.3.2 with both `pyarraw==9.0.0` and `pyarrow==8.0.0`.)

```python
from datasets import load_dataset


ds_mc4_ja = load_dataset("mc4", "ja")  # This will take 6+ hours... perhaps test it with a toy dataset instead?
ds_mc4_ja_2020 = ds_mc4_ja.filter(
  lambda example: example["timestamp"][:4] == "2020",
  batched=True,
)
```

## Expected results
No error

## Actual results
```python
---------------------------------------------------------------------------
RemoteTraceback                           Traceback (most recent call last)
RemoteTraceback: 
"""
Traceback (most recent call last):
  File "/opt/conda/lib/python3.7/site-packages/multiprocess/pool.py", line 121, in worker
    result = (True, func(*args, **kwds))
  File "/opt/conda/lib/python3.7/site-packages/datasets/arrow_dataset.py", line 557, in wrapper
    out: Union["Dataset", "DatasetDict"] = func(self, *args, **kwargs)
  File "/opt/conda/lib/python3.7/site-packages/datasets/arrow_dataset.py", line 524, in wrapper
    out: Union["Dataset", "DatasetDict"] = func(self, *args, **kwargs)
  File "/opt/conda/lib/python3.7/site-packages/datasets/fingerprint.py", line 480, in wrapper
    out = func(self, *args, **kwargs)
  File "/opt/conda/lib/python3.7/site-packages/datasets/arrow_dataset.py", line 2779, in _map_single
    offset=offset,
  File "/opt/conda/lib/python3.7/site-packages/datasets/arrow_dataset.py", line 2655, in apply_function_on_filtered_inputs
    processed_inputs = function(*fn_args, *additional_args, **fn_kwargs)
  File "/opt/conda/lib/python3.7/site-packages/datasets/arrow_dataset.py", line 2347, in decorated
    result = f(decorated_item, *args, **kwargs)
  File "/opt/conda/lib/python3.7/site-packages/datasets/arrow_dataset.py", line 4946, in get_indices_from_mask_function
    indices_array = [i for i, to_keep in zip(indices, mask) if to_keep]
TypeError: zip argument #2 must support iteration
"""

The above exception was the direct cause of the following exception:

TypeError                                 Traceback (most recent call last)
/tmp/ipykernel_51348/2345782281.py in <module>
      7     batched=True,
      8     # batch_size=10_000,
----> 9     num_proc=111,
     10 )
     11 # ds_mc4_ja_clean_2020 = ds_mc4_ja.filter(

/opt/conda/lib/python3.7/site-packages/datasets/dataset_dict.py in filter(self, function, with_indices, input_columns, batched, batch_size, keep_in_memory, load_from_cache_file, cache_file_names, writer_batch_size, fn_kwargs, num_proc, desc)
    878                     desc=desc,
    879                 )
--> 880                 for k, dataset in self.items()
    881             }
    882         )

/opt/conda/lib/python3.7/site-packages/datasets/dataset_dict.py in <dictcomp>(.0)
    878                     desc=desc,
    879                 )
--> 880                 for k, dataset in self.items()
    881             }
    882         )

/opt/conda/lib/python3.7/site-packages/datasets/arrow_dataset.py in wrapper(*args, **kwargs)
    522         }
    523         # apply actual function
--> 524         out: Union["Dataset", "DatasetDict"] = func(self, *args, **kwargs)
    525         datasets: List["Dataset"] = list(out.values()) if isinstance(out, dict) else [out]
    526         # re-apply format to the output

/opt/conda/lib/python3.7/site-packages/datasets/fingerprint.py in wrapper(*args, **kwargs)
    478             # Call actual function
    479 
--> 480             out = func(self, *args, **kwargs)
    481 
    482             # Update fingerprint of in-place transforms + update in-place history of transforms

/opt/conda/lib/python3.7/site-packages/datasets/arrow_dataset.py in filter(self, function, with_indices, input_columns, batched, batch_size, keep_in_memory, load_from_cache_file, cache_file_name, writer_batch_size, fn_kwargs, num_proc, suffix_template, new_fingerprint, desc)
   2920             new_fingerprint=new_fingerprint,
   2921             input_columns=input_columns,
-> 2922             desc=desc,
   2923         )
   2924         new_dataset = copy.deepcopy(self)

/opt/conda/lib/python3.7/site-packages/datasets/arrow_dataset.py in map(self, function, with_indices, with_rank, input_columns, batched, batch_size, drop_last_batch, remove_columns, keep_in_memory, load_from_cache_file, cache_file_name, writer_batch_size, features, disable_nullable, fn_kwargs, num_proc, suffix_template, new_fingerprint, desc)
   2498 
   2499                     for index, async_result in results.items():
-> 2500                         transformed_shards[index] = async_result.get()
   2501 
   2502             assert (

/opt/conda/lib/python3.7/site-packages/multiprocess/pool.py in get(self, timeout)
    655             return self._value
    656         else:
--> 657             raise self._value
    658 
    659     def _set(self, i, obj):

TypeError: zip argument #2 must support iteration
```

## Environment info
<!-- You can run the command `datasets-cli env` and copy-and-paste its output below. -->
- `datasets` version: 2.4.0
- Platform: Linux-4.19.0-21-cloud-amd64-x86_64-with-debian-10.12
- Python version: 3.7.12
- PyArrow version: 9.0.0
- Pandas version: 1.3.5

(I've tried 2.4.0 and 2.3.2 with both `pyarraw==9.0.0` and `pyarrow==8.0.0`.)

Comments:
Hi ! When `batched=True`, you filter function must take a batch as input, and return a list of booleans.

In your case, something like
```python
from datasets import load_dataset


ds_mc4_ja = load_dataset("mc4", "ja")  # This will take 6+ hours... perhaps test it with a toy dataset instead?
ds_mc4_ja_2020 = ds_mc4_ja.filter(
  lambda batch: [timestamp[:4] == "2020" for timestamp in batch["timestamp"]],
  batched=True,
)
```

Let me know if it helps !

> Hi ! When `batched=True`, you filter function must take a batch as input, and return a list of booleans.
> [...]
> Let me know if it helps !

Hi @lhoestq,

Ah, my bad, I totally forgot that part...
Sorry for the trouble and thank you for the kind help!