`DatasetDict` cannot be exported to parquet if the splits have different features
Opened by LysandreJik (closed)

## Describe the bug

I'm trying to use parquet as a means of serialization for both `Dataset` and `DatasetDict` objects. Using `to_parquet` alongside `from_parquet` or `load_dataset` for a `Dataset` works perfectly. 

For `DatasetDict`, I use `to_parquet` on each split to save the parquet files in individual folders representing individual splits. This works too, as long as the splits have identical features. If a split has different features to neighboring splits, then loading the dataset will fail: a single schema is used to load both splits, resulting in a failure to load the second parquet file.

## Steps to reproduce the bug

The following works as expected:

```python
from datasets import load_dataset

ds = load_dataset("lhoestq/custom_squad")

ds['train'].to_parquet("./ds/train/split.parquet")
ds['validation'].to_parquet("./ds/validation/split.parquet")

brand_new_dataset = load_dataset("ds")
```

Modifying a single split to add a new feature ends up in a crash:

```python
from datasets import load_dataset

ds = load_dataset("lhoestq/custom_squad")


def identical_answers(e):
    e['identical_answers'] = len(set(e['answers']['text'])) == 1
    return e


ds['validation'] = ds['validation'].map(identical_answers)
ds['train'].to_parquet("./ds/train/split.parquet")
ds['validation'].to_parquet("./ds/validation/split.parquet")

brand_new_dataset = load_dataset("ds")
```
```
  File "/home/lysandre/.config/JetBrains/PyCharm2021.2/scratches/datasets/upload_dataset.py", line 26, in <module>
    brand_new_dataset = load_dataset("ds")
  File "/home/lysandre/Workspaces/Python/datasets/src/datasets/load.py", line 1151, in load_dataset
    builder_instance.download_and_prepare(
  File "/home/lysandre/Workspaces/Python/datasets/src/datasets/builder.py", line 642, in download_and_prepare
    self._download_and_prepare(
  File "/home/lysandre/Workspaces/Python/datasets/src/datasets/builder.py", line 732, in _download_and_prepare
    self._prepare_split(split_generator, **prepare_split_kwargs)
  File "/home/lysandre/Workspaces/Python/datasets/src/datasets/builder.py", line 1194, in _prepare_split
    writer.write_table(table)
  File "/home/lysandre/Workspaces/Python/datasets/src/datasets/arrow_writer.py", line 428, in write_table
    pa_table = pa.Table.from_arrays([pa_table[name] for name in self._schema.names], schema=self._schema)
  File "/home/lysandre/Workspaces/Python/datasets/src/datasets/arrow_writer.py", line 428, in <listcomp>
    pa_table = pa.Table.from_arrays([pa_table[name] for name in self._schema.names], schema=self._schema)
  File "pyarrow/table.pxi", line 1257, in pyarrow.lib.Table.__getitem__
  File "pyarrow/table.pxi", line 1833, in pyarrow.lib.Table.column
  File "pyarrow/table.pxi", line 1808, in pyarrow.lib.Table._ensure_integer_index
KeyError: 'Field "identical_answers" does not exist in table schema'
```

It does work, however, to use the `save_to_disk` and `load_from_disk` methods:

```py
from datasets import load_from_disk

ds = load_dataset("lhoestq/custom_squad")


def identical_answers(e):
    e['identical_answers'] = len(set(e['answers']['text'])) == 1
    return e


ds['validation'] = ds['validation'].map(identical_answers)

ds.save_to_disk("local_path")
brand_new_dataset = load_from_disk("local_path")
```

## Expected results

The saving works correctly - but the loading fails. I would expect either an error when saving or an error-less instantiation of the dataset through the parquet files.

If it's helpful, I've traced a possible patch to the `write_table` method here: 

https://github.com/huggingface/datasets/blob/26ff41aa3a642e46489db9e95be1e9a8c4e64bea/src/datasets/arrow_writer.py#L424-L425

The writer is built only if the parquet writer is `None`, but I expect we would want to build a new writer as the table schema has changed. Furthermore, it relies on having the property `update_features` set to `True` in order to update the features:

https://github.com/huggingface/datasets/blob/26ff41aa3a642e46489db9e95be1e9a8c4e64bea/src/datasets/arrow_writer.py#L254-L255

but the `ArrowWriter` is instantiated without that option in the `_prepare_split` method of the `ArrowBasedBuilder`:

https://github.com/huggingface/datasets/blob/26ff41aa3a642e46489db9e95be1e9a8c4e64bea/src/datasets/builder.py#L1190

Updating these two parts to recreate a schema on each split results in an error that is, unfortunately, out of my expertise:

```
  File "/home/lysandre/.config/JetBrains/PyCharm2021.2/scratches/datasets/upload_dataset.py", line 27, in <module>
    brand_new_dataset = load_dataset("ds")
  File "/home/lysandre/Workspaces/Python/datasets/src/datasets/load.py", line 1163, in load_dataset
    ds = builder_instance.as_dataset(split=split, ignore_verifications=ignore_verifications, in_memory=keep_in_memory)
  File "/home/lysandre/Workspaces/Python/datasets/src/datasets/builder.py", line 819, in as_dataset
    datasets = utils.map_nested(
  File "/home/lysandre/Workspaces/Python/datasets/src/datasets/utils/py_utils.py", line 207, in map_nested
    mapped = [
  File "/home/lysandre/Workspaces/Python/datasets/src/datasets/utils/py_utils.py", line 208, in <listcomp>
    _single_map_nested((function, obj, types, None, True))
  File "/home/lysandre/Workspaces/Python/datasets/src/datasets/utils/py_utils.py", line 143, in _single_map_nested
    return function(data_struct)
  File "/home/lysandre/Workspaces/Python/datasets/src/datasets/builder.py", line 850, in _build_single_dataset
    ds = self._as_dataset(
  File "/home/lysandre/Workspaces/Python/datasets/src/datasets/builder.py", line 920, in _as_dataset
    dataset_kwargs = ArrowReader(self._cache_dir, self.info).read(
  File "/home/lysandre/Workspaces/Python/datasets/src/datasets/arrow_reader.py", line 217, in read
    return self.read_files(files=files, original_instructions=instructions, in_memory=in_memory)
  File "/home/lysandre/Workspaces/Python/datasets/src/datasets/arrow_reader.py", line 238, in read_files
    pa_table = self._read_files(files, in_memory=in_memory)
  File "/home/lysandre/Workspaces/Python/datasets/src/datasets/arrow_reader.py", line 173, in _read_files
    pa_table: Table = self._get_table_from_filename(f_dict, in_memory=in_memory)
  File "/home/lysandre/Workspaces/Python/datasets/src/datasets/arrow_reader.py", line 308, in _get_table_from_filename
    table = ArrowReader.read_table(filename, in_memory=in_memory)
  File "/home/lysandre/Workspaces/Python/datasets/src/datasets/arrow_reader.py", line 327, in read_table
    return table_cls.from_file(filename)
  File "/home/lysandre/Workspaces/Python/datasets/src/datasets/table.py", line 458, in from_file
    table = _memory_mapped_arrow_table_from_file(filename)
  File "/home/lysandre/Workspaces/Python/datasets/src/datasets/table.py", line 45, in _memory_mapped_arrow_table_from_file
    pa_table = opened_stream.read_all()
  File "pyarrow/ipc.pxi", line 563, in pyarrow.lib.RecordBatchReader.read_all
  File "pyarrow/error.pxi", line 114, in pyarrow.lib.check_status
OSError: Header-type of flatbuffer-encoded Message is not RecordBatch.
```
## Environment info

- `datasets` version: 1.12.2.dev0
- Platform: Linux-5.14.7-arch1-1-x86_64-with-glibc2.33
- Python version: 3.9.7
- PyArrow version: 5.0.0




Comments:
This is because you have to specify which split corresponds to what file:
```python
data_files = {"train": "train/split.parquet", "validation": "validation/split.parquet"}
brand_new_dataset_2 = load_dataset("ds", data_files=data_files)
```

Otherwise it tries to concatenate the two splits, and it fails because they don't have the same features.

It works with save_to_disk/load_from_disk because it also stores json files that contain the information about which files goes into which split

Wonderful, thanks for the help!

I may be mistaken but I think the following doesn't work either:

```python
from datasets import load_dataset

ds = load_dataset("lhoestq/custom_squad")


def identical_answers(e):
    e['identical_answers'] = len(set(e['answers']['text'])) == 1
    return e


ds['validation'] = ds['validation'].map(identical_answers)
ds['train'].to_parquet("./ds/train/split.parquet")
ds['validation'].to_parquet("./ds/validation/split.parquet")

data_files = {"train": "train/split.parquet", "validation": "validation/split.parquet"}
brand_new_dataset_2 = load_dataset("ds", data_files=data_files)
```

It works on my side as soon as the directories named `ds/train` and `ds/validation` exist (otherwise it returns a FileNotFoundError). What error are you getting ?

Also we may introduce a default mapping for the data files:
```python
{
    "train": ["*train*"],
    "test": ["*test*"],
    "validation": ["*dev*", "valid"],
}
```
this way if you name your files according to the splits you won't have to specify the data_files parameter. What do you think ?

I moved this discussion to #3027 

I'm getting the following error:

```
Downloading and preparing dataset custom_squad/plain_text to /home/lysandre/.cache/huggingface/datasets/lhoestq___custom_squad)/plain_text/1.0.0/397916d1ae99584877e0fb4f5b8b6f01e66fcbbeff4d178afb30c933a8d0d93a...
100%|██████████| 2/2 [00:00<00:00, 7760.04it/s]
100%|██████████| 2/2 [00:00<00:00, 2020.38it/s]
  0%|          | 0/2 [00:00<?, ?it/s]Traceback (most recent call last):
  File "<input>", line 1, in <module>
  File "/opt/pycharm-professional/plugins/python/helpers/pydev/_pydev_bundle/pydev_umd.py", line 198, in runfile
    pydev_imports.execfile(filename, global_vars, local_vars)  # execute the script
  File "/opt/pycharm-professional/plugins/python/helpers/pydev/_pydev_imps/_pydev_execfile.py", line 18, in execfile
    exec(compile(contents+"\n", file, 'exec'), glob, loc)
  File "/home/lysandre/.config/JetBrains/PyCharm2021.2/scratches/datasets/upload_dataset.py", line 12, in <module>
    ds = load_dataset("lhoestq/custom_squad")
  File "/home/lysandre/Workspaces/Python/datasets/src/datasets/load.py", line 1207, in load_dataset
    ds = builder_instance.as_dataset(split=split, ignore_verifications=ignore_verifications, in_memory=keep_in_memory)
  File "/home/lysandre/Workspaces/Python/datasets/src/datasets/builder.py", line 823, in as_dataset
    datasets = utils.map_nested(
  File "/home/lysandre/Workspaces/Python/datasets/src/datasets/utils/py_utils.py", line 207, in map_nested
    mapped = [
  File "/home/lysandre/Workspaces/Python/datasets/src/datasets/utils/py_utils.py", line 208, in <listcomp>
    _single_map_nested((function, obj, types, None, True))
  File "/home/lysandre/Workspaces/Python/datasets/src/datasets/utils/py_utils.py", line 143, in _single_map_nested
    return function(data_struct)
  File "/home/lysandre/Workspaces/Python/datasets/src/datasets/builder.py", line 854, in _build_single_dataset
    ds = self._as_dataset(
  File "/home/lysandre/Workspaces/Python/datasets/src/datasets/builder.py", line 924, in _as_dataset
    dataset_kwargs = ArrowReader(self._cache_dir, self.info).read(
  File "/home/lysandre/Workspaces/Python/datasets/src/datasets/arrow_reader.py", line 217, in read
    return self.read_files(files=files, original_instructions=instructions, in_memory=in_memory)
  File "/home/lysandre/Workspaces/Python/datasets/src/datasets/arrow_reader.py", line 238, in read_files
    pa_table = self._read_files(files, in_memory=in_memory)
  File "/home/lysandre/Workspaces/Python/datasets/src/datasets/arrow_reader.py", line 173, in _read_files
    pa_table: Table = self._get_table_from_filename(f_dict, in_memory=in_memory)
  File "/home/lysandre/Workspaces/Python/datasets/src/datasets/arrow_reader.py", line 308, in _get_table_from_filename
    table = ArrowReader.read_table(filename, in_memory=in_memory)
  File "/home/lysandre/Workspaces/Python/datasets/src/datasets/arrow_reader.py", line 327, in read_table
    return table_cls.from_file(filename)
  File "/home/lysandre/Workspaces/Python/datasets/src/datasets/table.py", line 458, in from_file
    table = _memory_mapped_arrow_table_from_file(filename)
  File "/home/lysandre/Workspaces/Python/datasets/src/datasets/table.py", line 45, in _memory_mapped_arrow_table_from_file
    pa_table = opened_stream.read_all()
  File "pyarrow/ipc.pxi", line 563, in pyarrow.lib.RecordBatchReader.read_all
  File "pyarrow/error.pxi", line 114, in pyarrow.lib.check_status
OSError: Header-type of flatbuffer-encoded Message is not RecordBatch.
```

Tried on current master, after updating latest dependencies and obtained the same result

The proposal in #3027 sounds good to me!

I just tried again on colab by installing `datasets` from source with pyarrow 3.0.0 and didn't get any error.

You error seems to happen when doing
```python
ds = load_dataset("lhoestq/custom_squad")
```

More specifically it fails when trying to read the arrow file that just got generated. I haven't issues like this before. Can you make sure you have a recent version of `pyarrow` ? Maybe it was an old version that wrote the arrow file and some header was missing.

Thank you for your pointer! This seems to have been linked to Python 3.9.7: it works flawlessly with Python 3.8.6. This can be closed, thanks a lot for your help.