More control over lazy data creation (chunking)

[TomekTrzeciak]

Currently, it is not possible to control how the lazy data gets chunked. It is also not possible to change that afterwards (dask rechunk function does not change the original chunking, it only adds additional split/merge operations on top of it). While the default choice of chunking might be OK in some cases, in other it might be unsuitable and it would be useful to allow for user choice in this respect.

[pp-mo]

Hi @TomekTrzeciak . Thanks for suggesting.

In what operations do you want to control chunking ?

You can already pass a pre-created dask array to the cube constructor, or assign it into cube.data.
We also have chunking control when saving to netcdf : iris.fileformats.netcdf.save supports a 'chunksize' keyword, so that can also appear as a kwarg in an iris.save call.

[TomekTrzeciak]

In what operations do you want to control chunking ?

iris.load & co.

You can already pass a pre-created dask array to the cube constructor, or assign it into cube.data.

Constructing the cube directly is not very convenient. I guess reassigning cube.data could be an option, but still rather awkward to write something like this:

cube = iris.load_cube(filepath)
cube.data = dask.array.from_array(netCDF4.Dataset(filepath)[cube.var_name], chunks=-1)

instead of just:

cube = iris.load_cube(filepath, chunks=-1)

[pp-mo]

Regarding load, options will depend on the source format.

The "field-based" file formats (FF, PP, GRIB) deal only in 2D fields, and they don't have any efficient access to subregions of a 2d field (i.e. the format code can only load a whole field then extract from it). Thus, the natural chunksize is the whole field, and I don't think there will ever be any practical use for chunking differently in those cases.

But I guess you are talking about netCDF ???
What we currently have there is basically determined by this code and this code .
In summary :

• take a whole variable as a chunk,
• .. or use file chunking if set
• .. and reduce the earlier dimensions if resulting chunks are impracticably large

There is certainly scope for controlling that : for instance, the chunk reduction assumes c-order contiguity, so will be worst-case if earlier dimensions vary faster in the file.

So I think we are talking about adding a chunk-control keyword to the netcdf loader.
That could be pretty simple, as the underlying interfaces already have suitable controls in them.
Is that what you mean ?

[TomekTrzeciak]

So I think we are talking about adding a chunk-control keyword to the netcdf loader.
That could be pretty simple, as the underlying interfaces already have suitable controls in them.
Is that what you mean ?

Yes, I think extra keyword passed through from load api to netcdf loader would be all that's needed.

[pp-mo]

extra keyword passed through from load api to netcdf loader

I had a quick look. Unfortunately, there is no support for additional args/kwargs in the generic load functions, iris.load/load_cube/load_cubes/load_raw.
I guess that's because it's hard to associate controls with individual load items, unlike save where it makes sense to provide additional controls on a per-cube basis.
So, I hope this is ok for your expectations, I think it will be much harder to extend the iris.load API.

adding a chunk-control keyword to the netcdf loader

In the netcdf-specific loader, we currently have iris.fileformats.netcdf.load_cubes(filenames, callback=None).
I'd expect to add a new facility there.
How about something like ...
iris.fileformats.netcdf.load_cubes(filenames, callback=None, filepath_varname_chunks={})
Then use like

from iris.fileformats.netcdf import load_cubes as nc_load_cubes
cubes = nc_load_cubes(filename, {(filename, 'var1'): (32, -1, 16)})

Will this work for your purposes?
And generally, what does your usecase look like -- we'll be needing a testcase anyway ...
Feel free to propose something.

[TomekTrzeciak]

@pp-mo, exposing chunks in iris.fileformats.netcdf.load_cubes could be one way, but this api is quite low level (returns a cube generator, so one needs to be careful how this gets passed around in user code, probably best to convert to something like CubeList immediately).

I think iris.load/load_cube/load_cubes/load_raw could support additional kwargs the same way that iris.save does - the backends that don't support certain kwargs could simply ignore them. But I can see how this leads to passing arguments through several levels of APIs, so arguably not the nicest pattern.

An alternative could be to use a context manager to set/pass backend options without bloating top level APIs. I've noticed that there already exists iris.config.netcdf context manager, so that could be a possible place to add this (or another one could be added specifically for dask if that's preferable). Then, the user code could be like this:

with iris.config.netcdf.context(chunks=-1):
    cube = iris.load_cube('dataset.nc')

[pp-mo]

Hi @TomekTrzeciak thanks for sticking with this.

I think iris.load/load_cube/load_cubes/load_raw could support additional kwargs the same way that iris.save does

I don't think there is any serious reason to oppose additional load controls. I just thought it sounded like more trouble to get such a change agreed.
I myself certainly don't object in principle to passing keywords down from the main 'load' functions. I think that is probably the most obvious way to do it.

My concern is that, to be useful, I think we need to be able to specify chunking of individual file variables (see why below...). This means that the controls can't be expressed in terms of core Iris concepts such as cube identity, which then looks rather different to the 'save' case.
As I said...

it's hard to associate controls with individual load items, unlike save where it makes sense

I don't see any sensible way around this, as you can't easily predict what a given load will produce, or which Iris objects relate to which parts of a source file -- because Iris itself doesn't make any simple guarantees about those behaviours : If data changes, you can't reliably know beforehand how it will merge, how many cubes are returned, or in what order -- see for example #3314.

The reason I think we need a flexible control solution is that we do need to cope with large AuxCoords -- often larger than the data variables. That is exactly why we implemented lazy loading for AuxCoords. So it means we will want to control chunking of those variables too. In the near future, we also expect to be dealing with large unstructured grids, which will present the same problem.

I think it could be fine, if we can design a default behaviour that enables us to simplify the simple cases. I'm just a bit wary, as it isn't immediately obvious to me how that can work.

[pp-mo]

Hints of progress ?
#3357 (comment)

[pp-mo]

Cross-copied from #3357
I'm now trying to retrospectively separate these 2 discussions ! ...

@TomekTrzeciak Xarray has a nice idea there to accept chunks given as a mapping of dimension name to chunk size.

@pp-mo A really nice spot 💐, but it still has the same problem of being tied to the low-level file encoding, which must then be both known (to the user) and stable : In this case we need dimension names, which don't actually appear anywhere in the iris/CF data representation.
If we could use names of dimension coordinates, that would help loads with this. But I think we can only do that if we first "pre-load", then "re-load" the file. 😢

Also, technically, as the occurrence and order of dimensions will differ between variables, you might still want dimensions divided differently for different variables. But I agree that is probably an obscure case : In principle you could even have an X(r:74, t:20000, y:1500, x:2000) and a Y(x:2000, t:20000, r:74), but I think it would be really rare -- maybe a usecase in ocean data ??