xarray icon indicating copy to clipboard operation
xarray copied to clipboard

Pass variable name to `encode_zarr_variable`

Open slevang opened this issue 11 months ago • 0 comments

  • [x] Closes https://github.com/xarray-contrib/xeofs/issues/148
  • [x] Tests added
  • [ ] User visible changes (including notable bug fixes) are documented in whats-new.rst

The change from https://github.com/pydata/xarray/pull/8672 mostly fixed the issue of serializing a reset multiindex in the backends, but there was an additional niche issue that turned up in xeofs that was causing serialization to still fail on the zarr backend.

The issue is that zarr is the only backend that uses a custom version of encode_cf_variable called encode_zarr_variable, and the way this gets called we don't pass through the name of the variable before running ensure_not_multiindex.

As a minimal fix, this PR just passes name through as an additional arg to the general encode_variable function. See @benbovy's comment that maybe we should actually unwrap the level coordinate in reset_index and clean up the checks in ensure_not_multiindex, but I wasn't able to get that working easily.

The exact workflow this turned up in involves DataTree and looks like this:

import numpy as np
import xarray as xr
from datatree import DataTree

# ND DataArray that gets stacked along a multiindex
da = xr.DataArray(np.ones((3, 3)), coords={"dim1": [1, 2, 3], "dim2": [4, 5, 6]})
da = da.stack(feature=["dim1", "dim2"])

# Extract just the stacked coordinates for saving in a dataset
ds = xr.Dataset(data_vars={"feature": da.feature})

# Reset the multiindex, which should make things serializable
ds = ds.reset_index("feature")
dt1 = DataTree()
dt2 = DataTree(name="feature", data=ds)
dt1["foo"] = dt2
# Somehow in this step, dt1.foo.feature.dim1.variable becomes an IndexVariable again
print(type(dt1.foo.feature.dim1.variable))

# Works
dt1.to_netcdf("test.nc", mode="w")
# Fails
dt1.to_zarr("test.zarr", mode="w")

But we can reproduce in xarray with the test added here.

slevang avatar Mar 06 '24 16:03 slevang