fix: register object-dtype and out-of-ns-range coordinates (#223)

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 ----------------------------------------
 

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"

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"
 

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,26 +490,40 @@ 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.
         min_val = coord_values.min()
         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()