Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 19 additions & 0 deletions tests/test_sql.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
"""SQL functionality tests for xarray-sql using pytest."""

import numpy as np
import pytest
import xarray as xr

Expand Down Expand Up @@ -127,3 +128,21 @@ def test_cross_join(air_and_stations):
"SELECT COUNT(*) AS total FROM air_data CROSS JOIN stations"
).to_pandas()
assert result["total"].iloc[0] > 0


def test_string_coordinates():

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks, good to have this test, and it's clear to read.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the review!!

"""String-typed coordinates should not crash during registration."""
ds = xr.Dataset(
{"score": (["student", "subject"], np.random.rand(3, 2))},
coords={
"student": ["alice", "bob", "charlie"],
"subject": ["math", "science"],
},
)
ctx = XarrayContext()
ctx.from_dataset("scores", ds.chunk({"student": 3, "subject": 2}))
result = ctx.sql("SELECT * FROM scores").to_pandas()
assert len(result) == 6
assert "student" in result.columns
assert "subject" in result.columns
assert "score" in result.columns
44 changes: 25 additions & 19 deletions xarray_sql/df.py
Original file line number Diff line number Diff line change
Expand Up @@ -323,25 +323,31 @@ def _block_metadata(coord_arrays: dict, block: Block) -> PartitionBounds:
ranges: PartitionBounds = {}
for dim, slc in block.items():
coord_values = coord_arrays[str(dim)][slc]
if len(coord_values) > 0:
# 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")
elif hasattr(min_val, "item"):
min_val = min_val.item()
max_val = max_val.item()
dtype = "float64" if isinstance(min_val, float) else "int64"
ranges[str(dim)] = (min_val, max_val, dtype)
else:
dtype = "float64" if isinstance(min_val, float) else "int64"
ranges[str(dim)] = (min_val, max_val, dtype)
if len(coord_values) == 0:
continue
Comment on lines +326 to +327

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks, this is a good idea.

# 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
# 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")
elif hasattr(min_val, "item"):
min_val = min_val.item()
max_val = max_val.item()
dtype = "float64" if isinstance(min_val, float) else "int64"
ranges[str(dim)] = (min_val, max_val, dtype)
else:
dtype = "float64" if isinstance(min_val, float) else "int64"
ranges[str(dim)] = (min_val, max_val, dtype)
return ranges


Expand Down
Loading