GitHub issue #2968: `DatasetDict` cannot be exported to parquet if the splits have different features
State: closed | Created: 2021-09-25T22:18:39Z

## 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