Differences in `to_netcdf` for dask and numpy backed arrays

[slevang]

What is your issue?

I make use of fsspec to quickly open netcdf files in the cloud and pull out slices of data without needing to read the entire file. Quick and dirty is just ds = xr.open_dataset(fs.open("gs://...")).

This works great, in that a many GB file can be lazy-loaded as a dataset in a few hundred milliseconds, by only parsing the netcdf headers with under-the-hood byte range requests. But, only if the netcdf is written from dask-backed arrays. Somehow, writing from numpy-backed arrays produces a different netcdf that requires reading deeper into the file to parse as a dataset.

I spent some time digging into the backends and see xarray is ultimately passing off the store write to dask.array here. A look at ncdump and Dataset.encoding didn't reveal any obvious differences between these files, but there is clearly something. Anyone know why the straight xarray store methods would produce a different netcdf structure, despite the underlying data and encoding being identical?

This should work as an MCVE:

import os
import string
import fsspec
import numpy as np
import xarray as xr

fs = fsspec.filesystem("gs")
bucket = "gs://<your-bucket>"

# create a ~160MB dataset with 20 variables
variables = {v: (["x", "y"], np.random.random(size=(1000, 1000))) for v in string.ascii_letters[:20]}
ds = xr.Dataset(variables)

# Save one version from numpy backed arrays and one from dask backed arrays
ds.compute().to_netcdf("numpy.nc")
ds.chunk().to_netcdf("dask.nc")

# Copy these to a bucket of your choice
fs.put("numpy.nc", bucket)
fs.put("dask.nc", bucket)

Then time reading in these files as datasets with fsspec:

%timeit xr.open_dataset(fs.open(os.path.join(bucket, "numpy.nc")))
# 2.15 s ± 40.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit xr.open_dataset(fs.open(os.path.join(bucket, "dask.nc")))
# 187 ms ± 26.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

[dcherian]

Hm... for me the first call in a session is always slow which suggests its some fsspec auth/setup thing.

[slevang]

I did try many loops and different order of operations to make sure this isn't a caching or auth issue. You can see the std dev of the timeit calls above is pretty consistent.

For my actual use case, the difference is very apparent, with open_dataset taking about 9 seconds on the numpy-saved file and <1 second on the dask-saved one. I can also clearly see when monitoring network traffic that the slow version has to read in hundreds of MB of data to open the dataset, while the fast one only reads the tiny headers.

I also inspected the actual header bytes of these two files and see they are indeed different.

[dcherian]

Hmmm open_dataset should really only read metadata

Can you run a line profiler and see what's requesting the reads please.

[slevang]

The slow file:

And the fast file:

[dcherian]

Do you think this is tied to fsspec somehow?

Perhaps @martindurant can take a look.

[martindurant]

I do generally recommend cache_type="first" for reading HDF5 files, because they tend to have most of the metadata in the header area of the file, with short pieces of metadata "elsewhere"; so the default readahead doesn't perform very well.

As to what the two writers might be doing differently, I only have guesses. I imagine xarray leaves it entirely to HDF to make whatever choices it likes. Dask does not write in parallel, since HDF does not support that, but it may order the writes more logically. It does set up the whole set of variables as a initialisation stage before writing any data - I don't know if xarray does this.

[slevang]

Yeah that seems to be it. Dask's write neatly packs all the needed metadata at the beginning of the file, since we can scale this up to a many GB file with dozens of variables and still read in ~100ms. While xarray is doing a less well organized write of the metadata and we have to go seeking in the middle of the byte range. cache_type="first" does provide some improvement but still not as good as on the dask-written file.

FWIW, I inspected the actual bytes of the dask and xarray written files and they are identical for a single variable, but diverge when multiple variables are being written. So, the important differences are probably associated with this step:

It does set up the whole set of variables as a initialisation stage before writing any data - I don't know if xarray does this.