More robust None handling

[mariosasko]

PyArrow has explicit support for null values, so it makes sense to support Nones on our side as well.

Colab Notebook with examples

Changes:

• allow None for the features types with special encoding (ClassLabel, TranslationVariableLanguages, Value, _ArrayXD)
• handle None in class_encode_column (also there is an option to stringify Nones and treat them as a class)
• support None sorting in sort (use pandas for that)
• handle None in align_labels_with_mapping
• support for None in ArrayXD (converts None to np.nan to align the behavior with PyArrow)
• support for None in the Audio/Image feature
• allow promotion when concatenating tables (pa.concat_tables(table_list, promote=True)) and null row/column broadcasting similar to pandas

Additional notes:

• use null instead of none for function arguments for consistency with existing disable_nullable
• fixes a bug with the update_metadata_with_features call in Dataset.rename_columns
• had to update some tests, let me know if that's ok

TODO:

• check how the Audio features behaves with Nones
• Better None handling in concatenate_datasets/add_item
• Fix formatting with Nones
• Add Colab with examples
• Tests

TODOs for subsequent PRs:

• Mention None handling in the docs
• Add drop_null/fill_null to Dataset/DatasetDict

Fix #3181 #3253

[lhoestq]

I also created a PR regarding disable_nullable that must be always False by default, in order to always allow None values
#3211

[lhoestq]

Looking good !
I just have a question about padding when concatenating datasets or adding a column, and some nitpicks:

[mariosasko]

@lhoestq I addressed your comments, added tests, did some refactoring to make the implementation cleaner and added support for None values in map transforms when the feature type is ArrayXD (previously, I only implemented None decoding).

My only concern is that during decoding ArrayXD arrays with None values will be auto-casted to float64 to allow np.nan insertion and this might be unexpected if dtype is not float, so one option would be to allow None values only if the storage type is float32 or float64. Let me know WDYT would be the most consistent behavior here.

[lhoestq]

Cool ! :D

My only concern is that during decoding ArrayXD arrays with None values will be auto-casted to float64 to allow np.nan insertion and this might be unexpected if dtype is not float, so one option would be to allow None values only if the storage type is float32 or float64. Let me know WDYT would be the most consistent behavior here.

Yes that makes sense to only fill with nan if the type is compatible

[mariosasko]

After some more experimenting, I think we can keep auto-cast to float because PyArrow also does it:

import pyarrow as pa
arr = pa.array([1, 2, 3, 4, None], type=pa.int32()).to_numpy(zero_copy_only=False) # None present - int32 -> float64
assert arr.dtype == np.float64

Additional changes:

• fixes a bug in the _is_zero_copy_only implementation for the ArraXD types. Previously, _is_zero_copy_only would always return False for these types. Still have to see if it's possible to optimize copying of the non-extension types (Sequence, ...), but I plan to work on that in a separate PR.
• Allow dynamic first dimension for ArrayXD #2891 introduced a bug where the dtype of ArrayXD wouldn't be preserved due to to_pylist call in NumPy Formatter (np.array(np.array(..).tolist()) doesn't necessarily preserve dtype of the initial array), so I'm also fixing that.

[lhoestq]

Very nice ! My final comments:

[lhoestq]

LGTM ! Thank you :)

[lhoestq]

The CI fail for windows is unrelated to this PR, merging