diff --git a/Cargo.toml b/Cargo.toml index e3454bec..9828703e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -11,6 +11,10 @@ crate-type = ["cdylib", "rlib"] [dependencies] pyo3 = { version = "0.27.1", features = ["abi3-py311"] } zarrs = { version = "0.23.6", features = ["async", "zlib", "pcodec", "bz2"] } +# `AsyncStoragePartialDecoder` is `pub` in `zarrs_codec` but not re-exported by +# `zarrs`; depend on it directly. Cargo unifies it to the same version `zarrs` +# uses, so the codec types are identical to those re-exported via `zarrs`. +zarrs_codec = { version = "0.2", features = ["async"] } rayon_iter_concurrent_limit = "0.2.0" rayon = "1.10.0" # fix for https://stackoverflow.com/questions/76593417/package-openssl-was-not-found-in-the-pkg-config-search-path @@ -20,9 +24,11 @@ unsafe_cell_slice = "0.2.0" serde_json = "1.0.128" pyo3-stub-gen = "0.17.1" opendal = { version = "0.55.0", features = ["services-http"] } -tokio = { version = "1.41.1", features = ["rt-multi-thread"] } +tokio = { version = "1.41.1", features = ["rt-multi-thread", "sync"] } +pyo3-async-runtimes = { version = "0.27.0", features = ["tokio-runtime"] } zarrs_opendal = "0.10.0" itertools = "0.14.0" +futures = "0.3.31" bytemuck = { version = "1.24.0", features = ["must_cast"] } pyo3-object_store = "0.7.0" # object_store 0.12 zarrs_object_store = "0.5.0" # object_store 0.12 diff --git a/python/zarrs/_internal.pyi b/python/zarrs/_internal.pyi index 2f676818..efce9e84 100644 --- a/python/zarrs/_internal.pyi +++ b/python/zarrs/_internal.pyi @@ -2,13 +2,53 @@ # ruff: noqa: E501, F401 import builtins +import collections.abc import typing +import numpy import numpy.typing import zarr.abc.store +@typing.final +class AsyncCodecPipelineImpl: + def __new__( + cls, + array_metadata: builtins.str, + store_config: zarr.abc.store.Store, + *, + validate_checksums: builtins.bool = False, + chunk_concurrent_minimum: builtins.int | None = None, + chunk_concurrent_maximum: builtins.int | None = None, + num_threads: builtins.int | None = None, + direct_io: builtins.bool = False, + file_handle_cache_size: builtins.int = 0, + ) -> AsyncCodecPipelineImpl: ... + def retrieve_chunks( + self, chunk_descriptions: typing.Sequence[ChunkItem] + ) -> builtins.list[collections.abc.Awaitable[numpy.typing.NDArray[numpy.uint8]]]: + r""" + One awaitable per chunk, each resolving to that chunk's decoded subset as a + flat `uint8` array; the caller copies it into the output. + """ + def store_chunks_with_indices( + self, + chunk_descriptions: typing.Sequence[ChunkItem], + value: numpy.typing.NDArray[typing.Any], + write_empty_chunks: builtins.bool, + ) -> collections.abc.Awaitable[None]: ... + @typing.final class ChunkItem: + @property + def io_shape(self) -> builtins.list[builtins.int]: + r""" + The shape of the input/output array, with dropped axes reinstated as 1. + """ + @property + def io_selection(self) -> builtins.list[slice]: + r""" + The selection of this chunk's subset within an array of [`Self::io_shape`]. + """ def __new__( cls, key: builtins.str, diff --git a/python/zarrs/pipeline.py b/python/zarrs/pipeline.py index b3368219..465dcd8a 100644 --- a/python/zarrs/pipeline.py +++ b/python/zarrs/pipeline.py @@ -12,9 +12,17 @@ from zarr.core import BatchedCodecPipeline from zarr.core.config import config from zarr.core.metadata import ArrayMetadata, ArrayV2Metadata, ArrayV3Metadata +from zarr.storage import FsspecStore + +try: + from zarr.storage import ObjectStore + + _ASYNC_STORES: tuple[type, ...] = (FsspecStore, ObjectStore) +except ImportError: + _ASYNC_STORES = (FsspecStore,) if TYPE_CHECKING: - from collections.abc import Iterable, Iterator + from collections.abc import Awaitable, Iterable, Iterator from typing import Self from zarr.abc.store import ByteGetter, ByteSetter, Store @@ -24,7 +32,7 @@ from zarr.core.indexing import SelectorTuple from zarr.dtype import ZDType -from ._internal import CodecPipelineImpl +from ._internal import AsyncCodecPipelineImpl, ChunkItem, CodecPipelineImpl from .utils import ( DiscontiguousArrayError, FillValueNoneError, @@ -41,16 +49,27 @@ class UnsupportedMetadataError(Exception): pass +def _supports_async_pipeline(store: Store) -> bool: + # Async-backed stores issue many requests concurrently instead of blocking a + # worker thread per request; everything else (e.g. ``LocalStore``) uses the + # synchronous pipeline. Mirrors the stores accepted by the Rust + # ``AsyncReadableWritableListableStorage`` conversion (see ``src/store.rs``). + return isinstance(store, _ASYNC_STORES) + + def get_codec_pipeline_impl( metadata: ArrayMetadata, store: Store, *, strict: bool -) -> CodecPipelineImpl | None: +) -> CodecPipelineImpl | AsyncCodecPipelineImpl | None: + pipeline_class = ( + AsyncCodecPipelineImpl if _supports_async_pipeline(store) else CodecPipelineImpl + ) try: array_metadata_json = json.dumps(metadata.to_dict()) # Maintain old behavior: https://github.com/zarrs/zarrs-python/tree/b36ba797cafec77f5f41a25316be02c718a2b4f8?tab=readme-ov-file#configuration validate_checksums = config.get("codec_pipeline.validate_checksums", True) if validate_checksums is None: validate_checksums = True - return CodecPipelineImpl( + return pipeline_class( array_metadata_json, store_config=store, validate_checksums=validate_checksums, @@ -104,7 +123,7 @@ def array_metadata_to_codecs(metadata: ArrayMetadata) -> list[Codec]: class ZarrsCodecPipeline(CodecPipeline): metadata: ArrayMetadata store: Store - impl: CodecPipelineImpl | None + impl: CodecPipelineImpl | AsyncCodecPipelineImpl | None python_impl: BatchedCodecPipeline | None def __getstate__(self) -> ZarrsCodecPipelineState: @@ -167,6 +186,22 @@ async def encode( ) -> Iterable[Buffer | None]: raise NotImplementedError("encode") + async def _retrieve_async( + self, impl: AsyncCodecPipelineImpl, chunks: list[ChunkItem], out: np.ndarray + ) -> None: + # Rust decodes each chunk on its tokio runtime; the copy into ``out`` + # happens here as each one lands, so decoded chunks are freed eagerly. + if not out.flags.c_contiguous: + raise ValueError("output array must be a C contiguous array") + + async def fill(chunk: ChunkItem, decoded: Awaitable[np.ndarray]) -> None: + dst = out.reshape(chunk.io_shape)[tuple(chunk.io_selection)] + dst[...] = (await decoded).view(out.dtype).reshape(dst.shape) + + await asyncio.gather( + *map(fill, chunks, impl.retrieve_chunks(chunks)), + ) + async def read( self, batch_info: Iterable[ @@ -198,11 +233,16 @@ async def read( return None else: out: NDArrayLike = out.as_ndarray_like() - await asyncio.to_thread( - self.impl.retrieve_chunks_and_apply_index, - chunks_desc.chunk_info_with_indices, - out, - ) + if isinstance(self.impl, AsyncCodecPipelineImpl): + await self._retrieve_async( + self.impl, chunks_desc.chunk_info_with_indices, out + ) + else: + await asyncio.to_thread( + self.impl.retrieve_chunks_and_apply_index, + chunks_desc.chunk_info_with_indices, + out, + ) return None async def write( @@ -240,12 +280,15 @@ async def write( ) elif not value_np.flags.c_contiguous: value_np = np.ascontiguousarray(value_np) - await asyncio.to_thread( - self.impl.store_chunks_with_indices, + args = ( chunks_desc.chunk_info_with_indices, value_np, chunks_desc.write_empty_chunks, ) + if isinstance(self.impl, AsyncCodecPipelineImpl): + await self.impl.store_chunks_with_indices(*args) + else: + await asyncio.to_thread(self.impl.store_chunks_with_indices, *args) return None def _raise_error_on_unsupported_batch_dtype( diff --git a/src/async_pipeline.rs b/src/async_pipeline.rs new file mode 100644 index 00000000..57acfce6 --- /dev/null +++ b/src/async_pipeline.rs @@ -0,0 +1,432 @@ +//! Asynchronous codec pipeline. +//! +//! [`AsyncCodecPipelineImpl`] mirrors [`crate::CodecPipelineImpl`], but holds an +//! [`AsyncReadableWritableListableStorage`] and talks to it with `async` +//! `get`/`set`/`erase` calls. +//! +//! Its methods return Python awaitables (via [`pyo3_async_runtimes`]) driven on +//! the shared Tokio runtime, so the asyncio event loop is never blocked. +//! +//! Those futures must be `Send + 'static`, so they cannot borrow the output +//! `numpy` buffer. Instead, [`AsyncCodecPipelineImpl::retrieve_chunks`] returns +//! one awaitable per chunk resolving to that chunk's decoded subset, and Python +//! copies each into the output as soon as it arrives. A decoded subset is +//! dropped right after its copy, so peak memory is bounded by the chunks in +//! flight rather than by the whole output. +//! +//! The write path only *reads* the input buffer; its future keeps the array +//! alive via a [`Py`] reference and reads it through a raw pointer. + +use std::borrow::Cow; +use std::collections::HashMap; +use std::sync::Arc; + +use futures::stream::{self, TryStreamExt}; +use itertools::Itertools; +use numpy::{IntoPyArray, PyArray1, PyUntypedArray, PyUntypedArrayMethods}; +use pyo3::exceptions::{PyRuntimeError, PyTypeError, PyValueError}; +use pyo3::prelude::*; +use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods}; +use tokio::sync::OnceCell; +use zarrs::array::{ + ArrayBytes, ArrayToBytesCodecTraits, AsyncArrayPartialDecoderTraits, CodecChain, CodecOptions, + DataType, FillValue, update_array_bytes, +}; +use zarrs::storage::{ + AsyncReadableStorage, AsyncReadableStorageTraits, AsyncReadableWritableListableStorage, + AsyncWritableStorageTraits, StorageHandle, +}; +// Not re-exported by `zarrs` (only the async traits are); see Cargo.toml. +use zarrs_codec::AsyncStoragePartialDecoder; + +use crate::chunk_item::ChunkItem; +use crate::concurrency::{ChunkConcurrentLimitAndCodecOptions, CodecPipelineConfig}; +use crate::store::StoreConfig; +use crate::utils::{PyCodecErrExt, PyErrExt as _, is_whole_chunk}; +use crate::{ArrayConfig, CodecPipelineImpl, parse_array_config}; + +type PartialDecoderCell = Arc>>; + +/// A decoded chunk subset, converted to a 1D `uint8` numpy array without copying. +struct DecodedBytes(Vec); + +impl<'py> IntoPyObject<'py> for DecodedBytes { + type Target = PyArray1; + type Output = Bound<'py, PyArray1>; + type Error = std::convert::Infallible; + + fn into_pyobject(self, py: Python<'py>) -> Result { + Ok(self.0.into_pyarray(py)) + } +} + +#[derive(Clone)] +#[gen_stub_pyclass] +#[pyclass] +pub(crate) struct AsyncCodecPipelineImpl { + /// Read handle, deliberately the READ-ONLY type; see [`CodecPipelineImpl`]. + pub(crate) readable_store: AsyncReadableStorage, + /// The writable handle -- `None` when zarr-python opened the store read-only. + pub(crate) writable_store: Option, + pub(crate) codec_chain: Arc, + pub(crate) codec_options: CodecOptions, + pub(crate) chunk_concurrent_minimum: usize, + pub(crate) chunk_concurrent_maximum: usize, + pub(crate) num_threads: usize, + pub(crate) fill_value: FillValue, + pub(crate) data_type: DataType, +} + +impl CodecPipelineConfig for AsyncCodecPipelineImpl { + fn codec_chain(&self) -> &CodecChain { + &self.codec_chain + } + fn data_type(&self) -> &DataType { + &self.data_type + } + fn codec_options(&self) -> &CodecOptions { + &self.codec_options + } + fn chunk_concurrent_minimum(&self) -> usize { + self.chunk_concurrent_minimum + } + fn chunk_concurrent_maximum(&self) -> usize { + self.chunk_concurrent_maximum + } + fn num_threads(&self) -> usize { + self.num_threads + } +} + +impl AsyncCodecPipelineImpl { + async fn retrieve_chunk_bytes<'a>( + &self, + item: &ChunkItem, + codec_options: &CodecOptions, + ) -> PyResult> { + let value_encoded = self + .readable_store + .get(&item.key) + .await + .map_py_err::()?; + let value_decoded = if let Some(value_encoded) = value_encoded { + let value_encoded: Vec = value_encoded.into(); // zero-copy in this case + self.codec_chain + .decode( + value_encoded.into(), + &item.shape, + &self.data_type, + &self.fill_value, + codec_options, + ) + .map_codec_err()? + } else { + ArrayBytes::new_fill_value(&self.data_type, item.num_elements, &self.fill_value) + .map_py_err::()? + }; + Ok(value_decoded) + } + + /// The writable store, or zarr-python's own refusal, verbatim. + fn writable(&self) -> PyResult<&AsyncReadableWritableListableStorage> { + self.writable_store.as_ref().ok_or_else(|| { + PyValueError::new_err("store was opened in read-only mode and does not support writing") + }) + } + + async fn store_chunk_bytes( + &self, + item: &ChunkItem, + value_decoded: ArrayBytes<'_>, + codec_options: &CodecOptions, + ) -> PyResult<()> { + value_decoded + .validate(item.num_elements, &self.data_type) + .map_codec_err()?; + + let store = self.writable()?; + + if value_decoded.is_fill_value(&self.fill_value) { + store.erase(&item.key).await.map_py_err::() + } else { + let value_encoded = self + .codec_chain + .encode( + value_decoded, + &item.shape, + &self.data_type, + &self.fill_value, + codec_options, + ) + .map(Cow::into_owned) + .map_codec_err()?; + + // Store the encoded chunk + store + .set(&item.key, value_encoded.into()) + .await + .map_py_err::() + } + } + + async fn store_chunk_subset_bytes( + &self, + item: &ChunkItem, + chunk_subset_bytes: ArrayBytes<'_>, + codec_options: &CodecOptions, + ) -> PyResult<()> { + let array_shape = &item.shape; + let chunk_subset = &item.chunk_subset; + if !chunk_subset.inbounds_shape(bytemuck::must_cast_slice(array_shape)) { + return Err(PyErr::new::(format!( + "chunk subset ({chunk_subset}) is out of bounds for array shape ({array_shape:?})" + ))); + } + let data_type_size = self.data_type.size(); + + if is_whole_chunk(item) { + // Fast path if the chunk subset spans the entire chunk, no read required + self.store_chunk_bytes(item, chunk_subset_bytes, codec_options) + .await + } else { + // Validate the chunk subset bytes + chunk_subset_bytes + .validate(chunk_subset.num_elements(), &self.data_type) + .map_codec_err()?; + + // Retrieve the chunk + let chunk_bytes_old = self.retrieve_chunk_bytes(item, codec_options).await?; + + // Update the chunk + let chunk_bytes_new = update_array_bytes( + chunk_bytes_old, + bytemuck::must_cast_slice(array_shape), + chunk_subset, + &chunk_subset_bytes, + data_type_size, + ) + .map_codec_err()?; + + // Store the updated chunk + self.store_chunk_bytes(item, chunk_bytes_new, codec_options) + .await + } + } +} + +#[gen_stub_pymethods] +#[pymethods] +impl AsyncCodecPipelineImpl { + #[pyo3(signature = ( + array_metadata, + store_config, + *, + validate_checksums=false, + chunk_concurrent_minimum=None, + chunk_concurrent_maximum=None, + num_threads=None, + direct_io=false, + file_handle_cache_size=0, + ))] + #[new] + fn new( + array_metadata: &str, + mut store_config: StoreConfig, + validate_checksums: bool, + chunk_concurrent_minimum: Option, + chunk_concurrent_maximum: Option, + num_threads: Option, + direct_io: bool, + file_handle_cache_size: usize, + ) -> PyResult { + // `direct_io` and `file_handle_cache_size` only affect the (synchronous) + // filesystem store, which the async pipeline does not support; kept for + // signature parity. + store_config.direct_io(direct_io); + store_config.file_handle_cache_size(file_handle_cache_size); + let ArrayConfig { + codec_chain, + codec_options, + chunk_concurrent_minimum, + chunk_concurrent_maximum, + num_threads, + fill_value, + data_type, + } = parse_array_config( + array_metadata, + validate_checksums, + chunk_concurrent_minimum, + chunk_concurrent_maximum, + num_threads, + )?; + + let store: AsyncReadableWritableListableStorage = + (&store_config).try_into().map_py_err::()?; + let writable_store = (!store_config.read_only).then(|| store.clone()); + let readable_store: AsyncReadableStorage = store.readable(); + + Ok(Self { + readable_store, + writable_store, + codec_chain, + codec_options, + chunk_concurrent_minimum, + chunk_concurrent_maximum, + num_threads, + fill_value, + data_type, + }) + } + + /// One awaitable per chunk, each resolving to that chunk's decoded subset as a + /// flat `uint8` array; the caller copies it into the output. + #[gen_stub(override_return_type( + type_repr = "builtins.list[collections.abc.Awaitable[numpy.typing.NDArray[numpy.uint8]]]", + imports = ("builtins", "collections.abc", "numpy", "numpy.typing") + ))] + fn retrieve_chunks<'py>( + &self, + py: Python<'py>, + chunk_descriptions: Vec, + ) -> PyResult>> { + // Only the codec options are needed here; chunk-level parallelism is + // delegated to the tokio runtime rather than a manual concurrency limit. + let Some((_chunk_concurrent_limit, codec_options)) = + chunk_descriptions.get_chunk_concurrent_limit_and_codec_options(self)? + else { + return Ok(vec![]); + }; + // Cheap (all `Arc`s / small values); the futures must own it. + let pipeline = Arc::new(self.clone()); + let codec_options = Arc::new(codec_options); + + // Subsets sharing a (partial) chunk share one lazily-built partial decoder. + let partial_decoders: HashMap<_, PartialDecoderCell> = chunk_descriptions + .iter() + .filter(|item| !is_whole_chunk(item)) + .map(|item| item.key.clone()) + .unique() + .map(|key| (key, Arc::default())) + .collect(); + + chunk_descriptions + .into_iter() + .map(|item| { + let p = pipeline.clone(); + let codec_options = codec_options.clone(); + let partial_decoder = partial_decoders.get(&item.key).cloned(); + pyo3_async_runtimes::tokio::future_into_py(py, async move { + let bytes = if let Some(cell) = partial_decoder { + let partial_decoder = cell + .get_or_try_init(|| async { + let input_handle = Arc::new(AsyncStoragePartialDecoder::new( + Arc::new(StorageHandle::new(p.readable_store.clone())), + item.key.clone(), + )); + p.codec_chain + .clone() + .async_partial_decoder( + input_handle, + &item.shape, + &p.data_type, + &p.fill_value, + &codec_options, + ) + .await + .map_codec_err() + }) + .await?; + partial_decoder + .partial_decode(&item.chunk_subset, &codec_options) + .await + .map_codec_err()? + .into_owned() + } else { + p.retrieve_chunk_bytes(&item, &codec_options).await? + }; + // FIXME: only supports fixed length data types. + let bytes = bytes.into_fixed().map_py_err::()?; + Ok(DecodedBytes(bytes.into_owned())) + }) + }) + .collect() + } + + #[gen_stub(override_return_type( + type_repr = "collections.abc.Awaitable[None]", + imports = ("collections.abc") + ))] + fn store_chunks_with_indices<'py>( + &self, + py: Python<'py>, + chunk_descriptions: Vec, + value: &Bound<'py, PyUntypedArray>, + write_empty_chunks: bool, + ) -> PyResult> { + // Fail before decoding anything; the write site checks again by construction. + self.writable()?; + + enum InputValue { + /// Address and length of `value`'s buffer, kept alive by `value_ref`. + Array(usize, usize), + Constant(FillValue), + } + + // Get input array + let input_slice = CodecPipelineImpl::nparray_to_slice(value)?; + let input = if value.ndim() > 0 { + InputValue::Array(input_slice.as_ptr() as usize, input_slice.len()) + } else { + InputValue::Constant(FillValue::new(input_slice.to_vec())) + }; + let value_ref = value.clone().unbind(); + + // Adjust the concurrency based on the codec chain and the first chunk description + let Some((chunk_concurrent_limit, mut codec_options)) = + chunk_descriptions.get_chunk_concurrent_limit_and_codec_options(self)? + else { + return pyo3_async_runtimes::tokio::future_into_py(py, async { Ok(()) }); + }; + codec_options.set_store_empty_chunks(write_empty_chunks); + + let p = self.clone(); + pyo3_async_runtimes::tokio::future_into_py(py, async move { + let _keep_alive = value_ref; + let p = &p; + let input = match input { + InputValue::Array(ptr, len) => { + // SAFETY: `value_ref` keeps the array alive until this future + // completes. As in the sync pipeline, the caller must not + // mutate it while the write is in flight. + // FIXME: Handle variable length data types, convert value to bytes and offsets + let input_slice = unsafe { std::slice::from_raw_parts(ptr as *const u8, len) }; + Ok(ArrayBytes::new_flen(Cow::Borrowed(input_slice))) + } + InputValue::Constant(constant_value) => Err(constant_value), + }; + let input = &input; + let codec_options = &codec_options; + stream::iter(chunk_descriptions.into_iter().map(Ok::)) + .try_for_each_concurrent(chunk_concurrent_limit, move |item| async move { + let chunk_subset_bytes = match input { + Ok(input) => input + .extract_array_subset( + &item.subset, + bytemuck::must_cast_slice(&item.array_shape), + &p.data_type, + ) + .map_codec_err()?, + Err(constant_value) => ArrayBytes::new_fill_value( + &p.data_type, + item.chunk_subset.num_elements(), + constant_value, + ) + .map_py_err::()?, + }; + p.store_chunk_subset_bytes(&item, chunk_subset_bytes, codec_options) + .await + }) + .await + }) + } +} diff --git a/src/chunk_item.rs b/src/chunk_item.rs index 8a9aa0ad..0cd774af 100644 --- a/src/chunk_item.rs +++ b/src/chunk_item.rs @@ -1,7 +1,7 @@ use std::num::NonZeroU64; use pyo3::{ - Bound, PyErr, PyResult, + Bound, PyErr, PyResult, Python, exceptions::{PyIndexError, PyValueError}, pyclass, pymethods, types::{PySlice, PySliceMethods as _}, @@ -69,6 +69,29 @@ impl ChunkItem { array_shape: shape_nonzero_u64, }) } + + /// The shape of the input/output array, with dropped axes reinstated as 1. + #[getter] + fn io_shape(&self) -> Vec { + self.array_shape.iter().map(|d| d.get()).collect() + } + + /// The selection of this chunk's subset within an array of [`Self::io_shape`]. + #[getter] + fn io_selection<'py>(&self, py: Python<'py>) -> PyResult>> { + self.subset + .to_ranges() + .into_iter() + .map(|r| { + Ok(PySlice::new( + py, + isize::try_from(r.start)?, + isize::try_from(r.end)?, + 1, + )) + }) + .collect() + } } fn slice_to_range(slice: &Bound<'_, PySlice>, length: isize) -> PyResult> { diff --git a/src/concurrency.rs b/src/concurrency.rs index 08b9dca1..ec4c29d7 100644 --- a/src/concurrency.rs +++ b/src/concurrency.rs @@ -1,44 +1,56 @@ use pyo3::PyResult; use zarrs::array::{ - ArrayCodecTraits, CodecOptions, RecommendedConcurrency, + ArrayCodecTraits, CodecChain, CodecOptions, DataType, RecommendedConcurrency, concurrency::calc_concurrency_outer_inner, }; -use crate::{CodecPipelineImpl, chunk_item::ChunkItem, utils::PyCodecErrExt as _}; +use crate::{chunk_item::ChunkItem, utils::PyCodecErrExt as _}; + +/// The pieces of pipeline configuration needed to compute concurrency limits. +/// +/// Implemented by both the synchronous [`crate::CodecPipelineImpl`] and the +/// asynchronous [`crate::async_pipeline::AsyncCodecPipelineImpl`], so the +/// concurrency calculation can be shared between them. +pub trait CodecPipelineConfig { + fn codec_chain(&self) -> &CodecChain; + fn data_type(&self) -> &DataType; + fn codec_options(&self) -> &CodecOptions; + fn chunk_concurrent_minimum(&self) -> usize; + fn chunk_concurrent_maximum(&self) -> usize; + fn num_threads(&self) -> usize; +} pub trait ChunkConcurrentLimitAndCodecOptions { - fn get_chunk_concurrent_limit_and_codec_options( + fn get_chunk_concurrent_limit_and_codec_options( &self, - codec_pipeline_impl: &CodecPipelineImpl, + config: &C, ) -> PyResult>; } impl ChunkConcurrentLimitAndCodecOptions for Vec { - fn get_chunk_concurrent_limit_and_codec_options( + fn get_chunk_concurrent_limit_and_codec_options( &self, - codec_pipeline_impl: &CodecPipelineImpl, + config: &C, ) -> PyResult> { let num_chunks = self.len(); let Some(item) = self.first() else { return Ok(None); }; - let codec_concurrency = codec_pipeline_impl - .codec_chain - .recommended_concurrency(&item.shape, &codec_pipeline_impl.data_type) + let codec_concurrency = config + .codec_chain() + .recommended_concurrency(&item.shape, config.data_type()) .map_codec_err()?; - let min_concurrent_chunks = - std::cmp::min(codec_pipeline_impl.chunk_concurrent_minimum, num_chunks); - let max_concurrent_chunks = - std::cmp::max(codec_pipeline_impl.chunk_concurrent_maximum, num_chunks); + let min_concurrent_chunks = std::cmp::min(config.chunk_concurrent_minimum(), num_chunks); + let max_concurrent_chunks = std::cmp::max(config.chunk_concurrent_maximum(), num_chunks); let (chunk_concurrent_limit, codec_concurrent_limit) = calc_concurrency_outer_inner( - codec_pipeline_impl.num_threads, + config.num_threads(), &RecommendedConcurrency::new(min_concurrent_chunks..max_concurrent_chunks), &codec_concurrency, ); - let codec_options = codec_pipeline_impl - .codec_options + let codec_options = config + .codec_options() .with_concurrent_target(codec_concurrent_limit); Ok(Some((chunk_concurrent_limit, codec_options))) } diff --git a/src/lib.rs b/src/lib.rs index c53e73de..a1cb932f 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -2,54 +2,32 @@ #![allow(clippy::module_name_repetitions)] use std::borrow::Cow; -use std::collections::HashMap; -use std::ptr::NonNull; use std::sync::Arc; -use chunk_item::ChunkItem; -use itertools::Itertools; -use numpy::npyffi::PyArrayObject; -use numpy::{PyArrayDescrMethods, PyUntypedArray, PyUntypedArrayMethods}; -use pyo3::exceptions::{PyRuntimeError, PyTypeError, PyValueError}; +use pyo3::exceptions::PyTypeError; use pyo3::prelude::*; use pyo3_stub_gen::define_stub_info_gatherer; -use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods}; -use rayon::iter::{IntoParallelIterator, ParallelIterator}; -use rayon_iter_concurrent_limit::iter_concurrent_limit; -use unsafe_cell_slice::UnsafeCellSlice; -use utils::is_whole_chunk; -use zarrs::array::{ - ArrayBytes, ArrayBytesDecodeIntoTarget, ArrayBytesFixedDisjointView, ArrayMetadata, - ArrayPartialDecoderTraits, ArrayToBytesCodecTraits, CodecChain, CodecOptions, DataType, - FillValue, StoragePartialDecoder, copy_fill_value_into, update_array_bytes, -}; +use zarrs::array::{ArrayMetadata, CodecChain, CodecOptions, DataType, FillValue}; use zarrs::config::global_config; use zarrs::convert::array_metadata_v2_to_v3; use zarrs::plugin::ZarrVersion; -use zarrs::storage::{ReadableStorage, ReadableWritableListableStorage, StorageHandle, StoreKey}; +mod async_pipeline; mod chunk_item; mod concurrency; -mod runtime; mod store; +mod sync_pipeline; #[cfg(test)] mod tests; mod utils; -use crate::concurrency::ChunkConcurrentLimitAndCodecOptions; -use crate::store::StoreConfig; -use crate::utils::{PyCodecErrExt, PyErrExt as _}; +use crate::async_pipeline::AsyncCodecPipelineImpl; +use crate::sync_pipeline::CodecPipelineImpl; +use crate::utils::PyErrExt as _; -// TODO: Use a OnceLock for store with get_or_try_init when stabilised? -#[gen_stub_pyclass] -#[pyclass] -pub(crate) struct CodecPipelineImpl { - /// Read handle, deliberately the READ-ONLY type: `set`/`erase` are not in its interface, - /// so the only way to reach a write is `writable()`. - pub(crate) readable_store: ReadableStorage, - /// The writable handle -- `None` when zarr-python opened the store read-only. Same object - /// as `readable_store`; the read-only case simply never keeps a writable view of it. - pub(crate) writable_store: Option, +/// Configuration parsed from the array metadata, shared by the synchronous and +/// asynchronous codec pipelines (everything except the store, which differs). +pub(crate) struct ArrayConfig { pub(crate) codec_chain: Arc, pub(crate) codec_options: CodecOptions, pub(crate) chunk_concurrent_minimum: usize, @@ -59,435 +37,53 @@ pub(crate) struct CodecPipelineImpl { pub(crate) data_type: DataType, } -impl CodecPipelineImpl { - fn retrieve_chunk_bytes<'a>( - &self, - item: &ChunkItem, - codec_chain: &CodecChain, - codec_options: &CodecOptions, - ) -> PyResult> { - let value_encoded = self - .readable_store - .get(&item.key) - .map_py_err::()?; - let value_decoded = if let Some(value_encoded) = value_encoded { - let value_encoded: Vec = value_encoded.into(); // zero-copy in this case - codec_chain - .decode( - value_encoded.into(), - &item.shape, - &self.data_type, - &self.fill_value, - codec_options, - ) - .map_codec_err()? - } else { - ArrayBytes::new_fill_value(&self.data_type, item.num_elements, &self.fill_value) - .map_py_err::()? - }; - Ok(value_decoded) - } - - /// The writable store, or zarr-python's own refusal, verbatim. - fn writable(&self) -> PyResult<&ReadableWritableListableStorage> { - self.writable_store.as_ref().ok_or_else(|| { - PyValueError::new_err("store was opened in read-only mode and does not support writing") - }) - } - - fn store_chunk_bytes( - &self, - item: &ChunkItem, - codec_chain: &CodecChain, - value_decoded: ArrayBytes, - codec_options: &CodecOptions, - ) -> PyResult<()> { - value_decoded - .validate(item.num_elements, &self.data_type) - .map_codec_err()?; - - let store = self.writable()?; - - if value_decoded.is_fill_value(&self.fill_value) { - store.erase(&item.key).map_py_err::() - } else { - let value_encoded = codec_chain - .encode( - value_decoded, - &item.shape, - &self.data_type, - &self.fill_value, - codec_options, - ) - .map(Cow::into_owned) - .map_codec_err()?; - - // Store the encoded chunk - store - .set(&item.key, value_encoded.into()) - .map_py_err::() - } - } - - fn store_chunk_subset_bytes( - &self, - item: &ChunkItem, - codec_chain: &CodecChain, - chunk_subset_bytes: ArrayBytes, - codec_options: &CodecOptions, - ) -> PyResult<()> { - let array_shape = &item.shape; - let chunk_subset = &item.chunk_subset; - if !chunk_subset.inbounds_shape(bytemuck::must_cast_slice(array_shape)) { - return Err(PyErr::new::(format!( - "chunk subset ({chunk_subset}) is out of bounds for array shape ({array_shape:?})" - ))); - } - let data_type_size = self.data_type.size(); - - if is_whole_chunk(item) { - // Fast path if the chunk subset spans the entire chunk, no read required - self.store_chunk_bytes(item, codec_chain, chunk_subset_bytes, codec_options) - } else { - // Validate the chunk subset bytes - chunk_subset_bytes - .validate(chunk_subset.num_elements(), &self.data_type) - .map_codec_err()?; - - // Retrieve the chunk - let chunk_bytes_old = self.retrieve_chunk_bytes(item, codec_chain, codec_options)?; - - // Update the chunk - let chunk_bytes_new = update_array_bytes( - chunk_bytes_old, - bytemuck::must_cast_slice(array_shape), - chunk_subset, - &chunk_subset_bytes, - data_type_size, - ) - .map_codec_err()?; - - // Store the updated chunk - self.store_chunk_bytes(item, codec_chain, chunk_bytes_new, codec_options) - } - } - - fn py_untyped_array_to_array_object<'a>( - value: &'a Bound<'_, PyUntypedArray>, - ) -> &'a PyArrayObject { - // TODO: Upstream a PyUntypedArray.as_array_ref()? - // https://github.com/zarrs/zarrs-python/pull/80/files/75be39184905d688ac04a5f8bca08c5241c458cd#r1918365296 - let array_object_ptr: NonNull = NonNull::new(value.as_array_ptr()) - .expect("bug in numpy crate: Bound<'_, PyUntypedArray>::as_array_ptr unexpectedly returned a null pointer"); - let array_object: &'a PyArrayObject = unsafe { - // SAFETY: the array object pointed to by array_object_ptr is valid for 'a - array_object_ptr.as_ref() - }; - array_object - } - - fn nparray_to_slice<'a>(value: &'a Bound<'_, PyUntypedArray>) -> Result<&'a [u8], PyErr> { - if !value.is_c_contiguous() { - return Err(PyErr::new::( - "input array must be a C contiguous array".to_string(), - )); +pub(crate) fn parse_array_config( + array_metadata: &str, + validate_checksums: bool, + chunk_concurrent_minimum: Option, + chunk_concurrent_maximum: Option, + num_threads: Option, +) -> PyResult { + let metadata = serde_json::from_str(array_metadata).map_py_err::()?; + let metadata_v3 = match &metadata { + ArrayMetadata::V2(v2) => { + Cow::Owned(array_metadata_v2_to_v3(v2).map_py_err::()?) } - let array_object: &PyArrayObject = Self::py_untyped_array_to_array_object(value); - let array_data = array_object.data.cast::(); - let array_len = value.len() * value.dtype().itemsize(); - let slice = unsafe { - // SAFETY: array_data is a valid pointer to a u8 array of length array_len - debug_assert!(!array_data.is_null()); - std::slice::from_raw_parts(array_data, array_len) - }; - Ok(slice) - } - - fn nparray_to_unsafe_cell_slice<'a>( - value: &'a Bound<'_, PyUntypedArray>, - ) -> Result, PyErr> { - if !value.is_c_contiguous() { - return Err(PyErr::new::( - "input array must be a C contiguous array".to_string(), - )); - } - let array_object: &PyArrayObject = Self::py_untyped_array_to_array_object(value); - let array_data = array_object.data.cast::(); - let array_len = value.len() * value.dtype().itemsize(); - let output = unsafe { - // SAFETY: array_data is a valid pointer to a u8 array of length array_len - debug_assert!(!array_data.is_null()); - std::slice::from_raw_parts_mut(array_data, array_len) - }; - Ok(UnsafeCellSlice::new(output)) - } -} - -#[gen_stub_pymethods] -#[pymethods] -impl CodecPipelineImpl { - #[pyo3(signature = ( - array_metadata, - store_config, - *, - validate_checksums=false, - chunk_concurrent_minimum=None, - chunk_concurrent_maximum=None, - num_threads=None, - direct_io=false, - file_handle_cache_size=0, - ))] - #[new] - fn new( - array_metadata: &str, - mut store_config: StoreConfig, - validate_checksums: bool, - chunk_concurrent_minimum: Option, - chunk_concurrent_maximum: Option, - num_threads: Option, - direct_io: bool, - file_handle_cache_size: usize, - ) -> PyResult { - store_config.direct_io(direct_io); - store_config.file_handle_cache_size(file_handle_cache_size); - let metadata = serde_json::from_str(array_metadata).map_py_err::()?; - let metadata_v3 = match &metadata { - ArrayMetadata::V2(v2) => { - Cow::Owned(array_metadata_v2_to_v3(v2).map_py_err::()?) - } - ArrayMetadata::V3(v3) => Cow::Borrowed(v3), - }; - let codec_chain = - Arc::new(CodecChain::from_metadata(&metadata_v3.codecs).map_py_err::()?); - let codec_options = CodecOptions::default().with_validate_checksums(validate_checksums); - - let chunk_concurrent_minimum = - chunk_concurrent_minimum.unwrap_or(global_config().chunk_concurrent_minimum()); - let chunk_concurrent_maximum = - chunk_concurrent_maximum.unwrap_or(rayon::current_num_threads()); - let num_threads = num_threads.unwrap_or(rayon::current_num_threads()); - - let store: ReadableWritableListableStorage = - (&store_config).try_into().map_py_err::()?; - let writable_store = (!store_config.read_only).then(|| store.clone()); - let readable_store: ReadableStorage = store.readable(); - - let data_type = - DataType::from_metadata(&metadata_v3.data_type).map_py_err::()?; - let fill_value = data_type - .fill_value(&metadata_v3.fill_value, ZarrVersion::V3) - .map_err(|_| match &metadata { - ArrayMetadata::V2(metadata) => format!( - "incompatible fill value metadata: dtype={}, fill_value={}", - metadata.dtype, metadata.fill_value - ), - ArrayMetadata::V3(metadata) => format!( - "incompatible fill value metadata: data_type={}, fill_value={}", - metadata.data_type, metadata.fill_value - ), - }) - .map_py_err::()?; - - Ok(Self { - readable_store, - codec_chain, - codec_options, - chunk_concurrent_minimum, - chunk_concurrent_maximum, - num_threads, - fill_value, - data_type, - writable_store, + ArrayMetadata::V3(v3) => Cow::Borrowed(v3), + }; + let codec_chain = + Arc::new(CodecChain::from_metadata(&metadata_v3.codecs).map_py_err::()?); + let codec_options = CodecOptions::default().with_validate_checksums(validate_checksums); + + let chunk_concurrent_minimum = + chunk_concurrent_minimum.unwrap_or(global_config().chunk_concurrent_minimum()); + let chunk_concurrent_maximum = chunk_concurrent_maximum.unwrap_or(rayon::current_num_threads()); + let num_threads = num_threads.unwrap_or(rayon::current_num_threads()); + + let data_type = DataType::from_metadata(&metadata_v3.data_type).map_py_err::()?; + let fill_value = data_type + .fill_value(&metadata_v3.fill_value, ZarrVersion::V3) + .map_err(|_| match &metadata { + ArrayMetadata::V2(metadata) => format!( + "incompatible fill value metadata: dtype={}, fill_value={}", + metadata.dtype, metadata.fill_value + ), + ArrayMetadata::V3(metadata) => format!( + "incompatible fill value metadata: data_type={}, fill_value={}", + metadata.data_type, metadata.fill_value + ), }) - } - - fn retrieve_chunks_and_apply_index( - &self, - py: Python, - chunk_descriptions: Vec, // FIXME: Ref / iterable? - value: &Bound<'_, PyUntypedArray>, - ) -> PyResult<()> { - // Get input array - let output = Self::nparray_to_unsafe_cell_slice(value)?; - - // Adjust the concurrency based on the codec chain and the first chunk description - let Some((chunk_concurrent_limit, codec_options)) = - chunk_descriptions.get_chunk_concurrent_limit_and_codec_options(self)? - else { - return Ok(()); - }; - - // Assemble partial decoders ahead of time and in parallel - let partial_chunk_items = chunk_descriptions - .iter() - .filter(|item| !(is_whole_chunk(item))) - .unique_by(|item| item.key.clone()) - .collect::>(); - let mut partial_decoder_cache: HashMap> = - HashMap::new(); - if !partial_chunk_items.is_empty() { - let key_decoder_pairs = - iter_concurrent_limit!(chunk_concurrent_limit, partial_chunk_items, map, |item| { - let storage_handle = Arc::new(StorageHandle::new(self.readable_store.clone())); - let input_handle = StoragePartialDecoder::new(storage_handle, item.key.clone()); - let partial_decoder = self - .codec_chain - .clone() - .partial_decoder( - Arc::new(input_handle), - &item.shape, - &self.data_type, - &self.fill_value, - &codec_options, - ) - .map_codec_err()?; - Ok((item.key.clone(), partial_decoder)) - }) - .collect::>>()?; - partial_decoder_cache.extend(key_decoder_pairs); - } - - py.detach(move || { - // FIXME: the `decode_into` methods only support fixed length data types. - // For variable length data types, need a codepath with non `_into` methods. - // Collect all the subsets and copy into value on the Python side? - let update_chunk_subset = |item: ChunkItem| { - let mut output_view = unsafe { - // TODO: Is the following correct? - // can we guarantee that when this function is called from Python with arbitrary arguments? - // SAFETY: chunks represent disjoint array subsets - ArrayBytesFixedDisjointView::new( - output, - // TODO: why is data_type in `item`, it should be derived from `output`, no? - self.data_type - .fixed_size() - .ok_or("variable length data type not supported") - .map_py_err::()?, - bytemuck::must_cast_slice(&item.array_shape), - item.subset.clone(), - ) - .map_py_err::()? - }; - let target = ArrayBytesDecodeIntoTarget::Fixed(&mut output_view); - // See zarrs::array::Array::retrieve_chunk_subset_into - if is_whole_chunk(&item) { - // See zarrs::array::Array::retrieve_chunk_into - if let Some(chunk_encoded) = self - .readable_store - .get(&item.key) - .map_py_err::()? - { - // Decode the encoded data into the output buffer - let chunk_encoded: Vec = chunk_encoded.into(); - self.codec_chain.decode_into( - Cow::Owned(chunk_encoded), - &item.shape, - &self.data_type, - &self.fill_value, - target, - &codec_options, - ) - } else { - // The chunk is missing, write the fill value - copy_fill_value_into(&self.data_type, &self.fill_value, target) - } - } else { - let key = &item.key; - let partial_decoder = partial_decoder_cache.get(key).ok_or_else(|| { - PyRuntimeError::new_err(format!("Partial decoder not found for key: {key}")) - })?; - partial_decoder.partial_decode_into(&item.chunk_subset, target, &codec_options) - } - .map_codec_err() - }; - - iter_concurrent_limit!( - chunk_concurrent_limit, - chunk_descriptions, - try_for_each, - update_chunk_subset - )?; - - Ok(()) - }) - } - - fn store_chunks_with_indices( - &self, - py: Python, - chunk_descriptions: Vec, - value: &Bound<'_, PyUntypedArray>, - write_empty_chunks: bool, - ) -> PyResult<()> { - // Fail before decoding anything; the write site checks again by construction. - self.writable()?; - - enum InputValue<'a> { - Array(ArrayBytes<'a>), - Constant(FillValue), - } - - // Get input array - let input_slice = Self::nparray_to_slice(value)?; - let input = if value.ndim() > 0 { - // FIXME: Handle variable length data types, convert value to bytes and offsets - InputValue::Array(ArrayBytes::new_flen(Cow::Borrowed(input_slice))) - } else { - InputValue::Constant(FillValue::new(input_slice.to_vec())) - }; - - // Adjust the concurrency based on the codec chain and the first chunk description - let Some((chunk_concurrent_limit, mut codec_options)) = - chunk_descriptions.get_chunk_concurrent_limit_and_codec_options(self)? - else { - return Ok(()); - }; - codec_options.set_store_empty_chunks(write_empty_chunks); - - py.detach(move || { - let store_chunk = |item: ChunkItem| match &input { - InputValue::Array(input) => { - let chunk_subset_bytes = input - .extract_array_subset( - &item.subset, - bytemuck::must_cast_slice(&item.array_shape), - &self.data_type, - ) - .map_codec_err()?; - self.store_chunk_subset_bytes( - &item, - &self.codec_chain, - chunk_subset_bytes, - &codec_options, - ) - } - InputValue::Constant(constant_value) => { - let chunk_subset_bytes = ArrayBytes::new_fill_value( - &self.data_type, - item.chunk_subset.num_elements(), - constant_value, - ) - .map_py_err::()?; - - self.store_chunk_subset_bytes( - &item, - &self.codec_chain, - chunk_subset_bytes, - &codec_options, - ) - } - }; - - iter_concurrent_limit!( - chunk_concurrent_limit, - chunk_descriptions, - try_for_each, - store_chunk - )?; - - Ok(()) - }) - } + .map_py_err::()?; + + Ok(ArrayConfig { + codec_chain, + codec_options, + chunk_concurrent_minimum, + chunk_concurrent_maximum, + num_threads, + fill_value, + data_type, + }) } /// A Python module implemented in Rust. @@ -497,6 +93,8 @@ pub mod _internal { #[allow(non_upper_case_globals)] const __version__: &str = env!("CARGO_PKG_VERSION"); #[pymodule_export] + use super::AsyncCodecPipelineImpl; + #[pymodule_export] use super::CodecPipelineImpl; #[pymodule_export] use super::chunk_item::ChunkItem; diff --git a/src/runtime.rs b/src/runtime.rs deleted file mode 100644 index 161db085..00000000 --- a/src/runtime.rs +++ /dev/null @@ -1,18 +0,0 @@ -use std::sync::OnceLock; -use tokio::runtime::Runtime; -use zarrs::storage::storage_adapter::async_to_sync::AsyncToSyncBlockOn; - -static RUNTIME: OnceLock = OnceLock::new(); - -pub struct TokioBlockOn(tokio::runtime::Handle); - -impl AsyncToSyncBlockOn for TokioBlockOn { - fn block_on(&self, future: F) -> F::Output { - self.0.block_on(future) - } -} - -pub fn tokio_block_on() -> TokioBlockOn { - let runtime = RUNTIME.get_or_init(|| Runtime::new().expect("Failed to create Tokio runtime")); - TokioBlockOn(runtime.handle().clone()) -} diff --git a/src/store.rs b/src/store.rs index b0ad2ff5..7ac0d48d 100644 --- a/src/store.rs +++ b/src/store.rs @@ -7,11 +7,9 @@ use pyo3::{ types::{PyAnyMethods, PyStringMethods, PyTypeMethods}, }; use pyo3_object_store::PyExternalObjectStore; -use zarrs::storage::{ - ReadableWritableListableStorage, storage_adapter::async_to_sync::AsyncToSyncStorageAdapter, -}; +use zarrs::storage::{AsyncReadableWritableListableStorage, ReadableWritableListableStorage}; -use crate::{runtime::tokio_block_on, utils::PyErrExt}; +use crate::utils::PyErrExt; mod filesystem; mod http; @@ -118,19 +116,44 @@ impl TryFrom<&StoreConfig> for ReadableWritableListableStorage { fn try_from(value: &StoreConfig) -> Result { match &value.kind { StoreKind::Filesystem(config) => config.try_into(), + // HTTP and object-store backends are asynchronous and are routed to the + // asynchronous pipeline instead (see `python/zarrs/pipeline.py`). + StoreKind::Http(_) | StoreKind::ObStore(_) => { + Err(PyErr::new::( + "the synchronous codec pipeline only supports the filesystem store; \ + use the asynchronous pipeline instead" + .to_string(), + )) + } + } + } +} + +impl TryFrom<&StoreConfig> for AsyncReadableWritableListableStorage { + type Error = PyErr; + + fn try_from(value: &StoreConfig) -> Result { + match &value.kind { + // The filesystem store (`zarrs::filesystem::FilesystemStore`) only + // implements the synchronous storage traits, so it cannot back the + // asynchronous pipeline. + StoreKind::Filesystem(_) => Err(PyErr::new::( + "the asynchronous codec pipeline does not support the filesystem store \ + (zarrs's FilesystemStore is synchronous); use the synchronous pipeline instead" + .to_string(), + )), StoreKind::Http(config) => config.try_into(), StoreKind::ObStore(config) => config.try_into(), } } } -fn opendal_builder_to_sync_store( +fn opendal_builder_to_async_store( builder: B, -) -> PyResult { +) -> PyResult { let operator = opendal::Operator::new(builder) .map_py_err::()? .finish(); let store = Arc::new(zarrs_opendal::AsyncOpendalStore::new(operator)); - let store = Arc::new(AsyncToSyncStorageAdapter::new(store, tokio_block_on())); Ok(store) } diff --git a/src/store/http.rs b/src/store/http.rs index 04e75f54..65d46e32 100644 --- a/src/store/http.rs +++ b/src/store/http.rs @@ -1,9 +1,9 @@ use std::collections::HashMap; use pyo3::{Bound, PyAny, PyErr, PyResult, exceptions::PyValueError}; -use zarrs::storage::ReadableWritableListableStorage; +use zarrs::storage::AsyncReadableWritableListableStorage; -use super::opendal_builder_to_sync_store; +use super::opendal_builder_to_async_store; #[derive(Debug, Clone)] pub struct HttpStoreConfig { @@ -32,11 +32,11 @@ impl HttpStoreConfig { } } -impl TryInto for &HttpStoreConfig { +impl TryInto for &HttpStoreConfig { type Error = PyErr; - fn try_into(self) -> Result { + fn try_into(self) -> Result { let builder = opendal::services::Http::default().endpoint(&self.endpoint); - opendal_builder_to_sync_store(builder) + opendal_builder_to_async_store(builder) } } diff --git a/src/store/obstore.rs b/src/store/obstore.rs index 029b1878..ecb5eeb4 100644 --- a/src/store/obstore.rs +++ b/src/store/obstore.rs @@ -1,13 +1,9 @@ use std::sync::Arc; use pyo3::PyErr; -use zarrs::storage::{ - ReadableWritableListableStorage, storage_adapter::async_to_sync::AsyncToSyncStorageAdapter, -}; +use zarrs::storage::AsyncReadableWritableListableStorage; use zarrs_object_store::{AsyncObjectStore, object_store::ObjectStore}; -use crate::runtime::tokio_block_on; - #[derive(Debug, Clone)] pub struct ObStoreConfig { store: Arc, @@ -19,15 +15,11 @@ impl ObStoreConfig { } } -impl TryInto for &ObStoreConfig { +impl TryInto for &ObStoreConfig { type Error = PyErr; - fn try_into(self) -> Result { + fn try_into(self) -> Result { let async_store = Arc::new(AsyncObjectStore::new(self.store.clone())); - let sync_store = Arc::new(AsyncToSyncStorageAdapter::new( - async_store, - tokio_block_on(), - )); - Ok(sync_store) + Ok(async_store) } } diff --git a/src/sync_pipeline.rs b/src/sync_pipeline.rs new file mode 100644 index 00000000..4618d6f8 --- /dev/null +++ b/src/sync_pipeline.rs @@ -0,0 +1,490 @@ +//! Synchronous codec pipeline. +//! +//! [`CodecPipelineImpl`] holds a [`ReadableStorage`] (and, unless opened +//! read-only, a [`ReadableWritableListableStorage`]) and drives +//! it with synchronous `get`/`set`/`erase` calls, decoding and encoding chunks +//! in parallel with `rayon`. Its asynchronous counterpart is +//! [`crate::async_pipeline::AsyncCodecPipelineImpl`]. + +use std::borrow::Cow; +use std::collections::HashMap; +use std::ptr::NonNull; +use std::sync::Arc; + +use itertools::Itertools; +use numpy::npyffi::PyArrayObject; +use numpy::{PyArrayDescrMethods, PyUntypedArray, PyUntypedArrayMethods}; +use pyo3::exceptions::{PyRuntimeError, PyTypeError, PyValueError}; +use pyo3::prelude::*; +use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods}; +use rayon::iter::{IntoParallelIterator, ParallelIterator}; +use rayon_iter_concurrent_limit::iter_concurrent_limit; +use unsafe_cell_slice::UnsafeCellSlice; +use zarrs::array::{ + ArrayBytes, ArrayBytesDecodeIntoTarget, ArrayBytesFixedDisjointView, ArrayPartialDecoderTraits, + ArrayToBytesCodecTraits, CodecChain, CodecOptions, DataType, FillValue, StoragePartialDecoder, + copy_fill_value_into, update_array_bytes, +}; +use zarrs::storage::{ReadableStorage, ReadableWritableListableStorage, StorageHandle, StoreKey}; + +use crate::chunk_item::{self, ChunkItem}; +use crate::concurrency::{ChunkConcurrentLimitAndCodecOptions, CodecPipelineConfig}; +use crate::store::StoreConfig; +use crate::utils::{PyCodecErrExt, PyErrExt as _, is_whole_chunk}; +use crate::{ArrayConfig, parse_array_config}; + +// TODO: Use a OnceLock for store with get_or_try_init when stabilised? +#[gen_stub_pyclass] +#[pyclass] +pub(crate) struct CodecPipelineImpl { + /// Read handle, deliberately the READ-ONLY type: `set`/`erase` are not in its interface, + /// so the only way to reach a write is `writable()`. + pub(crate) readable_store: ReadableStorage, + /// The writable handle -- `None` when zarr-python opened the store read-only. Same object + /// as `readable_store`; the read-only case simply never keeps a writable view of it. + pub(crate) writable_store: Option, + pub(crate) codec_chain: Arc, + pub(crate) codec_options: CodecOptions, + pub(crate) chunk_concurrent_minimum: usize, + pub(crate) chunk_concurrent_maximum: usize, + pub(crate) num_threads: usize, + pub(crate) fill_value: FillValue, + pub(crate) data_type: DataType, +} + +impl CodecPipelineConfig for CodecPipelineImpl { + fn codec_chain(&self) -> &CodecChain { + &self.codec_chain + } + fn data_type(&self) -> &DataType { + &self.data_type + } + fn codec_options(&self) -> &CodecOptions { + &self.codec_options + } + fn chunk_concurrent_minimum(&self) -> usize { + self.chunk_concurrent_minimum + } + fn chunk_concurrent_maximum(&self) -> usize { + self.chunk_concurrent_maximum + } + fn num_threads(&self) -> usize { + self.num_threads + } +} + +impl CodecPipelineImpl { + fn retrieve_chunk_bytes<'a>( + &self, + item: &ChunkItem, + codec_chain: &CodecChain, + codec_options: &CodecOptions, + ) -> PyResult> { + let value_encoded = self + .readable_store + .get(&item.key) + .map_py_err::()?; + let value_decoded = if let Some(value_encoded) = value_encoded { + let value_encoded: Vec = value_encoded.into(); // zero-copy in this case + codec_chain + .decode( + value_encoded.into(), + &item.shape, + &self.data_type, + &self.fill_value, + codec_options, + ) + .map_codec_err()? + } else { + ArrayBytes::new_fill_value(&self.data_type, item.num_elements, &self.fill_value) + .map_py_err::()? + }; + Ok(value_decoded) + } + + /// The writable store, or zarr-python's own refusal, verbatim. + fn writable(&self) -> PyResult<&ReadableWritableListableStorage> { + self.writable_store.as_ref().ok_or_else(|| { + PyValueError::new_err("store was opened in read-only mode and does not support writing") + }) + } + + fn store_chunk_bytes( + &self, + item: &ChunkItem, + codec_chain: &CodecChain, + value_decoded: ArrayBytes, + codec_options: &CodecOptions, + ) -> PyResult<()> { + value_decoded + .validate(item.num_elements, &self.data_type) + .map_codec_err()?; + + let store = self.writable()?; + + if value_decoded.is_fill_value(&self.fill_value) { + store.erase(&item.key).map_py_err::() + } else { + let value_encoded = codec_chain + .encode( + value_decoded, + &item.shape, + &self.data_type, + &self.fill_value, + codec_options, + ) + .map(Cow::into_owned) + .map_codec_err()?; + + // Store the encoded chunk + store + .set(&item.key, value_encoded.into()) + .map_py_err::() + } + } + + fn store_chunk_subset_bytes( + &self, + item: &ChunkItem, + codec_chain: &CodecChain, + chunk_subset_bytes: ArrayBytes, + codec_options: &CodecOptions, + ) -> PyResult<()> { + let array_shape = &item.shape; + let chunk_subset = &item.chunk_subset; + if !chunk_subset.inbounds_shape(bytemuck::must_cast_slice(array_shape)) { + return Err(PyErr::new::(format!( + "chunk subset ({chunk_subset}) is out of bounds for array shape ({array_shape:?})" + ))); + } + let data_type_size = self.data_type.size(); + + if is_whole_chunk(item) { + // Fast path if the chunk subset spans the entire chunk, no read required + self.store_chunk_bytes(item, codec_chain, chunk_subset_bytes, codec_options) + } else { + // Validate the chunk subset bytes + chunk_subset_bytes + .validate(chunk_subset.num_elements(), &self.data_type) + .map_codec_err()?; + + // Retrieve the chunk + let chunk_bytes_old = self.retrieve_chunk_bytes(item, codec_chain, codec_options)?; + + // Update the chunk + let chunk_bytes_new = update_array_bytes( + chunk_bytes_old, + bytemuck::must_cast_slice(array_shape), + chunk_subset, + &chunk_subset_bytes, + data_type_size, + ) + .map_codec_err()?; + + // Store the updated chunk + self.store_chunk_bytes(item, codec_chain, chunk_bytes_new, codec_options) + } + } + + fn py_untyped_array_to_array_object<'a>( + value: &'a Bound<'_, PyUntypedArray>, + ) -> &'a PyArrayObject { + // TODO: Upstream a PyUntypedArray.as_array_ref()? + // https://github.com/zarrs/zarrs-python/pull/80/files/75be39184905d688ac04a5f8bca08c5241c458cd#r1918365296 + let array_object_ptr: NonNull = NonNull::new(value.as_array_ptr()) + .expect("bug in numpy crate: Bound<'_, PyUntypedArray>::as_array_ptr unexpectedly returned a null pointer"); + let array_object: &'a PyArrayObject = unsafe { + // SAFETY: the array object pointed to by array_object_ptr is valid for 'a + array_object_ptr.as_ref() + }; + array_object + } + + pub(crate) fn nparray_to_slice<'a>( + value: &'a Bound<'_, PyUntypedArray>, + ) -> Result<&'a [u8], PyErr> { + if !value.is_c_contiguous() { + return Err(PyErr::new::( + "input array must be a C contiguous array".to_string(), + )); + } + let array_object: &PyArrayObject = Self::py_untyped_array_to_array_object(value); + let array_data = array_object.data.cast::(); + let array_len = value.len() * value.dtype().itemsize(); + let slice = unsafe { + // SAFETY: array_data is a valid pointer to a u8 array of length array_len + debug_assert!(!array_data.is_null()); + std::slice::from_raw_parts(array_data, array_len) + }; + Ok(slice) + } + + pub(crate) fn nparray_to_unsafe_cell_slice<'a>( + value: &'a Bound<'_, PyUntypedArray>, + ) -> Result, PyErr> { + if !value.is_c_contiguous() { + return Err(PyErr::new::( + "input array must be a C contiguous array".to_string(), + )); + } + let array_object: &PyArrayObject = Self::py_untyped_array_to_array_object(value); + let array_data = array_object.data.cast::(); + let array_len = value.len() * value.dtype().itemsize(); + let output = unsafe { + // SAFETY: array_data is a valid pointer to a u8 array of length array_len + debug_assert!(!array_data.is_null()); + std::slice::from_raw_parts_mut(array_data, array_len) + }; + Ok(UnsafeCellSlice::new(output)) + } +} + +#[gen_stub_pymethods] +#[pymethods] +impl CodecPipelineImpl { + #[pyo3(signature = ( + array_metadata, + store_config, + *, + validate_checksums=false, + chunk_concurrent_minimum=None, + chunk_concurrent_maximum=None, + num_threads=None, + direct_io=false, + file_handle_cache_size=0, + ))] + #[new] + fn new( + array_metadata: &str, + mut store_config: StoreConfig, + validate_checksums: bool, + chunk_concurrent_minimum: Option, + chunk_concurrent_maximum: Option, + num_threads: Option, + direct_io: bool, + file_handle_cache_size: usize, + ) -> PyResult { + store_config.direct_io(direct_io); + store_config.file_handle_cache_size(file_handle_cache_size); + let ArrayConfig { + codec_chain, + codec_options, + chunk_concurrent_minimum, + chunk_concurrent_maximum, + num_threads, + fill_value, + data_type, + } = parse_array_config( + array_metadata, + validate_checksums, + chunk_concurrent_minimum, + chunk_concurrent_maximum, + num_threads, + )?; + + let store: ReadableWritableListableStorage = + (&store_config).try_into().map_py_err::()?; + let writable_store = (!store_config.read_only).then(|| store.clone()); + let readable_store: ReadableStorage = store.readable(); + + Ok(Self { + readable_store, + codec_chain, + codec_options, + chunk_concurrent_minimum, + chunk_concurrent_maximum, + num_threads, + fill_value, + data_type, + writable_store, + }) + } + + fn retrieve_chunks_and_apply_index( + &self, + py: Python, + chunk_descriptions: Vec, // FIXME: Ref / iterable? + value: &Bound<'_, PyUntypedArray>, + ) -> PyResult<()> { + // Get input array + let output = Self::nparray_to_unsafe_cell_slice(value)?; + + // Adjust the concurrency based on the codec chain and the first chunk description + let Some((chunk_concurrent_limit, codec_options)) = + chunk_descriptions.get_chunk_concurrent_limit_and_codec_options(self)? + else { + return Ok(()); + }; + + // Assemble partial decoders ahead of time and in parallel + let partial_chunk_items = chunk_descriptions + .iter() + .filter(|item| !(is_whole_chunk(item))) + .unique_by(|item| item.key.clone()) + .collect::>(); + let mut partial_decoder_cache: HashMap> = + HashMap::new(); + if !partial_chunk_items.is_empty() { + let key_decoder_pairs = + iter_concurrent_limit!(chunk_concurrent_limit, partial_chunk_items, map, |item| { + let storage_handle = Arc::new(StorageHandle::new(self.readable_store.clone())); + let input_handle = StoragePartialDecoder::new(storage_handle, item.key.clone()); + let partial_decoder = self + .codec_chain + .clone() + .partial_decoder( + Arc::new(input_handle), + &item.shape, + &self.data_type, + &self.fill_value, + &codec_options, + ) + .map_codec_err()?; + Ok((item.key.clone(), partial_decoder)) + }) + .collect::>>()?; + partial_decoder_cache.extend(key_decoder_pairs); + } + + py.detach(move || { + // FIXME: the `decode_into` methods only support fixed length data types. + // For variable length data types, need a codepath with non `_into` methods. + // Collect all the subsets and copy into value on the Python side? + let update_chunk_subset = |item: ChunkItem| { + let mut output_view = unsafe { + // TODO: Is the following correct? + // can we guarantee that when this function is called from Python with arbitrary arguments? + // SAFETY: chunks represent disjoint array subsets + ArrayBytesFixedDisjointView::new( + output, + // TODO: why is data_type in `item`, it should be derived from `output`, no? + self.data_type + .fixed_size() + .ok_or("variable length data type not supported") + .map_py_err::()?, + bytemuck::must_cast_slice(&item.array_shape), + item.subset.clone(), + ) + .map_py_err::()? + }; + let target = ArrayBytesDecodeIntoTarget::Fixed(&mut output_view); + // See zarrs::array::Array::retrieve_chunk_subset_into + if is_whole_chunk(&item) { + // See zarrs::array::Array::retrieve_chunk_into + if let Some(chunk_encoded) = self + .readable_store + .get(&item.key) + .map_py_err::()? + { + // Decode the encoded data into the output buffer + let chunk_encoded: Vec = chunk_encoded.into(); + self.codec_chain.decode_into( + Cow::Owned(chunk_encoded), + &item.shape, + &self.data_type, + &self.fill_value, + target, + &codec_options, + ) + } else { + // The chunk is missing, write the fill value + copy_fill_value_into(&self.data_type, &self.fill_value, target) + } + } else { + let key = &item.key; + let partial_decoder = partial_decoder_cache.get(key).ok_or_else(|| { + PyRuntimeError::new_err(format!("Partial decoder not found for key: {key}")) + })?; + partial_decoder.partial_decode_into(&item.chunk_subset, target, &codec_options) + } + .map_codec_err() + }; + + iter_concurrent_limit!( + chunk_concurrent_limit, + chunk_descriptions, + try_for_each, + update_chunk_subset + )?; + + Ok(()) + }) + } + + fn store_chunks_with_indices( + &self, + py: Python, + chunk_descriptions: Vec, + value: &Bound<'_, PyUntypedArray>, + write_empty_chunks: bool, + ) -> PyResult<()> { + // Fail before decoding anything; the write site checks again by construction. + self.writable()?; + + enum InputValue<'a> { + Array(ArrayBytes<'a>), + Constant(FillValue), + } + + // Get input array + let input_slice = Self::nparray_to_slice(value)?; + let input = if value.ndim() > 0 { + // FIXME: Handle variable length data types, convert value to bytes and offsets + InputValue::Array(ArrayBytes::new_flen(Cow::Borrowed(input_slice))) + } else { + InputValue::Constant(FillValue::new(input_slice.to_vec())) + }; + + // Adjust the concurrency based on the codec chain and the first chunk description + let Some((chunk_concurrent_limit, mut codec_options)) = + chunk_descriptions.get_chunk_concurrent_limit_and_codec_options(self)? + else { + return Ok(()); + }; + codec_options.set_store_empty_chunks(write_empty_chunks); + + py.detach(move || { + let store_chunk = |item: ChunkItem| match &input { + InputValue::Array(input) => { + let chunk_subset_bytes = input + .extract_array_subset( + &item.subset, + bytemuck::must_cast_slice(&item.array_shape), + &self.data_type, + ) + .map_codec_err()?; + self.store_chunk_subset_bytes( + &item, + &self.codec_chain, + chunk_subset_bytes, + &codec_options, + ) + } + InputValue::Constant(constant_value) => { + let chunk_subset_bytes = ArrayBytes::new_fill_value( + &self.data_type, + item.chunk_subset.num_elements(), + constant_value, + ) + .map_py_err::()?; + + self.store_chunk_subset_bytes( + &item, + &self.codec_chain, + chunk_subset_bytes, + &codec_options, + ) + } + }; + + iter_concurrent_limit!( + chunk_concurrent_limit, + chunk_descriptions, + try_for_each, + store_chunk + )?; + + Ok(()) + }) + } +} diff --git a/src/utils.rs b/src/utils.rs index 3ededf06..97d5d0fb 100644 --- a/src/utils.rs +++ b/src/utils.rs @@ -3,7 +3,7 @@ use std::fmt::Display; use pyo3::{PyErr, PyResult, PyTypeInfo}; use zarrs::array::CodecError; -use crate::ChunkItem; +use crate::chunk_item::ChunkItem; pub(crate) trait PyErrExt { fn map_py_err(self) -> PyResult;