feat: create an `Array3D` column from a list of arrays of dimension 2
State: closed
Labels: description: New feature or request
name: enhancement

**Is your feature request related to a problem? Please describe.**

It is possible to create an `Array2D` column from a list of arrays of dimension 1. Similarly, I think it might be nice to be able to create a `Array3D` column from a list of lists of arrays of dimension 1.

To illustrate my proposal, let's take the following toy dataset t:
```python
import numpy as np
from datasets import Dataset, features

data_map = {
    1: np.array([[0.2, 0,4],[0.19, 0,3]]),
    2: np.array([[0.1, 0,4],[0.19, 0,3]]),
}

def create_toy_ds():
    my_dict = {"id":[1, 2]}
    return Dataset.from_dict(my_dict)

ds = create_toy_ds()
```

The following 2D processing works without any errors raised:
```python
def prepare_dataset_2D(batch):
    batch["pixel_values"] = [data_map[index] for index in batch["id"]]
    return batch
    
ds_2D = ds.map(
    prepare_dataset_2D, 
    batched=True, 
    remove_columns=ds.column_names, 
    features=features.Features({"pixel_values": features.Array2D(shape=(2, 3), dtype="float32")})
)
```

The following 3D processing doesn't work:
```python
def prepare_dataset_3D(batch):
    batch["pixel_values"] = [[data_map[index]] for index in batch["id"]]
    return batch
    
ds_3D = ds.map(
    prepare_dataset_3D, 
    batched=True, 
    remove_columns=ds.column_names, 
    features=features.Features({"pixel_values": features.Array3D(shape=(1,  2, 3, dtype="float32")})
)
```
The error raised is:
```
---------------------------------------------------------------------------
ArrowInvalid                              Traceback (most recent call last)
[<ipython-input-6-676547e4cd41>](https://localhost:8080/#) in <module>()
      3     batched=True,
      4     remove_columns=ds.column_names,
----> 5     features=features.Features({"pixel_values": features.Array3D(shape=(1, 2, 3), dtype="float32")})
      6 )

12 frames
[/usr/local/lib/python3.7/dist-packages/datasets/arrow_dataset.py](https://localhost:8080/#) in map(self, function, with_indices, with_rank, input_columns, batched, batch_size, drop_last_batch, remove_columns, keep_in_memory, load_from_cache_file, cache_file_name, writer_batch_size, features, disable_nullable, fn_kwargs, num_proc, suffix_template, new_fingerprint, desc)
   1971                 new_fingerprint=new_fingerprint,
   1972                 disable_tqdm=disable_tqdm,
-> 1973                 desc=desc,
   1974             )
   1975         else:

[/usr/local/lib/python3.7/dist-packages/datasets/arrow_dataset.py](https://localhost:8080/#) in wrapper(*args, **kwargs)
    518             self: "Dataset" = kwargs.pop("self")
    519         # apply actual function
--> 520         out: Union["Dataset", "DatasetDict"] = func(self, *args, **kwargs)
    521         datasets: List["Dataset"] = list(out.values()) if isinstance(out, dict) else [out]
    522         for dataset in datasets:

[/usr/local/lib/python3.7/dist-packages/datasets/arrow_dataset.py](https://localhost:8080/#) in wrapper(*args, **kwargs)
    485         }
    486         # apply actual function
--> 487         out: Union["Dataset", "DatasetDict"] = func(self, *args, **kwargs)
    488         datasets: List["Dataset"] = list(out.values()) if isinstance(out, dict) else [out]
    489         # re-apply format to the output

[/usr/local/lib/python3.7/dist-packages/datasets/fingerprint.py](https://localhost:8080/#) in wrapper(*args, **kwargs)
    456             # Call actual function
    457 
--> 458             out = func(self, *args, **kwargs)
    459 
    460             # Update fingerprint of in-place transforms + update in-place history of transforms

[/usr/local/lib/python3.7/dist-packages/datasets/arrow_dataset.py](https://localhost:8080/#) in _map_single(self, function, with_indices, with_rank, input_columns, batched, batch_size, drop_last_batch, remove_columns, keep_in_memory, load_from_cache_file, cache_file_name, writer_batch_size, features, disable_nullable, fn_kwargs, new_fingerprint, rank, offset, disable_tqdm, desc, cache_only)
   2354                                 writer.write_table(batch)
   2355                             else:
-> 2356                                 writer.write_batch(batch)
   2357                 if update_data and writer is not None:
   2358                     writer.finalize()  # close_stream=bool(buf_writer is None))  # We only close if we are writing in a file

[/usr/local/lib/python3.7/dist-packages/datasets/arrow_writer.py](https://localhost:8080/#) in write_batch(self, batch_examples, writer_batch_size)
    505             col_try_type = try_features[col] if try_features is not None and col in try_features else None
    506             typed_sequence = OptimizedTypedSequence(batch_examples[col], type=col_type, try_type=col_try_type, col=col)
--> 507             arrays.append(pa.array(typed_sequence))
    508             inferred_features[col] = typed_sequence.get_inferred_type()
    509         schema = inferred_features.arrow_schema if self.pa_writer is None else self.schema

/usr/local/lib/python3.7/dist-packages/pyarrow/array.pxi in pyarrow.lib.array()

/usr/local/lib/python3.7/dist-packages/pyarrow/array.pxi in pyarrow.lib._handle_arrow_array_protocol()

[/usr/local/lib/python3.7/dist-packages/datasets/arrow_writer.py](https://localhost:8080/#) in __arrow_array__(self, type)
    175                     storage = list_of_np_array_to_pyarrow_listarray(data, type=pa_type.value_type)
    176                 else:
--> 177                     storage = pa.array(data, pa_type.storage_dtype)
    178                 return pa.ExtensionArray.from_storage(pa_type, storage)
    179 

/usr/local/lib/python3.7/dist-packages/pyarrow/array.pxi in pyarrow.lib.array()

/usr/local/lib/python3.7/dist-packages/pyarrow/array.pxi in pyarrow.lib._sequence_to_array()

/usr/local/lib/python3.7/dist-packages/pyarrow/error.pxi in pyarrow.lib.pyarrow_internal_check_status()

/usr/local/lib/python3.7/dist-packages/pyarrow/error.pxi in pyarrow.lib.check_status()

ArrowInvalid: Can only convert 1-dimensional array values
```

**Describe the solution you'd like**
No error in the second scenario and an identical result to the following snippets.

**Describe alternatives you've considered**
There are other alternatives that work such as:
```python

def prepare_dataset_3D_bis(batch):
    batch["pixel_values"] = [[data_map[index].tolist()] for index in batch["id"]]
    return batch

ds_3D_bis = ds.map(
    prepare_dataset_3D_bis, 
    batched=True, 
    remove_columns=ds.column_names, 
    features=features.Features({"pixel_values": features.Array3D(shape=(1, 2, 3), dtype="float32")})
)
```
or
```python
def prepare_dataset_3D_ter(batch):
    batch["pixel_values"] = [data_map[index][np.newaxis, :, :] for index in batch["id"]]
    return batch

ds_3D_ter = ds.map(
    prepare_dataset_3D_ter, 
    batched=True, 
    remove_columns=ds.column_names, 
    features=features.Features({"pixel_values": features.Array3D(shape=(1,  2, 3), dtype="float32")})
)
```
But both solutions require the user to be aware that `data_map[index]` is an `np.array` type.

cc @lhoestq as we discuss this offline :smile: 

Comments:
Hi @SaulLu, thanks for your proposal.

Just I got a bit confused about the dimensions...
- For the 2D case, you mention it is possible to create an `Array2D` from a list of arrays of dimension 1
- However, you give an example of creating an `Array2D` from arrays of dimension 2:
  - the values of `data_map` are arrays of dimension 2
  - the outer list in `prepare_dataset_2D` should not be taken into account in the dimension counting, as it is used because in `map` you pass `batched=True`

Note that for the 3D alternatives you mention:
- In `prepare_dataset_3D_ter`, you create an `Array3D` from arrays of dimension 3:
  - the array `data_map[index][np.newaxis, :, :]` has dimension 3
  - the outer list in `prepare_dataset_3D_ter` is the one used by `batched=True`
- In `prepare_dataset_3D_bis`, you create an `Array3D` from a list of list of lists:
  - the value of `data_map[index].tolist()` is a list of lists
  - it is enclosed by another list `[data_map[index].tolist()]`, thus giving a list of list of lists
  - the outer list is the one used by `batched=True`

Therefore, if I understand correctly, your request would be to be able to create an `Array3D` from a list of an array of dimension 2:
- In `prepare_dataset_3D`, `data_map[index]` is an array of dimension 2
- it is enclosed by a list `[data_map[index]]`, thus giving a list of an array of dimension 2
- the outer list is the one used by `batched=True`

Please, feel free to tell me if I did not understand you correctly.

Hi @albertvillanova ,

Indeed my message was confusing and you guessed right :smile: : I think would be interesting to be able to create an Array3D from a list of an array of dimension 2. 

For the 2D case I should have given as a "similar" example:
```python

data_map_1D = {
    1: np.array([0.2, 0.4]),
    2: np.array([0.1, 0.4]),
}

def prepare_dataset_2D(batch):
    batch["pixel_values"] = [[data_map_1D[index]] for index in batch["id"]]
    return batch
    
ds_2D = ds.map(
    prepare_dataset_2D, 
    batched=True, 
    remove_columns=ds.column_names, 
    features=features.Features({"pixel_values": features.Array2D(shape=(1, 2), dtype="float32")})
)
```