Add Dataset.take and Dataset.skip

[lhoestq]

...to be aligned with IterableDataset.take and IterableDataset.skip

[Cadene]

LGTM, thanks!

[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.

[albertvillanova]

Thanks.

[mariosasko]

Nice!

[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.005153 / 0.011353 (-0.006200)	0.003560 / 0.011008 (-0.007448)	0.063142 / 0.038508 (0.024634)	0.030799 / 0.023109 (0.007690)	0.241754 / 0.275898 (-0.034144)	0.264874 / 0.323480 (-0.058606)	0.003099 / 0.007986 (-0.004887)	0.002629 / 0.004328 (-0.001700)	0.049006 / 0.004250 (0.044756)	0.044831 / 0.037052 (0.007779)	0.258961 / 0.258489 (0.000472)	0.286939 / 0.293841 (-0.006902)	0.026756 / 0.128546 (-0.101791)	0.010443 / 0.075646 (-0.065204)	0.207264 / 0.419271 (-0.212007)	0.035242 / 0.043533 (-0.008291)	0.250440 / 0.255139 (-0.004699)	0.265405 / 0.283200 (-0.017794)	0.018924 / 0.141683 (-0.122759)	1.138607 / 1.452155 (-0.313547)	1.203017 / 1.492716 (-0.289700)

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.091293 / 0.018006 (0.073286)	0.303937 / 0.000490 (0.303447)	0.000266 / 0.000200 (0.000066)	0.000056 / 0.000054 (0.000002)

Benchmark: benchmark_indices_mapping.json

metric	select	shard	shuffle	sort	train_test_split
new / old (diff)	0.018667 / 0.037411 (-0.018744)	0.061310 / 0.014526 (0.046784)	0.073565 / 0.176557 (-0.102991)	0.119044 / 0.737135 (-0.618091)	0.074484 / 0.296338 (-0.221854)

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.286324 / 0.215209 (0.071114)	2.836637 / 2.077655 (0.758982)	1.458531 / 1.504120 (-0.045589)	1.333081 / 1.541195 (-0.208114)	1.328398 / 1.468490 (-0.140092)	0.571467 / 4.584777 (-4.013310)	2.409869 / 3.745712 (-1.335843)	2.760241 / 5.269862 (-2.509621)	1.728153 / 4.565676 (-2.837523)	0.063008 / 0.424275 (-0.361267)	0.005375 / 0.007607 (-0.002232)	0.338574 / 0.226044 (0.112530)	3.355485 / 2.268929 (1.086556)	1.812741 / 55.444624 (-53.631884)	1.507435 / 6.876477 (-5.369041)	1.516957 / 2.142072 (-0.625116)	0.643790 / 4.805227 (-4.161437)	0.117465 / 6.500664 (-6.383199)	0.041960 / 0.075469 (-0.033509)

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.993787 / 1.841788 (-0.848001)	11.439076 / 8.074308 (3.364768)	9.636815 / 10.191392 (-0.554577)	0.131292 / 0.680424 (-0.549132)	0.014916 / 0.534201 (-0.519285)	0.287309 / 0.579283 (-0.291974)	0.261971 / 0.434364 (-0.172392)	0.324453 / 0.540337 (-0.215885)	0.420306 / 1.386936 (-0.966630)

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.005138 / 0.011353 (-0.006215)	0.003719 / 0.011008 (-0.007289)	0.050411 / 0.038508 (0.011903)	0.031334 / 0.023109 (0.008225)	0.281752 / 0.275898 (0.005854)	0.299445 / 0.323480 (-0.024035)	0.004194 / 0.007986 (-0.003792)	0.002737 / 0.004328 (-0.001591)	0.048527 / 0.004250 (0.044277)	0.040294 / 0.037052 (0.003242)	0.291763 / 0.258489 (0.033274)	0.317597 / 0.293841 (0.023757)	0.029014 / 0.128546 (-0.099532)	0.010372 / 0.075646 (-0.065274)	0.058704 / 0.419271 (-0.360568)	0.033259 / 0.043533 (-0.010273)	0.278109 / 0.255139 (0.022970)	0.299593 / 0.283200 (0.016393)	0.018048 / 0.141683 (-0.123635)	1.185558 / 1.452155 (-0.266597)	1.203481 / 1.492716 (-0.289236)

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.091149 / 0.018006 (0.073143)	0.306152 / 0.000490 (0.305662)	0.000246 / 0.000200 (0.000046)	0.000052 / 0.000054 (-0.000003)

Benchmark: benchmark_indices_mapping.json

metric	select	shard	shuffle	sort	train_test_split
new / old (diff)	0.022082 / 0.037411 (-0.015330)	0.074487 / 0.014526 (0.059961)	0.086112 / 0.176557 (-0.090444)	0.124303 / 0.737135 (-0.612832)	0.088831 / 0.296338 (-0.207508)

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.291745 / 0.215209 (0.076536)	2.878397 / 2.077655 (0.800742)	1.606920 / 1.504120 (0.102801)	1.492352 / 1.541195 (-0.048843)	1.509725 / 1.468490 (0.041235)	0.567087 / 4.584777 (-4.017690)	2.436423 / 3.745712 (-1.309290)	2.793930 / 5.269862 (-2.475932)	1.748329 / 4.565676 (-2.817347)	0.063424 / 0.424275 (-0.360851)	0.005476 / 0.007607 (-0.002131)	0.346211 / 0.226044 (0.120167)	3.461288 / 2.268929 (1.192360)	1.979362 / 55.444624 (-53.465262)	1.702877 / 6.876477 (-5.173600)	1.699087 / 2.142072 (-0.442985)	0.645116 / 4.805227 (-4.160112)	0.116186 / 6.500664 (-6.384478)	0.041246 / 0.075469 (-0.034223)

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.017540 / 1.841788 (-0.824248)	12.016640 / 8.074308 (3.942332)	10.234085 / 10.191392 (0.042693)	0.147558 / 0.680424 (-0.532866)	0.015096 / 0.534201 (-0.519105)	0.288077 / 0.579283 (-0.291206)	0.274629 / 0.434364 (-0.159735)	0.334097 / 0.540337 (-0.206241)	0.425476 / 1.386936 (-0.961460)