Prefer metadata over store reads in _decode_cf_datetime_dtype

doc/whats-new.rst

@@ -61,10 +61,16 @@ New Features
   variables of a section at once, and shows the expanded attributes and data
   more compactly and indented below their variable (:pull:`11284`).
   By `Michael Niklas <https://github.com/headtr1ck>`_.
+- :py:class:`~xarray.coders.CFDatetimeCoder` accepts ``infer_dtype``. With
+  ``infer_dtype="metadata"`` the decoded dtype is inferred from ``units``,
+  ``calendar``, ``use_cftime`` and ``time_unit`` instead of reading the first
+  and last values of each time variable, which makes opening remote datasets
+  with many time variables much faster. The default ``infer_dtype="data"``
+  keeps the current behavior (:issue:`11303`, :pull:`11304`).
+  By `Alfonso Ladino <https://github.com/aladinor>`_.
 - Use ``check_string_dtype()`` to detect variable-length strings in h5netcdf backend (:pull:`11789`).
   By `Kai Mühlbauer <https://github.com/kmuehlbauer>`_.
 
-
 Breaking Changes
 ~~~~~~~~~~~~~~~~
 - The ``interpolation`` argument of ``quantile`` on :py:class:`Dataset`,

xarray/coding/times.py

@@ -6,7 +6,7 @@
 from collections.abc import Callable, Hashable
 from datetime import datetime, timedelta
 from functools import partial
-from typing import TYPE_CHECKING, Union, cast, override
+from typing import TYPE_CHECKING, Literal, Union, cast, override
 
 import numpy as np
 import pandas as pd
@@ -336,21 +336,13 @@ def _unpack_time_units_and_ref_date_cftime(units: str, calendar: str):
     return time_units, ref_date
 
 
-def _decode_cf_datetime_dtype(
-    data,
+def _decoded_dtype(
+    example_value: np.ndarray,
     units: str,
     calendar: str | None,
     use_cftime: bool | None,
-    time_unit: PDDatetimeUnitOptions = "ns",
+    time_unit: PDDatetimeUnitOptions,
 ) -> np.dtype:
-    # Verify that at least the first and last date can be decoded
-    # successfully. Otherwise, tracebacks end up swallowed by
-    # Dataset.__repr__ when users try to view their lazily decoded array.
-    values = indexing.ImplicitToExplicitIndexingAdapter(indexing.as_indexable(data))
-    example_value = np.concatenate(
-        [to_numpy(first_n_items(values, 1)), to_numpy(last_item(values))]
-    )
-
     try:
         result = decode_cf_datetime(
             example_value, units, calendar, use_cftime, time_unit
@@ -371,6 +363,68 @@ def _decode_cf_datetime_dtype(
     return dtype
 
 
+def _decode_cf_datetime_dtype(
+    data,
+    units: str,
+    calendar: str | None,
+    use_cftime: bool | None,
+    time_unit: PDDatetimeUnitOptions = "ns",
+) -> np.dtype:
+    # Verify that at least the first and last date can be decoded
+    # successfully. Otherwise, tracebacks end up swallowed by
+    # Dataset.__repr__ when users try to view their lazily decoded array.
+    values = indexing.ImplicitToExplicitIndexingAdapter(indexing.as_indexable(data))
+    example_value = np.concatenate(
+        [to_numpy(first_n_items(values, 1)), to_numpy(last_item(values))]
+    )
+    return _decoded_dtype(example_value, units, calendar, use_cftime, time_unit)
+
+
+def _decode_cf_datetime_dtype_from_metadata(
+    units: str,
+    calendar: str | None,
+    use_cftime: bool | None,
+    time_unit: PDDatetimeUnitOptions,
+) -> np.dtype | None:
+    """Infer the decoded dtype from a synthetic ``[0, 1]``, or None if unsure."""
+    try:
+        with warnings.catch_warnings():
+            # this warning would describe the synthetic values, not the data
+            warnings.filterwarnings(
+                "ignore", "Unable to decode time axis", SerializationWarning
+            )
+            dtype = _decoded_dtype(
+                np.array([0, 1]), units, calendar, use_cftime, time_unit
+            )
+    except ValueError:
+        # e.g. "days since 2262-04-11": the synthetic dates overflow although
+        # the stored ones may not
+        return None
+    if (
+        dtype == np.dtype("object")
+        and use_cftime is None
+        and (calendar is None or _is_standard_calendar(calendar))
+    ):
+        # The reference date cannot be represented as datetime64 (e.g.
+        # "hours since 1-1-1"), so dates near it say nothing about the stored
+        # dates, which may well be in 2020.
+        return None
+    return dtype
+
+
+def _decode_cf_datetime_as(num_dates, *, dtype: np.dtype, **kwargs) -> np.ndarray:
+    """decode_cf_datetime, raising if the result does not have ``dtype``."""
+    result = decode_cf_datetime(num_dates, **kwargs)
+    if result.dtype != dtype:
+        raise ValueError(
+            f"time values with units {kwargs['units']!r} decode to "
+            f"{result.dtype}, not to {dtype} as inferred from metadata. Open "
+            "the dataset with decode_times=CFDatetimeCoder(infer_dtype='data'), "
+            "or set use_cftime or time_unit to match the stored values."
+        )
+    return result
+
+
 def _decode_datetime_with_cftime(
     num_dates: np.ndarray, units: str, calendar: str
 ) -> np.ndarray:
@@ -1381,15 +1435,49 @@ class CFDatetimeCoder(VariableCoder):
         May not be supported by all the backends.
     time_unit : PDDatetimeUnitOptions
           Target resolution when decoding dates. Defaults to "ns".
+    infer_dtype : {"data", "metadata"}, default: "data"
+        How to infer the dtype of the decoded variable:
+
+        - ``"data"`` (default): read the first and last values from the
+          store to verify they decode successfully. Preserves
+          backwards-compatible behavior, including the auto-fallback to
+          ``cftime.datetime`` for values outside the
+          ``np.datetime64[ns]`` range.
+        - ``"metadata"``: infer the dtype from ``units``, ``calendar``,
+          ``use_cftime`` and ``time_unit`` alone, assuming the stored dates
+          lie near the reference date in ``units``. Nothing is read from
+          the store, which speeds up opening remote stores with many
+          time-encoded variables. The first and last values are read, as
+          with ``"data"``, only when metadata cannot tell: when, with
+          ``use_cftime=None`` and a standard calendar, dates near the
+          reference date need ``cftime`` (e.g. ``"hours since 1-1-1"`` or a
+          reference date before 1582-10-15), or when decoding dates near
+          the reference date fails.
+
+          The stored values are checked when they are decoded instead, and a
+          ``ValueError`` is raised then if they need a different dtype than
+          inferred. This happens for dates outside the
+          ``np.datetime64[time_unit]`` range (pass ``use_cftime=True`` or a
+          coarser ``time_unit`` such as ``"s"``) and for floating point
+          values with fractions finer than ``units``, e.g. ``0.5`` with
+          ``"seconds since ..."`` (pass a ``time_unit`` fine enough to
+          represent them). Index coordinates are decoded when the dataset
+          is opened, other variables when they are loaded.
     """
 
     def __init__(
         self,
         use_cftime: bool | None = None,
         time_unit: PDDatetimeUnitOptions = "ns",
+        infer_dtype: Literal["data", "metadata"] = "data",
     ) -> None:
+        if infer_dtype not in ("data", "metadata"):
+            raise ValueError(
+                f"infer_dtype must be 'data' or 'metadata', got {infer_dtype!r}"
+            )
         self.use_cftime = use_cftime
         self.time_unit = time_unit
+        self.infer_dtype = infer_dtype
 
     @override
     def encode(self, variable: Variable, name: T_Name = None) -> Variable:
@@ -1424,11 +1512,22 @@ def decode(self, variable: Variable, name: T_Name = None) -> Variable:
 
             units = pop_to(attrs, encoding, "units")
             calendar = pop_to(attrs, encoding, "calendar")
-            dtype = _decode_cf_datetime_dtype(
-                data, units, calendar, self.use_cftime, self.time_unit
-            )
+            dtype = None
+            if self.infer_dtype == "metadata":
+                dtype = _decode_cf_datetime_dtype_from_metadata(
+                    units, calendar, self.use_cftime, self.time_unit
+                )
+            if dtype is None:
+                dtype = _decode_cf_datetime_dtype(
+                    data, units, calendar, self.use_cftime, self.time_unit
+                )
+                decode = decode_cf_datetime
+            else:
+                # nothing was read, so check every decoded chunk against the
+                # inferred dtype instead of letting it be cast silently
+                decode = partial(_decode_cf_datetime_as, dtype=dtype)
             transform = partial(
-                decode_cf_datetime,
+                decode,
                 units=units,
                 calendar=calendar,
                 use_cftime=self.use_cftime,

xarray/tests/backends/test_zarr.py

@@ -2756,6 +2756,74 @@ def test_zarr_create_default_indexes(tmp_path, create_default_indexes) -> None:
             assert len(loaded_ds.xindexes) == 0
 
 
+def _write_encoded_times(store_path, values, units, zarr_format) -> None:
+    # "time" is an index coordinate (decoded at open), "valid_time" is not
+    # (decoded on load); both are stored with one value per chunk
+    encoded = {"units": units}
+    xr.Dataset(
+        {"valid_time": ("time", np.array(values), encoded)},
+        coords={"time": ("time", np.array(values), encoded)},
+    ).to_zarr(
+        store_path,
+        mode="w",
+        zarr_format=zarr_format,
+        encoding={"time": {"chunks": (1,)}, "valid_time": {"chunks": (1,)}},
+    )
+
+
+@requires_zarr
+@requires_dask
+@requires_cftime
+@parametrize_zarr_format
+@pytest.mark.filterwarnings("ignore:Ambiguous reference date string")
+@pytest.mark.parametrize(
+    ("units", "values"),
+    [
+        ("hours since 2000-01-01", [0.0, 1.0, 2.0]),
+        # dates near the reference date need cftime, but these are in 2020
+        ("hours since 1-1-1 00:00:0.0", [17_700_000.0, 17_700_001.0, 17_700_002.0]),
+    ],
+)
+def test_zarr_infer_dtype_metadata_round_trip(
+    tmp_path, zarr_format, units, values
+) -> None:
+    store_path = tmp_path / "tmp.zarr"
+    _write_encoded_times(store_path, values, units, zarr_format)
+
+    with (
+        open_dataset(store_path, engine="zarr", chunks={}) as expected,
+        open_dataset(
+            store_path,
+            engine="zarr",
+            chunks={},
+            decode_times=CFDatetimeCoder(infer_dtype="metadata"),
+        ) as actual,
+    ):
+        assert actual.time.dtype == expected.time.dtype
+        assert actual.valid_time.dtype == expected.valid_time.dtype
+        assert_identical(actual.compute(), expected.compute())
+
+
+@requires_zarr
+@requires_dask
+@requires_cftime
+@parametrize_zarr_format
+@pytest.mark.filterwarnings("ignore:Unable to decode time axis")
+def test_zarr_infer_dtype_metadata_wrong_dtype_raises(tmp_path, zarr_format) -> None:
+    # dates in 2821, beyond datetime64[ns]: metadata infers datetime64[ns]
+    store_path = tmp_path / "tmp.zarr"
+    _write_encoded_times(store_path, [0, 300_000], "days since 2000-01-01", zarr_format)
+
+    with pytest.raises(ValueError, match="as inferred from metadata"):
+        # the index coordinate is decoded when the dataset is opened
+        open_dataset(
+            store_path,
+            engine="zarr",
+            chunks={},
+            decode_times=CFDatetimeCoder(infer_dtype="metadata"),
+        )
+
+
 @requires_zarr
 @pytest.mark.usefixtures("default_zarr_format")
 def test_raises_key_error_on_invalid_zarr_store(tmp_path):