diff --git a/tests/test_cft.py b/tests/test_cft.py index 24b362bb..87a2413e 100644 --- a/tests/test_cft.py +++ b/tests/test_cft.py @@ -67,12 +67,16 @@ def test_noleap_is_gregorian_like(self): assert cft.is_gregorian_like("noleap") assert cft.is_gregorian_like("standard") assert cft.is_gregorian_like("proleptic_gregorian") - assert cft.is_gregorian_like("all_leap") def test_360_day_is_not_gregorian_like(self): assert not cft.is_gregorian_like("360_day") assert not cft.is_gregorian_like("julian") + def test_all_leap_is_not_gregorian_like(self): + # Feb 29 of every year has no Gregorian timestamp + assert not cft.is_gregorian_like("all_leap") + assert not cft.is_gregorian_like("366_day") + # -- Numeric conversion ----------------------------------------------------- @@ -90,6 +94,12 @@ def test_to_microseconds_length_matches(self, rasm_ds): values = rasm_ds.coords["time"].values assert len(cft.to_microseconds(values)) == len(values) + def test_to_microseconds_is_exact_past_float64_range(self): + import cftime + + values = np.array([cftime.DatetimeNoLeap(2300, 1, 1, 0, 0, 0, 1)]) + assert cft.to_microseconds(values)[0] % 2 == 1 + def test_to_offsets_returns_int64(self, ds_360day): values = ds_360day.coords["time"].values offsets = cft.to_offsets(values, cft.DEFAULT_UNITS, "360_day") diff --git a/tests/test_sql.py b/tests/test_sql.py index 56645401..b65288d1 100644 --- a/tests/test_sql.py +++ b/tests/test_sql.py @@ -319,6 +319,64 @@ def test_row_count_matches_xarray(self, rasm_ds): ) assert result["cnt"].iloc[0] == expected + def test_noleap_dates_keep_their_fields(self): + """A noleap date reads as the same calendar date, not time elapsed.""" + times = xr.date_range( + "2050-01-01", + periods=3, + freq="D", + calendar="noleap", + use_cftime=True, + ) + ds = xr.Dataset( + {"tas": ("time", [1.0, 2.0, 3.0])}, coords={"time": times} + ) + ctx = XarrayContext() + ctx.from_dataset("tas", ds) + result = ctx.sql( + "SELECT time, EXTRACT(MONTH FROM time) AS month FROM tas ORDER BY time" + ).to_pandas() + assert result["time"].tolist() == list( + pd.date_range("2050-01-01", periods=3, freq="D") + ) + assert result["month"].tolist() == [1, 1, 1] + + def test_noleap_filter_matches_calendar_dates(self): + times = xr.date_range( + "2050-01-01", + periods=60, + freq="D", + calendar="noleap", + use_cftime=True, + ) + ds = xr.Dataset( + {"tas": ("time", np.arange(60.0))}, coords={"time": times} + ) + ctx = XarrayContext() + ctx.from_dataset("tas", ds, chunks={"time": 10}) + result = ctx.sql( + "SELECT COUNT(*) AS cnt FROM tas WHERE time >= '2050-02-01'" + ).to_pandas() + assert result["cnt"].iloc[0] == 29 + + def test_noleap_dates_across_gregorian_leap_day(self): + times = xr.date_range( + "2052-02-28", + periods=3, + freq="D", + calendar="noleap", + use_cftime=True, + ) + ds = xr.Dataset( + {"tas": ("time", [1.0, 2.0, 3.0])}, coords={"time": times} + ) + ctx = XarrayContext() + ctx.from_dataset("tas", ds) + result = ctx.sql("SELECT time FROM tas ORDER BY time").to_pandas() + assert result["time"].tolist() == list( + pd.to_datetime(["2052-02-28", "2052-03-01", "2052-03-02"]) + ) + class TestCftimeNonGregorian: """Tests for non-Gregorian cftime calendars (360_day, julian). @@ -371,6 +429,47 @@ def test_360day_integer_filter(self, ds_360day): ).to_pandas() assert 0 < result["cnt"].iloc[0] < 12 + def test_all_leap_sub_day_times_survive_day_units(self): + times = xr.date_range( + "2001-02-28T12:00", + periods=4, + freq="12h", + calendar="all_leap", + use_cftime=True, + ) + ds = xr.Dataset( + {"tas": ("time", np.arange(4.0))}, coords={"time": times} + ) + ds.time.encoding["units"] = "days since 2001-01-01" + ctx = XarrayContext() + ctx.from_dataset("tas", ds) + result = ctx.sql( + "SELECT COUNT(DISTINCT time) AS n, " + "COUNT(*) FILTER (WHERE time >= cftime('2001-02-29')) AS after " + "FROM tas" + ).to_pandas() + assert result["n"].iloc[0] == 4 + assert result["after"].iloc[0] == 3 + + def test_pruning_matches_encoded_day_units(self): + times = xr.date_range( + "2001-01-01", + periods=120, + freq="D", + calendar="360_day", + use_cftime=True, + ) + ds = xr.Dataset( + {"tas": ("time", np.arange(120.0))}, coords={"time": times} + ).chunk(time=30) + ds.time.encoding["units"] = "days since 2001-01-01" + ctx = XarrayContext() + ctx.from_dataset("tas", ds, chunks={"time": 30}) + result = ctx.sql( + "SELECT COUNT(*) AS cnt FROM tas WHERE time < cftime('2001-03-01')" + ).to_pandas() + assert result["cnt"].iloc[0] == 60 + def test_360day_cftime_udf_registered(self, ds_360day): """from_dataset should auto-register a cftime() UDF for 360-day calendars.""" ctx = XarrayContext() diff --git a/xarray_sql/cftime.py b/xarray_sql/cftime.py index 96591303..b5f5698b 100644 --- a/xarray_sql/cftime.py +++ b/xarray_sql/cftime.py @@ -6,16 +6,19 @@ this module handles the conversion in two tiers: * **Gregorian-like calendars** (standard, gregorian, proleptic_gregorian, - noleap/365_day, all_leap/366_day): mapped to ``pa.timestamp('us')`` so - that string-based SQL filters like ``WHERE time > '1980-01-01'`` work - naturally. Microsecond resolution avoids the 1678–2262 overflow of - nanoseconds while preserving sub-second precision. - -* **Non-Gregorian calendars** (360_day, julian): mapped to ``pa.int64()`` - with ``xarray:units`` and ``xarray:calendar`` metadata on the Arrow field. - This preserves the original CF-convention encoding losslessly. A - ``cftime()`` DataFusion UDF (registered automatically by - ``Context.from_dataset``) provides ergonomic SQL filtering. + noleap/365_day): mapped to ``pa.timestamp('us')`` so that string-based + SQL filters like ``WHERE time > '1980-01-01'`` work naturally. Every + noleap date is also a Gregorian date, so noleap dates map by their + fields: 2050-01-01 stays 2050-01-01. Microsecond resolution avoids the + 1678–2262 overflow of nanoseconds while preserving sub-second precision. + +* **Non-Gregorian calendars** (360_day, julian, all_leap/366_day): mapped + to ``pa.int64()`` with ``xarray:units`` and ``xarray:calendar`` metadata + on the Arrow field. all_leap has a Feb 29 every year, which no Gregorian + timestamp can hold. Offsets are microseconds in the coordinate's own + calendar, which is lossless. A ``cftime()`` DataFusion UDF (registered + automatically by ``Context.from_dataset``) provides ergonomic SQL + filtering. """ from __future__ import annotations @@ -38,8 +41,6 @@ "proleptic_gregorian", "noleap", "365_day", - "all_leap", - "366_day", } ) """Calendars close enough to proleptic Gregorian for ``pa.timestamp('us')``.""" @@ -114,12 +115,16 @@ def calendar(ds: xr.Dataset, coord_name: str) -> str | None: def encoding(ds: xr.Dataset, coord_name: str) -> tuple[str, str]: """Return ``(units, calendar)`` for a cftime coordinate. - Reads xarray ``.encoding`` metadata (from the originating NetCDF file) - first, falling back to [DEFAULT_UNITS][xarray_sql.cftime.DEFAULT_UNITS]. + Gregorian-like calendars keep the units from xarray ``.encoding`` + (the originating NetCDF file) as metadata. Non-Gregorian calendars + always use [DEFAULT_UNITS][xarray_sql.cftime.DEFAULT_UNITS]: their + int64 offsets must share units with the partition bounds, and coarser + encoded units such as days would truncate times like noon. """ cal = calendar(ds, coord_name) or "standard" - enc = ds.coords[coord_name].encoding - units = enc.get("units", DEFAULT_UNITS) + units = DEFAULT_UNITS + if is_gregorian_like(cal): + units = ds.coords[coord_name].encoding.get("units", DEFAULT_UNITS) return units, cal @@ -136,12 +141,14 @@ def to_microseconds(values) -> np.ndarray: """ import cftime as _cftime - us = _cftime.date2num( - values.ravel(), - units=DEFAULT_UNITS, - calendar=values.ravel()[0].calendar, - ) - return cast(np.ndarray, np.asarray(us, dtype=np.float64).astype(np.int64)) + cal = values.ravel()[0].calendar + # Map noleap dates by their fields, not by the time elapsed in a calendar + # with no leap days, which would read 2050-01-01 as 2049-12-12. + if cal in ("noleap", "365_day"): + cal = "proleptic_gregorian" + us = _cftime.date2num(values.ravel(), units=DEFAULT_UNITS, calendar=cal) + # No float64 detour: it rounds microsecond offsets past 2^53 (~2255). + return cast(np.ndarray, np.asarray(us).astype(np.int64)) def to_offsets(values, units: str, cal: str) -> np.ndarray: diff --git a/xarray_sql/df.py b/xarray_sql/df.py index d26da7f0..4c74db56 100644 --- a/xarray_sql/df.py +++ b/xarray_sql/df.py @@ -630,9 +630,9 @@ def _parse_schema(ds: xr.Dataset) -> pa.Schema: cftime coordinates are mapped to one of two Arrow types: - * **Gregorian-like calendars** (standard, noleap, all_leap, etc.): + * **Gregorian-like calendars** (standard, noleap, etc.): ``pa.timestamp('us')`` so string-based SQL filters work naturally. - * **Non-Gregorian calendars** (360_day, julian): + * **Non-Gregorian calendars** (360_day, julian, all_leap): ``pa.int64()`` with ``xarray:units`` / ``xarray:calendar`` metadata on the field, preserving lossless CF-convention encoding. """