diff --git a/doc/whats-new.rst b/doc/whats-new.rst index a9e8feb905c..21f3671188a 100644 --- a/doc/whats-new.rst +++ b/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 `_. +- :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 `_. - Use ``check_string_dtype()`` to detect variable-length strings in h5netcdf backend (:pull:`11789`). By `Kai Mühlbauer `_. - Breaking Changes ~~~~~~~~~~~~~~~~ - The ``interpolation`` argument of ``quantile`` on :py:class:`Dataset`, diff --git a/xarray/coding/times.py b/xarray/coding/times.py index a29a93f2ebf..bfeba988fb2 100644 --- a/xarray/coding/times.py +++ b/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, diff --git a/xarray/tests/backends/test_zarr.py b/xarray/tests/backends/test_zarr.py index 5f181949f01..346c81d569b 100644 --- a/xarray/tests/backends/test_zarr.py +++ b/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): diff --git a/xarray/tests/test_coding_times.py b/xarray/tests/test_coding_times.py index 8234d5edd06..8bdb86de9e2 100644 --- a/xarray/tests/test_coding_times.py +++ b/xarray/tests/test_coding_times.py @@ -1,5 +1,6 @@ from __future__ import annotations +import re import warnings from datetime import datetime, timedelta from itertools import product, starmap @@ -46,6 +47,8 @@ _STANDARD_CALENDARS, DuckArrayWrapper, FirstElementAccessibleArray, + InaccessibleArray, + UnexpectedDataAccess, _all_cftime_date_types, arm_xfail, assert_array_equal, @@ -2224,3 +2227,194 @@ def test_roundtrip_empty_datetime64_array(time_unit: PDDatetimeUnitOptions) -> N ) assert_identical(variable, roundtripped) assert roundtripped.dtype == variable.dtype + + +def _time_variable(data, units: str, calendar: str | None = None) -> Variable: + attrs = {"units": units} | ({} if calendar is None else {"calendar": calendar}) + return Variable(("t",), data, attrs=attrs) + + +def _assert_metadata_matches_data(make_variable, **coder_kwargs) -> None: + # infer_dtype="metadata" must predict what infer_dtype="data" produces + # without reading the data: the same dtype, or the same failure. + def decode(infer_dtype): + coder = CFDatetimeCoder(infer_dtype=infer_dtype, **coder_kwargs) + return coder.decode(make_variable()) + + try: + expected = decode("data").dtype + except ValueError as err: + with pytest.raises(ValueError, match=re.escape(str(err))): + decode("metadata") + else: + decoded = decode("metadata") + assert decoded.dtype == expected + # the declared (lazy) dtype is also what loading produces + assert decoded.load().dtype == expected + + +class TestInferDtypeOptIn: + """``CFDatetimeCoder(infer_dtype="metadata")``: dtype without data reads.""" + + @pytest.mark.parametrize( + ("use_cftime", "calendar", "kind"), + [ + (None, "standard", "M"), + pytest.param(True, "standard", "O", marks=requires_cftime), + pytest.param(None, "noleap", "O", marks=requires_cftime), + ], + ) + def test_no_reads(self, use_cftime, calendar, kind, time_unit) -> None: + variable = _time_variable( + InaccessibleArray(np.zeros(3, dtype="int64")), + "days since 2000-01-01", + calendar, + ) + coder = CFDatetimeCoder( + use_cftime=use_cftime, time_unit=time_unit, infer_dtype="metadata" + ) + assert coder.decode(variable).dtype.kind == kind + + def test_data_default_reads_first_and_last(self) -> None: + variable = _time_variable( + InaccessibleArray(np.zeros(3, dtype="int64")), "days since 2000-01-01" + ) + with pytest.raises(UnexpectedDataAccess): + CFDatetimeCoder().decode(variable) + + def test_invalid_infer_dtype_raises(self) -> None: + with pytest.raises( + ValueError, match="infer_dtype must be 'data' or 'metadata'" + ): + CFDatetimeCoder(infer_dtype="bogus") # type: ignore[arg-type] + + @requires_cftime + @pytest.mark.parametrize("calendar", [None, "standard", "noleap", "360_day"]) + @pytest.mark.parametrize("use_cftime", [None, True, False]) + @pytest.mark.parametrize("time_unit", ["s", "ns"]) + @pytest.mark.parametrize( + "units", + [ + "days since 2000-01-01", + "seconds since 2000-01-01", + "milliseconds since 2000-01-01", + "nanoseconds since 2000-01-01", + "seconds since 2000-01-01 00:00:00.001", + "nonsense since whenever", + ], + ) + @pytest.mark.parametrize("dtype", ["int64", "float64"]) + def test_matches_data_path( + self, calendar, use_cftime, time_unit, units, dtype + ) -> None: + _assert_metadata_matches_data( + lambda: _time_variable(np.array([0, 1, 2], dtype=dtype), units, calendar), + use_cftime=use_cftime, + time_unit=time_unit, + ) + + @requires_cftime + @pytest.mark.filterwarnings("ignore:Ambiguous reference date string") + @pytest.mark.parametrize("calendar", [None, *_STANDARD_CALENDAR_NAMES]) + @pytest.mark.parametrize("time_unit", ["s", "ns"]) + @pytest.mark.parametrize( + ("units", "values"), + [ + # dates near the reference date need cftime, but the stored dates + # are in 2020 (or decoding dates near it fails) + ("hours since 1-1-1 00:00:0.0", [17_700_000.0, 17_700_024.0]), + ("days since 1500-01-01", [190_000, 190_001]), + ("days since 1677-01-01", [125_000, 125_001]), + ("days since 3000-01-01", [-357_000, -356_999]), + ("days since 2262-04-11", [-5, -1]), + ], + ) + def test_far_reference_date_matches_data_path( + self, calendar, time_unit, units, values + ) -> None: + _assert_metadata_matches_data( + lambda: _time_variable(np.array(values), units, calendar), + time_unit=time_unit, + ) + + @pytest.mark.parametrize( + ("values", "units", "use_cftime"), + [ + # the synthetic day 1 overflows datetime64[ns], the stored dates don't + ([-5, -1], "days since 2262-04-11", False), + # cftime rejects these units, but there is nothing to decode + pytest.param( + np.array([], dtype="int64"), + "nanoseconds since 2000-01-01", + True, + marks=requires_cftime, + ), + ], + ) + def test_synthetic_decode_failure_matches_data_path( + self, values, units, use_cftime + ) -> None: + _assert_metadata_matches_data( + lambda: _time_variable(np.array(values), units), use_cftime=use_cftime + ) + + @requires_cftime + @pytest.mark.filterwarnings("ignore:Ambiguous reference date string") + def test_far_reference_date_reads_data(self) -> None: + variable = _time_variable( + InaccessibleArray(np.zeros(3)), "hours since 1-1-1 00:00:0.0" + ) + with pytest.raises(UnexpectedDataAccess): + CFDatetimeCoder(infer_dtype="metadata").decode(variable) + + def test_fractional_values_with_fine_time_unit(self) -> None: + # fractions finer than the units need a fine enough time_unit + variable = _time_variable(np.array([0.0, 0.5, 1.0]), "seconds since 2000-01-01") + decoded = CFDatetimeCoder(time_unit="ms", infer_dtype="metadata").decode( + variable + ) + expected = CFDatetimeCoder(time_unit="ms").decode(variable) + assert decoded.dtype == np.dtype("datetime64[ms]") + assert_identical(decoded.load(), expected.load()) + + @pytest.mark.filterwarnings("ignore:Unable to decode time axis") + @pytest.mark.filterwarnings("ignore:Can't decode floating point datetimes") + @pytest.mark.parametrize( + "chunked", [False, pytest.param(True, marks=requires_dask)] + ) + @pytest.mark.parametrize( + ("values", "units", "coder_kwargs"), + [ + # dates beyond datetime64[ns] + pytest.param( + [0, 300_000], "days since 2000-01-01", {}, marks=requires_cftime + ), + # fractions finer than time_unit + ([0.0, 0.5], "seconds since 2000-01-01", {"time_unit": "s"}), + ], + ) + def test_wrong_inferred_dtype_raises_on_load( + self, chunked, values, units, coder_kwargs + ) -> None: + variable = _time_variable(np.array(values), units) + if chunked: + variable = variable.chunk(t=1) + coder = CFDatetimeCoder(infer_dtype="metadata", **coder_kwargs) + decoded = coder.decode(variable) + with pytest.raises(ValueError, match="as inferred from metadata"): + decoded.compute() + + @requires_dask + def test_dask_dtype_matches_computed(self) -> None: + # map_blocks casts every chunk to the declared dtype, so a wrong + # declaration would silently change the values + variable = _time_variable( + np.array([0, 1, 2]), "milliseconds since 2000-01-01" + ).chunk(t=1) + coder = CFDatetimeCoder(use_cftime=False, time_unit="s", infer_dtype="metadata") + decoded = coder.decode(variable) + expected = CFDatetimeCoder(use_cftime=False, time_unit="s").decode( + variable.compute() + ) + assert decoded.dtype == np.dtype("datetime64[ms]") + assert_identical(decoded.compute(), expected.load())