Cache sequential native preparation reads on the xarray handle

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.

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:

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])

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)