Repository: zarr-developers/VirtualiZarr

Issue:
Option to interpret variables using cftime
https://github.com/TomNicholas/VirtualiZarr/issues/68#issuecomment-2113983642 explains how we don't currently have support for virtualizing encoded time variables. 

To handle this it looks like we need to an another option to `open_virtual_dataset`, a `cftime_vars` kwarg, which specifies which variables should be interpreted using cftime. (I think `cftime_vars` should be a subset of `loadable_vars`.) This suggestion comes from how kerchunk deals with this problem:

https://fsspec.github.io/kerchunk/reference.html#kerchunk.combine.MultiZarrToZarr

```
”cf:{var}”, interpret the value of var using cftime, returning a datetime. 
These will be automatically re-encoded with cftime, unless you specify an “M8[*]” 
dtype for the coordinate, in which case a conversion will be attempted.
```

_Originally posted by @dcherian in https://github.com/TomNicholas/VirtualiZarr/issues/73#issuecomment-2047989180_
            

Discussion:
jbusecke: If folks here have a minute, I just tested this PR over at https://github.com/jbusecke/esgf-virtual-zarr-data-access/pull/10 with real world CMIP data. I think it generally works, but the attributes of the resulting time coordinate are different. Would love to get some input on that.
jsignell: Thanks for the review! I think I have addressed all the comments. 
jsignell: Ok @TomNicholas I am pausing. Let me know how this is looking to you

TomNicholas: Thanks for going so deep into this @jsignell! To add more confusion - do you think what I'm seeing in https://github.com/zarr-developers/VirtualiZarr/pull/154#discussion_r1650170676 is related to your close-but-not-equal issue?
TomNicholas: @jsignell you might want to merge in #154 and #158.

Also sorry to push the finish line farther away, but what do you think about the idea in https://github.com/zarr-developers/VirtualiZarr/issues/117#issuecomment-2187406127?
jsignell: > @jsignell you might want to merge in #154 and #158.

✅ 

> Also sorry to push the finish line farther away, but what do you think about the idea in [#117 (comment)](https://github.com/zarr-developers/VirtualiZarr/issues/117#issuecomment-2187406127)?

No worries, better to do this the right way. I modeled the approach in this PR on kerchunk, but I will take a look at what is going on in xarray.

TomNicholas: > No worries, better to do this the right way. I modeled the approach in this PR on kerchunk, but I will take a look at what is going on in xarray.

❤️ 

So I was looking at this in more detail last night, and I think we might just want to replace your `encode_cftime` and `recode_cftime` with importing [xarray's `CFDatetimeCoder`](https://github.com/pydata/xarray/blob/b5180749d351f8b85fd39677bf137caaa90288a7/xarray/coding/times.py#L958) (and possibly the [`CFTimeDeltaEncoder`](https://github.com/pydata/xarray/blob/b5180749d351f8b85fd39677bf137caaa90288a7/xarray/coding/times.py#L1001) too). As far as I can tell (and please tell me if I'm wrong) those are basically more mature versions of the functions you've written here. Those classes are private but stable interfaces, and using them here could be an argument for factoring them out of xarray at some point (or exposing them more publicly).

Those encoders are called by `xr.decode_cf` [here](https://github.com/pydata/xarray/blob/56209bd9a3192e4f1e82c21e5ffcf4c3bacaaae3/xarray/conventions.py#L288), but replacing everything with `xr.decode_cf` seems like a longer-term idea that is beyond the scope of this PR (but using xarray's encoders would build towards it).
jsignell: Ok I think I've finally got it working.
jsignell: Let me know if you are looking for any more changes or feedback! Otherwise this is done from my perspective 😸 

Patch:
diff --git a/docs/releases.rst b/docs/releases.rst
index 1594b8c..8400a16 100644
--- a/docs/releases.rst
+++ b/docs/releases.rst
@@ -11,6 +11,8 @@ New Features
 
 - Added a `.rename_paths` convenience method to rename paths in a manifest according to a function.
   (:pull:`152`) By `Tom Nicholas <https://github.com/TomNicholas>`_.
+- New ``cftime_variables`` option on ``open_virtual_dataset`` enables encoding/decoding time.
+  (:pull:`122`) By `Julia Signell <https://github.com/jsignell>`_.
 
 Breaking changes
 ~~~~~~~~~~~~~~~~
diff --git a/virtualizarr/kerchunk.py b/virtualizarr/kerchunk.py
index 46b4f2b..c659804 100644
--- a/virtualizarr/kerchunk.py
+++ b/virtualizarr/kerchunk.py
@@ -7,6 +7,7 @@ from typing import Any, NewType, Optional, cast
 import numpy as np
 import ujson  # type: ignore
 import xarray as xr
+from xarray.coding.times import CFDatetimeCoder
 
 from virtualizarr.manifests.manifest import join
 from virtualizarr.utils import _fsspec_openfile_from_filepath
@@ -49,6 +50,8 @@ class NumpyEncoder(json.JSONEncoder):
             return obj.tolist()  # Convert NumPy array to Python list
         elif isinstance(obj, np.generic):
             return obj.item()  # Convert NumPy scalar to Python scalar
+        elif isinstance(obj, np.dtype):
+            return str(obj)
         return json.JSONEncoder.default(self, obj)
 
 
@@ -274,9 +277,7 @@ def variable_to_kerchunk_arr_refs(var: xr.Variable, var_name: str) -> KerchunkAr
                     f"Cannot serialize loaded variable {var_name}, as it is encoded with an offset"
                 )
             if "calendar" in var.encoding:
-                raise NotImplementedError(
-                    f"Cannot serialize loaded variable {var_name}, as it is encoded with a calendar"
-                )
+                np_arr = CFDatetimeCoder().encode(var.copy(), name=var_name).values
 
         # This encoding is what kerchunk does when it "inlines" data, see https://github.com/fsspec/kerchunk/blob/a0c4f3b828d37f6d07995925b324595af68c4a19/kerchunk/hdf.py#L472
         byte_data = np_arr.tobytes()
@@ -297,7 +298,7 @@ def variable_to_kerchunk_arr_refs(var: xr.Variable, var_name: str) -> KerchunkAr
     zarray_dict = zarray.to_kerchunk_json()
     arr_refs[".zarray"] = zarray_dict
 
-    zattrs = var.attrs
+    zattrs = {**var.attrs, **var.encoding}
     zattrs["_ARRAY_DIMENSIONS"] = list(var.dims)
     arr_refs[".zattrs"] = json.dumps(zattrs, separators=(",", ":"), cls=NumpyEncoder)
 
diff --git a/virtualizarr/xarray.py b/virtualizarr/xarray.py
index be2bf09..a90f28e 100644
--- a/virtualizarr/xarray.py
+++ b/virtualizarr/xarray.py
@@ -11,6 +11,7 @@ import ujson  # type: ignore
 import xarray as xr
 from xarray import register_dataset_accessor
 from xarray.backends import BackendArray
+from xarray.coding.times import CFDatetimeCoder
 from xarray.core.indexes import Index, PandasIndex
 from xarray.core.variable import IndexVariable
 
@@ -33,9 +34,11 @@ class ManifestBackendArray(ManifestArray, BackendArray):
 
 def open_virtual_dataset(
     filepath: str,
+    *,
     filetype: FileType | None = None,
     drop_variables: Iterable[str] | None = None,
     loadable_variables: Iterable[str] | None = None,
+    cftime_variables: Iterable[str] | None = None,
     indexes: Mapping[str, Index] | None = None,
     virtual_array_class=ManifestArray,
     reader_options: Optional[dict] = {
@@ -69,6 +72,10 @@ def open_virtual_dataset(
     virtual_array_class
         Virtual array class to use to represent the references to the chunks in each on-disk array.
         Currently can only be ManifestArray, but once VirtualZarrArray is implemented the default should be changed to that.
+    cftime_variables : list[str], default is None
+        Interpret the value of specified vars using cftime, returning a datetime.
+        These will be automatically re-encoded with cftime. This list must be a subset
+        of ``loadable_variables``.
     reader_options: dict, default {'storage_options':{'key':'', 'secret':'', 'anon':True}}
         Dict passed into Kerchunk file readers. Note: Each Kerchunk file reader has distinct arguments,
         so ensure reader_options match selected Kerchunk reader arguments.
@@ -95,6 +102,20 @@ def open_virtual_dataset(
     if common:
         raise ValueError(f"Cannot both load and drop variables {common}")
 
+    if cftime_variables is None:
+        cftime_variables = []
+    elif isinstance(cftime_variables, str):
+        cftime_variables = [cftime_variables]
+    else:
+        cftime_variables = list(cftime_variables)
+
+    if diff := (set(cftime_variables) - set(loadable_variables)):
+        missing_str = ", ".join([f"'{v}'" for v in diff])
+        raise ValueError(
+            "All ``cftime_variables`` must be included in ``loadable_variables`` "
+            f"({missing_str} not in ``loadable_variables``)"
+        )
+
     if virtual_array_class is not ManifestArray:
         raise NotImplementedError()
 
@@ -127,7 +148,9 @@ def open_virtual_dataset(
                 filepath=filepath, reader_options=reader_options
             )
 
-            ds = xr.open_dataset(fpath, drop_variables=drop_variables)
+            ds = xr.open_dataset(
+                fpath, drop_variables=drop_variables, decode_times=False
+            )
 
             if indexes is None:
                 # add default indexes by reading data from file
@@ -144,6 +167,10 @@ def open_virtual_dataset(
                 if name in loadable_variables
             }
 
+            for name in cftime_variables:
+                var = loadable_vars[name]
+                loadable_vars[name] = CFDatetimeCoder().decode(var, name=name)
+
             # if we only read the indexes we can just close the file right away as nothing is lazy
             if loadable_vars == {}:
                 ds.close()
@@ -326,7 +353,7 @@ def separate_coords(
             # use workaround to avoid creating IndexVariables described here https://github.com/pydata/xarray/pull/8107#discussion_r1311214263
             if len(var.dims) == 1:
                 dim1d, *_ = var.dims
-                coord_vars[name] = (dim1d, var.data, var.attrs)
+                coord_vars[name] = (dim1d, var.data, var.attrs, var.encoding)
 
                 if isinstance(var, IndexVariable):
                     # unless variable actually already is a loaded IndexVariable,


Test patch:
diff --git a/virtualizarr/tests/test_integration.py b/virtualizarr/tests/test_integration.py
index ec6a98a..8b84b28 100644
--- a/virtualizarr/tests/test_integration.py
+++ b/virtualizarr/tests/test_integration.py
@@ -74,9 +74,10 @@ class TestKerchunkRoundtrip:
         # assert identical to original dataset
         xrt.assert_identical(roundtrip, ds)
 
-    def test_kerchunk_roundtrip_concat(self, tmpdir, format):
+    @pytest.mark.parametrize("decode_times,time_vars", [(False, []), (True, ["time"])])
+    def test_kerchunk_roundtrip_concat(self, tmpdir, format, decode_times, time_vars):
         # set up example xarray dataset
-        ds = xr.tutorial.open_dataset("air_temperature", decode_times=False)
+        ds = xr.tutorial.open_dataset("air_temperature", decode_times=decode_times)
 
         # split into two datasets
         ds1, ds2 = ds.isel(time=slice(None, 1460)), ds.isel(time=slice(1460, None))
@@ -86,8 +87,25 @@ class TestKerchunkRoundtrip:
         ds2.to_netcdf(f"{tmpdir}/air2.nc")
 
         # use open_dataset_via_kerchunk to read it as references
-        vds1 = open_virtual_dataset(f"{tmpdir}/air1.nc", indexes={})
-        vds2 = open_virtual_dataset(f"{tmpdir}/air2.nc", indexes={})
+        vds1 = open_virtual_dataset(
+            f"{tmpdir}/air1.nc",
+            indexes={},
+            loadable_variables=time_vars,
+            cftime_variables=time_vars,
+        )
+        vds2 = open_virtual_dataset(
+            f"{tmpdir}/air2.nc",
+            indexes={},
+            loadable_variables=time_vars,
+            cftime_variables=time_vars,
+        )
+
+        if decode_times is False:
+            assert vds1.time.dtype == np.dtype("float32")
+        else:
+            assert vds1.time.dtype == np.dtype("<M8[ns]")
+            assert "units" in vds1.time.encoding
+            assert "calendar" in vds1.time.encoding
 
         # concatenate virtually along time
         vds = xr.concat([vds1, vds2], dim="time", coords="minimal", compat="override")
@@ -97,18 +115,26 @@ class TestKerchunkRoundtrip:
             ds_refs = vds.virtualize.to_kerchunk(format=format)
 
             # use fsspec to read the dataset from the kerchunk references dict
-            roundtrip = xr.open_dataset(ds_refs, engine="kerchunk", decode_times=False)
+            roundtrip = xr.open_dataset(
+                ds_refs, engine="kerchunk", decode_times=decode_times
+            )
         else:
             # write those references to disk as kerchunk references format
             vds.virtualize.to_kerchunk(f"{tmpdir}/refs.{format}", format=format)
 
             # use fsspec to read the dataset from disk via the kerchunk references
             roundtrip = xr.open_dataset(
-                f"{tmpdir}/refs.{format}", engine="kerchunk", decode_times=False
+                f"{tmpdir}/refs.{format}", engine="kerchunk", decode_times=decode_times
             )
-
-        # assert identical to original dataset
-        xrt.assert_identical(roundtrip, ds)
+        if decode_times is False:
+            # assert identical to original dataset
+            xrt.assert_identical(roundtrip, ds)
+        else:
+            # they are very very close! But assert_allclose doesn't seem to work on datetimes
+            assert (roundtrip.time - ds.time).sum() == 0
+            assert roundtrip.time.dtype == ds.time.dtype
+            assert roundtrip.time.encoding["units"] == ds.time.encoding["units"]
+            assert roundtrip.time.encoding["calendar"] == ds.time.encoding["calendar"]
 
     def test_non_dimension_coordinates(self, tmpdir, format):
         # regression test for GH issue #105
diff --git a/virtualizarr/tests/test_xarray.py b/virtualizarr/tests/test_xarray.py
index 5a0f1d3..36323c9 100644
--- a/virtualizarr/tests/test_xarray.py
+++ b/virtualizarr/tests/test_xarray.py
@@ -252,7 +252,7 @@ class TestOpenVirtualDatasetIndexes:
 
     def test_create_default_indexes(self, netcdf4_file):
         vds = open_virtual_dataset(netcdf4_file, indexes=None)
-        ds = xr.open_dataset(netcdf4_file)
+        ds = xr.open_dataset(netcdf4_file, decode_times=False)
 
         # TODO use xr.testing.assert_identical(vds.indexes, ds.indexes) instead once class supported by assertion comparison, see https://github.com/pydata/xarray/issues/5812
         assert index_mappings_equal(vds.xindexes, ds.xindexes)
@@ -344,7 +344,7 @@ class TestLoadVirtualDataset:
             else:
                 assert isinstance(vds[name].data, ManifestArray), name
 
-        full_ds = xr.open_dataset(netcdf4_file)
+        full_ds = xr.open_dataset(netcdf4_file, decode_times=False)
 
         for name in full_ds.variables:
             if name in vars_to_load:
@@ -413,3 +413,10 @@ class TestRenamePaths:
             == "s3://bucket/air.nc"
         )
         assert isinstance(renamed_vds["lat"].data, np.ndarray)
+
+
+def test_cftime_variables_must_be_in_loadable_variables(tmpdir):
+    ds = xr.Dataset(data_vars={"time": ["2024-06-21"]})
+    ds.to_netcdf(f"{tmpdir}/scalar.nc")
+    with pytest.raises(ValueError, match="'time' not in"):
+        open_virtual_dataset(f"{tmpdir}/scalar.nc", cftime_variables=["time"])