xarray_sql 0.3.2__tar.gz → 0.3.3__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {xarray_sql-0.3.2 → xarray_sql-0.3.3}/Cargo.lock +1 -1
- {xarray_sql-0.3.2 → xarray_sql-0.3.3}/Cargo.toml +1 -1
- {xarray_sql-0.3.2 → xarray_sql-0.3.3}/PKG-INFO +11 -3
- {xarray_sql-0.3.2 → xarray_sql-0.3.3}/README.md +7 -2
- {xarray_sql-0.3.2 → xarray_sql-0.3.3}/benchmarks/geospatial/07_reproject_udf.py +20 -59
- {xarray_sql-0.3.2 → xarray_sql-0.3.3}/benchmarks/geospatial/09_warp.py +6 -37
- {xarray_sql-0.3.2 → xarray_sql-0.3.3}/benchmarks/geospatial/README.md +1 -1
- {xarray_sql-0.3.2 → xarray_sql-0.3.3}/docs/geospatial.md +18 -8
- {xarray_sql-0.3.2 → xarray_sql-0.3.3}/pyproject.toml +4 -0
- {xarray_sql-0.3.2 → xarray_sql-0.3.3}/tests/test_cft.py +8 -0
- {xarray_sql-0.3.2 → xarray_sql-0.3.3}/tests/test_df.py +161 -0
- xarray_sql-0.3.3/tests/test_proj.py +158 -0
- {xarray_sql-0.3.2 → xarray_sql-0.3.3}/tests/test_to_dataset_perf.py +18 -9
- {xarray_sql-0.3.2 → xarray_sql-0.3.3}/xarray_sql/cftime.py +11 -2
- {xarray_sql-0.3.2 → xarray_sql-0.3.3}/xarray_sql/df.py +60 -15
- xarray_sql-0.3.3/xarray_sql/proj.py +235 -0
- {xarray_sql-0.3.2 → xarray_sql-0.3.3}/xarray_sql/sql.py +13 -0
- {xarray_sql-0.3.2 → xarray_sql-0.3.3}/.gitignore +0 -0
- {xarray_sql-0.3.2 → xarray_sql-0.3.3}/AGENTS.md +0 -0
- {xarray_sql-0.3.2 → xarray_sql-0.3.3}/LICENSE +0 -0
- {xarray_sql-0.3.2 → xarray_sql-0.3.3}/benchmarks/geospatial/01_ndvi.py +0 -0
- {xarray_sql-0.3.2 → xarray_sql-0.3.3}/benchmarks/geospatial/02_climatology.py +0 -0
- {xarray_sql-0.3.2 → xarray_sql-0.3.3}/benchmarks/geospatial/03_zonal_mean.py +0 -0
- {xarray_sql-0.3.2 → xarray_sql-0.3.3}/benchmarks/geospatial/04_anomaly.py +0 -0
- {xarray_sql-0.3.2 → xarray_sql-0.3.3}/benchmarks/geospatial/05_forecast_skill.py +0 -0
- {xarray_sql-0.3.2 → xarray_sql-0.3.3}/benchmarks/geospatial/06_zonal_vector.py +0 -0
- {xarray_sql-0.3.2 → xarray_sql-0.3.3}/benchmarks/geospatial/08_regrid_weights.py +0 -0
- {xarray_sql-0.3.2 → xarray_sql-0.3.3}/benchmarks/geospatial/_harness.py +0 -0
- {xarray_sql-0.3.2 → xarray_sql-0.3.3}/benchmarks/geospatial/perf_summary.py +0 -0
- {xarray_sql-0.3.2 → xarray_sql-0.3.3}/benchmarks/geospatial/run_all.sh +0 -0
- {xarray_sql-0.3.2 → xarray_sql-0.3.3}/benchmarks/geospatial/run_perf.sh +0 -0
- {xarray_sql-0.3.2 → xarray_sql-0.3.3}/docs/assets/logo.svg +0 -0
- {xarray_sql-0.3.2 → xarray_sql-0.3.3}/docs/contributing.md +0 -0
- {xarray_sql-0.3.2 → xarray_sql-0.3.3}/docs/examples.md +0 -0
- {xarray_sql-0.3.2 → xarray_sql-0.3.3}/docs/index.md +0 -0
- {xarray_sql-0.3.2 → xarray_sql-0.3.3}/docs/reference/xarray_sql.md +0 -0
- {xarray_sql-0.3.2 → xarray_sql-0.3.3}/src/lib.rs +0 -0
- {xarray_sql-0.3.2 → xarray_sql-0.3.3}/tests/__init__.py +0 -0
- {xarray_sql-0.3.2 → xarray_sql-0.3.3}/tests/conftest.py +0 -0
- {xarray_sql-0.3.2 → xarray_sql-0.3.3}/tests/test_ds.py +0 -0
- {xarray_sql-0.3.2 → xarray_sql-0.3.3}/tests/test_reader.py +0 -0
- {xarray_sql-0.3.2 → xarray_sql-0.3.3}/tests/test_sql.py +0 -0
- {xarray_sql-0.3.2 → xarray_sql-0.3.3}/tests/test_stats.py +0 -0
- {xarray_sql-0.3.2 → xarray_sql-0.3.3}/xarray_sql/__init__.py +0 -0
- {xarray_sql-0.3.2 → xarray_sql-0.3.3}/xarray_sql/core.py +0 -0
- {xarray_sql-0.3.2 → xarray_sql-0.3.3}/xarray_sql/ds.py +0 -0
- {xarray_sql-0.3.2 → xarray_sql-0.3.3}/xarray_sql/reader.py +0 -0
- {xarray_sql-0.3.2 → xarray_sql-0.3.3}/zensical.toml +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: xarray_sql
|
|
3
|
-
Version: 0.3.
|
|
3
|
+
Version: 0.3.3
|
|
4
4
|
Classifier: Development Status :: 4 - Beta
|
|
5
5
|
Classifier: Intended Audience :: Science/Research
|
|
6
6
|
Classifier: Intended Audience :: Developers
|
|
@@ -25,12 +25,15 @@ Requires-Dist: pytest ; extra == 'dev'
|
|
|
25
25
|
Requires-Dist: watchfiles ; extra == 'dev'
|
|
26
26
|
Requires-Dist: zensical ; extra == 'docs'
|
|
27
27
|
Requires-Dist: mkdocstrings[python] ; extra == 'docs'
|
|
28
|
+
Requires-Dist: pyproj ; extra == 'geo'
|
|
28
29
|
Requires-Dist: cftime ; extra == 'test'
|
|
30
|
+
Requires-Dist: xarray-sql[geo] ; extra == 'test'
|
|
29
31
|
Requires-Dist: pytest ; extra == 'test'
|
|
30
32
|
Requires-Dist: xarray[io] ; extra == 'test'
|
|
31
33
|
Requires-Dist: gcsfs ; extra == 'test'
|
|
32
34
|
Provides-Extra: dev
|
|
33
35
|
Provides-Extra: docs
|
|
36
|
+
Provides-Extra: geo
|
|
34
37
|
Provides-Extra: test
|
|
35
38
|
License-File: LICENSE
|
|
36
39
|
Summary: Querry Xarray with SQL.
|
|
@@ -45,10 +48,13 @@ Project-URL: Issues, https://github.com/alxmrs/xarray-sql/issues
|
|
|
45
48
|
|
|
46
49
|
_Query [Xarray](https://xarray.dev/) with SQL_
|
|
47
50
|
|
|
51
|
+

|
|
48
52
|
[](https://github.com/alxmrs/xarray-sql/actions/workflows/ci.yml)
|
|
49
53
|
[](https://github.com/alxmrs/xarray-sql/actions/workflows/lint.yml)
|
|
50
54
|
[](https://github.com/alxmrs/xarray-sql/actions/workflows/ci-build.yml)
|
|
51
55
|
[](https://github.com/alxmrs/xarray-sql/actions/workflows/ci-rust.yml)
|
|
56
|
+
[](https://pepy.tech/projects/xarray-sql)
|
|
57
|
+
[](https://pepy.tech/projects/xarray-sql)
|
|
52
58
|
|
|
53
59
|
```shell
|
|
54
60
|
pip install xarray-sql
|
|
@@ -253,8 +259,10 @@ against an xarray/array reference** to floating-point tolerance:
|
|
|
253
259
|
reproduces the published result that GraphCast beats Pangu at every lead.
|
|
254
260
|
* **Raster × vector zonal stats** — a range `JOIN` of the ERA5 grid against a
|
|
255
261
|
table of regions.
|
|
256
|
-
* **Reprojection and regridding** — a
|
|
257
|
-
|
|
262
|
+
* **Reprojection and regridding** — a `reproject(x, y, src_crs, dst_crs)`
|
|
263
|
+
scalar PROJ UDF, shipped as the optional geo extension
|
|
264
|
+
(`pip install xarray-sql[geo]`, validated against Earth Engine's own
|
|
265
|
+
geodesy via [Xee](https://github.com/google/Xee)) and a
|
|
258
266
|
sparse-weight-table `JOIN` (regridding real SRTM terrain).
|
|
259
267
|
|
|
260
268
|
Every case matches its array reference. The headline finding: these operations
|
|
@@ -2,10 +2,13 @@
|
|
|
2
2
|
|
|
3
3
|
_Query [Xarray](https://xarray.dev/) with SQL_
|
|
4
4
|
|
|
5
|
+

|
|
5
6
|
[](https://github.com/alxmrs/xarray-sql/actions/workflows/ci.yml)
|
|
6
7
|
[](https://github.com/alxmrs/xarray-sql/actions/workflows/lint.yml)
|
|
7
8
|
[](https://github.com/alxmrs/xarray-sql/actions/workflows/ci-build.yml)
|
|
8
9
|
[](https://github.com/alxmrs/xarray-sql/actions/workflows/ci-rust.yml)
|
|
10
|
+
[](https://pepy.tech/projects/xarray-sql)
|
|
11
|
+
[](https://pepy.tech/projects/xarray-sql)
|
|
9
12
|
|
|
10
13
|
```shell
|
|
11
14
|
pip install xarray-sql
|
|
@@ -210,8 +213,10 @@ against an xarray/array reference** to floating-point tolerance:
|
|
|
210
213
|
reproduces the published result that GraphCast beats Pangu at every lead.
|
|
211
214
|
* **Raster × vector zonal stats** — a range `JOIN` of the ERA5 grid against a
|
|
212
215
|
table of regions.
|
|
213
|
-
* **Reprojection and regridding** — a
|
|
214
|
-
|
|
216
|
+
* **Reprojection and regridding** — a `reproject(x, y, src_crs, dst_crs)`
|
|
217
|
+
scalar PROJ UDF, shipped as the optional geo extension
|
|
218
|
+
(`pip install xarray-sql[geo]`, validated against Earth Engine's own
|
|
219
|
+
geodesy via [Xee](https://github.com/google/Xee)) and a
|
|
215
220
|
sparse-weight-table `JOIN` (regridding real SRTM terrain).
|
|
216
221
|
|
|
217
222
|
Every case matches its array reference. The headline finding: these operations
|
|
@@ -2,10 +2,9 @@
|
|
|
2
2
|
# /// script
|
|
3
3
|
# requires-python = ">=3.11"
|
|
4
4
|
# dependencies = [
|
|
5
|
-
# "xarray-sql",
|
|
5
|
+
# "xarray-sql[geo]",
|
|
6
6
|
# "xarray",
|
|
7
7
|
# "numpy",
|
|
8
|
-
# "pyproj",
|
|
9
8
|
# "pyarrow",
|
|
10
9
|
# "xee",
|
|
11
10
|
# "earthengine-api",
|
|
@@ -24,9 +23,13 @@ the shape of a SQL *scalar UDF*, and it is precisely how the geospatial SQL
|
|
|
24
23
|
world already does it — PostGIS ``ST_Transform`` and DuckDB-spatial
|
|
25
24
|
``ST_Transform`` are scalar PROJ wrappers.
|
|
26
25
|
|
|
27
|
-
|
|
26
|
+
xarray-sql ships that UDF as its geo extension (``xarray_sql.proj``):
|
|
27
|
+
with pyproj installed, every ``XarrayContext`` speaks CRS out of the box,
|
|
28
|
+
and the CRS pair is part of the query rather than baked into the UDF::
|
|
28
29
|
|
|
29
|
-
SELECT x, y,
|
|
30
|
+
SELECT x, y,
|
|
31
|
+
reproject(x, y, 'EPSG:32610', 'EPSG:4326')['x'] AS lon,
|
|
32
|
+
reproject(x, y, 'EPSG:32610', 'EPSG:4326')['y'] AS lat
|
|
30
33
|
FROM grid
|
|
31
34
|
|
|
32
35
|
**The reference is Earth Engine itself.** There is *one* dataset: a single UTM
|
|
@@ -39,9 +42,9 @@ UDF and check the lon/lat it returns against EE's independently-computed lon/lat
|
|
|
39
42
|
for the *same* pixels. The reference is a different geodesy engine, not PROJ
|
|
40
43
|
again, and they agree to sub-metre precision.
|
|
41
44
|
|
|
42
|
-
|
|
43
|
-
|
|
44
|
-
|
|
45
|
+
The extension returns *both* coordinates from one struct-returning call
|
|
46
|
+
(one PROJ transform per row) and runs all PROJ work on its own worker
|
|
47
|
+
pool, so the query parallelizes across partitions safely.
|
|
45
48
|
|
|
46
49
|
Requires Earth Engine access: ``earthengine authenticate`` once, then an
|
|
47
50
|
initialized project (set ``EARTHENGINE_PROJECT``). Skips cleanly otherwise.
|
|
@@ -49,11 +52,7 @@ initialized project (set ``EARTHENGINE_PROJECT``). Skips cleanly otherwise.
|
|
|
49
52
|
|
|
50
53
|
from __future__ import annotations
|
|
51
54
|
|
|
52
|
-
import numpy as np
|
|
53
|
-
import pyarrow as pa
|
|
54
|
-
import pyproj
|
|
55
55
|
import xarray as xr
|
|
56
|
-
from datafusion import udf
|
|
57
56
|
|
|
58
57
|
import xarray_sql as xql
|
|
59
58
|
|
|
@@ -73,46 +72,6 @@ _AOI = (-122.6, 37.4, -121.6, 38.4)
|
|
|
73
72
|
_SCALE_M = 2_000 # 2 km pixels → a ~50×60 grid
|
|
74
73
|
|
|
75
74
|
|
|
76
|
-
def register_reproject_udf(
|
|
77
|
-
ctx, src_crs: str, dst_crs: str, name: str = "reproject"
|
|
78
|
-
) -> None:
|
|
79
|
-
"""Register a ``reproject(x, y) -> {lon, lat}`` PROJ scalar UDF.
|
|
80
|
-
|
|
81
|
-
Mirrors ``xarray_sql.cftime.make_cftime_udf``: a vectorized scalar UDF over
|
|
82
|
-
Arrow arrays. ``always_xy=True`` keeps argument order (easting, northing) →
|
|
83
|
-
(lon, lat) regardless of CRS axis conventions. Like PostGIS/DuckDB
|
|
84
|
-
``ST_Transform``, it returns *both* output coordinates from one call — here
|
|
85
|
-
as an Arrow struct, so callers write ``reproject(x, y)['lon']``.
|
|
86
|
-
|
|
87
|
-
Returning a struct (rather than two separate UDFs) is deliberate: PROJ's
|
|
88
|
-
context is not thread-safe, and DataFusion evaluates independent projection
|
|
89
|
-
expressions concurrently — two PROJ UDFs in one SELECT race and crash. One
|
|
90
|
-
struct-returning UDF does the transform exactly once per row, on one thread.
|
|
91
|
-
"""
|
|
92
|
-
ret = pa.struct([("lon", pa.float64()), ("lat", pa.float64())])
|
|
93
|
-
|
|
94
|
-
def _fn(x: pa.Array, y: pa.Array) -> pa.Array:
|
|
95
|
-
# Build the Transformer inside the call so it lives on the worker
|
|
96
|
-
# thread that uses it (PROJ contexts are thread-bound).
|
|
97
|
-
transformer = pyproj.Transformer.from_crs(
|
|
98
|
-
src_crs, dst_crs, always_xy=True
|
|
99
|
-
)
|
|
100
|
-
xs = np.asarray(x.to_numpy(zero_copy_only=False), dtype="float64")
|
|
101
|
-
ys = np.asarray(y.to_numpy(zero_copy_only=False), dtype="float64")
|
|
102
|
-
lon, lat = transformer.transform(xs, ys)
|
|
103
|
-
return pa.StructArray.from_arrays(
|
|
104
|
-
[
|
|
105
|
-
pa.array(np.asarray(lon, "float64")),
|
|
106
|
-
pa.array(np.asarray(lat, "float64")),
|
|
107
|
-
],
|
|
108
|
-
names=["lon", "lat"],
|
|
109
|
-
)
|
|
110
|
-
|
|
111
|
-
ctx.register_udf(
|
|
112
|
-
udf(_fn, [pa.float64(), pa.float64()], ret, "immutable", name)
|
|
113
|
-
)
|
|
114
|
-
|
|
115
|
-
|
|
116
75
|
def _open_ee_lonlat_grid() -> xr.Dataset:
|
|
117
76
|
"""Open ``ee.Image.pixelLonLat()`` on a UTM grid via Xee.
|
|
118
77
|
|
|
@@ -153,17 +112,19 @@ def main() -> None:
|
|
|
153
112
|
f"{_SRC_CRS} → {_DST_CRS}"
|
|
154
113
|
)
|
|
155
114
|
|
|
115
|
+
# XarrayContext registers reproject() automatically (the geo
|
|
116
|
+
# extension). The chunking deliberately splits the ~60-row grid into
|
|
117
|
+
# 15-row chunks → 4 partitions, forcing DataFusion to evaluate the UDF
|
|
118
|
+
# concurrently: the extension runs PROJ on its own worker pool, so
|
|
119
|
+
# parallel partitions are safe (previously this required one chunk →
|
|
120
|
+
# one partition → a serial UDF).
|
|
156
121
|
ctx = xql.XarrayContext()
|
|
157
|
-
|
|
158
|
-
ctx.from_dataset(
|
|
159
|
-
"grid", ds, chunks={"y": ds.sizes["y"], "x": ds.sizes["x"]}
|
|
160
|
-
)
|
|
161
|
-
register_reproject_udf(ctx, _SRC_CRS, _DST_CRS)
|
|
122
|
+
ctx.from_dataset("grid", ds, chunks={"y": 15, "x": ds.sizes["x"]})
|
|
162
123
|
|
|
163
|
-
sql = """
|
|
124
|
+
sql = f"""
|
|
164
125
|
SELECT x, y,
|
|
165
|
-
reproject(x, y)['
|
|
166
|
-
reproject(x, y)['
|
|
126
|
+
reproject(x, y, '{_SRC_CRS}', '{_DST_CRS}')['x'] AS lon,
|
|
127
|
+
reproject(x, y, '{_SRC_CRS}', '{_DST_CRS}')['y'] AS lat
|
|
167
128
|
FROM grid
|
|
168
129
|
ORDER BY y, x
|
|
169
130
|
"""
|
|
@@ -2,10 +2,9 @@
|
|
|
2
2
|
# /// script
|
|
3
3
|
# requires-python = ">=3.11"
|
|
4
4
|
# dependencies = [
|
|
5
|
-
# "xarray-sql",
|
|
5
|
+
# "xarray-sql[geo]",
|
|
6
6
|
# "xarray",
|
|
7
7
|
# "numpy",
|
|
8
|
-
# "pyproj",
|
|
9
8
|
# "pyarrow",
|
|
10
9
|
# "scipy",
|
|
11
10
|
# "xee",
|
|
@@ -55,11 +54,9 @@ cross-check). Requires Earth Engine access; skips cleanly otherwise.
|
|
|
55
54
|
from __future__ import annotations
|
|
56
55
|
|
|
57
56
|
import numpy as np
|
|
58
|
-
import pyarrow as pa
|
|
59
57
|
import pyproj
|
|
60
58
|
import shapely.geometry as sgeom
|
|
61
59
|
import xarray as xr
|
|
62
|
-
from datafusion import udf
|
|
63
60
|
|
|
64
61
|
import xarray_sql as xql
|
|
65
62
|
|
|
@@ -81,35 +78,6 @@ _SRC_SCALE_M = 2_000.0 # ~2 km source pixels
|
|
|
81
78
|
_DST_SCALE_DEG = 0.02 # ~2 km target cells
|
|
82
79
|
|
|
83
80
|
|
|
84
|
-
def _register_reproject_udf(ctx, src_crs, dst_crs, name="reproject"):
|
|
85
|
-
"""Register ``reproject(a, b) -> {x, y}`` — case 07's PROJ scalar UDF.
|
|
86
|
-
|
|
87
|
-
Vectorized over each Arrow batch; ``always_xy=True`` keeps (easting, northing)
|
|
88
|
-
/(lon, lat) order. Returns both output coordinates from one struct-returning
|
|
89
|
-
call (PROJ contexts are not thread-safe, so one UDF, evaluated serially).
|
|
90
|
-
"""
|
|
91
|
-
ret = pa.struct([("x", pa.float64()), ("y", pa.float64())])
|
|
92
|
-
|
|
93
|
-
def _fn(a: pa.Array, b: pa.Array) -> pa.Array:
|
|
94
|
-
transformer = pyproj.Transformer.from_crs(
|
|
95
|
-
src_crs, dst_crs, always_xy=True
|
|
96
|
-
)
|
|
97
|
-
xs = np.asarray(a.to_numpy(zero_copy_only=False), dtype="float64")
|
|
98
|
-
ys = np.asarray(b.to_numpy(zero_copy_only=False), dtype="float64")
|
|
99
|
-
ox, oy = transformer.transform(xs, ys)
|
|
100
|
-
return pa.StructArray.from_arrays(
|
|
101
|
-
[
|
|
102
|
-
pa.array(np.asarray(ox, "float64")),
|
|
103
|
-
pa.array(np.asarray(oy, "float64")),
|
|
104
|
-
],
|
|
105
|
-
names=["x", "y"],
|
|
106
|
-
)
|
|
107
|
-
|
|
108
|
-
ctx.register_udf(
|
|
109
|
-
udf(_fn, [pa.float64(), pa.float64()], ret, "immutable", name)
|
|
110
|
-
)
|
|
111
|
-
|
|
112
|
-
|
|
113
81
|
def _open_srtm(
|
|
114
82
|
grid_crs: str, scale: tuple[float, float], xy_names
|
|
115
83
|
) -> xr.DataArray:
|
|
@@ -212,8 +180,9 @@ def main() -> None:
|
|
|
212
180
|
f"{len(tlat)}×{len(tlon)} ({_SRC_CRS} → {_DST_CRS})"
|
|
213
181
|
)
|
|
214
182
|
|
|
183
|
+
# XarrayContext registers reproject() automatically (the geo
|
|
184
|
+
# extension) — the direction is spelled in the query itself.
|
|
215
185
|
ctx = xql.XarrayContext()
|
|
216
|
-
_register_reproject_udf(ctx, _DST_CRS, _SRC_CRS)
|
|
217
186
|
|
|
218
187
|
# The target grid as a (dst_lat, dst_lon) table.
|
|
219
188
|
LON, LAT = np.meshgrid(tlon, tlat)
|
|
@@ -227,10 +196,10 @@ def main() -> None:
|
|
|
227
196
|
ctx.from_dataset("target", target, chunks={"cell": LON.size})
|
|
228
197
|
|
|
229
198
|
# 1) SQL reprojects the target grid into the source CRS (case 07's UDF).
|
|
230
|
-
reproj_sql = """
|
|
199
|
+
reproj_sql = f"""
|
|
231
200
|
SELECT dst_lat, dst_lon,
|
|
232
|
-
reproject(dst_lon, dst_lat)['x'] AS sx,
|
|
233
|
-
reproject(dst_lon, dst_lat)['y'] AS sy
|
|
201
|
+
reproject(dst_lon, dst_lat, '{_DST_CRS}', '{_SRC_CRS}')['x'] AS sx,
|
|
202
|
+
reproject(dst_lon, dst_lat, '{_DST_CRS}', '{_SRC_CRS}')['y'] AS sy
|
|
234
203
|
FROM target
|
|
235
204
|
"""
|
|
236
205
|
show_sql(reproj_sql, label="SQL — reproject target grid (PROJ UDF)")
|
|
@@ -21,7 +21,7 @@ plain-English definition of the operation, and computes the same numbers.
|
|
|
21
21
|
| 04 | `04_anomaly.py` | climatology broadcast-subtract | climatology CTE self-`JOIN` |
|
|
22
22
|
| 05 | `05_forecast_skill.py` | align valid/init/lead, reduce | forecast↔truth `JOIN` on `valid_time` + aggregate |
|
|
23
23
|
| 06 | `06_zonal_vector.py` | rasterize + mask per region | range `JOIN` raster↔regions |
|
|
24
|
-
| 07 | `07_reproject_udf.py` | per-pixel CRS transform | scalar **UDF** (`reproject()`), à la PostGIS `ST_Transform` |
|
|
24
|
+
| 07 | `07_reproject_udf.py` | per-pixel CRS transform | scalar **UDF** (`reproject()` from the geo extension), à la PostGIS `ST_Transform` |
|
|
25
25
|
| 08 | `08_regrid_weights.py` | interpolation to a new grid | sparse-weight table `JOIN` + weighted `GROUP BY` |
|
|
26
26
|
| 09 | `09_warp.py` | reproject **and** resample (warp) | reproject **UDF** (07) → weight table `JOIN` (08) |
|
|
27
27
|
|
|
@@ -212,12 +212,16 @@ paradigm. They split cleanly along one line: **is the operation row-independent?
|
|
|
212
212
|
|
|
213
213
|
**Reprojection is.** Moving a coordinate from one CRS to another depends only on
|
|
214
214
|
that coordinate, so it is a *scalar function* — exactly what PostGIS and
|
|
215
|
-
DuckDB-spatial already ship as `ST_Transform`.
|
|
216
|
-
|
|
217
|
-
|
|
215
|
+
DuckDB-spatial already ship as `ST_Transform`. xarray-sql ships it as an
|
|
216
|
+
optional geo extension (`pip install xarray-sql[geo]`): with pyproj
|
|
217
|
+
installed, every `XarrayContext` registers a PROJ-backed
|
|
218
|
+
`reproject(x, y, src_crs, dst_crs)` scalar UDF, so the CRS pair — any CRS
|
|
219
|
+
pyproj understands — is part of the query rather than baked into the function:
|
|
218
220
|
|
|
219
221
|
```sql
|
|
220
|
-
SELECT x, y,
|
|
222
|
+
SELECT x, y,
|
|
223
|
+
reproject(x, y, 'EPSG:32610', 'EPSG:4326')['x'] AS lon,
|
|
224
|
+
reproject(x, y, 'EPSG:32610', 'EPSG:4326')['y'] AS lat
|
|
221
225
|
FROM grid
|
|
222
226
|
```
|
|
223
227
|
|
|
@@ -226,8 +230,13 @@ this against **Earth Engine itself**: it opens a UTM grid through
|
|
|
226
230
|
[Xee](https://github.com/google/Xee) carrying `ee.Image.pixelLonLat()`, so EE's
|
|
227
231
|
own geodesy engine reports the true lon/lat of every pixel — an *independent*
|
|
228
232
|
reprojection reference, not PROJ-vs-PROJ. The SQL UDF and EE agree to sub-metre
|
|
229
|
-
precision.
|
|
230
|
-
|
|
233
|
+
precision. There is one practical gotcha — released pyproj mishandles threads
|
|
234
|
+
not created by Python, like DataFusion's runtime workers
|
|
235
|
+
([pyproj#1541](https://github.com/pyproj4/pyproj/pull/1541)), so the extension
|
|
236
|
+
runs all PROJ work on its own pool of Python threads, which also caches
|
|
237
|
+
transformers per thread,
|
|
238
|
+
keeping the UDF safe (and parallel) under DataFusion's concurrent partitions —
|
|
239
|
+
but the caveat that matters here is conceptual: reprojection
|
|
231
240
|
moves the coordinates without resampling the data onto a new grid — and *that* is
|
|
232
241
|
the next operation.
|
|
233
242
|
|
|
@@ -257,8 +266,9 @@ into bilinear weights, and the 08 `JOIN` applies them:
|
|
|
257
266
|
|
|
258
267
|
```sql
|
|
259
268
|
-- 1. reproject the target grid into source coordinates (the 07 UDF)
|
|
260
|
-
SELECT dst_lat, dst_lon,
|
|
261
|
-
|
|
269
|
+
SELECT dst_lat, dst_lon,
|
|
270
|
+
reproject(dst_lon, dst_lat, 'EPSG:4326', 'EPSG:32610')['x'] AS sx,
|
|
271
|
+
reproject(dst_lon, dst_lat, 'EPSG:4326', 'EPSG:32610')['y'] AS sy
|
|
262
272
|
FROM target
|
|
263
273
|
-- 2. apply the bilinear weights built from those points (the 08 JOIN)
|
|
264
274
|
SELECT w.dst_lat AS lat, w.dst_lon AS lon, SUM(s.value * w.weight) AS warped
|
|
@@ -145,6 +145,14 @@ class TestPartitionBounds:
|
|
|
145
145
|
assert tag == "int64"
|
|
146
146
|
assert lo < hi
|
|
147
147
|
|
|
148
|
+
def test_out_of_int64_range_returns_none(self):
|
|
149
|
+
# Year-1 gregorian dates exceed the int64 nanosecond range, so no
|
|
150
|
+
# pruning bound can be reported; the caller skips the dimension.
|
|
151
|
+
values = xr.date_range(
|
|
152
|
+
"0001-01-01", periods=3, freq="100YS", use_cftime=True
|
|
153
|
+
).values
|
|
154
|
+
assert cft.partition_bounds(values) is None
|
|
155
|
+
|
|
148
156
|
|
|
149
157
|
# -- Integration with _parse_schema ----------------------------------------
|
|
150
158
|
|
|
@@ -8,6 +8,7 @@ import xarray as xr
|
|
|
8
8
|
|
|
9
9
|
from xarray_sql.df import (
|
|
10
10
|
DEFAULT_BATCH_SIZE,
|
|
11
|
+
_ensure_default_indexes,
|
|
11
12
|
_parse_schema,
|
|
12
13
|
block_slices,
|
|
13
14
|
compute_chunks,
|
|
@@ -16,6 +17,7 @@ from xarray_sql.df import (
|
|
|
16
17
|
from_map,
|
|
17
18
|
from_map_batched,
|
|
18
19
|
iter_record_batches,
|
|
20
|
+
partition_metadata,
|
|
19
21
|
pivot,
|
|
20
22
|
)
|
|
21
23
|
from xarray_sql.reader import read_xarray, read_xarray_table
|
|
@@ -179,6 +181,30 @@ def test_iter_record_batches_matches_dataset_to_record_batch(air_small):
|
|
|
179
181
|
pd.testing.assert_frame_equal(actual_df, expected_df)
|
|
180
182
|
|
|
181
183
|
|
|
184
|
+
def test_iter_record_batches_projection_drops_cftime_dim():
|
|
185
|
+
"""A projection that drops a cftime dim (e.g. time under GROUP BY level)
|
|
186
|
+
must not call schema.field() for it. The dim is absent from the projected
|
|
187
|
+
schema, and cftime coords take the convert_for_field(schema.field(name))
|
|
188
|
+
path, so an unguarded lookup raised KeyError during batch reading."""
|
|
189
|
+
cftime = pytest.importorskip("cftime")
|
|
190
|
+
times = np.array(
|
|
191
|
+
[cftime.DatetimeGregorian(2020, m, 1) for m in (1, 2, 3)], dtype=object
|
|
192
|
+
)
|
|
193
|
+
ds = xr.Dataset(
|
|
194
|
+
{"air": (["time", "lat"], np.arange(3 * 2, dtype=float).reshape(3, 2))},
|
|
195
|
+
coords={"time": times, "lat": [0.0, 1.0]},
|
|
196
|
+
)
|
|
197
|
+
full = _parse_schema(ds)
|
|
198
|
+
projected = pa.schema(
|
|
199
|
+
[full.field("lat"), full.field("air")]
|
|
200
|
+
) # time dropped
|
|
201
|
+
table = pa.Table.from_batches(
|
|
202
|
+
list(iter_record_batches(ds, projected, batch_size=16)), projected
|
|
203
|
+
)
|
|
204
|
+
assert table.schema.names == ["lat", "air"]
|
|
205
|
+
assert table.num_rows == 6
|
|
206
|
+
|
|
207
|
+
|
|
182
208
|
def test_iter_record_batches_default_batch_size():
|
|
183
209
|
"""A single-batch partition (rows <= DEFAULT_BATCH_SIZE) yields exactly one batch."""
|
|
184
210
|
ds = xr.tutorial.open_dataset("air_temperature").isel(time=slice(0, 2))
|
|
@@ -532,3 +558,138 @@ def test_compute_chunks_tuples_sum_to_dim_size():
|
|
|
532
558
|
result = compute_chunks(ds, {"a": 3, "b": 4, "c": 5})
|
|
533
559
|
for dim, tup in result.items():
|
|
534
560
|
assert sum(tup) == ds.sizes[dim]
|
|
561
|
+
|
|
562
|
+
|
|
563
|
+
# -- Object-dtype and out-of-ns-range coordinate support --------------------
|
|
564
|
+
|
|
565
|
+
|
|
566
|
+
def _field_type(schema, name):
|
|
567
|
+
return schema.field(name).type
|
|
568
|
+
|
|
569
|
+
|
|
570
|
+
def test_parse_schema_maps_object_string_data_var_to_string():
|
|
571
|
+
# A string variable arrives as numpy object dtype; _parse_schema must not
|
|
572
|
+
# hand it to pa.from_numpy_dtype (which raises "Unsupported numpy type 17").
|
|
573
|
+
ds = xr.Dataset(
|
|
574
|
+
{"label": (["x"], np.array(["a", "b"], dtype=object))},
|
|
575
|
+
coords={"x": [1, 2]},
|
|
576
|
+
)
|
|
577
|
+
schema = _parse_schema(_ensure_default_indexes(ds))
|
|
578
|
+
assert _field_type(schema, "label") == pa.string()
|
|
579
|
+
|
|
580
|
+
|
|
581
|
+
def test_parse_schema_maps_object_string_coord_to_string():
|
|
582
|
+
# A string dimension coordinate (e.g. station names) is object dtype too.
|
|
583
|
+
ds = xr.Dataset(
|
|
584
|
+
{"v": (["station"], [1.0, 2.0])},
|
|
585
|
+
coords={"station": np.array(["A", "B"], dtype=object)},
|
|
586
|
+
)
|
|
587
|
+
schema = _parse_schema(_ensure_default_indexes(ds))
|
|
588
|
+
assert _field_type(schema, "station") == pa.string()
|
|
589
|
+
|
|
590
|
+
|
|
591
|
+
def test_partition_metadata_skips_out_of_ns_datetime():
|
|
592
|
+
# datetime64 coordinates outside the datetime64[ns] range (pre-1678 /
|
|
593
|
+
# post-2262) cannot be represented as int64 nanoseconds, so partition
|
|
594
|
+
# pruning must be skipped for that dimension rather than raising
|
|
595
|
+
# OverflowError. Registration must still succeed.
|
|
596
|
+
times = xr.date_range(
|
|
597
|
+
"0001-01-01", periods=3, freq="100YS", use_cftime=True
|
|
598
|
+
).to_datetimeindex(time_unit="us", unsafe=True)
|
|
599
|
+
ds = _ensure_default_indexes(
|
|
600
|
+
xr.Dataset({"v": (["time"], np.arange(3.0))}, coords={"time": times})
|
|
601
|
+
)
|
|
602
|
+
blocks = list(block_slices(ds, chunks={"time": 2}))
|
|
603
|
+
|
|
604
|
+
meta = partition_metadata(ds, blocks) # must not raise
|
|
605
|
+
|
|
606
|
+
assert len(meta) == len(blocks)
|
|
607
|
+
# "time" is unpruneable here, so it is omitted from every partition.
|
|
608
|
+
assert all("time" not in m for m in meta)
|
|
609
|
+
|
|
610
|
+
|
|
611
|
+
def test_parse_schema_all_null_object_var_stays_null():
|
|
612
|
+
# An all-null object column has no data to infer a type from; let null be
|
|
613
|
+
# null rather than coercing it to a string column.
|
|
614
|
+
ds = _ensure_default_indexes(
|
|
615
|
+
xr.Dataset(
|
|
616
|
+
{"label": (["x"], np.array([None, None], dtype=object))},
|
|
617
|
+
coords={"x": [1, 2]},
|
|
618
|
+
)
|
|
619
|
+
)
|
|
620
|
+
schema = _parse_schema(ds)
|
|
621
|
+
assert pa.types.is_null(schema.field("label").type)
|
|
622
|
+
|
|
623
|
+
|
|
624
|
+
def test_partition_metadata_prunes_cftime_coord():
|
|
625
|
+
# cftime dimension coordinates must produce pruning bounds; previously the
|
|
626
|
+
# object-dtype skip shadowed the cftime branch, silently disabling pruning.
|
|
627
|
+
times = xr.date_range(
|
|
628
|
+
"2000-01-01", periods=4, freq="1D", calendar="noleap", use_cftime=True
|
|
629
|
+
)
|
|
630
|
+
ds = _ensure_default_indexes(
|
|
631
|
+
xr.Dataset({"v": (["time"], np.arange(4.0))}, coords={"time": times})
|
|
632
|
+
)
|
|
633
|
+
blocks = list(block_slices(ds, chunks={"time": 2}))
|
|
634
|
+
|
|
635
|
+
meta = partition_metadata(ds, blocks)
|
|
636
|
+
|
|
637
|
+
assert all("time" in m for m in meta)
|
|
638
|
+
for m in meta:
|
|
639
|
+
_, _, tag = m["time"]
|
|
640
|
+
assert tag == "timestamp_ns"
|
|
641
|
+
|
|
642
|
+
|
|
643
|
+
def test_partition_metadata_skips_ancient_cftime():
|
|
644
|
+
# Ancient gregorian cftime dates overflow the int64 nanosecond range, so
|
|
645
|
+
# pruning must be skipped for that dim (no raise, dim omitted).
|
|
646
|
+
times = xr.date_range(
|
|
647
|
+
"0001-01-01", periods=3, freq="100YS", use_cftime=True
|
|
648
|
+
)
|
|
649
|
+
ds = _ensure_default_indexes(
|
|
650
|
+
xr.Dataset({"v": (["time"], np.arange(3.0))}, coords={"time": times})
|
|
651
|
+
)
|
|
652
|
+
blocks = list(block_slices(ds, chunks={"time": 2}))
|
|
653
|
+
|
|
654
|
+
meta = partition_metadata(ds, blocks) # must not raise
|
|
655
|
+
|
|
656
|
+
assert all("time" not in m for m in meta)
|
|
657
|
+
|
|
658
|
+
|
|
659
|
+
def test_string_dataset_round_trips_through_record_batch():
|
|
660
|
+
# The schema fix must also flow through the batch builders: a string
|
|
661
|
+
# column has to materialize as an Arrow string array, not error out.
|
|
662
|
+
ds = _ensure_default_indexes(
|
|
663
|
+
xr.Dataset(
|
|
664
|
+
{"label": (["x"], np.array(["a", "b", "c", "d"], dtype=object))},
|
|
665
|
+
coords={"x": [10, 20, 30, 40]},
|
|
666
|
+
)
|
|
667
|
+
)
|
|
668
|
+
schema = _parse_schema(ds)
|
|
669
|
+
|
|
670
|
+
batch = dataset_to_record_batch(ds, schema)
|
|
671
|
+
assert batch.schema.field("label").type == pa.string()
|
|
672
|
+
assert batch.column("label").to_pylist() == ["a", "b", "c", "d"]
|
|
673
|
+
|
|
674
|
+
# The streaming path must agree with the one-shot path.
|
|
675
|
+
streamed = pa.Table.from_batches(
|
|
676
|
+
list(iter_record_batches(ds, schema, batch_size=2)), schema=schema
|
|
677
|
+
)
|
|
678
|
+
assert streamed.column("label").to_pylist() == ["a", "b", "c", "d"]
|
|
679
|
+
|
|
680
|
+
|
|
681
|
+
def test_partition_metadata_in_range_datetime_still_pruned():
|
|
682
|
+
# Regression guard: ordinary datetimes must keep producing timestamp_ns
|
|
683
|
+
# bounds so filter pushdown still works after the overflow fix.
|
|
684
|
+
times = pd.date_range("2000-01-01", periods=4, freq="D")
|
|
685
|
+
ds = _ensure_default_indexes(
|
|
686
|
+
xr.Dataset({"v": (["time"], np.arange(4.0))}, coords={"time": times})
|
|
687
|
+
)
|
|
688
|
+
blocks = list(block_slices(ds, chunks={"time": 2}))
|
|
689
|
+
|
|
690
|
+
meta = partition_metadata(ds, blocks)
|
|
691
|
+
|
|
692
|
+
assert all("time" in m for m in meta)
|
|
693
|
+
for m in meta:
|
|
694
|
+
_, _, tag = m["time"]
|
|
695
|
+
assert tag == "timestamp_ns"
|