fix: avoid unnecessary copy by using PyBackedBytes for fsspec

Cargo.toml

@@ -24,7 +24,7 @@ async-lock = "3.4"
 numpy = "0.29"
 num-traits = "0.2"
 delegate = "0.13"
-omfiles-rs = { git = "https://github.com/open-meteo/rust-omfiles", rev = "cd6f08c097931100943b3677b6eaeac0b37327ee", package = "omfiles", features = ["metadata-tree"] }
+omfiles-rs = { git = "https://github.com/open-meteo/rust-omfiles", rev = "00c582d71b005f5bf3585eda0d8f245eec6f6311", package = "omfiles", features = ["metadata-tree"] }
 om-file-format-sys = { git = "https://github.com/open-meteo/om-file-format", rev = "18bee8578d953bd95ad7de1ae688faa7cf421f7c" }
 
 [features]

src/fsspec_backend.rs

@@ -1,6 +1,7 @@
 use omfiles_rs::traits::{OmFileReaderBackend, OmFileReaderBackendAsync, OmFileWriterBackend};
 use omfiles_rs::OmFilesError;
 use pyo3::prelude::*;
+use pyo3::pybacked::PyBackedBytes;
 use pyo3::types::PyDict;
 use pyo3::Python;
 use pyo3_async_runtimes::async_std::into_future;
@@ -41,14 +42,16 @@ impl AsyncFsSpecBackend {
 }
 
 impl OmFileReaderBackendAsync for AsyncFsSpecBackend {
+    type Bytes = PyBackedBytes;
+
     fn count_async(&self) -> usize {
         self.file_size as usize
     }
 
     // This function calls an async read_bytes method on the Python file object
     // and transforms it into a future that can be awaited
     // This allows us to execute multiple asynchronous operations concurrently
-    async fn get_bytes_async(&self, offset: u64, count: u64) -> Result<Vec<u8>, OmFilesError> {
+    async fn get_bytes_async(&self, offset: u64, count: u64) -> Result<Self::Bytes, OmFilesError> {
         let fut = Python::attach(|py| {
             let bound_fs = self.fs.bind(py);
             // We only use named parameters here, because positional arguments can
@@ -66,9 +69,17 @@ impl OmFileReaderBackendAsync for AsyncFsSpecBackend {
             .await
             .map_err(|e| OmFilesError::DecoderError(format!("Python I/O error {}", e)))?;
 
-        let bytes = Python::attach(|py| bytes_obj.extract::<Vec<u8>>(py))
-            .map_err(|e| OmFilesError::DecoderError(format!("Python I/O error: {}", e)));
-        bytes
+        let bytes = Python::attach(|py| -> PyResult<PyBackedBytes> {
+            bytes_obj.extract::<PyBackedBytes>(py).map_err(Into::into)
+        })
+        .map_err(|e| OmFilesError::DecoderError(format!("Python I/O error: {}", e)))?;
+
+        if bytes.len() != count as usize {
+            return Err(OmFilesError::DecoderError(
+                "Obtained unexpected number of bytes from fsspec".to_string(),
+            ));
+        }
+        Ok(bytes)
     }
 }
 
@@ -101,7 +112,7 @@ impl FsSpecBackend {
 }
 
 impl OmFileReaderBackend for FsSpecBackend {
-    type Bytes<'a> = Vec<u8>;
+    type Bytes<'a> = PyBackedBytes;
 
     fn count(&self) -> usize {
         self.file_size as usize
@@ -117,7 +128,7 @@ impl OmFileReaderBackend for FsSpecBackend {
         offset: u64,
         count: u64,
     ) -> Result<Self::Bytes<'_>, omfiles_rs::OmFilesError> {
-        let bytes = Python::attach(|py| {
+        let bytes = Python::attach(|py| -> PyResult<PyBackedBytes> {
             let bound_fs = self.fs.bind(py);
             // We only use named parameters here, because positional arguments can
             // be different between different implementations of the super class!
@@ -127,7 +138,8 @@ impl OmFileReaderBackend for FsSpecBackend {
             kwargs.set_item("path", &self.path)?;
             bound_fs
                 .call_method("cat_file", (), Some(&kwargs))?
-                .extract::<Vec<u8>>()
+                .extract::<PyBackedBytes>()
+                .map_err(Into::into)
         })
         .map_err(|e| OmFilesError::DecoderError(format!("Python I/O error {}", e)))?;
 

src/reader.rs

@@ -20,15 +20,15 @@ use omfiles_rs::{
 use pyo3::{
     exceptions::{PyRuntimeError, PyValueError},
     prelude::*,
+    pybacked::PyBackedBytes,
     types::PyTuple,
     BoundObject,
 };
 use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods};
 use std::{
-    borrow::Cow,
     collections::HashMap,
     fs::File,
-    ops::Range,
+    ops::{Deref, Range},
     sync::{Arc, RwLock},
 };
 
@@ -687,24 +687,35 @@ enum ReaderBackendImpl {
     FsSpec(FsSpecBackend),
 }
 
+enum ReaderBytes<'a> {
+    Mmap(&'a [u8]),
+    FsSpec(PyBackedBytes),
+}
+
+impl Deref for ReaderBytes<'_> {
+    type Target = [u8];
+
+    fn deref(&self) -> &Self::Target {
+        match self {
+            Self::Mmap(bytes) => bytes,
+            Self::FsSpec(bytes) => bytes,
+        }
+    }
+}
+
 impl OmFileReaderBackend for ReaderBackendImpl {
-    // `Cow` can hold either a borrowed slice or an owned Vec, and it
-    // also implements `Deref<Target=[u8]>`, `Send`, and `Sync`,
-    // satisfying all our trait bounds.
-    type Bytes<'a> = Cow<'a, [u8]>;
+    type Bytes<'a> = ReaderBytes<'a>;
 
     // We must implement `get_bytes` manually to handle the type unification.
     fn get_bytes(&self, offset: u64, count: u64) -> Result<Self::Bytes<'_>, OmFilesError> {
         match self {
             ReaderBackendImpl::Mmap(backend) => {
-                // The mmap backend returns a `&[u8]`. We wrap it in `Cow::Borrowed`.
                 let slice = backend.get_bytes(offset, count)?;
-                Ok(Cow::Borrowed(slice))
+                Ok(ReaderBytes::Mmap(slice))
             }
             ReaderBackendImpl::FsSpec(backend) => {
-                // The fsspec backend returns a `Vec<u8>`. We wrap it in `Cow::Owned`.
-                let vec = backend.get_bytes(offset, count)?;
-                Ok(Cow::Owned(vec))
+                let bytes = backend.get_bytes(offset, count)?;
+                Ok(ReaderBytes::FsSpec(bytes))
             }
         }
     }

src/reader_async.rs

@@ -4,7 +4,6 @@ use crate::{
     typed_array::OmFileTypedArray,
 };
 use async_lock::RwLock;
-use delegate::delegate;
 use num_traits::Zero;
 use numpy::{
     ndarray::{self},
@@ -21,11 +20,16 @@ use omfiles_rs::{
 use pyo3::{
     exceptions::{PyRuntimeError, PyValueError},
     prelude::*,
+    pybacked::PyBackedBytes,
     types::PyTuple,
     BoundObject,
 };
 use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods};
-use std::{fs::File, ops::Range, sync::Arc};
+use std::{
+    fs::File,
+    ops::{Deref, Range},
+    sync::Arc,
+};
 
 /// An OmFileReaderAsync class for reading .om files asynchronously.
 ///
@@ -617,14 +621,46 @@ enum AsyncReaderBackendImpl {
     Mmap(MmapFile),
 }
 
+enum AsyncReaderBytes {
+    FsSpec(PyBackedBytes),
+    Mmap(Vec<u8>),
+}
+
+impl Deref for AsyncReaderBytes {
+    type Target = [u8];
+
+    fn deref(&self) -> &Self::Target {
+        match self {
+            Self::FsSpec(bytes) => bytes,
+            Self::Mmap(bytes) => bytes,
+        }
+    }
+}
+
 impl OmFileReaderBackendAsync for AsyncReaderBackendImpl {
-    delegate! {
-        to match self {
-            AsyncReaderBackendImpl::Mmap(backend) => backend,
-            AsyncReaderBackendImpl::FsSpec(backend) => backend,
-        } {
-            fn count_async(&self) -> usize;
-            async fn get_bytes_async(&self, offset: u64, count: u64) -> Result<Vec<u8>, omfiles_rs::OmFilesError>;
+    type Bytes = AsyncReaderBytes;
+
+    fn count_async(&self) -> usize {
+        match self {
+            Self::FsSpec(backend) => backend.count_async(),
+            Self::Mmap(backend) => backend.count_async(),
+        }
+    }
+
+    async fn get_bytes_async(
+        &self,
+        offset: u64,
+        count: u64,
+    ) -> Result<Self::Bytes, omfiles_rs::OmFilesError> {
+        match self {
+            Self::FsSpec(backend) => backend
+                .get_bytes_async(offset, count)
+                .await
+                .map(AsyncReaderBytes::FsSpec),
+            Self::Mmap(backend) => backend
+                .get_bytes_async(offset, count)
+                .await
+                .map(AsyncReaderBytes::Mmap),
         }
     }
 }

tests/test_fsspec.py

@@ -2,6 +2,7 @@
 import os
 import tempfile
 import threading
+from pathlib import Path
 
 import fsspec
 import numpy as np