From fa66500b51e85feed4e557c4e8a64952eb4715c4 Mon Sep 17 00:00:00 2001 From: Ben Smith Date: Sat, 5 Sep 2026 18:35:03 +0000 Subject: [PATCH] Read remote netCDF4 rasters chunk-wise through h5py from_nc() on an s3:// URI downloaded the whole object to return one tile: netCDF4.Dataset takes a filename or an in-memory buffer but never a file object, so a 60 km window out of a 2.9 GiB ATL15 granule cost 2.9 GiB. These files are HDF5 underneath and h5py does take a file object, so the remote case now goes through h5py and range requests, and a windowed read costs the chunks the window touches. grid/nc_h5.py presents an h5py.File through the netCDF4 API from_nc() uses, so the body of from_nc() is unchanged. Four differences between the libraries live there: * netCDF4 applies scale_factor/add_offset unless set_auto_scale(False) is called, and from_nc() only turns off masking, so the adapter replicates that arithmetic -- otherwise packed variables silently change value. * h5py rejects the negative-step slices select_slices() emits for a descending y axis, and broadcasts multiple index arrays. * netCDF dimensions that are not netCDF variables, and the HDF5 bookkeeping attributes netCDF4 hides, have to stay hidden. * netCDF4 lets a scalar variable be indexed with [:]. nc_open() is now resolve-source / open rather than a remote x compression matrix, and takes engine ('auto', 'h5py', 'netcdf4'), block_size and rdcc_nbytes. A file h5py cannot open (NETCDF3/classic) falls back to netCDF4. Local uncompressed reads take exactly the path they did before. External whole-file compression still decompresses whole: a gzip or bzip2 stream carries no index, so there is nothing to read chunk-wise. Block size is the difference between a win and a loss, not a tuning knob: fsspec caches one block, so scattered chunk reads at the 5 MiB default re-fetch enough to lose to downloading the file outright. Measured on a 29 MiB gzip-chunked granule over https, a 501x501 window fetches 91.1 MiB at 5 MiB blocks, 7.4 MiB at the 256 KiB default, 3.9 MiB at 64 KiB. io_utils.open_remote() passes block_size per file rather than to the session, so it reaches an earthaccess DAAC session too, whose constructor takes no such argument. h5_open(), from_h5() and io_utils.open_h5() gained the same parameter. Also fixes a fill-value bug this surfaced: from_nc() compared scaled data against a raw _FillValue, so a packed variable's fills were never converted. On a real GOES granule that left 91% of a window reading as 179.9976 K instead of NaN. Unpacked variables are unaffected. Tests are network-free -- engine='h5py' reads a local file through the same reader, and a stand-in filesystem maps s3:// URIs onto local files -- and assert the two readers return identical arrays under bounds, bands, skip, group, meta_only and t_axis, on to_nc() fixtures and on a netCDF4-written packed fixture. tests/test_nc_h5_remote.py measures bytes actually transferred against real granules; it is opt-in (PC_TEST_REMOTE for the https test, PC_TEST_S3 for the NSIDC ones, which need us-west-2). Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_016WPhgmqfUv5vYpBfnm592f --- pointCollection/grid/data.py | 153 ++++++++--- pointCollection/grid/nc_h5.py | 473 ++++++++++++++++++++++++++++++++++ pointCollection/io_utils.py | 49 +++- tests/test_nc_h5.py | 467 +++++++++++++++++++++++++++++++++ tests/test_nc_h5_remote.py | 181 +++++++++++++ 5 files changed, 1287 insertions(+), 36 deletions(-) create mode 100644 pointCollection/grid/nc_h5.py create mode 100644 tests/test_nc_h5.py create mode 100644 tests/test_nc_h5_remote.py diff --git a/pointCollection/grid/data.py b/pointCollection/grid/data.py index f7a3398..7fb0540 100755 --- a/pointCollection/grid/data.py +++ b/pointCollection/grid/data.py @@ -685,7 +685,8 @@ def from_gdal(self, ds, field='z', bands=None, bounds=None, extent=None, self.__update_size_and_shape__() return self - def h5_open(self, h5_file, mode='r', compression=None, fs=None): + def h5_open(self, h5_file, mode='r', compression=None, fs=None, + block_size=None): """ Open an HDF5 file with or without external compression. @@ -707,15 +708,19 @@ def h5_open(self, h5_file, mode='r', compression=None, fs=None): are ancillary data we own, not DAAC holdings, so an earthaccess session would not grant access to them. Pass an explicit fs to read a raster out of a DAAC bucket. + block_size: int or NoneType, default None + Bytes fetched per range request if h5_file is remote. None leaves + the filesystem's own default (5 MiB for s3fs); a windowed read of + a chunked file wants a much smaller one, e.g. + io_utils.DEFAULT_REMOTE_BLOCK_SIZE. Ignored for a local file. """ # lazy import of h5py import h5py # a remote file is opened as a file object and handed to h5py/bz2/gzip, # all three of which accept one in place of a name if pc.io_utils.is_remote_path(h5_file): - if fs is None: - fs = pc.io_utils.get_s3fs(daac=None) - source = fs.open(h5_file, 'rb') + source = pc.io_utils.open_remote(h5_file, fs=fs, + block_size=block_size, daac=None) else: source = h5_file if (compression is None): @@ -880,7 +885,7 @@ def from_h5(self, h5_file, field_mapping=None, group='/', \ bounds=None, skip=1, fill_value=None, t_axis=None, t_range=None, bands=None, compression=None, swap_xy=False, source_fillvalue=None, - coord_mapping=None, fs=None): + coord_mapping=None, fs=None, block_size=None): """ Read a raster from an HDF5 file. @@ -923,6 +928,9 @@ def from_h5(self, h5_file, field_mapping=None, group='/', \ fs: s3fs.S3FileSystem or NoneType, default None Filesystem to use if h5_file is a URI rather than a local path. See h5_open(). + block_size: int or NoneType, default None + Bytes fetched per range request if h5_file is remote. + See h5_open(). swap_xy: bool, default False Swap the orientation of x and y variables in the grid @@ -954,7 +962,8 @@ def from_h5(self, h5_file, field_mapping=None, group='/', \ #default yorient=1 - with self.h5_open(h5_file, mode='r', compression=compression, fs=fs) as h5f: + with self.h5_open(h5_file, mode='r', compression=compression, fs=fs, + block_size=block_size) as h5f: x = np.array(h5f[group][xname]).ravel() y = np.array(h5f[group][yname]).ravel() for time_var_name in set(['time','t', timename]): @@ -1109,10 +1118,20 @@ def from_h5(self, h5_file, field_mapping=None, group='/', \ self.__update_size_and_shape__() return self - def nc_open(self, nc_file, mode='r', compression=None, fs=None): + def nc_open(self, nc_file, mode='r', compression=None, fs=None, + engine='auto', block_size=None, rdcc_nbytes=None): """ Open a netCDF4 file with or without external compression. + A remote file is read through h5py and range requests rather than + netCDF4: netCDF4.Dataset takes a filename or an in-memory buffer but + never a file object, so reading a window out of a remote granule with + netCDF4 means downloading the whole thing. netCDF4 files are HDF5 + underneath and h5py does take a file object, so the same read costs + only the chunks the window touches. What comes back is then a + grid.nc_h5.H5Dataset, which presents the netCDF4 API that from_nc() + uses; a file h5py cannot open (NETCDF3/classic) falls back to netCDF4. + Parameters ---------- nc_file: str @@ -1120,7 +1139,10 @@ def nc_open(self, nc_file, mode='r', compression=None, fs=None): mode: str, default 'r' Mode of opening the netCDF4 file compression, str or NoneType, default None - Compression format for the netCDF4 file + External compression format for the netCDF4 file -- the whole file + wrapped in a compressed stream, as distinct from the per-chunk + compression inside a netCDF4 file, which is transparent and is + what makes a windowed read cheap. - ``None``: file is not externally compressed - ``'bzip'`` @@ -1128,32 +1150,82 @@ def nc_open(self, nc_file, mode='r', compression=None, fs=None): fs: s3fs.S3FileSystem or NoneType, default None Filesystem used if nc_file is remote. If None, a session built from the default AWS credential chain is used; see h5_open(). + engine: str, default 'auto' + Library used to read the file. + + - ``'auto'``: h5py for a remote file, netCDF4 for a local one + - ``'h5py'``: h5py in both cases (what the tests compare against) + - ``'netcdf4'``: netCDF4 in both cases, reading a remote file whole + block_size: int or NoneType, default None + Bytes fetched per range request for a remote file. If None, + io_utils.DEFAULT_REMOTE_BLOCK_SIZE is used for the h5py path, + which fetches far less than the fsspec default for a windowed + read. Ignored for a local file. + rdcc_nbytes: int or NoneType, default None + Size of the HDF5 chunk cache, in bytes. Worth raising for files + whose chunks are larger than the 1 MiB default. + + Returns + ------- + netCDF4.Dataset or grid.nc_h5.H5Dataset """ import netCDF4 - if pc.io_utils.is_remote_path(nc_file): - # netCDF4.Dataset takes a name or an in-memory buffer, but not a - # file object, so a remote file has to be read whole. That is the - # same thing the compression branches below already do, and it is - # why remote reads suit ancillary grids rather than large rasters. - if fs is None: - fs = pc.io_utils.get_s3fs(daac=None) - with fs.open(nc_file, 'rb') as fd: - buffer = fd.read() - if (compression == 'bzip'): - buffer = bz2.decompress(buffer) + from . import nc_h5 + + if engine not in ('auto', 'h5py', 'netcdf4'): + raise ValueError(f"unrecognized engine: {engine}") + remote = pc.io_utils.is_remote_path(nc_file) + + if compression is not None: + # A gzip or bzip2 stream carries no index, so byte N is only + # reachable by inflating everything before it: there is nothing to + # read chunk-wise, and the file is decompressed whole as it always + # has been. + if remote: + with pc.io_utils.open_remote(nc_file, fs=fs, daac=None) as fd: + buffer = fd.read() + if (compression == 'bzip'): + buffer = bz2.decompress(buffer) + elif (compression == 'gzip'): + buffer = gzip.decompress(buffer) + return netCDF4.Dataset(uuid.uuid4().hex, mode=mode, memory=buffer) + elif (compression == 'bzip'): + # read bytes from bzip compressed file + with bz2.BZ2File(nc_file) as fd: + return netCDF4.Dataset(uuid.uuid4().hex, mode=mode, memory=fd.read()) elif (compression == 'gzip'): - buffer = gzip.decompress(buffer) + # read bytes from gzip compressed file + with gzip.open(nc_file) as fd: + return netCDF4.Dataset(uuid.uuid4().hex, mode=mode, memory=fd.read()) + raise ValueError(f"unrecognized compression: {compression}") + + if engine == 'h5py' or (engine == 'auto' and remote): + if remote: + if block_size is None: + block_size = pc.io_utils.DEFAULT_REMOTE_BLOCK_SIZE + source = pc.io_utils.open_remote(nc_file, fs=fs, + block_size=block_size, daac=None) + else: + source = nc_file + try: + return nc_h5.open_nc_as_h5(source, mode=mode, rdcc_nbytes=rdcc_nbytes) + except Exception as exception: + if remote: + source.close() + # An OSError here means h5py could not read the file -- a + # NETCDF3/classic file, say. netCDF4 reads those, so fall + # through to it rather than failing. + if engine == 'h5py' or not isinstance(exception, OSError): + raise + + if remote: + # netCDF4 takes no file object, so the file has to be read whole. + # That is a sequential read, so it wants the filesystem's own + # block size rather than the small one a windowed read wants. + with pc.io_utils.open_remote(nc_file, fs=fs, daac=None) as fd: + buffer = fd.read() return netCDF4.Dataset(uuid.uuid4().hex, mode=mode, memory=buffer) - if (compression is None): - return netCDF4.Dataset(nc_file, mode=mode) - elif (compression == 'bzip'): - # read bytes from bzipcompressed file - with bz2.BZ2File(nc_file) as fd: - return netCDF4.Dataset(uuid.uuid4().hex, mode=mode, memory=fd.read()) - elif (compression == 'gzip'): - # read bytes from gzip compressed file - with gzip.open(nc_file) as fd: - return netCDF4.Dataset(uuid.uuid4().hex, mode=mode, memory=fd.read()) + return netCDF4.Dataset(nc_file, mode=mode) def from_nc(self, nc_file, field_mapping=None, group='', fields=None, @@ -1162,7 +1234,8 @@ def from_nc(self, nc_file, field_mapping=None, group='', bands=None, skip=1, fill_value=None, meta_only=False, - t_axis=None, compression=None, fs=None): + t_axis=None, compression=None, fs=None, + engine='auto', block_size=None, rdcc_nbytes=None): """ Read a raster from a netCDF4 file. @@ -1207,6 +1280,13 @@ def from_nc(self, nc_file, field_mapping=None, group='', fs: s3fs.S3FileSystem or NoneType, default None Filesystem to use if nc_file is a URI rather than a local path. See nc_open(). + engine: str, default 'auto' + Library used to read the file: 'auto', 'h5py' or 'netcdf4'. + See nc_open(). + block_size: int or NoneType, default None + Bytes fetched per range request for a remote file. See nc_open(). + rdcc_nbytes: int or NoneType, default None + Size of the HDF5 chunk cache, in bytes. See nc_open(). Returns ------- @@ -1232,7 +1312,10 @@ def from_nc(self, nc_file, field_mapping=None, group='', dims=[xname, yname, 't', 'time'] + list(self.coordinates or []) t=None grid_mapping_name = None - with self.nc_open(nc_file,mode='r',compression=compression, fs=fs) as fileID: + from . import nc_h5 + with self.nc_open(nc_file, mode='r', compression=compression, fs=fs, + engine=engine, block_size=block_size, + rdcc_nbytes=rdcc_nbytes) as fileID: # set automasking fileID.set_auto_mask(False) # check if reading from root group or sub-group @@ -1301,9 +1384,11 @@ def from_nc(self, nc_file, field_mapping=None, group='', this_slice = tuple([slices[dim] for dim in this_dim_order]) z = np.array(f_field[this_slice]) - # replace invalid values with nan + # replace invalid values with nan. A packed variable's data + # arrive scaled while its _FillValue is stored raw, so the + # fill value has to be put on the same scale to match. if hasattr(f_field, '_FillValue'): - fill_value = f_field.getncattr('_FillValue') + fill_value = nc_h5.scaled_fill_value(f_field) try: z[z == fill_value] = self.fill_value except ValueError: diff --git a/pointCollection/grid/nc_h5.py b/pointCollection/grid/nc_h5.py new file mode 100644 index 0000000..968952c --- /dev/null +++ b/pointCollection/grid/nc_h5.py @@ -0,0 +1,473 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" +Read a netCDF4 file through h5py, presenting the netCDF4 API that +grid.data.from_nc() uses. + +netCDF4.Dataset accepts a filename or an in-memory buffer, but never a file +object, so a remote netCDF4 file can only be read whole. netCDF4 files are +HDF5 underneath, though, and h5py does take a file object -- so pointing h5py +at an fsspec/s3fs file gives chunk-wise reads over range requests, and a +windowed read costs the chunks the window touches rather than the whole +granule. + +The classes here wrap an h5py.File so that from_nc() cannot tell the +difference: .groups, .variables, var.shape, var.dimensions, var.ncattrs(), +var.getncattr(), attribute access (hasattr(var, '_FillValue')) and the +set_auto_*() no-ops. Three differences between the libraries are papered +over here rather than in from_nc(): + + * netCDF4 applies scale_factor/add_offset unless set_auto_scale(False) is + called, and from_nc() only turns off masking. h5py applies neither, so + _scale() replicates what netCDF4 does -- otherwise packed variables come + back with different values and no error. + * h5py rejects negative-step slices, which select_slices() emits whenever + the y axis descends, and it broadcasts multiple index arrays the way + numpy does. _plan() rewrites an index tuple into a read h5py accepts + plus per-axis fixups applied afterwards. + * HDF5 datasets that are netCDF dimensions without being netCDF variables + (dimensions with no coordinate variable) are hidden, as are the HDF5 + bookkeeping attributes that netCDF4 keeps out of ncattrs(). +""" +import posixpath + +import numpy as np + +# attributes netCDF4 does not report through ncattrs() +HIDDEN_ATTRIBUTES = {'CLASS', 'NAME', 'DIMENSION_LIST', 'REFERENCE_LIST', + 'DIMENSION_LABELS', '_Netcdf4Dimid', '_Netcdf4Coordinates', + '_nc3_strict'} + +# NAME attribute of an HDF5 dimension scale that is not a netCDF variable +PHONY_DIMENSION_PREFIX = 'This is a netCDF dimension but not a netCDF variable' + + +def _as_str(value): + """decode an HDF5 string (bytes) the way netCDF4 hands it back""" + if isinstance(value, bytes): + return value.decode('utf-8', errors='replace') + return value + + +def _attribute_value(value): + """ + Convert an h5py attribute value into what netCDF4 would have returned: + a scalar for a single-valued attribute, str for text. + """ + if isinstance(value, np.ndarray): + if value.size == 1: + return _attribute_value(value.reshape(-1)[0]) + if value.dtype.kind in ('S', 'O'): + return [_as_str(item) for item in value] + return value + return _as_str(value) + + +def _is_phony_dimension(dset): + """True if dset is a netCDF dimension that is not also a netCDF variable""" + name = _as_str(dset.attrs.get('NAME')) + return isinstance(name, str) and name.startswith(PHONY_DIMENSION_PREFIX) + + +def _is_dimension_scale(dset): + return _as_str(dset.attrs.get('CLASS')) == 'DIMENSION_SCALE' + + +def apply_scaling(data, scale=None, offset=None): + """ + Apply scale_factor/add_offset to data the way netCDF4 does. + + Mirrors netCDF4.Variable.__getitem__, including its 1.0/0.0 special + cases, so that reading through h5py returns the same numbers. + """ + if scale is None and offset is None: + return data + if scale is not None and offset is not None: + if offset != 0.0 or scale != 1.0: + return data * scale + offset + return data.astype(np.asarray(scale).dtype) + if scale is not None: + return data * scale if scale != 1.0 else data + return data + offset if offset != 0.0 else data + + +def scaled_fill_value(variable): + """ + A variable's _FillValue on the scale of the data that comes back from it. + + netCDF4 scales packed data on the way out but reports _FillValue as it is + stored, so comparing the two directly never matches: a packed variable's + fill values survive into the output unconverted. Putting the fill value + through the same arithmetic as the data makes the comparison meaningful. + Works for a netCDF4.Variable as well as for an H5Variable. + + Parameters + ---------- + variable : netCDF4.Variable or H5Variable + + Returns + ------- + The fill value, scaled if the variable is packed, or None if the variable + has no _FillValue. + """ + attributes = variable.ncattrs() + if '_FillValue' not in attributes: + return None + fill = variable.getncattr('_FillValue') + scale = variable.getncattr('scale_factor') if 'scale_factor' in attributes else None + offset = variable.getncattr('add_offset') if 'add_offset' in attributes else None + if scale is None and offset is None: + return fill + # as stored, so that the arithmetic matches what the data went through + return apply_scaling(np.asarray(fill, dtype=variable.dtype), scale, offset)[()] + + +class H5Variable: + """an h5py.Dataset presented as a netCDF4.Variable""" + + def __init__(self, dset): + self._dset = dset + + # -- identity --------------------------------------------------------- + @property + def name(self): + return posixpath.basename(self._dset.name) + + @property + def shape(self): + return self._dset.shape + + @property + def dtype(self): + return self._dset.dtype + + @property + def ndim(self): + return self._dset.ndim + + def __len__(self): + return len(self._dset) + + @property + def dimensions(self): + """ + Names of the variable's dimensions, as netCDF4 reports them. + + In HDF5 these are the dimension scales attached to each axis. A + coordinate variable is its own dimension scale, and a scale is not + attached to itself, so that case falls back to the dataset's own name. + """ + dset = self._dset + names = [] + for axis, dim in enumerate(dset.dims): + if len(dim): + names.append(posixpath.basename(dim[0].name)) + elif dim.label: + names.append(posixpath.basename(_as_str(dim.label))) + elif dset.ndim == 1 and _is_dimension_scale(dset): + names.append(self.name) + else: + names.append(f'phony_dim_{axis}') + return tuple(names) + + # -- attributes ------------------------------------------------------- + def ncattrs(self): + return [name for name in self._dset.attrs + if name not in HIDDEN_ATTRIBUTES] + + def getncattr(self, name): + if name in HIDDEN_ATTRIBUTES: + raise AttributeError(name) + try: + return _attribute_value(self._dset.attrs[name]) + except KeyError: + raise AttributeError(name) + + def __getattr__(self, name): + # reached only when normal lookup fails, so this is the netCDF4 + # spelling of attribute access: var._FillValue, var.grid_mapping + if name.startswith('_') and name.endswith('__'): + raise AttributeError(name) + try: + dset = self.__dict__['_dset'] + except KeyError: + raise AttributeError(name) + if name in HIDDEN_ATTRIBUTES or name not in dset.attrs: + raise AttributeError(name) + return _attribute_value(dset.attrs[name]) + + # -- reading ---------------------------------------------------------- + def _plan(self, key): + """ + Rewrite an index tuple into one h5py accepts, plus the per-axis + fixups to apply to the array that comes back. + + h5py takes only forward slices, and index arrays on more than one + axis would broadcast rather than select an outer product, so index + arrays are read as the range they span and selected afterwards. + """ + if any(index is Ellipsis for index in key): + # expand to explicit slices so each entry maps to one axis + where = [i for i, index in enumerate(key) if index is Ellipsis][0] + fill = (slice(None),) * (self._dset.ndim - (len(key) - 1)) + key = key[:where] + fill + key[where + 1:] + read = [] + fixups = [] + for index in key: + if isinstance(index, (int, np.integer)): + # an integer index drops the axis, so it needs no fixup + read.append(int(index)) + continue + size = self._dset.shape[len(read)] if len(read) < self._dset.ndim else None + if isinstance(index, slice): + step = 1 if index.step is None else index.step + if step > 0: + read.append(index) + fixups.append(None) + continue + # negative step: read forwards, reverse afterwards + positions = range(*index.indices(size)) + if len(positions) == 0: + read.append(slice(0, 0)) + fixups.append(None) + continue + read.append(slice(positions[-1], positions[0] + 1, -step)) + fixups.append(slice(None, None, -1)) + continue + # a sequence of indices (bands, typically) + indices = np.asarray(index) + if indices.dtype == bool: + indices = np.flatnonzero(indices) + indices = np.asarray(indices, dtype=np.int64) + indices = np.where(indices < 0, indices + size, indices) + if indices.size == 0: + read.append(slice(0, 0)) + fixups.append(None) + continue + start, stop = int(indices.min()), int(indices.max()) + 1 + steps = np.unique(np.diff(indices)) if indices.size > 1 else np.array([1]) + if steps.size == 1 and steps[0] > 0: + # evenly spaced and increasing: a slice reads it exactly + read.append(slice(start, stop, int(steps[0]))) + fixups.append(None) + else: + # read the span the indices cover, then select within it + read.append(slice(start, stop)) + fixups.append(indices - start) + return tuple(read), fixups + + def __getitem__(self, key): + if not isinstance(key, tuple): + key = (key,) + if self._dset.ndim == 0: + # netCDF4 lets a scalar variable be indexed with [:], [...] or + # [()], and hands back a 0-d array; h5py accepts only the last two + return self._scale(np.asarray(self._dset[()])) + read, fixups = self._plan(key) + data = self._dset[read] + # one axis at a time: a tuple of index arrays would broadcast + for axis, fixup in enumerate(fixups): + if fixup is None: + continue + data = data[(slice(None),) * axis + (fixup,)] + return self._scale(data) + + def _scale(self, data): + """ + Apply scale_factor/add_offset the way netCDF4 does. + + from_nc() calls set_auto_mask(False) but not set_auto_scale(False), + so netCDF4 scales packed data on the way out. + """ + attrs = self._dset.attrs + return apply_scaling( + data, + _attribute_value(attrs['scale_factor']) if 'scale_factor' in attrs else None, + _attribute_value(attrs['add_offset']) if 'add_offset' in attrs else None) + + +class _Mapping: + """ + dict-like view of the datasets or groups in an h5py group. + + Group lookup is more forgiving than netCDF4's: a path with separators or + a leading '/' resolves too, since from_nc() takes the group name from its + caller. + """ + def __init__(self, group, wrap, want_group): + self._group = group + self._wrap = wrap + self._want_group = want_group + + def _members(self): + import h5py + members = {} + for name, item in self._group.items(): + if self._want_group: + if isinstance(item, h5py.Group): + members[name] = item + elif isinstance(item, h5py.Dataset) and not _is_phony_dimension(item): + members[name] = item + return members + + def __getitem__(self, name): + members = self._members() + if name in members: + return self._wrap(members[name]) + if self._want_group: + # nested or '/'-prefixed group path + key = name.strip('/') + if key in ('', '.'): + return self._wrap(self._group) + try: + item = self._group[key] + except KeyError: + raise KeyError(name) + return self._wrap(item) + raise KeyError(name) + + def __contains__(self, name): + try: + self[name] + except KeyError: + return False + return True + + def keys(self): + return self._members().keys() + + def values(self): + return [self._wrap(item) for item in self._members().values()] + + def items(self): + return [(name, self._wrap(item)) for name, item in self._members().items()] + + def __iter__(self): + return iter(self._members()) + + def __len__(self): + return len(self._members()) + + +class H5Group: + """an h5py.Group presented as a netCDF4.Group""" + + def __init__(self, group): + self._group = group + + @property + def variables(self): + return _Mapping(self._group, H5Variable, want_group=False) + + @property + def groups(self): + return _Mapping(self._group, H5Group, want_group=True) + + def __getitem__(self, name): + import h5py + item = self._group[name] + return H5Group(item) if isinstance(item, h5py.Group) else H5Variable(item) + + def __contains__(self, name): + return name in self._group + + def ncattrs(self): + return [name for name in self._group.attrs + if name not in HIDDEN_ATTRIBUTES] + + def getncattr(self, name): + try: + return _attribute_value(self._group.attrs[name]) + except KeyError: + raise AttributeError(name) + + def __getattr__(self, name): + if name.startswith('_') and name.endswith('__'): + raise AttributeError(name) + try: + group = self.__dict__['_group'] + except KeyError: + raise AttributeError(name) + if name in HIDDEN_ATTRIBUTES or name not in group.attrs: + raise AttributeError(name) + return _attribute_value(group.attrs[name]) + + # netCDF4's masking and scaling switches. Masking is off either way in + # h5py; scaling is always applied, matching netCDF4's default. + def set_auto_mask(self, value): + pass + + def set_auto_scale(self, value): + if not value: + raise NotImplementedError( + 'the h5py netCDF4 reader always applies scale_factor/add_offset') + + def set_auto_maskandscale(self, value): + self.set_auto_scale(value) + + def set_always_mask(self, value): + pass + + +class H5Dataset(H5Group): + """ + an h5py.File presented as a netCDF4.Dataset. + + Closing this closes the byte source underneath as well: closing an + h5py.File does not close the fsspec file object it was opened from. + """ + + def __init__(self, h5f, source=None): + super().__init__(h5f) + self._h5f = h5f + self._source = source + + @property + def filepath(self): + return self._h5f.filename + + def close(self): + try: + self._h5f.close() + finally: + if self._source is not None and hasattr(self._source, 'close'): + self._source.close() + + def __enter__(self): + return self + + def __exit__(self, *args): + self.close() + + +def open_nc_as_h5(source, mode='r', rdcc_nbytes=None, rdcc_nslots=None): + """ + Open a netCDF4 file with h5py and present it as a netCDF4.Dataset. + + Parameters + ---------- + source : str or file-like + Local path, or an open file object (e.g. from s3fs) to read through + range requests. + mode : str, default 'r' + rdcc_nbytes : int or NoneType + Size of the HDF5 chunk cache, in bytes. Worth raising above the 1 MiB + default for files with chunks larger than that -- ATL15's delta_h + chunks are (8, 686, 386) float32, 8.5 MiB apiece. + rdcc_nslots : int or NoneType + Number of chunk slots in that cache. + + Returns + ------- + H5Dataset + """ + import h5py + kwargs = {} + if rdcc_nbytes is not None: + kwargs['rdcc_nbytes'] = rdcc_nbytes + # HDF5 wants roughly 10 slots per cacheable chunk, and a prime + kwargs['rdcc_nslots'] = rdcc_nslots if rdcc_nslots is not None else 5003 + elif rdcc_nslots is not None: + kwargs['rdcc_nslots'] = rdcc_nslots + h5f = h5py.File(source, mode=mode, **kwargs) + return H5Dataset(h5f, source=None if isinstance(source, str) else source) diff --git a/pointCollection/io_utils.py b/pointCollection/io_utils.py index 8900386..ce3d5e4 100644 --- a/pointCollection/io_utils.py +++ b/pointCollection/io_utils.py @@ -10,6 +10,13 @@ _S3FS_CACHE = {} +# Block size for remote reads that pull windows out of a large file. fsspec's +# 5 MiB default is sized for reading a file end to end; a windowed read of a +# chunked HDF5 file touches scattered chunks, and the read-ahead is then mostly +# waste -- a 60 km window out of ATL14 fetches 35 MiB in 5 MiB blocks and +# 6.6 MiB in 256 KiB ones. +DEFAULT_REMOTE_BLOCK_SIZE = 256 * 1024 + # pc.indexedH5 is the class that reads and writes this format, so 'indexedH5' # is its canonical name. geoIndex files written before that spelling was # settled on, and calling code following the geoIndex file_type convention, @@ -88,6 +95,41 @@ def get_s3fs(daac='NSIDC', **kwargs): _S3FS_CACHE[key] = earthaccess.get_s3fs_session(daac=daac, **kwargs) return _S3FS_CACHE[key] +def open_remote(filename, mode='rb', fs=None, block_size=None, daac='NSIDC'): + """ + Open a remote (e.g. s3://) file as a file object. + + Parameters + ---------- + filename : str + mode : str, default 'rb' + fs : s3fs.S3FileSystem or NoneType, default None + Filesystem to open with. If None, a cached session is obtained via + get_s3fs(daac=daac). + block_size : int or NoneType, default None + Bytes fetched per range request. None leaves the filesystem's own + default (5 MiB for s3fs) in place; see DEFAULT_REMOTE_BLOCK_SIZE for + why a windowed read wants a smaller one. Passed per file rather than + to the session, so it applies to a caller-supplied fs too -- including + an earthaccess DAAC session, whose constructor takes no such argument. + daac : str or NoneType, default 'NSIDC' + DAAC whose credentials are needed, if fs is None. None selects the + default AWS credential chain; see get_s3fs(). + + Returns + ------- + file object + """ + if fs is None: + fs = get_s3fs(daac=daac) + if block_size is None: + return fs.open(filename, mode) + try: + return fs.open(filename, mode, block_size=block_size) + except TypeError: + # a filesystem (or a stand-in) whose open() takes no block_size + return fs.open(filename, mode) + def as_gdal_path(filename): """ Translate a URI into the /vsi... path GDAL uses for the same object. @@ -133,7 +175,7 @@ def path_exists(filename, fs=None, assume_remote_exists=True): return (fs or get_s3fs()).exists(filename) return os.path.isfile(filename) -def open_h5(filename, mode='r', fs=None): +def open_h5(filename, mode='r', fs=None, block_size=None): """ Open an HDF5 file for reading, whether local or remote. @@ -144,6 +186,9 @@ def open_h5(filename, mode='r', fs=None): fs : s3fs.S3FileSystem, optional filesystem to use for a remote open. If None, a cached session is obtained via get_s3fs(). + block_size : int, optional + bytes fetched per range request for a remote file. None leaves the + filesystem default; see open_remote(). Returns ------- @@ -151,5 +196,5 @@ def open_h5(filename, mode='r', fs=None): """ import h5py if is_remote_path(filename): - return h5py.File((fs or get_s3fs()).open(filename, 'rb'), mode) + return h5py.File(open_remote(filename, fs=fs, block_size=block_size), mode) return h5py.File(filename, mode) diff --git a/tests/test_nc_h5.py b/tests/test_nc_h5.py new file mode 100644 index 0000000..522c39d --- /dev/null +++ b/tests/test_nc_h5.py @@ -0,0 +1,467 @@ +""" +Tests for reading netCDF4 files through h5py (grid.nc_h5), the path a remote +from_nc() takes so that a windowed read costs the chunks the window touches +rather than the whole granule. + +No network access is required. engine='h5py' exercises the same reader +against a local file, and a stand-in filesystem maps s3:// URIs onto local +files for the remote branch, so what is under test is the plumbing and the +netCDF4 API the adapter presents, not S3 itself. +""" +import io +import numpy as np +import pytest +import pointCollection as pc + +netCDF4 = pytest.importorskip('netCDF4') + + +# --------------------------------------------------------------------------- +# fixtures +# --------------------------------------------------------------------------- + +def write_packed_nc(path, ydesc=True, group=None, nt=4, ny=7, nx=9): + """ + A file to_nc() cannot write: int16 data with scale_factor/add_offset and + a _FillValue, a grid_mapping variable, and a dimension with no coordinate + variable. netCDF4 applies the scaling on the way out even with masking + off, and h5py applies none of it, so this is the fixture that catches a + reader that quietly returns different numbers. + """ + with netCDF4.Dataset(path, 'w') as ds: + root = ds.createGroup(group) if group else ds + root.createDimension('x', nx) + root.createDimension('y', ny) + root.createDimension('time', nt) + root.createDimension('nv', 2) # netCDF dimension, not a variable + root.createVariable('x', 'f8', ('x',))[:] = np.arange(nx) * 100. + yv = np.arange(ny) * 100. + root.createVariable('y', 'f8', ('y',))[:] = yv[::-1] if ydesc else yv + root.createVariable('time', 'f8', ('time',))[:] = 2019. + np.arange(nt) + z = root.createVariable('z', 'i2', ('time', 'y', 'x'), + fill_value=np.int16(-9999), zlib=True) + z.scale_factor = np.float32(0.01) + z.add_offset = np.float32(5.) + z.grid_mapping = 'crs' + # write storage values directly: assigning through netCDF4's packing + # would turn a fill value into (fill - add_offset)/scale_factor + z.set_auto_maskandscale(False) + raw = np.arange(nt * ny * nx, dtype='i2').reshape(nt, ny, nx) + raw[0, 1, 1] = -9999 + raw[2, 3, 4] = -9999 + z[:] = raw + crs = root.createVariable('crs', 'S1', ()) + crs.spatial_epsg = '3413' + crs.false_easting = 0. + return path + + +@pytest.fixture +def nc_2d(tmp_path): + path = str(tmp_path / 'grid_2d.nc') + pc.grid.data().from_dict({ + 'x': np.arange(9) * 100., + 'y': np.arange(7) * 100., + 'z': np.arange(63, dtype=float).reshape((7, 9))}).to_nc(path, replace=True) + return path + + +@pytest.fixture +def nc_3d(tmp_path): + path = str(tmp_path / 'grid_3d.nc') + pc.grid.data().from_dict({ + 'x': np.arange(9) * 100., + 'y': np.arange(7) * 100., + 't': 2019. + np.arange(4), + 'z': np.arange(7 * 9 * 4, dtype=float).reshape((7, 9, 4))}).to_nc(path, replace=True) + return path + + +class FakeS3FS: + """ + Stand-in for s3fs.S3FileSystem that maps s3:// URIs onto local files and + records what was opened, and with what block size. + """ + def __init__(self, mapping, wrapper=None): + self.mapping = dict(mapping) + self.opened = [] + self.block_sizes = [] + self.wrapper = wrapper + + def open(self, path, mode='rb', block_size=None): + self.opened.append((path, mode)) + self.block_sizes.append(block_size) + fd = open(self.mapping[path], mode) + return self.wrapper(fd) if self.wrapper else fd + + def exists(self, path): + return path in self.mapping + + +class CountingFile(io.RawIOBase): + """file object that records how many bytes were actually read from it""" + def __init__(self, fd): + self.fd = fd + self.bytes_read = 0 + + def read(self, size=-1): + data = self.fd.read(size) + self.bytes_read += len(data) + return data + + def readinto(self, buffer): + n = self.fd.readinto(buffer) + self.bytes_read += n or 0 + return n + + def seek(self, offset, whence=0): + return self.fd.seek(offset, whence) + + def tell(self): + return self.fd.tell() + + def seekable(self): + return True + + def readable(self): + return True + + def close(self): + self.fd.close() + + +def assert_same_read(path, **kwargs): + """the h5py reader and the netCDF4 reader must return the same grid""" + expected = pc.grid.data().from_nc(path, engine='netcdf4', **kwargs) + got = pc.grid.data().from_nc(path, engine='h5py', **kwargs) + assert sorted(got.fields) == sorted(expected.fields) + for field in expected.fields: + a, b = getattr(expected, field), getattr(got, field) + assert a.shape == b.shape, field + assert a.dtype == b.dtype, field + assert np.array_equal(a, b, equal_nan=True), field + for coord in ('x', 'y', 't', 'time'): + a, b = getattr(expected, coord, None), getattr(got, coord, None) + assert (a is None) == (b is None), coord + if a is not None: + assert np.array_equal(np.asarray(a), np.asarray(b)), coord + assert getattr(got, 'projection', None) == getattr(expected, 'projection', None) + assert got.extent == expected.extent + return got + + +# --------------------------------------------------------------------------- +# the h5py reader returns what the netCDF4 reader returns +# --------------------------------------------------------------------------- + +READS = { + 'whole': {}, + 'bounds': {'bounds': [[100, 500], [200, 600]]}, + 'skip': {'skip': 2}, + 'bounds_and_skip': {'bounds': [[0, 800], [0, 600]], 'skip': 2}, + 'meta_only': {'meta_only': True}, + 'field': {'field': 'z'}, + 'fill_value': {'fill_value': -999.}, +} + + +@pytest.mark.parametrize('case', sorted(READS)) +def test_h5py_matches_netcdf4_2d(nc_2d, case): + assert_same_read(nc_2d, **READS[case]) + + +@pytest.mark.parametrize('case', sorted(READS)) +def test_h5py_matches_netcdf4_3d(nc_3d, case): + assert_same_read(nc_3d, **READS[case]) + + +def test_bounds_outside_the_raster_behaves_the_same(nc_2d): + """bounds that miss the raster raise out of select_slices(), either way""" + for engine in ('netcdf4', 'h5py'): + with pytest.raises(IndexError): + pc.grid.data().from_nc(nc_2d, engine=engine, + bounds=[[1e5, 2e5], [1e5, 2e5]]) + + +@pytest.mark.parametrize('kwargs', [ + {'bands': [1, 3]}, # scattered bands: not a plain slice + {'bands': [0, 1, 2]}, # contiguous bands + {'bands': [2]}, + {'t_range': [2020, 2021]}, + {'t_axis': 0}, + {'t_axis': 0, 'bounds': [[100, 500], [200, 600]]}, + {'bands': [0, 3], 'skip': 2}, +]) +def test_h5py_matches_netcdf4_bands(nc_3d, kwargs): + assert_same_read(nc_3d, **kwargs) + + +@pytest.mark.parametrize('ydesc', [True, False]) +@pytest.mark.parametrize('group', [None, 'grp']) +@pytest.mark.parametrize('kwargs', [ + {}, {'bounds': [[100, 500], [200, 600]]}, {'skip': 2}, {'bands': [0, 2, 3]}]) +def test_h5py_matches_netcdf4_packed(tmp_path, ydesc, group, kwargs): + """ + scale_factor/add_offset, _FillValue, grid_mapping, a descending y axis + (which makes select_slices() emit the negative-step slices h5py rejects) + and a dimension with no coordinate variable. + """ + path = write_packed_nc(str(tmp_path / 'packed.nc'), ydesc=ydesc, group=group) + if group: + kwargs = dict(kwargs, group=group) + out = assert_same_read(path, **kwargs) + # scaling really is being applied, and 'nv'/'crs' are not fields + assert out.z.dtype == np.float32 + assert sorted(out.fields) == ['z'] + assert out.projection['spatial_epsg'] == '3413' + + +def test_packed_fill_value_is_converted(tmp_path): + """ + A packed variable's data come back scaled while its _FillValue is stored + raw, so the two only match once the fill value is put on the same scale. + """ + path = write_packed_nc(str(tmp_path / 'packed.nc'), ydesc=False) + for engine in ('netcdf4', 'h5py'): + out = pc.grid.data().from_nc(path, engine=engine) + # two fill values were written, at (time, y, x) = (0,1,1) and (2,3,4) + assert np.isnan(out.z[1, 1, 0]) + assert np.isnan(out.z[3, 4, 2]) + assert np.count_nonzero(np.isnan(out.z)) == 2 + # and nothing else was swept up: -9999 unscaled is not a value here + assert np.isfinite(out.z[0, 0, 0]) + + +def test_packed_fill_matches_netcdf4_masking(tmp_path): + """ + netCDF4's own masking (which from_nc turns off) knows which values are + fills, so its mask is an independent check on which values from_nc + converted. + """ + path = write_packed_nc(str(tmp_path / 'packed.nc'), ydesc=False) + with netCDF4.Dataset(path) as ds: # auto-mask on: netCDF4 finds fills + masked = np.transpose(ds.variables['z'][:], [1, 2, 0]) + for engine in ('netcdf4', 'h5py'): + out = pc.grid.data().from_nc(path, engine=engine) + assert np.array_equal(np.isnan(out.z), np.ma.getmaskarray(masked)) + + +def test_packed_fill_value_respects_fill_value_argument(tmp_path): + path = write_packed_nc(str(tmp_path / 'packed.nc'), ydesc=False) + for engine in ('netcdf4', 'h5py'): + out = pc.grid.data().from_nc(path, engine=engine, fill_value=-999.) + assert out.z[1, 1, 0] == -999. + assert np.count_nonzero(out.z == -999.) == 2 + + +def test_unpacked_fill_value_is_unchanged(tmp_path): + """a variable with no scaling must behave exactly as it did before""" + path = str(tmp_path / 'plain.nc') + with netCDF4.Dataset(path, 'w') as ds: + ds.createDimension('x', 4) + ds.createDimension('y', 3) + ds.createVariable('x', 'f8', ('x',))[:] = np.arange(4) * 100. + ds.createVariable('y', 'f8', ('y',))[:] = np.arange(3) * 100. + z = ds.createVariable('z', 'f4', ('y', 'x'), fill_value=np.float32(-9999.)) + z.set_auto_maskandscale(False) + values = np.arange(12, dtype='f4').reshape(3, 4) + values[1, 2] = -9999. + z[:] = values + for engine in ('netcdf4', 'h5py'): + out = pc.grid.data().from_nc(path, engine=engine) + assert np.isnan(out.z[1, 2]) + assert np.count_nonzero(np.isnan(out.z)) == 1 + + +def test_h5py_applies_scale_factor(tmp_path): + """guard against the silent-wrong-numbers failure directly""" + path = write_packed_nc(str(tmp_path / 'packed.nc'), ydesc=False) + out = pc.grid.data().from_nc(path, engine='h5py') + with netCDF4.Dataset(path) as ds: + ds.set_auto_mask(False) + expected = np.transpose(ds.variables['z'][:], [1, 2, 0]) + # everything but the fill values, which from_nc turns into fill_value + valid = ~np.isnan(out.z) + assert np.count_nonzero(valid) == out.z.size - 2 + assert np.array_equal(out.z[valid], expected[valid]) + # unscaled data would be off by a factor of 100 and an offset + with pc.io_utils.open_h5(path) as h5f: + raw = np.transpose(h5f['z'][:], [1, 2, 0]) + assert not np.array_equal(out.z[valid], raw[valid]) + + +# --------------------------------------------------------------------------- +# remote reads go through h5py, and only fetch what they need +# --------------------------------------------------------------------------- + +def test_remote_read_matches_local(nc_3d): + fs = FakeS3FS({'s3://bucket/grid.nc': nc_3d}) + remote = pc.grid.data().from_nc('s3://bucket/grid.nc', fs=fs, + bounds=[[100, 500], [200, 600]]) + local = pc.grid.data().from_nc(nc_3d, bounds=[[100, 500], [200, 600]]) + assert np.array_equal(remote.z, local.z) + assert fs.opened == [('s3://bucket/grid.nc', 'rb')] + + +def test_remote_read_with_netcdf4_engine(nc_3d): + """engine='netcdf4' keeps the old behaviour: read the file whole""" + fs = FakeS3FS({'s3://bucket/grid.nc': nc_3d}) + remote = pc.grid.data().from_nc('s3://bucket/grid.nc', fs=fs, engine='netcdf4') + assert np.array_equal(remote.z, pc.grid.data().from_nc(nc_3d).z) + assert fs.block_sizes == [None] # a whole-file read wants big blocks + + +def test_remote_read_passes_block_size(nc_3d): + fs = FakeS3FS({'s3://bucket/grid.nc': nc_3d}) + pc.grid.data().from_nc('s3://bucket/grid.nc', fs=fs) + # FakeS3FS.open() accepts block_size, so the default must have reached it + assert fs.block_sizes == [pc.io_utils.DEFAULT_REMOTE_BLOCK_SIZE] + + fs = FakeS3FS({'s3://bucket/grid.nc': nc_3d}) + pc.grid.data().from_nc('s3://bucket/grid.nc', fs=fs, block_size=64 * 1024) + assert fs.block_sizes == [64 * 1024] + + +def test_remote_read_survives_fs_without_block_size(nc_3d): + """an fs whose open() takes no block_size must still work""" + class PlainFS(FakeS3FS): + def open(self, path, mode='rb'): + self.opened.append((path, mode)) + return open(self.mapping[path], mode) + + fs = PlainFS({'s3://bucket/grid.nc': nc_3d}) + out = pc.grid.data().from_nc('s3://bucket/grid.nc', fs=fs) + assert out.z is not None + assert fs.opened == [('s3://bucket/grid.nc', 'rb')] + + +def test_windowed_remote_read_touches_a_fraction_of_the_file(tmp_path): + """ + The point of the change: a bounded read must not pull the whole file. + """ + import os + path = str(tmp_path / 'chunked.nc') + n, chunk = 400, 50 + rng = np.random.default_rng(0) + with netCDF4.Dataset(path, 'w') as ds: + ds.createDimension('x', n) + ds.createDimension('y', n) + ds.createVariable('x', 'f8', ('x',))[:] = np.arange(n) * 100. + ds.createVariable('y', 'f8', ('y',))[:] = np.arange(n) * 100. + z = ds.createVariable('z', 'f4', ('y', 'x'), zlib=True, + chunksizes=(chunk, chunk)) + z[:] = rng.random((n, n), dtype='f4') + size = os.path.getsize(path) + + counters = [] + + def wrapper(fd): + counters.append(CountingFile(fd)) + return counters[-1] + + fs = FakeS3FS({'s3://bucket/chunked.nc': path}, wrapper=wrapper) + window = pc.grid.data().from_nc('s3://bucket/chunked.nc', fs=fs, + bounds=[[10e3, 15e3], [10e3, 15e3]]) + whole = pc.grid.data().from_nc(path) + read = counters[0].bytes_read + assert window.z.shape == (51, 51) + assert np.array_equal(window.z, whole.z[100:151, 100:151]) + # the window spans 2x2 chunks out of 64; allow generous slack for metadata + assert read < size / 4, f'read {read} of {size} bytes' + + +# --------------------------------------------------------------------------- +# files h5py cannot open +# --------------------------------------------------------------------------- + +def make_netcdf3(path): + with netCDF4.Dataset(path, 'w', format='NETCDF3_CLASSIC') as ds: + ds.createDimension('x', 9) + ds.createDimension('y', 7) + ds.createVariable('x', 'f8', ('x',))[:] = np.arange(9) * 100. + ds.createVariable('y', 'f8', ('y',))[:] = np.arange(7) * 100. + ds.createVariable('z', 'f8', ('y', 'x'))[:] = np.arange(63.).reshape(7, 9) + return path + + +def test_netcdf3_remote_falls_back_to_netcdf4(tmp_path): + path = make_netcdf3(str(tmp_path / 'classic.nc')) + fs = FakeS3FS({'s3://bucket/classic.nc': path}) + remote = pc.grid.data().from_nc('s3://bucket/classic.nc', fs=fs) + assert np.array_equal(remote.z, pc.grid.data().from_nc(path).z) + # h5py is tried first, then the file is read whole for netCDF4 + assert fs.opened == [('s3://bucket/classic.nc', 'rb')] * 2 + assert fs.block_sizes == [pc.io_utils.DEFAULT_REMOTE_BLOCK_SIZE, None] + + +def test_netcdf3_explicit_h5py_engine_raises(tmp_path): + path = make_netcdf3(str(tmp_path / 'classic.nc')) + with pytest.raises(OSError): + pc.grid.data().from_nc(path, engine='h5py') + + +def test_unknown_engine_rejected(nc_2d): + with pytest.raises(ValueError): + pc.grid.data().from_nc(nc_2d, engine='zarr') + + +# --------------------------------------------------------------------------- +# the netCDF4 API the adapter presents +# --------------------------------------------------------------------------- + +def test_adapter_presents_netcdf4_api(tmp_path): + from pointCollection.grid import nc_h5 + path = write_packed_nc(str(tmp_path / 'packed.nc'), ydesc=False, group='grp') + with pc.grid.data().nc_open(path, engine='h5py') as ds: + grp = ds.groups['grp'] + assert set(grp.variables.keys()) == {'x', 'y', 'time', 'z', 'crs'} + assert 'nv' not in grp.variables # a dimension, not a variable + z = grp.variables['z'] + assert z.dimensions == ('time', 'y', 'x') + assert z.shape == (4, 7, 9) + assert hasattr(z, '_FillValue') + assert z.getncattr('_FillValue') == -9999 + assert z.getncattr('grid_mapping') == 'crs' # decoded, not bytes + assert not hasattr(z, 'units') + # HDF5 bookkeeping attributes stay hidden, as in netCDF4 + assert 'DIMENSION_LIST' not in z.ncattrs() + assert '_Netcdf4Dimid' not in z.ncattrs() + assert set(z.ncattrs()) == {'_FillValue', 'scale_factor', 'add_offset', + 'grid_mapping'} + assert grp.variables['x'].dimensions == ('x',) + assert isinstance(ds, nc_h5.H5Dataset) + ds.set_auto_mask(False) # no-op, as netCDF4's is here + + +def test_adapter_slicing(tmp_path): + path = write_packed_nc(str(tmp_path / 'packed.nc'), ydesc=False) + with pc.grid.data().nc_open(path, engine='h5py') as ds: + z = ds.variables['z'] + full = z[:] + # negative steps, which h5py itself rejects + assert np.array_equal(z[:, ::-1, :], full[:, ::-1, :]) + assert np.array_equal(z[::-1, ::-2, ::-1], full[::-1, ::-2, ::-1]) + assert np.array_equal(z[:, 5:1:-1, :], full[:, 5:1:-1, :]) + # index lists, contiguous and scattered + assert np.array_equal(z[[0, 1, 2], :, :], full[[0, 1, 2], :, :]) + assert np.array_equal(z[[0, 3], :, :], full[[0, 3], :, :]) + assert np.array_equal(z[np.array([3, 0]), :, :], full[[3, 0], :, :]) + # integer index drops its axis; short tuples and Ellipsis + assert np.array_equal(z[2], full[2]) + assert np.array_equal(z[2, ::-1], full[2, ::-1]) + assert np.array_equal(z[..., ::-1], full[..., ::-1]) + assert np.array_equal(z[0:2], full[0:2]) + # an empty selection + assert z[:, 0:0, :].shape == (4, 0, 9) + # a scalar variable, which netCDF4 lets you index with [:] + crs = ds.variables['crs'] + assert crs.shape == () + assert np.shape(crs[:]) == () and np.shape(crs[...]) == () + + +def test_adapter_closes_the_byte_source(nc_2d): + fs = FakeS3FS({'s3://bucket/grid.nc': nc_2d}) + ds = pc.grid.data().nc_open('s3://bucket/grid.nc', fs=fs) + source = ds._source + ds.close() + assert source.closed, 'the remote file object outlived the dataset' diff --git a/tests/test_nc_h5_remote.py b/tests/test_nc_h5_remote.py new file mode 100644 index 0000000..eecd4c6 --- /dev/null +++ b/tests/test_nc_h5_remote.py @@ -0,0 +1,181 @@ +""" +Opt-in tests that read real remote files, to confirm that a windowed +from_nc() moves a small fraction of a large granule. + +Skipped by default: set PC_TEST_REMOTE=1 for the https test (no credentials +needed) and PC_TEST_S3=1 for the NSIDC direct-S3 tests, which need earthaccess +credentials and only work in us-west-2. Outside that region, earthaccess.open() +over https supports range requests and works the same way. +""" +import os +import time +import numpy as np +import pytest +import pointCollection as pc + + +def fetched_bytes(file_object): + """bytes an fsspec file object actually pulled, if it tracks that""" + cache = getattr(file_object, 'cache', None) + return getattr(cache, 'total_requested_bytes', None) + + +class RecordedFile: + """ + Proxy that snapshots the byte counter on close: fsspec drops the cache + (and the counter with it) when a file is closed, and from_nc() closes the + file it read from. + """ + def __init__(self, fd): + self.fd = fd + self.fetched = None + + def close(self): + if self.fetched is None: + self.fetched = fetched_bytes(self.fd) + self.fd.close() + + def __getattr__(self, name): + return getattr(self.fd, name) + + +class RecordingFS: + """wraps a filesystem so a test can look at the file objects afterwards""" + def __init__(self, fs): + self.fs = fs + self.files = [] + + def open(self, path, mode='rb', **kwargs): + fd = RecordedFile(self.fs.open(path, mode, **kwargs)) + self.files.append(fd) + return fd + + def __getattr__(self, name): + return getattr(self.fs, name) + + +@pytest.mark.skipif(not os.environ.get('PC_TEST_REMOTE'), + reason='set PC_TEST_REMOTE=1 to read over the network') +def test_windowed_https_read_is_chunkwise(): + """ + A public, gzip-chunked netCDF4 file read over plain https range requests: + no credentials, so this covers the remote plumbing anywhere. Reading it + at two block sizes also shows the block size doing what it is there for. + """ + fsspec = pytest.importorskip('fsspec') + url = ('https://noaa-goes16.s3.amazonaws.com/ABI-L2-SSTF/2020/001/00/' + 'OR_ABI-L2-SSTF-M6_G16_s20200010000216_e20200010059524_' + 'c20200010106082.nc') + fs = RecordingFS(fsspec.filesystem('https')) + size = fs.size(url) + + # x and y are packed int16 with a scale_factor, so take them as from_nc + # sees them rather than raw out of h5py + grid = pc.grid.data().from_nc(url, fs=fs, meta_only=True) + x, y = np.sort(grid.x), np.sort(grid.y) + bounds = [[x[2000], x[2500]], [y[2000], y[2500]]] + + reads = {} + for block_size in (5 * 2**20, pc.io_utils.DEFAULT_REMOTE_BLOCK_SIZE, + 64 * 1024): + start = time.time() + window = pc.grid.data().from_nc(url, fields=['SST'], fs=fs, + bounds=bounds, block_size=block_size) + elapsed = time.time() - start + reads[block_size] = fs.files[-1].fetched + assert window.SST.shape == (501, 501) + print(f'\n{window.SST.shape} window, {block_size//1024} KiB blocks: ' + f'{reads[block_size]/2**20:.1f} MiB of a {size/2**20:.0f} MiB ' + f'file in {elapsed:.1f} s') + + # the whole point: a window costs a fraction of the granule. This file + # is only 29 MiB, so the fraction is modest -- the win grows with the + # granule, which is what the ATL14 test below measures. + tuned = reads[pc.io_utils.DEFAULT_REMOTE_BLOCK_SIZE] + assert tuned is not None and tuned < size / 3 + # ... but only once the block size is sane. fsspec caches one block, so + # scattered chunk reads at 5 MiB a block re-fetch enough to beat + # downloading the file outright -- which is why block_size is a parameter. + assert tuned < reads[5 * 2**20] / 4 + + +def centered_window(grid, half_width=30e3): + """a 60 km window in the middle of a grid read with meta_only=True""" + x0 = 0.5 * (grid.x[0] + grid.x[-1]) + y0 = 0.5 * (grid.y[0] + grid.y[-1]) + return [[x0 - half_width, x0 + half_width], + [y0 - half_width, y0 + half_width]] + + +def granule_url(short_name, version='005', must_contain=()): + """ + Direct-S3 URL of one granule. ATL14/ATL15 collections carry the older + 0328 cycle range under the same short_name, so the caller filters on + '_0329_': picking up the wrong one is a silent wrong-data bug, not an + error. + """ + earthaccess = pytest.importorskip('earthaccess') + earthaccess.login() + urls = [] + for result in earthaccess.search_data(short_name=short_name, version=version): + links = result.data_links(access='direct') + if links and all(token in links[0] for token in must_contain): + urls.append(links[0]) + if not urls: + pytest.skip(f'no {short_name} granule matching {must_contain}') + return urls[0] + + +@pytest.mark.skipif(not os.environ.get('PC_TEST_S3'), + reason='set PC_TEST_S3=1 to read from NSIDC (us-west-2 only)') +def test_atl14_window_is_chunkwise(): + """ + ~1.4 GiB granule, 2-D float32, chunks (2491, 1401) gzip. A 60 km window + should move single-digit-to-tens of MiB and take about a second. + """ + url = granule_url('ATL14', must_contain=('_0329_', '_GL_')) + fs = RecordingFS(pc.io_utils.get_s3fs(daac='NSIDC')) + size = fs.size(url) + + bounds = centered_window(pc.grid.data().from_nc(url, fs=fs, meta_only=True)) + start = time.time() + window = pc.grid.data().from_nc(url, fields=['h'], fs=fs, bounds=bounds) + elapsed = time.time() - start + + read = fs.files[-1].fetched + print(f'\nATL14 {window.h.shape} window: {read/2**20:.1f} MiB of a ' + f'{size/2**20:.0f} MiB file in {elapsed:.1f} s') + assert np.isfinite(window.h).any() + assert read is not None and read < size / 10 + + +@pytest.mark.skipif(not os.environ.get('PC_TEST_S3'), + reason='set PC_TEST_S3=1 to read from NSIDC (us-west-2 only)') +def test_atl15_banded_window_is_chunkwise(): + """ + ATL15 delta_h is the awkward case: chunks (8, 686, 386) gzip, so the band + slicing has to be right or the read pulls whole time chunks it never uses. + """ + url = granule_url('ATL15', must_contain=('_0329_', 'ATL15_A1_01km')) + fs = RecordingFS(pc.io_utils.get_s3fs(daac='NSIDC')) + size = fs.size(url) + + meta = pc.grid.data().from_nc(url, group='delta_h', fs=fs, meta_only=True) + bounds = centered_window(meta) + start = time.time() + window = pc.grid.data().from_nc(url, group='delta_h', fields=['delta_h'], + fs=fs, bounds=bounds, + rdcc_nbytes=64 * 2**20) + elapsed = time.time() - start + + read = fs.files[-1].fetched + print(f'\nATL15 {window.delta_h.shape} window: {read/2**20:.1f} MiB of a ' + f'{size/2**20:.0f} MiB file in {elapsed:.1f} s') + assert window.delta_h.shape[:2] == (window.y.size, window.x.size) + assert read is not None and read < size / 10 + + # a band subset must not cost more than the whole time series + banded = pc.grid.data().from_nc(url, group='delta_h', fields=['delta_h'], + fs=fs, bounds=bounds, bands=[0, 1]) + assert banded.delta_h.shape[2] == 2 + assert fs.files[-1].fetched <= read