Avoid memory copy in obstore write (#2972)

* Avoid memory copy in obstore write

* Add as_bytes_like method to Buffer

* Add changelog entry

* No need to take unsigned bytes view following #2738

* Change method name to `as_buffer_like`

---------

Co-authored-by: jakirkham <[email protected]>
Co-authored-by: Davis Bennett <[email protected]>

Author: 3 people

changes/2972.misc.rst

@@ -0,0 +1 @@
+Avoid an unnecessary memory copy when writing Zarr with obstore

src/zarr/core/buffer/core.py

@@ -255,6 +255,19 @@ def as_numpy_array(self) -> npt.NDArray[Any]:
         """
         ...
 
+    def as_buffer_like(self) -> BytesLike:
+        """Returns the buffer as an object that implements the Python buffer protocol.
+
+        Notes
+        -----
+        Might have to copy data, since the implementation uses `.as_numpy_array()`.
+
+        Returns
+        -------
+            An object that implements the Python buffer protocol
+        """
+        return memoryview(self.as_numpy_array())  # type: ignore[arg-type]
+
     def to_bytes(self) -> bytes:
         """Returns the buffer as `bytes` (host memory).
 

src/zarr/storage/_local.py

@@ -52,10 +52,10 @@ def _put(
         with path.open("r+b") as f:
             f.seek(start)
             # write takes any object supporting the buffer protocol
-            f.write(value.as_numpy_array())  # type: ignore[arg-type]
+            f.write(value.as_buffer_like())
         return None
     else:
-        view = memoryview(value.as_numpy_array())  # type: ignore[arg-type]
+        view = value.as_buffer_like()
         if exclusive:
             mode = "xb"
         else:

src/zarr/storage/_obstore.py

@@ -161,15 +161,15 @@ async def set(self, key: str, value: Buffer) -> None:
 
         self._check_writable()
 
-        buf = value.to_bytes()
+        buf = value.as_buffer_like()
         await obs.put_async(self.store, key, buf)
 
     async def set_if_not_exists(self, key: str, value: Buffer) -> None:
         # docstring inherited
         import obstore as obs
 
         self._check_writable()
-        buf = value.to_bytes()
+        buf = value.as_buffer_like()
         with contextlib.suppress(obs.exceptions.AlreadyExistsError):
             await obs.put_async(self.store, key, buf, mode="create")