diff --git a/tests/test_cft.py b/tests/test_cft.py index 6560d72..24b362b 100644 --- a/tests/test_cft.py +++ b/tests/test_cft.py @@ -145,6 +145,14 @@ def test_non_gregorian_returns_int64_tag(self, ds_360day): assert tag == "int64" assert lo < hi + def test_out_of_int64_range_returns_none(self): + # Year-1 gregorian dates exceed the int64 nanosecond range, so no + # pruning bound can be reported; the caller skips the dimension. + values = xr.date_range( + "0001-01-01", periods=3, freq="100YS", use_cftime=True + ).values + assert cft.partition_bounds(values) is None + # -- Integration with _parse_schema ---------------------------------------- diff --git a/tests/test_df.py b/tests/test_df.py index 2c81144..5ba5b3f 100644 --- a/tests/test_df.py +++ b/tests/test_df.py @@ -8,6 +8,7 @@ from xarray_sql.df import ( DEFAULT_BATCH_SIZE, + _ensure_default_indexes, _parse_schema, block_slices, compute_chunks, @@ -16,6 +17,7 @@ from_map, from_map_batched, iter_record_batches, + partition_metadata, pivot, ) from xarray_sql.reader import read_xarray, read_xarray_table @@ -532,3 +534,138 @@ def test_compute_chunks_tuples_sum_to_dim_size(): result = compute_chunks(ds, {"a": 3, "b": 4, "c": 5}) for dim, tup in result.items(): assert sum(tup) == ds.sizes[dim] + + +# -- Object-dtype and out-of-ns-range coordinate support -------------------- + + +def _field_type(schema, name): + return schema.field(name).type + + +def test_parse_schema_maps_object_string_data_var_to_string(): + # A string variable arrives as numpy object dtype; _parse_schema must not + # hand it to pa.from_numpy_dtype (which raises "Unsupported numpy type 17"). + ds = xr.Dataset( + {"label": (["x"], np.array(["a", "b"], dtype=object))}, + coords={"x": [1, 2]}, + ) + schema = _parse_schema(_ensure_default_indexes(ds)) + assert _field_type(schema, "label") == pa.string() + + +def test_parse_schema_maps_object_string_coord_to_string(): + # A string dimension coordinate (e.g. station names) is object dtype too. + ds = xr.Dataset( + {"v": (["station"], [1.0, 2.0])}, + coords={"station": np.array(["A", "B"], dtype=object)}, + ) + schema = _parse_schema(_ensure_default_indexes(ds)) + assert _field_type(schema, "station") == pa.string() + + +def test_partition_metadata_skips_out_of_ns_datetime(): + # datetime64 coordinates outside the datetime64[ns] range (pre-1678 / + # post-2262) cannot be represented as int64 nanoseconds, so partition + # pruning must be skipped for that dimension rather than raising + # OverflowError. Registration must still succeed. + times = xr.date_range( + "0001-01-01", periods=3, freq="100YS", use_cftime=True + ).to_datetimeindex(time_unit="us", unsafe=True) + ds = _ensure_default_indexes( + xr.Dataset({"v": (["time"], np.arange(3.0))}, coords={"time": times}) + ) + blocks = list(block_slices(ds, chunks={"time": 2})) + + meta = partition_metadata(ds, blocks) # must not raise + + assert len(meta) == len(blocks) + # "time" is unpruneable here, so it is omitted from every partition. + assert all("time" not in m for m in meta) + + +def test_parse_schema_all_null_object_var_stays_null(): + # An all-null object column has no data to infer a type from; let null be + # null rather than coercing it to a string column. + ds = _ensure_default_indexes( + xr.Dataset( + {"label": (["x"], np.array([None, None], dtype=object))}, + coords={"x": [1, 2]}, + ) + ) + schema = _parse_schema(ds) + assert pa.types.is_null(schema.field("label").type) + + +def test_partition_metadata_prunes_cftime_coord(): + # cftime dimension coordinates must produce pruning bounds; previously the + # object-dtype skip shadowed the cftime branch, silently disabling pruning. + times = xr.date_range( + "2000-01-01", periods=4, freq="1D", calendar="noleap", use_cftime=True + ) + ds = _ensure_default_indexes( + xr.Dataset({"v": (["time"], np.arange(4.0))}, coords={"time": times}) + ) + blocks = list(block_slices(ds, chunks={"time": 2})) + + meta = partition_metadata(ds, blocks) + + assert all("time" in m for m in meta) + for m in meta: + _, _, tag = m["time"] + assert tag == "timestamp_ns" + + +def test_partition_metadata_skips_ancient_cftime(): + # Ancient gregorian cftime dates overflow the int64 nanosecond range, so + # pruning must be skipped for that dim (no raise, dim omitted). + times = xr.date_range( + "0001-01-01", periods=3, freq="100YS", use_cftime=True + ) + ds = _ensure_default_indexes( + xr.Dataset({"v": (["time"], np.arange(3.0))}, coords={"time": times}) + ) + blocks = list(block_slices(ds, chunks={"time": 2})) + + meta = partition_metadata(ds, blocks) # must not raise + + assert all("time" not in m for m in meta) + + +def test_string_dataset_round_trips_through_record_batch(): + # The schema fix must also flow through the batch builders: a string + # column has to materialize as an Arrow string array, not error out. + ds = _ensure_default_indexes( + xr.Dataset( + {"label": (["x"], np.array(["a", "b", "c", "d"], dtype=object))}, + coords={"x": [10, 20, 30, 40]}, + ) + ) + schema = _parse_schema(ds) + + batch = dataset_to_record_batch(ds, schema) + assert batch.schema.field("label").type == pa.string() + assert batch.column("label").to_pylist() == ["a", "b", "c", "d"] + + # The streaming path must agree with the one-shot path. + streamed = pa.Table.from_batches( + list(iter_record_batches(ds, schema, batch_size=2)), schema=schema + ) + assert streamed.column("label").to_pylist() == ["a", "b", "c", "d"] + + +def test_partition_metadata_in_range_datetime_still_pruned(): + # Regression guard: ordinary datetimes must keep producing timestamp_ns + # bounds so filter pushdown still works after the overflow fix. + times = pd.date_range("2000-01-01", periods=4, freq="D") + ds = _ensure_default_indexes( + xr.Dataset({"v": (["time"], np.arange(4.0))}, coords={"time": times}) + ) + blocks = list(block_slices(ds, chunks={"time": 2})) + + meta = partition_metadata(ds, blocks) + + assert all("time" in m for m in meta) + for m in meta: + _, _, tag = m["time"] + assert tag == "timestamp_ns" diff --git a/xarray_sql/cftime.py b/xarray_sql/cftime.py index 6ce7b4e..9acaa94 100644 --- a/xarray_sql/cftime.py +++ b/xarray_sql/cftime.py @@ -172,18 +172,27 @@ def convert_for_field(values, field: pa.Field) -> np.ndarray: def partition_bounds( values, -) -> tuple[int, int, str]: +) -> tuple[int, int, str] | None: """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"``. + + Returns ``None`` when the nanosecond bound falls outside the int64 range + (e.g. paleoclimate dates before ~1678), signalling the caller to skip + pruning for that dimension rather than emit a bound the Rust layer would + reject. """ 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" + lo, hi = int(us.min()) * 1_000, int(us.max()) * 1_000 + int64 = np.iinfo(np.int64) + if lo < int64.min or hi > int64.max: + return None + return lo, hi, "timestamp_ns" offsets = to_offsets(values, DEFAULT_UNITS, cal) return int(offsets.min()), int(offsets.max()), "int64" diff --git a/xarray_sql/df.py b/xarray_sql/df.py index ab80056..88aa06f 100644 --- a/xarray_sql/df.py +++ b/xarray_sql/df.py @@ -386,6 +386,20 @@ def iter_record_batches( yield pa.RecordBatch.from_arrays(arrays, schema=schema) +def _arrow_type_for_object(values: np.ndarray) -> pa.DataType: + """Infer an Arrow type for a non-cftime object-dtype array. + + ``pa.from_numpy_dtype`` cannot map numpy object dtype, so let pyarrow infer + the type from the data instead: strings become ``pa.string()``, bytes + ``pa.binary()``, and other representable Python scalars their Arrow + equivalent. An all-null array stays ``pa.null()``, and a column mixing + incompatible types (e.g. str and int) raises, surfacing a clear error + rather than a silent coercion. Object-dtype arrays are never Dask/Zarr + backed, so this triggers no remote I/O. + """ + return pa.array(np.asarray(values).ravel()).type + + def _parse_schema(ds: xr.Dataset) -> pa.Schema: """Extracts a `pa.Schema` from the Dataset, treating dims and data_vars as columns. @@ -414,19 +428,29 @@ def _parse_schema(ds: xr.Dataset) -> pa.Schema: if cft.is_cftime_index(ds, coord_name): units, calendar = cft.encoding(ds, coord_name) columns.append(cft.arrow_field(coord_name, units, calendar)) + elif coord_var.dtype == np.dtype("O"): + # Object dtype that isn't cftime (e.g. string station names). + arrow_type = _arrow_type_for_object(coord_var.values) + columns.append(pa.field(coord_name, arrow_type)) 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(): - # 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)) + # An object-dtype data variable may hold cftime objects (encode it like + # a cftime coordinate) or strings/other Python scalars (infer the Arrow + # type from the data). The dtype check keeps the common numeric path off + # the object branch. + if var.dtype == np.dtype("O"): + if cft.is_cftime(var.values): + # Encode with the same units/calendar as a cftime coordinate. + cal = var.values.ravel()[0].calendar + columns.append( + cft.arrow_field(var_name, cft.DEFAULT_UNITS, cal) + ) + else: + arrow_type = _arrow_type_for_object(var.values) + columns.append(pa.field(var_name, arrow_type)) else: pa_type = pa.from_numpy_dtype(var.dtype) columns.append(pa.field(var_name, pa_type)) @@ -466,16 +490,21 @@ def _block_metadata( coord_values = coord_arrays[str(dim)][slc] if len(coord_values) == 0: continue + # cftime coordinates are object dtype but carry their own bound + # encoding, so they must be handled before the string/object skip + # below (otherwise pruning is silently disabled for them). + # partition_bounds returns None when the bound overflows int64. + if cft.is_cftime(coord_values): + bounds = cft.partition_bounds(coord_values) + if bounds is not None: + ranges[str(dim)] = bounds + continue # String/object dtypes are not representable as ScalarBound # (Int64/Float64/TimestampNanos) and numpy min/max ufuncs do not # 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. @@ -483,9 +512,18 @@ def _block_metadata( max_val = coord_values.max() if isinstance(min_val, (np.datetime64, pd.Timestamp)): - min_val = int(pd.Timestamp(min_val).value) - max_val = int(pd.Timestamp(max_val).value) - ranges[str(dim)] = (min_val, max_val, "timestamp_ns") + # The Rust pruning layer only accepts int64 nanosecond bounds + # (ScalarBound::TimestampNanos). Dates outside the + # datetime64[ns] range (pre-1678 / post-2262) cannot be + # represented, so skip pruning for this dimension rather than + # raising -- registration still succeeds and the Rust pruner + # treats a missing dimension conservatively (never prunes on it). + try: + min_ns = int(pd.Timestamp(min_val).value) + max_ns = int(pd.Timestamp(max_val).value) + except (OverflowError, pd.errors.OutOfBoundsDatetime): + continue + ranges[str(dim)] = (min_ns, max_ns, "timestamp_ns") elif hasattr(min_val, "item"): min_val = min_val.item() max_val = max_val.item()