Refactor MultiZarrToZarr into multiple functions

@ivirshup

Problem

MultiZarrToZarr is extremely powerful but rather hard to use.

This is important - kerchunk has been transformative, so we increasingly recommend it as the best way to ingest large amounts of data into the pangeo ecosystem's tools. However that means we should make sure the kerchunk user experience is smooth, so that new users don't get stuck early on.

Part of the problem is that this one MultiZarrToZarr function can do many different things. Contrast with xarray - when combining multiple datasets into one, xarray takes some care to distinguish between a few common cases/concepts (we even have a glossary):

Concatenation along a single existing dimension. Achieved by xr.concat where dim is a str
Concatenation along a single new dimension (optionally providing new coordinates to use along that new dimension). Achieved by xr.concat where dim is a set of values
Merging of multiple variables which already share dimensions, first aligned according to their coordinates. Achieved by xr.merge
"Combining" by order given, which means some ordered combination of concatenation along one or more dimensions and/or merging. Achieved by xr.combine_nested
"Combining" by coordinate order, which again means some ordered combination of concatenation along one or more dimensions and/or merging, but the order is specified by information in the datasets' coordinates. Achieved by xr.combine_by_coords

In kerchunk it seems that the recommended way to handle operations resembling all 5 of these cases is through MultiZarrToZarr. It also cannot currently easily handle certain types of multi-dimensional concatenation.

Suggestion

Break up MultiZarrToZarr by defining a set of functions similar to xarray's merge/concat/combine/unify_chunks that consume and produce VirtualZarrStore objects (EDIT: see #375).

Advantages

We can replace/deprecate the heavily overloaded and unituitive coo_map kwarg (it has 10 possible input types!). Perhaps giving simply an ordered list of coordinate values would be sufficient, and just make it easier for the user to extract the values they want from the VirtualZarrStore objects they want to concatenate.
If users need to do something really unusual they can more easily break their problem up into concatenating each array separately (e.g. for concatenating on staggered grids)
Might generalise to later ZEPs more easily (e.g. understanding variable-length chunks, cc @ivirshup, see Concatenate arrays with varchunks #374)
Can think of as a refactoring to move some pangeo-forge functionality upstream, reducing redundancy. We shouldn't have 3 completely different designs for multidimensional concatenation in adjacent libraries in the stack.
These new functions would be more useful as basic primitives for parallelization frameworks to call (e.g. doing tree reduction via dask, beam, or cubed), rather than trying to wrap calls to those frameworks within kerchunk (like kerchunk.combine.auto_dask does).

Questions

How close can these functions be to xarray's version of merge/concat/combine? And what can we learn from the design decisions in pangeo-forge-recipes FilePattern? (@cisaacstern @rabernat )
How close are kerchunk's existing combine.merge_vars and combine.concatenate_arrays functions to providing this functionality? If the answer is "pretty close", then how much of this issue could be solved via documentation?

I've been thinking of something similar. I'm not sure we can get all the flexibility we want out of just two functions, but think a couple classes may do the trick. I've been thinking of an API like:

g = kerchunk.KerchunkGroup()
# issubclass(KerchunkGroup, MutableMapping[str, KerchunkGroup | KerchunkArray])

# Concatenation + assignment, following array-api
g["x"]: kerchunk.KerchunkArray = kerchunk.concat([...], axis=0)

# Write a manifest
g.to_json(...)

# Getting a zarr group:
g.to_zarr()

# Reading from a manifest
g2 = kerchunk.KerchunkGroup.from_json(...)

# Assignment
g2["new_group"] = g

# Single level merging
g.update(g2)

# Add an array from a different store:
g.create_group("my").create_group("new")
g["my"]["new"]["array"] = zarr.open("matrix.zarr"): zarr.Array

Then leave anything more complicated to something like xarray, which could re-use its combination logic on a collection of KerchunkArrays.

This suggestion has a lot of the functionality already in concatenate_arrays and merge_vars, but:

Makes manipulation of the hierarchy much easier (the other two you are manually prepending/ appending to the keys)
User friendly and easy to introspect interface
Possible performance benefits by avoiding repeated json string serialization

This idea may be more of a zarr thing than a kerchunk thing though.

Also, I think you may be missing a reference to #375 under the "Suggestion" heading.

Right now we are using Kerchunk as a workaround for things that Zarr itself does not support, by implementing various types of indirection on chunk keys at the Store level. I think this is a great way to prototype new possibilities. The downside is that only kerchunk-aware clients can ever hope to read these datasets. Also, from a design point of view, stores are technically ignorant of the Zarr data model (arrays, groups, etc.); they are formally just low-level key value stores.

So when it comes time to refactor those prototypes into more stable interfaces, I think we should seriously consider upstreaming some of these features into Zarr itself, possibly via new Zarr extensions. The advantage of that approach is that

We can implement these features via a specification, with an aim for multi-language support
We can leverage Zarr's understanding of the data model (arrays, groups, chunks, dimensions, etc.) to make the implementation cleaner

I believe that all 5 of the features above could be implemented in this a way. In particular, the concept of a "chunk manifest": a document which stores information about where to find the chunks for an array (rather than assuming they follow the Zarr chunk encoding scheme) would get us very far. The outcome of this would be a Zarr array with much more powerful virtualization capabilities, something between Xarray's duck array types and the current Kerchunk situation.

We would be eager to collaborate on this and have some work in this direction planned for this quarter.

To be more concrete about what I mean, we could imagine writing code like this

import zarr

a1 = zarr.open("array1.zarr")  # zarr.Array
a2 = zarr.open("array2.zarr")  # zarr.Array 

# this would create a Zarr array with the correct extensions to support indirection to a1 and a2
# (need to figure out the right API)
a3 = zarr.concat([a1, a2], axis=0, store="array3.zarr")

a3[0] = 1  # write some data...automatically goes into the correct underlying array

@manzt

We can implement these features via a specification, with an aim for multi-language support

We definitely want to see multi language support (cc: @manzt)

Do you have thoughts on what extensions to zarr would need to be made? Definitely some sort of specification for being able to pass around manifests. But maybe this means needing to define some set of stores as well?

Does mapping an HDF5 file to zarr fit as a zarr extension? As high priority of one?

we could imagine writing code like this

a3 = zarr.concat([a1, a2], axis=0, store="array3.zarr")

My idea currently would be that you don't need to assign to a store on concatenation. I think this could be a purely in memory object and you resolve everything at access or storage time.

I'd also even say I'm not sure a "reference" based array should be an instance or subclass of zarr.Array. Maybe a sibling class, or just a shared interface would work. I'm thinking of use cases like creating virtual arrays which don't use the exact same codecs (like HDF5 virtual datasets allow).

a3[0] = 1  # write some data...automatically goes into the correct underlying array

I've mostly been thinking about read access at the moment. Not sure how this complicates things.

We would be eager to collaborate on this and have some work in this direction planned for this quarter.

Me too!

I've mostly been thinking about read access at the moment. Not sure how this complicates things.

read-only is of course easier, but also just a special case of read-write. Our arraylake friends have a well-defined model of where and how writes should go regardless of where the original mapper/chunks were defined.

@manzt

We definitely want to see multi language support (cc: @manzt)

Awesome. Yah, something that should be very possible with zarrita.js.

Repeating https://observablehq.com/@manzt/ome-tiff-as-filesystemreference using some of these newer ideas would be awesome :)
(and note that that original example was not, of course, a standard xarray-compatible dataset, but a subscale pyramid)

Sorry I've been out of the kerchunk loop for a bit but I've been trying to stay abreast of changes. zarrita.js has an HTTP-based reference store (and will try to map references with s3/gcs protocols to http endpoints). This higher-level virtualization of zarr is very exciting.

(sorry I haven't yet had the chance to review this)

@rabernat

Thanks @rabernat - I think the proposal of defining how arrays should be concatenated at the Zarr level is really cool, and I would be interested to help. I also think that working this out consistently at the Zarr level would help expose the ways in which the current kerchunk API brushes subtleties under the rug.

One of these subtleties actually just came up for me in practice in #386 - I got confused by whether concatenation in Kerchunk was supposed to follow xarray rules or lower-level array-like rules instead.

We therefore have to be really clear about our type systems, as we now have at least 3 different data structures for which concatenation would be defined differently:

numpy-like arrays - no dimension names, no index-backed coordinates, so concatentation just follows axis number,
zarr (v3) arrays - has dimension_names, but no index-backed coordinates, so how should concatentation work?
xarray objects - has dimension names and index-backed coordinates, so concatenation follows dimension name but has to make exceptions for index-backed coordinates (see [BUG] MultiZarrToZarr fails to concatenate dimension coordinate variables #386 (comment)).

If we could get this straight it would be really powerful though. Also then ultimately kerchunk might not need to deal with combining at all, it would become exclusively a collection of zarr readers for different legacy file formats.

zarr (v3) arrays - has dimension_names

Aren't dimensions name optional? In which case these should just behave like numpy arrays?

From https://zarr-specs.readthedocs.io/en/latest/v3/core/v3.0.html#dimension-names

If specified, must be an array of strings or null objects with the same length as shape. An unnamed dimension is indicated by the null object. If dimension_names is not specified, all dimensions are unnamed.

Aren't dimensions name optional?

Oh I guess so!

In which case these should just behave like numpy arrays?

But obviously propagating the dimension_names in a transparent and consistent way. Doesn't that potentially raise problems already? Like

z1 = zarr.Array(1d_data, dim_names=['t'])
z1 = zarr.Array(1d_data, dim_names=['x'])

zarr.concat([z1, z2], axis=0)  # well-defined numpy-like operation, should return a 1D array, but what dim_name should the result have?

I'm just saying that having Zarr array type(s) that explicitly encode these propagation rules would be far better than implicitly doing this stuff in the depths of kerchunk.

EDIT: It reminds me of this suggestion data-apis/array-api#698 (comment)

I am weakly against extracting the logic and putting it in zarr. Zarr-python has proven to be very slow moving compared to the needs of kerchunk, and I feel like we probably will have a wider set of cases we want to cover than zarr wants. I am aware that to some extent we end up replicating stuff in xarray (like the encoding stuff), but I don't mind that for the same of being explicit. What we really need, is a good set of test cases to describe expected behaviour.

I would also point out that kerchunk aims to be useful beyond zarr, and that it will be a place to test out cutting edge zarr features (like var-chunks) too.

I wonder how we can get together kerchunk-interested people? Do we need a dedicated slack or discord or something?

But obviously propagating the dimension_names in a transparent and consistent way. Doesn't that potentially raise problems already?

I think it's reasonable to let people do concatenation like numpy arrays regardless of what the dimension names are. array API as much as possible would be a good idea, just to reduce the amount of new things one needs to learn.

Could something like the merge logic in xarray be used for handling dimension names and other metadata (i.e. anything in attrs)? E.g. rules like "unique", "first"?

I've been thinking about operations other than concatenation that are achievable with kerchunk:

Transpose/ permute dims (re-order chunks + add a transpose codec)
Rename dimensions
Modify metadata
Limited subsetting

It would be nice to have these operations exposed in a way that doesn't require rewriting actual chunk data.

I wonder how we can get together kerchunk-interested people? Do we need a dedicated slack or discord or something?

If there was a zarr zulip/ discord/ slack maybe kerchunk could fit in there?

I think it's reasonable to let people do concatenation like numpy arrays regardless of what the dimension names are. array API as much as possible would be a good idea, just to reduce the amount of new things one needs to learn.

I second this, especially given that not all HDF or other files will be netCDF compliant anyway, and it would be good to have an opinion on what to do with them.

(My responses)

What MZZ does cannot be cleanly decomposed into a set of steps like that.

I disagree - I think that what MZZ does can and should be decomposed into a series of steps just like that, you just also need to include the determination of the the total ordering in the series of steps. Xarray can either be told this ordering (via the order of inputs to concat/combine_nested) or it can deduce it by reading coordinates (via combine_by_coords). In both cases the ordering happens first, then atomic concat/merge functions are called to do the actual combining.

(None of this is to say that having concat/stack/merge functions or array-standard methods, which can be applied tree-wise is a bad idea; but the workflow would be quite different and put a lot of up-front work on the caller)

The workflow would be different if the user needs to go down to the level of combine_nested/concat/merge, but it would be different in a good way: it would hand over more power to them when they need it to express their desired result. However most users wouldn't even need to put in any more up-front work, because combine_by_coords (the default for open_mfdataset) suffices for many cases.

or it can deduce [the ordering] it by reading coordinates (via combine_by_coords).

I would say, that this is almost entirely what MZZ does right now, and I would be very interested to see how it could be done independently for a second stage of tree-like merge/combine.concat. Would the decomposition you are after be better described as

find the coordinates of each chunk of each input, and the total set of coordinate values for each coordinate of interest
do the combine with merge/combine/concat

Would the decomposition you are after be better described as

find the coordinates of each chunk of each input, and the total set of coordinate values for each coordinate of interest

do the combine with merge/combine/concat

Yes - this is exactly what xarray already does - you can see it clearly by how in the combine module there are two internal functions: _infer_concat_order_from_positions and _infer_concat_order_from_coords. The former corresponds to xr.combine_nested and the latter to xr.combine_by_coords. Once the concatenation order is known, both functions use the same code to perform the (multi-dimensional) series of concat/merge operations (see _combine_nd).

EDIT: And all of this nesting is conveniently hidden from the user when they just call xr.open_mfdataset.

OK, so all this time I have been asking about how one might use xarray's machinery to serialise the set of positions of the input data, and now I know that _infer_concat_order_from_coords exists...

The remaining differences that would need to be figured out:

we support many ways to find the coordinated of a chunk where these cannot be had from its internal data, particularly based on. filenames
we need to know not just the location of each data input, but each of its chunks.

It sounds like you already have ideas for all of this?

we support many ways to find the coordinated of a chunk where these cannot be had from its internal data, particularly based on filenames

This is just a special case of specifying the total ordering. The simplest way is just to ask the user to deduce the ordering and pass a list of datasets of the correct order to combine_nested. A more automatic way would be for the user to write a preprocess function for open_mfdataset that adds a new dimension coordinate to the dataset, which can then be used by the combine_by_coords part of open_mfdataset. ~~One subtlety is that then preprocess would need to have access to the filename, but that seems solvable.~~

EDIT: Apparently we can use preprocess as-is because xarray automatically saves the filename into ds.encoding["source"]. So MZZ'ing files that are ordered by filename could just be

def add_time_dim_coord_from_filename(ds: xr.Dataset) -> xr.Dataset:
    """Adds a new dimension coordinate containing the time value read from the datasets' filename."""
    filename = ds.encoding["source"]
    
    # actually derive the time value from the filename string
    subst = re.search(r"\d{4}-\d{2}", filename)[0]
    time_val = cftime.DatetimeNoLeap.strptime(subst, '%Y-%m', calendar='noleap', has_year_zero=True)
    
    time_dim_coord = xr.DataArray(name='time', data=[time_val], dims='time')
    return ds.assign_coords(time=time_dim_coord)


ds = xr.open_mfdataset(
    '/my/files*.nc',
    engine='kerchunk',  # kerchunk registers an xarray IO backend that returns zarr.Array objects
    combine='by_coords',  # 'by_coords' would require actually reading coordinate data
    preprocess=add_time_dim_coord_from_filename,
)
ds  # now wraps a bunch of zarr.Array / kerchunk.Array objects directly
ds.kerchunk.to_zarr(store='out.zarr')

we need to know not just the location of each data input, but each of its chunks.

Keeping track of the locations of the chunks is a separate requirement, which should be the responsibility of the Array object to be concatenated. It could be handled by a kerchunk.Array (as I tried to show in my notebook above), but I actually now think if we're going to that much effort we should just go the whole way to a zarr chunk manifest instead.

Right, we definitely need to store that manifest (or whatever it is called/however structured), we cannot be opening all of the input files each time we want access to the dataset. So open_mf is not realistic - even listing files can be expensive - we need to use those low level functions to figure out where things go and save that, ideally (I think) in the existing kerchunk format.

we cannot be opening all of the input files each time we want access to the dataset. So open_mf is not realistic - even listing files can be expensive

I am not suggesting that open_mfdataset get called every time the kerchunked dataset is accessed. I am suggesting that open_mfdataset get called once (with a special set of arguments), in the same way that MultiZarrToZarr currently gets called once.

I think of what kerchunk is doing as effectively caching the result of all the ordering, checks, and concatenation that open_mfdataset does, so that you only need to do that work once, and thereafter you can open instantly just using open_dataset('store.zarr'). The actual cached information is the kerchunk format / chunk manifest file in the zarr store written to disk.

[that information should be saved] ideally (I think) in the existing kerchunk format.

That's a separate question from the whole "split everything into array operations that can be wrapped and applied with xarray API" suggestion. It's the same thing as asking do we combine KerchunkArray objects that are then serialized via the existing kerchunk format, or do we combine new VirtualZarrArray objects that are then serialized into valid zarr stores (requiring extensions to the Zarr spec).

I am not suggesting that open_mfdataset get called every time the kerchunked dataset is accessed.

Sorry, misunderstood you

we support many ways to find the coordinated of a chunk where these cannot be had from its internal data, particularly based on. filenames

I wonder if it would be better to do this at the reference extraction step (in SingleHdf5ToZarr or scan_grib)? That way we wouldn't have to pass the filenames along with the references to the next step (this has come up in pangeo-forge/pangeo-forge-recipes#583)

(None of this is to say that having concat/stack/merge functions or array-standard methods, which can be applied tree-wise is a bad idea; but the workflow would be quite different and put a lot of up-front work on the caller)

Will some of these cases trigger computation? Or do you mean it will generally be fine in the netcdf-esque situation?

For mergers where you need the coordinate values, you will always need to compute something at some point, but the main data arrays should not need to be touched.

General reference to: https://github.com/TomNicholas/VirtualiZarr (which I have not yet had chance to look through)

Thanks Martin...the above a github rec this week, so that was good. Have only looked at briefly.

@jhamman

Yep - I was going to wait until I had the MVP fully working before commenting here, but basically I've been working on a new package "VirtualiZarr" that aims to implement all the ideas detailed above.

The rationale is laid out in the readme, but basically instead of manipulating references as nested dicts that represent entire files, instead we just use kerchunk to read the byte ranges of chunks from files, but then immediately hide the details inside array-like ManifestArray objects. Each ManifestArray essentially corresponds to one zarr array, but with the references to chunks stored inside a dict-like ChunkManifest object, one per ManifestArray.

The ChunkManifest class is inspired by @jhamman's "Chunk Manifest" ZEP proposal, and is basically a in-memory container storing a mapping from zarr chunk keys to references, i.e:

{
    "0.0.0": {"path": "s3://bucket/foo.nc", "offset": 100, "length": 100},
    "0.0.1": {"path": "s3://bucket/foo.nc", "offset": 200, "length": 100},
    "0.1.0": {"path": "s3://bucket/foo.nc", "offset": 300, "length": 100},
    "0.1.1": {"path": "s3://bucket/foo.nc", "offset": 400, "length": 100},
}

The ManifestArray objects are array-like enough that they can be wrapped and manipulated by xarray, so concatenation problems can be expressed through xarray operations rather than having to use MultiZarrToZarr. Some parts of the metadata are stored in the xarray object's attributes, for example the _ARRAY_DIMENSIONS are just the xarray objects' .dims.

The concatenation itself is defined at the level of the Manifest, where it can be performed by just renaming keys and merging dicts.

Once you've combined / concatenated / merged the virtual representations of your files using xarray, you write the complete reference sets to disk using an xarray accessor, as

ds.virtualize.to_kerchunk(outpath, format='json')

Now you should be able to open these using fsspec as you normally would open the results of using MultiZarrToZarr.

Because this design now separates the reading of byte ranges from the representation of them in memory from the concatenation of them from the on-disk format we save them to, it's easier to see how upstream changes in zarr (e.g. chunk manifests in zarr, virtual array concatenation in zarr, and variable chunks in zarr) could plug into VirtualiZarr as they come.

Navigation Menu

Repository navigation

Refactor MultiZarrToZarr into multiple functions #377

Problem

Suggestion

Advantages

Questions

27 remaining items

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Refactor MultiZarrToZarr into multiple functions #377

Description

Problem

Suggestion

Advantages

Questions

Activity

ivirshup commented on Oct 16, 2023

rabernat commented on Oct 16, 2023

rabernat commented on Oct 17, 2023

ivirshup commented on Oct 18, 2023

martindurant commented on Oct 18, 2023

manzt commented on Oct 18, 2023

martindurant commented on Oct 18, 2023

manzt commented on Oct 18, 2023

martindurant commented on Oct 26, 2023

TomNicholas commented on Nov 3, 2023

dcherian commented on Nov 3, 2023

TomNicholas commented on Nov 3, 2023

martindurant commented on Nov 6, 2023

ivirshup commented on Nov 7, 2023

martindurant commented on Nov 7, 2023

27 remaining items

TomNicholas commented on Mar 7, 2024

martindurant commented on Mar 7, 2024

TomNicholas commented on Mar 7, 2024

martindurant commented on Mar 7, 2024

TomNicholas commented on Mar 7, 2024

martindurant commented on Mar 7, 2024

TomNicholas commented on Mar 7, 2024

TomNicholas commented on Mar 7, 2024

martindurant commented on Mar 8, 2024

keewis commented on Mar 8, 2024

RichardScottOZ commented on Mar 13, 2024

martindurant commented on Mar 13, 2024

martindurant commented on Mar 15, 2024

RichardScottOZ commented on Mar 15, 2024

TomNicholas commented on Mar 15, 2024

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions