fix 3D array chunking

Author: kaizhang

anndata/src/container/collection.rs

@@ -5,7 +5,7 @@ use crate::{
     data::*,
 };
 
-use anyhow::{Result, bail, ensure};
+use anyhow::{Context, Result, bail, ensure};
 use itertools::Itertools;
 use log::warn;
 use parking_lot::{Mutex, MutexGuard};
@@ -351,7 +351,10 @@ impl<B: Backend> InnerAxisArrays<B> {
         if let Some(elem) = self.get(key) {
             elem.clear()?;
         }
-        let elem = ArrayElem::try_from(ArrayChunk::write_by_chunk(data, &self.container, key)?)?;
+        let elem = ArrayChunk::write_by_chunk(data, &self.container, key).with_context(|| {
+            format!("failed to write data to AxisArrays with key: '{}'", key)
+        })?;
+        let elem = ArrayElem::try_from(elem)?;
 
         let shape = { elem.inner().shape().clone() };
         match self.axis {

anndata/src/data/array/chunks.rs

@@ -1,6 +1,6 @@
-use crate::{ArrayElem, Selectable};
 use crate::backend::{AttributeOp, Backend, BackendData, DataContainer, GroupOp, ScalarType};
 use crate::data::{ArrayData, array::DynArray, array::utils::ExtendableDataset};
+use crate::{ArrayElem, Selectable};
 
 use super::{CsrNonCanonical, DynCscMatrix, DynCsrMatrix, DynCsrNonCanonical};
 use anyhow::{Context, Result, bail};
@@ -329,12 +329,13 @@ impl<D: RemoveAxis, T: BackendData> ArrayChunk for Array<T, D> {
         G: GroupOp<B>,
     {
         let mut iter = iter.peekable();
-        let chunk_size = if let Some(n) = D::NDIM {
-            vec![1000; n].into()
-        } else {
-            let n = iter.peek().unwrap().ndim();
-            vec![1000; n].into()
-        };
+        let chunk_size = iter
+            .peek()
+            .unwrap()
+            .shape()
+            .iter()
+            .map(|&x| x.min(1000))
+            .collect();
         let mut data: ExtendableDataset<B, T> =
             ExtendableDataset::with_capacity(location, name, chunk_size)?;
 

anndata/src/data/array/utils.rs

@@ -25,7 +25,7 @@ impl<B: Backend, T: BackendData> ExtendableDataset<B, T> {
     where
         G: GroupOp<B>,
     {
-        let block_size = vec![1000; capacity.ndim()].into();
+        let block_size = alloc_block_size_with_shape(&capacity, 16384);
         let dataset = group.new_empty_dataset::<T>(
             name,
             &capacity,
@@ -584,4 +584,29 @@ pub(crate) fn array_major_minor_index_default<T: Default + Clone>(
 ) -> Array2<T>
 {
     array_major_minor_index(major_idx, minor_idx, data, &T::default())
+}
+
+pub(crate) fn alloc_block_size_with_shape(shape: &Shape, total: usize) -> Shape {
+    let mut block_size = vec![0; shape.ndim()];
+    let mut n = shape.ndim();
+
+    let mut bs = get_block_size(n, total);
+    let mut visit_order: Vec<_> = (0..n).collect();
+    visit_order.sort_by_key(|&i| shape[i]);
+    for i in visit_order {
+        let s = shape[i];
+        if s < bs {
+            block_size[i] = s;
+            n -= 1;
+            bs = get_block_size(n, total / s);
+        } else {
+            block_size[i] = bs;
+        }
+    }
+
+    block_size.into()
+}
+
+fn get_block_size(n: usize, total: usize) -> usize {
+    (total as f64).powf(1.0 / n as f64).ceil() as usize
 }