(fix): integer indexing

ilan-gold · Aug 23, 2024 · cd8d1a8 · cd8d1a8
1 parent 0bbf681
commit cd8d1a8
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 5 deletions.
diff --git a/src/array.rs b/src/array.rs
@@ -187,13 +187,11 @@ impl ZarrsPythonArray {
             if self.is_selection_numpy_array(chunk_coords_and_selections, 1) {
                 let selections_extracted = self.extract_selection_to_vec_indices(chunk_coords_and_selections, 1)?;
                 let borrowed_selections = &selections_extracted;
-                println!("hereeees");
                 {
                     let output =
                         UnsafeCellSlice::new_from_vec_with_spare_capacity(&mut output);
                     let retrieve_chunk = |chunk: NdArrayChunk| {
-                        println!("{:?} {:?}", cartesian_product(chunk.selection), out_shape_extracted);
-                        let indices: Vec<u64> = cartesian_product(chunk.selection).iter().map(|x| x.iter().enumerate().fold(0, |acc, (ind, x)| {acc + (*x as u64) * out_shape_extracted[1..].iter().product::<u64>()})).collect();
+                        let indices: Vec<u64> = cartesian_product(chunk.selection).iter().map(|x| x.iter().enumerate().fold(0, |acc, (ind, x)| {acc + (*x as u64) * if (ind + 1 == chunk.selection.len()) { 1 } else { self.arr.chunk_shape(&chunk.index).unwrap()[(ind + 1)..].iter().map(|x| x.get() as u64).product::<u64>()}})).collect();
                         let chunk_subset_bytes = self.arr.retrieve_chunk(&chunk.index).map_err(|x| PyErr::new::<PyTypeError, _>(x.to_string()))?;
                         update_bytes_flen_with_indexer(
                             unsafe { output.get() },
@@ -213,6 +211,8 @@ impl ZarrsPythonArray {
                         retrieve_chunk
                     )?;
                 }
+                unsafe { output.set_len(size_output) };
+                return Ok(ManagerCtx::new(PyZarrArr{ shape: out_shape_extracted, arr: output, dtype: chunk_representation.data_type().clone() }));
             }
             let selections_extracted = self.extract_selection_to_array_subset(chunk_coords_and_selections, 1)?;
             let out_selections_extracted = &self.extract_selection_to_array_subset(chunk_coords_and_selections, 2)?;

diff --git a/src/utils.rs b/src/utils.rs
@@ -16,13 +16,12 @@ pub fn update_bytes_flen_with_indexer(
         unsafe { subset.contiguous_linearised_indices_unchecked(output_shape) };
     // TODO: Par iteration?
     let mut indexer_index = 0;
-    println!("{:?}, {:?}", subset_bytes.len(), indexer);
     for (array_subset_element_index, _num_elements) in contiguous_indices.iter() {
         let mut output_offset = usize::try_from(array_subset_element_index).unwrap() * data_type_size;
         for _num_elem in 0.._num_elements {
             let decoded_offset = (indexer[indexer_index] as usize) * data_type_size;
             debug_assert!((output_offset + data_type_size) <= output_bytes.len());
-            debug_assert!((decoded_offset + data_type_size) <= subset_bytes.len());
+            debug_assert!((decoded_offset + data_type_size) <= subset_bytes.len(), "Failed subset check: decoded_offset: {:?}, data_type_size: {:?}, subset_bytes.len(): {:?}", decoded_offset, data_type_size, subset_bytes.len());
             output_bytes[output_offset..output_offset + data_type_size]
                 .copy_from_slice(&subset_bytes[decoded_offset..decoded_offset + data_type_size]);
             indexer_index += 1;