title: feat: create an `Array3D` column from a list of arrays of dimension 2
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")})
  )
  ```
body: **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: 