Skip to content

Commit c700c2a

Browse files
committed
Skip string coordinates in partition bound computation
String-typed coordinates (dtype U, S, O) crash in _block_metadata() because numpy min/max ufuncs do not support string comparison. Since ScalarBound only supports Int64/Float64/TimestampNanos, skip string coordinates so the Rust pruning logic treats the dimension conservatively (never prunes on it). Partial fix for #121.
1 parent 67f075a commit c700c2a

2 files changed

Lines changed: 44 additions & 19 deletions

File tree

tests/test_sql.py

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
"""SQL functionality tests for xarray-sql using pytest."""
22

3+
import numpy as np
34
import pytest
45
import xarray as xr
56

@@ -127,3 +128,21 @@ def test_cross_join(air_and_stations):
127128
"SELECT COUNT(*) AS total FROM air_data CROSS JOIN stations"
128129
).to_pandas()
129130
assert result["total"].iloc[0] > 0
131+
132+
133+
def test_string_coordinates():
134+
"""String-typed coordinates should not crash during registration."""
135+
ds = xr.Dataset(
136+
{"score": (["student", "subject"], np.random.rand(3, 2))},
137+
coords={
138+
"student": ["alice", "bob", "charlie"],
139+
"subject": ["math", "science"],
140+
},
141+
)
142+
ctx = XarrayContext()
143+
ctx.from_dataset("scores", ds.chunk({"student": 3, "subject": 2}))
144+
result = ctx.sql("SELECT * FROM scores").to_pandas()
145+
assert len(result) == 6
146+
assert "student" in result.columns
147+
assert "subject" in result.columns
148+
assert "score" in result.columns

xarray_sql/df.py

Lines changed: 25 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -323,25 +323,31 @@ def _block_metadata(coord_arrays: dict, block: Block) -> PartitionBounds:
323323
ranges: PartitionBounds = {}
324324
for dim, slc in block.items():
325325
coord_values = coord_arrays[str(dim)][slc]
326-
if len(coord_values) > 0:
327-
# Use actual min/max rather than first/last so that non-monotonic
328-
# coordinate axes (e.g. descending latitude 90→-90) are handled
329-
# correctly. np.min/max work for both numeric and datetime64 arrays.
330-
min_val = coord_values.min()
331-
max_val = coord_values.max()
332-
333-
if isinstance(min_val, (np.datetime64, pd.Timestamp)):
334-
min_val = int(pd.Timestamp(min_val).value)
335-
max_val = int(pd.Timestamp(max_val).value)
336-
ranges[str(dim)] = (min_val, max_val, "timestamp_ns")
337-
elif hasattr(min_val, "item"):
338-
min_val = min_val.item()
339-
max_val = max_val.item()
340-
dtype = "float64" if isinstance(min_val, float) else "int64"
341-
ranges[str(dim)] = (min_val, max_val, dtype)
342-
else:
343-
dtype = "float64" if isinstance(min_val, float) else "int64"
344-
ranges[str(dim)] = (min_val, max_val, dtype)
326+
if len(coord_values) == 0:
327+
continue
328+
# String/object dtypes are not representable as ScalarBound
329+
# (Int64/Float64/TimestampNanos) and numpy min/max ufuncs do not
330+
# support them. Skip so pruning treats the dimension conservatively.
331+
if coord_values.dtype.kind in ('U', 'S', 'O'):
332+
continue
333+
# Use actual min/max rather than first/last so that non-monotonic
334+
# coordinate axes (e.g. descending latitude 90→-90) are handled
335+
# correctly. np.min/max work for both numeric and datetime64 arrays.
336+
min_val = coord_values.min()
337+
max_val = coord_values.max()
338+
339+
if isinstance(min_val, (np.datetime64, pd.Timestamp)):
340+
min_val = int(pd.Timestamp(min_val).value)
341+
max_val = int(pd.Timestamp(max_val).value)
342+
ranges[str(dim)] = (min_val, max_val, "timestamp_ns")
343+
elif hasattr(min_val, "item"):
344+
min_val = min_val.item()
345+
max_val = max_val.item()
346+
dtype = "float64" if isinstance(min_val, float) else "int64"
347+
ranges[str(dim)] = (min_val, max_val, dtype)
348+
else:
349+
dtype = "float64" if isinstance(min_val, float) else "int64"
350+
ranges[str(dim)] = (min_val, max_val, dtype)
345351
return ranges
346352

347353

0 commit comments

Comments
 (0)