Include all V2 codecs as filters (#713)

Author: maxrjones

conftest.py

@@ -29,6 +29,11 @@ def pytest_addoption(parser):
         action="store_true",
         help="runs tests requiring a network connection",
     )
+    parser.addoption(
+        "--run-slow-tests",
+        action="store_true",
+        help="runs slow tests",
+    )
     parser.addoption(
         "--run-minio-tests",
         action="store_true",
@@ -44,6 +49,8 @@ def pytest_runtest_setup(item):
         )
     if "minio" in item.keywords and not item.config.getoption("--run-minio-tests"):
         pytest.skip("set --run-minio-tests to run tests requiring docker and minio")
+    if "slow" in item.keywords and not item.config.getoption("--run-slow-tests"):
+        pytest.skip("set --run-slow-tests to run slow tests")
 
 
 def _xarray_subset():

docs/index.md

@@ -103,7 +103,7 @@ vds = open_virtual_mfdataset(urls, parser = parser, registry = registry)
 print(vds)
 ```
 
-The magic of VirtualiZarr is that you can persist the virtual dataset to disk in a chunk references format such as [Icechunk][https://icechunk.io/],
+The magic of VirtualiZarr is that you can persist the virtual dataset to disk in a chunk references format such as [Icechunk](https://icechunk.io/),
 meaning that the work of constructing the single coherent dataset only needs to happen once.
 For subsequent data access, you can use [xarray.open_zarr][] to open that Icechunk store, which on object storage is
 far faster than using [xarray.open_mfdataset][] to open the the original non-cloud-optimized files.

pyproject.toml

@@ -178,6 +178,7 @@ run-mypy = { cmd = "mypy virtualizarr" }
 # Using '--dist loadscope' (rather than default of '--dist load' when '-n auto'
 # is used), reduces test hangs that appear to be macOS-related.
 run-tests = { cmd = "pytest -n auto --dist loadscope --run-network-tests --verbose --durations=10" }
+run-tests-including-slow = { cmd = "pytest -n auto --dist loadscope --run-network-tests --run-slow-tests --verbose --durations=10" }
 run-tests-no-network = { cmd = "pytest -n auto --verbose" }
 run-tests-cov = { cmd = "pytest -n auto --run-network-tests --verbose --cov=virtualizarr --cov=term-missing" }
 run-tests-xml-cov = { cmd = "pytest -n auto --run-network-tests --verbose --cov=virtualizarr --cov-report=xml" }
@@ -224,6 +225,7 @@ show_error_codes = true
 module = [
     "docker",
     "fsspec.*",
+    "s3fs.*",
     "h5py",
     "kerchunk.*",
     "minio",
@@ -305,6 +307,7 @@ markers = [
     # this warning: "PytestUnknownMarkWarning: Unknown pytest.mark.flaky"
     "flaky: flaky tests",
     "network: marks test requiring internet (select with '--run-network-tests')",
+    "slow: marks test as slow (select with '--run-slow-tests')",
     "minio: marks test requiring docker and minio (select with '--run-minio-tests')",
 ]
 filterwarnings = [

virtualizarr/codecs.py

@@ -20,10 +20,11 @@
 ]
 
 
-def numcodec_config_to_configurable(num_codec: dict) -> dict:
+def zarr_codec_config_to_v3(num_codec: dict) -> dict:
     """
     Convert a numcodecs codec into a zarr v3 configurable.
     """
+    # TODO: Special case Blosc codec
     if num_codec["id"].startswith("numcodecs."):
         return num_codec
 
@@ -32,6 +33,19 @@ def numcodec_config_to_configurable(num_codec: dict) -> dict:
     return {"name": name, "configuration": num_codec_copy}
 
 
+def zarr_codec_config_to_v2(num_codec: dict) -> dict:
+    """
+    Convert a numcodecs codec into a zarr v2 configurable.
+    """
+    # TODO: Special case Blosc codec
+    if name := num_codec.get("name", None):
+        return {"id": name, **num_codec["configuration"]}
+    elif num_codec.get("id", None):
+        return num_codec
+    else:
+        raise ValueError(f"Expected a valid Zarr V2 or V3 codec dict, got {num_codec}")
+
+
 def extract_codecs(
     codecs: CodecPipeline,
 ) -> DeconstructedCodecPipeline:

virtualizarr/parsers/hdf/hdf.py

@@ -8,7 +8,7 @@
 
 import numpy as np
 
-from virtualizarr.codecs import numcodec_config_to_configurable
+from virtualizarr.codecs import zarr_codec_config_to_v3
 from virtualizarr.manifests import (
     ChunkEntry,
     ChunkManifest,
@@ -71,9 +71,7 @@ def _construct_manifest_array(
         encoded_cf_fill_value = encode_cf_fill_value(attrs["_FillValue"], dtype)
         attrs["_FillValue"] = encoded_cf_fill_value
 
-    codec_configs = [
-        numcodec_config_to_configurable(codec.get_config()) for codec in codecs
-    ]
+    codec_configs = [zarr_codec_config_to_v3(codec.get_config()) for codec in codecs]
 
     fill_value = dataset.fillvalue.item()
     dims = tuple(_dataset_dims(dataset, group=group))

virtualizarr/parsers/kerchunk/translator.py

@@ -9,7 +9,7 @@
 from zarr.core.metadata import ArrayV3Metadata
 
 from virtualizarr.codecs import (
-    numcodec_config_to_configurable,
+    zarr_codec_config_to_v3,
 )
 from virtualizarr.manifests import (
     ChunkManifest,
@@ -65,9 +65,7 @@ def from_kerchunk_refs(decoded_arr_refs_zarray, zattrs) -> "ArrayV3Metadata":
 
     # Ensure compressor is a list before unpacking
     codec_configs = [*filters, *(compressor if compressor is not None else [])]
-    numcodec_configs = [
-        numcodec_config_to_configurable(config) for config in codec_configs
-    ]
+    numcodec_configs = [zarr_codec_config_to_v3(config) for config in codec_configs]
     dimension_names = decoded_arr_refs_zarray["dimension_names"]
     return create_v3_array_metadata(
         chunk_shape=tuple(decoded_arr_refs_zarray["chunks"]),

virtualizarr/tests/__init__.py

@@ -5,6 +5,7 @@
 
 requires_network = pytest.mark.network
 requires_minio = pytest.mark.minio
+slow_test = pytest.mark.slow
 
 
 def _importorskip(

virtualizarr/tests/test_integration.py

@@ -7,7 +7,7 @@
 import pytest
 import xarray as xr
 import xarray.testing as xrt
-from obstore.store import LocalStore
+from obstore.store import LocalStore, from_url
 
 from conftest import ARRAYBYTES_CODEC, ZLIB_CODEC
 from virtualizarr import open_virtual_dataset
@@ -25,7 +25,9 @@
     has_icechunk,
     has_kerchunk,
     requires_kerchunk,
+    requires_network,
     requires_zarr_python,
+    slow_test,
 )
 
 icechunk = pytest.importorskip("icechunk")
@@ -519,3 +521,35 @@ def test_convert_relative_paths_to_urls(self, netcdf4_file, local_registry):
             path = manifest["0.0.0"]["path"]
 
             assert path == expected_path
+
+
+@requires_kerchunk
+@requires_network
+@slow_test
+def test_roundtrip_dataset_with_multiple_compressors():
+    # Regression test to make sure we can load data with a compression and a shuffle codec
+    # TODO: Simplify this test to not require network access
+    import s3fs
+
+    bucket = "s3://nex-gddp-cmip6"
+    path = "NEX-GDDP-CMIP6/ACCESS-CM2/ssp126/r1i1p1f1/tasmax/tasmax_day_ACCESS-CM2_ssp126_r1i1p1f1_gn_2015_v2.0.nc"
+    url = f"{bucket}/{path}"
+    store = from_url(bucket, region="us-west-2", skip_signature=True)
+    registry = ObjectStoreRegistry({bucket: store})
+    parser = HDFParser()
+    vds = open_virtual_dataset(
+        url=url, parser=parser, registry=registry, loadable_variables=[]
+    )
+
+    ds_refs = vds.vz.to_kerchunk(format="dict")
+    fs = s3fs.S3FileSystem(anon=True)
+    with (
+        xr.open_dataset(fs.open(url), engine="h5netcdf", decode_times=True) as expected,
+        xr.open_dataset(
+            ds_refs,
+            decode_times=True,
+            engine="kerchunk",
+            storage_options={"remote_options": {"anon": True}},
+        ) as observed,
+    ):
+        xr.testing.assert_allclose(expected, observed)

virtualizarr/tests/test_parsers/test_dmrpp.py

@@ -10,7 +10,7 @@
 from virtualizarr.parsers import DMRPPParser, HDFParser
 from virtualizarr.parsers.dmrpp import DMRParser
 from virtualizarr.registry import ObjectStoreRegistry
-from virtualizarr.tests import requires_network
+from virtualizarr.tests import requires_network, slow_test
 from virtualizarr.tests.utils import obstore_local, obstore_s3
 from virtualizarr.xarray import open_virtual_dataset
 
@@ -346,6 +346,7 @@ def dmrparser(dmrpp_xml_str: str, filepath: str) -> DMRParser:
     return DMRParser(root=ET.fromstring(dmrpp_xml_str), data_filepath=filepath)
 
 
+@slow_test
 @requires_network
 @pytest.mark.parametrize("data_url, dmrpp_url", urls)
 def test_NASA_dmrpp(data_url, dmrpp_url):
@@ -373,6 +374,7 @@ def test_NASA_dmrpp(data_url, dmrpp_url):
 
 
 @requires_network
+@slow_test
 @pytest.mark.parametrize("data_url, dmrpp_url", urls)
 def test_NASA_dmrpp_load(data_url, dmrpp_url):
     store = obstore_s3(

virtualizarr/tests/test_writers/test_kerchunk.py

@@ -214,7 +214,7 @@ def testconvert_v3_to_v2_metadata(array_v3_metadata):
     assert v2_metadata.dtype.to_native_dtype() == np.dtype("int32")
     assert v2_metadata.chunks == chunks
     assert v2_metadata.fill_value == 0
-    compressor_config = v2_metadata.compressor.get_config()
+    compressor_config = v2_metadata.filters[1].get_config()
     assert compressor_config["id"] == "blosc"
     assert compressor_config["cname"] == "zstd"
     assert compressor_config["clevel"] == 5