Allow dynamic first dimension for ArrayXD
Opened by rpowalski (closed)

Add support for dynamic first dimension for ArrayXD features. See issue [#887](https://github.com/huggingface/datasets/issues/887).
Following changes allow for `to_pylist` method of `ArrayExtensionArray` to return a list of numpy arrays where fist dimension can vary.

@lhoestq Could you suggest how you want to extend test suit. For now I added only very limited testing.

Comments:
@lhoestq, thanks for your review.

I added test for `to_pylist`, I didn't do that for `to_numpy` because this method shouldn't be called for dynamic dimension ArrayXD - this method will try to make a single numpy array for the whole column which cannot be done for dynamic arrays.

I dig into `to_pandas()` functionality and I found it quite difficult to implement. `PandasArrayExtensionArray` takes single np.array as an argument. It might be a bit of changes to make it work with the list of arrays. Do you mind if we exclude this work from this PR. I added an error message for the case if somebody tries to use dynamic arrays with `to_pandas`

@lhoestq, I just fixed all the tests. Let me know if there is anything else to add.

@lhoestq, any chance you had some time to check out this PR?


Hi ! Sorry for the delay

It looks good to me ! I think the only thing missing is the support for passing a list of numpy arrays to `map` when the first dimension is dynamic.

Currently it raises an error:
```python
from datasets import *
import numpy as np

features=  Features({"a": Array3D(shape=(None, 5, 2), dtype="int32")})
d = Dataset.from_dict({"a": [np.zeros((5,5,2)), np.zeros((2,5,2))]}, features=features)
d = d.map(lambda a: {"a": np.concatenate([a]*2)}, input_columns="a")
print(d[0])
```
raises
```python
Traceback (most recent call last):
  File "playground/ttest.py", line 6, in <module>
    d = d.map(lambda x: {"a": np.concatenate([x]*2)}, input_columns="a")
  File "/home/truent/hf/datasets/src/datasets/arrow_dataset.py", line 1932, in map
    return self._map_single(
  File "/home/truent/hf/datasets/src/datasets/arrow_dataset.py", line 426, in wrapper
    out: Union["Dataset", "DatasetDict"] = func(self, *args, **kwargs)
  File "/home/truent/hf/datasets/src/datasets/fingerprint.py", line 406, in wrapper
    out = func(self, *args, **kwargs)
  File "/home/truent/hf/datasets/src/datasets/arrow_dataset.py", line 2317, in _map_single
    writer.finalize()
  File "/home/truent/hf/datasets/src/datasets/arrow_writer.py", line 443, in finalize
    self.write_examples_on_file()
  File "/home/truent/hf/datasets/src/datasets/arrow_writer.py", line 312, in write_examples_on_file
    pa_array = pa.array(typed_sequence)
  File "pyarrow/array.pxi", line 222, in pyarrow.lib.array
  File "pyarrow/array.pxi", line 110, in pyarrow.lib._handle_arrow_array_protocol
  File "/home/truent/hf/datasets/src/datasets/arrow_writer.py", line 108, in __arrow_array__
    storage = pa.array(self.data, type.storage_dtype)
  File "pyarrow/array.pxi", line 305, in pyarrow.lib.array
  File "pyarrow/array.pxi", line 39, in pyarrow.lib._sequence_to_array
  File "pyarrow/error.pxi", line 122, in pyarrow.lib.pyarrow_internal_check_status
  File "pyarrow/error.pxi", line 84, in pyarrow.lib.check_status
pyarrow.lib.ArrowInvalid: Can only convert 1-dimensional array values
```

I think the issue is that here we don't cover the case where self.data is a list of numpy arrays:

https://github.com/huggingface/datasets/blob/55fd140a63b8f03a0e72985647e498f1fc799d3f/src/datasets/arrow_writer.py#L104-L109

We should remove the `isinstance(self.data[0], np.ndarray)` part and add these lines to cover this case:

https://github.com/huggingface/datasets/blob/55fd140a63b8f03a0e72985647e498f1fc799d3f/src/datasets/arrow_writer.py#L112-L113

@lhoestq, thanks, good catch!
Are you able to run this check with fixed dimension ArrayXD?
for below example
```
import numpy as np
from datasets import *

features = Features({"a": Array3D(shape=(2, 5, 2), dtype="int32")})
d = Dataset.from_dict({"a": [np.zeros((2, 5, 2)), np.zeros((2, 5, 2))]}, features=features)
d = d.map(lambda a: {"a": np.array(a) + 1}, input_columns="a")
print(d[0])
```

I am getting:
```
  File "/home/ib/datasets/src/datasets/arrow_writer.py", line 116, in __arrow_array__
    if trying_type and out[0].as_py() != self.data[0]:
ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()
```

Nevertheless, I tried to fix that. Let me know if that works.

@lhoestq, just resolved the conflicts. Let me know if there is anything left to do with this PR

Hi, thanks a lot for your comments.
Agree, happy to contribute to this topic in future PRs

Hi @rpowalski, thanks for adding this feature! 

I wanted to check if you are still interested in documenting this, otherwise I'd be happy to help with it :)