Skip to content

Commit 32e956a

Browse files
committed
Introduced new test_read_netcdf.py function for
directly converting netCDF files into XarrayContext.
1 parent 9609a7d commit 32e956a

3 files changed

Lines changed: 62 additions & 1 deletion

File tree

test_read_netcdf.py

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,11 @@
1+
from metpy.cbook import get_test_data
2+
from xarray_sql import XarrayContext
3+
4+
path = get_test_data('irma_gfs_example.nc', False)
5+
print(type(path), path)
6+
7+
ctx = XarrayContext.read_netcdf(get_test_data('irma_gfs_example.nc', False))
8+
print(ctx._registered_datasets.keys())
9+
10+
result = ctx.sql("SELECT * FROM irma_gfs_example.time1_isobaric1_latitude_longitude LIMIT 5")
11+
print(result)

xarray_sql/_native.pyd

58.2 MB
Binary file not shown.

xarray_sql/sql.py

Lines changed: 51 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,8 @@
77
from .df import Chunks
88
from .ds import XarrayDataFrame
99
from .reader import read_xarray_table
10-
10+
from pathlib import Path
11+
from matplotlib.dates import date2num
1112

1213
class XarrayContext(SessionContext):
1314
"""A datafusion `SessionContext` that also supports `xarray.Dataset`s."""
@@ -156,6 +157,55 @@ def _maybe_register_cftime_udf(self, ds: xr.Dataset) -> None:
156157
self.register_udf(cft.make_cftime_udf(units, cal))
157158
break # One UDF per context is enough.
158159

160+
@classmethod
161+
def read_netcdf(
162+
cls,
163+
path: str | Path,
164+
table_name: str | None = None,
165+
chunks: dict | None = None,
166+
engine: str | None = "netcdf4",
167+
**open_kwargs,
168+
):
169+
"""
170+
Open a NetCDF file and register it in an XarrayContext.
171+
172+
Parameters
173+
----------
174+
path : str | Path
175+
Path to the .nc file.
176+
table_name : str | None
177+
SQL table (or schema, for mixed-dimension datasets) name.
178+
Defaults to the file stem (e.g. 'irma_gfs_example').
179+
chunks : dict | None
180+
Dask chunks. Defaults to {'time': 24} if a 'time' dim exists,
181+
otherwise opens all dims as a single chunk.
182+
engine : str | None
183+
xarray backend engine. Default is 'netcdf4'.
184+
**open_kwargs
185+
Additional kwargs forwarded to xr.open_dataset.
186+
187+
Returns
188+
-------
189+
XarrayContext
190+
A context with the dataset already registered.
191+
"""
192+
path = Path(path)
193+
194+
ds = xr.open_dataset(path, engine=engine, **open_kwargs)
195+
196+
if not isinstance(ds, xr.Dataset):
197+
raise TypeError(f"Expected xr.Dataset, got {type(ds).__name__}")
198+
199+
if table_name is None:
200+
table_name = path.stem
201+
202+
if chunks is None:
203+
chunks = {"time": 24} if "time" in ds.dims else {d: -1 for d in ds.dims}
204+
205+
ctx = cls()
206+
ctx.from_dataset(table_name, ds, chunks=chunks)
207+
return ctx
208+
159209
def sql(self, query: str, *args, **kwargs) -> XarrayDataFrame:
160210
"""Run a SQL query, returning an :class:`XarrayDataFrame` wrapper.
161211

0 commit comments

Comments
 (0)