clawhdf5/crates/clawhdf5/tests/v4_chunk_index_selection.rs

//! Partial hyperslab reads of every v4 (`libver='latest'`) chunk index type,
//! compared element for element with h5py.
//!
//! `Dataset::read_selection` takes two routes: `partial_read` materialises
//! the selection's bounding box when it covers at most half the dataset, and
//! `data_read::read_raw_data_selection` handles the rest. The second route
//! once passed the layout's full chunk dimensions (which carry the element
//! size as an extra, last dimension) to the implicit-index chunk generator,
//! which then indexed past the dataset's rank and panicked. So every case
//! below reads both a small window and one covering most of the dataset.
//!
//! Each case asserts which chunk index the file actually uses (parsed from
//! the layout message), so a change in how h5py lays the file out can't turn
//! this into a test of the wrong index.
//!
//! Skipped when python3 with h5py is unavailable, unless
//! `CLAWHDF5_REQUIRE_INTEROP=1`.

use std::process::Command;

use clawhdf5::File;
use clawhdf5_format::data_layout::DataLayout;
use clawhdf5_format::message_type::MessageType;
use clawhdf5_format::object_header::ObjectHeader;
use clawhdf5_format::selection::Selection;
use clawhdf5_format::superblock::Superblock;

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)
}

fn run_python(script: &str) -> String {
    let output = Command::new(python())
        .args(["-c", script])
        .output()
        .expect("failed to run python");
    assert!(
        output.status.success(),
        "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()
}

/// The v4 chunk index type recorded in `name`'s layout message
/// (1 single chunk, 2 implicit, 3 fixed array, 4 extensible array, 5 B-tree v2).
fn chunk_index_type(path: &std::path::Path, name: &str) -> u8 {
    let data = std::fs::read(path).unwrap();
    let sb = Superblock::parse(&data, 0).unwrap();
    let addr = clawhdf5_format::group_v2::resolve_path_any(&data, &sb, name).unwrap();
    let hdr = ObjectHeader::parse(&data, addr as usize, sb.offset_size, sb.length_size).unwrap();
    let msg = hdr
        .messages
        .iter()
        .find(|m| m.msg_type == MessageType::DataLayout)
        .expect("layout message");
    match DataLayout::parse(&msg.data, sb.offset_size, sb.length_size).unwrap() {
        DataLayout::Chunked {
            version: 4,
            chunk_index_type: Some(t),
            ..
        } => t,
        other => panic!("{name}: expected a v4 chunked layout, got {other:?}"),
    }
}

struct Case {
    name: &'static str,
    /// Python keyword arguments to `create_dataset` besides `data`.
    kwargs: &'static str,
    index_type: u8,
}

const SHAPE: [u64; 2] = [37, 23];

const CASES: &[Case] = &[
    Case {
        name: "implicit",
        // Early allocation, no filters, fixed maximum: the implicit index.
        kwargs: "chunks=(5, 4), dcpl=early()",
        index_type: 2,
    },
    Case {
        name: "fixed_array",
        kwargs: "chunks=(5, 4), compression='gzip'",
        index_type: 3,
    },
    Case {
        name: "extensible_array",
        kwargs: "chunks=(5, 4), maxshape=(None, 23), compression='gzip'",
        index_type: 4,
    },
    Case {
        name: "btree2",
        kwargs: "chunks=(5, 4), maxshape=(None, None), compression='gzip'",
        index_type: 5,
    },
    Case {
        name: "single_chunk",
        kwargs: "chunks=(37, 23), compression='gzip'",
        index_type: 1,
    },
    Case {
        name: "single_chunk_unfiltered",
        kwargs: "chunks=(37, 23)",
        index_type: 1,
    },
];

/// `(start, stride, count, block)` per dimension; the first few stay below
/// half the dataset (bounding-box path), the rest exceed it (full path).
fn selections() -> Vec<([u64; 2], [u64; 2], [u64; 2], [u64; 2])> {
    vec![
        ([0, 0], [1, 1], [3, 23], [1, 1]),  // ds[0:3]
        ([7, 3], [1, 1], [9, 6], [1, 1]),   // interior window across chunks
        ([36, 22], [1, 1], [1, 1], [1, 1]), // last element (edge chunk)
        ([2, 1], [3, 4], [4, 3], [1, 1]),   // strided, small
        ([0, 0], [1, 1], [30, 23], [1, 1]), // most rows
        ([1, 0], [2, 1], [18, 23], [1, 1]), // every other row, spanning all
        ([0, 2], [1, 1], [37, 20], [1, 1]), // columns 2..22 of every row
        ([3, 1], [5, 3], [7, 7], [2, 2]),   // strided blocks over everything
    ]
}

fn py_slice(start: u64, stride: u64, count: u64, block: u64) -> String {
    // Each case is expressible as a numpy index when block == 1; with a block
    // the selected indices are listed explicitly.
    let idx: Vec<String> = (0..count)
        .flat_map(|c| (0..block).map(move |b| start + c * stride + b))
        .map(|i| i.to_string())
        .collect();
    format!("[{}]", idx.join(","))
}

#[test]
fn partial_hyperslabs_of_every_v4_chunk_index_match_h5py() {
    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;
    }
    let dir = tempfile::tempdir().unwrap();
    let path = dir.path().join("v4_index_selection.h5");
    let path_str = path.display().to_string();

    // h5py writes the file, then reads every selection back and prints the
    // values, one line per (case, selection).
    let mut script = format!(
        "import h5py, numpy as np\n\
         def early():\n\
         \x20   p = h5py.h5p.create(h5py.h5p.DATASET_CREATE)\n\
         \x20   p.set_alloc_time(h5py.h5d.ALLOC_TIME_EARLY)\n\
         \x20   return p\n\
         data = (np.arange({n}, dtype='<i4') * 7 - 1000).reshape({r}, {c})\n\
         with h5py.File(r'{path_str}', 'w', libver='latest') as f:\n",
        n = SHAPE[0] * SHAPE[1],
        r = SHAPE[0],
        c = SHAPE[1],
    );
    for case in CASES {
        script += &format!(
            "    f.create_dataset('{}', data=data, {})\n",
            case.name, case.kwargs
        );
    }
    script += &format!("with h5py.File(r'{path_str}', 'r') as f:\n");
    for case in CASES {
        for (start, stride, count, block) in selections() {
            let rows = py_slice(start[0], stride[0], count[0], block[0]);
            let cols = py_slice(start[1], stride[1], count[1], block[1]);
            script += &format!(
                "    print(' '.join(map(str, f['{}'][{rows}][:, {cols}].ravel())))\n",
                case.name
            );
        }
    }
    let out = run_python(&script);
    let mut expected = out.lines();

    let file = File::open(&path).unwrap();
    for case in CASES {
        assert_eq!(
            chunk_index_type(&path, case.name),
            case.index_type,
            "{}: h5py did not produce the intended chunk index",
            case.name
        );
        let ds = file.dataset(case.name).unwrap();
        assert_eq!(ds.shape().unwrap(), SHAPE);
        for (start, stride, count, block) in selections() {
            let sel = Selection::Hyperslab {
                start: start.to_vec(),
                stride: stride.to_vec(),
                count: count.to_vec(),
                block: block.to_vec(),
            };
            let want: Vec<i32> = expected
                .next()
                .expect("h5py printed too few lines")
                .split_whitespace()
                .map(|v| v.parse().unwrap())
                .collect();
            let raw = ds
                .read_selection(&sel)
                .unwrap_or_else(|e| panic!("{}: read_selection {sel:?} failed: {e}", case.name));
            let got: Vec<i32> = raw
                .chunks_exact(4)
                .map(|b| i32::from_le_bytes(b.try_into().unwrap()))
                .collect();
            assert_eq!(got, want, "{}: selection {sel:?}", case.name);
            assert_eq!(
                ds.read_i32_selection(&sel).unwrap(),
                want,
                "{}: read_i32_selection {sel:?}",
                case.name
            );
        }
    }
    assert!(expected.next().is_none(), "h5py printed extra lines");
}