Repository: zarr-developers/VirtualiZarr

Issue:
KeyError: '.zarray' with HDF5 data
I'm getting a KeyError Traceback for several different HDF5 files. 

For example, trying to use this test file from kerchunk:
https://github.com/fsspec/kerchunk/blob/ae692fead51a216691e4db9a67c99194c5ba8e14/kerchunk/tests/test_hdf.py#L307

```python
# (or local path)
url = 'https://github.com/fsspec/kerchunk/raw/ae692fead51a216691e4db9a67c99194c5ba8e14/kerchunk/tests/NEONDSTowerTemperatureData.hdf5'
kerchunk.hdf.SingleHdf5ToZarr(url).translate()

# side note: also not sure why `reader_options={}` is required to read from the URL
vds = open_virtual_dataset(url, reader_options={})
```

```pytb
File ~/GitHub/VirtualiZarr/virtualizarr/xarray.py:114, in open_virtual_dataset(filepath, filetype, drop_variables, loadable_variables, indexes, virtual_array_class, reader_options)
    106 else:
    107     # this is the only place we actually always need to use kerchunk directly
    108     # TODO avoid even reading byte ranges for variables that will be dropped later anyway?
    109     vds_refs = kerchunk.read_kerchunk_references_from_file(
    110         filepath=filepath,
    111         filetype=filetype,
    112         reader_options=reader_options,
    113     )
--> 114     virtual_vars = virtual_vars_from_kerchunk_refs(
    115         vds_refs,
    116         drop_variables=drop_variables + loadable_variables,
    117         virtual_array_class=virtual_array_class,
    118     )
    119     ds_attrs = kerchunk.fully_decode_arr_refs(vds_refs["refs"]).get(".zattrs", {})
    120     coord_names = ds_attrs.pop("coordinates", [])

File ~/GitHub/VirtualiZarr/virtualizarr/xarray.py:247, in virtual_vars_from_kerchunk_refs(refs, drop_variables, virtual_array_class)
    241     drop_variables = []
    242 var_names_to_keep = [
    243     var_name for var_name in var_names if var_name not in drop_variables
    244 ]
    246 vars = {
--> 247     var_name: variable_from_kerchunk_refs(refs, var_name, virtual_array_class)
    248     for var_name in var_names_to_keep
    249 }
    250 return vars

File ~/GitHub/VirtualiZarr/virtualizarr/xarray.py:293, in variable_from_kerchunk_refs(refs, var_name, virtual_array_class)
    290 """Create a single xarray Variable by reading specific keys of a kerchunk references dict."""
    292 arr_refs = kerchunk.extract_array_refs(refs, var_name)
--> 293 chunk_dict, zarray, zattrs = kerchunk.parse_array_refs(arr_refs)
    295 manifest = ChunkManifest._from_kerchunk_chunk_dict(chunk_dict)
    297 # we want to remove the _ARRAY_DIMENSIONS from the final variables' .attrs

File ~/GitHub/VirtualiZarr/virtualizarr/kerchunk.py:186, in parse_array_refs(arr_refs)
    183 def parse_array_refs(
    184     arr_refs: KerchunkArrRefs,
    185 ) -> tuple[dict, ZArray, ZAttrs]:
--> 186     zarray = ZArray.from_kerchunk_refs(arr_refs.pop(".zarray"))
    187     zattrs = arr_refs.pop(".zattrs", {})
    188     chunk_dict = arr_refs

KeyError: '.zarray'
```



Discussion:
scottyhq: > One other thing is that the group kwarg could also be supported when reading references from Zarr, but that could also just be flagged as an issue and left for a later PR, as the logic will be independent of the kerchunk-parsing logic here. 

Yeah, sounds like a plan. I think this is now refactored and minimally functional. I'm sure working with other example datasets  with groups (https://github.com/pydata/xarray-data/blob/master/cmip6.nc) will uncover improvements or the need for additional tests.

Patch:
diff --git a/docs/releases.rst b/docs/releases.rst
index 39d517b..216b3b8 100644
--- a/docs/releases.rst
+++ b/docs/releases.rst
@@ -8,6 +8,8 @@ v1.0.1 (unreleased)
 
 New Features
 ~~~~~~~~~~~~
+- New ``group`` option on ``open_virtual_dataset`` enables extracting specific HDF Groups.
+  (:pull:`165`) By `Scott Henderson <https://github.com/scottyhq>`_.
 
 - Adds `decode_times` to open_virtual_dataset (:pull:`232`)
   By `Raphael Hagen <https://github.com/norlandrhagen>`_.
diff --git a/virtualizarr/backend.py b/virtualizarr/backend.py
index b155c2c..73fabe7 100644
--- a/virtualizarr/backend.py
+++ b/virtualizarr/backend.py
@@ -51,6 +51,7 @@ def open_virtual_dataset(
     filepath: str,
     *,
     filetype: FileType | None = None,
+    group: str | None = None,
     drop_variables: Iterable[str] | None = None,
     loadable_variables: Iterable[str] | None = None,
     decode_times: bool | None = None,
@@ -74,6 +75,8 @@ def open_virtual_dataset(
         Type of file to be opened. Used to determine which kerchunk file format backend to use.
         Can be one of {'netCDF3', 'netCDF4', 'HDF', 'TIFF', 'GRIB', 'FITS', 'zarr_v3'}.
         If not provided will attempt to automatically infer the correct filetype from header bytes.
+    group : str, default is None
+        Path to the HDF5/netCDF4 group in the given file to open. Given as a str, supported by filetypes “netcdf4” and “hdf5”.
     drop_variables: list[str], default is None
         Variables in the file to drop before returning.
     loadable_variables: list[str], default is None
@@ -171,6 +174,7 @@ def open_virtual_dataset(
         vds_refs = read_kerchunk_references_from_file(
             filepath=filepath,
             filetype=filetype,
+            group=group,
             reader_options=reader_options,
         )
         virtual_vars = virtual_vars_from_kerchunk_refs(
@@ -195,6 +199,7 @@ def open_virtual_dataset(
             ds = xr.open_dataset(
                 cast(XArrayOpenT, fpath),
                 drop_variables=drop_variables,
+                group=group,
                 decode_times=decode_times,
             )
 
diff --git a/virtualizarr/readers/kerchunk.py b/virtualizarr/readers/kerchunk.py
index 4686ce9..19a8c28 100644
--- a/virtualizarr/readers/kerchunk.py
+++ b/virtualizarr/readers/kerchunk.py
@@ -57,6 +57,7 @@ def _automatically_determine_filetype(
 def read_kerchunk_references_from_file(
     filepath: str,
     filetype: FileType | None,
+    group: str | None,
     reader_options: Optional[dict[str, Any]] = None,
 ) -> KerchunkStoreRefs:
     """
@@ -69,7 +70,8 @@ def read_kerchunk_references_from_file(
     filetype : FileType, default: None
         Type of file to be opened. Used to determine which kerchunk file format backend to use.
         If not provided will attempt to automatically infer the correct filetype from the the filepath's extension.
-    reader_options: dict, default {}
+    group : str, default is None
+        Path to the HDF5/netCDF4 group in the given file to open. Given as a str, supported by filetypes “netcdf4” and “hdf5”.
         Dict passed into Kerchunk file readers. Note: Each Kerchunk file reader has distinct arguments,
         so ensure reader_options match selected Kerchunk reader arguments.
     """
@@ -96,6 +98,9 @@ def read_kerchunk_references_from_file(
         refs = SingleHdf5ToZarr(
             filepath, inline_threshold=0, **reader_options
         ).translate()
+
+        refs = extract_group(refs, group)
+
     elif filetype.name.lower() == "grib":
         # TODO Grib files should be handled as a DataTree object
         # see https://github.com/TomNicholas/VirtualiZarr/issues/11
@@ -123,6 +128,44 @@ def read_kerchunk_references_from_file(
     return refs
 
 
+def extract_group(vds_refs: KerchunkStoreRefs, group: str | None) -> KerchunkStoreRefs:
+    """Extract only the part of the kerchunk reference dict that is relevant to a single HDF group"""
+    hdf_groups = [
+        k.removesuffix(".zgroup") for k in vds_refs["refs"].keys() if ".zgroup" in k
+    ]
+    if len(hdf_groups) == 1:
+        return vds_refs
+    else:
+        if group is None:
+            raise ValueError(
+                f"Multiple HDF Groups found. Must specify group= keyword to select one of {hdf_groups}"
+            )
+        else:
+            # Ensure supplied group kwarg is consistent with kerchunk keys
+            if not group.endswith("/"):
+                group += "/"
+            if group.startswith("/"):
+                group = group.removeprefix("/")
+
+        if group not in hdf_groups:
+            raise ValueError(f'Group "{group}" not found in {hdf_groups}')
+
+        # Filter by group prefix and remove prefix from all keys
+        groupdict = {
+            k.removeprefix(group): v
+            for k, v in vds_refs["refs"].items()
+            if k.startswith(group)
+        }
+        # Also remove group prefix from _ARRAY_DIMENSIONS
+        for k, v in groupdict.items():
+            if isinstance(v, str):
+                groupdict[k] = v.replace("\\/", "/").replace(group, "")
+
+        vds_refs["refs"] = groupdict
+
+        return KerchunkStoreRefs(vds_refs)
+
+
 def virtual_vars_from_kerchunk_refs(
     refs: KerchunkStoreRefs,
     drop_variables: list[str] | None = None,
@@ -214,6 +257,7 @@ def find_var_names(ds_reference_dict: KerchunkStoreRefs) -> list[str]:
 
     refs = ds_reference_dict["refs"]
     found_var_names = {key.split("/")[0] for key in refs.keys() if "/" in key}
+
     return list(found_var_names)
 
 
@@ -235,6 +279,7 @@ def extract_array_refs(
         }
 
         return fully_decode_arr_refs(arr_refs)
+
     else:
         raise KeyError(
             f"Could not find zarr array variable name {var_name}, only {found_var_names}"
diff --git a/virtualizarr/types/kerchunk.py b/virtualizarr/types/kerchunk.py
index e8dada2..d124cca 100644
--- a/virtualizarr/types/kerchunk.py
+++ b/virtualizarr/types/kerchunk.py
@@ -4,9 +4,10 @@ from typing import NewType
 # (idea from https://kobzol.github.io/rust/python/2023/05/20/writing-python-like-its-rust.html)
 # TODO I would prefer to be more specific about these types
 KerchunkStoreRefs = NewType(
-    "KerchunkStoreRefs", dict
-)  # top-level dict with keys for 'version', 'refs'
+    "KerchunkStoreRefs",
+    dict,  # dict_keys(['version', 'refs'])
+)  # top-level dict containing kerchunk version and 'refs' dictionary which assumes single '.zgroup' key and multiple KerchunkArrRefs
 KerchunkArrRefs = NewType(
     "KerchunkArrRefs",
-    dict,
-)  # lower-level dict containing just the information for one zarr array
+    dict,  # dict_keys(['.zarray', '.zattrs', '0.0', '0.1', ...)
+)  # lower-level dict defining a single Zarr Array, with keys for '.zarray', '.zattrs', and every chunk


Test patch:
diff --git a/conftest.py b/conftest.py
index b558abf..32b3581 100644
--- a/conftest.py
+++ b/conftest.py
@@ -33,6 +33,19 @@ def netcdf4_file(tmpdir):
     return filepath
 
 
+@pytest.fixture
+def hdf5_groups_file(tmpdir):
+    # Set up example xarray dataset
+    ds = xr.tutorial.open_dataset("air_temperature")
+
+    # Save it to disk as netCDF (in temporary directory)
+    filepath = f"{tmpdir}/air.nc"
+    ds.to_netcdf(filepath, format="NETCDF4", group="test/group")
+    ds.close()
+
+    return filepath
+
+
 @pytest.fixture
 def netcdf4_files(tmpdir):
     # Set up example xarray dataset
diff --git a/virtualizarr/tests/test_backend.py b/virtualizarr/tests/test_backend.py
index c300158..e42ad9a 100644
--- a/virtualizarr/tests/test_backend.py
+++ b/virtualizarr/tests/test_backend.py
@@ -1,6 +1,7 @@
 from collections.abc import Mapping
 from unittest.mock import patch
 
+import fsspec
 import numpy as np
 import pytest
 import xarray as xr
@@ -192,6 +193,10 @@ class TestReadFromURL:
                 "hdf4",
                 "https://github.com/corteva/rioxarray/raw/master/test/test_data/input/MOD09GA.A2008296.h14v17.006.2015181011753.hdf",
             ),
+            (
+                "hdf5",
+                "https://nisar.asf.earthdatacloud.nasa.gov/NISAR-SAMPLE-DATA/GCOV/ALOS1_Rosamond_20081012/NISAR_L2_PR_GCOV_001_005_A_219_4020_SHNA_A_20081012T060910_20081012T060926_P01101_F_N_J_001.h5",
+            ),
             # https://github.com/zarr-developers/VirtualiZarr/issues/159
             # ("hdf5", "https://github.com/fsspec/kerchunk/raw/main/kerchunk/tests/NEONDSTowerTemperatureData.hdf5"),
             pytest.param(
@@ -218,10 +223,48 @@ class TestReadFromURL:
         if filetype in ["grib", "jpg", "hdf4"]:
             with pytest.raises(NotImplementedError):
                 vds = open_virtual_dataset(url, reader_options={}, indexes={})
+        elif filetype == "hdf5":
+            vds = open_virtual_dataset(
+                url,
+                group="science/LSAR/GCOV/grids/frequencyA",
+                drop_variables=["listOfCovarianceTerms", "listOfPolarizations"],
+                indexes={},
+                reader_options={},
+            )
+            assert isinstance(vds, xr.Dataset)
         else:
             vds = open_virtual_dataset(url, indexes={})
             assert isinstance(vds, xr.Dataset)
 
+    def test_virtualizarr_vs_local_nisar(self):
+        # Open group directly from locally cached file with xarray
+        url = "https://nisar.asf.earthdatacloud.nasa.gov/NISAR-SAMPLE-DATA/GCOV/ALOS1_Rosamond_20081012/NISAR_L2_PR_GCOV_001_005_A_219_4020_SHNA_A_20081012T060910_20081012T060926_P01101_F_N_J_001.h5"
+        tmpfile = fsspec.open_local(
+            f"filecache::{url}", filecache=dict(cache_storage="/tmp", same_names=True)
+        )
+        hdf_group = "science/LSAR/GCOV/grids/frequencyA"
+        dsXR = xr.open_dataset(
+            tmpfile,
+            engine="h5netcdf",
+            group=hdf_group,
+            drop_variables=["listOfCovarianceTerms", "listOfPolarizations"],
+            phony_dims="access",
+        )
+
+        # save group reference file via virtualizarr, then open with engine="kerchunk"
+        vds = open_virtual_dataset(
+            tmpfile,
+            group=hdf_group,
+            indexes={},
+            drop_variables=["listOfCovarianceTerms", "listOfPolarizations"],
+        )
+        tmpref = "/tmp/cmip6.json"
+        vds.virtualize.to_kerchunk(tmpref, format="json")
+        dsV = xr.open_dataset(tmpref, engine="kerchunk")
+
+        # xrt.assert_identical(dsXR, dsV) #Attribute order changes
+        xrt.assert_equal(dsXR, dsV)
+
 
 class TestLoadVirtualDataset:
     def test_loadable_variables(self, netcdf4_file):
@@ -249,6 +292,27 @@ class TestLoadVirtualDataset:
         with pytest.raises(NotImplementedError):
             open_virtual_dataset(netcdf4_file, filetype="grib")
 
+    def test_group_kwarg(self, hdf5_groups_file):
+        with pytest.raises(ValueError, match="Multiple HDF Groups found"):
+            open_virtual_dataset(hdf5_groups_file)
+        with pytest.raises(ValueError, match="not found in"):
+            open_virtual_dataset(hdf5_groups_file, group="doesnt_exist")
+
+        vars_to_load = ["air", "time"]
+        vds = open_virtual_dataset(
+            hdf5_groups_file,
+            group="test/group",
+            loadable_variables=vars_to_load,
+            indexes={},
+        )
+        full_ds = xr.open_dataset(
+            hdf5_groups_file,
+            group="test/group",
+        )
+        for name in full_ds.variables:
+            if name in vars_to_load:
+                xrt.assert_identical(vds.variables[name], full_ds.variables[name])
+
     @patch("virtualizarr.readers.kerchunk.read_kerchunk_references_from_file")
     def test_open_virtual_dataset_passes_expected_args(
         self, mock_read_kerchunk, netcdf4_file
@@ -258,6 +322,7 @@ class TestLoadVirtualDataset:
         args = {
             "filepath": netcdf4_file,
             "filetype": None,
+            "group": None,
             "reader_options": reader_options,
         }
         mock_read_kerchunk.assert_called_once_with(**args)