diff --git a/conftest.py b/conftest.py index 656400e54..2d88f3289 100644 --- a/conftest.py +++ b/conftest.py @@ -14,12 +14,14 @@ from obspec_utils.registry import ObjectStoreRegistry from obstore.store import LocalStore from xarray.core.variable import Variable +from zarr.core.metadata.v3 import ArrayV3Metadata +from zarr.dtype import parse_data_type # Local imports +from virtualizarr.codecs import convert_to_codec_pipeline from virtualizarr.manifests import ChunkManifest, ManifestArray from virtualizarr.manifests.manifest import join from virtualizarr.manifests.utils import create_v3_array_metadata -from virtualizarr.utils import ceildiv # Pytest configuration @@ -165,10 +167,9 @@ def _generate_chunk_entries( dict Mapping of chunk keys to entry dictionaries """ - chunk_grid_shape = tuple( - ceildiv(axis_length, chunk_length) - for axis_length, chunk_length in zip(shape, chunks) - ) + from zarr.experimental import ChunkGrid + + chunk_grid_shape = ChunkGrid.from_sizes(shape, chunks).grid_shape if chunk_grid_shape == (): return {"0": entry_generator((0,), (), itemsize)} @@ -451,6 +452,51 @@ def _manifest_array( return _manifest_array +@pytest.fixture +def array_v3_metadata_rectilinear(): + """ + Create V3 array metadata with a rectilinear (variable-length) chunk grid. + + Unlike ``array_v3_metadata``, ``chunk_shapes`` gives explicit per-axis chunk-edge + lengths (one sequence per axis, e.g. ``((10, 20, 30), (50, 50))``) rather than a + single chunk shape - every axis must be spelled out this way, even a uniformly + chunked one, because zarr's rectilinear chunk grid has no per-axis "regular" shorthand. + ``create_v3_array_metadata`` only builds regular chunk grids, so this constructs + ``ArrayV3Metadata`` directly instead of reusing it. + """ + + def _create_metadata( + shape: tuple = (5, 5), + chunk_shapes: tuple = ((2, 2, 1), (5,)), + data_type: np.dtype = np.dtype("int32"), + codecs: list[dict] | None = None, + fill_value: int | float | None = None, + attributes: dict | None = None, + dimension_names: Iterable[str] | None = None, + ): + codecs = codecs or [{"configuration": {"endian": "little"}, "name": "bytes"}] + zdtype = parse_data_type(data_type, zarr_format=3) + return ArrayV3Metadata( + shape=shape, + data_type=zdtype, + chunk_grid={ + "name": "rectilinear", + "configuration": { + "kind": "inline", + "chunk_shapes": [list(edges) for edges in chunk_shapes], + }, + }, + chunk_key_encoding={"name": "default"}, + fill_value=zdtype.default_scalar() if fill_value is None else fill_value, + codecs=convert_to_codec_pipeline(codecs=codecs, dtype=data_type), + attributes=attributes or {}, + dimension_names=dimension_names, + storage_transformers=None, + ) + + return _create_metadata + + @pytest.fixture def virtual_variable(array_v3_metadata: Callable) -> Callable: """Generate a virtual variable with configurable parameters.""" diff --git a/docs/about/releases.md b/docs/about/releases.md index a2592e9d4..50d5a3b4f 100644 --- a/docs/about/releases.md +++ b/docs/about/releases.md @@ -4,6 +4,18 @@ ### New Features +- Added experimental support for **rectilinear (variable-length) chunk grids**. Concatenating or + stacking virtual datasets whose chunk sizes genuinely differ along the join axis - previously + one of the restrictions on what data could be virtualized - now produces a rectilinear chunk + grid instead of raising, and the result can be written to an [Icechunk](https://icechunk.io/) store. + This is gated behind zarr-python's own experimental `array.rectilinear_chunks` config flag; + with it disabled, the same operations raise a clear error explaining how to enable it. + Appending is supported too, promoting an existing array to rectilinear where the appended + chunk sizes differ; region writes to a rectilinear array are not yet supported. + See [Rectilinear chunk grids](../explanation/data_structures.md#rectilinear-chunk-grids) for details. + Note that reading such stores back requires the latest version of Xarray (v2026.09.0). + By [Max Jones](https://github.com/maxrjones) and [Tom Nicholas](https://github.com/TomNicholas). + ### Breaking changes ### Bug fixes diff --git a/docs/explanation/data_structures.md b/docs/explanation/data_structures.md index 9fc3f3138..a642bba51 100644 --- a/docs/explanation/data_structures.md +++ b/docs/explanation/data_structures.md @@ -231,6 +231,57 @@ This concatenation property is what allows us to combine the data from multiple Implementing this feature will require a more abstract and general notion of concatenation, see [GH issue #5](https://github.com/zarr-developers/VirtualiZarr/issues/5). See the [FAQ](faq.md#can-my-specific-data-be-virtualized) for other restrictions on what data can be virtualized. +### Rectilinear chunk grids + +By default a Zarr array's chunk grid is **regular** - every chunk along a given axis has the same declared size (except the final chunk, which may be smaller). +`ManifestArray` also supports **rectilinear** (variable-length) chunk grids, where the chunk sizes along an axis are an explicit list rather than a single uniform value. +This is what allows concatenating archival files that were chunked differently along the concatenation axis, relaxing one of the [restrictions on what data can be virtualized](faq.md#can-my-specific-data-be-virtualized). + +Rectilinear chunk grids are an experimental Zarr feature, so support for them is opt-in. Enable it with: + +```python +import zarr + +zarr.config.set({"array.rectilinear_chunks": True}) +``` + +or the `ZARR_ARRAY__RECTILINEAR_CHUNKS` environment variable, before concatenating or stacking any `ManifestArray`s whose chunk sizes might not match exactly. + +The `ManifestArray.chunk_grid` property exposes the chunk grid directly, regardless of whether it's regular or rectilinear: + +```python +marr.chunk_grid.is_regular +``` + +``` +True +``` + +Concatenating or stacking `ManifestArray`s whose chunk sizes along the join axis differ automatically produces a rectilinear result, once the feature above is enabled: + +```python +import numpy as np + +# marr_10s is chunked in blocks of 10 along axis 0; marr_15s in blocks of 15 +concatenated = np.concatenate([marr_10s, marr_15s], axis=0) +concatenated.chunk_grid.chunk_sizes +``` + +``` +((10, 15),) +``` + +If rectilinear chunk grids are not enabled, the same call raises a `ValueError`, rather than silently producing metadata that most Zarr tooling can't yet read. + +!!! warning + Rectilinear chunk grid support is still limited. Concatenation, stacking, and appending are supported, and the result can be written to an [Icechunk](https://icechunk.io/) store. But **region writes** to a rectilinear-chunked array are not yet implemented - region alignment is checked against a single chunk size per axis, which has no equivalent for irregular chunk boundaries. + + Note also that some formats (including Zarr and Icechunk stores) could in principle already be rectilinear on disk, but VirtualiZarr can't read those back yet either. + +Note that reading such stores back with Xarray requires the latest version of Xarray (v2026.09.0). + +### Loading values + Remember that you cannot load values from a `ManifestArray` directly. ```python diff --git a/docs/explanation/faq.md b/docs/explanation/faq.md index b6ce2b035..2b06f8e6b 100644 --- a/docs/explanation/faq.md +++ b/docs/explanation/faq.md @@ -72,7 +72,7 @@ This means that if your data contains anything that cannot be represented within When virtualizing multi-file datasets, it is sometimes the case that it is possible to virtualize one file, but not possible to virtualize all the files together as part of one datacube, because of inconsistencies _between_ the files. The following restrictions apply across every file in the datacube you wish to create! - **Arrays** - The zarr data model is one of a set of arrays, so your data must be decodable as a set of arrays, each of which will map to single zarr array (via the `ManifestArray` class). If your data cannot be directly mapped to an array, for example because it has inconsistent lengths along a common dimension (known as "ragged data"), then it cannot be virtualized. -- **Homogeneous chunk shapes** - The zarr data model assumes that every chunk of data in a single array has the same chunk shape. For multi-file datasets each chunk often corresponds to (part of) one file, so if all your files do not have consistent chunking your data cannot be virtualized. This is a big restriction, and there are plans to relax it in future, by adding support for variable-length chunks to the zarr data model. +- **Homogeneous chunk shapes (now relaxed!)** - The zarr data model traditionally assumed that every chunk of data in a single array has the same chunk shape. For multi-file datasets each chunk often corresponds to (part of) one file, so if all your files do not have consistent chunking your data could not be virtualized, until recently! **This restriction is now relaxed** - if your chunk sizes a non-uniform along the concatenation axis, VirtualiZarr can combine them into a single **rectilinear** (variable-length) chunk grid - see [Rectilinear chunk grids](data_structures.md#rectilinear-chunk-grids). - **Homogeneous codecs** - The zarr data model assumes that every chunk of data in a single array uses the same set of codecs for compression etc. For multi-file datasets each chunk often corresponds to (part of) one file, so if all your files do not have consistent compression or other codecs your data cannot be virtualized. This is another big restriction, and there are also plans to relax it in the future. - **Registered codecs** - The codecs needed to decompress and deserialize your data must be known to zarr. This might require defining and registering a new zarr codec. - **Homogeneous data types** - The zarr data model assumes that every chunk of data in a single array decodes to the same data type (i.e. dtype). For multi-file datasets each chunk often corresponds to (part of) one file, so if all your files do not have consistent data types your data cannot be virtualized. This is arguably inherent to the concept of what an array is. diff --git a/docs/how_to/usage.md b/docs/how_to/usage.md index 0b0da9ad0..ea42cb6d7 100644 --- a/docs/how_to/usage.md +++ b/docs/how_to/usage.md @@ -287,8 +287,12 @@ In general we should be able to combine all the datasets from our archival files For combining along multiple dimensions in one call we also have [xarray.combine_nested][] and [xarray.combine_by_coords][]. If you're not familiar with any of these functions we recommend you skim through [xarray's docs on combining](https://docs.xarray.dev/en/stable/user-guide/combining.html). + +!!! note + If your datasets have different chunk sizes along the concatenation axis, concatenating them produces a **rectilinear** chunk grid instead of raising - but only once you opt in with `zarr.config.set({"array.rectilinear_chunks": True})`, since this is still an experimental zarr feature. See [Rectilinear chunk grids](../explanation/data_structures.md#rectilinear-chunk-grids) for details and current limitations. + !!! important - Currently the virtual approach requires the same chunking and encoding across datasets. See the [FAQ](../explanation/faq.md#can-my-specific-data-be-virtualized) for more details. + Currently the virtual approach requires the same encoding across datasets. See the [FAQ](../explanation/faq.md#can-my-specific-data-be-virtualized) for more details. !!! warning CF encoding attributes (such as `scale_factor` and `add_offset`) must also be consistent across files. Unlike the chunking/codec/dtype requirements above, a mismatch here will _not_ raise an error — xarray's default attribute-merging behaviour can silently drop conflicting values, leaving you with a combined dataset that is decoded incorrectly on read. See [issue #1004](https://github.com/zarr-developers/VirtualiZarr/issues/1004) and the [FAQ](../explanation/faq.md#can-my-specific-data-be-virtualized) for details. @@ -459,7 +463,7 @@ You can append a virtual dataset to an existing Icechunk store using the `append This option is designed to behave similarly to the `append_dim` option to xarray's [xarray.Dataset.to_zarr][] method, and is especially useful for datasets that grow over time. !!! important - Note again that the virtual Zarr approach requires the same chunking and encoding across datasets. This including when appending to an existing Icechunk-backed Zarr store. See the [FAQ](../explanation/faq.md#can-my-specific-data-be-virtualized) for more details. + Note again that the virtual Zarr approach requires the same encoding across datasets, and the same chunking too - unless the chunk sizes only differ along the append axis, in which case appending promotes the array to a **rectilinear** chunk grid instead of raising, provided you've opted in (see [Rectilinear chunk grids](../explanation/data_structures.md#rectilinear-chunk-grids)). This including when appending to an existing Icechunk-backed Zarr store. See the [FAQ](../explanation/faq.md#can-my-specific-data-be-virtualized) for more details. ```python exec="on" session="usage" source="material-block" result="code" # write the virtual dataset to the session with the IcechunkStore diff --git a/pyproject.toml b/pyproject.toml index 2fb0216ba..4b1a2bc6e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -27,7 +27,7 @@ dependencies = [ "numcodecs>=0.15.1", "ujson", "packaging", - "zarr>=3.1.6", # earlier versions mis-list nested group contents, silently dropping arrays (zarr-python#3657) + "zarr>=3.2.1", # first release exposing zarr.experimental.ChunkGrid, used for regular and rectilinear chunk grid handling "obstore>=0.7.0", "obspec_utils>=0.9.0", ] @@ -139,6 +139,7 @@ dev = [ "s3fs", "lithops", "dask", + "tifffile>=2026.2.16,<2027" ] [project.urls] @@ -187,6 +188,11 @@ h5netcdf = ">=1.5.0,<2" [tool.pixi.feature.icechunk-dev.dependencies] rust = "*" +[tool.pixi.feature.rectilinear.pypi-dependencies] +zarr = { git = "https://github.com/zarr-developers/zarr-python.git", branch = "main" } +xarray = { git = "https://github.com/maxrjones/xarray.git", branch = "poc/unified-zarr-chunk-grid" } +virtual-tiff = { git = "https://github.com/virtual-zarr/virtual-tiff.git", branch = "poc/unified-chunk-grid" } + # `pip` is needed by the Upstream CI job to overlay dev versions from # ci/upstream-overrides.txt on top of the pixi-resolved env. See issue #995. [tool.pixi.feature.upstream-overlay.dependencies] @@ -199,7 +205,7 @@ install-upstream-overrides = { cmd = "pip install --upgrade --no-deps -r ci/upst xarray = "==2025.6.0" numpy = "==2.1.0" numcodecs = "==0.15.1" -zarr = "==3.1.6" +zarr = "==3.2.1" obstore = "==0.7.0" icechunk = "==2.0.3" @@ -227,6 +233,7 @@ test-py314 = ["dev", "test", "remote", "hdf", "netcdf3", "fits", "hdf4", "icechu minio = ["dev", "remote", "hdf", "netcdf3", "fits", "hdf4", "icechunk", "kerchunk", "hdf5-lib", "py314", "zarr", "minio"] minimum-versions = ["dev", "test", "remote", "hdf", "netcdf3", "fits", "hdf4", "tiff", "grib", "icechunk", "kerchunk", "kerchunk_parquet", "hdf5-lib", "zarr","minimum-versions", "py312"] upstream = ["dev", "test", "hdf", "hdf5-lib", "netcdf3", "fits", "icechunk-dev", "upstream-overlay", "zarr", "py313"] +rectilinear = ["dev", "test", "tiff", "hdf5-lib", "rectilinear", "py313"] all = ["dev", "test", "remote", "hdf", "netcdf3", "fits", "hdf4", "icechunk", "kerchunk", "kerchunk_parquet", "hdf5-lib", "tiff", "grib", "zarr", "all_parsers", "all_writers", "py313"] docs = ["docs", "dev", "remote", "hdf", "netcdf3", "fits", "hdf4", "icechunk", "kerchunk", "kerchunk_parquet", "hdf5-lib", "tiff", "grib","zarr", "py313"] diff --git a/virtualizarr/manifests/array.py b/virtualizarr/manifests/array.py index a0150b399..80f2ef0d6 100644 --- a/virtualizarr/manifests/array.py +++ b/virtualizarr/manifests/array.py @@ -1,10 +1,11 @@ import dataclasses import warnings -from typing import TYPE_CHECKING, Any, Callable, Union, cast +from typing import Any, Callable, Union, cast import numpy as np import xarray as xr from zarr.core.metadata.v3 import ArrayV3Metadata +from zarr.experimental import ChunkGrid import virtualizarr.manifests.utils as utils from virtualizarr.manifests.array_api import ( @@ -16,16 +17,6 @@ from virtualizarr.manifests.utils import ChunkKeySeparator from virtualizarr.utils import determine_chunk_grid_shape -if TYPE_CHECKING: - from zarr.core.metadata.v3 import RegularChunkGridMetadata -else: - try: - from zarr.core.metadata.v3 import RegularChunkGridMetadata # zarr-python>3.1.6 - except ImportError: - from zarr.core.metadata.v3 import ( - RegularChunkGrid as RegularChunkGridMetadata, # zarr-python<=3.1.6 - ) - class ManifestArray: """ @@ -62,11 +53,6 @@ def __init__( # try unpacking the dict _metadata = ArrayV3Metadata(**metadata) - if not isinstance(_metadata.chunk_grid, RegularChunkGridMetadata): - raise NotImplementedError( - f"Only RegularChunkGrid is currently supported for chunk size, but got type {type(_metadata.chunk_grid)}" - ) - if isinstance(chunkmanifest, ChunkManifest): _chunkmanifest = chunkmanifest elif isinstance(chunkmanifest, dict): @@ -94,6 +80,11 @@ def manifest(self) -> ChunkManifest: def metadata(self) -> ArrayV3Metadata: return self._metadata + @property + def chunk_grid(self) -> ChunkGrid: + """Behavioral chunk grid bound to this array's shape.""" + return ChunkGrid.from_metadata(self._metadata) + @property def dtype(self) -> np.dtype: """The native dtype of the data (typically a numpy dtype)""" diff --git a/virtualizarr/manifests/array_api.py b/virtualizarr/manifests/array_api.py index c89408aaa..e74f0d019 100644 --- a/virtualizarr/manifests/array_api.py +++ b/virtualizarr/manifests/array_api.py @@ -2,8 +2,7 @@ from typing import TYPE_CHECKING, Any, Callable, Union, cast import numpy as np - -from virtualizarr.utils import determine_chunk_grid_shape +from zarr.experimental import ChunkGrid from .manifest import MISSING_CHUNK_PATH, ChunkManifest from .utils import ( @@ -12,8 +11,11 @@ check_same_ndims, check_same_shapes, check_same_shapes_except_on_concat_axis, + chunk_grid_sizes, copy_and_replace_metadata, + full_chunk_edges, manifest_chunk_shape, + require_rectilinear_chunks_enabled, ) if TYPE_CHECKING: @@ -97,6 +99,23 @@ def where(condition, x, y, /): ) +def _chunk_sizes( + arr: "ManifestArray", +) -> tuple[int, ...] | tuple[tuple[int, ...], ...]: + """ + Per-axis chunk size(s) of a ManifestArray. + + For a regular grid this is a tuple of ints (e.g. ``(30, 50)``); for a rectilinear + grid it's a tuple of per-axis chunk-edge tuples (e.g. ``((10, 20, 30), (50, 50))``). + + Deliberately not exposed as ``ManifestArray.chunks`` - xarray's ``is_chunked_array`` + duck-types on ``hasattr(x, "chunks")`` and would misclassify a virtual array as a + computable dask-like array (see #1016). + """ + grid = arr.chunk_grid + return grid.chunk_shape if grid.is_regular else grid.chunk_sizes + + def _missing_element_mask(marr: "ManifestArray") -> np.ndarray: """Boolean element-mask (shape == marr.shape), True at missing (null) chunks.""" mask = marr.manifest._paths == MISSING_CHUNK_PATH @@ -127,9 +146,6 @@ def concatenate( elif not isinstance(axis, int): raise TypeError() - # ensure dtypes, shapes, codecs etc. are consistent - check_combinable_zarr_arrays(arrays) - check_same_ndims([arr.ndim for arr in arrays]) # Ensure we handle axis being passed as a negative integer @@ -137,9 +153,19 @@ def concatenate( if axis < 0: axis = axis % first_arr.ndim + # Check shapes are consistent before chunk shapes: a mismatched array shape on a + # non-concat axis can also change that axis's boundary-truncated chunk edges, + # which would otherwise surface as a confusing "needs a rectilinear chunk grid" + # error instead of the more direct "differing shapes" one. arr_shapes = [arr.shape for arr in arrays] - arr_chunks = [manifest_chunk_shape(arr.metadata) for arr in arrays] check_same_shapes_except_on_concat_axis(arr_shapes, axis) + + # ensure dtypes, codecs and chunk shapes are consistent (chunk sizes along the + # concat axis itself are allowed to differ - that's what a rectilinear chunk + # grid is for) + check_combinable_zarr_arrays(arrays, exclude_axis=axis) + + arr_chunks = [chunk_grid_sizes(arr.metadata) for arr in arrays] check_no_partial_chunks_on_concat_axis(arr_shapes, arr_chunks, axis) # find what new array shape must be @@ -153,8 +179,28 @@ def concatenate( [arr.manifest for arr in arrays], axis=axis ) + # The result stays a regular grid only if every input is itself regular and they + # all declare the same chunk size along the concat axis. Otherwise the concat + # axis's real per-chunk edges (which may already differ, or may only differ once + # merged) have to be spelled out explicitly, promoting the result to a rectilinear + # chunk grid. + stays_regular = all(arr.chunk_grid.is_regular for arr in arrays) and ( + len({arr.chunk_grid.chunk_shape[axis] for arr in arrays}) == 1 + ) + + new_chunks = None + if not stays_regular: + require_rectilinear_chunks_enabled( + f"Concatenating these arrays along axis {axis}" + ) + new_chunks = list(full_chunk_edges(first_arr.metadata)) + concat_edges: tuple[int, ...] = () + for arr in arrays: + concat_edges = concat_edges + full_chunk_edges(arr.metadata)[axis] + new_chunks[axis] = concat_edges + new_metadata = copy_and_replace_metadata( - old_metadata=first_arr.metadata, new_shape=new_shape + old_metadata=first_arr.metadata, new_shape=new_shape, new_chunks=new_chunks ) return ManifestArray(chunkmanifest=concatenated_manifest, metadata=new_metadata) @@ -199,10 +245,17 @@ def stack( # do stacking of entries in manifest stacked_manifest = _stack_manifests([arr.manifest for arr in arrays], axis=axis) - # chunk shape has changed because a length-1 axis has been inserted - old_chunks = manifest_chunk_shape(first_arr.metadata) + # chunk shape has changed because a new axis has been inserted, with one + # length-1 chunk per stacked array + old_chunks = _chunk_sizes(first_arr) new_chunks = list(old_chunks) - new_chunks.insert(axis, 1) + # For rectilinear grids, each element is a sequence of edges rather than a + # single chunk size, so the new axis needs one size-1 edge per stacked array + if not first_arr.chunk_grid.is_regular: + require_rectilinear_chunks_enabled("Stacking these arrays") + new_chunks.insert(axis, (1,) * length_along_new_stacked_axis) + else: + new_chunks.insert(axis, 1) new_metadata = copy_and_replace_metadata( old_metadata=first_arr.metadata, new_shape=new_shape, new_chunks=new_chunks @@ -238,22 +291,21 @@ def broadcast_to(x: "ManifestArray", /, shape: tuple[int, ...]) -> "ManifestArra # new chunk_shape is old chunk_shape with singleton dimensions prepended # (chunk shape can never change by more than adding length-1 axes because each chunk represents a fixed number of array elements) - old_chunk_shape = manifest_chunk_shape(x.metadata) + # broadcast_to only applies to regular chunk grids + old_chunk_shape = x.chunk_grid.chunk_shape new_chunk_shape = _prepend_singleton_dimensions( old_chunk_shape, ndim=len(new_shape) ) - # find new chunk grid shape by dividing new array shape by new chunk shape - new_chunk_grid_shape = determine_chunk_grid_shape(new_shape, new_chunk_shape) - - # do broadcasting of entries in manifest - broadcasted_manifest = _broadcast_manifest(x.manifest, shape=new_chunk_grid_shape) - new_metadata = copy_and_replace_metadata( old_metadata=x.metadata, new_shape=list(new_shape), new_chunks=list(new_chunk_shape), ) + new_chunk_grid_shape = ChunkGrid.from_metadata(new_metadata).grid_shape + + # do broadcasting of entries in manifest + broadcasted_manifest = _broadcast_manifest(x.manifest, shape=new_chunk_grid_shape) return ManifestArray(chunkmanifest=broadcasted_manifest, metadata=new_metadata) diff --git a/virtualizarr/manifests/utils.py b/virtualizarr/manifests/utils.py index bcfe972a0..1397dfa41 100644 --- a/virtualizarr/manifests/utils.py +++ b/virtualizarr/manifests/utils.py @@ -2,10 +2,13 @@ import functools import re import typing +from collections.abc import Sequence from typing import TYPE_CHECKING, Any, Dict, Iterable, Literal, Optional, Union, cast import numpy as np +import zarr from zarr import Array +from zarr.core.chunk_grids import _is_rectilinear_chunks from zarr.core.chunk_key_encodings import ChunkKeyEncodingLike from zarr.core.metadata.v2 import ArrayV2Metadata from zarr.core.metadata.v3 import ( @@ -14,6 +17,7 @@ parse_shapelike, ) from zarr.dtype import parse_data_type +from zarr.experimental import ChunkGrid from virtualizarr.codecs import convert_to_codec_pipeline, get_codecs @@ -241,12 +245,41 @@ def check_same_codecs(codecs: list[Any]) -> None: ) -def check_same_chunk_shapes(chunks_list: list[tuple[int, ...]]) -> None: - """Check all the chunk shapes are the same""" +def check_same_chunk_shapes( + chunks_list: list[Sequence], exclude_axis: int | None = None +) -> None: + """ + Check all the chunk shapes are the same. + + Parameters + ---------- + chunks_list + Each array's per-axis chunk size(s), as returned by + [chunk_grid_sizes][virtualizarr.manifests.utils.chunk_grid_sizes]. + exclude_axis + An axis to ignore when comparing, e.g. the concat axis - along which chunk + sizes may differ (that's exactly what a rectilinear chunk grid is for). + """ + + def _normalize(size: Union[int, Sequence[int]]) -> tuple[int, ...]: + # A bare int (a regular axis) and a 1-tuple of that same value (an axis + # zarr's ChunkGrid classified as regular despite the grid being declared + # rectilinear, since every axis of a rectilinear grid with only uniform + # axes still reports as regular overall) describe the same chunking, and + # must compare equal regardless of which form either side happens to use. + return (size,) if isinstance(size, int) else tuple(size) + + def _comparable(chunks: Sequence) -> tuple: + axes = ( + tuple(chunks) + if exclude_axis is None + else _remove_element_at_position(tuple(chunks), exclude_axis) + ) + return tuple(_normalize(axis) for axis in axes) first_chunks, *other_chunks_list = chunks_list for other_chunks in other_chunks_list: - if other_chunks != first_chunks: + if _comparable(other_chunks) != _comparable(first_chunks): raise ValueError( f"Cannot concatenate arrays with inconsistent chunk shapes: {other_chunks} vs {first_chunks} ." "Requires ZEP003 (Variable-length Chunks)." @@ -284,11 +317,17 @@ def _remove_elements_at_positions( def check_no_partial_chunks_on_concat_axis( - shapes: list[tuple[int, ...]], chunks: list[tuple[int, ...]], axis: int + shapes: list[tuple[int, ...]], chunks: list, axis: int ): - """Check that there are no partial chunks along the concatenation axis""" - # loop over the arrays to be concatenated + """Check that there are no partial chunks along the concatenation axis. + + Only applies to regular chunk grids; rectilinear grids explicitly encode + variable chunk sizes so partial-chunk checks are not needed. + """ for i, (shape, chunk_shape) in enumerate(zip(shapes, chunks)): + # Rectilinear grids have sequences along each axis; skip the check + if _is_rectilinear_chunks(chunk_shape): + continue if shape[axis] % chunk_shape[axis] > 0: raise ValueError( "Cannot concatenate arrays with partial chunks because only regular chunk grids are currently supported. " @@ -331,19 +370,28 @@ def check_same_shapes_except_on_concat_axis(shapes: list[tuple[int, ...]], axis: def check_combinable_zarr_arrays( arrays: Iterable[Union["ManifestArray", "Array"]], + exclude_axis: int | None = None, ) -> None: """ The downside of the ManifestArray approach compared to the VirtualZarrArray concatenation proposal is that the result must also be a single valid zarr array, implying that the inputs must have the same dtype, codec etc. + + Parameters + ---------- + exclude_axis + An axis to ignore when comparing chunk shapes, e.g. the concat axis - passed + through to [check_same_chunk_shapes][virtualizarr.manifests.utils.check_same_chunk_shapes]. """ + arrays = list(arrays) check_same_dtypes([arr.dtype for arr in arrays]) # Can't combine different codecs in one manifest # see https://github.com/zarr-developers/zarr-specs/issues/288 check_same_codecs([get_codecs(arr) for arr in arrays]) - # Would require variable-length chunks ZEP - check_same_chunk_shapes([manifest_chunk_shape(arr.metadata) for arr in arrays]) + check_same_chunk_shapes( + [chunk_grid_sizes(arr.metadata) for arr in arrays], exclude_axis=exclude_axis + ) def check_compatible_arrays( @@ -369,13 +417,20 @@ def manifest_chunk_shape( ---------- metadata Metadata of the array whose manifest unit is wanted. Zarr V2 metadata is accepted - because [check_combinable_zarr_arrays][virtualizarr.manifests.utils.check_combinable_zarr_arrays] - may be handed a V2 `zarr.Array` alongside `ManifestArray`s. + because [chunk_grid_sizes][virtualizarr.manifests.utils.chunk_grid_sizes] may be + handed a V2 `zarr.Array` alongside `ManifestArray`s. Returns ------- The shape covered by one manifest entry: the shard shape if `metadata` has a sharding codec, else the chunk shape. + + Raises + ------ + AttributeError + If `metadata` has a rectilinear chunk grid, which has no single chunk shape. + Use [chunk_grid_sizes][virtualizarr.manifests.utils.chunk_grid_sizes] instead + where a rectilinear grid needs to be tolerated rather than rejected. """ if not isinstance(metadata, ArrayV3Metadata): # Zarr V2 has no sharding, so a chunk is the manifest's unit @@ -383,6 +438,94 @@ def manifest_chunk_shape( return tuple(cast("RegularChunkGridMetadata", metadata.chunk_grid).chunk_shape) +def chunk_grid_sizes( + metadata: Union[ArrayV3Metadata, "ArrayV2Metadata"], +) -> tuple: + """ + Per-axis chunk size(s) of `metadata`'s chunk grid, tolerating a rectilinear grid. + + Unlike [manifest_chunk_shape][virtualizarr.manifests.utils.manifest_chunk_shape], + this doesn't assume a regular grid: for a rectilinear grid each axis comes back as + its declared tuple of per-chunk edge lengths (e.g. `((10, 20, 30), (50, 50))`) + instead of raising, while a regular axis stays a bare int. Reports each axis as + declared, *not* expanded to account for a boundary-truncated final chunk, so a + genuinely regular axis with a partial last chunk stays distinguishable from a + rectilinear one - see + [check_no_partial_chunks_on_concat_axis][virtualizarr.manifests.utils.check_no_partial_chunks_on_concat_axis]. + For a form that's safe to compare for equality across grids that may or may not + have been simplified to regular, see + [full_chunk_edges][virtualizarr.manifests.utils.full_chunk_edges]. + + Parameters + ---------- + metadata + Metadata of the array whose chunk grid is wanted. Zarr V2 metadata is accepted + because [check_combinable_zarr_arrays][virtualizarr.manifests.utils.check_combinable_zarr_arrays] + may be handed a V2 `zarr.Array` alongside `ManifestArray`s. + """ + if not isinstance(metadata, ArrayV3Metadata): + return tuple(metadata.chunks) + grid = ChunkGrid.from_metadata(metadata) + return grid.chunk_shape if grid.is_regular else grid.chunk_sizes + + +def _axis_edges(extent: int, chunk_size: int) -> tuple[int, ...]: + """The real per-chunk lengths a regular axis breaks `extent` into, including a + truncated final chunk.""" + n_full, remainder = divmod(extent, chunk_size) + edges = (chunk_size,) * n_full + return edges + (remainder,) if remainder else edges + + +def full_chunk_edges( + metadata: Union[ArrayV3Metadata, "ArrayV2Metadata"], +) -> tuple[tuple[int, ...], ...]: + """ + Every axis's real per-chunk edge lengths, tolerating a rectilinear grid. + + Unlike [chunk_grid_sizes][virtualizarr.manifests.utils.chunk_grid_sizes], always + returns one explicit tuple of edge lengths per axis (e.g. + ``((10, 20, 30), (50, 50))``) - including for a regular grid's uniform axes, and + including a boundary-truncated final chunk - rather than a bare chunk size. zarr's + `ChunkGrid` classifies a whole grid as regular or rectilinear based on its *actual* + edge values, so a grid declared rectilinear but with only uniform axes still + reports as regular; always expanding to edge tuples keeps two grids comparable + regardless of that whole-grid classification. Used to compare grids for + compatibility, or to build a concat axis's merged edges - not where a boundary + chunk still needs to be told apart from a genuinely rectilinear one. + + Parameters + ---------- + metadata + Metadata of the array whose chunk grid is wanted. Zarr V2 metadata is accepted + because [check_combinable_zarr_arrays][virtualizarr.manifests.utils.check_combinable_zarr_arrays] + may be handed a V2 `zarr.Array` alongside `ManifestArray`s. + """ + if not isinstance(metadata, ArrayV3Metadata): + return tuple( + _axis_edges(extent, size) + for extent, size in zip(metadata.shape, metadata.chunks) + ) + grid = ChunkGrid.from_metadata(metadata) + if grid.is_regular: + return tuple( + _axis_edges(extent, size) + for extent, size in zip(metadata.shape, grid.chunk_shape) + ) + return grid.chunk_sizes + + +def require_rectilinear_chunks_enabled(context: str) -> None: + """Raise a clear, actionable error unless rectilinear chunk grids are enabled.""" + if not zarr.config.get("array.rectilinear_chunks"): + raise ValueError( + f"{context} would require a rectilinear (variable-length) chunk grid. " + "Rectilinear chunk grids are an experimental zarr-python feature; enable " + "them with zarr.config.set({'array.rectilinear_chunks': True}) or the " + "ZARR_ARRAY__RECTILINEAR_CHUNKS environment variable." + ) + + def _realign_inner_chunk_shape( old_chunks: tuple[int, ...], new_chunks: tuple[int, ...], @@ -504,7 +647,7 @@ def _realign_sharding_codecs( def copy_and_replace_metadata( old_metadata: ArrayV3Metadata, new_shape: list[int] | None = None, - new_chunks: list[int] | None = None, + new_chunks: list | None = None, new_dimension_names: Iterable[str] | None | Literal["default"] = "default", new_attributes: dict | None = None, ) -> ArrayV3Metadata: @@ -548,20 +691,30 @@ def copy_and_replace_metadata( if new_shape is not None: metadata_copy["shape"] = parse_shapelike(new_shape) # type: ignore[assignment] if new_chunks is not None: - old_chunks = manifest_chunk_shape(old_metadata) - new_chunks = list(new_chunks) - metadata_copy["chunk_grid"] = { - "name": "regular", - "configuration": {"chunk_shape": tuple(new_chunks)}, - } - if len(new_chunks) != len(old_chunks): - # a sharding codec's inner chunk_shape must match the array's ndim, so it has - # to gain or lose the same length-1 axes the outer chunk shape just did - metadata_copy["codecs"] = _realign_sharding_codecs( - cast(list[dict[str, Any]], metadata_copy["codecs"]), - old_chunks, - tuple(new_chunks), - ) + if _is_rectilinear_chunks(new_chunks): + metadata_copy["chunk_grid"] = { + "name": "rectilinear", + "configuration": { + "kind": "inline", + "chunk_shapes": [list(c) for c in new_chunks], + }, + } + else: + new_chunks = list(new_chunks) + metadata_copy["chunk_grid"] = { + "name": "regular", + "configuration": {"chunk_shape": tuple(new_chunks)}, + } + if isinstance(old_metadata.chunk_grid, RegularChunkGridMetadata): + old_chunks = manifest_chunk_shape(old_metadata) + if len(new_chunks) != len(old_chunks): + # a sharding codec's inner chunk_shape must match the array's ndim, so it has + # to gain or lose the same length-1 axes the outer chunk shape just did + metadata_copy["codecs"] = _realign_sharding_codecs( + cast(list[dict[str, Any]], metadata_copy["codecs"]), + old_chunks, + tuple(new_chunks), + ) if new_dimension_names != "default": # need the option to use the literal string "default" as a sentinel value because None is a valid choice for zarr dimension_names metadata_copy["dimension_names"] = parse_dimension_names(new_dimension_names) diff --git a/virtualizarr/parsers/kerchunk/translator.py b/virtualizarr/parsers/kerchunk/translator.py index afd01f15c..f25b7cccc 100644 --- a/virtualizarr/parsers/kerchunk/translator.py +++ b/virtualizarr/parsers/kerchunk/translator.py @@ -9,6 +9,7 @@ import ujson from zarr.core.common import JSON from zarr.core.metadata import ArrayV3Metadata +from zarr.experimental import ChunkGrid from virtualizarr.codecs import ( zarr_codec_config_to_v3, @@ -24,7 +25,6 @@ KerchunkArrRefs, KerchunkStoreRefs, ) -from virtualizarr.utils import determine_chunk_grid_shape def from_kerchunk_refs(decoded_arr_refs_zarray, zattrs) -> "ArrayV3Metadata": @@ -198,10 +198,7 @@ def manifestarray_from_kerchunk_refs( chunk_dict, metadata, zattrs = parse_array_refs(arr_refs) # we want to remove the _ARRAY_DIMENSIONS from the final variables' .attrs if chunk_dict: - chunk_grid_shape = determine_chunk_grid_shape( - metadata.shape, - metadata.chunks, - ) + chunk_grid_shape = ChunkGrid.from_metadata(metadata).grid_shape manifest = manifest_from_kerchunk_chunk_dict( chunk_dict, fs_root=fs_root, shape=chunk_grid_shape ) @@ -210,10 +207,7 @@ def manifestarray_from_kerchunk_refs( # empty variables don't have physical chunks, but zarray shows that the variable # is at least 1D - shape = determine_chunk_grid_shape( - metadata.shape, - metadata.chunks, - ) + shape = ChunkGrid.from_metadata(metadata).grid_shape manifest = ChunkManifest(entries={}, shape=shape) marr = ManifestArray(metadata=metadata, chunkmanifest=manifest) else: diff --git a/virtualizarr/parsers/zarr/common.py b/virtualizarr/parsers/zarr/common.py index fe564ca6c..dbddf34e0 100644 --- a/virtualizarr/parsers/zarr/common.py +++ b/virtualizarr/parsers/zarr/common.py @@ -7,25 +7,16 @@ import math from collections.abc import Coroutine, Sequence from enum import Enum -from typing import TYPE_CHECKING, Any, TypeVar, cast +from typing import Any, TypeVar, cast import numpy as np import zarr from zarr.core.metadata import ArrayV2Metadata, ArrayV3Metadata +from zarr.core.metadata.v3 import RegularChunkGridMetadata +from zarr.experimental import ChunkGrid from virtualizarr.manifests import ChunkManifest from virtualizarr.manifests.utils import ChunkKeySeparator -from virtualizarr.utils import determine_chunk_grid_shape - -if TYPE_CHECKING: - from zarr.core.metadata.v3 import RegularChunkGridMetadata -else: - try: - from zarr.core.metadata.v3 import RegularChunkGridMetadata # zarr-python>3.1.6 - except ImportError: - from zarr.core.metadata.v3 import ( - RegularChunkGrid as RegularChunkGridMetadata, # zarr-python<=3.1.6 - ) # obstore doesn't export a public base type for stores, so we use Any for now. ObstoreStore = Any @@ -176,7 +167,7 @@ def parse_array_layout( if not isinstance(metadata.chunk_grid, RegularChunkGridMetadata): raise NotImplementedError( - f"Only RegularChunkGrid is supported, but array {zarr_array.path} " + f"Only RegularChunkGridMetadata is supported, but array {zarr_array.path} " f"uses {type(metadata.chunk_grid).__name__}." ) @@ -189,12 +180,9 @@ def parse_array_layout( else cast(ArrayV2Metadata, zarr_array.metadata).dimension_separator ) - # For sharded arrays, chunk_grid.chunk_shape is the shard shape (not the inner - # chunk shape, which lives inside the ShardingCodec config). So this grid describes - # the number of shard files on disk, which is exactly what we want for the manifest. - chunk_grid_shape = determine_chunk_grid_shape( - metadata.shape, cast(RegularChunkGridMetadata, metadata.chunk_grid).chunk_shape - ) + # For sharded arrays, grid_shape reflects the number of shard files on disk, + # which is exactly what we want for the manifest. + chunk_grid_shape = ChunkGrid.from_metadata(metadata).grid_shape return metadata, on_disk_zarr_format, on_disk_separator, chunk_grid_shape diff --git a/virtualizarr/parsers/zarr/zarr.py b/virtualizarr/parsers/zarr/zarr.py index 37c17b235..21f9cd097 100644 --- a/virtualizarr/parsers/zarr/zarr.py +++ b/virtualizarr/parsers/zarr/zarr.py @@ -9,6 +9,7 @@ import zarr from obspec_utils.registry import ObjectStoreRegistry from zarr.core.metadata import ArrayV3Metadata +from zarr.experimental import ChunkGrid from zarr.storage import ObjectStore from virtualizarr.manifests import ( @@ -24,14 +25,12 @@ from virtualizarr.parsers.utils import construct_manifest_group_tree from virtualizarr.parsers.zarr.common import ( ObstoreStore, - RegularChunkGridMetadata, ZarrFormat, _run_async, chunk_entries_to_manifest, join_url, parse_array_layout, ) -from virtualizarr.utils import determine_chunk_grid_shape class ZarrParser: @@ -203,12 +202,9 @@ async def build_chunk_manifest( missing, Zarr will return the fill_value for those regions when the array is read. """ - # For sharded arrays, chunk_grid.chunk_shape is the shard shape (not the inner - # chunk shape, which lives inside the ShardingCodec config). So this grid describes - # the number of shard files on disk, which is exactly what we want for the manifest. - chunk_grid_shape = determine_chunk_grid_shape( - metadata.shape, cast(RegularChunkGridMetadata, metadata.chunk_grid).chunk_shape - ) + # For sharded arrays, grid_shape reflects the number of shard files on disk, + # which is exactly what we want for the manifest. + chunk_grid_shape = ChunkGrid.from_metadata(metadata).grid_shape # Handle scalar arrays if metadata.shape == (): diff --git a/virtualizarr/tests/conftest.py b/virtualizarr/tests/conftest.py index 55b08eee5..f78bf1789 100644 --- a/virtualizarr/tests/conftest.py +++ b/virtualizarr/tests/conftest.py @@ -2,6 +2,14 @@ import time import pytest +import zarr + + +@pytest.fixture(autouse=True) +def _enable_rectilinear_chunks(): + """Enable rectilinear chunks for all tests.""" + with zarr.config.set({"array.rectilinear_chunks": True}): + yield @pytest.fixture(scope="session") diff --git a/virtualizarr/tests/test_manifests/test_array.py b/virtualizarr/tests/test_manifests/test_array.py index d9b7a8ce6..03e13f9ec 100644 --- a/virtualizarr/tests/test_manifests/test_array.py +++ b/virtualizarr/tests/test_manifests/test_array.py @@ -1,5 +1,6 @@ import numpy as np import pytest +import zarr from zarr.core.metadata.v3 import ArrayV3Metadata from conftest import ( @@ -432,6 +433,42 @@ def test_broadcast_mixed_inlined_and_virtual(self, array_v3_metadata): } +class TestChunkGrid: + def test_chunk_grid_regular(self, manifest_array): + marr = manifest_array(shape=(10, 20), chunks=(5, 20)) + grid = marr.chunk_grid + assert grid.is_regular is True + assert grid.chunk_shape == (5, 20) + + def test_chunk_grid_rectilinear(self, array_v3_metadata_rectilinear): + metadata = array_v3_metadata_rectilinear( + shape=(5, 5), chunk_shapes=((2, 2, 1), (5,)) + ) + marr = ManifestArray( + metadata=metadata, + chunkmanifest=ChunkManifest(entries={}, shape=(3, 1)), + ) + grid = marr.chunk_grid + assert grid.is_regular is False + assert grid.chunk_sizes == ((2, 2, 1), (5,)) + + +class TestBroadcastRectilinear: + def test_broadcast_to_rectilinear_array_raises(self, array_v3_metadata_rectilinear): + # broadcast_to only supports regular chunk grids (a chunk's element count can + # only change by adding length-1 axes, which doesn't hold for a rectilinear grid) + metadata = array_v3_metadata_rectilinear( + shape=(60, 50), chunk_shapes=((10, 20, 30), (50,)) + ) + marr = ManifestArray( + metadata=metadata, + chunkmanifest=ChunkManifest(entries={}, shape=(3, 1)), + ) + + with pytest.raises(ValueError, match="only available for regular chunk grids"): + np.broadcast_to(marr, shape=(2, 60, 50)) + + # TODO we really need some kind of fixtures to generate useful example data # The hard part is having an alternative way to get to the expected result of concatenation class TestConcat: @@ -571,6 +608,154 @@ def test_concat_all_virtual_leaves_inlined_empty(self, array_v3_metadata): assert result.manifest._inlined == {} +class TestConcatRectilinear: + def test_concat_mismatched_shapes_on_non_concat_axis_raises_shape_error( + self, array_v3_metadata + ): + # regression test: both arrays declare the same chunk size (5), but shape + # (5, 6) truncates its last chunk on axis 1 differently than shape (5, 5) + # does. This must be caught as a shape mismatch, not misreported as + # needing a rectilinear chunk grid just because the boundary-truncated + # chunk edges happen to differ as a result. + metadata1 = array_v3_metadata(shape=(5, 5), chunks=(3, 3)) + marr1 = ManifestArray( + metadata=metadata1, chunkmanifest=ChunkManifest(entries={}, shape=(2, 2)) + ) + metadata2 = array_v3_metadata(shape=(5, 6), chunks=(3, 3)) + marr2 = ManifestArray( + metadata=metadata2, chunkmanifest=ChunkManifest(entries={}, shape=(2, 2)) + ) + + with pytest.raises(ValueError, match="Cannot concatenate arrays with shapes"): + np.concatenate([marr1, marr2], axis=0) + + def test_concat_regular_arrays_stays_regular(self, manifest_array): + # concatenating regular-grid arrays must not promote the result to a + # rectilinear chunk grid - the concat axis's chunk sizes are unchanged + marr1 = manifest_array(shape=(5, 2), chunks=(5, 2)) + marr2 = manifest_array(shape=(5, 2), chunks=(5, 2)) + + result = np.concatenate([marr1, marr2], axis=0) + + assert result.chunk_grid.is_regular is True + assert result.metadata.chunk_grid.to_dict()["name"] == "regular" + assert result.chunk_grid.chunk_shape == (5, 2) + + def test_concat_two_rectilinear_arrays_merges_chunk_edges_along_axis( + self, array_v3_metadata_rectilinear + ): + # array A: 60 elements along axis 0 chunked as (10, 20, 30); axis 1 uniform at 50 + metadata_a = array_v3_metadata_rectilinear( + shape=(60, 50), chunk_shapes=((10, 20, 30), (50,)) + ) + manifest_a = ChunkManifest( + entries={ + "0.0": {"path": "/a.nc", "offset": 0, "length": 100}, + "1.0": {"path": "/a.nc", "offset": 100, "length": 100}, + "2.0": {"path": "/a.nc", "offset": 200, "length": 100}, + } + ) + marr_a = ManifestArray(metadata=metadata_a, chunkmanifest=manifest_a) + + # array B: a single 15-element chunk along axis 0 + metadata_b = array_v3_metadata_rectilinear( + shape=(15, 50), chunk_shapes=((15,), (50,)) + ) + manifest_b = ChunkManifest( + entries={"0.0": {"path": "/b.nc", "offset": 0, "length": 100}} + ) + marr_b = ManifestArray(metadata=metadata_b, chunkmanifest=manifest_b) + + result = np.concatenate([marr_a, marr_b], axis=0) + + assert result.shape == (75, 50) + assert result.chunk_grid.is_regular is False + assert result.chunk_grid.chunk_sizes == ((10, 20, 30, 15), (50,)) + + def test_concat_rectilinear_arrays_along_uniform_axis( + self, array_v3_metadata_rectilinear + ): + # concatenating along an axis that happens to be uniformly chunked + # (within an otherwise-rectilinear array) just appends its edges + metadata = array_v3_metadata_rectilinear( + shape=(30, 50), chunk_shapes=((10, 20), (50,)) + ) + manifest1 = ChunkManifest( + entries={ + "0.0": {"path": "/a.nc", "offset": 0, "length": 100}, + "1.0": {"path": "/a.nc", "offset": 100, "length": 100}, + } + ) + marr1 = ManifestArray(metadata=metadata, chunkmanifest=manifest1) + manifest2 = ChunkManifest( + entries={ + "0.0": {"path": "/b.nc", "offset": 0, "length": 100}, + "1.0": {"path": "/b.nc", "offset": 100, "length": 100}, + } + ) + marr2 = ManifestArray(metadata=metadata, chunkmanifest=manifest2) + + result = np.concatenate([marr1, marr2], axis=1) + + assert result.shape == (30, 100) + assert result.chunk_grid.chunk_sizes == ((10, 20), (50, 50)) + + def test_concat_regular_arrays_with_different_chunk_sizes_raises_when_disabled( + self, array_v3_metadata + ): + # two regular arrays whose declared chunk sizes genuinely differ along the + # concat axis can only be combined as a rectilinear grid + metadata_a = array_v3_metadata(shape=(20,), chunks=(10,)) + manifest_a = ChunkManifest( + entries={ + "0": {"path": "/a.nc", "offset": 0, "length": 40}, + "1": {"path": "/a.nc", "offset": 40, "length": 40}, + } + ) + marr_a = ManifestArray(metadata=metadata_a, chunkmanifest=manifest_a) + + metadata_b = array_v3_metadata(shape=(15,), chunks=(15,)) + manifest_b = ChunkManifest( + entries={"0": {"path": "/b.nc", "offset": 0, "length": 60}} + ) + marr_b = ManifestArray(metadata=metadata_b, chunkmanifest=manifest_b) + + with zarr.config.set({"array.rectilinear_chunks": False}): + with pytest.raises(ValueError) as exc_info: + np.concatenate([marr_a, marr_b], axis=0) + + # error must clearly explain how to turn rectilinear chunks on + assert "rectilinear" in str(exc_info.value) + assert "zarr.config.set" in str(exc_info.value) + assert "ZARR_ARRAY__RECTILINEAR_CHUNKS" in str(exc_info.value) + + def test_concat_regular_arrays_with_different_chunk_sizes_succeeds_when_enabled( + self, array_v3_metadata + ): + metadata_a = array_v3_metadata(shape=(20,), chunks=(10,)) + manifest_a = ChunkManifest( + entries={ + "0": {"path": "/a.nc", "offset": 0, "length": 40}, + "1": {"path": "/a.nc", "offset": 40, "length": 40}, + } + ) + marr_a = ManifestArray(metadata=metadata_a, chunkmanifest=manifest_a) + + metadata_b = array_v3_metadata(shape=(15,), chunks=(15,)) + manifest_b = ChunkManifest( + entries={"0": {"path": "/b.nc", "offset": 0, "length": 60}} + ) + marr_b = ManifestArray(metadata=metadata_b, chunkmanifest=manifest_b) + + # flag is already enabled for all tests in this suite (see the autouse + # fixture), so this should succeed and produce a rectilinear result + result = np.concatenate([marr_a, marr_b], axis=0) + + assert result.shape == (35,) + assert result.chunk_grid.is_regular is False + assert result.chunk_grid.chunk_sizes == ((10, 10, 15),) + + class TestStack: def test_stack(self, array_v3_metadata): # both manifest arrays in this example have the same metadata @@ -904,6 +1089,62 @@ def test_stack_preserves_bytes_identity(self, array_v3_metadata): assert result.manifest._inlined[(1, 0)] is payload +class TestStackRectilinear: + def test_stack_regular_arrays_stays_regular(self, manifest_array): + marr1 = manifest_array(shape=(5, 2), chunks=(5, 2)) + marr2 = manifest_array(shape=(5, 2), chunks=(5, 2)) + + result = np.stack([marr1, marr2], axis=0) + + assert result.chunk_grid.is_regular is True + assert result.metadata.chunk_grid.to_dict()["name"] == "regular" + + def test_stack_rectilinear_arrays_inserts_singleton_edge( + self, array_v3_metadata_rectilinear + ): + metadata = array_v3_metadata_rectilinear( + shape=(60, 50), chunk_shapes=((10, 20, 30), (50,)) + ) + manifest = ChunkManifest( + entries={ + "0.0": {"path": "/a.nc", "offset": 0, "length": 100}, + "1.0": {"path": "/a.nc", "offset": 100, "length": 100}, + "2.0": {"path": "/a.nc", "offset": 200, "length": 100}, + } + ) + marr1 = ManifestArray(metadata=metadata, chunkmanifest=manifest) + marr2 = ManifestArray(metadata=metadata, chunkmanifest=manifest) + + result = np.stack([marr1, marr2], axis=0) + + assert result.shape == (2, 60, 50) + assert result.chunk_grid.is_regular is False + assert result.chunk_grid.chunk_sizes == ((1, 1), (10, 20, 30), (50,)) + + def test_stack_rectilinear_arrays_raises_when_disabled( + self, array_v3_metadata_rectilinear + ): + metadata = array_v3_metadata_rectilinear( + shape=(60, 50), chunk_shapes=((10, 20, 30), (50,)) + ) + manifest = ChunkManifest( + entries={ + "0.0": {"path": "/a.nc", "offset": 0, "length": 100}, + "1.0": {"path": "/a.nc", "offset": 100, "length": 100}, + "2.0": {"path": "/a.nc", "offset": 200, "length": 100}, + } + ) + marr1 = ManifestArray(metadata=metadata, chunkmanifest=manifest) + marr2 = ManifestArray(metadata=metadata, chunkmanifest=manifest) + + with zarr.config.set({"array.rectilinear_chunks": False}): + with pytest.raises(ValueError) as exc_info: + np.stack([marr1, marr2], axis=0) + + assert "rectilinear" in str(exc_info.value) + assert "zarr.config.set" in str(exc_info.value) + + class TestWithFillValueOnly: def test_returns_manifest_array_with_empty_manifest(self, array_v3_metadata): # with_fill_value_only produces a ManifestArray with the same schema @@ -983,6 +1224,20 @@ def test_preserves_existing_fill_value_when_passed(self, array_v3_metadata): assert result.metadata.to_dict() == marr.metadata.to_dict() assert result.manifest.dict() == {} + def test_rectilinear_array_raises(self, array_v3_metadata_rectilinear): + # documents a known gap: with_fill_value_only goes through + # manifest_chunk_shape, which assumes a regular chunk grid + metadata = array_v3_metadata_rectilinear( + shape=(60, 50), chunk_shapes=((10, 20, 30), (50,)) + ) + marr = ManifestArray( + metadata=metadata, + chunkmanifest=ChunkManifest(entries={}, shape=(3, 1)), + ) + + with pytest.raises(AttributeError): + marr.with_fill_value_only(0) + def test_refuse_combine(array_v3_metadata): # TODO test refusing to concatenate arrays that have conflicting shapes / chunk sizes @@ -1210,6 +1465,26 @@ def test_misaligned_with_chunks(self, manifest_array, in_shape, in_chunks, index marr[indexer] +class TestIndexingRectilinear: + def test_getitem_on_rectilinear_array_raises(self, array_v3_metadata_rectilinear): + # documents a known gap: indexing goes through manifest_chunk_shape, + # which assumes a regular chunk grid + metadata = array_v3_metadata_rectilinear( + shape=(60,), chunk_shapes=((10, 20, 30),) + ) + manifest = ChunkManifest( + entries={ + "0": {"path": "/a.nc", "offset": 0, "length": 100}, + "1": {"path": "/a.nc", "offset": 100, "length": 100}, + "2": {"path": "/a.nc", "offset": 200, "length": 100}, + } + ) + marr = ManifestArray(metadata=metadata, chunkmanifest=manifest) + + with pytest.raises(AttributeError): + marr[0:10] + + class TestSubChunkSlicingUncompressed: # For an uncompressed array, sub-chunk slicing along the axis with the largest byte # stride in storage can be expressed purely as a byte-offset/length adjustment into diff --git a/virtualizarr/tests/test_parsers/test_tiff.py b/virtualizarr/tests/test_parsers/test_tiff.py index cf0e110c6..32b3bef22 100644 --- a/virtualizarr/tests/test_parsers/test_tiff.py +++ b/virtualizarr/tests/test_parsers/test_tiff.py @@ -1,10 +1,19 @@ +import numpy as np import pytest +import xarray as xr from obspec_utils.registry import ObjectStoreRegistry -from obstore.store import S3Store +from obstore.store import LocalStore, S3Store from xarray import Dataset, DataTree from virtualizarr import open_virtual_dataset, open_virtual_datatree -from virtualizarr.tests import requires_network, requires_tiff + +try: + from zarr.core.metadata.v3 import RectilinearChunkGrid # noqa: F401 + + has_rectilinear_chunk_grid_support = True +except ImportError: + has_rectilinear_chunk_grid_support = False +from virtualizarr.tests import requires_network, requires_tiff, requires_tifffile virtual_tiff = pytest.importorskip("virtual_tiff") @@ -40,3 +49,62 @@ def test_virtual_tiff_dataset() -> None: var = vds["0"].variable assert var.sizes == {"y": 10980, "x": 10980} assert var.dtype == " None: + """Test concatenating two virtual TIFF datasets with rectilinear chunk grids. + + Creates stripped TIFFs where image_height is not evenly divisible by rows_per_strip, + producing rectilinear chunks, then verifies they can be concatenated. + """ + import tifffile + + # Create two stripped TIFFs where image_height (100) is not evenly divisible + # by rows_per_strip (30), creating rectilinear chunks: [[30, 30, 30, 10], [50]] + shape = (100, 50) + rows_per_strip = 30 + + filepath1 = tmp_path / "test1.tif" + filepath2 = tmp_path / "test2.tif" + + tifffile.imwrite( + str(filepath1), np.ones(shape, dtype=np.uint8), rowsperstrip=rows_per_strip + ) + tifffile.imwrite( + str(filepath2), np.ones(shape, dtype=np.uint8) * 2, rowsperstrip=rows_per_strip + ) + + parser = virtual_tiff.VirtualTIFF(ifd=0) + registry = ObjectStoreRegistry({"file://": LocalStore()}) + + with ( + open_virtual_dataset( + url=f"file://{filepath1}", parser=parser, registry=registry + ) as vds1, + open_virtual_dataset( + url=f"file://{filepath2}", parser=parser, registry=registry + ) as vds2, + ): + # Verify both datasets have the expected shape + assert vds1["0"].sizes == {"y": 100, "x": 50} + assert vds2["0"].sizes == {"y": 100, "x": 50} + + # Verify both datasets have rectilinear chunk grids + assert isinstance( + vds1["0"].variable.data.metadata.chunk_grid, RectilinearChunkGrid + ) + assert isinstance( + vds2["0"].variable.data.metadata.chunk_grid, RectilinearChunkGrid + ) + + # Concatenate along a new dimension + combined = xr.concat([vds1, vds2], dim="time") + + assert isinstance(combined, Dataset) + assert combined["0"].sizes == {"time": 2, "y": 100, "x": 50} diff --git a/virtualizarr/tests/test_writers/conftest.py b/virtualizarr/tests/test_writers/conftest.py index 3490618df..982ba6c70 100644 --- a/virtualizarr/tests/test_writers/conftest.py +++ b/virtualizarr/tests/test_writers/conftest.py @@ -76,6 +76,36 @@ def synthetic_vds(tmpdir: Path): return vds, arr +@pytest.fixture() +def synthetic_vds_rectilinear_grid(tmpdir: Path, array_v3_metadata_rectilinear): + """A 1D virtual dataset with a rectilinear (variable-length) chunk grid: three + chunks of sizes 2, 1, 3 covering a 6-element array.""" + filepath = f"{tmpdir}/data_chunk" + store = obstore.store.LocalStore() + arr = np.arange(6, dtype=" int: - return array.shape[axis] // array.chunks[axis] + sizes = chunk_grid_sizes(array.metadata)[axis] + if isinstance(sizes, int): + return array.shape[axis] // sizes + return len(sizes) def resize_array( @@ -460,7 +470,38 @@ def resize_array( ) -> None: new_shape = list(arr.shape) new_shape[append_axis] += manifest_array.shape[append_axis] - arr.resize(tuple(new_shape)) + + existing_grid = ChunkGrid.from_metadata(arr.metadata) + new_grid = manifest_array.chunk_grid + stays_regular = ( + existing_grid.is_regular + and new_grid.is_regular + and existing_grid.chunk_shape[append_axis] == new_grid.chunk_shape[append_axis] + ) + if stays_regular: + arr.resize(tuple(new_shape)) + return + + # The append axis's chunk sizes genuinely differ (or one side is already + # rectilinear), so the merged result needs a rectilinear chunk grid - exactly + # like concatenate(), but updating an existing on-disk array's metadata instead + # of building a fresh in-memory one. zarr's own Array.resize() only ever adds a + # single new edge covering the whole size increase, which would be wrong here if + # manifest_array itself has more than one chunk along this axis, so the merged + # edges are computed and written explicitly instead. + require_rectilinear_chunks_enabled(f"Appending along axis {append_axis}") + + old_edges = full_chunk_edges(arr.metadata) + new_edges = full_chunk_edges(manifest_array.metadata) + merged_chunks = list(old_edges) + merged_chunks[append_axis] = old_edges[append_axis] + new_edges[append_axis] + + new_metadata = copy_and_replace_metadata( + old_metadata=cast(ArrayV3Metadata, arr.metadata), + new_shape=new_shape, + new_chunks=merged_chunks, + ) + sync(save_metadata(arr.store_path, new_metadata)) def get_axis( @@ -481,14 +522,28 @@ def check_compatible_arrays( arrays: List[Union[ManifestArray, Array]] = [ma, existing_array] check_same_dtypes([arr.dtype for arr in arrays]) check_same_codecs([get_codecs(arr) for arr in arrays]) - check_same_chunk_shapes([arr.metadata.chunks for arr in arrays]) check_same_ndims([ma.ndim, existing_array.ndim]) + + # Check shapes before chunk shapes, so a genuine shape mismatch is reported as + # such rather than as a chunk-shape one. This matters for a region write in + # particular: ma there is deliberately a much smaller tile than existing_array + # (the full destination array), and their shapes are never expected to match. arr_shapes = [ma.shape, existing_array.shape] if append_axis is not None: check_same_shapes_except_on_concat_axis(arr_shapes, append_axis) if except_axes is not None: check_same_shapes_except_axes(arr_shapes, except_axes) + # Compare declared (not shape-expanded) chunk sizes: ma and existing_array can + # have very different shapes for a region write, so a shape-dependent form like + # full_chunk_edges would flag the same declared chunk size as a mismatch purely + # because it truncates differently at each array's own boundary. Chunk sizes are + # allowed to differ along the append axis - that's what lets resize_array() + # promote the result to a rectilinear chunk grid. + check_same_chunk_shapes( + [chunk_grid_sizes(arr.metadata) for arr in arrays], exclude_axis=append_axis + ) + def write_virtual_variable_to_icechunk( store: "IcechunkStore", @@ -510,8 +565,10 @@ def write_virtual_variable_to_icechunk( if append_dim and append_dim in dims: # TODO: MRP - zarr, or icechunk zarr, array assignment to a variable doesn't work to point to the same object # for example, if you resize an array, it resizes the array but not the bound variable. - if not isinstance(group[name], Array): + existing_arr = group[name] + if not isinstance(existing_arr, Array): raise ValueError("Expected existing array to be a zarr.core.Array") + append_axis = get_axis(dims, append_dim) # check if arrays can be concatenated @@ -539,6 +596,19 @@ def write_virtual_variable_to_icechunk( raise ValueError( f"Expected {name!r} to be a zarr.core.Array, got {type(existing_array)}" ) + + # Region alignment is checked against a single chunk_size per axis, which has + # no equivalent for a rectilinear axis's irregular chunk boundaries. + if ( + not ma.chunk_grid.is_regular + or not ChunkGrid.from_metadata(existing_array.metadata).is_regular + ): + raise NotImplementedError( + f"Cannot write variable {name!r} to icechunk region {region!r}: " + "region writes are not yet supported for arrays with a rectilinear " + "(variable-length) chunk grid." + ) + check_compatible_arrays( ma, existing_array, @@ -568,10 +638,19 @@ def write_virtual_variable_to_icechunk( else: chunk_offsets = [0 for _ in dims] filters, serializer, compressors = extract_codecs(metadata.inner_codecs) + try: + # For a sharded array, ArrayV3Metadata.chunks is the *inner* chunk shape + # (from the sharding codec) - the shape create_array expects for `chunks` + # alongside `shards`. chunk_grid_sizes gives the *outer*/shard shape + # instead (the manifest's unit), which is wrong here, so only fall back to + # it where .chunks itself doesn't apply (a rectilinear chunk grid). + chunks = metadata.chunks + except NotImplementedError: + chunks = chunk_grid_sizes(metadata) arr = group.require_array( name=name, shape=metadata.shape, - chunks=metadata.chunks, + chunks=chunks, shards=metadata.shards, dtype=metadata.data_type.to_native_dtype(), filters=filters,