Fix netcdf encoding (#95)

* the fix that didn't fix

* further work on the netcdf encoding issue

* use set2

* check for invalid groups in to_zarr

* add tests and comments for the future

Co-authored-by: Justin Magers <[email protected]>

Author: Joe Hamman

datatree/io.py

@@ -105,13 +105,6 @@ def _open_datatree_zarr(store, **kwargs) -> DataTree:
     return tree_root
 
 
-def _maybe_extract_group_kwargs(enc, group):
-    try:
-        return enc[group]
-    except KeyError:
-        return None
-
-
 def _create_empty_netcdf_group(filename, group, mode, engine):
     ncDataset = _get_nc_dataset_class(engine)
 
@@ -146,6 +139,14 @@ def _datatree_to_netcdf(
     if encoding is None:
         encoding = {}
 
+    # In the future, we may want to expand this check to insure all the provided encoding
+    # options are valid. For now, this simply checks that all provided encoding keys are
+    # groups in the datatree.
+    if set(encoding) - set(dt.groups):
+        raise ValueError(
+            f"unexpected encoding group name(s) provided: {set(encoding) - set(dt.groups)}"
+        )
+
     if unlimited_dims is None:
         unlimited_dims = {}
 
@@ -155,16 +156,15 @@ def _datatree_to_netcdf(
         if ds is None:
             _create_empty_netcdf_group(filepath, group_path, mode, engine)
         else:
-
             ds.to_netcdf(
                 filepath,
                 group=group_path,
                 mode=mode,
-                encoding=_maybe_extract_group_kwargs(encoding, dt.path),
-                unlimited_dims=_maybe_extract_group_kwargs(unlimited_dims, dt.path),
+                encoding=encoding.get(node.path),
+                unlimited_dims=unlimited_dims.get(node.path),
                 **kwargs,
             )
-        mode = "a"
+        mode = "r+"
 
 
 def _create_empty_zarr_group(store, group, mode):
@@ -196,6 +196,14 @@ def _datatree_to_zarr(
     if encoding is None:
         encoding = {}
 
+    # In the future, we may want to expand this check to insure all the provided encoding
+    # options are valid. For now, this simply checks that all provided encoding keys are
+    # groups in the datatree.
+    if set(encoding) - set(dt.groups):
+        raise ValueError(
+            f"unexpected encoding group name(s) provided: {set(encoding) - set(dt.groups)}"
+        )
+
     for node in dt.subtree:
         ds = node.ds
         group_path = node.path
@@ -206,7 +214,7 @@ def _datatree_to_zarr(
                 store,
                 group=group_path,
                 mode=mode,
-                encoding=_maybe_extract_group_kwargs(encoding, dt.path),
+                encoding=encoding.get(node.path),
                 consolidated=False,
                 **kwargs,
             )

datatree/tests/test_io.py

@@ -18,6 +18,27 @@ def test_to_netcdf(self, tmpdir):
         roundtrip_dt = open_datatree(filepath)
         assert_equal(original_dt, roundtrip_dt)
 
+    @requires_netCDF4
+    def test_netcdf_encoding(self, tmpdir):
+        filepath = str(
+            tmpdir / "test.nc"
+        )  # casting to str avoids a pathlib bug in xarray
+        original_dt = create_test_datatree()
+
+        # add compression
+        comp = dict(zlib=True, complevel=9)
+        enc = {"/set2": {var: comp for var in original_dt["/set2"].ds.data_vars}}
+
+        original_dt.to_netcdf(filepath, encoding=enc, engine="netcdf4")
+        roundtrip_dt = open_datatree(filepath)
+
+        assert roundtrip_dt["/set2/a"].encoding["zlib"] == comp["zlib"]
+        assert roundtrip_dt["/set2/a"].encoding["complevel"] == comp["complevel"]
+
+        enc["/not/a/group"] = {"foo": "bar"}
+        with pytest.raises(ValueError, match="unexpected encoding group.*"):
+            original_dt.to_netcdf(filepath, encoding=enc, engine="netcdf4")
+
     @requires_h5netcdf
     def test_to_h5netcdf(self, tmpdir):
         filepath = str(
@@ -40,6 +61,27 @@ def test_to_zarr(self, tmpdir):
         roundtrip_dt = open_datatree(filepath, engine="zarr")
         assert_equal(original_dt, roundtrip_dt)
 
+    @requires_zarr
+    def test_zarr_encoding(self, tmpdir):
+        import zarr
+
+        filepath = str(
+            tmpdir / "test.zarr"
+        )  # casting to str avoids a pathlib bug in xarray
+        original_dt = create_test_datatree()
+
+        comp = {"compressor": zarr.Blosc(cname="zstd", clevel=3, shuffle=2)}
+        enc = {"/set2": {var: comp for var in original_dt["/set2"].ds.data_vars}}
+        original_dt.to_zarr(filepath, encoding=enc)
+        roundtrip_dt = open_datatree(filepath, engine="zarr")
+
+        print(roundtrip_dt["/set2/a"].encoding)
+        assert roundtrip_dt["/set2/a"].encoding["compressor"] == comp["compressor"]
+
+        enc["/not/a/group"] = {"foo": "bar"}
+        with pytest.raises(ValueError, match="unexpected encoding group.*"):
+            original_dt.to_zarr(filepath, encoding=enc, engine="zarr")
+
     @requires_zarr
     def test_to_zarr_zip_store(self, tmpdir):
         from zarr.storage import ZipStore