Column of arrays (2D+) are using unreasonably high memory

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

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.

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

[vigsterkr]

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

[Slater-Victoroff]

That’s… unfortunate. I didn’t realize that.