implement async load for dataarray and dataset

Author: TomNicholas

xarray/core/dataarray.py

@@ -1160,6 +1160,14 @@ def load(self, **kwargs) -> Self:
         self._coords = new._coords
         return self
 
+    async def async_load(self, **kwargs) -> Self:
+        temp_ds = self._to_temp_dataset()
+        ds = await temp_ds.async_load(**kwargs)
+        new = self._from_temp_dataset(ds)
+        self._variable = new._variable
+        self._coords = new._coords
+        return self
+
     def compute(self, **kwargs) -> Self:
         """Manually trigger loading of this array's data from disk or a
         remote source into memory and return a new array.

xarray/core/dataset.py

@@ -552,6 +552,31 @@ def load(self, **kwargs) -> Self:
 
         return self
 
+    async def async_load(self, **kwargs) -> Self:
+        # this blocks on chunked arrays but not on lazily indexed arrays
+
+        # access .data to coerce everything to numpy or dask arrays
+        lazy_data = {
+            k: v._data for k, v in self.variables.items() if is_chunked_array(v._data)
+        }
+        if lazy_data:
+            chunkmanager = get_chunked_array_type(*lazy_data.values())
+
+            # evaluate all the chunked arrays simultaneously
+            evaluated_data: tuple[np.ndarray[Any, Any], ...] = chunkmanager.compute(
+                *lazy_data.values(), **kwargs
+            )
+
+            for k, data in zip(lazy_data, evaluated_data, strict=False):
+                self.variables[k].data = data
+
+        # load everything else sequentially
+        for k, v in self.variables.items():
+            if k not in lazy_data:
+                await v.async_load()
+
+        return self
+
     def __dask_tokenize__(self) -> object:
         from dask.base import normalize_token
 

xarray/tests/test_async.py

@@ -101,27 +101,70 @@ def memorystore() -> "MemoryStore":
 @requires_zarr_v3
 @pytest.mark.asyncio
 class TestAsyncLoad:
-    async def test_async_load_variable(self, memorystore):
-        N_LOADS = 5
-        LATENCY = 1.0
+    N_LOADS = 10
+    LATENCY = 1.0
 
-        latencystore = LatencyStore(memorystore, latency=LATENCY)
+    # TODO refactor these tests
+    async def test_async_load_variable(self, memorystore):
+        latencystore = LatencyStore(memorystore, latency=self.LATENCY)
         ds = xr.open_zarr(latencystore, zarr_format=3, consolidated=False, chunks=None)
 
         # TODO add async load to Dataset and DataArray as well as to Variable
         # TODO change the syntax to `.async.load()`?
         start_time = time.time()
-        tasks = [ds["foo"].variable.async_load() for _ in range(N_LOADS)]
+        tasks = [ds["foo"].variable.async_load() for _ in range(self.N_LOADS)]
         results = await asyncio.gather(*tasks)
         total_time = time.time() - start_time
 
         for result in results:
             xrt.assert_identical(result, ds["foo"].variable.load())
 
-        assert total_time > LATENCY  # Cannot possibly be quicker than this
+        assert total_time > self.LATENCY  # Cannot possibly be quicker than this
+        assert (
+            total_time < self.LATENCY * self.N_LOADS
+        )  # If this isn't true we're gaining nothing from async
+        assert (
+            abs(total_time - self.LATENCY) < 0.5
+        )  # Should take approximately LATENCY seconds, but allow some buffer
+
+    async def test_async_load_dataarray(self, memorystore):
+        latencystore = LatencyStore(memorystore, latency=self.LATENCY)
+        ds = xr.open_zarr(latencystore, zarr_format=3, consolidated=False, chunks=None)
+
+        # TODO change the syntax to `.async.load()`?
+        start_time = time.time()
+        tasks = [ds["foo"].async_load() for _ in range(self.N_LOADS)]
+        results = await asyncio.gather(*tasks)
+        total_time = time.time() - start_time
+
+        for result in results:
+            xrt.assert_identical(result, ds["foo"].load())
+
+        assert total_time > self.LATENCY  # Cannot possibly be quicker than this
+        assert (
+            total_time < self.LATENCY * self.N_LOADS
+        )  # If this isn't true we're gaining nothing from async
+        assert (
+            abs(total_time - self.LATENCY) < 0.5
+        )  # Should take approximately LATENCY seconds, but allow some buffer
+
+    async def test_async_load_dataset(self, memorystore):
+        latencystore = LatencyStore(memorystore, latency=self.LATENCY)
+        ds = xr.open_zarr(latencystore, zarr_format=3, consolidated=False, chunks=None)
+
+        # TODO change the syntax to `.async.load()`?
+        start_time = time.time()
+        tasks = [ds.async_load() for _ in range(self.N_LOADS)]
+        results = await asyncio.gather(*tasks)
+        total_time = time.time() - start_time
+
+        for result in results:
+            xrt.assert_identical(result, ds.load())
+
+        assert total_time > self.LATENCY  # Cannot possibly be quicker than this
         assert (
-            total_time < LATENCY * N_LOADS
+            total_time < self.LATENCY * self.N_LOADS
         )  # If this isn't true we're gaining nothing from async
         assert (
-            abs(total_time - LATENCY) < 0.5
+            abs(total_time - self.LATENCY) < 0.5
         )  # Should take approximately LATENCY seconds, but allow some buffer