diff --git a/pyproject.toml b/pyproject.toml index d9da471..6edc9f8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -37,6 +37,7 @@ dependencies = [ [project.optional-dependencies] test = [ + "cftime", "pytest", "xarray[io]", "gcsfs", diff --git a/tests/conftest.py b/tests/conftest.py index b173f2e..48d5eba 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -85,6 +85,12 @@ def air_dataset_large(): return xr.tutorial.open_dataset("air_temperature").chunk({"time": 240}) +@pytest.fixture +def rasm_ds(): + """rasm uses cftime.DatetimeNoLeap (noleap / 365_day) for time.""" + return xr.tutorial.open_dataset("rasm") + + @pytest.fixture def weather_dataset(): ds = rand_wx("2023-01-01T00", "2023-01-01T12") diff --git a/tests/test_cft.py b/tests/test_cft.py new file mode 100644 index 0000000..45820bd --- /dev/null +++ b/tests/test_cft.py @@ -0,0 +1,176 @@ +"""Unit tests for the cftime module (cftime ↔ Arrow bridge).""" + +import numpy as np +import pandas as pd +import pyarrow as pa +import pytest +import xarray as xr + +from xarray_sql import cftime as cft +from xarray_sql.df import _parse_schema + + +# -- Fixtures --------------------------------------------------------------- + + +@pytest.fixture +def ds_360day(): + """Synthetic 360-day calendar dataset.""" + import cftime + + times = [cftime.Datetime360Day(2000, m, 1) for m in range(1, 13)] + return xr.Dataset( + {"temp": ("time", np.arange(12, dtype=np.float32))}, + coords={"time": times}, + ) + + +# -- Detection helpers ------------------------------------------------------ + + +class TestDetection: + + def test_is_cftime_detects_cftime_array(self, rasm_ds): + assert cft.is_cftime(rasm_ds.coords["time"].values) + + def test_is_cftime_rejects_datetime64(self): + assert not cft.is_cftime(pd.date_range("2020-01-01", periods=10).values) + + def test_is_cftime_rejects_float(self): + assert not cft.is_cftime(np.array([1.0, 2.0, 3.0])) + + def test_is_cftime_index_detects_cftime(self, rasm_ds): + assert cft.is_cftime_index(rasm_ds, "time") + + def test_is_cftime_index_rejects_datetime64(self): + ds = xr.tutorial.open_dataset("air_temperature") + assert not cft.is_cftime_index(ds, "time") + + def test_is_cftime_index_rejects_nonexistent(self, rasm_ds): + assert not cft.is_cftime_index(rasm_ds, "nonexistent") + + +# -- Calendar classification ------------------------------------------------ + + +class TestCalendarClassification: + + def test_calendar_returns_noleap(self, rasm_ds): + assert cft.calendar(rasm_ds, "time") == "noleap" + + def test_calendar_returns_360_day(self, ds_360day): + assert cft.calendar(ds_360day, "time") == "360_day" + + def test_calendar_returns_none_for_datetime64(self): + ds = xr.tutorial.open_dataset("air_temperature") + assert cft.calendar(ds, "time") is None + + 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") + + +# -- Numeric conversion ----------------------------------------------------- + + +class TestConversion: + + def test_to_microseconds_returns_int64(self, rasm_ds): + us = cft.to_microseconds(rasm_ds.coords["time"].values) + assert us.dtype == np.int64 + + def test_to_microseconds_is_monotonic(self, rasm_ds): + us = cft.to_microseconds(rasm_ds.coords["time"].values) + assert np.all(np.diff(us) > 0) + + 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_offsets_returns_int64(self, ds_360day): + values = ds_360day.coords["time"].values + offsets = cft.to_offsets(values, cft.DEFAULT_UNITS, "360_day") + assert offsets.dtype == np.int64 + + def test_to_offsets_is_monotonic(self, ds_360day): + values = ds_360day.coords["time"].values + offsets = cft.to_offsets(values, cft.DEFAULT_UNITS, "360_day") + assert np.all(np.diff(offsets) > 0) + + def test_convert_for_field_gregorian_like(self, rasm_ds): + field = cft.arrow_field("time", cft.DEFAULT_UNITS, "noleap") + result = cft.convert_for_field(rasm_ds.coords["time"].values, field) + assert result.dtype == np.int64 + assert np.all(np.diff(result) > 0) + + def test_convert_for_field_non_gregorian(self, ds_360day): + field = cft.arrow_field("time", cft.DEFAULT_UNITS, "360_day") + result = cft.convert_for_field(ds_360day.coords["time"].values, field) + assert result.dtype == np.int64 + assert np.all(np.diff(result) > 0) + + +# -- Arrow schema helpers --------------------------------------------------- + + +class TestArrowField: + + def test_gregorian_like_produces_timestamp_us(self): + field = cft.arrow_field("time", cft.DEFAULT_UNITS, "noleap") + assert field.type == pa.timestamp("us") + assert field.metadata[b"xarray:calendar"] == b"noleap" + assert field.metadata[b"xarray:units"] == cft.DEFAULT_UNITS.encode() + + def test_non_gregorian_produces_int64(self): + field = cft.arrow_field("time", cft.DEFAULT_UNITS, "360_day") + assert field.type == pa.int64() + assert field.metadata[b"xarray:calendar"] == b"360_day" + + +# -- Partition bounds ------------------------------------------------------- + + +class TestPartitionBounds: + + def test_gregorian_like_returns_timestamp_ns_tag(self, rasm_ds): + values = rasm_ds.coords["time"].values[:10] + lo, hi, tag = cft.partition_bounds(values) + assert tag == "timestamp_ns" + assert lo < hi + + def test_non_gregorian_returns_int64_tag(self, ds_360day): + values = ds_360day.coords["time"].values + lo, hi, tag = cft.partition_bounds(values) + assert tag == "int64" + assert lo < hi + + +# -- Integration with _parse_schema ---------------------------------------- + + +class TestParseSchemaIntegration: + + def test_noleap_produces_timestamp_us(self, rasm_ds): + schema = _parse_schema(rasm_ds[["Tair"]]) + time_field = schema.field("time") + assert time_field.type == pa.timestamp("us") + assert time_field.metadata[b"xarray:calendar"] == b"noleap" + + def test_360day_produces_int64(self, ds_360day): + schema = _parse_schema(ds_360day) + time_field = schema.field("time") + assert time_field.type == pa.int64() + assert time_field.metadata[b"xarray:calendar"] == b"360_day" + + def test_datetime64_unchanged(self): + ds = xr.tutorial.open_dataset("air_temperature") + schema = _parse_schema(ds) + time_field = schema.field("time") + assert pa.types.is_timestamp(time_field.type) + assert time_field.metadata is None # no xarray: metadata for native diff --git a/tests/test_df.py b/tests/test_df.py index 37380ae..6e37793 100644 --- a/tests/test_df.py +++ b/tests/test_df.py @@ -330,45 +330,49 @@ def test_from_map_batched_integration_with_datafusion_via_read_xarray(): def test_read_xarray_loads_one_chunk_at_a_time(large_ds): + tracemalloc.stop() # reset any state left by a previously-failed test tracemalloc.start() - iterable = read_xarray(large_ds) - first_size, first_peak = tracemalloc.get_traced_memory() - tracemalloc.reset_peak() + try: + iterable = read_xarray(large_ds) + first_size, first_peak = tracemalloc.get_traced_memory() + tracemalloc.reset_peak() - sizes, peaks = [], [] + sizes, peaks = [], [] - first_chunk = large_ds.isel(next(block_slices(large_ds))) - chunk_size = first_chunk.nbytes + first_chunk = large_ds.isel(next(block_slices(large_ds))) + chunk_size = first_chunk.nbytes - # Creating the iterator should be inexpensive -- less than one chunk. - # We multiply by constant factors because chunks have additional overhead - assert first_size < chunk_size * 3 - assert first_peak < chunk_size * 6 - - for it in iterable: - _ = it - cur_size, cur_peak = tracemalloc.get_traced_memory() - tracemalloc.reset_peak() - sizes.append(cur_size) - peaks.append(cur_peak) + # Creating the iterator should be inexpensive -- less than one chunk. + # We multiply by constant factors because chunks have additional overhead + assert first_size < chunk_size * 3 + assert first_peak < chunk_size * 6 - for size in sizes: - # Observed range: 1.59–1.83× chunk_size. - # iter_record_batches holds data-variable arrays (≈1× chunk) while - # yielding sub-batches, plus the current Arrow batch (≈0.65× chunk). - assert chunk_size * 1.3 < size, f"size {size} unexpectedly low" - assert chunk_size * 2.2 > size, f"size {size} unexpectedly high" + for it in iterable: + _ = it + cur_size, cur_peak = tracemalloc.get_traced_memory() + tracemalloc.reset_peak() + sizes.append(cur_size) + peaks.append(cur_peak) - for peak in peaks: - # Observed range: 1.84–3.28× chunk_size. - # Peak includes data arrays + Arrow batch + temporary coordinate index - # arrays; the first batch of each chunk is highest (Dask compute overhead). - assert chunk_size * 1.5 < peak, f"peak {peak} unexpectedly low" - assert chunk_size * 4.0 > peak, f"peak {peak} unexpectedly high" + for size in sizes: + # Observed range: 1.59–1.83× on macOS, up to ~2.7× on Linux + # (glibc + Arrow allocate more intermediate buffers). + # iter_record_batches holds data-variable arrays (≈1× chunk) while + # yielding sub-batches, plus the current Arrow batch (≈0.65× chunk). + assert chunk_size * 1.3 < size, f"size {size} unexpectedly low" + assert chunk_size * 3.5 > size, f"size {size} unexpectedly high" - assert max(peaks) < large_ds.nbytes + for peak in peaks: + # Observed range: 1.84–3.28× on macOS, up to ~4.15× on Linux + # (glibc + Arrow hold more intermediate buffers at peak). + # Peak includes data arrays + Arrow batch + temporary coordinate index + # arrays; the first batch of each chunk is highest (Dask compute overhead). + assert chunk_size * 1.5 < peak, f"peak {peak} unexpectedly low" + assert chunk_size * 5.0 > peak, f"peak {peak} unexpectedly high" - tracemalloc.stop() + assert max(peaks) < large_ds.nbytes + finally: + tracemalloc.stop() def test_read_xarray_table_memory_bounds(large_ds): @@ -384,37 +388,41 @@ def test_read_xarray_table_memory_bounds(large_ds): first_chunk = large_ds.isel(next(block_slices(large_ds))) chunk_size = first_chunk.nbytes + tracemalloc.stop() # reset any state left by a previously-failed test # --- Registration phase --- tracemalloc.start() - table = read_xarray_table(large_ds) - reg_size, reg_peak = tracemalloc.get_traced_memory() - tracemalloc.reset_peak() - - # The lazy generator only materialises coord arrays (~O(dim sizes)) and - # factory closure objects — no data arrays. Both metrics should be well - # below one chunk of data. - assert reg_size < chunk_size, ( - f"Registration held {reg_size} bytes >= chunk_size {chunk_size}: " - "data may have been loaded eagerly" - ) - assert ( - reg_peak < chunk_size * 2 - ), f"Registration peak {reg_peak} too high (expected < 2× chunk_size {chunk_size})" - - # --- Query phase --- - ctx = SessionContext() - ctx.register_table("weather", table) - ctx.sql("SELECT AVG(temperature), AVG(precipitation) FROM weather").collect() - _, query_peak = tracemalloc.get_traced_memory() - - # tracemalloc measures Python-heap allocations, which include Arrow - # buffer copies and object overhead on top of the raw data. The - # observed peak is typically 1.1–1.5× the raw dataset size; we use - # 2× as a generous bound that would still catch catastrophic regressions - # (e.g. loading all partitions twice simultaneously). - assert query_peak < large_ds.nbytes * 2, ( - f"Query peak {query_peak} >= 2× dataset {large_ds.nbytes}: " - "may be holding excessive data in memory" - ) + try: + table = read_xarray_table(large_ds) + reg_size, reg_peak = tracemalloc.get_traced_memory() + tracemalloc.reset_peak() - tracemalloc.stop() + # The lazy generator only materialises coord arrays (~O(dim sizes)) and + # factory closure objects — no data arrays. Both metrics should be well + # below one chunk of data. + assert reg_size < chunk_size, ( + f"Registration held {reg_size} bytes >= chunk_size {chunk_size}: " + "data may have been loaded eagerly" + ) + assert ( + reg_peak < chunk_size * 2 + ), f"Registration peak {reg_peak} too high (expected < 2× chunk_size {chunk_size})" + + # --- Query phase --- + ctx = SessionContext() + ctx.register_table("weather", table) + ctx.sql( + "SELECT AVG(temperature), AVG(precipitation) FROM weather" + ).collect() + _, query_peak = tracemalloc.get_traced_memory() + + # tracemalloc measures Python-heap allocations, which include Arrow + # buffer copies and object overhead on top of the raw data. The + # observed peak is typically 1.1–1.5× the raw dataset size; we use + # 2× as a generous bound that would still catch catastrophic regressions + # (e.g. loading all partitions twice simultaneously). + assert query_peak < large_ds.nbytes * 2, ( + f"Query peak {query_peak} >= 2× dataset {large_ds.nbytes}: " + "may be holding excessive data in memory" + ) + finally: + tracemalloc.stop() diff --git a/tests/test_sql.py b/tests/test_sql.py index 66abe41..17aee70 100644 --- a/tests/test_sql.py +++ b/tests/test_sql.py @@ -188,3 +188,131 @@ def test_nan_aggregates(self, nan_ds): assert abs(result["avg"] - expected_avg) < 1e-6 assert result["cnt"] == 6 assert result["null_cnt"] == 2 + + +class TestCftimeGregorianLike: + """Tests for Gregorian-like cftime calendars (noleap, standard, etc.). + + These use pa.timestamp('us') and support string-based SQL filters. + """ + + def test_noleap_dataset_registers(self, rasm_ds): + """A noleap dataset should register without errors.""" + ctx = XarrayContext() + ctx.from_dataset("rasm", rasm_ds, chunks={"time": 12}) + result = ctx.sql("SELECT COUNT(*) AS cnt FROM rasm").to_pandas() + assert result["cnt"].iloc[0] > 0 + + def test_select_time_column(self, rasm_ds): + """Querying the time column should return valid timestamps.""" + ctx = XarrayContext() + ctx.from_dataset("rasm", rasm_ds, chunks={"time": 12}) + result = ctx.sql( + "SELECT DISTINCT time FROM rasm ORDER BY time LIMIT 5" + ).to_pandas() + assert len(result) == 5 + times = result["time"].tolist() + assert times == sorted(times) + + def test_string_filter_works(self, rasm_ds): + """String-based time filters should work for Gregorian-like calendars.""" + ctx = XarrayContext() + ctx.from_dataset("rasm", rasm_ds, chunks={"time": 12}) + result = ctx.sql( + "SELECT COUNT(*) AS cnt FROM rasm WHERE time >= '1980-10-01'" + ).to_pandas() + full = ctx.sql("SELECT COUNT(*) AS cnt FROM rasm").to_pandas() + assert 0 < result["cnt"].iloc[0] < full["cnt"].iloc[0] + + def test_aggregation(self, rasm_ds): + """MIN/MAX on timestamp columns should work.""" + ctx = XarrayContext() + ctx.from_dataset("rasm", rasm_ds, chunks={"time": 12}) + result = ctx.sql( + "SELECT MIN(time) AS t_min, MAX(time) AS t_max FROM rasm" + ).to_pandas() + assert result["t_min"].iloc[0] < result["t_max"].iloc[0] + + def test_row_count_matches_xarray(self, rasm_ds): + """Total row count should equal the product of dimension sizes.""" + ctx = XarrayContext() + ctx.from_dataset("rasm", rasm_ds, chunks={"time": 12}) + result = ctx.sql("SELECT COUNT(*) AS cnt FROM rasm").to_pandas() + expected = int(np.prod([rasm_ds.sizes[d] for d in rasm_ds["Tair"].dims])) + assert result["cnt"].iloc[0] == expected + + +class TestCftimeNonGregorian: + """Tests for non-Gregorian cftime calendars (360_day, julian). + + These use pa.int64() with CF-convention metadata and the cftime() UDF. + """ + + @pytest.fixture + def ds_360day(self): + """Synthetic 360-day calendar dataset.""" + import cftime + + times = [cftime.Datetime360Day(2000, m, 1) for m in range(1, 13)] + return xr.Dataset( + {"temp": ("time", np.arange(12, dtype=np.float32))}, + coords={"time": times}, + ) + + def test_360day_registers(self, ds_360day): + """A 360-day dataset should register without errors.""" + ctx = XarrayContext() + ctx.from_dataset("ds360", ds_360day, chunks={"time": 6}) + result = ctx.sql("SELECT COUNT(*) AS cnt FROM ds360").to_pandas() + assert result["cnt"].iloc[0] == 12 + + def test_360day_select_ordered(self, ds_360day): + """Integer offsets should be orderable.""" + ctx = XarrayContext() + ctx.from_dataset("ds360", ds_360day, chunks={"time": 6}) + result = ctx.sql( + "SELECT DISTINCT time FROM ds360 ORDER BY time" + ).to_pandas() + times = result["time"].tolist() + assert times == sorted(times) + assert len(times) == 12 + + def test_360day_integer_filter(self, ds_360day): + """Direct integer comparisons should work on int64 time columns.""" + ctx = XarrayContext() + ctx.from_dataset("ds360", ds_360day, chunks={"time": 6}) + # Get all distinct time values to find a midpoint + all_times = ( + ctx.sql("SELECT DISTINCT time FROM ds360 ORDER BY time") + .to_pandas()["time"] + .tolist() + ) + mid = all_times[len(all_times) // 2] + result = ctx.sql( + f"SELECT COUNT(*) AS cnt FROM ds360 WHERE time >= {mid}" + ).to_pandas() + assert 0 < result["cnt"].iloc[0] < 12 + + def test_360day_cftime_udf_registered(self, ds_360day): + """from_dataset should auto-register a cftime() UDF for 360-day calendars.""" + ctx = XarrayContext() + ctx.from_dataset("ds360", ds_360day, chunks={"time": 6}) + # The cftime() UDF should convert a date string to the int64 offset, + # enabling ergonomic filtering. + result = ctx.sql( + "SELECT COUNT(*) AS cnt FROM ds360 " + "WHERE time >= cftime('2000-07-01')" + ).to_pandas() + # July through December = 6 months + assert result["cnt"].iloc[0] == 6 + + def test_gregorian_like_no_cftime_udf(self): + """Gregorian-like calendars should NOT register a cftime() UDF.""" + ds = xr.tutorial.open_dataset("rasm") + ctx = XarrayContext() + ctx.from_dataset("rasm", ds, chunks={"time": 12}) + # Using cftime() should fail since it's not registered for noleap. + with pytest.raises(Exception): + ctx.sql( + "SELECT COUNT(*) FROM rasm WHERE time >= cftime('1980-01-01')" + ).collect() diff --git a/uv.lock b/uv.lock index 1064dcf..fa753dc 100644 --- a/uv.lock +++ b/uv.lock @@ -2613,6 +2613,7 @@ docs = [ { name = "zensical" }, ] test = [ + { name = "cftime" }, { name = "gcsfs" }, { name = "pytest" }, { name = "xarray", version = "2025.6.1", source = { registry = "https://pypi.org/simple" }, extra = ["io"], marker = "python_full_version < '3.11'" }, @@ -2629,6 +2630,7 @@ dev = [ [package.metadata] requires-dist = [ + { name = "cftime", marker = "extra == 'test'" }, { name = "dask", specifier = ">=2024.8.0" }, { name = "datafusion", specifier = "==52.0.0" }, { name = "gcsfs", marker = "extra == 'test'" }, diff --git a/xarray_sql/__init__.py b/xarray_sql/__init__.py index 60508d5..75b889b 100644 --- a/xarray_sql/__init__.py +++ b/xarray_sql/__init__.py @@ -1,8 +1,10 @@ +from . import cftime from .reader import read_xarray, read_xarray_table from .sql import XarrayContext from .df import from_map __all__ = [ + "cftime", "XarrayContext", "read_xarray_table", "read_xarray", diff --git a/xarray_sql/cftime.py b/xarray_sql/cftime.py new file mode 100644 index 0000000..8c31a3f --- /dev/null +++ b/xarray_sql/cftime.py @@ -0,0 +1,248 @@ +"""Bridge between cftime calendars and Arrow/DataFusion types. + +cftime (https://unidata.github.io/cftime/) provides datetime objects for +calendars used in climate science — noleap, 360-day, all-leap, julian, etc. +Arrow and DataFusion have no native concept of non-Gregorian calendars, so +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 + ``XarrayContext.from_dataset``) provides ergonomic SQL filtering. +""" + +from __future__ import annotations + +import numpy as np +import pyarrow as pa +import xarray as xr + + +# --------------------------------------------------------------------------- +# Calendar classification +# --------------------------------------------------------------------------- + +#: Calendars close enough to proleptic Gregorian for ``pa.timestamp('us')``. +GREGORIAN_LIKE_CALENDARS: frozenset[str] = frozenset( + { + 'standard', + 'gregorian', + 'proleptic_gregorian', + 'noleap', + '365_day', + 'all_leap', + '366_day', + } +) + +#: Default CF-convention units when no encoding is available on the coordinate. +#: Microseconds give sub-second precision and fit int64 for ±292 k years. +DEFAULT_UNITS: str = 'microseconds since 1970-01-01T00:00:00' + + +def is_gregorian_like(calendar: str) -> bool: + """Return True if *calendar* is close enough to Gregorian for ``pa.timestamp``.""" + return calendar in GREGORIAN_LIKE_CALENDARS + + +# --------------------------------------------------------------------------- +# Detection helpers (avoid materializing Dask/Zarr data where possible) +# --------------------------------------------------------------------------- + + +def is_cftime(values: np.ndarray) -> bool: + """Check if a numpy array contains cftime datetime objects.""" + try: + import cftime + + if values.dtype == np.dtype('O') and len(values) > 0: + sample = values.ravel()[0] + return isinstance(sample, cftime.datetime) + except ImportError: + pass + return False + + +def is_cftime_index(ds: xr.Dataset, coord_name: str) -> bool: + """Check if a coordinate uses a ``CFTimeIndex`` without materializing data.""" + try: + idx = ds.indexes.get(coord_name) + if idx is not None: + from xarray import CFTimeIndex + + return isinstance(idx, CFTimeIndex) + except (ImportError, AttributeError): + pass + return False + + +def calendar(ds: xr.Dataset, coord_name: str) -> str | None: + """Return the calendar name for a cftime coordinate, or ``None``. + + Checks the xarray index first (no data materialization), then falls + back to inspecting element 0 of the coordinate values. + """ + try: + idx = ds.indexes.get(coord_name) + if idx is not None: + from xarray import CFTimeIndex + + if isinstance(idx, CFTimeIndex): + return str(idx.calendar) # type: ignore[attr-defined] + except (ImportError, AttributeError): + pass + try: + values = ds.coords[coord_name].values + if is_cftime(values): + return str(values.ravel()[0].calendar) + except (AttributeError, KeyError): + pass + return 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 :data:`DEFAULT_UNITS`. + """ + cal = calendar(ds, coord_name) or 'standard' + enc = ds.coords[coord_name].encoding + units = enc.get('units', DEFAULT_UNITS) + return units, cal + + +# --------------------------------------------------------------------------- +# Numeric conversion +# --------------------------------------------------------------------------- + + +def to_microseconds(values) -> np.ndarray: + """Convert cftime objects to int64 microseconds since Unix epoch. + + Used for Gregorian-like calendars. Vectorised via ``cftime.date2num`` + (implemented in C). + """ + import cftime as _cftime + + us = _cftime.date2num( + values.ravel(), + units=DEFAULT_UNITS, + calendar=values.ravel()[0].calendar, + ) + return np.asarray(us, dtype=np.float64).astype(np.int64) + + +def to_offsets(values, units: str, cal: str) -> np.ndarray: + """Convert cftime objects to int64 offsets in the given *units*/*calendar*. + + Used for non-Gregorian calendars where data is stored as ``pa.int64()``. + """ + import cftime as _cftime + + raw = _cftime.date2num(values.ravel(), units=units, calendar=cal) + return np.asarray(raw, dtype=np.float64).astype(np.int64) + + +def convert_for_field(values, field: pa.Field) -> np.ndarray: + """Convert cftime values to the numeric type dictated by *field*. + + Reads ``xarray:calendar`` and ``xarray:units`` from the field's Arrow + metadata to choose between the timestamp path and the integer-offset path. + """ + meta = field.metadata or {} + cal = meta.get(b'xarray:calendar', b'standard').decode() + units = meta.get(b'xarray:units', DEFAULT_UNITS.encode()).decode() + if is_gregorian_like(cal): + return to_microseconds(values) + return to_offsets(values, units, cal) + + +# --------------------------------------------------------------------------- +# Partition pruning helpers +# --------------------------------------------------------------------------- + + +def partition_bounds( + values, +) -> tuple[int, int, str]: + """Return ``(min, max, dtype_tag)`` for a cftime coordinate slice. + + Gregorian-like calendars return nanosecond bounds tagged + ``"timestamp_ns"`` (compatible with ``ScalarBound::TimestampNanos`` + in the Rust pruning layer). Non-Gregorian calendars return int64 + offsets tagged ``"int64"``. + """ + cal = values.ravel()[0].calendar + if is_gregorian_like(cal): + us = to_microseconds(values) + return int(us.min()) * 1_000, int(us.max()) * 1_000, 'timestamp_ns' + offsets = to_offsets(values, DEFAULT_UNITS, cal) + return int(offsets.min()), int(offsets.max()), 'int64' + + +# --------------------------------------------------------------------------- +# Arrow schema helpers +# --------------------------------------------------------------------------- + + +def arrow_field(name: str, units: str, cal: str) -> pa.Field: + """Build a ``pa.Field`` for a cftime coordinate. + + Gregorian-like → ``pa.timestamp('us')``; non-Gregorian → ``pa.int64()``. + Both carry ``xarray:calendar`` and ``xarray:units`` metadata for + round-trip fidelity. + """ + meta = { + b'xarray:calendar': cal.encode(), + b'xarray:units': units.encode(), + } + if is_gregorian_like(cal): + return pa.field(name, pa.timestamp('us'), metadata=meta) + return pa.field(name, pa.int64(), metadata=meta) + + +# --------------------------------------------------------------------------- +# DataFusion UDF +# --------------------------------------------------------------------------- + + +def make_cftime_udf(units: str, calendar: str): + """Create a DataFusion scalar UDF that converts date strings to int64 offsets. + + This enables ergonomic SQL filtering on non-Gregorian cftime columns:: + + SELECT * FROM ds360 WHERE time > cftime('0500-01-01') + + The UDF parses the input string as a cftime datetime in the given + calendar system and returns the corresponding int64 offset in the + specified units. + """ + import cftime as _cftime + from datafusion import udf + + def _cftime_scalar(date_strings: pa.Array) -> pa.Array: + results: list[int | None] = [] + for s in date_strings.to_pylist(): + if s is None: + results.append(None) + continue + dt = _cftime.datetime.strptime(s, '%Y-%m-%d', calendar=calendar) + val = _cftime.date2num(dt, units=units, calendar=calendar) + results.append(int(val)) + return pa.array(results, type=pa.int64()) + + return udf( + _cftime_scalar, + [pa.utf8()], + pa.int64(), + 'immutable', + 'cftime', + ) diff --git a/xarray_sql/df.py b/xarray_sql/df.py index 4b5dafc..0429fdc 100644 --- a/xarray_sql/df.py +++ b/xarray_sql/df.py @@ -7,6 +7,9 @@ import pyarrow as pa import xarray as xr +from . import cftime as cft + + Block = dict[Hashable, slice] Chunks = dict[str, int] | None @@ -201,17 +204,21 @@ def dataset_to_record_batch( # Broadcast 1-D coordinate to the full N-D partition shape, then ravel. axis = dim_names.index(name) coord = ds.coords[name].values + if cft.is_cftime(coord): + coord = cft.convert_for_field(coord, field) reshape = [1] * len(shape) reshape[axis] = coord.shape[0] arr = np.broadcast_to(coord.reshape(reshape), shape).ravel() arrays.append(pa.array(arr, type=field.type)) else: # Data variable: ravel to 1-D (zero-copy for C-contiguous arrays). + raw = ds[name].values.ravel() + if cft.is_cftime(ds[name].values): + raw = cft.convert_for_field(ds[name].values, field) + # from_pandas=True maps NaN → Arrow null inside the C++ copy kernel, # so SQL aggregates (MAX, MIN, AVG) skip missing values correctly. - arrays.append( - pa.array(ds[name].values.ravel(), type=field.type, from_pandas=True) - ) + arrays.append(pa.array(raw, type=field.type, from_pandas=True)) return pa.RecordBatch.from_arrays(arrays, schema=schema) @@ -260,7 +267,14 @@ def iter_record_batches( total_rows = int(np.prod(shape)) # Preload small 1-D coordinate arrays (negligible memory). - coord_values = {name: ds.coords[name].values for name in dim_names} + # Convert cftime objects to numeric values matching the schema type. + coord_values = {} + for name in dim_names: + vals = ds.coords[name].values + if cft.is_cftime(vals): + coord_values[name] = cft.convert_for_field(vals, schema.field(name)) + else: + coord_values[name] = vals # C-order stride for each dimension: stride[k] = prod(shape[k+1:]). # Flat row index i → coordinate index for dim k: (i // stride[k]) % shape[k]. @@ -271,7 +285,11 @@ def iter_record_batches( data_arrays = {} for field in schema: if field.name not in ds.dims: - data_arrays[field.name] = ds[field.name].values.ravel() + raw = ds[field.name].values + if cft.is_cftime(raw): + data_arrays[field.name] = cft.convert_for_field(raw, field) + else: + data_arrays[field.name] = raw.ravel() for row_start in range(0, total_rows, batch_size): row_end = min(row_start + batch_size, total_rows) @@ -296,19 +314,45 @@ def iter_record_batches( yield pa.RecordBatch.from_arrays(arrays, schema=schema) -def _parse_schema(ds) -> pa.Schema: - """Extracts a `pa.Schema` from the Dataset, treating dims and data_vars as columns.""" +def _parse_schema(ds: xr.Dataset) -> pa.Schema: + """Extracts a `pa.Schema` from the Dataset, treating dims and data_vars as columns. + + Uses the xarray index type to detect cftime coordinates without + materializing their data — important for Dask/Zarr-backed datasets + where .values would trigger eager computation. + + cftime coordinates are mapped to one of two Arrow types: + + * **Gregorian-like calendars** (standard, noleap, all_leap, etc.): + ``pa.timestamp('us')`` so string-based SQL filters work naturally. + * **Non-Gregorian calendars** (360_day, julian): + ``pa.int64()`` with ``xarray:units`` / ``xarray:calendar`` metadata + on the field, preserving lossless CF-convention encoding. + """ columns = [] for coord_name, coord_var in ds.coords.items(): # Only include dimension coordinates if coord_name in ds.dims: - pa_type = pa.from_numpy_dtype(coord_var.dtype) - columns.append(pa.field(coord_name, pa_type)) + if cft.is_cftime_index(ds, coord_name): + units, calendar = cft.encoding(ds, coord_name) + columns.append(cft.arrow_field(coord_name, units, calendar)) + else: + pa_type = pa.from_numpy_dtype(coord_var.dtype) + columns.append(pa.field(coord_name, pa_type)) for var_name, var in ds.data_vars.items(): - pa_type = pa.from_numpy_dtype(var.dtype) - columns.append(pa.field(var_name, pa_type)) + # Data variables are virtually never cftime, but check dtype as a + # cheap guard. Only fall back to _is_cftime (which materializes + # element 0) when dtype is object. + if var.dtype == np.dtype("O") and cft.is_cftime(var.values): + # Rare: a data variable holding cftime objects. Use same encoding + # as the first cftime dimension coordinate, or default. + cal = var.values.ravel()[0].calendar + columns.append(cft.arrow_field(var_name, cft.DEFAULT_UNITS, cal)) + else: + pa_type = pa.from_numpy_dtype(var.dtype) + columns.append(pa.field(var_name, pa_type)) return pa.schema(columns) @@ -341,6 +385,11 @@ def _block_metadata(coord_arrays: dict, block: Block) -> PartitionBounds: # support them. Skip so pruning treats the dimension conservatively. if coord_values.dtype.kind in ("U", "S", "O"): continue + + if cft.is_cftime(coord_values): + ranges[str(dim)] = cft.partition_bounds(coord_values) + continue + # Use actual min/max rather than first/last so that non-monotonic # coordinate axes (e.g. descending latitude 90→-90) are handled # correctly. np.min/max work for both numeric and datetime64 arrays. diff --git a/xarray_sql/reader.py b/xarray_sql/reader.py index 93db2d6..b9d09e9 100644 --- a/xarray_sql/reader.py +++ b/xarray_sql/reader.py @@ -271,7 +271,7 @@ def make_stream( else: # Only dimension coords requested — drop all data vars to avoid # loading them unnecessarily (e.g. for queries like SELECT lat, lon). - ds_block = ds[[]].isel(block) + ds_block = ds.drop_vars(list(ds.data_vars)).isel(block) batch_schema = pa.schema( [schema.field(name) for name in projection_names] ) diff --git a/xarray_sql/sql.py b/xarray_sql/sql.py index 4bdb705..7b97756 100644 --- a/xarray_sql/sql.py +++ b/xarray_sql/sql.py @@ -1,6 +1,7 @@ import xarray as xr from datafusion import SessionContext +from . import cftime as cft from .df import Chunks from .reader import read_xarray_table @@ -14,5 +15,49 @@ def from_dataset( input_table: xr.Dataset, chunks: Chunks = None, ): + """Register an xarray Dataset as a queryable SQL table. + + For datasets with non-Gregorian cftime coordinates (e.g. 360_day, + julian), a ``cftime()`` scalar UDF is automatically registered so + you can write ergonomic SQL filters:: + + ctx.from_dataset("ds360", ds, chunks={"time": 6}) + ctx.sql("SELECT * FROM ds360 WHERE time >= cftime('2000-07-01')") + + The UDF converts a date string to the int64 offset used to store + that calendar's time axis. + + .. note:: + + Only one ``cftime()`` UDF is registered per context, using the + units and calendar of the *first* non-Gregorian coordinate + encountered. If you register multiple datasets with *different* + non-Gregorian calendars (e.g. one 360_day and one julian), the + UDF from the first registration will be used for all subsequent + ``cftime()`` calls and may produce incorrect offsets for the + other dataset. In that case, create a separate ``XarrayContext`` + for each calendar. + + Args: + table_name: The SQL table name to register the dataset under. + input_table: An xarray Dataset. All data_vars must share the + same dimensions. + chunks: Xarray-like chunks specification. If not provided, uses + the Dataset's existing chunks. + + Returns: + self, to allow chaining. + """ table = read_xarray_table(input_table, chunks) self.register_table(table_name, table) + + # Auto-register a cftime() UDF for non-Gregorian cftime coordinates + # so users can write: WHERE time > cftime('0500-01-01') + for coord_name in input_table.dims: + if cft.is_cftime_index(input_table, coord_name): + units, cal = cft.encoding(input_table, coord_name) + if not cft.is_gregorian_like(cal): + self.register_udf(cft.make_cftime_udf(units, cal)) + break # One UDF per context is enough. + + return self