From 4e6693ee91c286f0c8193dad921b21faa4cf7bd4 Mon Sep 17 00:00:00 2001 From: Yi Yu <72683152+yuyi13@users.noreply.github.com> Date: Sun, 18 Oct 2026 04:31:48 +1100 Subject: [PATCH] Cache sequential native preparation reads on the xarray handle --- docs/output-netcdf-layout.md | 3 ++ pysweb/io/netcdf_stream.py | 34 ++++++++++++++++++-- pysweb/swb/preprocess.py | 3 +- tests/io/test_netcdf_stream.py | 59 ++++++++++++++++++++++++++++++++-- 4 files changed, 93 insertions(+), 6 deletions(-) diff --git a/docs/output-netcdf-layout.md b/docs/output-netcdf-layout.md index ec5fc2b..9d2a5fe 100644 --- a/docs/output-netcdf-layout.md +++ b/docs/output-netcdf-layout.md @@ -28,5 +28,8 @@ blocks used to construct the control. Older source files with large temporal chunks remain readable. Streaming copies temporarily increase the current variable's read cache to retain chunk slabs, capped at 512 MiB per variable, then restore its original cache settings. +Daily preprocessing uses the same bounded cache on the NetCDF handle read by +xarray, while retaining xarray's CF decoding and the existing daily reprojection. +Only the selected component's cache changes; global NetCDF defaults stay intact. These storage choices preserve values, masks, coordinates, units and scientific provenance; they do not change model or forcing parameters. diff --git a/pysweb/io/netcdf_stream.py b/pysweb/io/netcdf_stream.py index 12e98c8..3c228cc 100644 --- a/pysweb/io/netcdf_stream.py +++ b/pysweb/io/netcdf_stream.py @@ -5,15 +5,43 @@ Created: 2026-10-17 Last updated: 2026-10-17 Inputs: Variable dimensions, shapes and existing NetCDF chunk metadata. -Outputs: Bounded chunk layouts and temporary per-variable read-cache settings. +Outputs: Bounded chunk layouts, temporary per-variable caches and decoded sequential readers. Usage: from pysweb.io.netcdf_stream import stream_chunks, sequential_read_cache -Dependencies: contextlib, math, numpy, netCDF4-compatible variables. +Dependencies: contextlib, math, numpy, netCDF4, xarray. """ -from contextlib import contextmanager +from contextlib import contextmanager, ExitStack from math import ceil, prod import numpy as np +@contextmanager +def open_stream_dataset(path, variables, *, max_cache_bytes=512*1024**2, **kwargs): + """Decode through xarray while caching sequential reads on its actual handle. + + Each selected variable receives a bounded cache. Global NetCDF settings and + other open datasets are untouched; values and CF decoding stay with xarray. + Callers must load retained data before leaving this context. + """ + import netCDF4 as nc + import xarray as xr + + source = nc.Dataset(path) + dataset = None + try: + dataset = xr.open_dataset(xr.backends.NetCDF4DataStore(source), **kwargs) + with ExitStack() as caches: + for name in dict.fromkeys(variables): + caches.enter_context(sequential_read_cache(source[name], + max_cache_bytes=max_cache_bytes)) + yield dataset + finally: + if source.isopen(): + if dataset is not None: + dataset.close() + else: + source.close() + + def stream_chunks(dimensions, shape, *, block_dim=None, block_size=32, spatial_size=64): """Keep compressed chunks within the slabs written by streaming callers.""" if block_size < 1 or spatial_size < 1: diff --git a/pysweb/swb/preprocess.py b/pysweb/swb/preprocess.py index 071d111..ded0420 100644 --- a/pysweb/swb/preprocess.py +++ b/pysweb/swb/preprocess.py @@ -44,6 +44,7 @@ from rasterio.transform import from_origin from pysweb.contracts import daily_data, month_bounds from pysweb.io.provenance import atomic_netcdf, input_manifest, atomic_json, sha256 +from pysweb.io.netcdf_stream import open_stream_dataset GSSM_ASSET = "users/qianrswaterr/GlobalSSM1km0509" GSSM_SCALE_FACTOR = 1000.0 GSSM_EXPORT_SCALE_M = 1000.0 @@ -599,7 +600,7 @@ def _prepare_et_file_component( ) -> xr.DataArray: expected_dates = pd.date_range(start = start_date, end = end_date, freq = "D") layers: List[xr.DataArray] = [] - with xr.open_dataset(path, decode_coords="all") as ds: + with open_stream_dataset(path, [var_name], decode_coords="all") as ds: if var_name not in ds: raise KeyError(f"Variable '{var_name}' not found in ET dataset: {path}") da = spatial(ds[var_name]) diff --git a/tests/io/test_netcdf_stream.py b/tests/io/test_netcdf_stream.py index 9fc416f..b4e5d45 100644 --- a/tests/io/test_netcdf_stream.py +++ b/tests/io/test_netcdf_stream.py @@ -7,15 +7,70 @@ Inputs: Small numeric NetCDF fixtures and unpopulated long-calendar variable metadata. Outputs: Streaming memory/layout and value-preservation assertions. Usage: python -m pytest tests/io/test_netcdf_stream.py -Dependencies: netCDF4, numpy, pytest, pysweb. +Dependencies: netCDF4, numpy, pytest, xarray, pysweb. """ import netCDF4 as nc import numpy as np import pytest -from pysweb.io.netcdf_stream import stream_chunks, sequential_read_cache +import xarray as xr +from pysweb.io.netcdf_stream import stream_chunks, sequential_read_cache, open_stream_dataset from pysweb.swb.paired import _copy_variable +@pytest.mark.parametrize('format', ['NETCDF3_CLASSIC', 'NETCDF4']) +@pytest.mark.parametrize('fail', [False, True]) +def test_stream_dataset_preserves_cf_decoding_and_closes_handle(tmp_path, monkeypatch, format, fail): + path = tmp_path/'packed.nc' + with nc.Dataset(path, 'w', format=format) as source: + source.createDimension('time', 3) + dates = source.createVariable('time', 'i4', ('time',)) + dates.units = 'days since 2014-01-01' + dates[:] = [0, 1, 2] + variable = source.createVariable('ET', 'i2', ('time',), fill_value=-9999) + variable.scale_factor, variable.add_offset = 0.1, 1. + variable.units = 'mm day-1' + variable[:] = np.ma.array([1., 2., 3.], mask=[False, True, False]) + with xr.open_dataset(path) as source: + expected = source.load() + handles = [] + original = xr.backends.NetCDF4DataStore + def record_handle(source): + handles.append(source) + return original(source) + monkeypatch.setattr(xr.backends, 'NetCDF4DataStore', record_handle) + global_cache = nc.get_chunk_cache() + try: + with open_stream_dataset(path, ['ET']) as actual: + xr.testing.assert_identical(actual.load(), expected) + if fail: + raise RuntimeError('consumer failed') + except RuntimeError as error: + assert fail and str(error) == 'consumer failed' + assert len(handles) == 1 and not handles[0].isopen() + assert nc.get_chunk_cache() == global_cache + + +def test_stream_dataset_changes_only_selected_variable_cache(tmp_path, monkeypatch): + path = tmp_path/'metadata.nc' + with nc.Dataset(path, 'w') as source: + for dim, size in [('time', 6000), ('y', 128), ('x', 128)]: + source.createDimension(dim, size) + for name in ['ET', 'unread']: + source.createVariable(name, 'f4', ('time', 'y', 'x'), + zlib=True, chunksizes=(1500, 32, 32)) + handles = [] + original = xr.backends.NetCDF4DataStore + def record_handle(source): + handles.append(source) + return original(source) + monkeypatch.setattr(xr.backends, 'NetCDF4DataStore', record_handle) + with open_stream_dataset(path, ['ET'], max_cache_bytes=128*1024**2) as dataset: + assert dataset.sizes['time'] == 6000 + assert handles[0]['ET'].get_var_chunk_cache()[0] == 128*1024**2 + assert handles[0]['unread'].get_var_chunk_cache()[0] == nc.get_chunk_cache()[0] + assert not handles[0].isopen() + + def test_long_calendar_chunks_follow_temporal_and_row_write_blocks(): assert stream_chunks(('time', 'y', 'x'), (5844, 167, 167)) == (32, 64, 64) assert stream_chunks(('scene_time', 'y', 'x'), (577, 202, 221)) == (32, 64, 64)