Loading big dataset raises pyarrow.lib.ArrowNotImplementedError 2
Opened by kopyl (closed)

### Describe the bug

I'm aware of the issue #5695 .

I'm using a modified SDXL trainer: https://github.com/kopyl/diffusers/blob/5e70f604155aeecee254a5c63c5e4236ad4a0d3d/examples/text_to_image/train_text_to_image_sdxl.py#L1027C16-L1027C16

So i
1. Map dataset
2. Save to disk
3. Try to upload:
```
import datasets
from datasets import load_from_disk


dataset = load_from_disk("ds")
datasets.config.DEFAULT_MAX_BATCH_SIZE = 1
dataset.push_to_hub("kopyl/ds", private=True, max_shard_size="500MB")
```

And i get this error:
`pyarrow.lib.ArrowNotImplementedError: Unhandled type for Arrow to Parquet schema conversion: halffloat`

Full traceback:
```
>>> dataset.push_to_hub("kopyl/3M_icons_monochrome_only_no_captioning_mapped-for-SDXL-2", private=True, max_shard_size="500MB")
Map: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1451/1451 [00:00<00:00, 6827.40 examples/s]
Uploading the dataset shards:   0%|                                                                                                                        | 0/2099 [00:00<?, ?it/s]
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/usr/local/lib/python3.10/dist-packages/datasets/dataset_dict.py", line 1705, in push_to_hub
    split_additions, uploaded_size, dataset_nbytes = self[split]._push_parquet_shards_to_hub(
  File "/usr/local/lib/python3.10/dist-packages/datasets/arrow_dataset.py", line 5208, in _push_parquet_shards_to_hub
    shard.to_parquet(buffer)
  File "/usr/local/lib/python3.10/dist-packages/datasets/arrow_dataset.py", line 4931, in to_parquet
    return ParquetDatasetWriter(self, path_or_buf, batch_size=batch_size, **parquet_writer_kwargs).write()
  File "/usr/local/lib/python3.10/dist-packages/datasets/io/parquet.py", line 129, in write
    written = self._write(file_obj=self.path_or_buf, batch_size=batch_size, **self.parquet_writer_kwargs)
  File "/usr/local/lib/python3.10/dist-packages/datasets/io/parquet.py", line 141, in _write
    writer = pq.ParquetWriter(file_obj, schema=schema, **parquet_writer_kwargs)
  File "/usr/local/lib/python3.10/dist-packages/pyarrow/parquet/core.py", line 1016, in __init__
    self.writer = _parquet.ParquetWriter(
  File "pyarrow/_parquet.pyx", line 1869, in pyarrow._parquet.ParquetWriter.__cinit__
  File "pyarrow/error.pxi", line 154, in pyarrow.lib.pyarrow_internal_check_status
  File "pyarrow/error.pxi", line 91, in pyarrow.lib.check_status
pyarrow.lib.ArrowNotImplementedError: Unhandled type for Arrow to Parquet schema conversion: halffloat
```

Smaller datasets with the same way of saving and pushing work wonders. Big ones are not.

I'm currently trying to upload dataset like this:
`HfApi().upload_folder...`
But i'm not sure that in this case "load_dataset" would work well.

This setting num_shards does not help too:
```
dataset.push_to_hub("kopyl/3M_icons_monochrome_only_no_captioning_mapped-for-SDXL-2", private=True, num_shards={'train': 500})
```
Tried 3000, 500, 478, 100

Also do you know if it's possible to push a dataset with multiple processes? It would take an eternity pushing 1TB...

### Steps to reproduce the bug

Described above

### Expected behavior

Should be able to upload...

### Environment info

Total dataset size: 978G
Amount of `.arrow` files: 2101
Each `.arrow` file size: 477M (i know 477 megabytes * 2101 does not equal 978G, but i just checked the size of a couple `.arrow` files, i don't know if some might have different size)

Some files:
- "ds/train/state.json": https://pastebin.com/tJ3ZLGAg
- "ds/train/dataset_info.json": https://pastebin.com/JdXMQ5ih

Comments:
Hi ! I think the issue comes from the "float16" features that are not supported yet in Parquet

Feel free to open an issue in `pyarrow` about this. In the meantime, I'd encourage you to use "float32" for your "pooled_prompt_embeds" and "prompt_embeds" features.

You can cast them to "float32" using

```python
from datasets import Value

ds = ds.cast_column("pooled_prompt_embeds", Value("float32"))
ds = ds.cast_column("prompt_embeds", Value("float32"))
```

@lhoestq hm. Thank you very much.

Do you think it won't have any impact on the training? That it won't break it or the quality won't degrade because of this?

I need to use it for [SDXL training](https://github.com/huggingface/diffusers/blob/main/examples/text_to_image/train_text_to_image_sdxl.py)

Increasing the precision should not degrade training (it only increases the precision), but make sure that it doesn't break your pytorch code (e.g. if it expects a float16 instead of a float32 somewhere)

@lhoestq just fyi pyarrow 15.0.0 (just released) supports float16 as the underlying parquetcpp does as well now :)

Oh that's amazing ! (and great timing ^^)

@kopyl can you try to update `pyarrow` and try again ?

Btw @assignUser there seems to be some casting implementations missing with float16 in 15.0.0, e.g.

```
ArrowNotImplementedError: Unsupported cast from int64 to halffloat using function cast_half_float
```

```
ArrowNotImplementedError: Unsupported cast from double to halffloat using function cast_half_float
```

Ah you are right casting is not implemented yet, it's even mentioned in the docs. This pr references the relevant issues if you'd like to track them
https://github.com/apache/arrow/pull/38494

Cool thank you :)

@lhoestq i just recently found out that it's supported in 15.0.0, but wanted to try it first before telling you...

Trying this right now and it seemingly works (although i need to wait till the end to make sure there is nothing wrong). Will update you when it's finished.

<img width="918" alt="image" src="https://github.com/huggingface/datasets/assets/17604849/4821e215-e782-4736-8c76-d06187078175">

A couple of questions though:

1. What does that missing casting implementation mean for my specific case and what does it mean in general?
2. Do you know how to `push_to_hub` with multiple processes?

@lhoestq also it's strange that there was no error for a dataset with the same features, same data type, but smaller (much smaller).

Altho i'm not sure about this, but chances are the dataset was loaded directly, not `load_from_disk`.... Maybe because of this.

> What does that missing casting implementation mean for my specific case and what does it mean in general?

Nothing for you, just that casting to float16 using `.cast_column("my_column_name", Value("float16"))` raises an error

> Do you know how to push_to_hub with multiple processes?

It's not possible (yet ?). Mostly because we haven't implemented yet how to do parallel uploads to the Hub from `datasets`.
Though if you want faster uploads you can already enable `hf_transfer` 

```
pip install hf_transfer
```

and setting `HF_HUB_ENABLE_HF_TRANSFER=1` as an environment variable

see https://huggingface.co/docs/huggingface_hub/guides/upload#tips-and-tricks-for-large-uploads

@lhoestq thank you very much.

That would be amazing, I need to create a feature request for this :)

By the way, in short, how does hf_transfer improves the upload speed under the hood?

@lhoestq i was just able to successfully upload without the dataset with the new pyarrow update and without increasing the precision :)

Awesome !

Regarding hf_transfer: it's been optimized in rust ;)

@lhoestq wow, cool :)