Add ManifestStore for loading data from ManifestArrays

ci/environment.yml

@@ -40,3 +40,4 @@ dependencies:
   - zarr>=3.0.2
   - pip:
       - imagecodecs-numcodecs==2024.6.1
+      - obstore @ git+https://github.com/developmentseed/obstore.git@main#subdirectory=obstore

docs/releases.rst

@@ -9,6 +9,8 @@ v1.3.3 (unreleased)
 New Features
 ~~~~~~~~~~~~
 
+- Added experimental VirtualObjectStore for loading data directly from virtual datasets.
+
 Breaking changes
 ~~~~~~~~~~~~~~~~
 

pyproject.toml

@@ -37,7 +37,9 @@ remote = [
     "aiohttp",
     "s3fs",
 ]
-
+obstore = [
+    "obstore @ git+https://github.com/developmentseed/obstore@main#subdirectory=obstore",
+]
 # non-kerchunk-based readers
 hdf = [
     "virtualizarr[remote]",
@@ -95,7 +97,8 @@ test = [
     "pytest",
     "ruff",
     "s3fs",
-    "virtualizarr[all]"
+    "virtualizarr[obstore]",
+    "virtualizarr[all]",
 ]
 
 [project.urls]

virtualizarr/accessor.py

@@ -1,15 +1,21 @@
+from __future__ import annotations
+
 from datetime import datetime
 from pathlib import Path
 from typing import TYPE_CHECKING, Callable, Literal, overload
 
-from xarray import Dataset, register_dataset_accessor
+from xarray import Dataset, open_dataset, register_dataset_accessor
 
 from virtualizarr.manifests import ManifestArray
+from virtualizarr.storage._obstore import VirtualObjectStore
 from virtualizarr.types.kerchunk import KerchunkStoreRefs
 from virtualizarr.writers.kerchunk import dataset_to_kerchunk_refs
 
 if TYPE_CHECKING:
     from icechunk import IcechunkStore  # type: ignore[import-not-found]
+    from obstore.store import (
+        ObjectStore as _UpstreamObjectStore,  # type: ignore[import-not-found]
+    )
 
 
 @register_dataset_accessor("virtualize")
@@ -88,6 +94,43 @@ def to_icechunk(
             last_updated_at=last_updated_at,
         )
 
+    def to_xarray(self, stores: dict[str, _UpstreamObjectStore], **kwargs) -> Dataset:
+        """
+        Convert the virtual dataset to an VirtualObjectStore-backed xarray dataset which can
+        be used to load data directly without first serializing using Icechunk or Kerchunk.
+
+        Parameters
+        ----------
+        stores: dict
+            Object stores to use for fetching data. The keys should provide a prefix that can
+            be matched to urls contained within the virtual dataset (e.g., `file://`,
+            `s3://mybucket`, `s3://mybucket1`) and the values should be properly configured
+            ObjectStore instances. A store must be provided for each storage type and
+            configuration set (e.g., LocalStore, public S3Store, private S3Store).
+
+        kwargs: dict, optional
+            Additional keyword-arguments to pass to xarray.open_dataset().
+
+
+        Examples
+        --------
+
+        >>> from obstore.store import LocalStore
+        >>> stores = {"file://": LocalStore()}
+        >>> vds.virtualize.to_xarray(stores=stores)  # doctest: +SKIP
+
+        Warnings
+        --------
+        The to_xarray() functionality is experimental.
+        """
+        store = VirtualObjectStore(self.ds, stores)
+        return open_dataset(
+            store,  # type: ignore[arg-type]
+            engine="zarr",
+            backend_kwargs={"zarr_format": 3, "consolidated": False},
+            **kwargs,
+        )
+
     @overload
     def to_kerchunk(
         self, filepath: None, format: Literal["dict"]

virtualizarr/storage/_obstore.py

@@ -0,0 +1,246 @@
+from __future__ import annotations
+
+import pickle
+from collections.abc import Iterable
+from typing import TYPE_CHECKING, Any
+
+from xarray import DataArray, Dataset
+from zarr.abc.store import (
+    ByteRequest,
+    OffsetByteRequest,
+    RangeByteRequest,
+    Store,
+    SuffixByteRequest,
+)
+from zarr.core.buffer import Buffer
+from zarr.core.buffer.core import BufferPrototype
+
+from virtualizarr.storage.common import (
+    find_matching_store,
+    get_zarr_metadata,
+    list_dir_from_xr_obj,
+    parse_manifest_index,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncGenerator, Iterable
+    from typing import Any
+
+    from obstore.store import ObjectStore as _UpstreamObjectStore
+    from zarr.core.buffer import BufferPrototype
+    from zarr.core.common import BytesLike
+
+
+__all__ = ["VirtualObjectStore"]
+
+_ALLOWED_EXCEPTIONS: tuple[type[Exception], ...] = (
+    FileNotFoundError,
+    IsADirectoryError,
+    NotADirectoryError,
+)
+
+
+class VirtualObjectStore(Store):
+    """A Zarr store that uses obstore for fast read/write from AWS, GCP, Azure.
+
+    Parameters
+    ----------
+    stores : dict[prefix, obstore.store.ObjectStore]
+        A mapping of url prefixes to obstore Store instances set up with the proper credentials.
+
+    Warnings
+    --------
+    VirtualObjectStore is experimental and subject to API changes without notice. Please
+    raise an issue with any comments/concerns about the store.
+
+    Notes
+    -----
+    Modified from https://github.com/zarr-developers/zarr-python/pull/1661
+    """
+
+    def __eq__(self, value: object):
+        NotImplementedError
+
+    def __init__(
+        self,
+        xr_obj: DataArray | Dataset,
+        stores: dict[str, _UpstreamObjectStore],
+    ) -> None:
+        import obstore as obs
+
+        # TODO: Support DataArray, Dataset, or DataTree across all methods
+        for store in stores.values():
+            if not isinstance(
+                store,
+                (
+                    obs.store.AzureStore,
+                    obs.store.GCSStore,
+                    obs.store.HTTPStore,
+                    obs.store.S3Store,
+                    obs.store.LocalStore,
+                    obs.store.MemoryStore,
+                ),
+            ):
+                raise TypeError(f"expected ObjectStore class, got {store!r}")
+        super().__init__(read_only=True)
+        self.stores = stores
+        self.xr_obj: DataArray | Dataset = xr_obj
+
+    def __str__(self) -> str:
+        return f"ManifesStore({self.xr_obj})"
+
+    def __getstate__(self) -> dict[Any, Any]:
+        state = self.__dict__.copy()
+        stores = state["stores"]
+        for k, v in stores:
+            stores[k] = pickle.dumps(v)
+        state["stores"] = stores
+        return state
+
+    def __setstate__(self, state: dict[Any, Any]) -> None:
+        stores = state["stores"]
+        for k, v in stores:
+            stores[k] = pickle.loads(v)
+        state["stores"] = stores
+        self.__dict__.update(state)
+
+    async def get(
+        self,
+        key: str,
+        prototype: BufferPrototype,
+        byte_range: ByteRequest | None = None,
+    ) -> Buffer | None:
+        # docstring inherited
+        import obstore as obs
+
+        if "zarr.json" in key:
+            return get_zarr_metadata(self.xr_obj, key)
+        manifest_index = parse_manifest_index(key)
+        # Get path, offset, and length matching this key from the ChunkManifest
+        if manifest_index.variable == "__xarray_dataarray_variable__":
+            path = self.xr_obj.data.manifest._paths[*manifest_index.indexes]
+            offset = self.xr_obj.data.manifest._offsets[*manifest_index.indexes]
+            length = self.xr_obj.data.manifest._lengths[*manifest_index.indexes]
+        else:
+            path = self.xr_obj[manifest_index.variable].data.manifest._paths[
+                *manifest_index.indexes
+            ]
+            offset = self.xr_obj[manifest_index.variable].data.manifest._offsets[
+                *manifest_index.indexes
+            ]
+            length = self.xr_obj[manifest_index.variable].data.manifest._lengths[
+                *manifest_index.indexes
+            ]
+        store_request = find_matching_store(stores=self.stores, request_key=path)
+        # Transform the input byte range to account for the chunk location in the file
+        chunk_end_exclusive = offset + length
+        byte_range = _transform_byte_range(
+            byte_range, chunk_start=offset, chunk_end_exclusive=chunk_end_exclusive
+        )
+        # Actually get the bytes
+        try:
+            bytes = await obs.get_range_async(
+                self.stores[store_request.store_id],
+                store_request.key,
+                start=byte_range.start,
+                end=byte_range.end,
+            )
+            return prototype.buffer.from_bytes(bytes)  # type: ignore[arg-type]
+        except _ALLOWED_EXCEPTIONS:
+            return None
+
+    async def get_partial_values(
+        self,
+        prototype: BufferPrototype,
+        key_ranges: Iterable[tuple[str, ByteRequest | None]],
+    ) -> list[Buffer | None]:
+        # docstring inherited
+        raise NotImplementedError
+
+    async def exists(self, key: str) -> bool:
+        # docstring inherited
+        raise NotImplementedError
+
+    @property
+    def supports_writes(self) -> bool:
+        # docstring inherited
+        return False
+
+    async def set(self, key: str, value: Buffer) -> None:
+        # docstring inherited
+        raise NotImplementedError
+
+    async def set_if_not_exists(self, key: str, value: Buffer) -> None:
+        # docstring inherited
+        raise NotImplementedError
+
+    @property
+    def supports_deletes(self) -> bool:
+        # docstring inherited
+        return False
+
+    async def delete(self, key: str) -> None:
+        raise NotImplementedError
+
+    @property
+    def supports_partial_writes(self) -> bool:
+        # docstring inherited
+        return False
+
+    async def set_partial_values(
+        self, key_start_values: Iterable[tuple[str, int, BytesLike]]
+    ) -> None:
+        # docstring inherited
+        raise NotImplementedError
+
+    @property
+    def supports_listing(self) -> bool:
+        # docstring inherited
+        return True
+
+    def list(self) -> AsyncGenerator[str, None]:
+        # docstring inherited
+        raise NotImplementedError
+
+    def list_prefix(self, prefix: str) -> AsyncGenerator[str, None]:
+        # docstring inherited
+        raise NotImplementedError
+
+    def list_dir(self, prefix: str) -> AsyncGenerator[str, None]:
+        # docstring inherited
+        if isinstance(self.xr_obj, (DataArray, Dataset)):
+            return list_dir_from_xr_obj(self.xr_obj, prefix)
+        else:
+            raise NotImplementedError(
+                "Only DataArray and Datasets are currently supported"
+            )
+
+
+def _transform_byte_range(
+    byte_range: ByteRequest | None, *, chunk_start: int, chunk_end_exclusive: int
+) -> RangeByteRequest:
+    """
+    Convert an incoming byte_range which assumes one chunk per file to a
+    virtual byte range that accounts for the location of a chunk within a file.
+    """
+    if byte_range is None:
+        byte_range = RangeByteRequest(chunk_start, chunk_end_exclusive)
+    elif isinstance(byte_range, RangeByteRequest):
+        if byte_range.end > chunk_end_exclusive:
+            raise ValueError(
+                f"Chunk ends before byte {chunk_end_exclusive} but request end was {byte_range.end}"
+            )
+        byte_range = RangeByteRequest(
+            chunk_start + byte_range.start, chunk_start + byte_range.end
+        )
+    elif isinstance(byte_range, OffsetByteRequest):
+        byte_range = RangeByteRequest(
+            chunk_start + byte_range.offset, chunk_end_exclusive
+        )  # type: ignore[arg-type]
+    elif isinstance(byte_range, SuffixByteRequest):
+        byte_range = RangeByteRequest(
+            chunk_end_exclusive - byte_range.suffix, chunk_end_exclusive
+        )  # type: ignore[arg-type]
+    else:
+        raise ValueError(f"Unexpected byte_range, got {byte_range}")
+    return byte_range