fix: handle dropped dims

Author: ilan-gold

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.6", 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

python/zarrs/utils.py

@@ -60,7 +60,7 @@ def make_chunk_info_for_rust_with_indices(
     drop_axes: tuple[int, ...],
     shape: tuple[int, ...],
 ) -> RustChunkInfo:
-    shape = shape if shape else (1,)  # constant array
+    is_shape_constant = shape == ()
     chunk_info_with_indices: list[ChunkItem] = []
     write_empty_chunks: bool = True
     for (
@@ -71,15 +71,28 @@ def make_chunk_info_for_rust_with_indices(
         _,
     ) in batch_info:
         write_empty_chunks = chunk_spec.config.write_empty_chunks
-        out_selection_as_slices = selector_tuple_to_slice_selection(out_selection)
         chunk_selection_as_slices = selector_tuple_to_slice_selection(chunk_selection)
+        out_selection_as_slices_iter = iter(selector_tuple_to_slice_selection(out_selection))
+        out_selection_as_slices_expanded = list(next(out_selection_as_slices_iter) if not isinstance(c, int) else slice(0, 1) for c in chunk_selection)
+        if is_shape_constant:
+            io_array_shape = (1,)
+        else:
+            shape_iter = iter(shape)
+            try:
+                io_array_shape = list(next(shape_iter) if not isinstance(c, int) else 1 for c in chunk_selection)
+            except RuntimeError as e:
+                if isinstance(e.__cause__, StopIteration):
+                    raise IndexError(
+                        f"the size of the chunk subset {chunk_selection} and input/output subset {shape} are incompatible"
+                    ) from None
+                raise
         chunk_info_with_indices.append(
             ChunkItem(
                 key=byte_getter.path,
                 chunk_subset=chunk_selection_as_slices,
                 chunk_shape=chunk_spec.shape,
-                subset=out_selection_as_slices,
-                shape=shape,
+                subset=out_selection_as_slices_expanded,
+                shape=io_array_shape,
             )
         )
     return RustChunkInfo(chunk_info_with_indices, write_empty_chunks)

src/chunk_item.rs

@@ -31,6 +31,7 @@ pub(crate) struct ChunkItem {
     pub chunk_subset: ArraySubset,
     pub subset: ArraySubset,
     pub shape: Vec<NonZeroU64>,
+    pub output_shape: Vec<NonZeroU64>,
     pub num_elements: u64,
 }
 
@@ -65,6 +66,7 @@ impl ChunkItem {
             subset,
             shape: chunk_shape_nonzero_u64,
             num_elements,
+            output_shape: shape_nonzero_u64
         })
     }
 }

src/lib.rs

@@ -38,7 +38,7 @@ mod utils;
 
 use crate::concurrency::ChunkConcurrentLimitAndCodecOptions;
 use crate::store::StoreConfig;
-use crate::utils::{PyCodecErrExt, PyErrExt as _, PyUntypedArrayExt as _};
+use crate::utils::{PyCodecErrExt, PyErrExt as _};
 
 // TODO: Use a OnceLock for store with get_or_try_init when stabilised?
 #[gen_stub_pyclass]
@@ -288,7 +288,6 @@ impl CodecPipelineImpl {
     ) -> PyResult<()> {
         // Get input array
         let output = Self::nparray_to_unsafe_cell_slice(value)?;
-        let output_shape: Vec<u64> = value.shape_zarr()?;
 
         // Adjust the concurrency based on the codec chain and the first chunk description
         let Some((chunk_concurrent_limit, codec_options)) =
@@ -343,7 +342,7 @@ impl CodecPipelineImpl {
                             .fixed_size()
                             .ok_or("variable length data type not supported")
                             .map_py_err::<PyTypeError>()?,
-                        &output_shape,
+                        bytemuck::must_cast_slice::<_, u64>(&item.output_shape),
                         item.subset.clone(),
                     )
                     .map_py_err::<PyRuntimeError>()?
@@ -410,7 +409,6 @@ impl CodecPipelineImpl {
         } else {
             InputValue::Constant(FillValue::new(input_slice.to_vec()))
         };
-        let input_shape: Vec<u64> = value.shape_zarr()?;
 
         // Adjust the concurrency based on the codec chain and the first chunk description
         let Some((chunk_concurrent_limit, mut codec_options)) =
@@ -424,7 +422,7 @@ impl CodecPipelineImpl {
             let store_chunk = |item: ChunkItem| match &input {
                 InputValue::Array(input) => {
                     let chunk_subset_bytes = input
-                        .extract_array_subset(&item.subset, &input_shape, &self.data_type)
+                        .extract_array_subset(&item.subset, bytemuck::must_cast_slice::<_, u64>(&item.output_shape), &self.data_type)
                         .map_codec_err()?;
                     self.store_chunk_subset_bytes(
                         &item,

src/utils.rs

@@ -1,7 +1,6 @@
 use std::fmt::Display;
 
-use numpy::{PyUntypedArray, PyUntypedArrayMethods};
-use pyo3::{Bound, PyErr, PyResult, PyTypeInfo};
+use pyo3::{PyErr, PyResult, PyTypeInfo};
 use zarrs::array::CodecError;
 
 use crate::ChunkItem;
@@ -38,23 +37,6 @@ impl<T> PyCodecErrExt<T> for Result<T, CodecError> {
     }
 }
 
-pub(crate) trait PyUntypedArrayExt {
-    fn shape_zarr(&self) -> PyResult<Vec<u64>>;
-}
-
-impl PyUntypedArrayExt for Bound<'_, PyUntypedArray> {
-    fn shape_zarr(&self) -> PyResult<Vec<u64>> {
-        Ok(if self.shape().is_empty() {
-            vec![1] // scalar value
-        } else {
-            self.shape()
-                .iter()
-                .map(|&i| u64::try_from(i))
-                .collect::<Result<_, _>>()?
-        })
-    }
-}
-
 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)