perf: use pinned memory with preload_to_gpu

docs/conf.py

@@ -100,7 +100,7 @@
     "scipy": ("https://docs.scipy.org/doc/scipy", None),
     "cupy": ("https://docs.cupy.dev/en/stable/", None),
     "zarrs": ("https://zarrs-python.readthedocs.io/en/latest/", None),
-    "pandas": ("https://pandas.pydata.org/pandas-docs/stable", None),
+    "pandas": ("https://pandas.pydata.org/pandas-docs/version/2.3", None),
     "h5py": ("https://docs.h5py.org/en/latest", None),
 }
 

src/annbatch/loader.py

@@ -554,6 +554,15 @@ def _slices_to_slices_with_array_index(
             dataset_index_to_slices_sorted[k] = dataset_index_to_slices[k]
         return dataset_index_to_slices_sorted
 
+    def _get_kwargs_for_zarr_fetching(self, z: zarr.Array, indexer_shape: tuple[int, ...]) -> dict:
+        buffer_prototype = zarr.core.buffer.default_buffer_prototype()
+        kwargs = {"prototype": buffer_prototype}
+        if self._preload_to_gpu:
+            import cupyx as cpx
+
+            kwargs["out"] = buffer_prototype.nd_buffer(cpx.empty_pinned(indexer_shape, z.dtype))
+        return kwargs
+
     @singledispatchmethod
     async def _fetch_data(self, dataset: ZarrArray | CSRDatasetElems, slices: list[slice]) -> InputInMemoryArray:
         """Fetch data from an on-disk store.
@@ -590,7 +599,9 @@ async def _fetch_data_dense(self, dataset: ZarrArray, slices: list[slice]) -> np
         )
         res = cast(
             "np.ndarray",
-            await dataset._async_array._get_selection(indexer, prototype=zarr.core.buffer.default_buffer_prototype()),
+            await dataset._async_array._get_selection(
+                indexer, **self._get_kwargs_for_zarr_fetching(dataset, indexer.shape)
+            ),
         )
         return res
 
@@ -665,9 +676,12 @@ async def _fetch_data_sparse(
                 for l in indptr_limits
             ]
         )
+
         data_np, indices_np = await asyncio.gather(
-            data._get_selection(indexer, prototype=zarr.core.buffer.default_buffer_prototype()),
-            indices._get_selection(indexer, prototype=zarr.core.buffer.default_buffer_prototype()),
+            *(
+                z._get_selection(indexer, **self._get_kwargs_for_zarr_fetching(z, indexer.shape))
+                for z in [data, indices]
+            )
         )
         gaps = (s1.start - s0.stop for s0, s1 in pairwise(indptr_limits))
         offsets = accumulate(chain([indptr_limits[0].start], gaps))