Improve skip take shuffling and distributed

[lhoestq]

set the right behavior of skip/take depending on whether it's called after or before shuffle/split_by_node

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[tianyu-l]

Thanks for implementing this so fast! Do you plan to land it soon?

[tianyu-l]

this looks good to me!

[tianyu-l]

I don't know the implementation of split_dataset_by_node very well, just trying to raise a concern on this test case and please correct me if I'm wrong.

Assume world_size = 2, count = 1. Assume underlying data is [0, 1, 2, ..., 9].
true_distributed_dataset on rank 0 would be [0, 2, 4, 6, 8]
true_distributed_dataset on rank 1 would be [1, 3, 5, 7, 9]
and after calling skip and split_dataset_by_node
distributed_dataset on rank 0 would be [1, 3, 5, 7, 9]
distributed_dataset on rank 1 would be [2, 4, 6, 8]

[lhoestq]

Yes that's correct. The test doesn't ensure that some examples are not skipped if count is not a factor of world_size at the moment, bu we can improve that later

[github-actions]

Show benchmarks

PyArrow==8.0.0

Show updated benchmarks!

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.005879 / 0.011353 (-0.005474)	0.004144 / 0.011008 (-0.006865)	0.063327 / 0.038508 (0.024819)	0.032577 / 0.023109 (0.009468)	0.242936 / 0.275898 (-0.032962)	0.269882 / 0.323480 (-0.053598)	0.003339 / 0.007986 (-0.004647)	0.002901 / 0.004328 (-0.001428)	0.049163 / 0.004250 (0.044912)	0.047072 / 0.037052 (0.010019)	0.261120 / 0.258489 (0.002631)	0.287857 / 0.293841 (-0.005984)	0.029688 / 0.128546 (-0.098858)	0.012702 / 0.075646 (-0.062944)	0.204040 / 0.419271 (-0.215231)	0.036012 / 0.043533 (-0.007521)	0.244210 / 0.255139 (-0.010929)	0.267600 / 0.283200 (-0.015599)	0.019627 / 0.141683 (-0.122056)	1.103770 / 1.452155 (-0.348385)	1.197710 / 1.492716 (-0.295006)

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.101683 / 0.018006 (0.083677)	0.311825 / 0.000490 (0.311335)	0.000236 / 0.000200 (0.000036)	0.000043 / 0.000054 (-0.000011)

Benchmark: benchmark_indices_mapping.json

metric	select	shard	shuffle	sort	train_test_split
new / old (diff)	0.019642 / 0.037411 (-0.017769)	0.061618 / 0.014526 (0.047092)	0.075237 / 0.176557 (-0.101320)	0.122250 / 0.737135 (-0.614886)	0.076087 / 0.296338 (-0.220251)

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.285120 / 0.215209 (0.069911)	2.811527 / 2.077655 (0.733872)	1.457961 / 1.504120 (-0.046159)	1.333819 / 1.541195 (-0.207376)	1.387863 / 1.468490 (-0.080627)	0.730828 / 4.584777 (-3.853949)	2.417224 / 3.745712 (-1.328488)	2.994842 / 5.269862 (-2.275020)	1.922079 / 4.565676 (-2.643598)	0.087486 / 0.424275 (-0.336789)	0.005211 / 0.007607 (-0.002396)	0.335585 / 0.226044 (0.109541)	3.297664 / 2.268929 (1.028735)	1.809391 / 55.444624 (-53.635233)	1.501646 / 6.876477 (-5.374831)	1.567573 / 2.142072 (-0.574500)	0.800816 / 4.805227 (-4.004411)	0.134204 / 6.500664 (-6.366460)	0.043156 / 0.075469 (-0.032313)

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.982955 / 1.841788 (-0.858833)	12.256850 / 8.074308 (4.182542)	9.821500 / 10.191392 (-0.369892)	0.143739 / 0.680424 (-0.536685)	0.014425 / 0.534201 (-0.519776)	0.302718 / 0.579283 (-0.276565)	0.267746 / 0.434364 (-0.166618)	0.340036 / 0.540337 (-0.200301)	0.436211 / 1.386936 (-0.950725)

PyArrow==latest

Show updated benchmarks!

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.006136 / 0.011353 (-0.005217)	0.004125 / 0.011008 (-0.006883)	0.050341 / 0.038508 (0.011833)	0.034547 / 0.023109 (0.011438)	0.270237 / 0.275898 (-0.005661)	0.294503 / 0.323480 (-0.028977)	0.004528 / 0.007986 (-0.003458)	0.003103 / 0.004328 (-0.001225)	0.048817 / 0.004250 (0.044566)	0.041301 / 0.037052 (0.004249)	0.279461 / 0.258489 (0.020972)	0.319376 / 0.293841 (0.025535)	0.032733 / 0.128546 (-0.095813)	0.012426 / 0.075646 (-0.063221)	0.060543 / 0.419271 (-0.358729)	0.034015 / 0.043533 (-0.009518)	0.267387 / 0.255139 (0.012248)	0.288590 / 0.283200 (0.005390)	0.019697 / 0.141683 (-0.121986)	1.145994 / 1.452155 (-0.306161)	1.198122 / 1.492716 (-0.294595)

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.099091 / 0.018006 (0.081085)	0.313767 / 0.000490 (0.313277)	0.000220 / 0.000200 (0.000020)	0.000054 / 0.000054 (-0.000001)

Benchmark: benchmark_indices_mapping.json

metric	select	shard	shuffle	sort	train_test_split
new / old (diff)	0.023219 / 0.037411 (-0.014192)	0.083609 / 0.014526 (0.069084)	0.089529 / 0.176557 (-0.087028)	0.131025 / 0.737135 (-0.606110)	0.091947 / 0.296338 (-0.204391)

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.283711 / 0.215209 (0.068502)	2.811702 / 2.077655 (0.734047)	1.577720 / 1.504120 (0.073600)	1.415700 / 1.541195 (-0.125495)	1.436097 / 1.468490 (-0.032393)	0.732090 / 4.584777 (-3.852687)	0.990552 / 3.745712 (-2.755160)	2.887319 / 5.269862 (-2.382543)	1.923707 / 4.565676 (-2.641969)	0.079361 / 0.424275 (-0.344915)	0.005520 / 0.007607 (-0.002087)	0.336684 / 0.226044 (0.110639)	3.325342 / 2.268929 (1.056413)	1.911853 / 55.444624 (-53.532771)	1.621450 / 6.876477 (-5.255027)	1.807964 / 2.142072 (-0.334109)	0.813958 / 4.805227 (-3.991269)	0.137564 / 6.500664 (-6.363100)	0.043151 / 0.075469 (-0.032318)

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.002775 / 1.841788 (-0.839013)	12.526367 / 8.074308 (4.452058)	10.426992 / 10.191392 (0.235600)	0.134902 / 0.680424 (-0.545522)	0.016726 / 0.534201 (-0.517475)	0.303549 / 0.579283 (-0.275734)	0.129334 / 0.434364 (-0.305030)	0.339254 / 0.540337 (-0.201084)	0.456845 / 1.386936 (-0.930091)