From 6fcd873421fd845bf25e70f0eba9185ceb4527b2 Mon Sep 17 00:00:00 2001
From: Deepak Cherian <deepak@cherian.net>
Date: Thu, 2 May 2024 15:53:17 -0600
Subject: [PATCH] Zarr: optimize appending

---
 xarray/backends/api.py  | 43 +++--------------------------------------
 xarray/backends/zarr.py | 32 ++++++++++++++++++++++++++++++
 2 files changed, 35 insertions(+), 40 deletions(-)

diff --git a/xarray/backends/api.py b/xarray/backends/api.py
index 1c0537b5d10..21163227e88 100644
--- a/xarray/backends/api.py
+++ b/xarray/backends/api.py
@@ -1521,42 +1521,6 @@ def save_mfdataset(
         )
 
 
-def _validate_datatypes_for_zarr_append(zstore, dataset):
-    """If variable exists in the store, confirm dtype of the data to append is compatible with
-    existing dtype.
-    """
-
-    existing_vars = zstore.get_variables()
-
-    def check_dtype(vname, var):
-        if (
-            vname not in existing_vars
-            or np.issubdtype(var.dtype, np.number)
-            or np.issubdtype(var.dtype, np.datetime64)
-            or np.issubdtype(var.dtype, np.bool_)
-            or var.dtype == object
-        ):
-            # We can skip dtype equality checks under two conditions: (1) if the var to append is
-            # new to the dataset, because in this case there is no existing var to compare it to;
-            # or (2) if var to append's dtype is known to be easy-to-append, because in this case
-            # we can be confident appending won't cause problems. Examples of dtypes which are not
-            # easy-to-append include length-specified strings of type `|S*` or `<U*` (where * is a
-            # positive integer character length). For these dtypes, appending dissimilar lengths
-            # can result in truncation of appended data. Therefore, variables which already exist
-            # in the dataset, and with dtypes which are not known to be easy-to-append, necessitate
-            # exact dtype equality, as checked below.
-            pass
-        elif not var.dtype == existing_vars[vname].dtype:
-            raise ValueError(
-                f"Mismatched dtypes for variable {vname} between Zarr store on disk "
-                f"and dataset to append. Store has dtype {existing_vars[vname].dtype} but "
-                f"dataset to append has dtype {var.dtype}."
-            )
-
-    for vname, var in dataset.data_vars.items():
-        check_dtype(vname, var)
-
-
 # compute=True returns ZarrStore
 @overload
 def to_zarr(
@@ -1721,17 +1685,16 @@ def to_zarr(
             )
 
     if mode in ["a", "a-", "r+"]:
-        _validate_datatypes_for_zarr_append(zstore, dataset)
-        if append_dim is not None:
+        existing_var_names = set(zstore.zarr_group.array_keys())
+        if append_dim is not None and append_dim not in existing_var_names:
             existing_dims = zstore.get_dimensions()
             if append_dim not in existing_dims:
                 raise ValueError(
                     f"append_dim={append_dim!r} does not match any existing "
                     f"dataset dimensions {existing_dims}"
                 )
-        existing_var_names = set(zstore.zarr_group.array_keys())
         for var_name in existing_var_names:
-            if var_name in encoding.keys():
+            if var_name in encoding:
                 raise ValueError(
                     f"variable {var_name!r} already exists, but encoding was provided"
                 )
diff --git a/xarray/backends/zarr.py b/xarray/backends/zarr.py
index 6f21881e0fc..ff8a4996702 100644
--- a/xarray/backends/zarr.py
+++ b/xarray/backends/zarr.py
@@ -324,6 +324,34 @@ def encode_zarr_variable(var, needs_copy=True, name=None):
     return var
 
 
+def _validate_datatypes_for_zarr_append(vname, existing_var, new_var):
+    """If variable exists in the store, confirm dtype of the data to append is compatible with
+    existing dtype.
+    """
+    if (
+        np.issubdtype(new_var.dtype, np.number)
+        or np.issubdtype(new_var.dtype, np.datetime64)
+        or np.issubdtype(new_var.dtype, np.bool_)
+        or new_var.dtype == object
+    ):
+        # We can skip dtype equality checks under two conditions: (1) if the var to append is
+        # new to the dataset, because in this case there is no existing var to compare it to;
+        # or (2) if var to append's dtype is known to be easy-to-append, because in this case
+        # we can be confident appending won't cause problems. Examples of dtypes which are not
+        # easy-to-append include length-specified strings of type `|S*` or `<U*` (where * is a
+        # positive integer character length). For these dtypes, appending dissimilar lengths
+        # can result in truncation of appended data. Therefore, variables which already exist
+        # in the dataset, and with dtypes which are not known to be easy-to-append, necessitate
+        # exact dtype equality, as checked below.
+        pass
+    elif not new_var.dtype == existing_var.dtype:
+        raise ValueError(
+            f"Mismatched dtypes for variable {vname} between Zarr store on disk "
+            f"and dataset to append. Store has dtype {existing_var.dtype} but "
+            f"dataset to append has dtype {new_var.dtype}."
+        )
+
+
 def _validate_and_transpose_existing_dims(
     var_name, new_var, existing_var, region, append_dim
 ):
@@ -632,6 +660,10 @@ def store(
             # Modified variables must use the same encoding as the store.
             vars_with_encoding = {}
             for vn in existing_variable_names:
+                if self._mode in ["a", "a-", "r+"]:
+                    _validate_datatypes_for_zarr_append(
+                        vn, existing_vars[vn], variables[vn]
+                    )
                 vars_with_encoding[vn] = variables[vn].copy(deep=False)
                 vars_with_encoding[vn].encoding = existing_vars[vn].encoding
             vars_with_encoding, _ = self.encode(vars_with_encoding, {})