Title: Use pandas ujson in JSON loader to improve performance
Author: albertvillanova (MEMBER)
State: closed
Created: 2024-05-06T12:01:27Z

Use pandas ujson in JSON loader to improve performance.

Note that `datasets` has `pandas` as required dependency. And `pandas` includes `ujson` in `pd.io.json.ujson_loads`.

Fix #6867.

CC: @natolambert

Comments:
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/datasets/pr_6874). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

Before pandas-2.2.0, the function `ujson_loads` was named `loads`: https://github.com/pandas-dev/pandas/blob/v2.1.0/pandas/io/json/__init__.py#L5
```python
import ujson_loads as loads
```

Thanks for your review, @lhoestq.

The performance gain depends on many factors, such as underlying data structures, file size...

In my benchmark, the performance gain was around 8.1%. 

<details>
<summary>Show benchmarks</summary>

PyArrow==8.0.0

<details>
<summary>Show updated benchmarks!</summary>

### Benchmark: benchmark_array_xd.json

| metric | read_batch_formatted_as_numpy after write_array2d | read_batch_formatted_as_numpy after write_flattened_sequence | read_batch_formatted_as_numpy after write_nested_sequence | read_batch_unformated after write_array2d | read_batch_unformated after write_flattened_sequence | read_batch_unformated after write_nested_sequence | read_col_formatted_as_numpy after write_array2d | read_col_formatted_as_numpy after write_flattened_sequence | read_col_formatted_as_numpy after write_nested_sequence | read_col_unformated after write_array2d | read_col_unformated after write_flattened_sequence | read_col_unformated after write_nested_sequence | read_formatted_as_numpy after write_array2d | read_formatted_as_numpy after write_flattened_sequence | read_formatted_as_numpy after write_nested_sequence | read_unformated after write_array2d | read_unformated after write_flattened_sequence | read_unformated after write_nested_sequence | write_array2d | write_flattened_sequence | write_nested_sequence |
|--------|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| new / old (diff) | 0.005428 / 0.011353 (-0.005925) | 0.003682 / 0.011008 (-0.007326) | 0.064360 / 0.038508 (0.025852) | 0.032044 / 0.023109 (0.008934) | 0.238281 / 0.275898 (-0.037617) | 0.267542 / 0.323480 (-0.055937) | 0.003152 / 0.007986 (-0.004834) | 0.003292 / 0.004328 (-0.001037) | 0.050157 / 0.004250 (0.045906) | 0.048311 / 0.037052 (0.011259) | 0.253743 / 0.258489 (-0.004746) | 0.282729 / 0.293841 (-0.011112) | 0.027271 / 0.128546 (-0.101275) | 0.010238 / 0.075646 (-0.065408) | 0.208179 / 0.419271 (-0.211092) | 0.035607 / 0.043533 (-0.007925) | 0.246750 / 0.255139 (-0.008389) | 0.263362 / 0.283200 (-0.019837) | 0.018475 / 0.141683 (-0.123208) | 1.152978 / 1.452155 (-0.299177) | 1.158545 / 1.492716 (-0.334171) |

### Benchmark: benchmark_getitem\_100B.json

| metric | get_batch_of\_1024\_random_rows | get_batch_of\_1024\_rows | get_first_row | get_last_row |
|--------|---|---|---|---|
| new / old (diff) | 0.096645 / 0.018006 (0.078639) | 0.313186 / 0.000490 (0.312696) | 0.000209 / 0.000200 (0.000009) | 0.000052 / 0.000054 (-0.000003) |

### Benchmark: benchmark_indices_mapping.json

| metric | select | shard | shuffle | sort | train_test_split |
|--------|---|---|---|---|---|
| new / old (diff) | 0.018800 / 0.037411 (-0.018612) | 0.065833 / 0.014526 (0.051307) | 0.073668 / 0.176557 (-0.102888) | 0.120608 / 0.737135 (-0.616527) | 0.074936 / 0.296338 (-0.221403) |

### Benchmark: benchmark_iterating.json

| metric | read 5000 | read 50000 | read_batch 50000 10 | read_batch 50000 100 | read_batch 50000 1000 | read_formatted numpy 5000 | read_formatted pandas 5000 | read_formatted tensorflow 5000 | read_formatted torch 5000 | read_formatted_batch numpy 5000 10 | read_formatted_batch numpy 5000 1000 | shuffled read 5000 | shuffled read 50000 | shuffled read_batch 50000 10 | shuffled read_batch 50000 100 | shuffled read_batch 50000 1000 | shuffled read_formatted numpy 5000 | shuffled read_formatted_batch numpy 5000 10 | shuffled read_formatted_batch numpy 5000 1000 |
|--------|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| new / old (diff) | 0.281596 / 0.215209 (0.066387) | 2.814537 / 2.077655 (0.736882) | 1.482781 / 1.504120 (-0.021338) | 1.349770 / 1.541195 (-0.191424) | 1.371571 / 1.468490 (-0.096919) | 0.555068 / 4.584777 (-4.029709) | 2.369588 / 3.745712 (-1.376124) | 2.742771 / 5.269862 (-2.527091) | 1.711519 / 4.565676 (-2.854158) | 0.060921 / 0.424275 (-0.363354) | 0.005263 / 0.007607 (-0.002344) | 0.333721 / 0.226044 (0.107677) | 3.329598 / 2.268929 (1.060669) | 1.806983 / 55.444624 (-53.637641) | 1.515730 / 6.876477 (-5.360746) | 1.557622 / 2.142072 (-0.584451) | 0.619564 / 4.805227 (-4.185663) | 0.115503 / 6.500664 (-6.385161) | 0.041728 / 0.075469 (-0.033741) |

### Benchmark: benchmark_map_filter.json

| metric | filter | map fast-tokenizer batched | map identity | map identity batched | map no-op batched | map no-op batched numpy | map no-op batched pandas | map no-op batched pytorch | map no-op batched tensorflow |
|--------|---|---|---|---|---|---|---|---|---|
| new / old (diff) | 0.967300 / 1.841788 (-0.874487) | 11.295081 / 8.074308 (3.220773) | 9.535119 / 10.191392 (-0.656273) | 0.140232 / 0.680424 (-0.540192) | 0.013774 / 0.534201 (-0.520427) | 0.281847 / 0.579283 (-0.297436) | 0.260076 / 0.434364 (-0.174288) | 0.323657 / 0.540337 (-0.216681) | 0.421116 / 1.386936 (-0.965820) |

</details>
PyArrow==latest

<details>
<summary>Show updated benchmarks!</summary>

### Benchmark: benchmark_array_xd.json

| metric | read_batch_formatted_as_numpy after write_array2d | read_batch_formatted_as_numpy after write_flattened_sequence | read_batch_formatted_as_numpy after write_nested_sequence | read_batch_unformated after write_array2d | read_batch_unformated after write_flattened_sequence | read_batch_unformated after write_nested_sequence | read_col_formatted_as_numpy after write_array2d | read_col_formatted_as_numpy after write_flattened_sequence | read_col_formatted_as_numpy after write_nested_sequence | read_col_unformated after write_array2d | read_col_unformated after write_flattened_sequence | read_col_unformated after write_nested_sequence | read_formatted_as_numpy after write_array2d | read_formatted_as_numpy after write_flattened_sequence | read_formatted_as_numpy after write_nested_sequence | read_unformated after write_array2d | read_unformated after write_flattened_sequence | read_unformated after write_nested_sequence | write_array2d | write_flattened_sequence | write_nested_sequence |
|--------|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| new / old (diff) | 0.005276 / 0.011353 (-0.006077) | 0.003639 / 0.011008 (-0.007370) | 0.050451 / 0.038508 (0.011943) | 0.032787 / 0.023109 (0.009678) | 0.267029 / 0.275898 (-0.008869) | 0.299899 / 0.323480 (-0.023581) | 0.004177 / 0.007986 (-0.003809) | 0.002697 / 0.004328 (-0.001631) | 0.049631 / 0.004250 (0.045380) | 0.041942 / 0.037052 (0.004889) | 0.279249 / 0.258489 (0.020760) | 0.306512 / 0.293841 (0.012671) | 0.029340 / 0.128546 (-0.099207) | 0.010118 / 0.075646 (-0.065528) | 0.058243 / 0.419271 (-0.361028) | 0.033871 / 0.043533 (-0.009662) | 0.265949 / 0.255139 (0.010810) | 0.284263 / 0.283200 (0.001064) | 0.017351 / 0.141683 (-0.124332) | 1.107081 / 1.452155 (-0.345074) | 1.184946 / 1.492716 (-0.307770) |

### Benchmark: benchmark_getitem\_100B.json

| metric | get_batch_of\_1024\_random_rows | get_batch_of\_1024\_rows | get_first_row | get_last_row |
|--------|---|---|---|---|
| new / old (diff) | 0.095621 / 0.018006 (0.077614) | 0.304758 / 0.000490 (0.304269) | 0.000204 / 0.000200 (0.000004) | 0.000052 / 0.000054 (-0.000003) |

### Benchmark: benchmark_indices_mapping.json

| metric | select | shard | shuffle | sort | train_test_split |
|--------|---|---|---|---|---|
| new / old (diff) | 0.022444 / 0.037411 (-0.014967) | 0.075894 / 0.014526 (0.061368) | 0.089077 / 0.176557 (-0.087480) | 0.126960 / 0.737135 (-0.610176) | 0.089120 / 0.296338 (-0.207218) |

### Benchmark: benchmark_iterating.json

| metric | read 5000 | read 50000 | read_batch 50000 10 | read_batch 50000 100 | read_batch 50000 1000 | read_formatted numpy 5000 | read_formatted pandas 5000 | read_formatted tensorflow 5000 | read_formatted torch 5000 | read_formatted_batch numpy 5000 10 | read_formatted_batch numpy 5000 1000 | shuffled read 5000 | shuffled read 50000 | shuffled read_batch 50000 10 | shuffled read_batch 50000 100 | shuffled read_batch 50000 1000 | shuffled read_formatted numpy 5000 | shuffled read_formatted_batch numpy 5000 10 | shuffled read_formatted_batch numpy 5000 1000 |
|--------|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| new / old (diff) | 0.289885 / 0.215209 (0.074676) | 2.843219 / 2.077655 (0.765565) | 1.582704 / 1.504120 (0.078584) | 1.426551 / 1.541195 (-0.114644) | 1.431591 / 1.468490 (-0.036899) | 0.577265 / 4.584777 (-4.007512) | 0.956040 / 3.745712 (-2.789673) | 2.753517 / 5.269862 (-2.516345) | 1.732503 / 4.565676 (-2.833173) | 0.063511 / 0.424275 (-0.360764) | 0.005089 / 0.007607 (-0.002518) | 0.339205 / 0.226044 (0.113160) | 3.339148 / 2.268929 (1.070219) | 1.901543 / 55.444624 (-53.543081) | 1.618392 / 6.876477 (-5.258084) | 1.612885 / 2.142072 (-0.529188) | 0.656563 / 4.805227 (-4.148664) | 0.116740 / 6.500664 (-6.383924) | 0.040497 / 0.075469 (-0.034973) |

### Benchmark: benchmark_map_filter.json

| metric | filter | map fast-tokenizer batched | map identity | map identity batched | map no-op batched | map no-op batched numpy | map no-op batched pandas | map no-op batched pytorch | map no-op batched tensorflow |
|--------|---|---|---|---|---|---|---|---|---|
| new / old (diff) | 1.005568 / 1.841788 (-0.836219) | 11.872770 / 8.074308 (3.798462) | 9.867118 / 10.191392 (-0.324274) | 0.130193 / 0.680424 (-0.550231) | 0.022857 / 0.534201 (-0.511344) | 0.281908 / 0.579283 (-0.297375) | 0.125978 / 0.434364 (-0.308386) | 0.382604 / 0.540337 (-0.157733) | 0.415078 / 1.386936 (-0.971858) |

</details>
</details>

![](https://cml.dev/watermark.png#1eabcfaf87368a5cbfa0341aa2223f457508b3e9 "CML watermark")