From c3e078ea0d2d6bbaad97f80e2e79c74ab86ca9ad Mon Sep 17 00:00:00 2001 From: terraputix Date: Wed, 22 Jul 2026 19:58:55 +0200 Subject: [PATCH 1/2] fix: avoid unneccessary copy by using PyBackedBytes for fsspec --- Cargo.lock | 2 +- Cargo.toml | 2 +- src/fsspec_backend.rs | 26 +++++++++++++++------ src/reader.rs | 33 +++++++++++++++++--------- src/reader_async.rs | 54 +++++++++++++++++++++++++++++++++++-------- tests/test_fsspec.py | 20 ++++++++++++++++ 6 files changed, 108 insertions(+), 29 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index ba2569cd..8a10daeb 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -810,7 +810,7 @@ dependencies = [ [[package]] name = "omfiles" version = "0.1.0" -source = "git+https://github.com/open-meteo/rust-omfiles?rev=cd6f08c097931100943b3677b6eaeac0b37327ee#cd6f08c097931100943b3677b6eaeac0b37327ee" +source = "git+https://github.com/open-meteo/rust-omfiles?rev=00c582d71b005f5bf3585eda0d8f245eec6f6311#00c582d71b005f5bf3585eda0d8f245eec6f6311" dependencies = [ "async-executor", "async-lock", diff --git a/Cargo.toml b/Cargo.toml index 8925d6ec..7f24dc52 100644 --- a/Cargo.toml +++ b/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] diff --git a/src/fsspec_backend.rs b/src/fsspec_backend.rs index 3aeb9a39..64ef8e49 100644 --- a/src/fsspec_backend.rs +++ b/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,6 +42,8 @@ impl AsyncFsSpecBackend { } impl OmFileReaderBackendAsync for AsyncFsSpecBackend { + type Bytes = PyBackedBytes; + fn count_async(&self) -> usize { self.file_size as usize } @@ -48,7 +51,7 @@ impl OmFileReaderBackendAsync for AsyncFsSpecBackend { // 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, OmFilesError> { + async fn get_bytes_async(&self, offset: u64, count: u64) -> Result { 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::>(py)) - .map_err(|e| OmFilesError::DecoderError(format!("Python I/O error: {}", e))); - bytes + let bytes = Python::attach(|py| -> PyResult { + bytes_obj.extract::(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; + 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, omfiles_rs::OmFilesError> { - let bytes = Python::attach(|py| { + let bytes = Python::attach(|py| -> PyResult { 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::>() + .extract::() + .map_err(Into::into) }) .map_err(|e| OmFilesError::DecoderError(format!("Python I/O error {}", e)))?; diff --git a/src/reader.rs b/src/reader.rs index 9e67a3d3..0aca831a 100644 --- a/src/reader.rs +++ b/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`, `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, 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`. 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)) } } } diff --git a/src/reader_async.rs b/src/reader_async.rs index 358d0816..72542261 100644 --- a/src/reader_async.rs +++ b/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), +} + +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, 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 { + 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), } } } diff --git a/tests/test_fsspec.py b/tests/test_fsspec.py index 31edebf6..31eec124 100644 --- a/tests/test_fsspec.py +++ b/tests/test_fsspec.py @@ -2,6 +2,7 @@ import os import tempfile import threading +from pathlib import Path import fsspec import numpy as np @@ -126,6 +127,25 @@ async def test_s3_read_async(s3_backend_async, s3_test_file): assert np.isfinite(data).all() +@pytest.mark.asyncio +async def test_async_fsspec_read_from_python_bytes(temp_om_file): + class AsyncBytesFileSystem: + def __init__(self, data: bytes): + self.data = data + + async def _size(self, path): + return len(self.data) + + async def _cat_file(self, path, start=None, end=None): + return self.data[start:end] + + fs = AsyncBytesFileSystem(Path(temp_om_file).read_bytes()) + reader = await omfiles.OmFileReaderAsync.from_fsspec(fs, "test.om") + data = await reader.read_array((slice(0, 5), slice(0, 5))) + + np.testing.assert_array_equal(data, np.arange(25).reshape(5, 5)) + + @filter_numpy_size_warning def test_s3_xarray(s3_spatial_test_file): # The way described in the xarray documentation does not really use the caching mechanism From d3430ac1875ce5b76d7dad00bf18aa166e9665b2 Mon Sep 17 00:00:00 2001 From: terraputix Date: Wed, 22 Jul 2026 20:01:17 +0200 Subject: [PATCH 2/2] remove useless test --- tests/test_fsspec.py | 19 ------------------- 1 file changed, 19 deletions(-) diff --git a/tests/test_fsspec.py b/tests/test_fsspec.py index 31eec124..aed5c87e 100644 --- a/tests/test_fsspec.py +++ b/tests/test_fsspec.py @@ -127,25 +127,6 @@ async def test_s3_read_async(s3_backend_async, s3_test_file): assert np.isfinite(data).all() -@pytest.mark.asyncio -async def test_async_fsspec_read_from_python_bytes(temp_om_file): - class AsyncBytesFileSystem: - def __init__(self, data: bytes): - self.data = data - - async def _size(self, path): - return len(self.data) - - async def _cat_file(self, path, start=None, end=None): - return self.data[start:end] - - fs = AsyncBytesFileSystem(Path(temp_om_file).read_bytes()) - reader = await omfiles.OmFileReaderAsync.from_fsspec(fs, "test.om") - data = await reader.read_array((slice(0, 5), slice(0, 5))) - - np.testing.assert_array_equal(data, np.arange(25).reshape(5, 5)) - - @filter_numpy_size_warning def test_s3_xarray(s3_spatial_test_file): # The way described in the xarray documentation does not really use the caching mechanism