Datasets: Caching processed dataset at wrong folder

[mrm8488]

Hi guys, I run this on my Colab (PRO):

from datasets import load_dataset
dataset = load_dataset('text', data_files='/content/corpus.txt', cache_dir='/content/drive/My Drive', split='train')

def encode(examples):
  return tokenizer(examples['text'], truncation=True, padding='max_length')

dataset = dataset.map(encode, batched=True)

The file is about 4 GB, so I cannot process it on the Colab HD because there is no enough space. So I decided to mount my Google Drive fs and do it on it.
The dataset is cached in the right place but by processing it (applying encode function) seems to use a different folder because Colab HD starts to grow and it crashes when it should be done in the Drive fs.

What gets me crazy, it prints it is processing/encoding the dataset in the right folder:

Testing the mapped function outputs
Testing finished, running the mapping function on the dataset
Caching processed dataset at /content/drive/My Drive/text/default-ad3e69d6242ee916/0.0.0/7e13bc0fa76783d4ef197f079dc8acfe54c3efda980f2c9adfab046ede2f0ff7/cache-b16341780a59747d.arrow

[lhoestq]

Thanks for reporting !
It uses a temporary file to write the data.
However it looks like the temporary file is not placed in the right directory during the processing

[lhoestq]

Well actually I just tested and the temporary file is placed in the same directory, so it should work as expected.
Which version of datasets are you using ?

[mrm8488]

datasets-1.0.1
Here you can reproduce it here:
https://colab.research.google.com/drive/1O0KcepTFsmpkBbrbLLMq42iwTKmQh8d5?usp=sharing

[lhoestq]

It looks like a pyarrow issue with google colab.
For some reason this code increases the disk usage of google colab while it actually writes into google drive:

import pyarrow as pa

stream = pa.OSFile("/content/drive/My Drive/path/to/file.arrow", "wb")
writer = pa.RecordBatchStreamWriter(stream, schema=pa.schema({"text": pa.string()}))
writer.write_table(pa.Table.from_pydict({"text": ["a"*511 + "\n"] * ((1 << 30) // 512)}))  # 1GiB
writer.close()
stream.close()

Moreover if I rm the file on google drive, it frees disk space on google colab.

[lhoestq]

It looks like replacing pa.OSFile by open fixes it, I'm going to open a PR

[mrm8488]

Ok. Thank you so much!

[lhoestq]

Actually I did more tests it doesn't >.<
I'll let you know if I find a way to fix that

[lhoestq]

Actually I also have the issue when writing a regular text file

f = open("/content/drive/My Drive/path/to/file", "w")
f.write(("a"*511 + "\n") * ((1 << 30) // 512))  # 1GiB
f.close()

Is that supposed to happen ?

[mrm8488]

The code you wrote should write a 1GB file in the Google Drive folder. Doesn't it?

[lhoestq]

Yes it does, but the disk usage of google colab also increases by 1GB

[mrm8488]

I could check it and as you say as I write to te Drive disk the colab disk also increases...

[mrm8488]

To reproduce it:

!df -h | grep sda1

```python
f = open("/content/drive/My Drive/test_to_remove.txt", "w")
f.write(("a"511 + "\n") * ((1 << 30) // 512)) # 1GiB
f.write(("a"511 + "\n") * ((1 << 30) // 512)) # 1GiB
f.close()

```bash
!ls -lh /content/drive/My\ Drive/test_to_remove.txt

!df -h | grep sda1

!rm -rf /content/drive/My\ Drive/test_to_remove.txt

Colab