GitHub issue #5817: Setting `num_proc` errors when `.map` returns additional items.
State: closed | Created: 2023-05-03T21:46:53Z

### Describe the bug

I'm using a map function that returns more rows than are passed in.

If I try to use `num_proc` I get:
```
  File "/home/davidg/.virtualenvs/learning/lib/python3.10/site-packages/datasets/arrow_dataset.py", line 563, in wrapper
    out: Union["Dataset", "DatasetDict"] = func(self, *args, **kwargs)
  File "/home/davidg/.virtualenvs/learning/lib/python3.10/site-packages/datasets/arrow_dataset.py", line 528, in wrapper
    out: Union["Dataset", "DatasetDict"] = func(self, *args, **kwargs)
  File "/home/davidg/.virtualenvs/learning/lib/python3.10/site-packages/datasets/arrow_dataset.py", line 3097, in map
    for rank, done, content in iflatmap_unordered(
  File "/home/davidg/.virtualenvs/learning/lib/python3.10/site-packages/datasets/utils/py_utils.py", line 1372, in iflatmap_unordered
    yield queue.get(timeout=0.05)
  File "<string>", line 2, in get
  File "/home/davidg/.virtualenvs/learning/lib/python3.10/site-packages/multiprocess/managers.py", line 818, in _callmethod
    kind, result = conn.recv()
  File "/home/davidg/.virtualenvs/learning/lib/python3.10/site-packages/multiprocess/connection.py", line 258, in recv
    buf = self._recv_bytes()
  File "/home/davidg/.virtualenvs/learning/lib/python3.10/site-packages/multiprocess/connection.py", line 422, in _recv_bytes
    buf = self._recv(4)
  File "/home/davidg/.virtualenvs/learning/lib/python3.10/site-packages/multiprocess/connection.py", line 391, in _recv
    raise EOFError
EOFError
```

### Steps to reproduce the bug

This is copied from the [Datasets docs](https://huggingface.co/docs/datasets/v2.12.0/en/process#batch-processing), with `num_proc` added, and will error.

```py
import datasets

dataset = ... # any old dataset


def chunk_examples(examples):
    chunks = []
    for sentence in examples["text"]:
        chunks += [sentence[i : i + 50] for i in range(0, len(sentence), 50)]
    return {"chunks": chunks}


chunked_dataset = dataset.map(
    chunk_examples,
    batched=True,
    remove_columns=dataset.column_names,
    num_proc=2,  # Remove and it works
)
```

### Expected behavior

Should work fine. On a related note, multi-processing also fails if there is a Meta class anywhere in scope (and there are plenty in the standard library). This is the fault of `dill` and is a long standing issue.

Have you considered using Loky for multiprocessing? I've found that the built-in `datasets` multi-processing breaks more than it works so have written my own function using `loky`, for reference:
```py
import datasets
import loky

def fast_loop(dataset: datasets.Dataset, func, num_proc=None):
    if num_proc is None:
        import os

        num_proc = len(os.sched_getaffinity(0))

    shards = [
        dataset.shard(num_shards=num_proc, index=i, contiguous=True)
        for i in range(num_proc)
    ]

    executor = loky.get_reusable_executor(max_workers=num_proc)
    results = executor.map(func, shards)
    return datasets.combine.concatenate_datasets(list(results))
```

### Environment info

- `datasets` version: 2.11.0
- Platform: Linux-5.15.90.1-microsoft-standard-WSL2-x86_64-with-glibc2.31
- Python version: 3.10.8
- Huggingface_hub version: 0.12.1
- PyArrow version: 11.0.0
- Pandas version: 2.0.1