Proof of concept: lazily loading datasets with VirtualiZarr

virtualizarr/__init__.py

@@ -4,7 +4,11 @@
     VirtualiZarrDatasetAccessor,
     VirtualiZarrDataTreeAccessor,
 )
-from virtualizarr.xarray import open_virtual_dataset, open_virtual_mfdataset
+from virtualizarr.xarray import (
+    open_dataset,
+    open_virtual_dataset,
+    open_virtual_mfdataset,
+)
 
 try:
     __version__ = _version("virtualizarr")
@@ -18,4 +22,5 @@
     "VirtualiZarrDataTreeAccessor",
     "open_virtual_dataset",
     "open_virtual_mfdataset",
+    "open_dataset",
 ]

virtualizarr/tests/test_xarray.py

@@ -11,6 +11,7 @@
 from xarray import Dataset, open_dataset
 from xarray.core.indexes import Index
 
+from virtualizarr import open_dataset as vz_open_dataset
 from virtualizarr import open_virtual_dataset, open_virtual_mfdataset
 from virtualizarr.manifests import ChunkManifest, ManifestArray
 from virtualizarr.parsers import HDFParser
@@ -24,6 +25,18 @@
 from virtualizarr.tests.utils import obstore_http, obstore_local, obstore_s3
 
 
+def test_open_dataset(netcdf4_file: str):
+    store = obstore_local(file_url=netcdf4_file)
+    parser = HDFParser()
+
+    observed = vz_open_dataset(netcdf4_file, object_store=store, parser=parser)
+    observed_value = observed.air.isel(lat=10, lon=10, time=10).data
+    expected = xr.open_dataset(netcdf4_file, decode_timedelta=True, engine="netcdf4")
+    expected_value = expected.air.isel(lat=10, lon=10, time=10).data
+    xr.testing.assert_allclose(observed, expected)
+    np.testing.assert_allclose(observed_value, expected_value)
+
+
 def test_wrapping(array_v3_metadata):
     chunks = (5, 10)
     shape = (5, 20)

virtualizarr/xarray.py

@@ -34,6 +34,31 @@
     )
 
 
+def open_dataset(
+    file_url: str,
+    object_store: ObjectStore,
+    parser: Parser,
+    drop_variables: Iterable[str] | None = None,
+    decode_times: bool | None = None,
+    cftime_variables: Iterable[str] | None = None,
+) -> xr.Dataset:
+    filepath = validate_and_normalize_path_to_uri(file_url, fs_root=Path.cwd().as_uri())
+
+    manifest_store = parser(
+        file_url=filepath,
+        object_store=object_store,
+    )
+
+    ds = xr.open_dataset(
+        manifest_store,
+        decode_times=decode_times,
+        engine="zarr",
+        consolidated=False,
+        zarr_format=3,
+    )
+    return ds.drop_vars(list(drop_variables or ()))
+
+
 def open_virtual_dataset(
     file_url: str,
     object_store: ObjectStore,