Implement capability to restore non-nullability in Features
State: closed | Author: BramVanroy (CONTRIBUTOR)

This PR attempts to keep track of non_nullable pyarrow fields when converting a `pa.Schema` to `Features`. At the same time, when outputting the `arrow_schema`, the original non-nullable fields are restored. This allows for more consistent behavior and avoids breaking behavior as illustrated in #7479.

I am by no means a pyarrow expert so some logic in `find_non_nullable_fields` may not perfect. Not sure if more logic (type checks) are needed for deep-checking a given schema. Maybe there are other pyarrow structures that need to be covered?

Tests are added, but again, these may not have sufficient coverage in terms of pyarrow structure types.

closes #7479 

Comments: Interestingly, this does not close #7479. The Features are not correctly maintained when calling `from_dict` with the custom Features.

Unfortunately this PR does not fix the reported issue. After more digging:

- when the dataset is created, nullability information is lost in Features;
- even with this PR, it will get lost eventually because of internal copying/recreation of the Features object without accounting for the nullable fields;
- even if that is also fixed, and Features.arrow_schema correctly holds the nullability info, [casting the arrow Table](https://github.com/huggingface/datasets/blob/5f8d2ad9a1b0bccfd962d998987228addfd5be9f/src/datasets/arrow_dataset.py#L677) with a less strict schema to a more strict one (with nullability) will fail (only on deeper structs, not on flat fields). 

Interestingly, passing custom Features does not immediately load the underlying data with the right arrow_schema. Instead, the workflow is like this:

- load pyarrow table with any of the methods (from_dict, from_pandas, etc.), which will always AUTO INFER rather than use a provided schema
- the loaded table with auto-schema will be used to initialize the `Dataset` class, and only during construction will [CAST](https://github.com/huggingface/datasets/blob/5f8d2ad9a1b0bccfd962d998987228addfd5be9f/src/datasets/arrow_dataset.py#L677) the table to the user-provided schema if needed, if it differs from the auto-inferred one.

So I figured, since many/all of the pyarrow [`Table.from_*`](https://arrow.apache.org/docs/python/generated/pyarrow.Table.html) methods have a `schema=` argument, we should already load the Table with the correct schema to begin with. As an example, I tried changing this line:

https://github.com/huggingface/datasets/blob/5f8d2ad9a1b0bccfd962d998987228addfd5be9f/src/datasets/arrow_dataset.py#L940

to include the arrow_schema, if provided:

```python
pa_table = InMemoryTable.from_pydict(mapping=mapping, schema=features.arrow_schema if features is not None else None)
```

But that leads to:

```
Traceback (most recent call last):
  File "/home/ampere/vanroy/datasets/scratch.py", line 33, in <module>
    ds = Dataset.from_dict(
         ^^^^^^^^^^^^^^^^^^
  File "/home/local/vanroy/datasets/src/datasets/arrow_dataset.py", line 957, in from_dict
    pa_table = InMemoryTable.from_pydict(mapping=mapping, schema=features.arrow_schema if features is not None else None)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/local/vanroy/datasets/src/datasets/table.py", line 758, in from_pydict
    return cls(pa.Table.from_pydict(*args, **kwargs))
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "pyarrow/table.pxi", line 1968, in pyarrow.lib._Tabular.from_pydict
  File "pyarrow/table.pxi", line 6354, in pyarrow.lib._from_pydict
  File "pyarrow/array.pxi", line 402, in pyarrow.lib.asarray
  File "pyarrow/array.pxi", line 252, in pyarrow.lib.array
  File "pyarrow/array.pxi", line 114, in pyarrow.lib._handle_arrow_array_protocol
  File "/home/local/vanroy/datasets/src/datasets/arrow_writer.py", line 201, in __arrow_array__
    raise ValueError("TypedSequence is supposed to be used with pa.array(typed_sequence, type=None)")
ValueError: TypedSequence is supposed to be used with pa.array(typed_sequence, type=None)
```

and I am not too familiar with pyarrow to solve this.

So ultimately I'm a bit at a loss here. I *think*, if we'd want to do this right, the automatic casting in init should be removed in favor of handling the logic inside `Dataset.from_*`, by passing the schema explicitly to `pa.Table.from_*(..., schema=schema)`. But I lack the knowledge of pyarrow to go further than what I've written about above.


It's indeed a bit more work to support nullable since in addition to your comments, there are unclear behavior when it comes to concatenating nullable with non-nullable, and maybe how to handle non-nullable lists and nested data.

But yup I agree having the `Dataset.from_*` function pass the `schema` to the `pa.Table.from*` would be the way.

Just one comment about this error: 

```
ValueError: TypedSequence is supposed to be used with pa.array(typed_sequence, type=None)
```

This happens because `Dataset.from_dict` uses `OptimizedTypedSequence` by default, which should only be used if the user doesn't specify a schema