Code shared from the Rust Playground

chunkshape.rs

pub type Ix = usize;

pub fn estimate_chunkshape(shape: &[Ix], itemsize: usize, base_kb: usize) -> Vec<Ix> {
    if shape.is_empty() {
        return vec![1];
    }
    let size: Ix = shape.iter().product();
    let size_mb = ((size * itemsize) as f64) / ((1 << 20) as f64);
    let basesize = base_kb << 10;
    let pow2 = size_mb.max(1.).min((1 << 23) as _).log10().trunc();
    let factor = 2usize.pow(pow2 as _);
    let chunksize = basesize * factor;
    let mut chunkshape = shape.to_vec();
    let rowsize: Ix = shape.iter().skip(1).product();
    chunkshape[0] = (chunksize / rowsize).max(1).min(shape[0]);
    let chunkbytes = itemsize * chunkshape[0] * rowsize;
    let max_chunk_size = 1 << 31; // BLOSC_MAX_BUFFERSIZE
    if chunkbytes > max_chunk_size {
        chunkshape[0] = (((chunkshape[0] as f64) * (max_chunk_size as f64) / (chunkbytes as f64))
            .trunc() as Ix)
            .max(1)
            .min(shape[0]);
    }
    chunkshape
}

/*
TESTS:

    (1024 * 1024, 1, 128, (128 * 1024,)),
    ((1024 * 1024,), 1, 128, (128 * 1024,)),
    (1024 * 1024 - 1, 1, 128, (128 * 1024,)),
    (1024 * 1024 + 1, 1, 128, (128 * 1024,)),
    (1024 * 1024 * 10, 1, 128, (256 * 1024,)),
    (1024 * 1024 * 10 - 1, 1, 128, (128 * 1024,)),
    (1024 * 1024 * 100, 1, 128, (512 * 1024,)),
    (1024 * 1024 * 100 - 1, 1, 128, (256 * 1024,)),
    (1024 * 1024 * 100, 4, 128, (128 * 1024,)),
    ((1024 * 1024 * 100, 2), 4, 128, (64 * 1024, 2)),
    (1024 * 1024, 1, 32, (32 * 1024,)),
    ((1024 * 1024 * 100, 2), 4, 32, (16 * 1024, 2)),
    ((1024 * 1024, 0, 2), 4, 32, None)
*/

fn main() {
    println!("{:?}", estimate_chunkshape(&[1024 * 1024 * 100, 2], 4, 32));
}