From 293e5fc5e36a280bdec210dc94ba68538db5f822 Mon Sep 17 00:00:00 2001 From: Lachlan Deakin Date: Sat, 7 Feb 2026 22:44:21 +1100 Subject: [PATCH 1/4] feat: add `ZarrsArray` --- Cargo.toml | 3 +- python/zarrs/__init__.py | 2 + python/zarrs/_internal.pyi | 34 +++ python/zarrs/array.py | 222 ++++++++++++++ src/array.rs | 258 ++++++++++++++++ src/lib.rs | 64 +--- src/utils.rs | 56 +++- tests/test_array.py | 590 +++++++++++++++++++++++++++++++++++++ 8 files changed, 1171 insertions(+), 58 deletions(-) create mode 100644 python/zarrs/array.py create mode 100644 src/array.rs create mode 100644 tests/test_array.py diff --git a/Cargo.toml b/Cargo.toml index 10f33fa..1c836cc 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -10,7 +10,7 @@ crate-type = ["cdylib", "rlib"] [dependencies] pyo3 = { version = "0.27.1", features = ["abi3-py311"] } -zarrs = { version = "0.23.0", features = ["async", "zlib", "pcodec", "bz2"] } +zarrs = { version = "0.23.1", features = ["async", "zlib", "pcodec", "bz2"] } 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 @@ -26,6 +26,7 @@ itertools = "0.14.0" 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 +mimalloc = { version = "0.1", default-features = false } [profile.release] lto = true diff --git a/python/zarrs/__init__.py b/python/zarrs/__init__.py index cf5b8bd..c45fe8f 100644 --- a/python/zarrs/__init__.py +++ b/python/zarrs/__init__.py @@ -1,4 +1,5 @@ from ._internal import __version__ +from .array import ZarrsArray from .pipeline import ZarrsCodecPipeline as _ZarrsCodecPipeline from .utils import CollapsedDimensionError, DiscontiguousArrayError @@ -9,6 +10,7 @@ class ZarrsCodecPipeline(_ZarrsCodecPipeline): __all__ = [ + "ZarrsArray", "ZarrsCodecPipeline", "DiscontiguousArrayError", "CollapsedDimensionError", diff --git a/python/zarrs/_internal.pyi b/python/zarrs/_internal.pyi index 4af635a..5e96304 100644 --- a/python/zarrs/_internal.pyi +++ b/python/zarrs/_internal.pyi @@ -7,6 +7,40 @@ import typing import numpy.typing import zarr.abc.store +@typing.final +class ArrayImpl: + @property + def shape(self) -> builtins.list[builtins.int]: ... + @property + def ndim(self) -> builtins.int: ... + @property + def dtype(self) -> builtins.str: ... + def __new__( + cls, + store_config: zarr.abc.store.Store, + path: builtins.str, + *, + validate_checksums: builtins.bool = False, + num_threads: builtins.int | None = None, + direct_io: builtins.bool = False, + ) -> ArrayImpl: ... + def retrieve( + self, + ranges: typing.Sequence[tuple[builtins.int, builtins.int]], + output: numpy.typing.NDArray[typing.Any], + ) -> None: ... + def store( + self, + ranges: typing.Sequence[tuple[builtins.int, builtins.int]], + input: numpy.typing.NDArray[typing.Any], + ) -> None: ... + def copy_from( + self, + source: ArrayImpl, + source_ranges: typing.Sequence[tuple[builtins.int, builtins.int]], + dest_ranges: typing.Sequence[tuple[builtins.int, builtins.int]], + ) -> None: ... + @typing.final class ChunkItem: def __new__( diff --git a/python/zarrs/array.py b/python/zarrs/array.py new file mode 100644 index 0000000..d017f72 --- /dev/null +++ b/python/zarrs/array.py @@ -0,0 +1,222 @@ +from __future__ import annotations + +import numpy as np +import zarr +from zarr.core.array import Array + +from ._internal import ArrayImpl + + +def _is_basic_indexing(key) -> bool: + """Check if key uses only int, step-1 slices, and/or a single Ellipsis.""" + if not isinstance(key, tuple): + key = (key,) + has_ellipsis = False + for k in key: + if isinstance(k, int): + continue + elif isinstance(k, slice): + if k.step is not None and k.step != 1: + return False + elif k is Ellipsis: + if has_ellipsis: + return False # multiple ellipses + has_ellipsis = True + else: + return False + return True + + +class _LazySlice: + """Lazy reference to a subset of a ZarrsArray (no I/O until consumed).""" + + __slots__ = ("_dtype", "_impl", "_ranges", "_region_shape", "_squeeze_dims") + + def __init__(self, impl_, ranges, region_shape, dtype, squeeze_dims): + self._impl = impl_ + self._ranges = ranges + self._region_shape = region_shape + self._dtype = dtype + self._squeeze_dims = squeeze_dims + + def __array__(self, dtype=None, copy=None) -> np.ndarray: + out = np.empty(self._region_shape, dtype=self._dtype) + if out.size > 0: + self._impl.retrieve(self._ranges, out) + if self._squeeze_dims: + out = out.squeeze(axis=tuple(self._squeeze_dims)) + if dtype is not None and out.dtype != dtype: + out = out.astype(dtype, copy=False) + return out + + +class _LazyIndexer: + """Proxy returned by ``ZarrsArray.lazy`` that captures indexing lazily.""" + + __slots__ = ("_pipeline",) + + def __init__(self, pipeline: ZarrsArray): + self._pipeline = pipeline + + def __getitem__(self, key: slice | int | tuple[slice | int, ...]) -> _LazySlice: + ranges, region_shape, squeeze_dims = self._pipeline._parse_key(key) + return _LazySlice( + self._pipeline._impl, + ranges, + region_shape, + self._pipeline.dtype, + squeeze_dims, + ) + + +class ZarrsArray(Array): + """zarr.Array subclass backed by zarrs for fast I/O. + + Supports all zarr.Array operations. Basic slice indexing (ints, step-1 + slices, ellipsis) is handled by the Rust fast path; advanced indexing + falls back to zarr.Array unless ``codec_pipeline.strict`` is set. + """ + + def __init__( + self, + array: Array, + *, + validate_checksums: bool = False, + chunk_concurrent_minimum: int | None = None, + num_threads: int | None = None, + direct_io: bool = False, + ) -> None: + super().__init__(array._async_array) + store = array.store_path.store + zarr_path = array.store_path.path + zarrs_path = "/" + zarr_path if zarr_path else "/" + self._impl = ArrayImpl( + store, + zarrs_path, + validate_checksums=validate_checksums, + chunk_concurrent_minimum=chunk_concurrent_minimum, + num_threads=num_threads, + direct_io=direct_io, + ) + + @property + def lazy(self) -> _LazyIndexer: + return _LazyIndexer(self) + + def _parse_key( + self, key: slice | int | tuple[slice | int, ...] + ) -> tuple[list[tuple[int, int]], list[int], list[int]]: + if not isinstance(key, tuple): + key = (key,) + + # Expand Ellipsis + if Ellipsis in key: + idx = key.index(Ellipsis) + n_explicit = len(key) - 1 # everything except the Ellipsis + n_expand = self.ndim - n_explicit + if n_expand < 0: + raise IndexError( + f"too many indices for array: " + f"array is {self.ndim}-dimensional, " + f"but {n_explicit} were indexed" + ) + key = key[:idx] + (slice(None),) * n_expand + key[idx + 1 :] + + if len(key) > self.ndim: + raise IndexError( + f"too many indices for array: " + f"array is {self.ndim}-dimensional, " + f"but {len(key)} were indexed" + ) + + # Pad missing dimensions with full slices + if len(key) < self.ndim: + key = key + (slice(None),) * (self.ndim - len(key)) + + ranges: list[tuple[int, int]] = [] + region_shape: list[int] = [] + squeeze_dims: list[int] = [] + + for i, (k, dim_size) in enumerate(zip(key, self.shape)): + if isinstance(k, int): + if k < 0: + k += dim_size + if k < 0 or k >= dim_size: + raise IndexError( + f"index {k} is out of bounds for axis {i} with size {dim_size}" + ) + ranges.append((k, k + 1)) + region_shape.append(1) + squeeze_dims.append(i) + elif isinstance(k, slice): + start, stop, step = k.indices(dim_size) + if step != 1: + raise IndexError("only step=1 slices are supported") + ranges.append((start, stop)) + region_shape.append(max(0, stop - start)) + else: + raise IndexError(f"unsupported index type: {type(k).__name__}") + + return ranges, region_shape, squeeze_dims + + def __getitem__(self, key: slice | int | tuple[slice | int, ...]) -> np.ndarray: + if _is_basic_indexing(key): + ranges, region_shape, squeeze_dims = self._parse_key(key) + out = np.empty(region_shape, dtype=self.dtype) + if out.size > 0: + self._impl.retrieve(ranges, out) + if squeeze_dims: + out = out.squeeze(axis=tuple(squeeze_dims)) + return out + + strict = zarr.config.get("codec_pipeline.strict", False) + if strict: + raise IndexError( + "ZarrsArray in strict mode does not support advanced indexing" + ) + return super().__getitem__(key) + + def __setitem__(self, key: slice | int | tuple[slice | int, ...], value) -> None: + if _is_basic_indexing(key): + ranges, region_shape, squeeze_dims = self._parse_key(key) + + if isinstance(value, _LazySlice): + if value._region_shape != region_shape: + raise ValueError( + f"could not broadcast input array from shape " + f"{tuple(value._region_shape)} " + f"into shape {tuple(region_shape)}" + ) + if all(s > 0 for s in region_shape): + self._impl.copy_from(value._impl, value._ranges, ranges) + return + + value = np.asarray(value, dtype=self.dtype) + + # Ensure native byte order + if not value.dtype.isnative: + value = value.byteswap().view(value.dtype.newbyteorder("=")) + + # Expand squeezed dimensions back + for dim in squeeze_dims: + value = np.expand_dims(value, axis=dim) + + if value.shape != tuple(region_shape): + raise ValueError( + f"could not broadcast input array from shape {value.shape} " + f"into shape {tuple(region_shape)}" + ) + + # Ensure C-contiguous before passing to Rust + value = np.ascontiguousarray(value) + + if value.size > 0: + self._impl.store(ranges, value) + return + + strict = zarr.config.get("codec_pipeline.strict", False) + if strict: + raise IndexError( + "ZarrsArray in strict mode does not support advanced indexing" + ) + super().__setitem__(key, value) diff --git a/src/array.rs b/src/array.rs new file mode 100644 index 0000000..2229d2a --- /dev/null +++ b/src/array.rs @@ -0,0 +1,258 @@ +use std::borrow::Cow; + +use numpy::{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 zarrs::array::concurrency::concurrency_chunks_and_codec; +use zarrs::array::{ + Array, ArrayBytes, ArrayBytesDecodeIntoTarget, ArrayBytesFixedDisjointView, ArrayCodecTraits, + ArrayIndicesTinyVec, ArraySubset, CodecOptions, +}; +use zarrs::plugin::{ExtensionName, ZarrVersion}; +use zarrs::storage::ReadableWritableListableStorage; + +use crate::store::StoreConfig; +use crate::utils::{PyErrExt as _, nparray_to_slice, nparray_to_unsafe_cell_slice}; + +fn data_type_to_numpy_str(dt: &zarrs::array::DataType) -> PyResult { + let name = dt + .name(ZarrVersion::V3) + .ok_or_else(|| PyErr::new::("unknown data type"))?; + match name.as_ref() { + "bool" | "int8" | "int16" | "int32" | "int64" | "uint8" | "uint16" | "uint32" + | "uint64" | "float16" | "float32" | "float64" | "complex64" | "complex128" => { + Ok(name.into_owned()) + } + _ => Err(PyErr::new::(format!( + "unsupported data type for numpy: {name}" + ))), + } +} + +#[gen_stub_pyclass] +#[pyclass] +pub struct ArrayImpl { + array: Array, + codec_options: CodecOptions, + data_type_size: usize, +} + +#[gen_stub_pymethods] +#[pymethods] +impl ArrayImpl { + #[pyo3(signature = ( + store_config, + path, + *, + validate_checksums=false, + chunk_concurrent_minimum=None, + num_threads=None, + direct_io=false, + ))] + #[new] + fn new( + mut store_config: StoreConfig, + path: &str, + validate_checksums: bool, + chunk_concurrent_minimum: Option, + num_threads: Option, + direct_io: bool, + ) -> PyResult { + store_config.direct_io(direct_io); + let store: ReadableWritableListableStorage = + (&store_config).try_into().map_py_err::()?; + let array = Array::open(store, path).map_py_err::()?; + let data_type_size = array.data_type().fixed_size().ok_or_else(|| { + PyErr::new::("variable length data type not supported") + })?; + let mut codec_options = CodecOptions::default() + .with_validate_checksums(validate_checksums) + .with_concurrent_target(num_threads.unwrap_or(rayon::current_num_threads())); + if let Some(chunk_concurrent_minimum) = chunk_concurrent_minimum { + codec_options.set_chunk_concurrent_minimum(chunk_concurrent_minimum); + } + Ok(Self { + array, + codec_options, + data_type_size, + }) + } + + #[getter] + fn shape(&self) -> Vec { + self.array.shape().to_vec() + } + + #[getter] + fn ndim(&self) -> usize { + self.array.dimensionality() + } + + #[getter] + fn dtype(&self) -> PyResult { + data_type_to_numpy_str(self.array.data_type()) + } + + fn retrieve( + &self, + py: Python, + ranges: Vec<(u64, u64)>, + output: &Bound<'_, PyUntypedArray>, + ) -> PyResult<()> { + let subset_ranges: Vec> = + ranges.iter().map(|&(start, stop)| start..stop).collect(); + let array_subset = ArraySubset::new_with_ranges(&subset_ranges); + + let output_cell_slice = nparray_to_unsafe_cell_slice(output)?; + let output_shape: Vec = output + .shape() + .iter() + .map(|&s| u64::try_from(s).map_err(|e| PyErr::new::(format!("{e}")))) + .collect::>>()?; + + let data_type_size = self.data_type_size; + let full_output_subset = ArraySubset::new_with_shape(output_shape.clone()); + + py.detach(move || { + let mut output_view = unsafe { + // SAFETY: we are the sole writer to this output array + ArrayBytesFixedDisjointView::new( + output_cell_slice, + data_type_size, + &output_shape, + full_output_subset, + ) + .map_py_err::()? + }; + let target = ArrayBytesDecodeIntoTarget::Fixed(&mut output_view); + self.array + .retrieve_array_subset_into_opt(&array_subset, target, &self.codec_options) + .map_py_err::() + }) + } + + fn store( + &self, + py: Python, + ranges: Vec<(u64, u64)>, + input: &Bound<'_, PyUntypedArray>, + ) -> PyResult<()> { + let subset_ranges: Vec> = + ranges.iter().map(|&(start, stop)| start..stop).collect(); + let array_subset = ArraySubset::new_with_ranges(&subset_ranges); + + let input_slice = nparray_to_slice(input)?; + let array_bytes = ArrayBytes::new_flen(Cow::Borrowed(input_slice)); + + py.detach(move || { + self.array + .store_array_subset_opt(&array_subset, array_bytes, &self.codec_options) + .map_py_err::() + }) + } + + fn copy_from( + &self, + py: Python, + source: &ArrayImpl, + source_ranges: Vec<(u64, u64)>, + dest_ranges: Vec<(u64, u64)>, + ) -> PyResult<()> { + let source_subset = ArraySubset::new_with_ranges( + &source_ranges.iter().map(|&(s, e)| s..e).collect::>(), + ); + let dest_subset = ArraySubset::new_with_ranges( + &dest_ranges.iter().map(|&(s, e)| s..e).collect::>(), + ); + + if source_subset.shape() != dest_subset.shape() { + return Err(PyErr::new::( + "source and destination region shapes must match", + )); + } + + py.detach(move || { + let chunks = self + .array + .chunks_in_array_subset(&dest_subset) + .map_py_err::()? + .ok_or_else(|| { + PyErr::new::("failed to compute overlapping chunks") + })?; + let num_chunks = chunks.num_elements_usize(); + + // Calculate chunk/codec concurrency + let chunk_shape = self + .array + .chunk_shape(&vec![0; self.array.dimensionality()]) + .map_py_err::()?; + let codec_concurrency = self + .array + .codecs() + .recommended_concurrency(&chunk_shape, self.array.data_type()) + .map_py_err::()?; + let (chunk_concurrent_limit, codec_options) = concurrency_chunks_and_codec( + self.codec_options.concurrent_target(), + num_chunks, + &self.codec_options, + &codec_concurrency, + ); + + let dest_start = dest_subset.start(); + let source_start = source_subset.start(); + // let source_cache = ArrayShardedReadableExtCache::new(&source.array); + + let copy_chunk = |chunk_indices: ArrayIndicesTinyVec| -> PyResult<()> { + let chunk_subset = self + .array + .chunk_subset_bounded(&chunk_indices) + .map_py_err::()?; + + let overlap = chunk_subset + .overlap(&dest_subset) + .map_py_err::()?; + if overlap.is_empty() { + return Ok(()); + } + + // Map overlap coordinates from dest space to source space + let source_overlap_ranges: Vec> = overlap + .to_ranges() + .iter() + .enumerate() + .map(|(i, range)| { + let offset = range.start - dest_start[i] + source_start[i]; + let len = range.end - range.start; + offset..(offset + len) + }) + .collect(); + let source_overlap = ArraySubset::new_with_ranges(&source_overlap_ranges); + + // NOTE: Could retrieve into a pre-allocated buffer (per thread) with `regular` grid + let data: ArrayBytes = source + .array + .retrieve_array_subset_opt(&source_overlap, &source.codec_options) + // .retrieve_array_subset_sharded_opt( + // &source_cache, + // &source_overlap, + // &source.codec_options, + // ) + .map_py_err::()?; + + self.array + .store_array_subset_opt(&overlap, data, &codec_options) + .map_py_err::()?; + + Ok(()) + }; + + let indices = chunks.indices(); + iter_concurrent_limit!(chunk_concurrent_limit, indices, try_for_each, copy_chunk)?; + + Ok(()) + }) + } +} diff --git a/src/lib.rs b/src/lib.rs index 1b7c08e..84aa8e6 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,22 +1,22 @@ #![warn(clippy::pedantic)] #![allow(clippy::module_name_repetitions)] +#[global_allocator] +static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc; + 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 numpy::{PyUntypedArray, PyUntypedArrayMethods}; use pyo3::exceptions::{PyRuntimeError, PyTypeError, PyValueError}; 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, @@ -28,6 +28,7 @@ use zarrs::convert::array_metadata_v2_to_v3; use zarrs::plugin::ZarrVersion; use zarrs::storage::{ReadableWritableListableStorage, StorageHandle, StoreKey}; +mod array; mod chunk_item; mod concurrency; mod runtime; @@ -154,56 +155,6 @@ impl CodecPipelineImpl { 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(), - )); - } - 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] @@ -287,7 +238,7 @@ impl CodecPipelineImpl { value: &Bound<'_, PyUntypedArray>, ) -> PyResult<()> { // Get input array - let output = Self::nparray_to_unsafe_cell_slice(value)?; + let output = utils::nparray_to_unsafe_cell_slice(value)?; let output_shape: Vec = value.shape_zarr()?; // Adjust the concurrency based on the codec chain and the first chunk description @@ -403,7 +354,7 @@ impl CodecPipelineImpl { } // Get input array - let input_slice = Self::nparray_to_slice(value)?; + let input_slice = utils::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))) @@ -468,6 +419,7 @@ fn _internal(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add("__version__", env!("CARGO_PKG_VERSION"))?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; Ok(()) } diff --git a/src/utils.rs b/src/utils.rs index d663b5c..984b097 100644 --- a/src/utils.rs +++ b/src/utils.rs @@ -1,7 +1,11 @@ use std::fmt::Display; +use std::ptr::NonNull; -use numpy::{PyUntypedArray, PyUntypedArrayMethods}; +use numpy::npyffi::PyArrayObject; +use numpy::{PyArrayDescrMethods, PyUntypedArray, PyUntypedArrayMethods}; +use pyo3::exceptions::PyValueError; use pyo3::{Bound, PyErr, PyResult, PyTypeInfo}; +use unsafe_cell_slice::UnsafeCellSlice; use zarrs::array::CodecError; use crate::ChunkItem; @@ -59,3 +63,53 @@ pub fn is_whole_chunk(item: &ChunkItem) -> bool { item.chunk_subset.start().iter().all(|&o| o == 0) && item.chunk_subset.shape() == bytemuck::must_cast_slice::<_, u64>(&item.shape) } + +pub 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 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 = 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 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 = 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)) +} diff --git a/tests/test_array.py b/tests/test_array.py new file mode 100644 index 0000000..ca5ee62 --- /dev/null +++ b/tests/test_array.py @@ -0,0 +1,590 @@ +from __future__ import annotations + +import numpy as np +import pytest +import zarr +from zarr.storage import LocalStore + +from zarrs import ZarrsArray + + +@pytest.fixture +def store(tmp_path): + return LocalStore(str(tmp_path / "test.zarr")) + + +def save_array(store, data, *, chunks=None): + z = zarr.open_array( + store, mode="w", shape=data.shape, dtype=data.dtype, chunks=chunks + ) + z[:] = data + + +def open_zarrs(store, **kwargs): + return ZarrsArray(zarr.open_array(store), **kwargs) + + +class TestProperties: + def test_shape_2d(self, store): + save_array(store, np.zeros((10, 20), dtype="float32")) + arr = open_zarrs(store) + assert arr.shape == (10, 20) + + def test_shape_1d(self, store): + save_array(store, np.zeros(50, dtype="int32")) + arr = open_zarrs(store) + assert arr.shape == (50,) + + def test_shape_3d(self, store): + save_array(store, np.zeros((3, 4, 5), dtype="uint8")) + arr = open_zarrs(store) + assert arr.shape == (3, 4, 5) + + def test_ndim(self, store): + save_array(store, np.zeros((3, 4, 5), dtype="float64")) + arr = open_zarrs(store) + assert arr.ndim == 3 + + def test_size(self, store): + save_array(store, np.zeros((3, 4, 5), dtype="float64")) + arr = open_zarrs(store) + assert arr.size == 60 + + def test_dtype(self, store): + save_array(store, np.zeros(5, dtype="float32")) + arr = open_zarrs(store) + assert arr.dtype == np.dtype("float32") + + +class TestDtypes: + @pytest.mark.parametrize( + "dtype", + [ + "int8", + "int16", + "int32", + "int64", + "uint8", + "uint16", + "uint32", + "uint64", + "float32", + "float64", + ], + ) + def test_dtype_roundtrip(self, tmp_path, dtype): + data = np.arange(12, dtype=dtype).reshape(3, 4) + store = LocalStore(str(tmp_path / "test.zarr")) + save_array(store, data) + arr = open_zarrs(store) + assert arr.dtype == np.dtype(dtype) + np.testing.assert_array_equal(arr[:], data) + + def test_dtype_bool(self, tmp_path): + data = np.array([[True, False, True], [False, True, False]]) + store = LocalStore(str(tmp_path / "test.zarr")) + save_array(store, data) + arr = open_zarrs(store) + assert arr.dtype == np.dtype("bool") + np.testing.assert_array_equal(arr[:], data) + + +class TestFullRead: + def test_full_read_1d(self, store): + data = np.arange(20, dtype="int64") + save_array(store, data) + arr = open_zarrs(store) + np.testing.assert_array_equal(arr[:], data) + + def test_full_read_2d(self, store): + data = np.arange(100, dtype="float32").reshape(10, 10) + save_array(store, data) + arr = open_zarrs(store) + np.testing.assert_array_equal(arr[:], data) + + def test_full_read_3d(self, store): + data = np.arange(60, dtype="float64").reshape(3, 4, 5) + save_array(store, data) + arr = open_zarrs(store) + np.testing.assert_array_equal(arr[:], data) + + +class TestSliceRead: + def test_slice_1d(self, store): + data = np.arange(20, dtype="int32") + save_array(store, data) + arr = open_zarrs(store) + np.testing.assert_array_equal(arr[5:10], data[5:10]) + + def test_slice_2d(self, store): + data = np.arange(100, dtype="float32").reshape(10, 10) + save_array(store, data) + arr = open_zarrs(store) + np.testing.assert_array_equal(arr[2:5, 3:7], data[2:5, 3:7]) + + def test_slice_3d(self, store): + data = np.arange(120, dtype="float64").reshape(4, 5, 6) + save_array(store, data) + arr = open_zarrs(store) + np.testing.assert_array_equal(arr[1:3, 2:4, 0:3], data[1:3, 2:4, 0:3]) + + def test_partial_dims(self, store): + data = np.arange(100, dtype="float32").reshape(10, 10) + save_array(store, data) + arr = open_zarrs(store) + np.testing.assert_array_equal(arr[2:4], data[2:4]) + + def test_across_chunks(self, store): + data = np.arange(100, dtype="int32").reshape(10, 10) + save_array(store, data, chunks=(3, 3)) + arr = open_zarrs(store) + np.testing.assert_array_equal(arr[1:8, 2:9], data[1:8, 2:9]) + + +class TestIntegerIndex: + def test_integer_2d(self, store): + data = np.arange(30, dtype="float64").reshape(5, 6) + save_array(store, data) + arr = open_zarrs(store) + result = arr[2] + np.testing.assert_array_equal(result, data[2]) + assert result.shape == (6,) + + def test_integer_3d(self, store): + data = np.arange(60, dtype="float64").reshape(3, 4, 5) + save_array(store, data) + arr = open_zarrs(store) + result = arr[1] + np.testing.assert_array_equal(result, data[1]) + assert result.shape == (4, 5) + + def test_integer_1d_scalar(self, store): + data = np.arange(10, dtype="int32") + save_array(store, data) + arr = open_zarrs(store) + result = arr[3] + np.testing.assert_array_equal(result, data[3]) + assert result.shape == () + + def test_mixed_int_and_slice(self, store): + data = np.arange(100, dtype="float32").reshape(10, 10) + save_array(store, data) + arr = open_zarrs(store) + result = arr[3, 2:5] + np.testing.assert_array_equal(result, data[3, 2:5]) + assert result.shape == (3,) + + def test_multiple_integers(self, store): + data = np.arange(60, dtype="float64").reshape(3, 4, 5) + save_array(store, data) + arr = open_zarrs(store) + result = arr[1, 2] + np.testing.assert_array_equal(result, data[1, 2]) + assert result.shape == (5,) + + +class TestNegativeIndex: + def test_negative_integer(self, store): + data = np.arange(10, dtype="int32") + save_array(store, data) + arr = open_zarrs(store) + np.testing.assert_array_equal(arr[-1], data[-1]) + + def test_negative_slice(self, store): + data = np.arange(10, dtype="int32") + save_array(store, data) + arr = open_zarrs(store) + np.testing.assert_array_equal(arr[-3:], data[-3:]) + + def test_negative_both(self, store): + data = np.arange(20, dtype="float64").reshape(4, 5) + save_array(store, data) + arr = open_zarrs(store) + np.testing.assert_array_equal(arr[-2:, -3:], data[-2:, -3:]) + + +class TestEmptySlice: + def test_empty_1d(self, store): + data = np.arange(10, dtype="int64") + save_array(store, data) + arr = open_zarrs(store) + result = arr[3:3] + assert result.shape == (0,) + assert result.dtype == np.int64 + + def test_empty_2d(self, store): + data = np.arange(20, dtype="float32").reshape(4, 5) + save_array(store, data) + arr = open_zarrs(store) + result = arr[2:2, :] + assert result.shape == (0, 5) + + +class TestEdgeCases: + def test_single_element(self, store): + data = np.array([42], dtype="int32") + save_array(store, data) + arr = open_zarrs(store) + np.testing.assert_array_equal(arr[:], data) + + def test_large_chunk_small_slice(self, store): + data = np.arange(1000, dtype="float32").reshape(10, 100) + save_array(store, data, chunks=(10, 100)) + arr = open_zarrs(store) + np.testing.assert_array_equal(arr[0:1, 0:1], data[0:1, 0:1]) + + +class TestErrors: + def test_too_many_indices(self, store): + save_array(store, np.zeros((3, 4), dtype="float32")) + arr = open_zarrs(store) + with pytest.raises(IndexError, match="too many indices"): + arr[0, 0, 0] + + def test_out_of_bounds_integer(self, store): + save_array(store, np.zeros(5, dtype="float32")) + arr = open_zarrs(store) + with pytest.raises(IndexError, match="out of bounds"): + arr[5] + + +class TestWrite: + def test_full_write_1d(self, store): + save_array(store, np.zeros(10, dtype="int32")) + arr = open_zarrs(store) + data = np.arange(10, dtype="int32") + arr[:] = data + np.testing.assert_array_equal(arr[:], data) + + def test_full_write_2d(self, store): + save_array(store, np.zeros((4, 5), dtype="float64")) + arr = open_zarrs(store) + data = np.arange(20, dtype="float64").reshape(4, 5) + arr[:] = data + np.testing.assert_array_equal(arr[:], data) + + def test_full_write_3d(self, store): + save_array(store, np.zeros((2, 3, 4), dtype="float32")) + arr = open_zarrs(store) + data = np.arange(24, dtype="float32").reshape(2, 3, 4) + arr[:] = data + np.testing.assert_array_equal(arr[:], data) + + +class TestWriteSlice: + def test_slice_1d(self, store): + data = np.zeros(10, dtype="int32") + save_array(store, data) + arr = open_zarrs(store) + arr[3:7] = np.array([10, 20, 30, 40], dtype="int32") + expected = data.copy() + expected[3:7] = [10, 20, 30, 40] + np.testing.assert_array_equal(arr[:], expected) + + def test_slice_2d(self, store): + data = np.zeros((5, 6), dtype="float64") + save_array(store, data) + arr = open_zarrs(store) + patch = np.ones((2, 3), dtype="float64") * 99 + arr[1:3, 2:5] = patch + expected = data.copy() + expected[1:3, 2:5] = 99 + np.testing.assert_array_equal(arr[:], expected) + + def test_partial_dims(self, store): + data = np.zeros((4, 5), dtype="int64") + save_array(store, data) + arr = open_zarrs(store) + patch = np.ones((2, 5), dtype="int64") * 7 + arr[1:3] = patch + expected = data.copy() + expected[1:3] = 7 + np.testing.assert_array_equal(arr[:], expected) + + def test_across_chunks(self, store): + data = np.zeros((10, 10), dtype="int32") + save_array(store, data, chunks=(3, 3)) + arr = open_zarrs(store) + patch = np.arange(49, dtype="int32").reshape(7, 7) + arr[1:8, 2:9] = patch + expected = data.copy() + expected[1:8, 2:9] = patch + np.testing.assert_array_equal(arr[:], expected) + + +class TestWriteIntegerIndex: + def test_integer_2d(self, store): + data = np.zeros((5, 6), dtype="float64") + save_array(store, data) + arr = open_zarrs(store) + row = np.arange(6, dtype="float64") + 1 + arr[2] = row + expected = data.copy() + expected[2] = row + np.testing.assert_array_equal(arr[:], expected) + + def test_integer_1d_scalar(self, store): + data = np.zeros(10, dtype="int32") + save_array(store, data) + arr = open_zarrs(store) + arr[3] = 42 + expected = data.copy() + expected[3] = 42 + np.testing.assert_array_equal(arr[:], expected) + + def test_mixed_int_and_slice(self, store): + data = np.zeros((10, 10), dtype="float32") + save_array(store, data) + arr = open_zarrs(store) + arr[3, 2:5] = np.array([10, 20, 30], dtype="float32") + expected = data.copy() + expected[3, 2:5] = [10, 20, 30] + np.testing.assert_array_equal(arr[:], expected) + + def test_multiple_integers(self, store): + data = np.zeros((3, 4, 5), dtype="float64") + save_array(store, data) + arr = open_zarrs(store) + arr[1, 2] = np.arange(5, dtype="float64") + expected = data.copy() + expected[1, 2] = np.arange(5, dtype="float64") + np.testing.assert_array_equal(arr[:], expected) + + +class TestWriteErrors: + def test_shape_mismatch(self, store): + save_array(store, np.zeros((4, 5), dtype="float32")) + arr = open_zarrs(store) + with pytest.raises(ValueError, match="could not broadcast"): + arr[0:2, 0:3] = np.zeros((3, 2), dtype="float32") + + def test_too_many_indices(self, store): + save_array(store, np.zeros((3, 4), dtype="float32")) + arr = open_zarrs(store) + with pytest.raises(IndexError, match="too many indices"): + arr[0, 0, 0] = 1.0 + + def test_out_of_bounds(self, store): + save_array(store, np.zeros(5, dtype="float32")) + arr = open_zarrs(store) + with pytest.raises(IndexError, match="out of bounds"): + arr[5] = 1.0 + + def test_dtype_coercion(self, store): + data = np.zeros(5, dtype="float64") + save_array(store, data) + arr = open_zarrs(store) + arr[0:3] = np.array([1, 2, 3], dtype="int32") + expected = data.copy() + expected[0:3] = [1, 2, 3] + np.testing.assert_array_equal(arr[:], expected) + + +class TestLazyCopy: + def test_full_copy(self, tmp_path): + data = np.arange(20, dtype="float64").reshape(4, 5) + src_store = LocalStore(str(tmp_path / "src.zarr")) + dst_store = LocalStore(str(tmp_path / "dst.zarr")) + save_array(src_store, data) + save_array(dst_store, np.zeros_like(data)) + src = open_zarrs(src_store) + dst = open_zarrs(dst_store) + dst[:] = src.lazy[:] + np.testing.assert_array_equal(dst[:], data) + + def test_slice_copy(self, tmp_path): + data = np.arange(100, dtype="int32").reshape(10, 10) + src_store = LocalStore(str(tmp_path / "src.zarr")) + dst_store = LocalStore(str(tmp_path / "dst.zarr")) + save_array(src_store, data) + save_array(dst_store, np.zeros_like(data)) + src = open_zarrs(src_store) + dst = open_zarrs(dst_store) + dst[2:5, 3:7] = src.lazy[2:5, 3:7] + expected = np.zeros_like(data) + expected[2:5, 3:7] = data[2:5, 3:7] + np.testing.assert_array_equal(dst[:], expected) + + def test_different_regions(self, tmp_path): + data = np.arange(100, dtype="float32").reshape(10, 10) + src_store = LocalStore(str(tmp_path / "src.zarr")) + dst_store = LocalStore(str(tmp_path / "dst.zarr")) + save_array(src_store, data) + save_array(dst_store, np.zeros_like(data)) + src = open_zarrs(src_store) + dst = open_zarrs(dst_store) + dst[0:3, 0:4] = src.lazy[5:8, 6:10] + expected = np.zeros_like(data) + expected[0:3, 0:4] = data[5:8, 6:10] + np.testing.assert_array_equal(dst[:], expected) + + def test_across_chunks(self, tmp_path): + data = np.arange(100, dtype="int32").reshape(10, 10) + src_store = LocalStore(str(tmp_path / "src.zarr")) + dst_store = LocalStore(str(tmp_path / "dst.zarr")) + save_array(src_store, data, chunks=(3, 3)) + save_array(dst_store, np.zeros_like(data), chunks=(4, 4)) + src = open_zarrs(src_store) + dst = open_zarrs(dst_store) + dst[1:8, 2:9] = src.lazy[1:8, 2:9] + expected = np.zeros_like(data) + expected[1:8, 2:9] = data[1:8, 2:9] + np.testing.assert_array_equal(dst[:], expected) + + def test_integer_index(self, tmp_path): + data = np.arange(30, dtype="float64").reshape(5, 6) + src_store = LocalStore(str(tmp_path / "src.zarr")) + dst_store = LocalStore(str(tmp_path / "dst.zarr")) + save_array(src_store, data) + save_array(dst_store, np.zeros_like(data)) + src = open_zarrs(src_store) + dst = open_zarrs(dst_store) + dst[2] = src.lazy[3] + expected = np.zeros_like(data) + expected[2] = data[3] + np.testing.assert_array_equal(dst[:], expected) + + def test_shape_mismatch(self, tmp_path): + src_store = LocalStore(str(tmp_path / "src.zarr")) + dst_store = LocalStore(str(tmp_path / "dst.zarr")) + save_array(src_store, np.zeros((10, 10), dtype="float32")) + save_array(dst_store, np.zeros((10, 10), dtype="float32")) + src = open_zarrs(src_store) + dst = open_zarrs(dst_store) + with pytest.raises(ValueError, match="could not broadcast"): + dst[0:2, 0:3] = src.lazy[0:3, 0:2] + + def test_empty_region(self, tmp_path): + src_store = LocalStore(str(tmp_path / "src.zarr")) + dst_store = LocalStore(str(tmp_path / "dst.zarr")) + save_array(src_store, np.zeros(10, dtype="int64")) + save_array(dst_store, np.zeros(10, dtype="int64")) + src = open_zarrs(src_store) + dst = open_zarrs(dst_store) + dst[3:3] = src.lazy[5:5] # empty region, should be a no-op + + def test_np_asarray_full(self, store): + data = np.arange(20, dtype="float64").reshape(4, 5) + save_array(store, data) + arr = open_zarrs(store) + result = np.asarray(arr.lazy[:]) + np.testing.assert_array_equal(result, data) + + def test_np_asarray_slice(self, store): + data = np.arange(100, dtype="int32").reshape(10, 10) + save_array(store, data) + arr = open_zarrs(store) + result = np.asarray(arr.lazy[2:5, 3:7]) + np.testing.assert_array_equal(result, data[2:5, 3:7]) + + def test_np_asarray_integer_index(self, store): + data = np.arange(30, dtype="float64").reshape(5, 6) + save_array(store, data) + arr = open_zarrs(store) + result = np.asarray(arr.lazy[2]) + np.testing.assert_array_equal(result, data[2]) + assert result.shape == (6,) + + def test_np_asarray_dtype_cast(self, store): + data = np.arange(10, dtype="int32") + save_array(store, data) + arr = open_zarrs(store) + result = np.asarray(arr.lazy[:], dtype="float64") + np.testing.assert_array_equal(result, data.astype("float64")) + assert result.dtype == np.float64 + + +class TestSubclass: + def test_isinstance(self, store): + save_array(store, np.zeros((3, 4), dtype="float32")) + arr = open_zarrs(store) + assert isinstance(arr, zarr.Array) + + def test_from_zarr_array(self, store): + data = np.arange(20, dtype="float64").reshape(4, 5) + save_array(store, data) + zarr_arr = zarr.open_array(store) + arr = ZarrsArray(zarr_arr) + np.testing.assert_array_equal(arr[:], data) + assert arr.shape == (4, 5) + + def test_zarr_properties(self, store): + save_array(store, np.zeros((3, 4), dtype="float32")) + arr = open_zarrs(store) + assert arr.store_path is not None + assert arr.metadata is not None + + +class TestEllipsis: + def test_ellipsis_full(self, store): + data = np.arange(12, dtype="float32").reshape(3, 4) + save_array(store, data) + arr = open_zarrs(store) + np.testing.assert_array_equal(arr[...], data) + + def test_ellipsis_leading(self, store): + data = np.arange(60, dtype="float64").reshape(3, 4, 5) + save_array(store, data) + arr = open_zarrs(store) + np.testing.assert_array_equal(arr[..., 1:3], data[..., 1:3]) + + def test_ellipsis_trailing(self, store): + data = np.arange(60, dtype="float64").reshape(3, 4, 5) + save_array(store, data) + arr = open_zarrs(store) + np.testing.assert_array_equal(arr[1, ...], data[1, ...]) + + def test_ellipsis_middle(self, store): + data = np.arange(60, dtype="float64").reshape(3, 4, 5) + save_array(store, data) + arr = open_zarrs(store) + np.testing.assert_array_equal(arr[1, ..., 2:4], data[1, ..., 2:4]) + + +class TestFallback: + def test_step_slicing(self, store): + data = np.arange(10, dtype="int32") + save_array(store, data) + arr = open_zarrs(store) + np.testing.assert_array_equal(arr[::2], data[::2]) + + def test_fancy_indexing(self, store): + data = np.arange(10, dtype="int32") + save_array(store, data) + arr = open_zarrs(store) + indices = np.array([1, 3, 5]) + np.testing.assert_array_equal(arr[indices], data[indices]) + + +class TestStrictMode: + def test_strict_rejects_step(self, store): + save_array(store, np.zeros(10, dtype="float32")) + arr = open_zarrs(store) + with ( + zarr.config.set({"codec_pipeline.strict": True}), + pytest.raises(IndexError, match="advanced indexing"), + ): + arr[::2] + + def test_strict_rejects_fancy(self, store): + save_array(store, np.zeros(10, dtype="float32")) + arr = open_zarrs(store) + with ( + zarr.config.set({"codec_pipeline.strict": True}), + pytest.raises(IndexError, match="advanced indexing"), + ): + arr[np.array([1, 2])] + + def test_strict_allows_basic(self, store): + data = np.arange(10, dtype="int32") + save_array(store, data) + arr = open_zarrs(store) + with zarr.config.set({"codec_pipeline.strict": True}): + np.testing.assert_array_equal(arr[2:5], data[2:5]) + + def test_strict_allows_ellipsis(self, store): + data = np.arange(12, dtype="float32").reshape(3, 4) + save_array(store, data) + arr = open_zarrs(store) + with zarr.config.set({"codec_pipeline.strict": True}): + np.testing.assert_array_equal(arr[...], data) From 195f0a28020734f8dc2e6dc102d510f8fe989614 Mon Sep 17 00:00:00 2001 From: Lachlan Deakin Date: Sun, 8 Feb 2026 15:11:24 +1100 Subject: [PATCH 2/4] fix: update test to use moved `nparray_to_unsafe_cell_slice` --- src/tests.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/tests.rs b/src/tests.rs index 9f1f8aa..384eaa9 100644 --- a/src/tests.rs +++ b/src/tests.rs @@ -6,7 +6,7 @@ use pyo3::{ types::{PyAnyMethods, PyModule}, }; -use crate::CodecPipelineImpl; +use crate::utils::nparray_to_unsafe_cell_slice; #[test] fn test_nparray_to_unsafe_cell_slice_empty() -> PyResult<()> { @@ -26,7 +26,7 @@ fn test_nparray_to_unsafe_cell_slice_empty() -> PyResult<()> { .call0()? .extract()?; - let slice = CodecPipelineImpl::nparray_to_unsafe_cell_slice(&arr)?; + let slice = nparray_to_unsafe_cell_slice(&arr)?; assert!(slice.is_empty()); Ok(()) }) From 8ca7cde7dca8e2437a9e353f62914e5a621c22ce Mon Sep 17 00:00:00 2001 From: Lachlan Deakin Date: Sun, 8 Feb 2026 15:13:01 +1100 Subject: [PATCH 3/4] fix: use of deprecated `pyo3` methods --- src/tests.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/tests.rs b/src/tests.rs index 384eaa9..714a361 100644 --- a/src/tests.rs +++ b/src/tests.rs @@ -10,8 +10,8 @@ use crate::utils::nparray_to_unsafe_cell_slice; #[test] fn test_nparray_to_unsafe_cell_slice_empty() -> PyResult<()> { - pyo3::prepare_freethreaded_python(); - Python::with_gil(|py| { + Python::initialize(); + Python::attach(|py| { let arr: Bound<'_, PyUntypedArray> = PyModule::from_code( py, c_str!( From 91e10c56dda66f3bd564c9db9c2e0c79eccc54df Mon Sep 17 00:00:00 2001 From: Lachlan Deakin Date: Sun, 8 Feb 2026 16:34:30 +1100 Subject: [PATCH 4/4] chore: update stubs --- python/zarrs/_internal.pyi | 1 + 1 file changed, 1 insertion(+) diff --git a/python/zarrs/_internal.pyi b/python/zarrs/_internal.pyi index 5e96304..d8545ce 100644 --- a/python/zarrs/_internal.pyi +++ b/python/zarrs/_internal.pyi @@ -21,6 +21,7 @@ class ArrayImpl: path: builtins.str, *, validate_checksums: builtins.bool = False, + chunk_concurrent_minimum: builtins.int | None = None, num_threads: builtins.int | None = None, direct_io: builtins.bool = False, ) -> ArrayImpl: ...