Add methods to create virtual dataset from manifeststore (#522)

TomNicholas · pre-commit-ci[bot] · web-flow · commit e5fff199d716 · 2025-04-03T14:29:11.000-07:00
* improve type hinting by just using Mapping * rename manifest array attribute * add metadata property * make ManifestGroup a Mapping * rename manifest_group -> group * typing * improve the repr * test new group API * sketch out conversion method * more sketching * add to_virtual_variable * add ManifestGroup.to_virtual_dataset() * linting * consolidate duplicated manifest_array fixtures * ensure coordinates are understood * move test to store-level instead of group-level * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * docs * Un-comment out icechunk upstream --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
diff --git a/conftest.py b/conftest.py
@@ -3,7 +3,7 @@
 # Standard library imports
 import itertools
 from pathlib import Path
-from typing import Any, Callable, Mapping, Optional
+from typing import Any, Callable, Iterable, Mapping, Optional
 
 # Third-party imports
 import h5py  # type: ignore[import]
@@ -266,6 +266,8 @@ def _create_metadata(
         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"}]
         return create_v3_array_metadata(
@@ -274,6 +276,8 @@ def _create_metadata(
             data_type=data_type,
             codecs=codecs,
             fill_value=fill_value or 0,
+            attributes=attributes,
+            dimension_names=dimension_names,
         )
 
     return _create_metadata
@@ -288,11 +292,14 @@ def manifest_array(array_v3_metadata):
     """
 
     def _manifest_array(
-        shape: tuple = (5, 5),
-        chunks: tuple = (5, 5),
+        shape: tuple = (5, 2),
+        chunks: tuple = (5, 2),
         codecs: list[dict] | None = [ARRAYBYTES_CODEC, ZLIB_CODEC],
+        dimension_names: Iterable[str] | None = None,
     ):
-        metadata = array_v3_metadata(shape=shape, chunks=chunks, codecs=codecs)
+        metadata = array_v3_metadata(
+            shape=shape, chunks=chunks, codecs=codecs, dimension_names=dimension_names
+        )
         entries = _generate_chunk_entries(shape, chunks, _entry_from_chunk_key)
         chunkmanifest = ChunkManifest(entries=entries)
         return ManifestArray(chunkmanifest=chunkmanifest, metadata=metadata)
diff --git a/docs/releases.rst b/docs/releases.rst
@@ -10,6 +10,8 @@ New Features
 ~~~~~~~~~~~~
 
 - Added experimental ManifestStore (:pull:`490`).
+- Added :py:meth:`ManifestStore.to_virtual_dataset()` method (:pull:`522`).
+  By `Tom Nicholas <https://github.com/TomNicholas>`_.
 - Added experimental :py:func:`open_virtual_mfdataset` function (:issue:`345`, :pull:`349`).
   By `Tom Nicholas <https://github.com/TomNicholas>`_.
 
diff --git a/virtualizarr/common.py b/virtualizarr/common.py
diff --git a/virtualizarr/manifests/array.py b/virtualizarr/manifests/array.py
@@ -2,6 +2,7 @@
 from typing import Any, Callable, Union
 
 import numpy as np
+import xarray as xr
 from zarr.core.metadata.v3 import ArrayV3Metadata, RegularChunkGrid
 
 import virtualizarr.manifests.utils as utils
@@ -277,6 +278,30 @@ def rename_paths(
         renamed_manifest = self.manifest.rename_paths(new)
         return ManifestArray(metadata=self.metadata, chunkmanifest=renamed_manifest)
 
+    def to_virtual_variable(self) -> xr.Variable:
+        """
+        Create a "virtual" xarray.Variable containing the contents of one zarr array.
+
+        The returned variable will be "virtual", i.e. it will wrap a single ManifestArray object.
+        """
+
+        # The xarray data model stores dimension names and arbitrary extra metadata outside of the wrapped array class,
+        # so to avoid that information being duplicated we strip it from the ManifestArray before wrapping it.
+        dims = self.metadata.dimension_names
+        attrs = self.metadata.attributes
+        stripped_metadata = utils.copy_and_replace_metadata(
+            self.metadata, new_dimension_names=None, new_attributes={}
+        )
+        stripped_marr = ManifestArray(
+            chunkmanifest=self.manifest, metadata=stripped_metadata
+        )
+
+        return xr.Variable(
+            data=stripped_marr,
+            dims=dims,
+            attrs=attrs,
+        )
+
 
 def _possibly_expand_trailing_ellipsis(key, ndim: int):
     if key[-1] == ...:
diff --git a/virtualizarr/manifests/group.py b/virtualizarr/manifests/group.py
@@ -3,6 +3,7 @@
 import textwrap
 from typing import Iterator, Mapping
 
+import xarray as xr
 from zarr.core.group import GroupMetadata
 
 from virtualizarr.manifests import ManifestArray
@@ -105,3 +106,28 @@ def __repr__(self) -> str:
             )
             """
         )
+
+    def to_virtual_dataset(self) -> xr.Dataset:
+        """
+        Create a "virtual" xarray.Dataset containing the contents of one zarr group.
+
+        All variables in the returned Dataset will be "virtual", i.e. they will wrap ManifestArray objects.
+        """
+
+        from virtualizarr.common import construct_fully_virtual_dataset
+
+        # The xarray data model stores coordinate names outside of the arbitrary extra metadata it can store on a Dataset,
+        # so to avoid that information being duplicated we strip it from the zarr group attributes before storing it.
+        metadata_dict = self.metadata.to_dict()
+        attributes = metadata_dict["attributes"]
+        coord_names = attributes.pop("coordinates", [])
+
+        virtual_vars = {
+            name: marr.to_virtual_variable() for name, marr in self.arrays.items()
+        }
+
+        return construct_fully_virtual_dataset(
+            virtual_vars=virtual_vars,
+            coord_names=coord_names,
+            attrs=attributes,
+        )
diff --git a/virtualizarr/manifests/store.py b/virtualizarr/manifests/store.py
@@ -5,6 +5,7 @@
 from typing import TYPE_CHECKING, Any, Mapping
 from urllib.parse import urlparse
 
+import xarray as xr
 from zarr.abc.store import (
     ByteRequest,
     OffsetByteRequest,
@@ -351,6 +352,32 @@ async def list_dir(self, prefix: str) -> AsyncGenerator[str, None]:
         for k in self._group.arrays.keys():
             yield k
 
+    def to_virtual_dataset(self, group="") -> xr.Dataset:
+        """
+        Create a "virtual" xarray dataset containing the contents of one zarr group.
+
+        All variables in the returned Dataset will be "virtual", i.e. they will wrap ManifestArray objects.
+
+        Will ignore the contents of any other groups in the store.
+
+        Parameters
+        ----------
+        group : str
+
+        Returns
+        -------
+        vds : xarray.Dataset
+        """
+
+        if group:
+            raise NotImplementedError(
+                "ManifestStore does not yet support nested groups"
+            )
+        else:
+            manifestgroup = self._group
+
+        return manifestgroup.to_virtual_dataset()
+
 
 def _transform_byte_range(
     byte_range: ByteRequest | None, *, chunk_start: int, chunk_end_exclusive: int
diff --git a/virtualizarr/manifests/utils.py b/virtualizarr/manifests/utils.py
@@ -1,9 +1,13 @@
-from typing import TYPE_CHECKING, Any, Dict, Iterable, Optional, Union
+from typing import TYPE_CHECKING, Any, Dict, Iterable, Literal, Optional, Union
 
 import numpy as np
 from zarr import Array
 from zarr.core.chunk_key_encodings import ChunkKeyEncodingLike
-from zarr.core.metadata.v3 import ArrayV3Metadata
+from zarr.core.metadata.v3 import (
+    ArrayV3Metadata,
+    parse_dimension_names,
+    parse_shapelike,
+)
 
 from virtualizarr.codecs import convert_to_codec_pipeline, get_codecs
 
@@ -19,7 +23,7 @@ def create_v3_array_metadata(
     fill_value: Any = None,
     codecs: Optional[list[Dict[str, Any]]] = None,
     attributes: Optional[Dict[str, Any]] = None,
-    dimension_names: Optional[tuple[str, ...]] = None,
+    dimension_names: Iterable[str] | None = None,
 ) -> ArrayV3Metadata:
     """
     Create an ArrayV3Metadata instance with standard configuration.
@@ -198,17 +202,29 @@ def copy_and_replace_metadata(
     old_metadata: ArrayV3Metadata,
     new_shape: list[int] | None = None,
     new_chunks: list[int] | None = None,
+    new_dimension_names: Iterable[str] | None | Literal["default"] = "default",
+    new_attributes: dict | None = None,
 ) -> ArrayV3Metadata:
     """
     Update metadata to reflect a new shape and/or chunk shape.
     """
+    # TODO this should really be upstreamed into zarr-python
+
     metadata_copy = old_metadata.to_dict().copy()
-    metadata_copy["shape"] = new_shape  # type: ignore[assignment]
+
+    if new_shape is not None:
+        metadata_copy["shape"] = parse_shapelike(new_shape)  # type: ignore[assignment]
     if new_chunks is not None:
         metadata_copy["chunk_grid"] = {
             "name": "regular",
             "configuration": {"chunk_shape": 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)
+    if new_attributes is not None:
+        metadata_copy["attributes"] = new_attributes
+
     # ArrayV3Metadata.from_dict removes extra keys zarr_format and node_type
     new_metadata = ArrayV3Metadata.from_dict(metadata_copy)
     return new_metadata
diff --git a/virtualizarr/readers/fits.py b/virtualizarr/readers/fits.py
@@ -3,10 +3,10 @@
 
 from xarray import Dataset, Index
 
+from virtualizarr.common import construct_fully_virtual_dataset
 from virtualizarr.readers.api import (
     VirtualBackend,
 )
-from virtualizarr.readers.common import construct_fully_virtual_dataset
 from virtualizarr.translators.kerchunk import (
     extract_group,
     virtual_vars_and_metadata_from_kerchunk_refs,
diff --git a/virtualizarr/readers/hdf/hdf.py b/virtualizarr/readers/hdf/hdf.py
@@ -19,6 +19,10 @@
 from xarray.backends.zarr import FillValueCoder
 
 from virtualizarr.codecs import numcodec_config_to_configurable
+from virtualizarr.common import (
+    construct_fully_virtual_dataset,
+    replace_virtual_with_loadable_vars,
+)
 from virtualizarr.manifests import (
     ChunkEntry,
     ChunkManifest,
@@ -27,10 +31,6 @@
 from virtualizarr.manifests.manifest import validate_and_normalize_path_to_uri
 from virtualizarr.manifests.utils import create_v3_array_metadata
 from virtualizarr.readers.api import VirtualBackend
-from virtualizarr.readers.common import (
-    construct_fully_virtual_dataset,
-    replace_virtual_with_loadable_vars,
-)
 from virtualizarr.readers.hdf.filters import cfcodec_from_dataset, codecs_from_dataset
 from virtualizarr.types import ChunkKey
 from virtualizarr.utils import _FsspecFSFromFilepath, soft_import
diff --git a/virtualizarr/readers/hdf5.py b/virtualizarr/readers/hdf5.py
@@ -3,11 +3,11 @@
 
 from xarray import Dataset, Index
 
-from virtualizarr.readers.api import VirtualBackend
-from virtualizarr.readers.common import (
+from virtualizarr.common import (
     construct_fully_virtual_dataset,
     replace_virtual_with_loadable_vars,
 )
+from virtualizarr.readers.api import VirtualBackend
 from virtualizarr.translators.kerchunk import (
     extract_group,
     virtual_vars_and_metadata_from_kerchunk_refs,
diff --git a/virtualizarr/readers/netcdf3.py b/virtualizarr/readers/netcdf3.py
@@ -3,11 +3,11 @@
 
 from xarray import Dataset, Index
 
-from virtualizarr.readers.api import VirtualBackend
-from virtualizarr.readers.common import (
+from virtualizarr.common import (
     construct_fully_virtual_dataset,
     replace_virtual_with_loadable_vars,
 )
+from virtualizarr.readers.api import VirtualBackend
 from virtualizarr.translators.kerchunk import (
     virtual_vars_and_metadata_from_kerchunk_refs,
 )
diff --git a/virtualizarr/readers/tiff.py b/virtualizarr/readers/tiff.py
@@ -4,11 +4,11 @@
 
 from xarray import Dataset, Index
 
-from virtualizarr.readers.api import VirtualBackend
-from virtualizarr.readers.common import (
+from virtualizarr.common import (
     construct_fully_virtual_dataset,
     replace_virtual_with_loadable_vars,
 )
+from virtualizarr.readers.api import VirtualBackend
 from virtualizarr.translators.kerchunk import (
     extract_group,
     virtual_vars_and_metadata_from_kerchunk_refs,
diff --git a/virtualizarr/tests/test_manifests/test_array.py b/virtualizarr/tests/test_manifests/test_array.py
@@ -384,3 +384,27 @@ def test_refuse_combine(array_v3_metadata):
     for func in [np.concatenate, np.stack]:
         with pytest.raises(ValueError, match="inconsistent dtypes"):
             func([marr1, marr2], axis=0)
+
+
+def test_to_xarray(array_v3_metadata):
+    chunks = (5, 10)
+    shape = (5, 20)
+    metadata = array_v3_metadata(
+        shape=shape,
+        chunks=chunks,
+        attributes={"ham": "sandwich"},
+        dimension_names=["x", "y"],
+    )
+    chunks_dict = {
+        "0.0": {"path": "/foo.nc", "offset": 100, "length": 100},
+        "0.1": {"path": "/foo.nc", "offset": 200, "length": 100},
+    }
+    manifest = ChunkManifest(entries=chunks_dict)
+    marr = ManifestArray(metadata=metadata, chunkmanifest=manifest)
+
+    vv = marr.to_virtual_variable()
+    assert isinstance(vv.data, ManifestArray)
+    assert vv.dims == ("x", "y")
+    assert vv.data.metadata.dimension_names is None
+    assert vv.attrs == {"ham": "sandwich"}
+    assert vv.data.metadata.attributes == {}
diff --git a/virtualizarr/tests/test_manifests/test_group.py b/virtualizarr/tests/test_manifests/test_group.py
@@ -3,30 +3,16 @@
 import pytest
 from zarr.core.group import GroupMetadata
 
-from virtualizarr.manifests import ChunkManifest, ManifestArray, ManifestGroup
-
-
-@pytest.fixture
-def manifest_array(array_v3_metadata):
-    chunk_dict = {
-        "0.0.0": {"path": "s3://bucket/foo.nc", "offset": 100, "length": 100},
-        "0.0.1": {"path": "s3://bucket/foo.nc", "offset": 200, "length": 100},
-        "0.1.0": {"path": "s3://bucket/foo.nc", "offset": 300, "length": 100},
-        "0.1.1": {"path": "s3://bucket/foo.nc", "offset": 400, "length": 100},
-    }
-    manifest = ChunkManifest(entries=chunk_dict)
-    chunks = (5, 1, 10)
-    shape = (5, 2, 20)
-    array_metadata = array_v3_metadata(shape=shape, chunks=chunks)
-    return ManifestArray(metadata=array_metadata, chunkmanifest=manifest)
+from virtualizarr.manifests import ManifestArray, ManifestGroup
 
 
 class TestManifestGroup:
     def test_group_containing_array(self, manifest_array):
         var = "foo"
-        manifest_group = ManifestGroup(arrays={var: manifest_array}, attributes={})
+        marr = manifest_array()
+        manifest_group = ManifestGroup(arrays={var: marr}, attributes={})
 
-        assert manifest_group.arrays == {var: manifest_array}
+        assert manifest_group.arrays == {var: marr}
         assert manifest_group.groups == {}
         assert isinstance(manifest_group[var], ManifestArray)
         with pytest.raises(KeyError):
@@ -36,11 +22,12 @@ def test_group_containing_array(self, manifest_array):
         assert list(manifest_group) == [var]
 
     def test_manifest_repr(self, manifest_array):
-        manifest_group = ManifestGroup(arrays={"foo": manifest_array}, attributes={})
+        marr = manifest_array(shape=(5, 2), chunks=(5, 2))
+        manifest_group = ManifestGroup(arrays={"foo": marr}, attributes={})
         expected_repr = textwrap.dedent(
             """
             ManifestGroup(
-                arrays={'foo': ManifestArray<shape=(5, 2, 20), dtype=int32, chunks=(5, 1, 10)>},
+                arrays={'foo': ManifestArray<shape=(5, 2), dtype=int32, chunks=(5, 2)>},
                 groups={},
                 metadata=GroupMetadata(attributes={}, zarr_format=3, consolidated_metadata=None, node_type='group'),
             )
diff --git a/virtualizarr/tests/test_manifests/test_store.py b/virtualizarr/tests/test_manifests/test_store.py
diff --git a/virtualizarr/tests/test_writers/test_icechunk.py b/virtualizarr/tests/test_writers/test_icechunk.py
diff --git a/virtualizarr/translators/kerchunk.py b/virtualizarr/translators/kerchunk.py
