`WritableCFDataStore` realizes variable data when loading with object stored Zarr store

[valeriupredoi]

What happened?

Hi folks, something that I just found out, please see issue below:

Problem class: WritableCFDataStore()

xarray/xarray/backends/common.py

Line 647 in 40119bf

class WritableCFDataStore(AbstractWritableDataStore):

Used case and MRE

I am loading a Zarr store from an S3 bucket; the file is about 500MB compressed (Blosc), but uncompressed data is about 1GB, then I'm passing it to WritableCFDataStore:

import xarray as xr

from xarray.backends.common import WritableCFDataStore


def encode_zarr_file():
    # file is about 500MB compressed Zarr3
    zarr_path = (
        "https://uor-aces-o.s3-ext.jc.rl.ac.uk/"
        "esmvaltool-zarr/cl_Amon_UKESM1-0-LL_ssp370SST-lowNTCF_r1i1p1f2.zarr3"
    )

    time_coder = xr.coders.CFDatetimeCoder(use_cftime=True)
    zarr_xr = xr.open_dataset(
        zarr_path,
        consolidated=True,
        decode_times=time_coder,
        engine="zarr",
        backend_kwargs={},
    )

    variables = zarr_xr.variables
    dts = WritableCFDataStore()
    dts.encode(variables, {})


encode_zarr_file()

MRE observed behaviour:

• Xarray Dataset loaded from Zarr has lazy data (mem about 1 few tens of MBs)
• encoding it with WritableCFDataStore realizes its data, and my memory consumption goes up to 1GB and change; passing attributes makes that even bigger but that's normal

MRE desired behaviour:

• Since the base class WritableCFDataStore is used quite a bit especially for Zarr -> NetCDF4 format conversions (the one I'm after hehe), it would be absolutely brilliant and much desired that the variables' data be kept lazy, and all the operations done by the CF converter/encoder be done on Dask arrays rather than Numpy arrays
• CF-compliance/formatting done lazily is possible, and that's currently implemented in such tools like @davidhassell cf-python

Can I help?

Most definitely! But my knowledge of Xarray is rather limited, so I think I'd best be suited to test a PR than to implement it.

Very many thanks in advance! Cheers 🍻

What did you expect to happen?

No response

Minimal Complete Verifiable Example

MVCE confirmation

• Minimal example — the example is as focused as reasonably possible to demonstrate the underlying issue in xarray.
• Complete example — the example is self-contained, including all data and the text of any traceback.
• Verifiable example — the example copy & pastes into an IPython prompt or Binder notebook, returning the result.
• New issue — a search of GitHub Issues suggests this is not a duplicate.
• Recent environment — the issue occurs with the latest version of xarray and its dependencies.

Relevant log output

Anything else we need to know?

No response

Environment

Details

[kmuehlbauer]

@valeriupredoi Cheers, nice to see you around.

I've tried to follow pp-mo/ncdata#139 to see where you coming from with this. So, I can reproduce this behaviour with your code example. The issue is, that in that case xarray's lazy classes take the data. If the data is finally requested from disc (eg for encoding) this will materialize immediately as you observed.

To really move this into the dask world you would need to add chunks={} or similar to the call of open_dataset.

I've checked:

import xarray as xr

from xarray.backends.common import WritableCFDataStore

def encode_zarr_file(chunks):
    # file is about 500MB compressed Zarr3
    zarr_path = (
        "https://uor-aces-o.s3-ext.jc.rl.ac.uk/"
        "esmvaltool-zarr/cl_Amon_UKESM1-0-LL_ssp370SST-lowNTCF_r1i1p1f2.zarr3"
    )

    time_coder = xr.coders.CFDatetimeCoder(use_cftime=True)
    zarr_xr = xr.open_dataset(
        zarr_path,
        consolidated=True,
        decode_times=time_coder,
        engine="zarr",
        chunks=chunks,
        backend_kwargs={},
    )

    variables = zarr_xr.variables
    dts = WritableCFDataStore()
    dts.encode(variables, {})

%load_ext memory_profiler
%memit encode_zarr_file({})
%memit encode_zarr_file(None)

peak memory: 347.28 MiB, increment: 0.00 MiB
peak memory: 1337.39 MiB, increment: 990.11 MiB

This does compare to your estimates, right?

[valeriupredoi]

@kmuehlbauer great to see you here too, mate! Ah that's brilliant, very very many thanks - I was not aware of the complete change of behaviour into lazy data realm (though I knew about passing the chunks kwarg, but had no idea it'd make such a difference). I owes you a pint for sure, I'll have a play with the optimal chunks, but even an empty dict drops the mem beautifully. Cheers 🍻

BTW my apologies for using the "bug" Issue template, it's not really a bug, but when I realized I'd already typed the issue 😄

[kmuehlbauer]

Great it works on your side!

No worries, we are always happy to receive nice formatted reports with many details and an MCVE. Thanks for that.

Would you think the documentation could be improved to make these differences more obvious?

[kmuehlbauer]

The documentation on these features are here:

https://docs.xarray.dev/en/latest/internals/internal-design.html#lazy-indexing-classes

There is another mention of these lazy classes in https://docs.xarray.dev/en/latest/user-guide/io.html#netcdf but the docstring of https://docs.xarray.dev/en/latest/generated/xarray.open_dataset.html#xarray.open_dataset does only reflect on the dask-part for the chunks-kwarg. Adding a line there for chunks=None linking to the above lazy-indexing-classes might help users to identify issues of that kind better.

[valeriupredoi]

hi @kmuehlbauer many thanks again! I tested with the same file as in the MRE, with different chunks values, and for the cases of remote (bucket) and local:

    ### chunks={} 325MB; 3.5s
    ### chunks="auto" doesn't work:
    # NotImplementedError: Can not use auto rechunking with object dtype. We are unable to estimate the size in bytes of object data
    ### chunks=-1 same as {}
    ### chunks=100 same as {}

    ### From Local
    ### chunks=None: 2.4GB; 6.4s
    ### chunks={}: 315MB; 3s
    ### chunks=-1: same as {}

chunks="auto" didn't work for this particular file, but that's no biggie since generic Dask {} should probably be used when not much is known about the actual storage chunking, and especially with our kind of Zarr files that may be chunked absolutely bonkers! I did, however, encounter this for chunks={}:

ValueError: Object has inconsistent chunks along dimension time. This can be fixed by calling unify_chunks().

which goes to show how nasty these files are chunked. But unify_chunks() did the trick.

Many thanks again! Sooo much better performance-wise this way 😁

About docs: I'd definitely add a warning here https://docs.xarray.dev/en/stable/generated/xarray.open_dataset.html in the case of chunks=None (default) a la "Warning: this will transfer the data and load it in memory as Numpy ndarrays" - but up to you folks. My problem is now fixed, and with very many thanks 🍻

[valeriupredoi]

BTW not sure what Xarray means by "small array" but for us, small array is eg (2352, 128, 256) and for that setting chunks={} does make a significant difference shoving about 700MB of memory - so I'd be careful with saying "chunks=None skips using dask, which is generally faster for small arrays." - maybe tiny instead of small 😆 That's me done bugging you, cheers again 🍺