Cache pre-existing Zarr arrays in Zarr backend

[d-v-b]

ZarrStore calls array_keys on the zarr.Group it wraps multiple times. Each of these array_keys calls performs IO, but if we model the names of the arrays as static, then they can be cached, then we can do less IO.

Accordingly this PR adds a new method to ZarrStore (get_array_keys), a new keyword argument for ZarrStore.__init__ (cache_array_keys), and a new attribute to ZarrStore (_cache). get_array_keys gets the array keys from the .zarr_group instance contained in the ZarrStore, either by reading the keys from the cache (if caching is enabled and the cache is populated), or by directly invoking ZarrStore.zarr_group.array_keys() (if caching is disabled, or the cache is empty).

Caching is controlled by the cache_array_keys keyword argument in ZarrStore.__init__.

• If ZarrStore.__init__ is called with with cache_array_keys=True, the _cache attribute of the ZarrStore is set to {"array_keys": None}. On the first invocation of get_array_keys, this empty-but-active cache is filled with a call to zarr_group.array_keys, and that value is returned. On subsequent invocations of get_array_keys, the value from the cache is retrieved and returned.
• If ZarrStore.__init__ is called with cache_array_keys=False, then the _cache attribute never gets a array_keys key inserted, and this indicates to get_array_keys that the cache is not being used. Calls to get_array_keys will return the result of zarr_group.array_keys.

Let me know if the _cache thing is a bit baroque -- my initial effort added a _cache_array_keys attribute as well as the cache, but I figured the state of the cache itself can convey whether the cache is meant to be used, which removes the need for a "use the cache or not" attribute on ZarrStore.

And I would like some feedback on the test -- is it in the right place?

closes #9853

[welcome]

Thank you for opening this pull request! It may take us a few days to respond here, so thank you for being patient.
If you have questions, some answers may be found in our contributing guidelines.

[dcherian]

If you swtich the equality to != we should see a decrease in list_prefix:

xarray/xarray/tests/test_backends.py

Line 3325 in eac5105

assert summary[k] <= expected[k], (k, summary)

[d-v-b]

@dcherian I have implemented caching both array keys and arrays, but for that reason this test is failing:

xarray/xarray/tests/test_backends.py

Lines 409 to 416 in eac5105

	def test_write_store(self) -> None:
	expected = create_test_data()
	with self.create_store() as store:
	expected.dump_to_store(store)
	# we need to cf decode the store because it has time and
	# non-dimension coordinates
	with xr.decode_cf(store) as actual:
	assert_allclose(expected, actual)

This test assumes that the ZarrStore will track changes to the underlying Zarr data which makes caching tricky. Does re-using the ZarrStore here reflect real-world usage, or is it just for testing convenience? and thus should I modify this test to create a new ZarrStore, or explicitly turn off the cache for this test?

[dcherian]

oh wow, I did not even know that there was a Dataset.dump_to_store method!

This code is from 10 years ago and I don't think anyone is using it like this. Today this test would be

with self.roundtrip(expected) as actual:
    assert_identical(actual, expected)

cc @shoyer

[d-v-b]

oh wow, I did not even know that there was a Dataset.dump_to_store method!

This code is from 10 years ago and I don't think anyone is using it like this. Today this test would be

with self.roundtrip(expected) as actual:
    assert_identical(actual, expected)

cc @shoyer

I altered the test as you recommended, let me know if that change should be in a separate PR

[dcherian]

does this need to be over-ridden still if the default is to turn the cache off

[d-v-b]

I think so, because we need to explicitly ensure that the cache is off, otherwise tests will fail. If we rely on the default behavior (which is defined very far away from this test), and if that default value changes, then tests will fail here and it will not be clear why.

[d-v-b]

another way to remove this code would be to change the tests so that they no longer rely on the old ZarrStore behavior; maybe this could be a separate PR.

[dcherian]

I'll push a commit. The create_store on this class seems pointless.

[dcherian]

Very nice! Clearly its pretty effective. I left some minor suggestions.

[jhamman]

Looks great @d-v-b. Left a comment on the default value.

Also, please add a what's new note in the docs.

[jhamman]

What do you think about changing the default to True? I'd like to know if we think there is a downside (risk, performance, etc.) to defaulting to this behavior?

[dcherian]

There are some weird tests that use open_store and dump_to_store that can't handle that default.

For nearly all users, the effective default is the value set in open_zarr and to_zarr (which is True, or should be after that one suggestion is merge)

[d-v-b]

i'm up for it! the main risk is whether anyone is actively using ZarrStore instances explicitly, with the assumption that they dynamically track the state of the underlying zarr Group. if nobody is doing that, then there's no risk. ofc I have no way of know if this use case is common, but my impression was that ZarrStore is internal enough that users are not expected to be using it directly.

[d-v-b]

CI is failing on mypy due to the following:

 Found 5 errors in 1 file (checked 170 source files)
xarray/tests/test_plot.py:98: error: Unused "type: ignore" comment  [unused-ignore]
xarray/tests/test_plot.py:105: error: Unused "type: ignore" comment  [unused-ignore]
xarray/tests/test_plot.py:113: error: Unused "type: ignore" comment  [unused-ignore]
xarray/tests/test_plot.py:125: error: Unused "type: ignore" comment  [unused-ignore]

I don't think that's from this PR?

[d-v-b]

Looks great @d-v-b. Left a comment on the default value.

Also, please add a what's new note in the docs.

"what's new" has been added and the default value is now True. CI is still red due to some mypy issues that I think are unrelated to my pr (see #9861 (comment))