Title: Column of arrays (2D+) are using unreasonably high memory
Author: vigsterkr (CONTRIBUTOR)
State: open
Created: 2022-09-10T13:07:22Z
Labels: bug

## Describe the bug
When trying to store `Array2D, Array3D, etc` as column values in a dataset, accessing that column (or creating depending on how you create it, see code below) will cause more than 10 fold of memory usage.

## Steps to reproduce the bug
```python
from datasets import Dataset, Features, Array2D, Array3D
import numpy as np

column_name = "a"
array_shape = (64, 64, 3)

data = np.random.random((10000,) + array_shape)
dataset = Dataset.from_dict({column_name: data}, features=Features({column_name: Array3D(shape=array_shape, dtype="float64")}))
```

the code above will use about 10Gb of RAM while constructing the `dataset` object.

The code below will use roughly the same amount of memory (and time) when trying to actually access the data itself of that column.
```python
from datasets import Dataset
import numpy as np

column_name = "a"
array_shape = (64, 64, 3)

data = np.random.random((10000,) + array_shape)
dataset = Dataset.from_dict({column_name: data})
dataset[column_name]
```

## Expected results
Some memory overhead, but not like as it is now and certainly not an overhead of such runtime that is currently happening.

## Actual results
Enormous memory- and runtime overhead.

## Environment info
- `datasets` version: 2.3.2
- Platform: macOS-12.5.1-arm64-arm-64bit
- Python version: 3.8.13
- PyArrow version: 9.0.0
- Pandas version: 1.4.4

Comments:
note i have tried the same code with `datasets` version 2.4.0, the outcome is the very same as described above.

Seems related to issues #4623 and #4802 so it would appear this issue has been around for a few months.

Hi ! `Dataset.from_dict` keeps the data in memory. You can write on disk and reload them with
```python
dataset.save_to_disk("path/to/local")
dataset = load_from_disk("path/to/local")
```
this way you'll end up with a dataset loaded from your disk using memory mapping, and it won't fill up your RAM :)

related to https://github.com/huggingface/datasets/issues/4861

@lhoestq thnx for getting back to me! i've tested the suggested method, but unfortunately the memory consumption is the very same:

```
from datasets import Dataset, Features, Array2D, Array3D, load_from_disk
import numpy as np

column_name = "a"
array_shape = (64, 64, 3)

data = np.random.random((10000,) + array_shape)
dataset = Dataset.from_dict({column_name: data}, features=Features({column_name: Array3D(shape=array_shape, dtype="float64")}))
dataset.save_to_disk("foo")

foo_db = load_from_disk("foo")
colum_value = foo_db[column_name]
```

the very same happens when you create the dataset, but dont specify the feature type.

i've tried running this on different envs (macOS, linux) and it's behaving the very same way.

When you call `colum_value = foo_db[column_name]`, you load the full column in memory.

If you want to avoid filling up your memory, you can access chunks of data instead
```python
embeddings = dataset[i:i + chunk_size]["embeddings"]
```

@lhoestq yeah that's intentional, i.e. i really want to load the whole column into the memory. but as said above there's an unreasonable amount of overhead for the memory. the np array itself is using about 1G of memory:
```
>>> getsizeof(data)/1024/1024
937.5001525878906
```
that accessing of column above is using 10x memory compared to the original numpy array.

The dataset must be twice as big because we use regular arrow ListArray under the hood and not FixedSizeListArray. Basically we store unnecessary offsets.

And this should affect performance as well. When we developed this, FixedSizeListArray still had some issues but they should be resolved on the PyArrow side now

A doubling would be fine. My very basic understanding of PyArrow is that using ListArray is probably related to the issue though. Using a multi-dimensional array in datasets is storing everything as strange nested 1d object arrays, which I imagine is creating the massive overhead.

I think it should be a PyArrow Tensor, no?

PyArrow tensors are not part of the Arrow format AFAIK:

> There is no direct support in the arrow columnar format to store Tensors as column values.

source: https://github.com/apache/arrow/issues/4802#issuecomment-508494694

That's... unfortunate. I didn't realize that.