fix: handle io array shape dimensionality

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/__init__.py

@@ -1,6 +1,6 @@
 from ._internal import __version__
 from .pipeline import ZarrsCodecPipeline as _ZarrsCodecPipeline
-from .utils import CollapsedDimensionError, DiscontiguousArrayError
+from .utils import UnsupportedVIndexingError, DiscontiguousArrayError
 
 
 # Need to do this redirection so people can access the pipeline as `zarrs.ZarrsCodecPipeline` instead of `zarrs.pipeline.ZarrsCodecPipeline`
@@ -11,6 +11,6 @@ class ZarrsCodecPipeline(_ZarrsCodecPipeline):
 __all__ = [
     "ZarrsCodecPipeline",
     "DiscontiguousArrayError",
-    "CollapsedDimensionError",
+    "UnsupportedVIndexingError",
     "__version__",
 ]

python/zarrs/pipeline.py

@@ -26,9 +26,9 @@
 
 from ._internal import CodecPipelineImpl
 from .utils import (
-    CollapsedDimensionError,
     DiscontiguousArrayError,
     FillValueNoneError,
+    UnsupportedVIndexingError,
     make_chunk_info_for_rust_with_indices,
 )
 
@@ -185,7 +185,7 @@ async def read(
         except (
             UnsupportedMetadataError,
             DiscontiguousArrayError,
-            CollapsedDimensionError,
+            UnsupportedVIndexingError,
             UnsupportedDataTypeError,
             FillValueNoneError,
         ):
@@ -220,7 +220,7 @@ async def write(
         except (
             UnsupportedMetadataError,
             DiscontiguousArrayError,
-            CollapsedDimensionError,
+            UnsupportedVIndexingError,
             UnsupportedDataTypeError,
             FillValueNoneError,
         ):

python/zarrs/utils.py

@@ -30,7 +30,7 @@ class DiscontiguousArrayError(Exception):
     pass
 
 
-class CollapsedDimensionError(Exception):
+class UnsupportedVIndexingError(Exception):
     pass
 
 
@@ -160,7 +160,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_constant = shape == ()
     chunk_info_with_indices: list[ChunkItem] = []
     write_empty_chunks: bool = True
     for (
@@ -182,17 +182,39 @@ def make_chunk_info_for_rust_with_indices(
         shape_chunk_selection = get_shape_for_selector(
             chunk_selection, chunk_spec.shape, pad=True, drop_axes=drop_axes
         )
-        if prod_op(shape_chunk_selection) != prod_op(shape_chunk_selection_slices):
-            raise CollapsedDimensionError(
+        if (chunk_size := prod_op(shape_chunk_selection)) != prod_op(shape_chunk_selection_slices):
+            raise UnsupportedVIndexingError(
                 f"{shape_chunk_selection} != {shape_chunk_selection_slices}"
             )
+        if not is_constant and chunk_size > prod_op(shape):
+            raise IndexError(
+                f"the size of the chunk subset {chunk_size} and input/output subset {prod_op(shape)} are incompatible"
+            )
+        # We need to have io_array_shape and out_selection_expanded with dimensionalities matching that of the underlying array.
+        # So if we detect that a dimension has been dropped (due to a singleton axis) when converting to slices, we update these two values.
+        if not is_constant and len(shape_chunk_selection) != len(shape_chunk_selection_slices):
+            shape_ctr = 0
+            io_array_shape = []
+            out_selection_expanded = []
+            for shape_chunk in shape_chunk_selection_slices:
+                # Append 1/size-1 slice if this dimension has been dropped on the io_array i.e., shape_chunk_selection has been exhausted so there is an extra 1-sized dimension at the end or has a mismatch with the "full" chunk shape `shape_chunk_selection_slices`.
+                if shape_chunk == 1 and (shape_ctr >= len(shape_chunk_selection) or shape_chunk != shape_chunk_selection[shape_ctr]):
+                    io_array_shape += [1]
+                    out_selection_expanded += [slice(0, 1)]
+                else:
+                    io_array_shape += [shape[shape_ctr]]
+                    out_selection_expanded += [out_selection_as_slices[shape_ctr]]
+                    shape_ctr += 1
+        else:
+            io_array_shape = shape
+            out_selection_expanded = out_selection_as_slices
         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_expanded,
+                shape=io_array_shape,
             )
         )
     return RustChunkInfo(chunk_info_with_indices, write_empty_chunks)

src/chunk_item.rs

@@ -32,6 +32,7 @@ pub(crate) struct ChunkItem {
     pub subset: ArraySubset,
     pub shape: Vec<NonZeroU64>,
     pub num_elements: u64,
+    pub array_shape: Vec<NonZeroU64>
 }
 
 #[gen_stub_pymethods]
@@ -65,6 +66,7 @@ impl ChunkItem {
             subset,
             shape: chunk_shape_nonzero_u64,
             num_elements,
+            array_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(&item.array_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(&item.array_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)