Title: Fix string-encoding, make `batch_size` optional, and minor improvements in `Dataset.to_tf_dataset`
Author: alvarobartt (MEMBER)
State: closed
Created: 2023-05-22T11:51:07Z

## What's in this PR?

This PR addresses some minor fixes and general improvements in the `to_tf_dataset` method of `datasets.Dataset`, to convert a 🤗HuggingFace Dataset as a TensorFlow Dataset.

The main bug solved in this PR comes with the string-encoding, since for safety purposes the internal conversion of `numpy.arrays` when `dtype` is unicode/string, is to convert it into `numpy.bytes`, more information in the docstring of https://github.com/tensorflow/tensorflow/blob/388d952114e59a1aeda440ed4737b29f8b7c6e8a/tensorflow/python/ops/script_ops.py#L210. That's triggered when using `tensorflow.numpy_function` as it's applying another type cast besides the one that `datasets` does, so the casting is applied at least twice per entry/batch. So this means that the definition of the `numpy.unicode_` dtype when the data in the batch is a string, is ignored, and replaced by `numpy.bytes_`.

Besides that, some other minor things have been fixed:

* Made `batch_size` an optional parameter in `to_tf_dataset`
* Map the `tensorflow` output dtypes just once, and not in every `tf.function` call during `map`
* Keep `numpy` formatting in the `datasets.Dataset` if already formatted like it, no need to format it again as `numpy`
* Docstring indentation in `dataset_to_tf` and `multiprocess_dataset_to_tf`

## What's missing in this PR?

I can include some integration tests if needed, to validate that `batch_size` is optional, and that the tensors in the TF-Dataset can be looped over with no issues as before.

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

To showcase the current issue, here's a Colab Gist, that shows that the `imdb` dataset cannot be read/iterated, since one or more samples contain a non-ascii character that is being converted to `numpy.bytes_`, and so on fails.

Colab Gist at https://gist.github.com/alvarobartt/1746959d1abb9a33e0c593f3bd82a2fb

Also, here's a quick sample of what's happening:

```python
from datasets import load_dataset

ds = load_dataset("imdb", split="train")
tfds = ds.to_tf_dataset(batch_size=16)
for batch in tfds:
  print(batch)
>>> UnicodeEncodeError: 'ascii' codec can't encode character '\xe9' in position 0: ordinal not in range(128)
```

A more detailed version of it:

```python
from datasets import Dataset

ds = Dataset.from_dict(
    {
        "a": [1],
        "b": ["é"],
    }
)
tfds = ds.to_tf_dataset(batch_size=1)
for batch in tfds:
    print(batch)
>>> UnicodeEncodeError: 'ascii' codec can't encode character '\xe9' in position 0: ordinal not in range(128)
```

The original issue comes from https://github.com/tensorflow/tensorflow/blob/388d952114e59a1aeda440ed4737b29f8b7c6e8a/tensorflow/python/ops/script_ops.py#LL234C4-L234C4, which could easily be solved by replacing that line with `return result.astype(np.unicode_)` but they are mentioning that it may lead to issues.

Even the following fails in `numpy`:

```python
import numpy as np

x = np.array(["é"]).astype(np.bytes_)
```

cc. @lhoestq :hugs:

cc @Rocketknight1 

> Nice ! Could you add some tests to make sure that batch_size=None works as expected ?

Sure, I'll add the tests for everything, including the string-encoding issue to make sure it's solved!

Thanks for the review @lhoestq and @Rocketknight1! I do understand that processing it in batches is always more efficient than processing it one-by-one, it was just to make `batch_size` optional. What we can do is default it to a certain batch size e.g. 16 as before, and that's it, but I think it can still remain optional.

@Rocketknight1 then I'll add the integration tests for the optional `batch_size` as well as for the encoding of non-ASCII compatible characters 😄 Do we set the default `batch_size` to 16 instead of `None`?

@alvarobartt  I think 16 is a reasonable default, yep!

I think default should be None, not 16.
Users won't expect to have it batched by default.

Then I'll leave it as is, and add the unit/integration tests, thanks @Rocketknight1 and @lhoestq 

Hi @Rocketknight1 @lhoestq! So the string-encoding issue is already solved, but I've got one doubt about the `batch_size` being optional in the multiprocessing approach, since in that case I assume the `batch_size` should be mandatory, for the moment I'm assuming it is/should be mandatory, but let me know if you want me to add a check to disallow `batch_size=None` when `num_workers>1`. Thanks!

> To showcase the current issue, here's a Colab Gist, that shows that the `imdb` dataset cannot be read/iterated, since one or more samples contain a non-ascii character that is being converted to `numpy.bytes_`, and so on fails.
> 
> Colab Gist at https://gist.github.com/alvarobartt/1746959d1abb9a33e0c593f3bd82a2fb

I've used the Colab shared above for testing purposes, and it works fine, plus the unit/integration tests are passing. I've also trained a `KerasNLP` model with incoming data from 🤗`datasets` with no issue at all!

> in the multiprocessing approach, since in that case I assume the batch_size should be mandatory,

No I think they're quite orthogonal, no need to have it mandatory

> No I think they're quite orthogonal, no need to have it mandatory

But it will break if `batch_size=None` as the multiprocessing approach will aim to prepare batches and distribute those to every worker, and assuming `batch_size=1` when `batch_size=None` I guess is not a good assumption, right?

Ah I see. Multiprocessing should support batch_size=None indeed. If you have ideas you can do it in this PR, or raise a NotImplementedError and we can see later

Sure @lhoestq, I can add a `NotImplementedError` for the moment, and prepare the next PR straight-away to tackle the multiprocessing approach with `batch_size=None`, but not sure if that may eventually collide with @Rocketknight1 PR at https://github.com/huggingface/datasets/pull/5863

Yes, let me merge the PR at #5863 after this one, and then we can open another to improve the behaviour with multiprocessing and `batch_size=None`!

Sure @Rocketknight1 makes complete sense to me! Do you want me to add the `raise NotImplementedError` and then we merge this PR? Or you prefer to directly merge the current?

`raise NotImplementedError` for now with an error telling the user that multiprocessing needs them to specify a batch size, I think!

Since you recently approved @Rocketknight1, are we ready to merge? Thanks 🤗

Ah actually it looks like `minimal_tf_collate_fn` doesn't support batch_size=None

Hi @lhoestq so I didn't include the call to `collate_fn`, as we won't need to collate the incoming data e.g. "str" should remain a "str" not a ["str"], and the `minimal_collate_fn` was indeed putting everything into a list, so the output was not un-batched, but batched with size 1

What if the user passes a collate_fn ? The torch DataLoader still applies it if batch_size=None for example.

Does my last change look of to you ? If so I think we can merge

> What if the user passes a collate_fn ? The torch DataLoader still applies it if batch_size=None for example.
> 
> Does my last change look of to you ? If so I think we can merge

I think we're good, since it won't batch it under the scenario of `str` being provided instead of `List[str]`, and the unit/integration tests are passing, so I'm OK to merge. Maybe we can double check with Matt? cc @Rocketknight1 

Yes, and sorry for the delay! I'm happy to merge.

<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.006555 / 0.011353 (-0.004798) | 0.004521 / 0.011008 (-0.006487) | 0.096633 / 0.038508 (0.058125) | 0.032859 / 0.023109 (0.009750) | 0.294632 / 0.275898 (0.018734) | 0.325140 / 0.323480 (0.001660) | 0.005676 / 0.007986 (-0.002310) | 0.005252 / 0.004328 (0.000924) | 0.074349 / 0.004250 (0.070099) | 0.045836 / 0.037052 (0.008784) | 0.302919 / 0.258489 (0.044430) | 0.340686 / 0.293841 (0.046845) | 0.028398 / 0.128546 (-0.100148) | 0.008942 / 0.075646 (-0.066704) | 0.326994 / 0.419271 (-0.092278) | 0.049556 / 0.043533 (0.006023) | 0.293883 / 0.255139 (0.038744) | 0.316522 / 0.283200 (0.033322) | 0.097385 / 0.141683 (-0.044298) | 1.405334 / 1.452155 (-0.046821) | 1.521529 / 1.492716 (0.028812) |

### 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.212269 / 0.018006 (0.194263) | 0.445692 / 0.000490 (0.445203) | 0.004930 / 0.000200 (0.004730) | 0.000093 / 0.000054 (0.000039) |

### Benchmark: benchmark_indices_mapping.json

| metric | select | shard | shuffle | sort | train_test_split |
|--------|---|---|---|---|---|
| new / old (diff) | 0.026907 / 0.037411 (-0.010504) | 0.108607 / 0.014526 (0.094081) | 0.116806 / 0.176557 (-0.059751) | 0.178428 / 0.737135 (-0.558707) | 0.122326 / 0.296338 (-0.174012) |

### 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.404211 / 0.215209 (0.189002) | 4.045374 / 2.077655 (1.967719) | 1.877237 / 1.504120 (0.373117) | 1.706276 / 1.541195 (0.165081) | 1.750610 / 1.468490 (0.282120) | 0.522331 / 4.584777 (-4.062446) | 3.742286 / 3.745712 (-0.003426) | 1.791285 / 5.269862 (-3.478577) | 1.043872 / 4.565676 (-3.521805) | 0.065176 / 0.424275 (-0.359099) | 0.011821 / 0.007607 (0.004214) | 0.507374 / 0.226044 (0.281329) | 5.088803 / 2.268929 (2.819875) | 2.282742 / 55.444624 (-53.161882) | 1.950737 / 6.876477 (-4.925740) | 2.042262 / 2.142072 (-0.099810) | 0.636525 / 4.805227 (-4.168702) | 0.140837 / 6.500664 (-6.359827) | 0.063223 / 0.075469 (-0.012246) |

### 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.188070 / 1.841788 (-0.653718) | 14.622681 / 8.074308 (6.548372) | 13.247988 / 10.191392 (3.056596) | 0.165858 / 0.680424 (-0.514566) | 0.017476 / 0.534201 (-0.516725) | 0.391973 / 0.579283 (-0.187310) | 0.433326 / 0.434364 (-0.001038) | 0.467163 / 0.540337 (-0.073175) | 0.568359 / 1.386936 (-0.818577) |

</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.006076 / 0.011353 (-0.005276) | 0.004439 / 0.011008 (-0.006570) | 0.074496 / 0.038508 (0.035988) | 0.031396 / 0.023109 (0.008287) | 0.372237 / 0.275898 (0.096339) | 0.403412 / 0.323480 (0.079932) | 0.005430 / 0.007986 (-0.002555) | 0.003846 / 0.004328 (-0.000483) | 0.074403 / 0.004250 (0.070153) | 0.045398 / 0.037052 (0.008346) | 0.394133 / 0.258489 (0.135644) | 0.421769 / 0.293841 (0.127928) | 0.027936 / 0.128546 (-0.100610) | 0.008962 / 0.075646 (-0.066685) | 0.083158 / 0.419271 (-0.336113) | 0.044863 / 0.043533 (0.001331) | 0.393834 / 0.255139 (0.138695) | 0.391537 / 0.283200 (0.108337) | 0.097971 / 0.141683 (-0.043712) | 1.496632 / 1.452155 (0.044477) | 1.585511 / 1.492716 (0.092795) |

### 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.010094 / 0.018006 (-0.007913) | 0.437811 / 0.000490 (0.437321) | 0.000963 / 0.000200 (0.000763) | 0.000084 / 0.000054 (0.000029) |

### Benchmark: benchmark_indices_mapping.json

| metric | select | shard | shuffle | sort | train_test_split |
|--------|---|---|---|---|---|
| new / old (diff) | 0.028864 / 0.037411 (-0.008547) | 0.112480 / 0.014526 (0.097954) | 0.120938 / 0.176557 (-0.055619) | 0.170888 / 0.737135 (-0.566247) | 0.125903 / 0.296338 (-0.170435) |

### 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.426716 / 0.215209 (0.211507) | 4.238380 / 2.077655 (2.160725) | 2.052889 / 1.504120 (0.548769) | 1.871043 / 1.541195 (0.329848) | 1.890405 / 1.468490 (0.421915) | 0.522059 / 4.584777 (-4.062718) | 3.813331 / 3.745712 (0.067619) | 2.891651 / 5.269862 (-2.378210) | 1.323836 / 4.565676 (-3.241841) | 0.065124 / 0.424275 (-0.359151) | 0.011498 / 0.007607 (0.003891) | 0.525102 / 0.226044 (0.299057) | 5.245190 / 2.268929 (2.976261) | 2.531149 / 55.444624 (-52.913476) | 2.197323 / 6.876477 (-4.679153) | 2.197314 / 2.142072 (0.055241) | 0.633423 / 4.805227 (-4.171804) | 0.140248 / 6.500664 (-6.360416) | 0.064432 / 0.075469 (-0.011037) |

### 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.270639 / 1.841788 (-0.571149) | 14.856678 / 8.074308 (6.782369) | 14.337631 / 10.191392 (4.146239) | 0.195319 / 0.680424 (-0.485105) | 0.017628 / 0.534201 (-0.516573) | 0.393984 / 0.579283 (-0.185299) | 0.421987 / 0.434364 (-0.012376) | 0.459245 / 0.540337 (-0.081092) | 0.557786 / 1.386936 (-0.829150) |

</details>
</details>

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


Will you eventually need help with your PR @Rocketknight1? I'll be happy to help if needed 😄 

<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.007577 / 0.011353 (-0.003776) | 0.004960 / 0.011008 (-0.006048) | 0.113622 / 0.038508 (0.075114) | 0.037981 / 0.023109 (0.014872) | 0.355312 / 0.275898 (0.079414) | 0.393384 / 0.323480 (0.069904) | 0.006575 / 0.007986 (-0.001411) | 0.005941 / 0.004328 (0.001612) | 0.085976 / 0.004250 (0.081726) | 0.053784 / 0.037052 (0.016732) | 0.369358 / 0.258489 (0.110869) | 0.399402 / 0.293841 (0.105561) | 0.032155 / 0.128546 (-0.096391) | 0.010448 / 0.075646 (-0.065199) | 0.389009 / 0.419271 (-0.030263) | 0.057377 / 0.043533 (0.013844) | 0.354968 / 0.255139 (0.099829) | 0.382404 / 0.283200 (0.099204) | 0.111056 / 0.141683 (-0.030627) | 1.807986 / 1.452155 (0.355832) | 1.866070 / 1.492716 (0.373354) |

### 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.244439 / 0.018006 (0.226432) | 0.491942 / 0.000490 (0.491452) | 0.001910 / 0.000200 (0.001710) | 0.000112 / 0.000054 (0.000058) |

### Benchmark: benchmark_indices_mapping.json

| metric | select | shard | shuffle | sort | train_test_split |
|--------|---|---|---|---|---|
| new / old (diff) | 0.031024 / 0.037411 (-0.006387) | 0.129674 / 0.014526 (0.115148) | 0.142974 / 0.176557 (-0.033583) | 0.213568 / 0.737135 (-0.523568) | 0.147794 / 0.296338 (-0.148545) |

### 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.480333 / 0.215209 (0.265124) | 4.792901 / 2.077655 (2.715246) | 2.233145 / 1.504120 (0.729025) | 2.036291 / 1.541195 (0.495096) | 2.109631 / 1.468490 (0.641140) | 0.624546 / 4.584777 (-3.960231) | 4.543511 / 3.745712 (0.797799) | 3.961345 / 5.269862 (-1.308517) | 1.903634 / 4.565676 (-2.662042) | 0.076584 / 0.424275 (-0.347691) | 0.014590 / 0.007607 (0.006983) | 0.593195 / 0.226044 (0.367151) | 5.928740 / 2.268929 (3.659811) | 2.781164 / 55.444624 (-52.663460) | 2.364303 / 6.876477 (-4.512173) | 2.510139 / 2.142072 (0.368067) | 0.770886 / 4.805227 (-4.034341) | 0.167995 / 6.500664 (-6.332669) | 0.076622 / 0.075469 (0.001153) |

### 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.402398 / 1.841788 (-0.439390) | 17.921233 / 8.074308 (9.846925) | 17.036738 / 10.191392 (6.845346) | 0.168997 / 0.680424 (-0.511427) | 0.020259 / 0.534201 (-0.513941) | 0.465322 / 0.579283 (-0.113962) | 0.500435 / 0.434364 (0.066071) | 0.546846 / 0.540337 (0.006509) | 0.658130 / 1.386936 (-0.728806) |

</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.007624 / 0.011353 (-0.003729) | 0.005265 / 0.011008 (-0.005744) | 0.086886 / 0.038508 (0.048377) | 0.038235 / 0.023109 (0.015126) | 0.463969 / 0.275898 (0.188071) | 0.502451 / 0.323480 (0.178971) | 0.006285 / 0.007986 (-0.001701) | 0.004525 / 0.004328 (0.000197) | 0.086557 / 0.004250 (0.082307) | 0.052414 / 0.037052 (0.015362) | 0.482167 / 0.258489 (0.223678) | 0.513684 / 0.293841 (0.219843) | 0.032929 / 0.128546 (-0.095618) | 0.010249 / 0.075646 (-0.065397) | 0.093377 / 0.419271 (-0.325895) | 0.054114 / 0.043533 (0.010582) | 0.466116 / 0.255139 (0.210977) | 0.488977 / 0.283200 (0.205777) | 0.115446 / 0.141683 (-0.026237) | 1.762912 / 1.452155 (0.310757) | 1.874191 / 1.492716 (0.381475) |

### 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.012666 / 0.018006 (-0.005341) | 0.485962 / 0.000490 (0.485473) | 0.002621 / 0.000200 (0.002421) | 0.000128 / 0.000054 (0.000074) |

### Benchmark: benchmark_indices_mapping.json

| metric | select | shard | shuffle | sort | train_test_split |
|--------|---|---|---|---|---|
| new / old (diff) | 0.033661 / 0.037411 (-0.003751) | 0.135395 / 0.014526 (0.120869) | 0.147230 / 0.176557 (-0.029326) | 0.205847 / 0.737135 (-0.531288) | 0.151496 / 0.296338 (-0.144842) |

### 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.514097 / 0.215209 (0.298887) | 5.134093 / 2.077655 (3.056438) | 2.496775 / 1.504120 (0.992655) | 2.268078 / 1.541195 (0.726883) | 2.342153 / 1.468490 (0.873663) | 0.623130 / 4.584777 (-3.961647) | 4.601787 / 3.745712 (0.856075) | 3.414249 / 5.269862 (-1.855613) | 1.849603 / 4.565676 (-2.716073) | 0.078350 / 0.424275 (-0.345925) | 0.013785 / 0.007607 (0.006178) | 0.638783 / 0.226044 (0.412739) | 6.378356 / 2.268929 (4.109427) | 3.072867 / 55.444624 (-52.371757) | 2.668123 / 6.876477 (-4.208354) | 2.693905 / 2.142072 (0.551833) | 0.764583 / 4.805227 (-4.040644) | 0.166854 / 6.500664 (-6.333810) | 0.076883 / 0.075469 (0.001414) |

### 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.502003 / 1.841788 (-0.339784) | 18.674205 / 8.074308 (10.599897) | 16.837759 / 10.191392 (6.646367) | 0.176995 / 0.680424 (-0.503428) | 0.020126 / 0.534201 (-0.514075) | 0.464480 / 0.579283 (-0.114803) | 0.516477 / 0.434364 (0.082113) | 0.549818 / 0.540337 (0.009481) | 0.659927 / 1.386936 (-0.727009) |

</details>
</details>

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


@alvarobartt Yes, I'll ping you for a review once it's ready!