Reductions with on disk operands#

In this benchmark we will compare ironArray’s performance with Zarr, HDF5, and TileDB. We will work with the data used in the Reductions tutorial but this time it will be on disk. We will compute the same operations as in the previous benchmark. We will use the default chunk shapes for every library; this is important to stress that defaults for every library are typically fine tuned for general performance; so this is a pretty fair way to compare them.

Let’s go:

[1]:
%load_ext memprofiler

import numpy as np
import iarray as ia
import os
import zarr
import h5py
import hdf5plugin as h5plugin
import tiledb
import dask.array as da
from numcodecs import blosc, Blosc
[2]:
!du -sh ../tutorials/precip-3m.iarr
809M    ../tutorials/precip-3m.iarr

ironArray#

[3]:
ia_precip_d = ia.open("../tutorials/precip-3m.iarr")
chunks = ia_precip_d.chunks
shape = ia_precip_d.shape
[4]:
!vmtouch -e "../tutorials/precip-3m.iarr"
           Files: 1
     Directories: 0
   Evicted Pages: 207031 (808M)
         Elapsed: 0.042612 seconds
[5]:
%%mprof_run 1.iarray::mean_disk
ia_reduc_d = ia.mean(ia_precip_d, axis=(3, 2, 0)).data
memprofiler: used 22.84 MiB RAM (peak of 26.87 MiB) in 0.8316 s, total RAM usage 244.72 MiB

Zarr#

[6]:
%%time
zarr_urlpath = "precip-3m.zarr"
if not os.path.exists(zarr_urlpath):
    precip_zarr_d = zarr.empty(ia_precip_d.shape, dtype=ia_precip_d.dtype, store=zarr_urlpath)
    ia_precip_d.copyto(precip_zarr_d)
CPU times: user 9 µs, sys: 4 µs, total: 13 µs
Wall time: 15.5 µs
[7]:
precip_zarr_d = zarr.open(zarr_urlpath)
print(precip_zarr_d.info)
precip_zdask_d = da.from_zarr(precip_zarr_d)
Type               : zarr.core.Array
Data type          : float32
Shape              : (3, 720, 721, 1440)
Chunk shape        : (1, 90, 91, 180)
Order              : C
Read-only          : False
Compressor         : Blosc(cname='lz4', clevel=5, shuffle=SHUFFLE, blocksize=0)
Store type         : zarr.storage.DirectoryStore
No. bytes          : 8970393600 (8.4G)
No. bytes stored   : 1021932469 (974.6M)
Storage ratio      : 8.8
Chunks initialized : 1536/1536

[8]:
!vmtouch -e $zarr_urlpath
           Files: 1537
     Directories: 1
   Evicted Pages: 250262 (977M)
         Elapsed: 0.042899 seconds
[9]:
%%mprof_run 2.zarr::mean_disk
zdask_reduc_d = da.mean(precip_zdask_d, axis=(3, 2, 0))
zarr_reduc_d = zarr.create(shape=ia_reduc_d.shape, dtype=ia_reduc_d.dtype)
da.to_zarr(zdask_reduc_d, zarr_reduc_d)
memprofiler: used 531.14 MiB RAM (peak of 631.49 MiB) in 1.4725 s, total RAM usage 776.56 MiB
[10]:
del precip_zarr_d
del precip_zdask_d
del zdask_reduc_d
del zarr_reduc_d

HDF5#

[11]:
h5_urlpath = "precip-3m.hdf5"
if not os.path.exists(h5_urlpath):
    precip_h5_d = h5py.File(h5_urlpath, "w")
    dset = f.create_dataset("h5_precip", data=precip_zarr_d.data)

h5_file_d = h5py.File(h5_urlpath, "r")
precip_h5_d = h5_file_d['h5_precip']
precip_h5dask = da.from_array(precip_h5_d)
[12]:
!vmtouch -e $h5_urlpath
           Files: 1
     Directories: 0
   Evicted Pages: 277358 (1G)
         Elapsed: 0.041419 seconds
[13]:
%%mprof_run 3.hdf5::mean_disk
h5dask_reduc_d = da.mean(precip_h5dask, axis=(3, 2, 0))

f = h5py.File("reduc_d.hdf5", "w", driver='core', backing_store=False)
h5_reduc_d = f.create_dataset(name="reduc_d.hdf5", shape=ia_reduc_d.shape, dtype=ia_reduc_d.dtype,
                              **h5plugin.Blosc())
da.to_hdf5("reduc_d.hdf5", '/x', h5dask_reduc_d)
memprofiler: used 379.99 MiB RAM (peak of 3205.50 MiB) in 5.9929 s, total RAM usage 1138.83 MiB
[14]:
del h5_file_d
del precip_h5_d
del precip_h5dask
del h5dask_reduc_d

TileDB#

[15]:
# TileDB does not come with defaults for chunk shape (tiles), so using ironArray's one
shape = ia_precip_d.shape
tiling= ia_precip_d.chunks
adom = tiledb.Domain(
        tiledb.Dim(name="rows", domain=(0, shape[0] - 1), dtype=np.int32, tile=chunks[0]),
        tiledb.Dim(name="cols", domain=(0, shape[1] - 1), dtype=np.int32, tile=chunks[1]),
        tiledb.Dim(name="else", domain=(0, shape[2] - 1), dtype=np.int32, tile=chunks[2]),
        tiledb.Dim(name="else2", domain=(0, shape[3] - 1), dtype=np.int32, tile=chunks[3]),
    )

filters = tiledb.FilterList([tiledb.ByteShuffleFilter(), tiledb.LZ4Filter(5)])
aschema = tiledb.ArraySchema(
    domain=adom,  sparse=False, attrs=[tiledb.Attr(name="a", dtype=np.float32, filters=filters)]
)

# Create the (empty) array on disk.
tiledb_urlpath = "precip-3m.tiledb"
if not os.path.exists(tiledb_urlpath):
    tiledb.DenseArray.create(tiledb_urlpath, aschema)
    with tiledb.DenseArray(tiledb_urlpath, mode="w") as A_opt:
        A_opt[:] = ia_precip_d.data

dask_tiledb = da.from_tiledb(tiledb_urlpath, attribute='a')
[16]:
!vmtouch -e $tiledb_urlpath
           Files: 5
     Directories: 4
   Evicted Pages: 275892 (1G)
         Elapsed: 0.044538 seconds
[17]:
%%mprof_run 4.tiledb::mean_disk
res_dask_opt = da.mean(dask_tiledb, axis=(3,2,0))
res_dask_opt.to_tiledb("mem://res_tiledb")
memprofiler: used 426.24 MiB RAM (peak of 441.90 MiB) in 4.8878 s, total RAM usage 1570.53 MiB
[18]:
del res_dask_opt
del dask_tiledb
tiledb.remove("mem://res_tiledb")

Results#

Finally, we will do a plot of the memory and time consumption with each different library.

[19]:
%mprof_plot .*::mean_disk