diff --git a/docs/how_to/examples.md b/docs/how_to/examples.md index 9c319716..700e9822 100644 --- a/docs/how_to/examples.md +++ b/docs/how_to/examples.md @@ -15,6 +15,10 @@ This notebook accompanies the [_Old format, no problem!: Cloud-optimizing the GO 1. [Ingesting the GOES-16 archive](https://github.com/zarr-developers/VirtualiZarr/blob/main/examples/V2/goes-16-ingest.ipynb) - End-to-end notebook building the virtual store from the netCDF archive. +### ITS_LIVE glacier-velocity mosaic + +1. [Mosaicking ITS_LIVE granules into a virtual cube](https://github.com/zarr-developers/VirtualiZarr/blob/main/examples/V2/its_live.ipynb) - Aligns granules that share a global grid but cover different regions using native xarray `concat(..., join="outer")`, stacks them along `time`, and writes the sparse virtual cube to Icechunk. + ## V1 Examples !!! note diff --git a/docs/how_to/grid_alignment.png b/docs/how_to/grid_alignment.png new file mode 100644 index 00000000..afee919e Binary files /dev/null and b/docs/how_to/grid_alignment.png differ diff --git a/docs/how_to/validation.md b/docs/how_to/validation.md index 4a9803d9..60de83e1 100644 --- a/docs/how_to/validation.md +++ b/docs/how_to/validation.md @@ -231,7 +231,6 @@ VirtualiZarr provides some utilities for creating virtual datasets with such spa The `ManifestArray.with_fill_value_only` method returns a new ManifestArray with the same schema (shape, chunks, codecs, dimension names, attributes) as a given ManifestArray, but with an empty chunk manifest and the given `fill_value`. This is useful for filling in missing files/variables, by creating virtual datasets containing manifestarrays which have no chunk references. -These empty manifestarrays can then be concatenated with the manifestarrays containing chunk references, to create a virtual dataset which a logical grid spanning regions with both present and missing data. ```python exec="on" session="usage" source="material-block" result="code" # a virtual variable from a file where the data is present @@ -255,6 +254,16 @@ print(combined) The combined variable is still a single `ManifestArray` of shape `(2, 3, 4)`, but only the first timestep references real chunks. Reading data from the second timestep would return the array's `fill_value`. +In the previously mentioned case when a dataset is "inherently sparse at a global level, whilst still being dense at a regional level", when users want to concatenate these sparse datasets they can use normal xarray concatenation patterns with outer join. + +``` +xr.concat([ds_1, ds_2], dim="time", join="outer") +``` +This will internally align the sparse datasets on a larger logical grid spanning both regions with null chunks for the sparse regions where no source chunks exist prior to concatenating them. + +![Aligning two regionally-dense but globally-sparse datasets onto a shared logical grid before concatenation along time](grid_alignment.png) + + ### Inhomogeneous Codecs Zarr arrays are by definition made up of a grid of chunks with identical codecs (e.g. for compression). diff --git a/examples/V2/README.md b/examples/V2/README.md index 269a6ad3..da461c49 100644 --- a/examples/V2/README.md +++ b/examples/V2/README.md @@ -20,6 +20,10 @@ Virtualizes the same GOES-16 file using `CachingReadableStore` and `SplittingRea uv run --script examples/V2/goes_with_caching_stores.py ``` +### [its_live.ipynb](its_live.ipynb) + +Mosaics two [ITS_LIVE](https://its-live.jpl.nasa.gov/) glacier-velocity granules that share a global grid but each cover only part of it. Uses native xarray `concat(..., join="outer")` to align them onto the shared grid (sparse, with fill where a granule has no data) and stack them along `time`, keeping the data variables virtual (`ManifestArray`) throughout, then writes the result to Icechunk without copying any pixel data. + ## Performance Comparison Run both scripts and compare the "Virtualization time" printed at the end of each: diff --git a/examples/V2/its_live.ipynb b/examples/V2/its_live.ipynb new file mode 100644 index 00000000..e53068a3 --- /dev/null +++ b/examples/V2/its_live.ipynb @@ -0,0 +1,288 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "its-live-intro", + "metadata": {}, + "source": [ + "# Mosaicking ITS\\_LIVE granules into a virtual cube\n", + "\n", + "This notebook virtualizes two [ITS\\_LIVE](https://its-live.jpl.nasa.gov/) glacier-velocity\n", + "granules that share a global grid but each cover only part of it, then uses native xarray\n", + "`concat(..., join=\"outer\")` to align them onto the shared grid (sparse, with fill where a\n", + "granule has no data) and stack them along `time`. The data variables stay virtual\n", + "(`ManifestArray`) throughout — only the coordinates are loaded — and the result is written to\n", + "Icechunk without copying any pixel data.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "adc1e1cd-7d0f-4345-a6b0-54a70ab044de", + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "\n", + "import numpy as np\n", + "import obstore\n", + "import xarray as xr\n", + "from obspec_utils.registry import ObjectStoreRegistry\n", + "\n", + "import virtualizarr as vz\n", + "from virtualizarr.parsers import HDFParser" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "f5cf7b30-577a-46c3-ac48-a6f62ba1ea86", + "metadata": {}, + "outputs": [], + "source": [ + "bucket = \"s3://its-live-data\"\n", + "key = \"test-space/virtual-cubes\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "356d1624-041a-4928-a45c-ae0070ce0b56", + "metadata": {}, + "outputs": [], + "source": [ + "granule_1 = \"LC08_L1GT_020121_20231013_20231102_02_T2_X_LC09_L1GT_020121_20231106_20231106_02_T2_G0120V02_P084.nc\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "42bca045-986d-4d17-834b-7526ac8afabb", + "metadata": {}, + "outputs": [], + "source": [ + "store = obstore.store.from_url(bucket, region=\"us-west-2\", skip_signature=True)\n", + "registry = ObjectStoreRegistry({bucket: store})\n", + "parser = HDFParser(drop_variables=[\"mapping\", \"img_pair_info\"])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ce09e8ef-8959-459f-8c1f-4676cb21a735", + "metadata": {}, + "outputs": [], + "source": [ + "vds_1 = vz.open_virtual_dataset(\n", + " url=os.path.join(bucket, key, granule_1),\n", + " parser=parser,\n", + " registry=registry,\n", + " # load x/y/time as real, indexed coordinates so xarray has an index on every\n", + " # dimension for alignment. Data variables stay virtual (ManifestArray).\n", + " loadable_variables=[\"time\", \"y\", \"x\"],\n", + " decode_times=True,\n", + ")\n", + "vds_1" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "884dffd7-f125-4c1f-add0-25075fc7ce14", + "metadata": {}, + "outputs": [], + "source": [ + "granule_2 = \"LC08_L1GT_020120_20201121_20210315_02_T2_X_LC08_L1GT_020120_20210124_20210305_02_T2_G0120V02_P051.nc\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2198e8ee-62fc-431d-986b-b4f59a9a4667", + "metadata": {}, + "outputs": [], + "source": [ + "vds_2 = vz.open_virtual_dataset(\n", + " url=os.path.join(bucket, key, granule_2),\n", + " parser=parser,\n", + " registry=registry,\n", + " loadable_variables=[\"time\", \"y\", \"x\"],\n", + " decode_times=True,\n", + ")\n", + "vds_2" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9e7d4099-bbfc-416b-b0e4-d8eb1bfb0b57", + "metadata": {}, + "outputs": [], + "source": [ + "# The ITS_LIVE data uses a descending y coordinate. Internally, xarray's alignment machinery requires ascending\n", + "# ordering so the y coordinate needs to be flipped prior to concatenation.\n", + "flip = lambda ds: ds.assign_coords(y=-ds.y) # descending y -> ascending -y\n", + "flip_1 = flip(vds_1)\n", + "flip_2 = flip(vds_2)\n", + "cube = xr.concat(\n", + " [flip_1, flip_2], dim=\"time\", join=\"outer\", combine_attrs=\"drop_conflicts\"\n", + ")\n", + "cube = flip(cube)\n", + "cube" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1d2b78bc-9560-4a0a-b55c-21933e11ff21", + "metadata": {}, + "outputs": [], + "source": [ + "import shutil\n", + "\n", + "import icechunk as ic\n", + "\n", + "# the referenced chunk data lives on this public, anonymous S3 bucket\n", + "url_prefix = \"s3://its-live-data/\"\n", + "store_path = \"its_live_cube.icechunk\"\n", + "\n", + "# start fresh so this cell is re-runnable\n", + "shutil.rmtree(store_path, ignore_errors=True)\n", + "\n", + "# register a virtual chunk container so icechunk knows how to read the s3 refs\n", + "config = ic.RepositoryConfig.default()\n", + "config.set_virtual_chunk_container(\n", + " ic.VirtualChunkContainer(\n", + " url_prefix, ic.s3_store(region=\"us-west-2\", anonymous=True)\n", + " )\n", + ")\n", + "repo = ic.Repository.create(\n", + " storage=ic.local_filesystem_storage(store_path),\n", + " config=config,\n", + " authorize_virtual_chunk_access=ic.containers_credentials(\n", + " {url_prefix: ic.s3_credentials(anonymous=True)}\n", + " ),\n", + ")\n", + "\n", + "\n", + "# icechunk metadata is strict JSON and cannot represent NaN/inf, but ITS_LIVE\n", + "# stores NaN in some attributes (e.g. 'stable_shift_stationary'); drop those.\n", + "def _drop_nonfinite_attrs(ds):\n", + " ds = ds.copy()\n", + "\n", + " def clean(attrs):\n", + " keep = {}\n", + " for k, v in attrs.items():\n", + " arr = np.asarray(v)\n", + " if arr.dtype.kind == \"f\" and not np.isfinite(arr).all():\n", + " continue\n", + " keep[k] = v\n", + " return keep\n", + "\n", + " ds.attrs = clean(ds.attrs)\n", + " for var in ds.variables.values():\n", + " var.attrs = clean(var.attrs)\n", + " return ds\n", + "\n", + "\n", + "# write the cube: only the virtual references are stored, no pixel data is copied\n", + "session = repo.writable_session(\"main\")\n", + "_drop_nonfinite_attrs(cube).vz.to_icechunk(session.store)\n", + "snapshot_id = session.commit(\n", + " \"its_live virtual cube: 2 granules mosaicked + stacked on time\"\n", + ")\n", + "print(\"committed snapshot\", snapshot_id)\n", + "\n", + "# reopen from the committed store\n", + "cube_roundtrip = xr.open_zarr(\n", + " repo.readonly_session(\"main\").store, consolidated=False, zarr_format=3\n", + ")\n", + "cube_roundtrip" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6aa95554-5a0c-481c-bc2a-945576737e00", + "metadata": {}, + "outputs": [], + "source": [ + "v = cube_roundtrip[\"v\"]\n", + "dv = v.isel(time=1) - v.isel(time=0)\n", + "dv" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "fd10354b", + "metadata": {}, + "outputs": [], + "source": [ + "import matplotlib.pyplot as plt\n", + "\n", + "# surface speed `v` at each time step, on a shared color scale\n", + "v = cube_roundtrip[\"v\"].load()\n", + "vmax = float(np.nanpercentile(v, 98))\n", + "\n", + "nt = v.sizes[\"time\"]\n", + "fig, axes = plt.subplots(1, nt, figsize=(7 * nt, 6), constrained_layout=True)\n", + "axes = np.atleast_1d(axes)\n", + "for ax, t in zip(axes, range(nt)):\n", + " im = v.isel(time=t).plot.imshow(\n", + " ax=ax, cmap=\"viridis\", vmin=0, vmax=vmax, add_colorbar=False\n", + " )\n", + " ax.set_aspect(\"equal\")\n", + " ax.set_title(str(v[\"time\"].values[t])[:10])\n", + "fig.colorbar(im, ax=list(axes), label=\"speed (m/yr)\", shrink=0.8)\n", + "plt.show()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1cc2f9b4-ce0b-4614-8b91-e2416390434e", + "metadata": {}, + "outputs": [], + "source": [ + "import matplotlib.pyplot as plt\n", + "\n", + "lim = np.nanpercentile(np.abs(dv), 98)\n", + "\n", + "# plot\n", + "fig, ax = plt.subplots(figsize=(9, 7))\n", + "dv.plot.imshow(\n", + " ax=ax,\n", + " cmap=\"RdBu_r\", # red = speedup, blue = slowdown\n", + " vmin=-lim,\n", + " vmax=lim, # symmetric, centered on zero\n", + " cbar_kwargs={\"label\": \"Δ speed (m/yr)\"},\n", + ")\n", + "ax.set_aspect(\"equal\")\n", + "plt.tight_layout()\n", + "plt.show()" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.13.2" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/virtualizarr/manifests/array.py b/virtualizarr/manifests/array.py index d7909b1d..2f519693 100644 --- a/virtualizarr/manifests/array.py +++ b/virtualizarr/manifests/array.py @@ -144,8 +144,11 @@ def __array_function__(self, func, types, args, kwargs) -> Any: return NotImplemented # Note: this allows subclasses that don't override - # __array_function__ to handle ManifestArray objects - if not all(issubclass(t, ManifestArray) for t in types): + # __array_function__ to handle ManifestArray objects. Plain ndarrays are + # also permitted among the argument types because some handled functions + # legitimately mix them with ManifestArrays — e.g. np.where, whose boolean + # condition arrives as an ndarray during xarray's reindex/alignment fill. + if not all(issubclass(t, (ManifestArray, np.ndarray)) for t in types): return NotImplemented return MANIFESTARRAY_HANDLED_ARRAY_FUNCTIONS[func](*args, **kwargs) @@ -333,6 +336,71 @@ def with_fill_value_only(self, fill_value: Any) -> "ManifestArray": ) return ManifestArray(metadata=new_metadata, chunkmanifest=empty_manifest) + def _reindex_axis( + self, axis: int, chunk_map: list[int | None], new_size: int + ) -> "ManifestArray": + """ + Return a new ManifestArray with the chunk grid along ``axis`` remapped. + + Each entry of ``chunk_map`` is the source chunk index to copy into that + target chunk slot, or ``None`` for a missing (null-path) chunk that reads + back as ``fill_value``. ``new_size`` is the new length of ``axis`` in + elements; the chunk size is unchanged. + """ + from virtualizarr.manifests.manifest import ( + MISSING_CHUNK_PATH, + ChunkManifest, + ) + from virtualizarr.manifests.utils import copy_and_replace_metadata + + manifest = self.manifest + src_paths = manifest._paths + src_offsets = manifest._offsets + src_lengths = manifest._lengths + + new_grid_shape = list(src_paths.shape) + new_grid_shape[axis] = len(chunk_map) + + new_paths = np.full( + new_grid_shape, MISSING_CHUNK_PATH, dtype=np.dtypes.StringDType() + ) + new_offsets = np.zeros(new_grid_shape, dtype=np.uint64) + new_lengths = np.zeros(new_grid_shape, dtype=np.uint64) + + new_inlined: dict[tuple[int, ...], bytes] = {} + for new_idx, src_chunk in enumerate(chunk_map): + if src_chunk is None: + continue # leave this slab as missing/null + src_slice: list[Any] = [slice(None)] * src_paths.ndim + src_slice[axis] = src_chunk + dst_slice: list[Any] = [slice(None)] * src_paths.ndim + dst_slice[axis] = new_idx + new_paths[tuple(dst_slice)] = src_paths[tuple(src_slice)] + new_offsets[tuple(dst_slice)] = src_offsets[tuple(src_slice)] + new_lengths[tuple(dst_slice)] = src_lengths[tuple(src_slice)] + # re-key any inlined chunks that lived in this source slab + for key, data in manifest._inlined.items(): + if key[axis] == src_chunk: + shifted = list(key) + shifted[axis] = new_idx + new_inlined[tuple(shifted)] = data + + new_manifest = ChunkManifest.from_arrays( + paths=new_paths, + offsets=new_offsets, + lengths=new_lengths, + validate_paths=False, + inlined=new_inlined if new_inlined else None, + ) + + new_shape = list(self.shape) + new_shape[axis] = new_size + new_metadata = copy_and_replace_metadata( + old_metadata=self.metadata, new_shape=new_shape + ) + + return ManifestArray(chunkmanifest=new_manifest, metadata=new_metadata) + def to_virtual_variable(self) -> xr.Variable: """ Create a "virtual" xarray.Variable containing the contents of one zarr array. diff --git a/virtualizarr/manifests/array_api.py b/virtualizarr/manifests/array_api.py index 164e6147..55e255d0 100644 --- a/virtualizarr/manifests/array_api.py +++ b/virtualizarr/manifests/array_api.py @@ -36,13 +36,28 @@ def decorator(func): @implements(np.result_type) def result_type(*arrays_and_dtypes: Union["ManifestArray", np.dtype]) -> np.dtype: - """Called by xarray to ensure all arguments to concat have the same dtype.""" + """ + Resolve a result dtype for ManifestArray arguments. + + Called by xarray both to check that concat/stack inputs share a dtype and, + during reindex/alignment, to combine a ManifestArray with a scalar fill + value. A ManifestArray's dtype is fixed by its metadata (the same metadata + that defines its fill value), so when exactly one ManifestArray is combined + with scalars/dtypes we return its dtype rather than promoting — keeping the + array's native dtype and declared fill value through a reindex. + """ from virtualizarr.manifests.array import ManifestArray - dtypes = ( + manifest_dtypes = [ + obj.dtype for obj in arrays_and_dtypes if isinstance(obj, ManifestArray) + ] + if len(manifest_dtypes) == 1: + return manifest_dtypes[0] + + dtypes = [ obj.dtype if isinstance(obj, ManifestArray) else np.dtype(obj) for obj in arrays_and_dtypes - ) + ] first_dtype, *other_dtypes = dtypes unique_dtypes = set(dtypes) for other_dtype in other_dtypes: @@ -54,6 +69,43 @@ def result_type(*arrays_and_dtypes: Union["ManifestArray", np.dtype]) -> np.dtyp return first_dtype +@implements(np.where) +def where(condition, x, y, /): + """ + Support xarray's reindex/alignment fill, which calls + ``where(~mask, gathered_array, fill_value)`` after gathering chunks. + + The gathered ManifestArray already carries null-path chunks (which read back + as ``fill_value``) at exactly the missing positions, so this is an identity + whenever the requested fill positions coincide with the array's missing + chunks. In that case ``x`` is returned unchanged and the manifest's own fill + value governs. Any other ``where`` usage (e.g. general boolean masking) would + require materializing values and is not supported. + """ + from virtualizarr.manifests.array import ManifestArray + + if isinstance(x, ManifestArray) and np.isscalar(y): + cond = np.asarray(condition, dtype=bool) + if cond.shape == x.shape and np.array_equal(~cond, _missing_element_mask(x)): + return x + + raise NotImplementedError( + "np.where on a ManifestArray is only supported for the reindex/alignment " + "fill pattern (filling an array's own missing chunks); general masking " + "would require materializing values." + ) + + +def _missing_element_mask(marr: "ManifestArray") -> np.ndarray: + """Boolean element-mask (shape == marr.shape), True at missing (null) chunks.""" + from virtualizarr.manifests.manifest import MISSING_CHUNK_PATH + + mask = marr.manifest._paths == MISSING_CHUNK_PATH + for axis, chunk_size in enumerate(marr.metadata.chunks): + mask = np.repeat(mask, chunk_size, axis=axis) + return mask[tuple(slice(0, length) for length in marr.shape)] + + @implements(np.concatenate) def concatenate( arrays: tuple["ManifestArray", ...] | list["ManifestArray"], diff --git a/virtualizarr/manifests/indexing.py b/virtualizarr/manifests/indexing.py index 18dc18db..435189cc 100644 --- a/virtualizarr/manifests/indexing.py +++ b/virtualizarr/manifests/indexing.py @@ -7,6 +7,7 @@ from virtualizarr.manifests.array_api import expand_dims from virtualizarr.manifests.manifest import ChunkManifest +from virtualizarr.manifests.reindex import chunk_map_from_indexer from virtualizarr.manifests.utils import copy_and_replace_metadata # indexer with only basic selectors, no new axes or ellipsis @@ -145,6 +146,89 @@ def apply_indexer( return output_arr +def _is_reindex_indexer(indexer_1d: T_BasicIndexer_1d) -> TypeGuard[np.ndarray]: + """True if this 1D indexer is an integer array (xarray's reindex indexer).""" + return isinstance(indexer_1d, np.ndarray) and indexer_1d.dtype.kind in ("i", "u") + + +def _collapse_outer_indexer( + indexer_1d: T_BasicIndexer_1d, axis: int +) -> T_BasicIndexer_1d: + """Collapse a broadcast reindex indexer back to its 1D per-axis form. + + When xarray reindexes more than one dimension at once it sends a single + *broadcast* (vectorized) indexer: one integer array per axis, shaped to the + array's rank with the real labels on its own axis and length 1 everywhere + else (e.g. ``(N, 1)`` for axis 0 and ``(1, M)`` for axis 1). Because reindex + indexing is orthogonal, each such array reshapes losslessly to the 1D + indexer for ``axis``. int/slice indexers and already-1D arrays pass through. + """ + if not (isinstance(indexer_1d, np.ndarray) and indexer_1d.ndim > 1): + return indexer_1d + if any(size != 1 for ax, size in enumerate(indexer_1d.shape) if ax != axis): + raise NotImplementedError( + "Pointwise (vectorized) fancy indexing is not supported on a " + f"ManifestArray; received a {indexer_1d.shape} indexer for axis {axis}." + ) + return indexer_1d.reshape(-1) + + +def _is_full_slice(indexer_1d: T_BasicIndexer_1d, length: int) -> bool: + """True if this slice selects the whole axis (a no-op).""" + if not isinstance(indexer_1d, slice): + return False + return ( + indexer_1d == slice(None) + or indexer_1d == slice(0, length, None) + or (indexer_1d.indices(length) == (0, length, 1)) + ) + + +def _apply_reindex( + marr: "ManifestArray", indexers: tuple[T_BasicIndexer_1d, ...] +) -> "ManifestArray": + """ + Remap the chunk grid for one or more integer-array (reindex) indexers. + + Each integer-array axis is reindexed by translating its indexer into a + chunk-grid map (null chunks for missing positions). Any other axis must be a + full-axis slice (the shape xarray's reindex/alignment produces); combining + reindexing with chunk subsetting in a single call is not supported. + """ + result = marr + for axis, raw_indexer_1d in enumerate(indexers): + indexer_1d = _collapse_outer_indexer(raw_indexer_1d, axis) + if _is_reindex_indexer(indexer_1d): + chunk_size = result.metadata.chunks[axis] + try: + chunk_map = chunk_map_from_indexer( + np.asarray(indexer_1d), chunk_size, result.shape[axis] + ) + except NotImplementedError as err: + # chunk_map_from_indexer is reached from deep inside xarray's + # reindex/alignment machinery, so re-raise with call-site context + # naming the operation, the axis, and what the user can do. + raise NotImplementedError( + f"Cannot align/reindex this virtual array along axis {axis} " + f"(chunk size {chunk_size}) without materializing data: the " + "target coordinate labels do not line up with chunk boundaries, " + "so filling them would require splitting a chunk. VirtualiZarr " + "only supports whole-chunk appends, inserts, and reorders. " + "Re-chunk along this dimension, or load the variable into memory, " + "if you need this alignment. See " + "https://github.com/zarr-developers/VirtualiZarr/issues/51." + ) from err + result = result._reindex_axis(axis, chunk_map, new_size=len(indexer_1d)) + elif _is_full_slice(indexer_1d, marr.shape[axis]): + continue + else: + raise NotImplementedError( + "Combining reindexing (an integer-array indexer) with chunk " + "subsetting in a single operation is not supported." + ) + return result + + def apply_selection( marr: "ManifestArray", indexer_without_newaxes: tuple[T_BasicIndexer_1d, ...] ) -> "ManifestArray": @@ -163,6 +247,14 @@ def apply_selection( # at this point there should be no ellipsis, no Nones, and one 1D indexer for each axis. assert len(indexer_without_newaxes) == marr.ndim + # An integer array indexer is xarray's reindex/alignment indexer (with -1 + # marking missing labels). Rather than reject it as fancy indexing, remap the + # chunk grid: missing positions become null-path chunks that read back as + # fill_value. This is what lets xarray's reindex/align/concat machinery work + # over ManifestArrays without materializing data. + if any(_is_reindex_indexer(ind) for ind in indexer_without_newaxes): + return _apply_reindex(marr, indexer_without_newaxes) + # validate types and reject anything we can never support per-axis narrowed_indexers: list[int | slice] = [] for indexer_1d in indexer_without_newaxes: diff --git a/virtualizarr/manifests/reindex.py b/virtualizarr/manifests/reindex.py new file mode 100644 index 00000000..74f015a3 --- /dev/null +++ b/virtualizarr/manifests/reindex.py @@ -0,0 +1,78 @@ +"""Translate xarray's reindex/alignment indexer into a chunk-grid map.""" + +from __future__ import annotations + +import numpy as np + +_SPLIT_MSG = ( + "Cannot reindex/align lazily: the requested labels would require splitting " + "or sub-chunk reordering of a source chunk, which VirtualiZarr does not do " + "(it would require reading chunk bytes). Only whole-chunk appends, inserts, " + "and reorders are supported. See https://github.com/zarr-developers/VirtualiZarr/issues/51." +) + + +def chunk_map_from_indexer( + indexer: np.ndarray, + chunk_size: int, + source_len: int, +) -> list[int | None]: + """ + Partition a reindex/alignment indexer into a chunk-grid map. + + xarray's reindex hands the backing array an integer indexer along the + reindexed axis, with ``-1`` marking positions absent from the source (the + would-be-fill positions). This partitions that indexer by chunk and returns, + for each target chunk slot, either the index of the source chunk to copy into + it, or ``None`` for an all-missing (null-path) chunk that reads back as + ``fill_value``. + + Each chunk-sized block of the indexer must be either entirely ``-1`` (→ null + chunk) or a contiguous, ascending, chunk-aligned run of source positions + (→ that source chunk). Anything else — a block mixing present and missing + positions, a sub-chunk reorder, or an unaligned start — raises, because it + cannot be expressed without splitting a chunk. + + Parameters + ---------- + indexer + Integer array (xarray's positional indexer), ``-1`` where missing. + chunk_size + Chunk size along this axis (number of elements per chunk). + source_len + Length of the source axis in elements (to size the trailing partial chunk). + + Returns + ------- + list of (int or None) + One entry per target chunk slot. + + Raises + ------ + NotImplementedError + If the indexer cannot be expressed at chunk granularity. + """ + pos = np.asarray(indexer) + chunk_map: list[int | None] = [] + for start in range(0, len(pos), chunk_size): + block = pos[start : start + chunk_size] + + if np.all(block == -1): + chunk_map.append(None) + continue + if np.any(block == -1): + raise NotImplementedError(_SPLIT_MSG) + + s = int(block[0]) + if not np.array_equal(block, np.arange(s, s + len(block))): + raise NotImplementedError(_SPLIT_MSG) + if s % chunk_size != 0: + raise NotImplementedError(_SPLIT_MSG) + src_chunk = s // chunk_size + src_chunk_size = min(chunk_size, source_len - src_chunk * chunk_size) + if len(block) != src_chunk_size: + raise NotImplementedError(_SPLIT_MSG) + + chunk_map.append(src_chunk) + + return chunk_map diff --git a/virtualizarr/tests/test_align.py b/virtualizarr/tests/test_align.py new file mode 100644 index 00000000..9656f7d3 --- /dev/null +++ b/virtualizarr/tests/test_align.py @@ -0,0 +1,261 @@ +""" +End-to-end tests that xarray alignment machinery works over ManifestArrays. + +Users call plain ``xr.reindex`` / ``xr.align`` / ``xr.concat`` -- no +VirtualiZarr-specific API -- and the fill stays virtual: null-path chunks that +read back as the array's ``fill_value``. + +Structure/laziness cases build ManifestArrays directly via the ``array_v3_metadata`` +fixture (as ``TestConcat`` in test_xarray.py does). Value-readback cases open real +tiled NetCDF files via the ``netcdf4_files_factory_2d`` fixture + a parser (as +``TestCombine`` does) and read the reindexed result back through a ManifestStore. +""" + +import numpy as np +import pytest +import xarray as xr + +from virtualizarr import open_virtual_dataset +from virtualizarr.manifests import ( + ChunkManifest, + ManifestArray, + ManifestGroup, + ManifestStore, +) +from virtualizarr.manifests.utils import copy_and_replace_metadata +from virtualizarr.parsers import HDFParser +from virtualizarr.tests import requires_hdf5plugin, requires_imagecodecs + +LOADABLE_COORDS = ["time", "lat", "lon"] + + +def _read_back(marr: ManifestArray, dims: list[str], registry) -> np.ndarray: + """Materialise a virtual ManifestArray's values by wrapping it in a + ManifestStore and opening it. + + Reindex preserves whatever ``dimension_names`` the source metadata carried; + parser-built arrays carry ``None``, which ``xr.open_zarr`` requires, so set + them here for the read. + """ + metadata = copy_and_replace_metadata(marr.metadata, new_dimension_names=dims) + ms = ManifestStore( + group=ManifestGroup( + arrays={"v": ManifestArray(metadata=metadata, chunkmanifest=marr.manifest)} + ), + registry=registry, + ) + return xr.open_zarr(ms, consolidated=False, zarr_format=3)["v"].values + + +class TestReindex: + def test_reindex_stays_lazy(self, array_v3_metadata): + metadata = array_v3_metadata(shape=(3,), chunks=(1,), dimension_names=["x"]) + marr = ManifestArray( + metadata=metadata, + chunkmanifest=ChunkManifest( + entries={ + "0": {"path": "/a.nc", "offset": 0, "length": 4}, + "1": {"path": "/a.nc", "offset": 4, "length": 4}, + "2": {"path": "/a.nc", "offset": 8, "length": 4}, + } + ), + ) + ds = xr.Dataset({"foo": ("x", marr)}, coords={"x": [0, 1, 2]}) + + result = ds.reindex(x=[0, 1, 2, 3, 4]) # plain xarray, no vz API + + assert isinstance(result["foo"].data, ManifestArray) + assert list(result.x.values) == [0, 1, 2, 3, 4] + # appended positions become null chunks (absent from the manifest) + assert result["foo"].data.manifest.dict() == { + "0": {"path": "file:///a.nc", "offset": 0, "length": 4}, + "1": {"path": "file:///a.nc", "offset": 4, "length": 4}, + "2": {"path": "file:///a.nc", "offset": 8, "length": 4}, + } + + def test_reindex_two_axes_simultaneously(self, array_v3_metadata): + # A 2D outer-join reindexes BOTH axes in one call, which xarray sends as a + # single broadcast (vectorized) indexer -- one N-D array per axis, shaped + # (Nt, 1) and (1, Nx) -- rather than two 1D indexers. The chunk-grid remap + # must collapse each back to its 1D per-axis indexer, else it misreads the + # array's rank (the real ITS_LIVE x+y mosaic case). + metadata = array_v3_metadata( + shape=(2, 2), chunks=(1, 1), dimension_names=["time", "x"] + ) + marr = ManifestArray( + metadata=metadata, + chunkmanifest=ChunkManifest( + entries={ + "0.0": {"path": "/a.nc", "offset": 0, "length": 4}, + "0.1": {"path": "/a.nc", "offset": 4, "length": 4}, + "1.0": {"path": "/a.nc", "offset": 8, "length": 4}, + "1.1": {"path": "/a.nc", "offset": 12, "length": 4}, + } + ), + ) + ds = xr.Dataset( + {"foo": (["time", "x"], marr)}, coords={"time": [0, 1], "x": [0, 1]} + ) + + result = ds.reindex(time=[0, 1, 2], x=[0, 1, 2]) + + assert isinstance(result["foo"].data, ManifestArray) + assert result["foo"].shape == (3, 3) + # the four real chunks are kept; the appended row and column are null + assert result["foo"].data.manifest.dict() == { + "0.0": {"path": "file:///a.nc", "offset": 0, "length": 4}, + "0.1": {"path": "file:///a.nc", "offset": 4, "length": 4}, + "1.0": {"path": "file:///a.nc", "offset": 8, "length": 4}, + "1.1": {"path": "file:///a.nc", "offset": 12, "length": 4}, + } + + def test_reindex_non_chunk_aligned_raises(self, array_v3_metadata): + metadata = array_v3_metadata(shape=(3,), chunks=(2,), dimension_names=["x"]) + marr = ManifestArray( + metadata=metadata, + chunkmanifest=ChunkManifest( + entries={ + "0": {"path": "/a.nc", "offset": 0, "length": 8}, + "1": {"path": "/a.nc", "offset": 8, "length": 4}, + } + ), + ) + ds = xr.Dataset({"foo": ("x", marr)}, coords={"x": [0, 1, 2]}) + # the deep error is wrapped with call-site context: the operation, the + # axis, and what to do -- and the low-level reason is chained. + with pytest.raises( + NotImplementedError, match="align/reindex this virtual array along axis 0" + ) as excinfo: + ds.reindex(x=[0, 1, 99, 2]) + assert "chunk boundaries" in str(excinfo.value) + assert isinstance(excinfo.value.__cause__, NotImplementedError) + + +class TestAlign: + def test_align_outer_union(self, array_v3_metadata): + def _ds(path, xs): + metadata = array_v3_metadata(shape=(3,), chunks=(1,), dimension_names=["x"]) + manifest = ChunkManifest( + entries={ + str(i): {"path": path, "offset": i * 4, "length": 4} + for i in range(3) + } + ) + marr = ManifestArray(metadata=metadata, chunkmanifest=manifest) + return xr.Dataset({"foo": ("x", marr)}, coords={"x": xs}) + + a = _ds("/a.nc", [0, 1, 2]) + b = _ds("/b.nc", [3, 4, 5]) + + ra, rb = xr.align(a, b, join="outer") + + assert isinstance(ra["foo"].data, ManifestArray) + assert isinstance(rb["foo"].data, ManifestArray) + assert list(ra.x.values) == [0, 1, 2, 3, 4, 5] + + +@requires_hdf5plugin +@requires_imagecodecs +class TestReindexReadback: + """Reindex fill reads back as the array's fill_value, verified over real files. + + Uses the air_temperature tiles from ``netcdf4_files_factory_2d`` (as + ``TestCombine`` does): two time slices x two lat bands. + + NOTE: air's ``lat`` is descending, so ``xr.align(join="outer")`` -- which sorts + the union ascending -- would require reversing within a chunk (not lazy). The + chunk-aligned, lazy path is an explicit ``reindex`` onto a descending union, + which is also the real ITS_LIVE mosaic pattern. + """ + + def test_spatial_reindex_then_concat_on_time_keeps_int_sentinel( + self, netcdf4_files_factory_2d, local_registry + ): + # air_temperature is stored int16 with sentinel fill -32767. Two tiles that + # differ in BOTH time slice and lat band (the ITS_LIVE granule shape): + # reindex each onto the union lat, then concat on time. + early_north, _, _, late_south = netcdf4_files_factory_2d() + parser = HDFParser() + with ( + open_virtual_dataset( + url=early_north, + registry=local_registry, + parser=parser, + loadable_variables=LOADABLE_COORDS, + ) as en, + open_virtual_dataset( + url=late_south, + registry=local_registry, + parser=parser, + loadable_variables=LOADABLE_COORDS, + ) as ls, + ): + union_lat = np.concatenate([en.lat.values, ls.lat.values]) # descending + cube = xr.concat( + [en.reindex(lat=union_lat), ls.reindex(lat=union_lat)], + dim="time", + join="exact", + ) + assert isinstance(cube["air"].data, ManifestArray) + assert cube["air"].dtype == np.dtype("int16") # not promoted to float + + nt, nlat = en.sizes["time"], en.sizes["lat"] + fill = cube["air"].data.metadata.fill_value + vals = _read_back(cube["air"].data, ["time", "lat", "lon"], local_registry) + en_vals = _read_back(en["air"].data, ["time", "lat", "lon"], local_registry) + ls_vals = _read_back(ls["air"].data, ["time", "lat", "lon"], local_registry) + + np.testing.assert_array_equal(vals[:nt, :nlat, :], en_vals) # early/north + np.testing.assert_array_equal(vals[nt:, nlat:, :], ls_vals) # late/south + assert (vals[:nt, nlat:, :] == fill).all() # early/south -> sentinel + assert (vals[nt:, :nlat, :] == fill).all() # late/north -> sentinel + + def test_float_fill_reads_back_nan(self, netcdf4_files_factory_2d, local_registry): + # write the tiles as float32/NaN-fill so the padded band reads back as NaN + encoding = {"air": {"dtype": "float32", "_FillValue": np.float32("nan")}} + north, _, south, _ = netcdf4_files_factory_2d(encoding=encoding) + parser = HDFParser() + with ( + open_virtual_dataset( + url=north, + registry=local_registry, + parser=parser, + loadable_variables=LOADABLE_COORDS, + ) as n, + open_virtual_dataset( + url=south, + registry=local_registry, + parser=parser, + loadable_variables=LOADABLE_COORDS, + ) as s, + ): + assert n["air"].dtype == np.dtype("float32") + union_lat = np.concatenate([n.lat.values, s.lat.values]) + reindexed = n.reindex(lat=union_lat) + assert isinstance(reindexed["air"].data, ManifestArray) + + nlat = n.sizes["lat"] + vals = _read_back( + reindexed["air"].data, ["time", "lat", "lon"], local_registry + ) + assert np.isfinite(vals[:, :nlat, :]).all() # real band + assert np.isnan(vals[:, nlat:, :]).all() # padded band -> NaN + + +class TestGeneralWhereStillRejected: + def test_boolean_masking_raises_clearly(self, array_v3_metadata): + # we must NOT silently no-op general where(); it should raise, since it + # would require materializing values. + metadata = array_v3_metadata(shape=(3,), chunks=(1,), dimension_names=["x"]) + marr = ManifestArray( + metadata=metadata, + chunkmanifest=ChunkManifest( + entries={ + str(i): {"path": "/a.nc", "offset": i * 4, "length": 4} + for i in range(3) + } + ), + ) + ds = xr.Dataset({"foo": ("x", marr)}, coords={"x": [0, 1, 2]}) + with pytest.raises(NotImplementedError, match="materializing|where"): + ds["foo"].where(ds["x"] > 0).compute() diff --git a/virtualizarr/tests/test_manifests/test_reindex.py b/virtualizarr/tests/test_manifests/test_reindex.py new file mode 100644 index 00000000..a79f82ae --- /dev/null +++ b/virtualizarr/tests/test_manifests/test_reindex.py @@ -0,0 +1,205 @@ +""" +Reindex/alignment behaviour exercised through the *public* ManifestArray API. + +xarray's reindex/alignment machinery conforms an array to new labels by indexing +it with an integer "gather" indexer that has ``-1`` at every missing position +(e.g. ``[0, 1, 2, -1, -1]``). ManifestArray supports that indexer at the array +level -- ``marr[indexer]`` -- by remapping the chunk grid: a run of ``-1`` that +covers a whole chunk becomes a null-path chunk (reads back as ``fill_value``), +and a chunk-aligned run of real indices keeps its chunk reference. Nothing that +would split a chunk is allowed. + +These tests drive that public indexing path directly rather than the private +``chunk_map_from_indexer``/``_reindex_axis`` helpers that implement it. +""" + +import numpy as np +import pytest + +from virtualizarr.manifests import ChunkManifest, ManifestArray +from virtualizarr.manifests.utils import create_v3_array_metadata + + +def _idx(*vals): + return np.array(vals, dtype="int64") + + +def _marr(shape, chunks, entries, dims): + metadata = create_v3_array_metadata( + shape=shape, + chunk_shape=chunks, + data_type=np.dtype("float32"), + codecs=[{"configuration": {"endian": "little"}, "name": "bytes"}], + fill_value=float("nan"), + dimension_names=dims, + ) + return ManifestArray( + metadata=metadata, chunkmanifest=ChunkManifest(entries=entries) + ) + + +class TestReindexGatherKeepsChunks: + """An integer gather indexer remaps the chunk grid, lazily, keeping refs.""" + + def test_append_null_chunks(self): + # chunk size 1: every missing (-1) position becomes its own null chunk + marr = _marr( + (3,), + (1,), + { + "0": {"path": "/a.nc", "offset": 0, "length": 4}, + "1": {"path": "/a.nc", "offset": 4, "length": 4}, + "2": {"path": "/a.nc", "offset": 8, "length": 4}, + }, + ["x"], + ) + result = marr[_idx(0, 1, 2, -1, -1)] + + assert result.shape == (5,) + assert result.metadata.chunks == (1,) + # appended positions are absent from the manifest -> null chunks + assert sorted(result.manifest.dict()) == ["0", "1", "2"] + + def test_chunked_append(self): + # source chunks [0,1],[2,3]; one all-missing chunk appended -> one null + marr = _marr( + (4,), + (2,), + { + "0": {"path": "/a.nc", "offset": 0, "length": 8}, + "1": {"path": "/b.nc", "offset": 8, "length": 8}, + }, + ["x"], + ) + result = marr[_idx(0, 1, 2, 3, -1, -1)] + + assert result.shape == (6,) + assert result.manifest.dict() == { + "0": {"path": "file:///a.nc", "offset": 0, "length": 8}, + "1": {"path": "file:///b.nc", "offset": 8, "length": 8}, + } + + def test_prepend_null_chunks(self): + marr = _marr( + (2,), + (1,), + { + "0": {"path": "/a.nc", "offset": 0, "length": 4}, + "1": {"path": "/a.nc", "offset": 4, "length": 4}, + }, + ["x"], + ) + result = marr[_idx(-1, -1, 0, 1)] + + assert result.shape == (4,) + # real chunks shifted to slots 2,3; slots 0,1 are null + assert sorted(result.manifest.dict()) == ["2", "3"] + + def test_insert_whole_chunk_gap(self): + # [0,1] [missing] [2,3]: a whole null chunk inserted between two real ones + marr = _marr( + (4,), + (2,), + { + "0": {"path": "/a.nc", "offset": 0, "length": 8}, + "1": {"path": "/b.nc", "offset": 8, "length": 8}, + }, + ["x"], + ) + result = marr[_idx(0, 1, -1, -1, 2, 3)] + + assert result.shape == (6,) + assert result.manifest.dict() == { + "0": {"path": "file:///a.nc", "offset": 0, "length": 8}, + "2": {"path": "file:///b.nc", "offset": 8, "length": 8}, + } + + def test_whole_chunk_reorder(self): + # reversing two whole chunks just swaps their references; no bytes read + marr = _marr( + (4,), + (2,), + { + "0": {"path": "/a.nc", "offset": 0, "length": 8}, + "1": {"path": "/b.nc", "offset": 8, "length": 8}, + }, + ["x"], + ) + result = marr[_idx(2, 3, 0, 1)] + + assert result.manifest.dict() == { + "0": {"path": "file:///b.nc", "offset": 8, "length": 8}, + "1": {"path": "file:///a.nc", "offset": 0, "length": 8}, + } + + def test_select_trailing_partial_chunk(self): + # source len 5, chunk 2 -> chunk sizes 2,2,1; selecting just the size-1 tail + marr = _marr( + (5,), + (2,), + { + "0": {"path": "/a.nc", "offset": 0, "length": 8}, + "1": {"path": "/a.nc", "offset": 8, "length": 8}, + "2": {"path": "/a.nc", "offset": 16, "length": 4}, + }, + ["x"], + ) + result = marr[_idx(4)] + + assert result.shape == (1,) + assert result.manifest.dict() == { + "0": {"path": "file:///a.nc", "offset": 16, "length": 4}, + } + + def test_only_indexed_axis_remapped_2d(self): + # a slice on axis 0 + gather on axis 1: only axis 1's grid changes + marr = _marr( + (2, 4), + (2, 2), + { + "0.0": {"path": "/a.nc", "offset": 0, "length": 16}, + "0.1": {"path": "/b.nc", "offset": 16, "length": 16}, + }, + ["y", "x"], + ) + result = marr[:, _idx(0, 1, 2, 3, -1, -1)] + + assert result.shape == (2, 6) + assert result.metadata.chunks == (2, 2) + assert result.manifest.dict() == { + "0.0": {"path": "file:///a.nc", "offset": 0, "length": 16}, + "0.1": {"path": "file:///b.nc", "offset": 16, "length": 16}, + } + + +class TestReindexGatherRejectsChunkSplits: + """Anything that would split a chunk raises, naming the offending axis.""" + + @pytest.fixture + def marr(self): + return _marr( + (4,), + (2,), + { + "0": {"path": "/a.nc", "offset": 0, "length": 8}, + "1": {"path": "/b.nc", "offset": 8, "length": 8}, + }, + ["x"], + ) + + def test_mixed_present_and_missing_in_a_chunk(self, marr): + with pytest.raises(NotImplementedError, match="chunk boundaries"): + marr[_idx(0, 1, 2, -1)] + + def test_sub_chunk_reorder(self, marr): + with pytest.raises(NotImplementedError, match="chunk boundaries"): + marr[_idx(1, 0, 2, 3)] + + def test_unaligned_start(self, marr): + with pytest.raises(NotImplementedError, match="chunk boundaries"): + marr[_idx(1, 2)] + + def test_partial_target_of_full_source(self, marr): + # taking only part of source chunk [2,3] would split it + with pytest.raises(NotImplementedError, match="chunk boundaries"): + marr[_idx(0, 1, 2)]