Passing dask array to zarr leads to a TypeError

[michael-sutherland]

Minimal, reproducible code sample, a copy-pastable example if possible

import numpy as np
import zarr
import os
import dask.array as da

from ome_zarr.io import parse_url
from ome_zarr.writer import write_image

path = "test_dask_image.zarr"
os.mkdir(path)

mean_val = 10
rng = np.random.default_rng(0)
data = rng.poisson(mean_val, size=(50, 128, 128)).astype(np.uint8)
data = da.from_array(data)  #  <--- !!!THIS IS WHAT BREAKS THE EXAMPLE!!!

# write the image data
store = parse_url(path, mode="w").store
root = zarr.group(store=store)
write_image(image=data, group=root, chunks=(1, 128, 128), axes="zyx")
root.attrs["omero"] = {
    "channels": [{
        "color": "00FFFF",
        "window": {"start": 0, "end": 20},
        "label": "random",
        "active": True,
    }]
}

Problem description

Get an error when trying to write a dask array to a zarr file. I took the proposed example from ome/ome-zarr-py#121 and wrapped the numpy array in a dask array. It appears that the dask array function astype doesn't support the parameter "order". The example code works when using a numpy array for me. Here's the resulting traceback:

Traceback (most recent call last):
  File "/home/sutherland/git/xray_napari/write_zarr_with_dask_failure.py", line 20, in <module>
    write_image(image=data, group=root, chunks=(1, 128, 128), axes="zyx")
  File "/home/sutherland/anaconda3/envs/napari2/lib/python3.9/site-packages/ome_zarr/writer.py", line 178, in write_image
    write_multiscale(image, group, chunks=chunks, fmt=fmt, axes=axes)
  File "/home/sutherland/anaconda3/envs/napari2/lib/python3.9/site-packages/ome_zarr/writer.py", line 105, in write_multiscale
    group.create_dataset(str(path), data=dataset, chunks=chunks)
  File "/home/sutherland/anaconda3/envs/napari2/lib/python3.9/site-packages/zarr/hierarchy.py", line 808, in create_dataset
    return self._write_op(self._create_dataset_nosync, name, **kwargs)
  File "/home/sutherland/anaconda3/envs/napari2/lib/python3.9/site-packages/zarr/hierarchy.py", line 661, in _write_op
    return f(*args, **kwargs)
  File "/home/sutherland/anaconda3/envs/napari2/lib/python3.9/site-packages/zarr/hierarchy.py", line 824, in _create_dataset_nosync
    a = array(data, store=self._store, path=path, chunk_store=self._chunk_store,
  File "/home/sutherland/anaconda3/envs/napari2/lib/python3.9/site-packages/zarr/creation.py", line 377, in array
    z[...] = data
  File "/home/sutherland/anaconda3/envs/napari2/lib/python3.9/site-packages/zarr/core.py", line 1224, in __setitem__
    self.set_basic_selection(selection, value, fields=fields)
  File "/home/sutherland/anaconda3/envs/napari2/lib/python3.9/site-packages/zarr/core.py", line 1319, in set_basic_selection
    return self._set_basic_selection_nd(selection, value, fields=fields)
  File "/home/sutherland/anaconda3/envs/napari2/lib/python3.9/site-packages/zarr/core.py", line 1610, in _set_basic_selection_nd
    self._set_selection(indexer, value, fields=fields)
  File "/home/sutherland/anaconda3/envs/napari2/lib/python3.9/site-packages/zarr/core.py", line 1682, in _set_selection
    self._chunk_setitems(lchunk_coords, lchunk_selection, chunk_values,
  File "/home/sutherland/anaconda3/envs/napari2/lib/python3.9/site-packages/zarr/core.py", line 1871, in _chunk_setitems
    cdatas = [self._process_for_setitem(key, sel, val, fields=fields)
  File "/home/sutherland/anaconda3/envs/napari2/lib/python3.9/site-packages/zarr/core.py", line 1871, in <listcomp>
    cdatas = [self._process_for_setitem(key, sel, val, fields=fields)
  File "/home/sutherland/anaconda3/envs/napari2/lib/python3.9/site-packages/zarr/core.py", line 1924, in _process_for_setitem
    chunk = value.astype(self._dtype, order=self._order, copy=False)
  File "/home/sutherland/anaconda3/envs/napari2/lib/python3.9/site-packages/dask/array/core.py", line 2086, in astype
    raise TypeError(
TypeError: astype does not take the following keyword arguments: ['order']

Version and installation information

Please provide the following:

• Value of zarr.__version__: 2.10.3
• Value of numcodecs.__version__: 0.9.1
• Version of Python interpreter: 3.9.7
• Operating system (Linux/Windows/Mac): RHEL 8
• How Zarr was installed (e.g., "using pip into virtual environment", or "using conda"): using pip into conda env

Also, if you think it might be relevant, please provide the output from pip freeze or
conda env export depending on which was used to install Zarr.

[joshmoore]

Hi @michael-sutherland. Can you include which version of dask?

[michael-sutherland]

dask: 2021.12.0

[michael-sutherland]

numpy docs for astype: https://numpy.org/doc/stable/reference/generated/numpy.ndarray.astype.html
dask docs for astype: https://docs.dask.org/en/latest/generated/dask.array.Array.astype.html

[joshmoore]

Thanks for the info, @michael-sutherland. Looking at the stacktrace again, any objections to migrating this to github.com/ome/ome-zarr-py ?

cc: @sbesson @will-moore

[michael-sutherland]

Oops! Sorry if I misidentified which library the problem was in. Would you like me to migrate the issue or would you prefer to do it?

[joshmoore]

No worries. Hmmm.... looks like I can't transfer automatically between orgs anyway. See ome/ome-zarr-py#169

[will-moore]

As mentioned on #962, the example can be simplified to remove ome_zarr and get the same exception:

import numpy as np
import zarr
from zarr.storage import FSStore
import os
import dask.array as da

path = "test_dask_image.zarr"
os.mkdir(path)

mean_val = 10
rng = np.random.default_rng(0)
data = rng.poisson(mean_val, size=(50, 128, 128)).astype(np.uint8)
data = da.from_array(data)  #  <--- !!!THIS IS WHAT BREAKS THE EXAMPLE!!!

# write the image data
store = FSStore(path, mode="w")
root = zarr.group(store=store)
root.create_dataset("0", data=data, chunks=(1, 128, 128))

[joshmoore]

Updated @will-moore's example to define mean_val = 10.

[joshmoore]

Ok. I assume this is related to @madsbk's #934 in the sense that typically dask arrays are assumed to wrap zarr arrays rather than vice versa.

cc: @jakirkham

[GenevieveBuckley]

Hi @michael-sutherland, is there a reason you don't want to use the dask.array.to_zarr() function?

I imagine there might be some cases where it makes sense, and this example could have lost that context when it was simplified to make debugging easier.

[GenevieveBuckley]

I poked around with this a little, and there seem to be two places where things go wrong: one in dask, and one in zarr.

1. In dask/array/core.py, in the astype method of the Array class

This line:
https://github.com/dask/dask/blob/cfaf931072ae079313761d4d61e25956029abf6c/dask/array/core.py#L2237

If it were changed to this, I think we'd fix the dask part of the problem:

extra = set(kwargs) - {"casting", "copy", "order"}

2. In zarr, zarr/core.py(2189)_encode_chunk()
   There is an ensure_ndarray function, which checks to see if the chunk can be written properly (this check happens with numcodecs). But, because chunk here is a chunk of the dask array instead of a numpy array, the numcodecs library chokes on it. What it sees in the buffer is not a numpy array but instead something different, and so it throws an error.

> /Users/genevieb/mambaforge/envs/dask-dev/lib/python3.9/site-packages/zarr/core.py(2189)_encode_chunk()
   2187
   2188         # check object encoding
-> 2189         if ensure_ndarray(chunk).dtype == object:
   2190             raise RuntimeError('cannot write object array without object codec')

Zarr could potentially work around this second issue by trying to .compute() the dask chunk in ensure_ndarray, before passing it along. But then there's some question about whether that's actually a sensible thing to do or not. Things would probably also become uglier in the zarr code, if now zarr needs to check what type of array it is handling, especially since numpy and dask are not the only two options for array types.

Josh says above that typically it is expected dask will wrap zarr, and not the other way around. Making changes to (2) above would be a bit in conflict with that expectation.

[vietnguyengit]

Hi guys, Google brings me here, I recently face the same error as well and I used dask.array.to_zarr() method: