From 83e69398d4e92804ff9bd822964a1cc8a8fc52dd Mon Sep 17 00:00:00 2001 From: ghostiee-11 Date: Sun, 5 Jul 2026 22:02:38 +0530 Subject: [PATCH 1/2] fix: register object-dtype and out-of-ns-range coordinates (#223) _parse_schema handed object dtype straight to pa.from_numpy_dtype, which raises "Unsupported numpy type 17" for string variables and coordinates. Infer the Arrow type from a bounded sample instead (strings map to pa.string(); an empty or all-null object column defaults to pa.string()). _block_metadata forced datetime bounds to int64 nanoseconds, overflowing for dates outside the datetime64[ns] range (before 1678 or after 2262). Skip pruning for such dimensions rather than raising, so registration still succeeds. Also revive cftime partition pruning: the string/object skip shadowed the cftime branch, silently disabling it. Check cftime first, and have partition_bounds return None when the nanosecond bound overflows int64 so that ancient dates are skipped rather than rejected by the Rust layer. Adds regression tests for each case plus an end-to-end pruning check. --- tests/test_cft.py | 8 +++ tests/test_df.py | 138 +++++++++++++++++++++++++++++++++++++++++++ xarray_sql/cftime.py | 18 +++++- xarray_sql/df.py | 74 +++++++++++++++++++---- 4 files changed, 224 insertions(+), 14 deletions(-) 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..bd0cf13 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,139 @@ 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 (issue #223) ------- + + +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_defaults_to_string(): + # An all-null (or empty) object column has no data to infer from; object + # dtype in xarray almost always means strings, so default to pa.string() + # rather than a null column. + ds = _ensure_default_indexes( + xr.Dataset( + {"label": (["x"], np.array([None, None], dtype=object))}, + coords={"x": [1, 2]}, + ) + ) + schema = _parse_schema(ds) + assert schema.field("label").type == pa.string() + + +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..3d2fe43 100644 --- a/xarray_sql/cftime.py +++ b/xarray_sql/cftime.py @@ -170,20 +170,34 @@ def convert_for_field(values, field: pa.Field) -> np.ndarray: # --------------------------------------------------------------------------- +#: Bounds of a signed 64-bit integer; a pruning bound outside this range +#: cannot be handed to the Rust ``ScalarBound`` layer. +_INT64_MIN: int = -(2**63) +_INT64_MAX: int = 2**63 - 1 + + 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 + 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..771de66 100644 --- a/xarray_sql/df.py +++ b/xarray_sql/df.py @@ -386,6 +386,31 @@ def iter_record_batches( yield pa.RecordBatch.from_arrays(arrays, schema=schema) +def _arrow_type_for_object( + values: np.ndarray, sample_size: int = 100 +) -> pa.DataType: + """Infer an Arrow type for a non-cftime object-dtype array. + + ``pa.from_numpy_dtype`` cannot map numpy object dtype and raises + ``ArrowNotImplementedError: Unsupported numpy type 17`` -- which is exactly + what a string variable or coordinate produces. Object-dtype arrays are + never Dask/Zarr-backed, so letting Arrow infer the type from the data with + ``pa.array`` (e.g. ``pa.string()`` for Python strings) is safe and triggers + no remote I/O. + + Only the first ``sample_size`` elements are inspected so a huge in-memory + object column is not fully copied just to read its type. An empty or + all-null sample yields ``pa.null()``, which is meaningless as a column + type; object dtype in xarray almost always means strings, so fall back to + ``pa.string()`` in that case. + """ + sample = np.asarray(values).ravel()[:sample_size] + arrow_type = pa.array(sample).type + if pa.types.is_null(arrow_type): + return pa.string() + return arrow_type + + def _parse_schema(ds: xr.Dataset) -> pa.Schema: """Extracts a `pa.Schema` from the Dataset, treating dims and data_vars as columns. @@ -414,6 +439,10 @@ 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)) @@ -422,11 +451,18 @@ def _parse_schema(ds: xr.Dataset) -> pa.Schema: # 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)) + if var.dtype == np.dtype("O"): + if 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: + # String / other object data variable. + 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 +502,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 +524,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() From f773dca8553c09ce09a2692e118c9b102cf45c61 Mon Sep 17 00:00:00 2001 From: ghostiee-11 Date: Tue, 7 Jul 2026 03:08:12 +0530 Subject: [PATCH 2/2] refactor: address review feedback on object-dtype handling Infer object-dtype columns directly with pyarrow instead of sampling and coercing: strings map to string, other representable scalars to their Arrow type, an all-null column stays null, and a column mixing incompatible types raises rather than being silently coerced to string. Derive the int64 pruning-bound limits from np.iinfo(np.int64) rather than hardcoded literals, and drop the inaccurate "rare" comments on the cftime data-variable path. --- tests/test_df.py | 11 +++++------ xarray_sql/cftime.py | 9 ++------- xarray_sql/df.py | 40 ++++++++++++++-------------------------- 3 files changed, 21 insertions(+), 39 deletions(-) diff --git a/tests/test_df.py b/tests/test_df.py index bd0cf13..5ba5b3f 100644 --- a/tests/test_df.py +++ b/tests/test_df.py @@ -536,7 +536,7 @@ def test_compute_chunks_tuples_sum_to_dim_size(): assert sum(tup) == ds.sizes[dim] -# -- Object-dtype and out-of-ns-range coordinate support (issue #223) ------- +# -- Object-dtype and out-of-ns-range coordinate support -------------------- def _field_type(schema, name): @@ -584,10 +584,9 @@ def test_partition_metadata_skips_out_of_ns_datetime(): assert all("time" not in m for m in meta) -def test_parse_schema_all_null_object_var_defaults_to_string(): - # An all-null (or empty) object column has no data to infer from; object - # dtype in xarray almost always means strings, so default to pa.string() - # rather than a null column. +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))}, @@ -595,7 +594,7 @@ def test_parse_schema_all_null_object_var_defaults_to_string(): ) ) schema = _parse_schema(ds) - assert schema.field("label").type == pa.string() + assert pa.types.is_null(schema.field("label").type) def test_partition_metadata_prunes_cftime_coord(): diff --git a/xarray_sql/cftime.py b/xarray_sql/cftime.py index 3d2fe43..9acaa94 100644 --- a/xarray_sql/cftime.py +++ b/xarray_sql/cftime.py @@ -170,12 +170,6 @@ def convert_for_field(values, field: pa.Field) -> np.ndarray: # --------------------------------------------------------------------------- -#: Bounds of a signed 64-bit integer; a pruning bound outside this range -#: cannot be handed to the Rust ``ScalarBound`` layer. -_INT64_MIN: int = -(2**63) -_INT64_MAX: int = 2**63 - 1 - - def partition_bounds( values, ) -> tuple[int, int, str] | None: @@ -195,7 +189,8 @@ def partition_bounds( if is_gregorian_like(cal): us = to_microseconds(values) lo, hi = int(us.min()) * 1_000, int(us.max()) * 1_000 - if lo < _INT64_MIN or hi > _INT64_MAX: + 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) diff --git a/xarray_sql/df.py b/xarray_sql/df.py index 771de66..88aa06f 100644 --- a/xarray_sql/df.py +++ b/xarray_sql/df.py @@ -386,29 +386,18 @@ def iter_record_batches( yield pa.RecordBatch.from_arrays(arrays, schema=schema) -def _arrow_type_for_object( - values: np.ndarray, sample_size: int = 100 -) -> pa.DataType: +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 and raises - ``ArrowNotImplementedError: Unsupported numpy type 17`` -- which is exactly - what a string variable or coordinate produces. Object-dtype arrays are - never Dask/Zarr-backed, so letting Arrow infer the type from the data with - ``pa.array`` (e.g. ``pa.string()`` for Python strings) is safe and triggers - no remote I/O. - - Only the first ``sample_size`` elements are inspected so a huge in-memory - object column is not fully copied just to read its type. An empty or - all-null sample yields ``pa.null()``, which is meaningless as a column - type; object dtype in xarray almost always means strings, so fall back to - ``pa.string()`` in that case. + ``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. """ - sample = np.asarray(values).ravel()[:sample_size] - arrow_type = pa.array(sample).type - if pa.types.is_null(arrow_type): - return pa.string() - return arrow_type + return pa.array(np.asarray(values).ravel()).type def _parse_schema(ds: xr.Dataset) -> pa.Schema: @@ -448,19 +437,18 @@ def _parse_schema(ds: xr.Dataset) -> pa.Schema: 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. + # 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): - # Rare: a data variable holding cftime objects. Use same - # encoding as the first cftime dimension coordinate, or default. + # 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: - # String / other object data variable. arrow_type = _arrow_type_for_object(var.values) columns.append(pa.field(var_name, arrow_type)) else: