Zarr: Optimize appending #8998

dcherian · 2024-05-03T22:21:44Z

~~Builds on #8997~~

Continues moving checks one level down to ZarrStore.store where we already have a bunch of checks, and a loop over existing variables in the store.

dcherian · 2024-05-03T22:22:18Z

xarray/backends/api.py

-    return region
-
-
-def _validate_datatypes_for_zarr_append(zstore, dataset):


This is moved to backends/zarr.py

dcherian · 2024-05-03T22:22:58Z

xarray/backends/api.py

-    existing dtype.
-    """
-
-    existing_vars = zstore.get_variables()


This PR removes this get_variables call by instead running the check in ZarrStore.store where we are already requesting existing variables, and doing checks.

dcherian · 2024-05-03T22:23:34Z

xarray/backends/zarr.py

+                if self._mode in ["a", "a-", "r+"]:
+                    _validate_datatypes_for_zarr_append(
+                        vn, existing_vars[vn], variables[vn]
+                    )


Now the check runs here where we have both the variables in the store, and the new variables to be written.

dcherian · 2024-05-07T15:51:51Z

xarray/backends/api.py

@@ -1721,28 +1685,12 @@ def to_zarr(
            )

    if mode in ["a", "a-", "r+"]:
-        _validate_datatypes_for_zarr_append(zstore, dataset)


moved to ZarrStore.store

dcherian · 2024-05-07T15:51:55Z

xarray/backends/api.py

                raise ValueError(
                    f"variable {var_name!r} already exists, but encoding was provided"
                )
-        if mode == "r+":


moved to ZarrStore.store

dcherian · 2024-05-07T15:52:13Z

xarray/backends/zarr.py

@@ -612,26 +640,58 @@ def store(
        import zarr

        existing_keys = tuple(self.zarr_group.array_keys())
+
+        if self._mode == "r+":


checks moved from backends/api.py

xarray/backends/zarr.py

dcherian · 2024-05-07T15:53:02Z

xarray/backends/zarr.py

@@ -696,7 +756,6 @@ def set_variables(self, variables, check_encoding_set, writer, unlimited_dims=No

        for vn, v in variables.items():
            name = _encode_variable_name(vn)
-            check = vn in check_encoding_set


reducing some indirection.

dcherian · 2024-05-07T15:56:31Z

xarray/backends/zarr.py

+                    "``mode='r+'``. To allow writing new variables, set ``mode='a'``."
+                )
+
+        if self._append_dim is not None and self._append_dim not in existing_keys:


self._append_dim not in existing_keys is a new addition. We needn't parse all arrays in the store for any append dimensions that we know to exist in the store.

dcherian · 2024-05-08T14:17:57Z

xarray/backends/api.py

-                    f"append_dim={append_dim!r} does not match any existing "
-                    f"dataset dimensions {existing_dims}"
-                )
+    if encoding and mode in ["a", "a-", "r+"]:


We can now skip this request for existing array names if encoding was not provided.

jhamman · 2024-05-09T06:34:17Z

@dcherian -- this looks great! The logic makes sense and things seem to work. Could you share some diagnostics on what this will mean in terms of traffic between Xarray and the Zarr store? Also, appreciating that this is just a refactor, do you think a new test would help avoid "uneccessary" traffic in the future?

max-sixty

Looks good! (Though I don't have that much context down at this level...)

dcherian · 2024-05-09T18:00:01Z

Pushed a regression test that covers #8997, #8893, and this PR.

It'd be nice to add this as an asv benchmark, but my concern is that we won't catch a regression since it doesn't always run

dcherian · 2024-05-09T18:00:39Z

xarray/tests/test_backends.py

+            modified.to_zarr(store, mode="a", append_dim="x")
+            # v2024.03.0: {'iter': 6, 'contains': 2, 'setitem': 5, 'getitem': 10, 'listdir': 6, 'list_prefix': 0}
+            # 6057128b: {'iter': 5, 'contains': 2, 'setitem': 5, 'getitem': 10, "listdir": 5, "list_prefix": 0}
+            expected = {


@jhamman 👁️ , compares release to current main (6057128)

dcherian · 2024-05-09T18:00:45Z

xarray/tests/test_backends.py

+            ds.to_zarr(store, region={"x": slice(None)})
+            # v2024.03.0: {'iter': 5, 'contains': 2, 'setitem': 1, 'getitem': 6, 'listdir': 5, 'list_prefix': 0}
+            # 6057128b: {'iter': 4, 'contains': 2, 'setitem': 1, 'getitem': 5, 'listdir': 4, 'list_prefix': 0}
+            expected = {


@jhamman 👁️

jhamman

Good stuff @dcherian!

* main: Add whatsnew entry for #8974 (#9022) Zarr: Optimize appending (#8998)

dcherian commented May 3, 2024

View reviewed changes

dcherian force-pushed the cleanup-zarr-append branch from 6fcd873 to b21f7f4 Compare May 7, 2024 15:51

dcherian commented May 7, 2024

View reviewed changes

xarray/backends/zarr.py Outdated Show resolved Hide resolved

dcherian commented May 7, 2024

View reviewed changes

dcherian marked this pull request as ready for review May 7, 2024 15:53

dcherian commented May 7, 2024

View reviewed changes

dcherian requested review from max-sixty and jhamman May 7, 2024 16:04

dcherian commented May 8, 2024

View reviewed changes

max-sixty approved these changes May 9, 2024

View reviewed changes

dcherian mentioned this pull request May 9, 2024

Release? #9018

Closed

TomNicholas added the topic-zarr Related to zarr storage library label May 9, 2024

dcherian and others added 3 commits May 9, 2024 11:56

Zarr: optimize appending

52891ca

Update xarray/backends/zarr.py

e38b6ee

Don't run encoding check if it wasn't provided.

b66452a

dcherian force-pushed the cleanup-zarr-append branch from 942e1dd to dc33846 Compare May 9, 2024 17:59

dcherian commented May 9, 2024

View reviewed changes

dcherian force-pushed the cleanup-zarr-append branch from dc33846 to f2d4ff4 Compare May 9, 2024 18:06

Add regression test

f76e725

dcherian force-pushed the cleanup-zarr-append branch from f2d4ff4 to f76e725 Compare May 9, 2024 18:12

dcherian added 2 commits May 9, 2024 12:21

fix types

86b6da6

fix test

5856c0c

jhamman approved these changes May 10, 2024

View reviewed changes

Use mock instead

2628aea

dcherian merged commit 6fe1234 into pydata:main May 10, 2024
25 of 28 checks passed

dcherian deleted the cleanup-zarr-append branch May 10, 2024 16:58

andersy005 added a commit that referenced this pull request May 12, 2024

Merge branch 'main' into backend-indexing

7d55345

* main: Add whatsnew entry for #8974 (#9022) Zarr: Optimize appending (#8998)

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Zarr: Optimize appending #8998

Zarr: Optimize appending #8998

dcherian commented May 3, 2024 •

edited

dcherian May 3, 2024

dcherian May 3, 2024

dcherian May 3, 2024

dcherian May 7, 2024

dcherian May 7, 2024

dcherian May 7, 2024

dcherian May 7, 2024

dcherian May 7, 2024

dcherian May 8, 2024

jhamman commented May 9, 2024

max-sixty left a comment

dcherian commented May 9, 2024 •

edited

dcherian May 9, 2024 •

edited

dcherian May 9, 2024

jhamman left a comment

		return region


		def _validate_datatypes_for_zarr_append(zstore, dataset):

Zarr: Optimize appending #8998

Zarr: Optimize appending #8998

Conversation

dcherian commented May 3, 2024 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jhamman commented May 9, 2024

max-sixty left a comment

Choose a reason for hiding this comment

dcherian commented May 9, 2024 • edited

dcherian May 9, 2024 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jhamman left a comment

Choose a reason for hiding this comment

dcherian commented May 3, 2024 •

edited

dcherian commented May 9, 2024 •

edited

dcherian May 9, 2024 •

edited