|
7 | 7 | from .df import Chunks |
8 | 8 | from .ds import XarrayDataFrame |
9 | 9 | from .reader import read_xarray_table |
10 | | - |
| 10 | +from pathlib import Path |
| 11 | +from matplotlib.dates import date2num |
11 | 12 |
|
12 | 13 | class XarrayContext(SessionContext): |
13 | 14 | """A datafusion `SessionContext` that also supports `xarray.Dataset`s.""" |
@@ -156,6 +157,55 @@ def _maybe_register_cftime_udf(self, ds: xr.Dataset) -> None: |
156 | 157 | self.register_udf(cft.make_cftime_udf(units, cal)) |
157 | 158 | break # One UDF per context is enough. |
158 | 159 |
|
| 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 | + |
159 | 209 | def sql(self, query: str, *args, **kwargs) -> XarrayDataFrame: |
160 | 210 | """Run a SQL query, returning an :class:`XarrayDataFrame` wrapper. |
161 | 211 |
|
|
0 commit comments