laion2B-en failed to load on Windows with PrefetchVirtualMemory failed
Opened by doctorpangloss (open)

### Describe the bug

I have downloaded laion2B-en, and I'm receiving the following error trying to load it:
```
Resolving data files: 100%|██████████| 128/128 [00:00<00:00, 1173.79it/s]
Traceback (most recent call last):
  File "D:\Art-Workspace\src\artworkspace\tokeneval\compute_frequencies.py", line 31, in <module>
    count = compute_frequencies()
            ^^^^^^^^^^^^^^^^^^^^^
  File "D:\Art-Workspace\src\artworkspace\tokeneval\compute_frequencies.py", line 17, in compute_frequencies
    laion2b_dataset = load_dataset("laion/laion2B-en", split="train", cache_dir=_CACHE_DIR, keep_in_memory=False)
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\bberman\Documents\Art-Workspace\venv\Lib\site-packages\datasets\load.py", line 2165, in load_dataset
    ds = builder_instance.as_dataset(split=split, verification_mode=verification_mode, in_memory=keep_in_memory)
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\bberman\Documents\Art-Workspace\venv\Lib\site-packages\datasets\builder.py", line 1187, in as_dataset
    datasets = map_nested(
               ^^^^^^^^^^^
  File "C:\Users\bberman\Documents\Art-Workspace\venv\Lib\site-packages\datasets\utils\py_utils.py", line 456, in map_nested
    return function(data_struct)
           ^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\bberman\Documents\Art-Workspace\venv\Lib\site-packages\datasets\builder.py", line 1217, in _build_single_dataset
    ds = self._as_dataset(
         ^^^^^^^^^^^^^^^^^
  File "C:\Users\bberman\Documents\Art-Workspace\venv\Lib\site-packages\datasets\builder.py", line 1291, in _as_dataset
    dataset_kwargs = ArrowReader(cache_dir, self.info).read(
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\bberman\Documents\Art-Workspace\venv\Lib\site-packages\datasets\arrow_reader.py", line 244, in read
    return self.read_files(files=files, original_instructions=instructions, in_memory=in_memory)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\bberman\Documents\Art-Workspace\venv\Lib\site-packages\datasets\arrow_reader.py", line 265, in read_files
    pa_table = self._read_files(files, in_memory=in_memory)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\bberman\Documents\Art-Workspace\venv\Lib\site-packages\datasets\arrow_reader.py", line 200, in _read_files
    pa_table: Table = self._get_table_from_filename(f_dict, in_memory=in_memory)
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\bberman\Documents\Art-Workspace\venv\Lib\site-packages\datasets\arrow_reader.py", line 336, in _get_table_from_filename
    table = ArrowReader.read_table(filename, in_memory=in_memory)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\bberman\Documents\Art-Workspace\venv\Lib\site-packages\datasets\arrow_reader.py", line 357, in read_table
    return table_cls.from_file(filename)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\bberman\Documents\Art-Workspace\venv\Lib\site-packages\datasets\table.py", line 1059, in from_file
    table = _memory_mapped_arrow_table_from_file(filename)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\bberman\Documents\Art-Workspace\venv\Lib\site-packages\datasets\table.py", line 66, in _memory_mapped_arrow_table_from_file
    pa_table = opened_stream.read_all()
               ^^^^^^^^^^^^^^^^^^^^^^^^
  File "pyarrow\ipc.pxi", line 757, in pyarrow.lib.RecordBatchReader.read_all
  File "pyarrow\error.pxi", line 91, in pyarrow.lib.check_status
OSError: [WinError 8] PrefetchVirtualMemory failed. Detail: [Windows error 8] Not enough memory resources are available to process this command.
```

This error is probably a red herring: https://stackoverflow.com/questions/50263929/numpy-memmap-returns-not-enough-memory-while-there-are-plenty-available In other words, the issue is related to asking for a memory mapping of length N > M the length of the file on Windows. This gracefully succeeds on Linux.

I have 1024 arrow files in my cache instead of 128 like in the repository for it. Probably related. I don't know why `datasets` reorganized/rewrote the dataset in my cache to be 1024 slices instead of the original 128.

### Steps to reproduce the bug

```
# as a huggingface developer, you may already have laion2B-en somewhere
_CACHE_DIR = "."

from datasets import load_dataset
load_dataset("laion/laion2B-en", split="train", cache_dir=_CACHE_DIR, keep_in_memory=False)
```

### Expected behavior

This should correctly load as a memory mapped Arrow dataset.

### Environment info


- `datasets` version: 2.15.0
- Platform: Windows-10-10.0.20348-SP0 (this is windows 2022)
- Python version: 3.11.4
- `huggingface_hub` version: 0.19.4
- PyArrow version: 14.0.1
- Pandas version: 2.1.2
- `fsspec` version: 2023.10.0


Comments: ~~You will see this error if the cache dir filepath contains relative `..` paths. Use `os.path.realpath(_CACHE_DIR)` before passing it to the `load_dataset` function.~~

This is a real issue and not related to paths.

Based on the StackOverflow answer, this causes the error to go away:
```diff
diff --git a/table.py b/table.py
--- a/table.py	
+++ b/table.py	(date 1701824849806)
@@ -47,7 +47,7 @@
 
 
 def _memory_mapped_record_batch_reader_from_file(filename: str) -> pa.RecordBatchStreamReader:
-    memory_mapped_stream = pa.memory_map(filename)
+    memory_mapped_stream = pa.memory_map(filename, "r+")
     return pa.ipc.open_stream(memory_mapped_stream)
```
But now loading the dataset goes very, very slowly, which is unexpected.

I don't really comprehend what it is that `datasets` gave me when it downloaded the laion2B-en dataset, because nothing can seemingly read these 1024 .arrow files it is retrieving. Not `polars`, not `pyarrow`, it's not an `ipc` file, it's not a `parquet` file...

Hi! 

Instead of generating one (potentially large) Arrow file, we shard the generated data into 500 MB shards because memory-mapping large Arrow files can be problematic on some systems. Maybe deleting the dataset's cache and increasing the shard size (controlled with the `datasets.config.MAX_SHARD_SIZE` variable; e.g. to "4GB") can fix the issue for you.

> I don't really comprehend what it is that `datasets` gave me when it downloaded the laion2B-en dataset, because nothing can seemingly read these 1024 .arrow files it is retrieving. Not `polars`, not `pyarrow`, it's not an `ipc` file, it's not a `parquet` file...

Our `.arrow` files are in the [Arrow streaming format](https://arrow.apache.org/docs/python/ipc.html#using-streams). To load them as a `polars` DataFrame, do the following:
```python
df = pl.from_arrow(Dataset.from_from(path_to_arrow_file).data.table)
```

We plan to switch to the IPC version eventually.


Hmm, I have a feeling this works fine on Linux, and is a real bug for however `datasets` is doing the sharding on Windows. I will follow up, but I think this is a real bug.