Arrow dataset builder to be able to load and stream Arrow datasets
Opened by mariusz-jachimowicz-83 (closed)

This adds a Arrow dataset builder to be able to load and stream from already preprocessed Arrow files.
It's related to https://github.com/huggingface/datasets/issues/3035

Comments:
_The documentation is not available anymore as the PR was closed or merged._

@lhoestq tips applied. Thanks for a review. :smile: It's a lot of fun to improve this project. 

Let's add some documentation in a subsequent PR :)

In particular @mariosasko  and I think it's important to note to users that local arrow data are copied to cache according to the way load_dataset works, but if they want they can use Dataset.from_file instead

<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.006384 / 0.011353 (-0.004969) | 0.003788 / 0.011008 (-0.007220) | 0.098524 / 0.038508 (0.060016) | 0.031786 / 0.023109 (0.008677) | 0.307799 / 0.275898 (0.031901) | 0.337329 / 0.323480 (0.013849) | 0.003650 / 0.007986 (-0.004336) | 0.003731 / 0.004328 (-0.000598) | 0.076816 / 0.004250 (0.072566) | 0.041888 / 0.037052 (0.004835) | 0.310702 / 0.258489 (0.052213) | 0.343846 / 0.293841 (0.050005) | 0.027841 / 0.128546 (-0.100705) | 0.008312 / 0.075646 (-0.067334) | 0.320230 / 0.419271 (-0.099042) | 0.047378 / 0.043533 (0.003845) | 0.308683 / 0.255139 (0.053544) | 0.335129 / 0.283200 (0.051930) | 0.096294 / 0.141683 (-0.045389) | 1.485521 / 1.452155 (0.033366) | 1.559868 / 1.492716 (0.067152) |

### 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.197376 / 0.018006 (0.179370) | 0.430461 / 0.000490 (0.429972) | 0.004152 / 0.000200 (0.003953) | 0.000068 / 0.000054 (0.000014) |

### Benchmark: benchmark_indices_mapping.json

| metric | select | shard | shuffle | sort | train_test_split |
|--------|---|---|---|---|---|
| new / old (diff) | 0.023660 / 0.037411 (-0.013751) | 0.103128 / 0.014526 (0.088602) | 0.107549 / 0.176557 (-0.069008) | 0.175934 / 0.737135 (-0.561201) | 0.112210 / 0.296338 (-0.184129) |

### 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.415804 / 0.215209 (0.200595) | 4.216333 / 2.077655 (2.138679) | 1.910354 / 1.504120 (0.406234) | 1.712689 / 1.541195 (0.171494) | 1.754705 / 1.468490 (0.286215) | 0.554647 / 4.584777 (-4.030130) | 3.393592 / 3.745712 (-0.352120) | 1.737504 / 5.269862 (-3.532358) | 1.021213 / 4.565676 (-3.544464) | 0.066908 / 0.424275 (-0.357367) | 0.011446 / 0.007607 (0.003839) | 0.524630 / 0.226044 (0.298585) | 5.243005 / 2.268929 (2.974077) | 2.349685 / 55.444624 (-53.094939) | 2.027457 / 6.876477 (-4.849020) | 2.131053 / 2.142072 (-0.011020) | 0.669070 / 4.805227 (-4.136157) | 0.136317 / 6.500664 (-6.364347) | 0.065924 / 0.075469 (-0.009545) |

### 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.254102 / 1.841788 (-0.587686) | 13.790492 / 8.074308 (5.716184) | 14.197772 / 10.191392 (4.006380) | 0.143989 / 0.680424 (-0.536434) | 0.016577 / 0.534201 (-0.517624) | 0.375437 / 0.579283 (-0.203846) | 0.398995 / 0.434364 (-0.035369) | 0.445287 / 0.540337 (-0.095050) | 0.538632 / 1.386936 (-0.848304) |

</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.006251 / 0.011353 (-0.005101) | 0.004019 / 0.011008 (-0.006989) | 0.077985 / 0.038508 (0.039477) | 0.028705 / 0.023109 (0.005596) | 0.417360 / 0.275898 (0.141462) | 0.463964 / 0.323480 (0.140484) | 0.003489 / 0.007986 (-0.004497) | 0.003032 / 0.004328 (-0.001296) | 0.077953 / 0.004250 (0.073702) | 0.040104 / 0.037052 (0.003051) | 0.405242 / 0.258489 (0.146753) | 0.475029 / 0.293841 (0.181188) | 0.028113 / 0.128546 (-0.100433) | 0.008610 / 0.075646 (-0.067036) | 0.084847 / 0.419271 (-0.334424) | 0.048227 / 0.043533 (0.004694) | 0.417235 / 0.255139 (0.162096) | 0.450470 / 0.283200 (0.167270) | 0.096978 / 0.141683 (-0.044705) | 1.514688 / 1.452155 (0.062533) | 1.560205 / 1.492716 (0.067488) |

### 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.235125 / 0.018006 (0.217119) | 0.409904 / 0.000490 (0.409414) | 0.002474 / 0.000200 (0.002275) | 0.000074 / 0.000054 (0.000020) |

### Benchmark: benchmark_indices_mapping.json

| metric | select | shard | shuffle | sort | train_test_split |
|--------|---|---|---|---|---|
| new / old (diff) | 0.025152 / 0.037411 (-0.012259) | 0.103517 / 0.014526 (0.088991) | 0.110154 / 0.176557 (-0.066402) | 0.161431 / 0.737135 (-0.575704) | 0.114891 / 0.296338 (-0.181448) |

### 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.456077 / 0.215209 (0.240868) | 4.541171 / 2.077655 (2.463517) | 2.297912 / 1.504120 (0.793792) | 2.079337 / 1.541195 (0.538143) | 2.121291 / 1.468490 (0.652801) | 0.560172 / 4.584777 (-4.024605) | 3.421122 / 3.745712 (-0.324590) | 1.764675 / 5.269862 (-3.505186) | 1.043482 / 4.565676 (-3.522195) | 0.067652 / 0.424275 (-0.356623) | 0.011181 / 0.007607 (0.003574) | 0.557232 / 0.226044 (0.331188) | 5.607851 / 2.268929 (3.338922) | 2.783715 / 55.444624 (-52.660909) | 2.380943 / 6.876477 (-4.495534) | 2.378316 / 2.142072 (0.236244) | 0.674356 / 4.805227 (-4.130871) | 0.135912 / 6.500664 (-6.364752) | 0.067009 / 0.075469 (-0.008460) |

### 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.309002 / 1.841788 (-0.532786) | 14.464073 / 8.074308 (6.389765) | 14.418727 / 10.191392 (4.227335) | 0.148486 / 0.680424 (-0.531938) | 0.016650 / 0.534201 (-0.517551) | 0.368786 / 0.579283 (-0.210497) | 0.395026 / 0.434364 (-0.039338) | 0.433565 / 0.540337 (-0.106772) | 0.526603 / 1.386936 (-0.860333) |

</details>
</details>

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