perf(netcdf): bound peak memory — stream from_xarray, container crop/to_crs/resample, and UGRID reads - #980
Merged
Merged
Conversation
… fully materialised from_xarray built the MEM container by calling np.asarray(var.values) on every variable up front, holding them all in RAM at once and forcing a full compute of a dask-backed input. Pass var.data (lazy) through and teach _build_multidim to write a dask variable block by block via a new _write_md_array_streamed helper -- each block is computed, written to its hyperslab, and released. NumPy inputs and the other _build_multidim caller (to_netcdf) keep the eager whole-array write unchanged. Closes #977.
The container fan-out for crop/to_crs/resample built an in-memory MEM container holding every transformed variable at once — peak memory scaled with the whole cube, not a single slab. Add an opt-in `path=` keyword to `crop`, `to_crs` and `resample`. On a root MDIM container, `path` streams every transformed variable straight to that .nc file one leading-dimension slab at a time via `open_streaming_multidim_netcdf`, so the whole result is never resident, and returns a file-backed NetCDF. `path=None` keeps today's in-memory behaviour unchanged. 3-D (band-dim, y, x) variables are written slab-by-slab; 2-D (y, x) variables and carried-through non-spatial auxiliary variables are written whole via a new `_StreamingMultidimWriter.write_whole`. Shapes the single-leading-dimension slab writer cannot represent — a variable with two or more band dimensions, a string/object auxiliary variable, or a curvilinear (2-D) coordinate grid — make `_stream_apply_to_file` return None and the caller falls back to the eager build-then-write path (still correct, just not bounded-memory for that rare shape). A single-variable (non-container) crop honours `path` by writing the raster and re-opening it file-backed.
MAfarrag
changed the title
Fixes the F821 ruff-check failure — the `path: str | Path | None` annotations resolve `Path` statically even under `from __future__ import annotations` — and applies ruff-format's line wrapping to the #976 changes.
…ream to_file (#982) UgridDataset deferred variable reads to first `.data` access, but then read the whole `(n_time, n_elements)` array where only a slice was needed, and duplicated the largest arrays with redundant copies. - `MeshVariable.sel_time` / `sel_time_range` and `UgridDataset.to_geodataframe` now read only the requested time slab straight from disk (`ReadAsArray(array_start_idx=, count=)`) when the array is not already cached and the time axis is axis 0, falling back to the full-load-then-slice path for a cached array, a negative/out-of-range index, or a non-axis-0 time dimension. A single-step selection no longer materialises (or caches) the whole temporal array. The source file is threaded onto each variable via a new `_source_path` field. - Dropped the redundant `.copy()` after `ReadAsArray()` in the lazy data loader, `Connectivity.from_gdal_array`, and `_write_connectivity_array`: `ReadAsArray` already returns a fresh array and the following `astype` copies again. - `to_file` streams each variable through a new `MeshVariable.load_array` that reads without populating the shared cache, so the write holds one variable at a time instead of the whole cube (and leaves the caller's variables lazy). No public API or behaviour change: the same values are returned/written, only the amount read into memory shrinks.
MAfarrag
changed the title
`open_streaming_multidim_netcdf` (via the module-level `_interop` alias) and `_read_attributes` are already imported at module top and there is no import cycle to break, so the inline imports violated the repo's no-inline-imports rule. (review L4)
…n't share a grid `_stream_apply_to_file` wrote every spatial variable against the first variable's y/x template grid. A container whose variables transform onto different grids (mixed native cell size / extent, or a reprojection yielding per-variable shapes) would mis-write a slab into a template-shaped MDArray and raise a raw GDAL error, where the eager fan-out handles it fine. Guard the shared-grid assumption: when any variable's transformed y/x shape differs from the template, return None so the caller falls back to the eager per-variable build. (review L1)
…er fallback When the streaming path returns None and _apply_to_all_variables recurses for the eager build, the recursive call re-ran the aux scan and warned about carried-through multi- dimensional variables a second time for one call. Suppress warnings on the recursive build since the outer call already emitted them. (review N1)
`_read_time_slab` left the per-selection reopened dataset to fall out of scope; close it in a finally so Windows does not keep a read handle on the file. The ReadAsArray result is a fresh numpy-owned array and stays valid after the dataset closes. (review N3)
…ount) `_can_window_time` required `stop` in [0, n_steps] but not `stop >= start`, so a file-backed `sel_time_range(2, 1)` took the windowed path and computed `count = -1`, which GDAL cannot service — raising a misleading 'no loaded data' error while the cached branch returned a consistent empty array. Exclude reversed/empty ranges so they fall back to the in-memory slice. (review L2)
Switching to `sel_time(0)` (#982) made `to_geodataframe` raise a `sel_time` 'no loaded data' error for a temporal variable that resolves to no data. Add a metadata-only `MeshVariable.has_data_source` and gate the temporal read on it (no forced load), and emit a length-correct null column when a variable has no data instead of raising — pandas rejects a scalar `None` column, so neither the old nor the new code produced a valid frame there. (review L3)
The streaming-op tests compared only array values, so a regression that dropped or garbled the reconstructed CRS/geotransform would pass. Extend _assert_same to check the container epsg and, for each genuinely spatial variable, its epsg and geotransform (relative tolerance, since the eager and streamed affines differ sub-pixel). Non-spatial bounds variables inherit the global CRS in the streamed file but not the eager path — a known harmless divergence, so they are excluded from the metadata comparison. (review L5)
…string Match the dominant docstring style and the sibling load_array docstring.
The reversed-range guard used `start <= stop`, still admitting the empty range `sel_time_range(k, k)` — which sent `count = 0` to `ReadAsArray` and raised `RuntimeError: count[0] = 0 is invalid` under gdal.UseExceptions(), while the cached path returned an empty array. Use a strict `start < stop` so empty ranges fall back to the in-memory empty slice. Adds empty-range and full-range (`stop == n_steps`) boundary tests. (review R2-M1)
…variable warning
The de-duplication wrapped the recursive eager rebuild in a blanket
`warnings.simplefilter('ignore')`, but that recursion is where the actual crop/reproject
work runs — so it masked *every* warning (e.g. real `invalid value encountered`
RuntimeWarnings) on the `path=` fallback route, not just the duplicate demoted-variable
warning. Thread a `warn_demoted=False` flag into the recursive call and gate only the
demoted-variable `warnings.warn` on it, so genuine transform warnings still surface.
(review R2-L1)
…_file `get_variable` builds a fresh classic-raster wrapper (a new GDAL dataset) per call, and the method re-fetched the same variable in the band-dim guard, the results comprehension, the var-specs loop, and the write loop — up to four opens each. Resolve them once into `spatial_var_objs` and reuse, cutting the transient GDAL handles on a wide container. (review R2-N1)
…mplexity SonarCloud S3776: the streaming builder had a cognitive complexity of 63. Extract the feasibility guard (_stream_feasible), the per-variable spatial and auxiliary spec builders (_add_spatial_var_spec / _add_aux_var_spec), the root-attribute builder (_stream_root_attrs), and the per-variable slab writer (_write_stream_variable), leaving _stream_apply_to_file a linear orchestration. Behaviour is unchanged (streaming tests green).
…ly_to_all_variables SonarCloud S3776: cognitive complexity 22. Extract the demoted-auxiliary-variable scan and warning into _warn_demoted_variables (stacklevel adjusted for the added frame) and the streaming-or-eager to-file routing into _to_file_via_stream_or_eager, leaving the eager fan-out loop untouched. Behaviour unchanged.
SonarCloud S3776: cognitive complexity 17. Move the data-variable body (unknown-dim check, dask-stream-vs-CF-encode branch, shape check, create + write + attrs) into a module-level _write_data_var, leaving _build_multidim's data-var loop a single call. Behaviour unchanged.
SonarCloud S3776: the earlier extraction left _apply_to_all_variables at complexity 16 (limit 15). Move the eager per-variable build loop into _fan_out_eager, so _apply_to_all_variables is a short guard-and-dispatch. Behaviour unchanged.
Dropping the `.copy()` (#982) left the loader returning GDAL's untyped `Any` from `ReadAsArray()`, which mypy rejects for a function declared to return `ndarray | None`. Assign through a typed local to coerce the type while keeping the no-copy read.
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Two write-side streaming changes to
NetCDFthat keep peak memory bounded to a single slab instead of the wholecube. Both are the write-side mirror of the lazy
to_xarray(chunks=)reader added in #864.from_xarraystreams a dask-backed input (#977)NetCDF.from_xarraybuilt the in-memory MEM container by callingnp.asarray(var.values)on every datavariable up front: all variables were held in RAM simultaneously, and a dask-backed input was fully
.compute()-dbefore a single byte was written.
Now
from_xarraypassesvar.data(the underlying array — lazy for a dask-backed variable) through to_build_multidim, and_build_multidimwrites a dask variable block by block via a new_write_md_array_streamedhelper: each block is computed, written to its hyperslab(
md_arr.Write(block, array_start_idx=…, count=…)), and released before the next. A lazily-loaded xarray variabletherefore never becomes fully resident, and multiple variables are no longer all materialised into a dict at once.
Backward-compatible:
_build_multidimcaller — the GDAL-nativewrite_multidim_netcdf/DatasetCollection.to_netcdfpath — passes NumPy arrays, so it is unchanged.
also takes the eager path.
crop/to_crs/resamplecan stream a container to a file (#976)The root-container fan-out for these three ops applied the op to every spatial variable and built an in-memory MEM
container holding all transformed variables at once, so peak memory scaled with the whole cube rather than a
single slab.
These methods now take an opt-in
path=keyword. On a root MDIM container,pathstreams every transformedvariable straight to that
.ncfile one leading-dimension slab at a time viaopen_streaming_multidim_netcdf(sothe whole result is never resident) and returns a file-backed
NetCDF.path=None(default) keeps today'sin-memory fan-out unchanged.
(band-dim, y, x)variables are written slab-by-slab; 2-D(y, x)variables and carried-throughnon-spatial auxiliary variables are written whole via a new
_StreamingMultidimWriter.write_whole.a string/object auxiliary variable, or a curvilinear (2-D) coordinate grid — make
_stream_apply_to_filereturnNone, and the caller falls back to the eager build-then-write path (still correct, just not bounded-memory forthat rare shape).
pathby writing the raster and re-opening it file-backed.No dependencies added.
UGRID eager reads: windowed time selection, dropped copies, streamed write (#982)
The same audit applied to
UgridDataset. Variable reads were already deferred to first.dataaccess, butseveral methods then read the whole
(n_time, n_elements)array where only a slice was needed, and a couple ofhot paths duplicated the largest arrays.
MeshVariable.sel_time/sel_time_rangeandto_geodataframenow read only the requested time slabstraight from disk (
ReadAsArray(array_start_idx=…, count=…)) when the array isn't already cached and the timeaxis is axis 0, falling back to the full-load-then-slice path for a cached array, a negative/out-of-range index,
or a non-axis-0 time dimension. A single-step selection no longer materialises (or caches) the whole temporal
cube. The source file is threaded onto each variable via a new
_source_pathfield..copy()afterReadAsArray()in the lazy data loader,Connectivity.from_gdal_array,and
_write_connectivity_array—ReadAsArrayalready returns a fresh array and the followingastypecopiesagain.
to_filestreams each variable through a newMeshVariable.load_arraythat reads without populating the sharedcache, so the write holds one variable at a time instead of the whole cube (and leaves the caller's variables
lazy).
to_crswas verified already lazy (it passes the variable dict through and only transforms coordinates).No dependencies added.
Issues
from_xarrayhalf of ARC-48 (theopen_mfdataset+to_xarray(chunks=)halvesshipped in refactor(netcdf): resolve arc-netcdf architecture-review findings #816 / feat(netcdf): lazy to_xarray(chunks=) and shared NETCDF: subdataset-prefix helper #864)
crop/to_crs/resampleUgridDataset(windowed time selection, redundant copies, writefan-out)
Type of change
crop/to_crs/resamplegain an opt-inpath=parameter (defaulting to the previous in-memory behaviour);from_xarray's signature and output are unchanged. No existing behaviour changes whenpathis omitted.How Has This Been Tested?
tests/netcdf/lazy/test_netcdf_interop.py::TestFromXarrayStreaming— a dask-backed input round-trips andagrees with the equivalent NumPy input;
_write_md_array_streamedissues one hyperslab write per dask blockand the blocks reconstruct the array; a NumPy array is written in a single whole write.
tests/netcdf/samples/test_stream_apply_to_file.py(6 tests) —crop/to_crs/resamplewithpath=write a file-backed container value-identical to the in-memory result; the 3-D variable is written one band
slab at a time (not a single whole write); a 4-D variable falls back to the eager path and still matches; a
single-variable crop honours
path.tests/ugrid/test_lazy_data_variables.py(5 new tests) —sel_time/sel_time_rangeissue exactly onewindowed slab read and match the full-array slice; a cached variable slices in memory; a negative index
falls back to a full read;
to_fileleaves the source variables uncached and the output round-trips.netcdfnon-plot suite: 2378 passed, 38 skipped; fullugridsuite: 258 passed.Reproduce:
pixi run -e dev pytest tests/netcdf/lazy/test_netcdf_interop.py tests/netcdf/samples/test_stream_apply_to_file.pyChecklist:
Versions and the change log are generated by commitizen in CI, so those three boxes are intentionally left for the
release automation.