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.
Files changed (48) hide show
  1. {xarray_sql-0.3.2 → xarray_sql-0.3.3}/Cargo.lock +1 -1
  2. {xarray_sql-0.3.2 → xarray_sql-0.3.3}/Cargo.toml +1 -1
  3. {xarray_sql-0.3.2 → xarray_sql-0.3.3}/PKG-INFO +11 -3
  4. {xarray_sql-0.3.2 → xarray_sql-0.3.3}/README.md +7 -2
  5. {xarray_sql-0.3.2 → xarray_sql-0.3.3}/benchmarks/geospatial/07_reproject_udf.py +20 -59
  6. {xarray_sql-0.3.2 → xarray_sql-0.3.3}/benchmarks/geospatial/09_warp.py +6 -37
  7. {xarray_sql-0.3.2 → xarray_sql-0.3.3}/benchmarks/geospatial/README.md +1 -1
  8. {xarray_sql-0.3.2 → xarray_sql-0.3.3}/docs/geospatial.md +18 -8
  9. {xarray_sql-0.3.2 → xarray_sql-0.3.3}/pyproject.toml +4 -0
  10. {xarray_sql-0.3.2 → xarray_sql-0.3.3}/tests/test_cft.py +8 -0
  11. {xarray_sql-0.3.2 → xarray_sql-0.3.3}/tests/test_df.py +161 -0
  12. xarray_sql-0.3.3/tests/test_proj.py +158 -0
  13. {xarray_sql-0.3.2 → xarray_sql-0.3.3}/tests/test_to_dataset_perf.py +18 -9
  14. {xarray_sql-0.3.2 → xarray_sql-0.3.3}/xarray_sql/cftime.py +11 -2
  15. {xarray_sql-0.3.2 → xarray_sql-0.3.3}/xarray_sql/df.py +60 -15
  16. xarray_sql-0.3.3/xarray_sql/proj.py +235 -0
  17. {xarray_sql-0.3.2 → xarray_sql-0.3.3}/xarray_sql/sql.py +13 -0
  18. {xarray_sql-0.3.2 → xarray_sql-0.3.3}/.gitignore +0 -0
  19. {xarray_sql-0.3.2 → xarray_sql-0.3.3}/AGENTS.md +0 -0
  20. {xarray_sql-0.3.2 → xarray_sql-0.3.3}/LICENSE +0 -0
  21. {xarray_sql-0.3.2 → xarray_sql-0.3.3}/benchmarks/geospatial/01_ndvi.py +0 -0
  22. {xarray_sql-0.3.2 → xarray_sql-0.3.3}/benchmarks/geospatial/02_climatology.py +0 -0
  23. {xarray_sql-0.3.2 → xarray_sql-0.3.3}/benchmarks/geospatial/03_zonal_mean.py +0 -0
  24. {xarray_sql-0.3.2 → xarray_sql-0.3.3}/benchmarks/geospatial/04_anomaly.py +0 -0
  25. {xarray_sql-0.3.2 → xarray_sql-0.3.3}/benchmarks/geospatial/05_forecast_skill.py +0 -0
  26. {xarray_sql-0.3.2 → xarray_sql-0.3.3}/benchmarks/geospatial/06_zonal_vector.py +0 -0
  27. {xarray_sql-0.3.2 → xarray_sql-0.3.3}/benchmarks/geospatial/08_regrid_weights.py +0 -0
  28. {xarray_sql-0.3.2 → xarray_sql-0.3.3}/benchmarks/geospatial/_harness.py +0 -0
  29. {xarray_sql-0.3.2 → xarray_sql-0.3.3}/benchmarks/geospatial/perf_summary.py +0 -0
  30. {xarray_sql-0.3.2 → xarray_sql-0.3.3}/benchmarks/geospatial/run_all.sh +0 -0
  31. {xarray_sql-0.3.2 → xarray_sql-0.3.3}/benchmarks/geospatial/run_perf.sh +0 -0
  32. {xarray_sql-0.3.2 → xarray_sql-0.3.3}/docs/assets/logo.svg +0 -0
  33. {xarray_sql-0.3.2 → xarray_sql-0.3.3}/docs/contributing.md +0 -0
  34. {xarray_sql-0.3.2 → xarray_sql-0.3.3}/docs/examples.md +0 -0
  35. {xarray_sql-0.3.2 → xarray_sql-0.3.3}/docs/index.md +0 -0
  36. {xarray_sql-0.3.2 → xarray_sql-0.3.3}/docs/reference/xarray_sql.md +0 -0
  37. {xarray_sql-0.3.2 → xarray_sql-0.3.3}/src/lib.rs +0 -0
  38. {xarray_sql-0.3.2 → xarray_sql-0.3.3}/tests/__init__.py +0 -0
  39. {xarray_sql-0.3.2 → xarray_sql-0.3.3}/tests/conftest.py +0 -0
  40. {xarray_sql-0.3.2 → xarray_sql-0.3.3}/tests/test_ds.py +0 -0
  41. {xarray_sql-0.3.2 → xarray_sql-0.3.3}/tests/test_reader.py +0 -0
  42. {xarray_sql-0.3.2 → xarray_sql-0.3.3}/tests/test_sql.py +0 -0
  43. {xarray_sql-0.3.2 → xarray_sql-0.3.3}/tests/test_stats.py +0 -0
  44. {xarray_sql-0.3.2 → xarray_sql-0.3.3}/xarray_sql/__init__.py +0 -0
  45. {xarray_sql-0.3.2 → xarray_sql-0.3.3}/xarray_sql/core.py +0 -0
  46. {xarray_sql-0.3.2 → xarray_sql-0.3.3}/xarray_sql/ds.py +0 -0
  47. {xarray_sql-0.3.2 → xarray_sql-0.3.3}/xarray_sql/reader.py +0 -0
  48. {xarray_sql-0.3.2 → xarray_sql-0.3.3}/zensical.toml +0 -0
@@ -3367,7 +3367,7 @@ checksum = "ea2f10b9bb0928dfb1b42b65e1f9e36f7f54dbdf08457afefb38afcdec4fa2bb"
3367
3367
 
3368
3368
  [[package]]
3369
3369
  name = "xarray_sql"
3370
- version = "0.3.2"
3370
+ version = "0.3.3"
3371
3371
  dependencies = [
3372
3372
  "arrow",
3373
3373
  "async-stream",
@@ -1,6 +1,6 @@
1
1
  [package]
2
2
  name = "xarray_sql"
3
- version = "0.3.2"
3
+ version = "0.3.3"
4
4
  authors = ["Alex Merose"]
5
5
  edition = "2021"
6
6
  exclude = [
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: xarray_sql
3
- Version: 0.3.2
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
+ ![PyPI Version](https://img.shields.io/pypi/v/xarray-sql?color=green)
48
52
  [![ci](https://github.com/alxmrs/xarray-sql/actions/workflows/ci.yml/badge.svg)](https://github.com/alxmrs/xarray-sql/actions/workflows/ci.yml)
49
53
  [![lint](https://github.com/alxmrs/xarray-sql/actions/workflows/lint.yml/badge.svg)](https://github.com/alxmrs/xarray-sql/actions/workflows/lint.yml)
50
54
  [![ci-build](https://github.com/alxmrs/xarray-sql/actions/workflows/ci-build.yml/badge.svg)](https://github.com/alxmrs/xarray-sql/actions/workflows/ci-build.yml)
51
55
  [![ci-rust](https://github.com/alxmrs/xarray-sql/actions/workflows/ci-rust.yml/badge.svg)](https://github.com/alxmrs/xarray-sql/actions/workflows/ci-rust.yml)
56
+ [![PyPI Downloads](https://static.pepy.tech/personalized-badge/xarray-sql?period=total&units=INTERNATIONAL_SYSTEM&left_color=BLACK&right_color=GREEN&left_text=downloads)](https://pepy.tech/projects/xarray-sql)
57
+ [![PyPI Downloads](https://static.pepy.tech/personalized-badge/xarray-sql?period=monthly&units=INTERNATIONAL_SYSTEM&left_color=BLACK&right_color=GREEN&left_text=downloads%2Fmonth)](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 scalar PROJ UDF (validated against Earth
257
- Engine's own geodesy via [Xee](https://github.com/google/Xee)) and a
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
+ ![PyPI Version](https://img.shields.io/pypi/v/xarray-sql?color=green)
5
6
  [![ci](https://github.com/alxmrs/xarray-sql/actions/workflows/ci.yml/badge.svg)](https://github.com/alxmrs/xarray-sql/actions/workflows/ci.yml)
6
7
  [![lint](https://github.com/alxmrs/xarray-sql/actions/workflows/lint.yml/badge.svg)](https://github.com/alxmrs/xarray-sql/actions/workflows/lint.yml)
7
8
  [![ci-build](https://github.com/alxmrs/xarray-sql/actions/workflows/ci-build.yml/badge.svg)](https://github.com/alxmrs/xarray-sql/actions/workflows/ci-build.yml)
8
9
  [![ci-rust](https://github.com/alxmrs/xarray-sql/actions/workflows/ci-rust.yml/badge.svg)](https://github.com/alxmrs/xarray-sql/actions/workflows/ci-rust.yml)
10
+ [![PyPI Downloads](https://static.pepy.tech/personalized-badge/xarray-sql?period=total&units=INTERNATIONAL_SYSTEM&left_color=BLACK&right_color=GREEN&left_text=downloads)](https://pepy.tech/projects/xarray-sql)
11
+ [![PyPI Downloads](https://static.pepy.tech/personalized-badge/xarray-sql?period=monthly&units=INTERNATIONAL_SYSTEM&left_color=BLACK&right_color=GREEN&left_text=downloads%2Fmonth)](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 scalar PROJ UDF (validated against Earth
214
- Engine's own geodesy via [Xee](https://github.com/google/Xee)) and a
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
- So we register a PROJ-backed scalar UDF and reproject in SQL::
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, reproject(x, y)['lon'] AS lon, reproject(x, y)['lat'] AS lat
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
- PROJ's context is not thread-safe and DataFusion evaluates projection
43
- expressions concurrently, so we return *both* coordinates from one
44
- struct-returning UDF and keep the source in a single chunk (one serial UDF).
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
- # Single chunk single partition → serial UDF (PROJ is not thread-safe).
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)['lon'] AS lon,
166
- reproject(x, y)['lat'] AS lat
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`. We register a PROJ-backed scalar
216
- UDF (mirroring the `cftime()` UDF already in `xarray_sql/cftime.py`) and
217
- reproject in SQL:
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, reproject(x, y)['lon'] AS lon, reproject(x, y)['lat'] AS lat
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. The script flags one practical gotcha (PROJ is not thread-safe, so the
230
- UDF runs serially), but the caveat that matters here is conceptual: reprojection
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, reproject(dst_lon, dst_lat)['x'] AS sx,
261
- reproject(dst_lon, dst_lat)['y'] AS sy
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
@@ -36,8 +36,12 @@ dependencies = [
36
36
  ]
37
37
 
38
38
  [project.optional-dependencies]
39
+ geo = [
40
+ "pyproj",
41
+ ]
39
42
  test = [
40
43
  "cftime",
44
+ "xarray-sql[geo]",
41
45
  "pytest",
42
46
  "xarray[io]",
43
47
  "gcsfs",
@@ -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"