Add initialize_zarr

[dcherian]

• Closes Add metadata_only param to .to_zarr? #8343
• Tests added
• User visible changes (including notable bug fixes) are documented in whats-new.rst
• New functions/methods are listed in api.rst

The intended pattern is:

    after_init = initialize_zarr(store, ds, region_dims=("x",))
    for i in range(ds.sizes["x"]):
        after_init.isel(x=[i]).to_zarr(store, region={"x": slice(i, i + 1)})

cc @slevang

[dcherian]

☝🏾 This is important. What do we mean by initialization?

1. All non-dask arrays will be written to disk. & Metadata for all dask arrays will be written to disk.
2. Alternatively, it could be all metadata for all variables will be written to disk. Reading that initialized store will get you fill values.

[dcherian]

As a concrete example we could do soemthing like this from pangeo-forge:
https://github.com/pangeo-forge/pangeo-forge-recipes/blob/5367f6363be91dd4826f3e5f21ffea61b60994cb/pangeo_forge_recipes/aggregation.py#L205

(thanks @rabernat)

[max-sixty]

I was originally thinking that it would be the former. It's at least important to write indexes, not just their metadata...

[rabernat]

If you don't write the dimension coordinates, the dataset may be unreadable by xarray. The fill values could decode to invalid results. This happens especially with time coordinates.

[slevang]

Nice! full_like().to_zarr(compute=False) being so much faster feels like a hack but I'll take it.

I do think having a version of this built into Dataset.to_zarr() as a kwarg (metadata_only, template, etc) would be nice for the most concise user code. Going that route would potentially require resolving the issue of needing to drop non-region containing variables that you address here though, or at least more documentation.

[slevang]

This is basically happening already on write in any case of region="auto", see here. Interestingly, if you try to write to a region along a single dimension of a multidimensional array, it will fail if you skip this step with region={"x": "auto"} but succeed with region="auto", because the latter now expands region to {"x: slice(...), "y": slice(...)}, so it doesn't trip this error.

This gets back to the discussion of "should we actually be writing any indices in region mode". For the vast majority of use cases I think not, and we could simplify a lot of these logic branches by just dropping the indices early on. But because there is still the available option to use mode="a", region=..., doing so would break for anyone trying to simultaneously write a region and grow the coordinate. Which also gets to @rabernat's proposal of reworking these modes entirely...

[slevang]

I think for this to fairly be called an "initialize" or "template" method, we should not be writing any data_vars, regardless of dask vs numpy, but we should write all the indices and metadata.

[dcherian]

Going that route would potentially require resolving the issue of needing to drop non-region containing variables that you address here though, or at least more documentation.

Yes, i was trying to make an easy recommended pattern for region writes. To do that, it seemed like to_zarr wasn't the right place.

[jhamman]

I just pushed a little experiment that could be a really nice path forward here: c0cf4ee.

Demo:

ds = xr.tutorial.open_dataset('rasm').chunk('5mb')

# baseline
mapper_a = fs.get_mapper("s3://foo/init-zarr-test-a")
%time xr.zeros_like(ds).to_zarr(mapper_a, mode='w', compute=False)
# CPU times: user 296 ms, sys: 7.16 ms, total: 303 ms
# Wall time: 2.75 s

# my most recent commit
mapper_b = fs.get_mapper("s3://foo/init-zarr-test-b")
%time store = init_zarr_v2(ds, store=mapper_b)
# CPU times: user 41.4 ms, sys: 8.19 ms, total: 49.6 ms
# Wall time: 144 ms

# it works!
a = xr.open_zarr(mapper_a)
b = xr.open_zarr(mapper_b)
xr.testing.assert_equal(a, b)

The key realization I made is that if we aren't going to write the arrays, we can just use the Zarr API to create the metadata objects and push them to store.

Some caveats about what I just pushed:

• Missing a bunch of feature flags, tests, etc.
• Currently makes an assumption that the target store is clean and ready to be initialized
• May also be missing some specialized encoding around strings / object dtypes. This is all in the Zarr backend so could likely be abstracted away to share a common code path.

[max-sixty]

Nice!!

---

(One small difference between that and the PR — we possibly should write any var that doesn't have a region dim, not only the indexes)

[max-sixty]

Not a necessary change, but one plausible way to frame this is

• default region_dims=[...]
• ...and then there's no need for the if — we default to writing the minimum and returning the maximum.

[dcherian]

Interesting suggestion!

[dcherian]

There's an edge case around scalars. We would write them if the user wants region writes, but not write them if they don't. So we need to distinguish between region_dims=(), region_dims=(...,) and region_dims=tuple(ds.dims). If those last two are identical then we can't handle the scalar case.

[max-sixty]

OK interesting point.

OOI, is this function helpful if we're not doing region writes?

[max-sixty]

Yeah, on reflection — is there a case where we really don't want to write scalars? (genuine question, very possibly there is...)

[max-sixty]

Suggested change

	2. Return a dataset with the remaining variables, which contain one or more of ``region_dims``.
	2. Return a dataset with the remaining variables and indexes, which contain one or more of ``region_dims``.

...is that right?

Should we drop the non-region indexes?

My mental model of this (but it's from a few months ago):

• For indexes that are part of the region, we want to keep them now that we have region="auto". When we use region="auto" in each process, it correctly won't rewrite those indexes
• For other indexes, we want to drop them, because otherwise each process will incorrectly attempt to write them, and this can lead to conflicts

[dcherian]

When we use region="auto" in each process, it correctly won't rewrite those indexes

Surprising to me, but you are correct! #8589

Currently we keep any indexes needed for any variable that has dimensions that overlap with region_dims. So for the test case, there is a variable with dims (x, y). If region_dims==("x",), then the index y is also returned. This works fine with region="auto" but requires adding "y": slice(None) when explicitly specifying region. I think this is OK, if you're specifying region, explicitly specify the whole thing as region="auto" does.

For other indexes, we want to drop them, because otherwise each process will incorrectly attempt to write them, and this can lead to conflicts

We do end up dropping those indexes. The test case has a variable with dimensions ("z",) but the regions only contain "x" and/or "y". The z index ends up getting dropped.

[max-sixty]

This is a big one!! I hadn't realized it was such a lift.

Thanks a lot.

I left a few comments

[max-sixty]

Obv don't change but FYI the None is default in python here

[slevang]

Sorry @dcherian I haven't had a chance to look at this in detail yet, but overall it seems really nice.

One random thought on the API: is there a way to configure options such that ds.to_zarr(..., compute=False) could just call this method, since it does effectively the same thing but much faster? This way users get an immediate performance improvement without having to change existing code, plus you can still use the standalone xr.initialize_zarr to get back a template dataset for region writes.

[petewa]

Hi @dcherian, great work! This has solved one of the performance bottlenecks for us.

We are finding Xarray fairly slow in writing small Zarrs to S3 with many variables/DataArrays.

It seems slow in:

• Creating the metadata (due to latency?), which your PR solves. I tested this.
• Writing many chunks to S3 (due to latency?)
  • It's significantly faster to write to disk/memory then a parallel upload to s3.

Are there any plans to progress this PR. It would be good to have this merged.

[slevang]

Yeah this looks great, what else needs to be done?

[slevang]

No expert, but maybe:
index_vars = set(ds.indexes.keys())

[slevang]

If #8904 went in we wouldn't need to drop index_vars here.

[rabernat]

I share the perspective that it would be nicer to implement this within the existing to_zarr API rather than add a new top-level function to Xarray. @dcherian do you see a path to doing that?

[dcherian]

do you see a path to doing that?

Yes i too agree now, but it'll take some work (=time). I'm thinking we add a initialize method to ZarrStore and use that in store (?).

[rabernat]

I'm thinking we add a initialize method to ZarrStore and use that in store (?).

This makes sense to me. There's no reason for this not to be the default for all Zarr dataset creation. It will massively speed things up in most cases.

Thinking bigger, it should be possible to upstream some of this to Zarr itself. But that can come later.

[max-sixty]

To the extent that it's helpful and not annoying to add "this would be really useful": this would be really useful! It seems fairly close IIUC (though easy for me to say...)

One reflection from encouraging more Zarr usage with colleagues is that it is a) really spectacular for reading data, b) still much more complicated to write data for the average dev relative to something like parquet. I think this would close a lot of the gap...

[rabernat]

I'm fine with getting this out into the wild, but I am hesitant to commit to it as permanent API. Could we say it is an experimental feature, don't depend on it in production, etc?

[max-sixty]

I'm fine with getting this out into the wild, but I am hesitant to commit to it as permanent API. Could we say it is an experimental feature, don't depend on it in production, etc?

Definitely +1 from me