Write virtual references to Icechunk

virtualizarr/tests/test_writers/conftest.py

@@ -0,0 +1,27 @@
+import numpy as np
+import pytest
+from xarray import Dataset
+from xarray.core.variable import Variable
+
+from virtualizarr.manifests import ChunkManifest, ManifestArray
+
+
+@pytest.fixture
+def vds_with_manifest_arrays() -> Dataset:
+    arr = ManifestArray(
+        chunkmanifest=ChunkManifest(
+            entries={"0.0": dict(path="test.nc", offset=6144, length=48)}
+        ),
+        zarray=dict(
+            shape=(2, 3),
+            dtype=np.dtype("<i8"),
+            chunks=(2, 3),
+            compressor={"id": "zlib", "level": 1},
+            filters=None,
+            fill_value=0,
+            order="C",
+            zarr_format=3,
+        ),
+    )
+    var = Variable(dims=["x", "y"], data=arr, attrs={"units": "km"})
+    return Dataset({"a": var}, attrs={"something": 0})

virtualizarr/tests/test_writers/test_icechunk.py

@@ -0,0 +1,120 @@
+import asyncio
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+import pytest
+
+pytest.importorskip("icechunk")
+
+import numpy as np
+import numpy.testing as npt
+from xarray import Dataset, open_dataset
+from xarray.core.variable import Variable
+from zarr import Array, Group, group
+
+from virtualizarr.manifests import ChunkManifest, ManifestArray
+from virtualizarr.writers.icechunk import dataset_to_icechunk
+from virtualizarr.zarr import ZArray
+
+if TYPE_CHECKING:
+    from icechunk import IcechunkStore
+
+
+@pytest.fixture
+def icechunk_filestore(tmpdir) -> "IcechunkStore":
+    from icechunk import IcechunkStore, StorageConfig
+
+    storage = StorageConfig.filesystem(str(tmpdir))
+
+    # TODO if icechunk exposed a synchronous version of .open then we wouldn't need to use asyncio.run here
+    # TODO is this the correct mode to use?
+    store = asyncio.run(IcechunkStore.open(storage=storage, mode="r+"))
+
+    # TODO instead yield store then store.close() ??
+    return store
+
+
+class TestWriteVirtualRefs:
+    def test_write_new_virtual_variable(
+        self, icechunk_filestore: "IcechunkStore", vds_with_manifest_arrays: Dataset
+    ):
+        vds = vds_with_manifest_arrays
+
+        dataset_to_icechunk(vds, icechunk_filestore)
+
+        # check attrs
+        root_group = group(store=icechunk_filestore)
+        assert isinstance(root_group, Group)
+        assert root_group.attrs == {"something": 0}
+
+        # TODO check against vds, then perhaps parametrize?
+
+        # check array exists
+        assert "a" in root_group
+        arr = root_group["a"]
+        assert isinstance(arr, Array)
+
+        # check array metadata
+        # TODO why doesn't a .zarr_format or .version attribute exist on zarr.Array?
+        # assert arr.zarr_format == 3
+        assert arr.shape == (2, 3)
+        assert arr.chunks == (2, 3)
+        assert arr.dtype == np.dtype("<i8")
+        assert arr.order == "C"
+        assert arr.fill_value == 0
+        # TODO check compressor, filters?
+
+        # check array attrs
+        # TODO somehow this is broken by setting the dimension names???
+        # assert dict(arr.attrs) == {"units": "km"}
+
+        # check dimensions
+        assert arr.attrs["DIMENSION_NAMES"] == ["x", "y"]
+
+    def test_set_single_virtual_ref(
+        self, icechunk_filestore: "IcechunkStore", netcdf4_file: Path
+    ):
+        # TODO kerchunk doesn't work with zarr-python v3 yet so we can't use open_virtual_dataset and icechunk together!
+        # vds = open_virtual_dataset(netcdf4_file, indexes={})
+
+        # instead for now just write out byte ranges explicitly
+        manifest = ChunkManifest(
+            {"0.0": {"path": netcdf4_file, "offset": 15419, "length": 7738000}}
+        )
+        zarray = ZArray(
+            shape=(2920, 25, 53),
+            chunks=(2920, 25, 53),
+            dtype=np.dtype("int16"),
+            compressor=None,
+            filters=None,
+            fill_value=None,
+        )
+        ma = ManifestArray(
+            chunkmanifest=manifest,
+            zarray=zarray,
+        )
+        air = Variable(data=ma, dims=["time", "lat", "lon"])
+        vds = Dataset(
+            {"air": air},
+        )
+
+        dataset_to_icechunk(vds, icechunk_filestore)
+
+        root_group = group(store=icechunk_filestore)
+        air_array = root_group["air"]
+        print(air_array)
+
+        # check chunk references
+        # TODO we can't explicitly check that the path/offset/length is correct because icechunk doesn't yet expose any get_virtual_refs method
+
+        expected_ds = open_dataset(netcdf4_file)
+        expected_air_array = expected_ds["air"].to_numpy()
+        npt.assert_equal(air_array, expected_air_array)
+
+        # note: we don't need to test that committing works, because now we have confirmed
+        # the refs are in the store (even uncommitted) it's icechunk's problem to manage them now.
+
+
+# TODO test writing loadable variables
+
+# TODO roundtripping tests - requires icechunk compatibility with xarray

virtualizarr/tests/test_writers/test_zarr.py

@@ -1,37 +1,14 @@
 import json
 
-import numpy as np
-import pytest
 import xarray.testing as xrt
 from xarray import Dataset
 
-from virtualizarr import ManifestArray, open_virtual_dataset
+from virtualizarr import open_virtual_dataset
 from virtualizarr.backend import FileType
-from virtualizarr.manifests.manifest import ChunkManifest
 from virtualizarr.readers.zarr import metadata_from_zarr_json
 from virtualizarr.writers.zarr import dataset_to_zarr
 
 
-@pytest.fixture
-def vds_with_manifest_arrays() -> Dataset:
-    arr = ManifestArray(
-        chunkmanifest=ChunkManifest(
-            entries={"0.0": dict(path="test.nc", offset=6144, length=48)}
-        ),
-        zarray=dict(
-            shape=(2, 3),
-            dtype=np.dtype("<i8"),
-            chunks=(2, 3),
-            compressor={"id": "zlib", "level": 1},
-            filters=None,
-            fill_value=0,
-            order="C",
-            zarr_format=3,
-        ),
-    )
-    return Dataset({"a": (["x", "y"], arr)}, attrs={"something": 0})
-
-
 def isconfigurable(value: dict) -> bool:
     """
     Several metadata attributes in ZarrV3 use a dictionary with keys "name" : str and "configuration" : dict

virtualizarr/writers/icechunk.py

@@ -0,0 +1,152 @@
+import asyncio
+from typing import TYPE_CHECKING
+
+import numpy as np
+from xarray import Dataset
+from xarray.core.variable import Variable
+from zarr import Group
+
+from virtualizarr.manifests import ChunkManifest, ManifestArray
+from virtualizarr.zarr import encode_dtype
+
+if TYPE_CHECKING:
+    from icechunk import IcechunkStore
+
+
+def dataset_to_icechunk(ds: Dataset, store: "IcechunkStore") -> None:
+    """
+    Write an xarray dataset whose variables wrap ManifestArrays to an Icechunk store.
+
+    Currently requires all variables to be backed by ManifestArray objects.
+
+    Parameters
+    ----------
+    ds: xr.Dataset
+    store: IcechunkStore
+    """
+    from icechunk import IcechunkStore
+
+    if not isinstance(store, IcechunkStore):
+        raise TypeError(f"expected type IcechunkStore, but got type {type(store)}")
+
+    # TODO should we check that the store supports writes at this point?
+
+    # TODO only supports writing to the root group currently
+    # TODO pass zarr_format kwarg?
+    root_group = Group.from_store(store=store)
+
+    # TODO this is Frozen, the API for setting attributes must be something else
+    # root_group.attrs = ds.attrs
+    for k, v in ds.attrs.items():
+        root_group.attrs[k] = v
+
+    # we should be able to write references for each variable concurrently
+    asyncio.run(
+        write_variables_to_icechunk_group(
+            ds.variables,
+            store=store,
+            group=root_group,
+        )
+    )
+
+
+async def write_variables_to_icechunk_group(
+    variables,
+    store,
+    group,
+):
+    await asyncio.gather(
+        *(
+            write_variable_to_icechunk(
+                store=store,
+                group=group,
+                arr_name=name,
+                var=var,
+            )
+            for name, var in variables.items()
+        )
+    )
+
+
+async def write_variable_to_icechunk(
+    store: "IcechunkStore",
+    group: Group,
+    arr_name: str,
+    var: Variable,
+) -> None:
+    if not isinstance(var.data, ManifestArray):
+        # TODO is writing loadable_variables just normal xarray ds.to_zarr?
+        raise NotImplementedError()
+
+    ma = var.data
+    zarray = ma.zarray
+
+    # TODO should I be checking that this array doesn't already exist? Or is that icechunks' job?
+    arr = group.create_array(
+        arr_name,
+        shape=zarray.shape,
+        chunk_shape=zarray.chunks,
+        dtype=encode_dtype(zarray.dtype),
+        # TODO fill_value?
+        # TODO order?
+        # TODO zarr format?
+        # TODO compressors?
+    )
+
+    # TODO it would be nice if we could assign directly to the .attrs property
+    for k, v in var.attrs.items():
+        arr.attrs[k] = v
+    # TODO we should probably be doing some encoding of those attributes?
+    arr.attrs["DIMENSION_NAMES"] = var.dims
+
+    await write_manifest_virtual_refs(
+        store=store,
+        group=group,
+        arr_name=arr_name,
+        manifest=ma.manifest,
+    )
+
+
+async def write_manifest_virtual_refs(
+    store: "IcechunkStore",
+    group: Group,
+    arr_name: str,
+    manifest: ChunkManifest,
+) -> None:
+    """Write all the virtual references for one array manifest at once."""
+
+    key_prefix = f"{group.name}{arr_name}"
+    if key_prefix.startswith("/"):
+        # remove preceding / character
+        # TODO unsure if this is correct
+        key_prefix = key_prefix[1:]
+
+    # loop over every reference in the ChunkManifest for that array
+    # TODO inefficient: this should be replaced with something that sets all (new) references for the array at once
+    # but Icechunk need to expose a suitable API first
+    it = np.nditer(
+        [manifest._paths, manifest._offsets, manifest._lengths],
+        flags=[
+            "refs_ok",
+            "multi_index",
+            "c_index",  # TODO is "c_index" correct? what's the convention for zarr chunk keys?
+        ],
+        op_flags=[["readonly"]] * 3,
+    )
+    for path, offset, length in it:
+        index = it.multi_index
+        chunk_key = "/".join(str(i) for i in index)
+
+        key = f"{key_prefix}/{chunk_key}"  # should be of form 'group/name/0/1/2'
+
+        print(key)
+
+        # TODO this needs to be awaited or something
+        # set each reference individually
+        await store.set_virtual_ref(
+            # TODO it would be marginally neater if I could pass the group and name as separate args
+            key=key,
+            location=path.item(),
+            offset=offset.item(),
+            length=length.item(),
+        )