partial shard reads appear slower than expected in some cases

[ilan-gold]

# /// script
# requires-python = ">=3.12"
# dependencies = [
#   "zarrs",
#   "numpy",
# ]
# ///
#
# This script automatically imports the development branch of zarr to check for issues

import zarr
import numpy as np
import time


import platform
import subprocess


def clear_cache():
    if platform.system() == "Darwin":
        subprocess.call(["sync", "&&", "sudo", "purge"])
    elif platform.system() == "Linux":
        subprocess.call(["sudo", "sh", "-c", "sync; echo 3 > /proc/sys/vm/drop_caches"])
    else:
        raise Exception("Unsupported platform")

zarr.config.set({"codec_pipeline.path": "zarrs.ZarrsCodecPipeline"})

z = zarr.create_array(
    "foo.zarr",
    shape=(8192, 4, 128, 128),
    shards=(4096, 4, 128, 128),
    chunks=(1, 1, 128, 128),
    dtype=np.float64,
    overwrite=True,
)

# Originally this was in the reproducer but it's clear that the behavior is somewhat dependent on compression
# z[...] = np.random.randn(8192, 4, 128, 128)
z[...] = np.ones((8192, 4, 128, 128))
clear_cache()
t = time.time()
z[...]
print("full read took: ", time.time() - t)
clear_cache()
t = time.time()
z[:4095, ...]
print("partial shard read took: ", time.time() - t)

On my mac, the partial read takes as long as the full read:

full read took:  2.4683969020843506
partial shard read took:  2.291616201400757

While I love sequential io just as much as the next guy, this feels wrong. I have a feeling this has to do with the way that concurrency is calculated in the presence of very few shards. As an instructive counter-example, increasing that 4095 to 4096 yields

full read took:  2.334256887435913
partial shard read took:  0.7899982929229736

Which feels much more reasonable.

(I am aware these chunks are comically small, so just setting that aside for now)

UPDATE: I've changed the np.random to be np.ones because the above repro was only working on mac's which have compressed file systems. With np.ones (i.e., non-random data), this issue should now be reproducible on linux machines

[ilan-gold]

I wrote another script (to be run after the above since the above creates the foo.zarr array) that tries to do something in rust at a higher level than what exists in the codec pipeline here:

#!/usr/bin/env -S cargo +nightly -Zscript

---cargo
[dependencies]
zarrs = { version = "0.23.3", features = ["zstd"] }
---

use std::sync::Arc;
use std::time::SystemTime;

use zarrs::array::ArraySubset;
use zarrs::filesystem::FilesystemStore;

use std::process::Command;
use std::io;


fn clear_page_cache() -> io::Result<()> {
    #[cfg(target_os = "linux")]
    {
        Command::new("sync").status()?;
        Command::new("sudo sh")
            .arg("-c")
            .arg("echo 3 > /proc/sys/vm/drop_caches")
            .status()?;
    }

    #[cfg(target_os = "macos")]
    {
        Command::new("sync").status()?;
        Command::new("sudo purge").status()?;
    }

    Ok(())
}

pub fn main() -> Result<(), Box<dyn std::error::Error>> {
    // Create the store
    let store = Arc::new(FilesystemStore::new("/Users/ilangold/Projects/Theis/zarrs-python/foo.zarr")?);
    let array = zarrs::array::Array::open(store, "/")?;
    // Full chunks
    let fast_subset = ArraySubset::new_with_ranges(&[(0..8192), (0..4), (0..128), (0..128)]);
    // Just under a full chunk
    let slow_subset = ArraySubset::new_with_ranges(&[(0..4095), (0..4), (0..128), (0..128)]);
    
    // Time full
    clear_page_cache();
    let start = SystemTime::now();
    array.retrieve_array_subset_ndarray::<f64>(&fast_subset)?;
    let end = SystemTime::now();
    let duration = end.duration_since(start).unwrap();
    println!("full took {} seconds", duration.as_millis());

    // Time partial
    clear_page_cache();
    let start = SystemTime::now();
    array.retrieve_array_subset_ndarray::<f64>(&slow_subset)?;
    let end = SystemTime::now();
    let duration = end.duration_since(start).unwrap();
    println!("partial took {} seconds", duration.as_millis());
    clear_page_cache();

    Ok(())
}

gives

full took 9387 seconds
patial took 3324 seconds

which makes sense to me.

[ilan-gold]

And an arguably-more-faithful example that uses the _into capabilities which are used inside zarrs-python:

#!/usr/bin/env -S cargo +nightly -Zscript

---cargo
[dependencies]
zarrs = { version = "0.23.3", features = ["zstd"] }
unsafe_cell_slice = { version = "0.2.0" }
---

use std::sync::Arc;
use std::time::SystemTime;

use unsafe_cell_slice::UnsafeCellSlice;
use zarrs::array::{Array, ArrayBytesDecodeIntoTarget, ArrayBytesFixedDisjointView, ArraySubset};
use zarrs::filesystem::FilesystemStore;
use zarrs::storage::ReadableStorageTraits;

use std::process::Command;
use std::io;


fn clear_page_cache() -> io::Result<()> {
    #[cfg(target_os = "linux")]
    {
        Command::new("sudo sync").status()?;
        Command::new("sh")
            .arg("-c")
            .arg("echo 3 > /proc/sys/vm/drop_caches")
            .status()?;
    }

    #[cfg(target_os = "macos")]
    {
        Command::new("sync").status()?;
        Command::new("sudo purge").status()?;
    }

    Ok(())
}

fn time_subsetting<T>(subset: ArraySubset, array: &Array<T>) -> Result<(), Box<dyn std::error::Error>> where T: ReadableStorageTraits + 'static {
    // Output buffer, 8 times the number of elements because data is f64.
    let mut buf = vec![0u8; subset.num_elements_usize() * 8];
    let output = UnsafeCellSlice::new(buf.as_mut_slice());
    let mut output_view = unsafe {
        ArrayBytesFixedDisjointView::new(
            output,
            array.data_type()
                .fixed_size()
                .unwrap(),
            subset.shape(),
            subset.clone(),
        )?
    };
    let target = ArrayBytesDecodeIntoTarget::Fixed(&mut output_view);
    clear_page_cache();
    let start = SystemTime::now();
    array.retrieve_array_subset_into(&subset, target)?;
    let end = SystemTime::now();
    let duration = end.duration_since(start).unwrap();
    println!("{} took {} seconds", subset, duration.as_millis());
    Ok(())
}

pub fn main() -> Result<(), Box<dyn std::error::Error>> {
    // Create the store
    let store = Arc::new(FilesystemStore::new("/Users/ilangold/Projects/Theis/zarrs-python/foo.zarr")?);
    let array = Array::open(store, "/")?;
    // Full read, should in theory be slower than the partial read
    let full_subset = ArraySubset::new_with_ranges(&[(0..8192), (0..4), (0..128), (0..128)]);
    // Partial read, should in theory be faster than the full read
    let partial_subset = ArraySubset::new_with_ranges(&[(0..4095), (0..4), (0..128), (0..128)]);


    // ----------
    // Time Full
    // ----------

    time_subsetting(full_subset, &array);
    

    // ----------
    // Time Partial
    // ----------
    time_subsetting(partial_subset, &array);
    
    Ok(())
}

gives

[0..8192, 0..4, 0..128, 0..128] took 6333 seconds
[0..4095, 0..4, 0..128, 0..128] took 2470 seconds

Exactly as expected.

[ilan-gold]

And using 4096 i.e., a full shard read

[0..4096, 0..4, 0..128, 0..128] took 1969 seconds
[0..4095, 0..4, 0..128, 0..128] took 1500 seconds

which again, makes sense - the partial shard is comparable to the full shard.

TL;DR, reads (just) under a full shard size in zarrs-python are very slow in a way that I cannot reproduce in pure rust.