NonMatchingSplitsSizesError when using data_dir
Opened by srehaag (closed)

### Describe the bug

Loading a dataset from with a data_dir argument generates a NonMatchingSplitsSizesError if there are multiple directories in the dataset. 

This appears to happen because the expected split is calculated based on the data in all the directories whereas the recorded split is calculated based on the data in the directory specified using the data_dir argument. 

This is recent behavior. Until the past few weeks loading using the data_dir argument worked without any issue.

### Steps to reproduce the bug

Simple test dataset available here: https://huggingface.co/datasets/srehaag/hf-bug-temp

The dataset contains two directories "data1" and "data2", each with a file called "train.parquet" with a 2 x 5 table.

from datasets import load_dataset
dataset = load_dataset("srehaag/hf-bug-temp", data_dir = "data1")

Generates:

---------------------------------------------------------------------------
NonMatchingSplitsSizesError               Traceback (most recent call last)
Cell In[3], <a href='vscode-notebook-cell:?execution_count=3&line=2'>line 2</a>
      <a href='vscode-notebook-cell:?execution_count=3&line=1'>1</a> from datasets import load_dataset
----> <a href='vscode-notebook-cell:?execution_count=3&line=2'>2</a> dataset = load_dataset("srehaag/hf-bug-temp", data_dir = "data1")

File ~/.python/current/lib/python3.10/site-packages/datasets/load.py:2609, in load_dataset(path, name, data_dir, data_files, split, cache_dir, features, download_config, download_mode, verification_mode, ignore_verifications, keep_in_memory, save_infos, revision, token, use_auth_token, task, streaming, num_proc, storage_options, trust_remote_code, **config_kwargs)
   <a href='~/.python/current/lib/python3.10/site-packages/datasets/load.py:2606'>2606</a>     return builder_instance.as_streaming_dataset(split=split)
   <a href='~/.python/current/lib/python3.10/site-packages/datasets/load.py:2608'>2608</a> # Download and prepare data
-> <a href='~/.python/current/lib/python3.10/site-packages/datasets/load.py:2609'>2609</a> builder_instance.download_and_prepare(
   <a href='~/.python/current/lib/python3.10/site-packages/datasets/load.py:2610'>2610</a>     download_config=download_config,
   <a href='~/.python/current/lib/python3.10/site-packages/datasets/load.py:2611'>2611</a>     download_mode=download_mode,
   <a href='~/.python/current/lib/python3.10/site-packages/datasets/load.py:2612'>2612</a>     verification_mode=verification_mode,
   <a href='~/.python/current/lib/python3.10/site-packages/datasets/load.py:2613'>2613</a>     num_proc=num_proc,
   <a href='~/.python/current/lib/python3.10/site-packages/datasets/load.py:2614'>2614</a>     storage_options=storage_options,
   <a href='~/.python/current/lib/python3.10/site-packages/datasets/load.py:2615'>2615</a> )
   <a href='~/.python/current/lib/python3.10/site-packages/datasets/load.py:2617'>2617</a> # Build dataset for splits
   <a href='~/.python/current/lib/python3.10/site-packages/datasets/load.py:2618'>2618</a> keep_in_memory = (
   <a href='~/.python/current/lib/python3.10/site-packages/datasets/load.py:2619'>2619</a>     keep_in_memory if keep_in_memory is not None else is_small_dataset(builder_instance.info.dataset_size)
   <a href='~/.python/current/lib/python3.10/site-packages/datasets/load.py:2620'>2620</a> )

File ~/.python/current/lib/python3.10/site-packages/datasets/builder.py:1027, in DatasetBuilder.download_and_prepare(self, output_dir, download_config, download_mode, verification_mode, ignore_verifications, try_from_hf_gcs, dl_manager, base_path, use_auth_token, file_format, max_shard_size, num_proc, storage_options, **download_and_prepare_kwargs)
   <a href='~/.python/current/lib/python3.10/site-packages/datasets/builder.py:1025'>1025</a>     if num_proc is not None:
   <a href='~/.python/current/lib/python3.10/site-packages/datasets/builder.py:1026'>1026</a>         prepare_split_kwargs["num_proc"] = num_proc
-> <a href='~/.python/current/lib/python3.10/site-packages/datasets/builder.py:1027'>1027</a>     self._download_and_prepare(
   <a href='~/.python/current/lib/python3.10/site-packages/datasets/builder.py:1028'>1028</a>         dl_manager=dl_manager,
   <a href='~/.python/current/lib/python3.10/site-packages/datasets/builder.py:1029'>1029</a>         verification_mode=verification_mode,
   <a href='~/.python/current/lib/python3.10/site-packages/datasets/builder.py:1030'>1030</a>         **prepare_split_kwargs,
   <a href='~/.python/current/lib/python3.10/site-packages/datasets/builder.py:1031'>1031</a>         **download_and_prepare_kwargs,
   <a href='~/.python/current/lib/python3.10/site-packages/datasets/builder.py:1032'>1032</a>     )
   <a href='~/.python/current/lib/python3.10/site-packages/datasets/builder.py:1033'>1033</a> # Sync info
   <a href='~/.python/current/lib/python3.10/site-packages/datasets/builder.py:1034'>1034</a> self.info.dataset_size = sum(split.num_bytes for split in self.info.splits.values())

File ~/.python/current/lib/python3.10/site-packages/datasets/builder.py:1140, in DatasetBuilder._download_and_prepare(self, dl_manager, verification_mode, **prepare_split_kwargs)
   <a href='~/.python/current/lib/python3.10/site-packages/datasets/builder.py:1137'>1137</a>     dl_manager.manage_extracted_files()
   <a href='~/.python/current/lib/python3.10/site-packages/datasets/builder.py:1139'>1139</a> if verification_mode == VerificationMode.BASIC_CHECKS or verification_mode == VerificationMode.ALL_CHECKS:
-> <a href='~/.python/current/lib/python3.10/site-packages/datasets/builder.py:1140'>1140</a>     verify_splits(self.info.splits, split_dict)
   <a href='~/.python/current/lib/python3.10/site-packages/datasets/builder.py:1142'>1142</a> # Update the info object with the splits.
   <a href='~/.python/current/lib/python3.10/site-packages/datasets/builder.py:1143'>1143</a> self.info.splits = split_dict

File ~/.python/current/lib/python3.10/site-packages/datasets/utils/info_utils.py:101, in verify_splits(expected_splits, recorded_splits)
     <a href='~/.python/current/lib/python3.10/site-packages/datasets/utils/info_utils.py:95'>95</a> bad_splits = [
     <a href='~/.python/current/lib/python3.10/site-packages/datasets/utils/info_utils.py:96'>96</a>     {"expected": expected_splits[name], "recorded": recorded_splits[name]}
     <a href='~/.python/current/lib/python3.10/site-packages/datasets/utils/info_utils.py:97'>97</a>     for name in expected_splits
     <a href='~/.python/current/lib/python3.10/site-packages/datasets/utils/info_utils.py:98'>98</a>     if expected_splits[name].num_examples != recorded_splits[name].num_examples
     <a href='~/.python/current/lib/python3.10/site-packages/datasets/utils/info_utils.py:99'>99</a> ]
    <a href='~/.python/current/lib/python3.10/site-packages/datasets/utils/info_utils.py:100'>100</a> if len(bad_splits) > 0:
--> <a href='~/.python/current/lib/python3.10/site-packages/datasets/utils/info_utils.py:101'>101</a>     raise NonMatchingSplitsSizesError(str(bad_splits))
    <a href='~/.python/current/lib/python3.10/site-packages/datasets/utils/info_utils.py:102'>102</a> logger.info("All the splits matched successfully.")

NonMatchingSplitsSizesError: [{'expected': SplitInfo(name='train', num_bytes=212, num_examples=10, shard_lengths=None, dataset_name=None), 'recorded': SplitInfo(name='train', num_bytes=106, num_examples=5, shard_lengths=None, dataset_name='hf-bug-temp')}]

__________

By contrast, this loads the data from both data1/train.parquet and data2/train.parquet without any error message:

from datasets import load_dataset
dataset = load_dataset("srehaag/hf-bug-temp")



### Expected behavior

Should load the 5 x 2 table from data1/train.parquet without error message.

### Environment info

Used Codespaces to simplify environment (see details below), but bug is present across various configurations.

- `datasets` version: 2.19.1
- Platform: Linux-6.5.0-1021-azure-x86_64-with-glibc2.31
- Python version: 3.10.13
- `huggingface_hub` version: 0.23.1
- PyArrow version: 16.1.0
- Pandas version: 2.2.2
- `fsspec` version: 2024.3.1

Comments:
Thanks for reporting, @srehaag.

We are investigating this issue.

I confirm there is a bug for data-based Hub datasets when the user passes `data_dir`, which was introduced by PR:
- #6714