Pass variable name to encode_zarr_variable

[slevang]

• Closes Serialization fails with xarray>=2024.1.0 xarray-contrib/xeofs#148
• Tests added
• User visible changes (including notable bug fixes) are documented in whats-new.rst

The change from #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.

[max-sixty]

While I don't know as much as others about the code, this looks good to me

[slevang]

Anyone else familiar with the backends have any objections to this, @dcherian @jhamman ?

[jhamman]

Thanks for this @slevang! One comment on backward compat but otherwise looks good.

[dcherian]

I think the real error is that resetting the MultiIndex doesn't clear out the index for the levels:

Using

ds = ds.reset_index(["feature", "dim1", "dim2"])

makes it all work. I don't think this change is needed.

See #8887. I don't understand why only Zarr errors out though.

[slevang]

Here's an approximate stack of function calls that leads to this being an issue, and why it only applies to zarr:

# ZarrStore inherits from:
class AbstractWritableDataStore:
    def encode():
        variables = {k: self.encode_variable(v) for k, v in variables.items()}
            encode_zarr_variable(variable) # <- NO NAME PASSED HERE
                conventions.encode_cf_variable(var, name=name) # <- name is now always None
                    ensure_not_multiindex(var, name=name) # <- fails

# NetCDF4DataStore/H5NetCDFArrayWrapper inherit from:
class WritableCFDataStore(AbstractWritableDataStore):
    # which overrides encode:
    def encode():
        conventions.cf_encoder(variables, attributes)
            new_vars = {k: encode_cf_variable(v, name=k) for k, v in variables.items()} # <- NAME IS PASSED HERE
                ensure_not_multiindex(var, name=name) # <- all good

I don't know the background on why this structure exists, but it does lead to a subtle but fundamental difference in the multiindex check specifically.

#8888 interestingly fixes the test I added here, but still doesn't fix the zarr serialization issue with DataTree in my original post. Not sure why.

[benbovy]

but still doesn't fix the zarr serialization issue with DataTree in my original post. Not sure why.

The bug likely happens when creating a new dataset by passing a DataArray object wrapping a multi-index (i.e., ds = xr.Dataset(data_vars={"feature": da.feature})), as explained in #8888 (comment).

I would advise against doing that if the goal is to pass the whole thing, i.e., the multi-index and both its dimension and level coordinate variables (I'd like to deprecate it in #8140). Instead I'd strongly suggest passing all multi-index coordinates explicitly with ds = xr.Dataset(coords=da.coords).

[slevang]

Ah that does it, thanks!

Happy to either leave this up or close it then - understand this was a pretty niche issue.

[slevang]

I'll close this since it's probably unnecessary but seems like we should try to push #8140 through to prevent this sort of confusion in the future.