clawhdf5/crates/clawhdf5/tests/chunk_index_interop.rs

//! Fixed Array / Extensible Array chunk-index interop with libhdf5 (via h5py).
//!
//! Both indexes place each chunk at a linear index computed from the
//! dataset's *maximum* dimensions, and the Extensible Array additionally
//! moves its unlimited dimension to the slowest-varying position. Getting
//! either wrong reads (or writes) every chunk after the first row in the
//! wrong place, silently, so these tests compare every value.
//!
//! Skipped when python3 with h5py is unavailable, unless
//! `CLAWHDF5_REQUIRE_INTEROP=1`.

use std::process::Command;

use clawhdf5::{File, FileBuilder};

fn python() -> String {
    std::env::var("CLAWHDF5_PYTHON").unwrap_or_else(|_| "python3".to_string())
}

fn interop_required() -> bool {
    std::env::var("CLAWHDF5_REQUIRE_INTEROP").is_ok_and(|v| v == "1")
}

fn python_available() -> bool {
    Command::new(python())
        .args(["-c", "import h5py"])
        .output()
        .map(|o| o.status.success())
        .unwrap_or(false)
}

macro_rules! skip_if_no_python {
    () => {
        if !python_available() {
            assert!(
                !interop_required(),
                "CLAWHDF5_REQUIRE_INTEROP=1 but python3 with h5py is not available"
            );
            eprintln!("SKIP: python3 with h5py not available");
            return;
        }
    };
}

fn run_python(script: &str) -> String {
    let output = Command::new(python())
        .args(["-c", script])
        .output()
        .expect("failed to run python");
    if !output.status.success() {
        panic!(
            "Python script failed:\nSTDOUT: {}\nSTDERR: {}",
            String::from_utf8_lossy(&output.stdout),
            String::from_utf8_lossy(&output.stderr)
        );
    }
    String::from_utf8_lossy(&output.stdout).trim().to_string()
}

/// Row-major `arange` of `shape`, cropped to `crop` (the current extent).
fn arange_cropped(full: &[usize], crop: &[usize]) -> Vec<i32> {
    let n: usize = crop.iter().product();
    let mut out = Vec::with_capacity(n);
    for flat in 0..n {
        let mut rem = flat;
        let mut src = 0usize;
        let mut stride = 1usize;
        let mut coords = vec![0usize; crop.len()];
        for d in (0..crop.len()).rev() {
            coords[d] = rem % crop[d];
            rem /= crop[d];
        }
        for d in (0..full.len()).rev() {
            src += coords[d] * stride;
            stride *= full[d];
        }
        out.push(src as i32);
    }
    out
}

/// One `i4` dataset, filled with `arange` over `full` and then resized to
/// `shape` (equal to `full` unless the case shrinks it).
struct Case {
    name: &'static str,
    full: Vec<usize>,
    shape: Vec<usize>,
    chunks: Vec<usize>,
    maxshape: &'static str,
    extra: &'static str,
    index: &'static str,
}

fn py_tuple(v: &[usize]) -> String {
    let parts: Vec<String> = v.iter().map(|x| x.to_string()).collect();
    format!("({},)", parts.join(","))
}

/// Have h5py (`libver="latest"`, so Fixed/Extensible Array indexes) write
/// every case to one file, then read each back and compare every value.
fn check_h5py_written(cases: &[Case]) {
    let dir = tempfile::tempdir().unwrap();
    let path = dir.path().join("h5py_chunk_index.h5");
    let path_str = path.display().to_string();

    let mut script =
        format!("import h5py, numpy as np\nf = h5py.File(r'{path_str}', 'w', libver='latest')\n");
    for c in cases {
        script += &format!(
            "d = f.create_dataset('{name}', data=np.arange({n}, dtype='i4').reshape({full}), \
             chunks={chunks}, maxshape={maxshape}{extra})\n\
             d.resize({shape})\n",
            name = c.name,
            n = c.full.iter().product::<usize>(),
            full = py_tuple(&c.full),
            chunks = py_tuple(&c.chunks),
            maxshape = c.maxshape,
            extra = c.extra,
            shape = py_tuple(&c.shape),
        );
    }
    script += "f.close()\n";
    run_python(&script);

    let file = File::open(&path).unwrap();
    for c in cases {
        let ds = file.dataset(c.name).unwrap();
        let shape: Vec<usize> = ds.shape().unwrap().iter().map(|&d| d as usize).collect();
        assert_eq!(shape, c.shape, "{}: shape", c.name);
        let got = ds.read_i32().unwrap();
        let want = arange_cropped(&c.full, &c.shape);
        let bad = got.iter().zip(&want).filter(|(a, b)| a != b).count();
        assert_eq!(
            got,
            want,
            "{}: {bad} of {} values differ (index {})",
            c.name,
            want.len(),
            c.index
        );
    }
}

/// h5py-written Extensible Array whose unlimited dimension is not the first,
/// with the current shape smaller than the finite maximum: the library
/// swizzles the unlimited dimension to the slowest position and strides the
/// rest by their maximum chunk counts.
#[test]
fn h5py_extensible_array_partial_extent_reads_correctly() {
    skip_if_no_python!();
    check_h5py_written(&[
        // The `ea_fa_partial.h5` repro from the conformance sweep.
        Case {
            name: "ea_10_none",
            full: vec![4, 6],
            shape: vec![4, 6],
            chunks: vec![2, 3],
            maxshape: "(10, None)",
            extra: "",
            index: "EA, unlimited dim 1",
        },
        Case {
            name: "ea_none_10",
            full: vec![4, 6],
            shape: vec![4, 6],
            chunks: vec![2, 3],
            maxshape: "(None, 10)",
            extra: "",
            index: "EA, unlimited dim 0",
        },
        Case {
            name: "ea_3d_mid",
            full: vec![3, 4, 5],
            shape: vec![3, 4, 5],
            chunks: vec![2, 3, 2],
            maxshape: "(5, None, 7)",
            extra: "",
            index: "EA, unlimited dim 1 of 3",
        },
        Case {
            name: "ea_3d_last_gzip",
            full: vec![3, 4, 5],
            shape: vec![3, 4, 5],
            chunks: vec![2, 3, 2],
            maxshape: "(5, 9, None)",
            extra: ", compression='gzip'",
            index: "EA, unlimited dim 2 of 3, filtered",
        },
        // Many chunks: crosses data blocks, super blocks and paging.
        Case {
            name: "ea_many",
            full: vec![3, 1500],
            shape: vec![3, 1500],
            chunks: vec![1, 1],
            maxshape: "(4, None)",
            extra: "",
            index: "EA, 4500 slots",
        },
        // Shrunk after writing: chunks beyond the extent must be ignored.
        Case {
            name: "ea_shrunk",
            full: vec![8, 9],
            shape: vec![3, 4],
            chunks: vec![2, 3],
            maxshape: "(10, None)",
            extra: "",
            index: "EA, shrunk",
        },
    ]);
}

/// h5py-written Fixed Array with the current shape smaller than a finite
/// maxshape: the index has one slot per chunk of the *maximum* extent.
#[test]
fn h5py_fixed_array_partial_extent_reads_correctly() {
    skip_if_no_python!();
    check_h5py_written(&[
        Case {
            name: "fa_20_10",
            full: vec![4, 6],
            shape: vec![4, 6],
            chunks: vec![2, 3],
            maxshape: "(20, 10)",
            extra: "",
            index: "FA",
        },
        Case {
            name: "fa_3d_gzip",
            full: vec![3, 4, 5],
            shape: vec![3, 4, 5],
            chunks: vec![2, 3, 2],
            maxshape: "(6, 8, 10)",
            extra: ", compression='gzip'",
            index: "FA, filtered",
        },
        // Paged (> 1024 slots) with most of them beyond the extent.
        Case {
            name: "fa_paged",
            full: vec![30, 50],
            shape: vec![30, 50],
            chunks: vec![1, 1],
            maxshape: "(40, 60)",
            extra: "",
            index: "FA, 2400 slots, paged",
        },
        Case {
            name: "fa_shrunk",
            full: vec![8, 9],
            shape: vec![5, 2],
            chunks: vec![2, 3],
            maxshape: "(20, 10)",
            extra: "",
            index: "FA, shrunk",
        },
    ]);
}

// ===========================================================================
// Files we write, read back by libhdf5 (h5py and h5dump) and by us
// ===========================================================================

/// One `i4` dataset we write, filled with `arange` over `shape`.
struct WriteCase {
    name: String,
    shape: Vec<u64>,
    chunks: Vec<u64>,
    maxshape: Option<Vec<u64>>,
    deflate: bool,
}

fn wcase(name: &str, shape: &[u64], chunks: &[u64], maxshape: Option<&[u64]>) -> WriteCase {
    WriteCase {
        name: name.to_string(),
        shape: shape.to_vec(),
        chunks: chunks.to_vec(),
        maxshape: maxshape.map(<[u64]>::to_vec),
        deflate: false,
    }
}

fn h5dump_available() -> bool {
    Command::new("h5dump")
        .arg("--version")
        .output()
        .map(|o| o.status.success())
        .unwrap_or(false)
}

/// Write every case into one file with our writer, then check that our own
/// reader, h5py and h5dump (when installed) all return every value. Only the
/// libhdf5 half is skipped without h5py.
fn check_we_write(cases: &[WriteCase]) {
    let dir = tempfile::tempdir().unwrap();
    let path = dir.path().join("ours_chunk_index.h5");
    let path_str = path.display().to_string();

    let mut b = FileBuilder::new();
    for c in cases {
        let n: u64 = c.shape.iter().product();
        let data: Vec<i32> = (0..n as i32).collect();
        let ds = b.create_dataset(&c.name);
        ds.with_i32_data(&data)
            .with_shape(&c.shape)
            .with_chunks(&c.chunks);
        if let Some(ms) = &c.maxshape {
            ds.with_maxshape(ms);
        }
        if c.deflate {
            ds.with_deflate(4);
        }
    }
    b.write(&path).unwrap();

    // Our reader.
    let file = File::open(&path).unwrap();
    for c in cases {
        let got = file.dataset(&c.name).unwrap().read_i32().unwrap();
        let n: u64 = c.shape.iter().product();
        let bad = got
            .iter()
            .enumerate()
            .filter(|&(i, &v)| v != i as i32)
            .count();
        assert!(
            got.len() == n as usize && bad == 0,
            "{}: our reader: {bad} of {n} values wrong",
            c.name
        );
    }

    // libhdf5 via h5py.
    skip_if_no_python!();
    let mut script =
        format!("import h5py, numpy as np\nbad = []\nf = h5py.File(r'{path_str}', 'r')\n");
    for c in cases {
        let shape: Vec<String> = c.shape.iter().map(u64::to_string).collect();
        let maxshape: Vec<String> = c
            .maxshape
            .as_ref()
            .unwrap_or(&c.shape)
            .iter()
            .map(|&d| {
                if d == u64::MAX {
                    "None".to_string()
                } else {
                    d.to_string()
                }
            })
            .collect();
        script += &format!(
            "d = f['{name}']\n\
             want = np.arange({n}, dtype='i4').reshape(({shape},))\n\
             got = d[()]\n\
             if d.maxshape != ({maxshape},): bad.append(('{name}', 'maxshape', d.maxshape))\n\
             elif not np.array_equal(got, want): \
             bad.append(('{name}', int((got != want).sum()), 'of', got.size))\n",
            name = c.name,
            n = c.shape.iter().product::<u64>(),
            shape = shape.join(","),
            maxshape = maxshape.join(","),
        );
    }
    script += "print(bad if bad else 'OK')\n";
    let out = run_python(&script);
    assert_eq!(out, "OK", "h5py disagrees");

    // libhdf5's own tool, when installed.
    if h5dump_available() {
        let o = Command::new("h5dump").arg(&path).output().unwrap();
        let stderr = String::from_utf8_lossy(&o.stderr);
        assert!(
            o.status.success() && !stderr.to_lowercase().contains("error"),
            "h5dump failed: {stderr}"
        );
    }
}

/// A Fixed Array with more than 1024 elements must be paged, or libhdf5
/// rejects the data block's checksum.
#[test]
fn we_write_paged_fixed_array() {
    let mut cases: Vec<WriteCase> = [1023u64, 1024, 1025, 2048, 5000]
        .iter()
        .map(|&n| wcase(&format!("fa_{n}"), &[n * 4], &[4], None))
        .collect();
    // Filtered elements are wider; a 2-D grid pages the same way.
    let mut filtered = wcase("fa_1500_deflate", &[1500 * 4], &[4], None);
    filtered.deflate = true;
    cases.push(filtered);
    cases.push(wcase("fa_2d_1100", &[110, 40], &[1, 4], None));
    check_we_write(&cases);
}

/// An Extensible Array holds 4 elements in its index block and 240 in the
/// data blocks the index block addresses; everything after that lives under
/// super blocks, and from ~131K elements on in paged data blocks. Chunks past
/// index 243 used to be written but never indexed (read back as fill by us
/// and by libhdf5).
#[test]
fn we_write_extensible_array_past_index_block() {
    let unl: &[u64] = &[u64::MAX];
    let mut cases: Vec<WriteCase> = [1u64, 4, 5, 243, 244, 245, 300, 1000, 5000]
        .iter()
        .map(|&n| wcase(&format!("ea_{n}"), &[n * 4], &[4], Some(unl)))
        .collect();
    let mut filtered = wcase("ea_300_deflate", &[300 * 4], &[4], Some(unl));
    filtered.deflate = true;
    cases.push(filtered);
    // Several super blocks and paged data blocks (level 13, the first with
    // data blocks over 1024 elements, starts at element 4 + 131056).
    cases.push(wcase("ea_140000", &[140_000], &[1], Some(unl)));
    check_we_write(&cases);
}