Repository: huggingface/datasets
Pull request: add multi-proc in `to_json`
Closes #2663. I've tried adding multiprocessing in `to_json`. Here's some benchmarking I did to compare the timings of current version (say v1) and multi-proc version (say v2). I did this with `cpu_count` 4 (2015 Macbook Air)

1. Dataset name: `ascent_kb` - 8.9M samples (all samples were used, reporting this for a single run)
v1- ~225 seconds for converting whole dataset to json
v2- ~200 seconds for converting whole dataset to json

2. Dataset name: `lama` - 1.3M samples (all samples were used, reporting this for 2 runs)
v1- ~26 seconds for converting whole dataset to json
v2- ~23.6 seconds for converting whole dataset to json

I think it's safe to say that v2 is 10% faster as compared to v1. Timings may improve further with better configuration.

The only bottleneck I feel is writing to file from the output list. If we can improve that aspect then timings may improve further. 

Let me know if any changes/improvements can be done in this @stas00, @lhoestq, @albertvillanova. @lhoestq even suggested to extend this work with other export methods as well like `csv` or `parquet`.

Problem statement:
[`to_json`] add multi-proc sharding support
As discussed on slack it appears that `to_json` is quite slow on huge datasets like OSCAR.

I implemented sharded saving, which is much much faster - but the tqdm bars all overwrite each other, so it's hard to make sense of the progress, so if possible ideally this multi-proc support could be implemented internally in `to_json` via `num_proc` argument. I guess `num_proc` will be the number of shards?

I think the user will need to use this feature wisely, since too many processes writing to say normal style HD is likely to be slower than one process.

I'm not sure whether the user should be responsible to concatenate the shards at the end  or `datasets`, either way works for my needs.

The code I was using:

```
from multiprocessing import cpu_count, Process, Queue

[...]

filtered_dataset = concat_dataset.map(filter_short_documents, batched=True, batch_size=256, num_proc=cpu_count())

DATASET_NAME = "oscar"
SHARDS = 10
def process_shard(idx):
    print(f"Sharding {idx}")
    ds_shard = filtered_dataset.shard(SHARDS, idx, contiguous=True)
    # ds_shard = ds_shard.shuffle() # remove contiguous=True above if shuffling
    print(f"Saving {DATASET_NAME}-{idx}.jsonl")
    ds_shard.to_json(f"{DATASET_NAME}-{idx}.jsonl", orient="records", lines=True, force_ascii=False)

queue = Queue()
processes = [Process(target=process_shard, args=(idx,)) for idx in range(SHARDS)]
for p in processes:
    p.start()

for p in processes:
    p.join()
```

Thank you!

@lhoestq

Hints from issue discussion:
Hi @stas00, 
I want to work on this issue and I was thinking why don't we use `imap` [in this loop](https://github.com/huggingface/datasets/blob/440b14d0dd428ae1b25881aa72ba7bbb8ad9ff84/src/datasets/io/json.py#L99)? This way, using offset (which is being used to slice the pyarrow table) we can convert  pyarrow table to `json` using multiprocessing. I've a small code snippet for some clarity:
```
result = list(
            pool.imap(self._apply_df, [(offset, batch_size) for offset in range(0, len(self.dataset), batch_size)])
        )
```
`_apply_df` is a function which will return `batch.to_pandas().to_json(path_or_buf=None, orient="records", lines=True)` which is basically json version of the batched pyarrow table. Later on we can concatenate it to form json file? 

I think the only downside here is to write file from `imap` output (output would be a list and we'll need to iterate over it and write in a file) which might add a little overhead cost. What do you think about this?
Followed up in https://github.com/huggingface/datasets/pull/2747

Commit message: flake8 datasets

Patch to review:
diff --git a/src/datasets/arrow_dataset.py b/src/datasets/arrow_dataset.py
index 86a8e75b9a6..08028d5fa5a 100644
--- a/src/datasets/arrow_dataset.py
+++ b/src/datasets/arrow_dataset.py
@@ -2869,6 +2869,7 @@ def to_json(
         self,
         path_or_buf: Union[PathLike, BinaryIO],
         batch_size: Optional[int] = None,
+        num_proc: Optional[int] = None,
         **to_json_kwargs,
     ) -> int:
         """Export the dataset to JSON Lines or JSON.
@@ -2877,6 +2878,8 @@ def to_json(
             path_or_buf (``PathLike`` or ``FileOrBuffer``): Either a path to a file or a BinaryIO.
             batch_size (:obj:`int`, optional): Size of the batch to load in memory and write at once.
                 Defaults to :obj:`datasets.config.DEFAULT_MAX_BATCH_SIZE`.
+            num_proc (:obj:`int`, optional): Number of processes for multiprocessing. By default it doesn't
+                use multiprocessing.
             lines (:obj:`bool`, default ``True``): Whether output JSON lines format.
                 Only possible if ``orient="records"`. It will throw ValueError with ``orient`` different from
                 ``"records"``, since the others are not list-like.
@@ -2897,7 +2900,7 @@ def to_json(
         # Dynamic import to avoid circular dependency
         from .io.json import JsonDatasetWriter
 
-        return JsonDatasetWriter(self, path_or_buf, batch_size=batch_size, **to_json_kwargs).write()
+        return JsonDatasetWriter(self, path_or_buf, batch_size=batch_size, num_proc=num_proc, **to_json_kwargs).write()
 
     def to_pandas(
         self, batch_size: Optional[int] = None, batched: bool = False
diff --git a/src/datasets/io/json.py b/src/datasets/io/json.py
index 377ff1d9c44..610916b70a2 100644
--- a/src/datasets/io/json.py
+++ b/src/datasets/io/json.py
@@ -1,3 +1,4 @@
+import multiprocessing
 import os
 from typing import BinaryIO, Optional, Union
 
@@ -63,30 +64,49 @@ def __init__(
         dataset: Dataset,
         path_or_buf: Union[PathLike, BinaryIO],
         batch_size: Optional[int] = None,
+        num_proc: Optional[int] = None,
         **to_json_kwargs,
     ):
         self.dataset = dataset
         self.path_or_buf = path_or_buf
-        self.batch_size = batch_size
+        if batch_size:
+            self.batch_size = batch_size
+        elif num_proc is not None:
+            self.batch_size = 100_000
+        else:
+            self.batch_size = config.DEFAULT_MAX_BATCH_SIZE
+        self.num_proc = num_proc
+        self.encoding = "utf-8"
         self.to_json_kwargs = to_json_kwargs
 
     def write(self) -> int:
-        batch_size = self.batch_size if self.batch_size else config.DEFAULT_MAX_BATCH_SIZE
+        _ = self.to_json_kwargs.pop("path_or_buf", None)
+        orient = self.to_json_kwargs.pop("orient", "records")
+        lines = self.to_json_kwargs.pop("lines", True)
 
         if isinstance(self.path_or_buf, (str, bytes, os.PathLike)):
             with open(self.path_or_buf, "wb+") as buffer:
-                written = self._write(file_obj=buffer, batch_size=batch_size, **self.to_json_kwargs)
+                written = self._write(file_obj=buffer, orient=orient, lines=lines, **self.to_json_kwargs)
         else:
-            written = self._write(file_obj=self.path_or_buf, batch_size=batch_size, **self.to_json_kwargs)
+            written = self._write(file_obj=self.path_or_buf, orient=orient, lines=lines, **self.to_json_kwargs)
         return written
 
+    def _batch_json(self, args):
+        offset, orient, lines = args
+
+        batch = query_table(
+            table=self.dataset.data,
+            key=slice(offset, offset + self.batch_size),
+            indices=self.dataset._indices,
+        )
+        var = batch.to_pandas().to_json(path_or_buf=None, orient=orient, lines=lines, **self.to_json_kwargs)
+        return var.encode(self.encoding)
+
     def _write(
         self,
         file_obj: BinaryIO,
-        batch_size: int,
-        encoding: str = "utf-8",
-        orient="records",
-        lines=True,
+        orient,
+        lines,
         **to_json_kwargs,
     ) -> int:
         """Writes the pyarrow table as JSON lines to a binary file handle.
@@ -94,18 +114,39 @@ def _write(
         Caller is responsible for opening and closing the handle.
         """
         written = 0
-        _ = to_json_kwargs.pop("path_or_buf", None)
-
-        for offset in utils.tqdm(
-            range(0, len(self.dataset), batch_size), unit="ba", disable=bool(logging.get_verbosity() == logging.NOTSET)
-        ):
-            batch = query_table(
-                table=self.dataset.data,
-                key=slice(offset, offset + batch_size),
-                indices=self.dataset._indices if self.dataset._indices is not None else None,
-            )
-            json_str = batch.to_pandas().to_json(path_or_buf=None, orient=orient, lines=lines, **to_json_kwargs)
-            if not json_str.endswith("\n"):
-                json_str += "\n"
-            written += file_obj.write(json_str.encode(encoding))
+
+        if self.num_proc is None or self.num_proc == 1:
+            for offset in utils.tqdm(
+                range(0, len(self.dataset), self.batch_size),
+                unit="ba",
+                disable=bool(logging.get_verbosity() == logging.NOTSET),
+                desc="Creating json from Arrow format",
+            ):
+                batch = query_table(
+                    table=self.dataset.data,
+                    key=slice(offset, offset + self.batch_size),
+                    indices=self.dataset._indices if self.dataset._indices is not None else None,
+                )
+                json_str = batch.to_pandas().to_json(path_or_buf=None, orient=orient, lines=lines, **to_json_kwargs)
+                if not json_str.endswith("\n"):
+                    json_str += "\n"
+                written += file_obj.write(json_str.encode(self.encoding))
+        else:
+            pool = multiprocessing.Pool(processes=self.num_proc)
+
+            for json_str in utils.tqdm(
+                pool.imap(
+                    self._batch_json,
+                    [(offset, orient, lines) for offset in range(0, len(self.dataset), self.batch_size)],
+                ),
+                total=(len(self.dataset) // self.batch_size) + 1,
+                unit="ba",
+                disable=bool(logging.get_verbosity() == logging.NOTSET),
+                desc="Creating json from Arrow format",
+            ):
+                written += file_obj.write((json_str))
+
+            pool.close()
+            pool.join()
+
         return written


Code review comments:
diff_hunk: @@ -63,49 +64,89 @@ def __init__(
         dataset: Dataset,
         path_or_buf: Union[PathLike, BinaryIO],
         batch_size: Optional[int] = None,
+        num_proc: Optional[int] = None,
         **to_json_kwargs,
     ):
         self.dataset = dataset
         self.path_or_buf = path_or_buf
-        self.batch_size = batch_size
+        if batch_size:
+            self.batch_size = batch_size
+        elif num_proc is not None:
+            self.batch_size = 100_000
+        else:
+            self.batch_size = config.DEFAULT_MAX_BATCH_SIZE
path: src/datasets/io/json.py
text: @user1:
Why not use the same default value here ?

@author:
I realised that each Process (In multiprocessing) is capable of handling a much higher batch size which makes it more faster as compared to current version. Same default value will give only 10% speed up. We can do that as well and mention in docstring to increase batch size? What do you think?

@user1:
It really depends on the data. Any batch size bigger than 100,000 can cause memory issues on small computers if the dataset has long examples.
IMO we should stick to the default `config.DEFAULT_MAX_BATCH_SIZE` even though it's not the most optimal for the text data you used. At least there is less chances to crash on processed vision/audio/time series/etc datasets x)
A note in the docstring to explain this would be very welcome !

@user1:
Hi @author , have you had a chance to consider using `config.DEFAULT_MAX_BATCH_SIZE` by default ?

@author:
Hey! Yeah, I've updated the docstring. Also, we need to have separate `test_dataset_to_json_lines` and `test_dataset_to_json_orient` for this right?
Edit:  Updated the tests as well!