calculate histogram of a hugh dask xarray dataarray

[Jing25]

Hi, I'm calculating the histogram of a huge 3D dataarray (day, lat, lon). If the number of bins is 1000, then the output would be a 1000 x 1 vector.

I load data from the cloud:

ds = xr.open_dataset(fsmapper,
                     engine='zarr',
                     decode_coords='all')

rechunk:
ds = ds.chunk({'y': 720, 'x': 720})

I tested on a small subset with a few things such as using apply_ufunc:

bins = np.linspace(0, 1, 100)
mm = xr.apply_ufunc(partial(da.histogram, bins), ds.isel(y = slice(0, 3*720), x = slice(0, 3*720)), dask='parallelized', 
                    output_dtypes='float32')

This doesn't work.

I tried us dask.array.histogram directly, and it doesn't work.

da.histogram(ds.isel(y = slice(0, 3*720), x = slice(0, 3*720)), np.linspace(0, 1, 100))

I'm wondering what a good way would be for me to calculate the histogram of a huge xaray dataarray ~(TB)?

[Jing25]

Here's the error message I got:

distributed.protocol.core - CRITICAL - Failed to Serialize
Traceback (most recent call last):
File "/anaconda/envs/rapids-21.10-py3.8/lib/python3.8/site-packages/distributed/protocol/core.py", line 76, in dumps
frames[0] = msgpack.dumps(msg, default=_encode_default, use_bin_type=True)
File "/anaconda/envs/rapids-21.10-py3.8/lib/python3.8/site-packages/msgpack/init.py", line 35, in packb
return Packer(**kwargs).pack(o)
File "msgpack/_packer.pyx", line 292, in msgpack._cmsgpack.Packer.pack
File "msgpack/_packer.pyx", line 298, in msgpack._cmsgpack.Packer.pack
File "msgpack/_packer.pyx", line 295, in msgpack._cmsgpack.Packer.pack
File "msgpack/_packer.pyx", line 264, in msgpack._cmsgpack.Packer._pack
File "msgpack/_packer.pyx", line 231, in msgpack._cmsgpack.Packer._pack
File "msgpack/_packer.pyx", line 231, in msgpack._cmsgpack.Packer._pack
File "msgpack/_packer.pyx", line 264, in msgpack._cmsgpack.Packer._pack
File "msgpack/_packer.pyx", line 231, in msgpack._cmsgpack.Packer._pack
File "msgpack/_packer.pyx", line 231, in msgpack._cmsgpack.Packer._pack
File "msgpack/_packer.pyx", line 231, in msgpack._cmsgpack.Packer._pack
File "msgpack/_packer.pyx", line 285, in msgpack._cmsgpack.Packer._pack
File "/anaconda/envs/rapids-21.10-py3.8/lib/python3.8/site-packages/distributed/protocol/core.py", line 57, in _encode_default
sub_header, sub_frames = serialize_and_split(
File "/anaconda/envs/rapids-21.10-py3.8/lib/python3.8/site-packages/distributed/protocol/serialize.py", line 425, in serialize_and_split
header, frames = serialize(x, serializers, on_error, context)
File "/anaconda/envs/rapids-21.10-py3.8/lib/python3.8/site-packages/distributed/protocol/serialize.py", line 252, in serialize
return serialize(
File "/anaconda/envs/rapids-21.10-py3.8/lib/python3.8/site-packages/distributed/protocol/serialize.py", line 349, in serialize
raise TypeError(msg, str(x)[:10000])
TypeError: ('Could not serialize object of type ImplicitToExplicitIndexingAdapter.', "ImplicitToExplicitIndexingAdapter(array=MemoryCachedArray(array=CopyOnWriteArray(array=LazilyIndexedArray(array=_ElementwiseFunctionArray(_ElementwiseFunctionArray(LazilyIndexedArray(array=<xarray.backends.zarr.ZarrArrayWrapper object at 0x7fa80296cfc0>, key=BasicIndexer((slice(None, None, None), slice(None, None, None), slice(None, None, None)))), func=functools.partial(<function _apply_mask at 0x7fa8a21338b0>, encoded_fill_values={0}, decoded_fill_value=nan, dtype=dtype('float32')), dtype=dtype('float32')), func=functools.partial(<function _scale_offset_decoding at 0x7fa8a2133940>, scale_factor=0.003937007874015748, add_offset=-0.003937007874015748, dtype=<class 'numpy.float32'>), dtype=dtype('float32')), key=BasicIndexer((slice(None, None, None), slice(None, None, None), slice(None, None, None)))))))")
distributed.comm.utils - ERROR - ('Could not serialize object of type ImplicitToExplicitIndexingAdapter.', "ImplicitToExplicitIndexingAdapter(array=MemoryCachedArray(array=CopyOnWriteArray(array=LazilyIndexedArray(array=_ElementwiseFunctionArray(_ElementwiseFunctionArray(LazilyIndexedArray(array=<xarray.backends.zarr.ZarrArrayWrapper object at 0x7fa80296cfc0>, key=BasicIndexer((slice(None, None, None), slice(None, None, None), slice(None, None, None)))), func=functools.partial(<function _apply_mask at 0x7fa8a21338b0>, encoded_fill_values={0}, decoded_fill_value=nan, dtype=dtype('float32')), dtype=dtype('float32')), func=functools.partial(<function _scale_offset_decoding at 0x7fa8a2133940>, scale_factor=0.003937007874015748, add_offset=-0.003937007874015748, dtype=<class 'numpy.float32'>), dtype=dtype('float32')), key=BasicIndexer((slice(None, None, None), slice(None, None, None), slice(None, None, None)))))))")
Traceback (most recent call last):
File "/anaconda/envs/rapids-21.10-py3.8/lib/python3.8/site-packages/distributed/comm/utils.py", line 33, in _to_frames
return list(protocol.dumps(msg, **kwargs))
File "/anaconda/envs/rapids-21.10-py3.8/lib/python3.8/site-packages/distributed/protocol/core.py", line 76, in dumps
frames[0] = msgpack.dumps(msg, default=_encode_default, use_bin_type=True)
File "/anaconda/envs/rapids-21.10-py3.8/lib/python3.8/site-packages/msgpack/init.py", line 35, in packb
return Packer(**kwargs).pack(o)
File "msgpack/_packer.pyx", line 292, in msgpack._cmsgpack.Packer.pack
File "msgpack/_packer.pyx", line 298, in msgpack._cmsgpack.Packer.pack
File "msgpack/_packer.pyx", line 295, in msgpack._cmsgpack.Packer.pack
File "msgpack/_packer.pyx", line 264, in msgpack._cmsgpack.Packer._pack
File "msgpack/_packer.pyx", line 231, in msgpack._cmsgpack.Packer._pack
File "msgpack/_packer.pyx", line 231, in msgpack._cmsgpack.Packer._pack
File "msgpack/_packer.pyx", line 264, in msgpack._cmsgpack.Packer._pack
File "msgpack/_packer.pyx", line 231, in msgpack._cmsgpack.Packer._pack
File "msgpack/_packer.pyx", line 231, in msgpack._cmsgpack.Packer._pack
File "msgpack/_packer.pyx", line 231, in msgpack._cmsgpack.Packer._pack
File "msgpack/_packer.pyx", line 285, in msgpack._cmsgpack.Packer._pack
File "/anaconda/envs/rapids-21.10-py3.8/lib/python3.8/site-packages/distributed/protocol/core.py", line 57, in _encode_default
sub_header, sub_frames = serialize_and_split(
File "/anaconda/envs/rapids-21.10-py3.8/lib/python3.8/site-packages/distributed/protocol/serialize.py", line 425, in serialize_and_split
header, frames = serialize(x, serializers, on_error, context)
File "/anaconda/envs/rapids-21.10-py3.8/lib/python3.8/site-packages/distributed/protocol/serialize.py", line 252, in serialize
return serialize(
File "/anaconda/envs/rapids-21.10-py3.8/lib/python3.8/site-packages/distributed/protocol/serialize.py", line 349, in serialize
raise TypeError(msg, str(x)[:10000])
TypeError: ('Could not serialize object of type ImplicitToExplicitIndexingAdapter.', "ImplicitToExplicitIndexingAdapter(array=MemoryCachedArray(array=CopyOnWriteArray(array=LazilyIndexedArray(array=_ElementwiseFunctionArray(_ElementwiseFunctionArray(LazilyIndexedArray(array=<xarray.backends.zarr.ZarrArrayWrapper object at 0x7fa80296cfc0>, key=BasicIndexer((slice(None, None, None), slice(None, None, None), slice(None, None, None)))), func=functools.partial(<function _apply_mask at 0x7fa8a21338b0>, encoded_fill_values={0}, decoded_fill_value=nan, dtype=dtype('float32')), dtype=dtype('float32')), func=functools.partial(<function _scale_offset_decoding at 0x7fa8a2133940>, scale_factor=0.003937007874015748, add_offset=-0.003937007874015748, dtype=<class 'numpy.float32'>), dtype=dtype('float32')), key=BasicIndexer((slice(None, None, None), slice(None, None, None), slice(None, None, None)))))))")
distributed.batched - ERROR - Error in batched write
Traceback (most recent call last):
File "/anaconda/envs/rapids-21.10-py3.8/lib/python3.8/site-packages/distributed/batched.py", line 93, in _background_send
nbytes = yield self.comm.write(
File "/anaconda/envs/rapids-21.10-py3.8/lib/python3.8/site-packages/tornado/gen.py", line 762, in run
value = future.result()
File "/anaconda/envs/rapids-21.10-py3.8/lib/python3.8/site-packages/distributed/comm/tcp.py", line 243, in write
frames = await to_frames(
File "/anaconda/envs/rapids-21.10-py3.8/lib/python3.8/site-packages/distributed/comm/utils.py", line 50, in to_frames
return _to_frames()
File "/anaconda/envs/rapids-21.10-py3.8/lib/python3.8/site-packages/distributed/comm/utils.py", line 33, in _to_frames
return list(protocol.dumps(msg, **kwargs))
File "/anaconda/envs/rapids-21.10-py3.8/lib/python3.8/site-packages/distributed/protocol/core.py", line 76, in dumps
frames[0] = msgpack.dumps(msg, default=_encode_default, use_bin_type=True)
File "/anaconda/envs/rapids-21.10-py3.8/lib/python3.8/site-packages/msgpack/init.py", line 35, in packb
return Packer(**kwargs).pack(o)
File "msgpack/_packer.pyx", line 292, in msgpack._cmsgpack.Packer.pack
File "msgpack/_packer.pyx", line 298, in msgpack._cmsgpack.Packer.pack
File "msgpack/_packer.pyx", line 295, in msgpack._cmsgpack.Packer.pack
File "msgpack/_packer.pyx", line 264, in msgpack._cmsgpack.Packer._pack
File "msgpack/_packer.pyx", line 231, in msgpack._cmsgpack.Packer._pack
File "msgpack/_packer.pyx", line 231, in msgpack._cmsgpack.Packer._pack
File "msgpack/_packer.pyx", line 264, in msgpack._cmsgpack.Packer._pack
File "msgpack/_packer.pyx", line 231, in msgpack._cmsgpack.Packer._pack
File "msgpack/_packer.pyx", line 231, in msgpack._cmsgpack.Packer._pack
File "msgpack/_packer.pyx", line 231, in msgpack._cmsgpack.Packer._pack
File "msgpack/_packer.pyx", line 285, in msgpack._cmsgpack.Packer._pack
File "/anaconda/envs/rapids-21.10-py3.8/lib/python3.8/site-packages/distributed/protocol/core.py", line 57, in _encode_default
sub_header, sub_frames = serialize_and_split(
File "/anaconda/envs/rapids-21.10-py3.8/lib/python3.8/site-packages/distributed/protocol/serialize.py", line 425, in serialize_and_split
header, frames = serialize(x, serializers, on_error, context)
File "/anaconda/envs/rapids-21.10-py3.8/lib/python3.8/site-packages/distributed/protocol/serialize.py", line 252, in serialize
return serialize(
File "/anaconda/envs/rapids-21.10-py3.8/lib/python3.8/site-packages/distributed/protocol/serialize.py", line 349, in serialize
raise TypeError(msg, str(x)[:10000])
TypeError: ('Could not serialize object of type ImplicitToExplicitIndexingAdapter.', "ImplicitToExplicitIndexingAdapter(array=MemoryCachedArray(array=CopyOnWriteArray(array=LazilyIndexedArray(array=_ElementwiseFunctionArray(_ElementwiseFunctionArray(LazilyIndexedArray(array=<xarray.backends.zarr.ZarrArrayWrapper object at 0x7fa80296cfc0>, key=BasicIndexer((slice(None, None, None), slice(None, None, None), slice(None, None, None)))), func=functools.partial(<function _apply_mask at 0x7fa8a21338b0>, encoded_fill_values={0}, decoded_fill_value=nan, dtype=dtype('float32')), dtype=dtype('float32')), func=functools.partial(<function _scale_offset_decoding at 0x7fa8a2133940>, scale_factor=0.003937007874015748, add_offset=-0.003937007874015748, dtype=<class 'numpy.float32'>), dtype=dtype('float32')), key=BasicIndexer((slice(None, None, None), slice(None, None, None), slice(None, None, None)))))))")

---

CancelledError Traceback (most recent call last)
/tmp/ipykernel_158854/709552324.py in
----> 1 mm[0].compute()

/anaconda/envs/rapids-21.10-py3.8/lib/python3.8/site-packages/dask/base.py in compute(self, **kwargs)
286 dask.base.compute
287 """
--> 288 (result,) = compute(self, traverse=False, **kwargs)
289 return result
290

/anaconda/envs/rapids-21.10-py3.8/lib/python3.8/site-packages/dask/base.py in compute(*args, **kwargs)
568 postcomputes.append(x.dask_postcompute())
569
--> 570 results = schedule(dsk, keys, **kwargs)
571 return repack([f(r, *a) for r, (f, a) in zip(results, postcomputes)])
572

/anaconda/envs/rapids-21.10-py3.8/lib/python3.8/site-packages/distributed/client.py in get(self, dsk, keys, workers, allow_other_workers, resources, sync, asynchronous, direct, retries, priority, fifo_timeout, actors, **kwargs)
2687 should_rejoin = False
2688 try:
-> 2689 results = self.gather(packed, asynchronous=asynchronous, direct=direct)
2690 finally:
2691 for f in futures.values():

/anaconda/envs/rapids-21.10-py3.8/lib/python3.8/site-packages/distributed/client.py in gather(self, futures, errors, direct, asynchronous)
1964 else:
1965 local_worker = None
-> 1966 return self.sync(
1967 self._gather,
1968 futures,

/anaconda/envs/rapids-21.10-py3.8/lib/python3.8/site-packages/distributed/client.py in sync(self, func, asynchronous, callback_timeout, *args, **kwargs)
858 return future
859 else:
--> 860 return sync(
861 self.loop, func, *args, callback_timeout=callback_timeout, **kwargs
862 )

/anaconda/envs/rapids-21.10-py3.8/lib/python3.8/site-packages/distributed/utils.py in sync(loop, func, callback_timeout, *args, **kwargs)
324 if error[0]:
325 typ, exc, tb = error[0]
--> 326 raise exc.with_traceback(tb)
327 else:
328 return result[0]

/anaconda/envs/rapids-21.10-py3.8/lib/python3.8/site-packages/distributed/utils.py in f()
307 if callback_timeout is not None:
308 future = asyncio.wait_for(future, callback_timeout)
--> 309 result[0] = yield future
310 except Exception:
311 error[0] = sys.exc_info()

/anaconda/envs/rapids-21.10-py3.8/lib/python3.8/site-packages/tornado/gen.py in run(self)
760
761 try:
--> 762 value = future.result()
763 except Exception:
764 exc_info = sys.exc_info()

/anaconda/envs/rapids-21.10-py3.8/lib/python3.8/site-packages/distributed/client.py in _gather(self, futures, errors, direct, local_worker)
1830 else:
1831 raise exception.with_traceback(traceback)
-> 1832 raise exc
1833 if errors == "skip":
1834 bad_keys.add(key)

CancelledError: ('sum-aggregate-bb4486815ed489a6c514809a10daad71', 0)

[dcherian]

Try https://github.com/xgcm/xhistogram

[Jing25]

Hi @dcherian thank you for your answer! I just tried Xhistogram. It looks very powerful. But I still cannot get it work in my case.

When I apply it on a dask xarray, it works fine without computing it:

If I do h.compute(), I got the same error:

I'm wondering if my dask dataarray is incorrect or?

[TomNicholas]

Have you tried computing a small part of your array, and running xhistogram on that? i.e. checking that it works on a small portion of your data as numpy arrays?

[Jing25]

Hi @TomNicholas, I've tested what you mentioned and a few other things. I found out that this error might come from the Azure blob credential. I loaded data from the Azure blob, and dask seems not compatible with the way I use credentials. So I can't actually use dask here.