diff --git a/Cargo.lock b/Cargo.lock index ba2569c..8a10dae 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 8925d6e..7f24dc5 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 3aeb9a3..64ef8e4 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 9e67a3d..0aca831 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 358d081..7254226 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 31edebf..aed5c87 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