Setting `num_proc` errors when `.map` returns additional items.
Opened by davidgilbertson (closed)

### 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

Comments:
Hi ! Unfortunately I couldn't reproduce on my side locally and with datasets 2.11 and python 3.10.11 on colab.
What version of `multiprocess` are you using ?

I've got `multiprocess` version `0.70.14`.

I've done some more testing and the error only occurs in PyCharm's Python Console. It seems to be [this PyCharm bug](https://youtrack.jetbrains.com/issue/PY-51922/Multiprocessing-bug.-Can-only-run-in-debugger.), I'll close this.

For other users facing this, my workaround is to conditionally set `num_proc` so I can work interactively in the PyCharm Python Console while developing, then when I'm ready to run on the whole dataset, run it as a script and use multiprocessing.

```py
mapped_ds = ds.map(
    my_map_function,
    batched=True,
    remove_columns=ds.column_names,
    num_proc=1 if "PYCHARM_HOSTED" in os.environ else 8,
)
```