What happened:

The following example produces four 10GB zarr files on disk, chunked into 20 chunks along two axes. The goal is to combine them into a single zarr file, along the first axis. This can be done by concatenating the arrays, then rechunking them before saving as a zarr file.

In each of the four input files, the final chunk on the first axis is smaller than the chunk size, so concatenating them produces non-uniform chunks, which is why the rechunk step is needed, as zarr requires uniform chunk sizes. This diagram gives an idea of what is happening:

The code to do this works in some cases, but uses a lot of memory, and in other cases does not complete since the system memory is overwhelmed.

What you expected to happen:

I would like a way of running this computation with modest memory requirements - e.g. a small multiple of the chunk size.

Minimal Complete Verifiable Example:

import dask
import dask.array as da
from dask.diagnostics import ProgressBar
import zarr

# create data on disk
x = da.zeros((20 * 500_000 + 100_000, 20 * 50), chunks=(500_000, 50))
da.to_zarr(x, "data/x1.zarr", overwrite=True)
da.to_zarr(x, "data/x2.zarr", overwrite=True)
da.to_zarr(x, "data/x3.zarr", overwrite=True)
da.to_zarr(x, "data/x4.zarr", overwrite=True)

# read
x1 = da.from_zarr("data/x1.zarr")
x2 = da.from_zarr("data/x2.zarr")
x3 = da.from_zarr("data/x3.zarr")
x4 = da.from_zarr("data/x4.zarr")

# concatenate
x = da.concatenate([x1, x2, x3, x4], axis=0)

# rechunk uniformly (so we can save as zarr)
x = x.rechunk(x.chunksize)

# create empty store to save as zarr
store = zarr.DirectoryStore("store/out_big")
root = zarr.group(store, overwrite=True)
dest = root.empty_like(name="dest", data=x, chunks=x.chunksize, overwrite=True)

# inline roots of graph, see https://github.com/dask/dask/issues/6668
def custom_optimize(dsk, keys):
    print("running custom optimize")
    dsk = dask.optimization.inline(dsk, inline_constants=True)
    return dask.array.optimization.optimize(dsk, keys)


with dask.config.set(array_optimize=custom_optimize):
    d = x.store(dest, lock=False, compute=False)
    with ProgressBar():
        d.compute()

This example uses well over 40GB of memory (on 4 cores), even though the chunks are 200MB in size. The following shows a portion of the task graph:

Produced with this code:

with dask.config.set(array_optimize=custom_optimize):
    x = x[-3*500_000:,:2*50]
    d = x.store(dest, lock=False, compute=False)
    dask.visualize(d, filename=f"min-example-order", optimize_graph=True, verbose=True, color="order", cmap="autumn", node_attr={"penwidth": "4"})

The numbers in circles are the order the tasks run in (red = earlier, yellow = later). It's clear that the lack of alignment in the chunking causes chains of tasks. So x[n][0] are all linked to each other in one chain, x[n][1] are in another, etc. These chains are processed concurrently (and with the example above there are 20 of them), and each chain will keep a (presumably small) number of zarr chunks in memory. The problem is that some chunking schemes can easily overwhelm memory.

Anything else we need to know?:

I came up with a workaround that uses a "pull" scheme, where the target array is filled in by pulling data from the relevant source chunks. The downside is that each zarr chunk has to be read twice (although it's no slower overall on local disks), but the main advantage is that everything is very parallelizable, and runs with modest memory requirements (<2GB for this example).

z1 = zarr.open("data/x1.zarr")
z2 = zarr.open("data/x2.zarr")
z3 = zarr.open("data/x3.zarr")
z4 = zarr.open("data/x4.zarr")

# concatenate and rechunk uniformly (so we can save as zarr)
x = concatenate_and_rechunk([z1, z2, z3, z4])

The code for concatenate_and_rechunk is here, however it works at the level of zarr files, making it hard to use with xarray for example (I'd like to be able to concatenate and rechunk whole datasets).

Is there a better way of doing this with Dask?

Environment:

• Dask version: 2.30.0
• Python version: 3.7
• Operating System: OS X
• Install method (conda, pip, source): pip