fix(format): selections of v4 implicit-index chunked data no longer panic
read_raw_data_selection's chunked fallback (taken when partial_read declines, e.g. a bounding box over half the dataset) handed the layout's chunk dimensions, element-size dimension included, to generate_implicit_chunks, which indexed past the dataset rank. It then decoded the whole dataset regardless, so the enumeration is gone: the arm decodes and extracts for every chunk index. The new test reads small and large hyperslabs of all five v4 indexes written by h5py and compares with h5py's values; it panicked before. Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>
This commit is contained in:
osobh
@@ -510,6 +510,10 @@ fn collect_chunk_info_inner(
|
||||
///
|
||||
/// Chunks are stored contiguously starting at `base_address`. No stored index;
|
||||
/// addresses are computed from the chunk position.
|
||||
///
|
||||
/// `chunk_dimensions` are the spatial chunk dimensions, one per entry of
|
||||
/// `dataset_dims` — not the layout message's list, which carries the element
|
||||
/// size as an extra last dimension.
|
||||
pub fn generate_implicit_chunks(
|
||||
base_address: u64,
|
||||
dataset_dims: &[u64],
|
||||
|
||||
@@ -356,85 +356,17 @@ pub fn read_raw_data_selection(
|
||||
}
|
||||
DataLayout::Chunked {
|
||||
chunk_dimensions,
|
||||
btree_address,
|
||||
version,
|
||||
chunk_index_type,
|
||||
..
|
||||
} => {
|
||||
// `partial_read` declined (a bounding box covering most of the
|
||||
// dataset, or a selection it doesn't box), so decode every chunk
|
||||
// and pick the selection out, whatever the chunk index. This arm
|
||||
// used to enumerate the chunks first — passing the layout's
|
||||
// chunk dimensions, element-size dimension included, to the
|
||||
// implicit-index generator, which then indexed past the rank and
|
||||
// panicked — only to decode the full dataset anyway.
|
||||
crate::chunked_read::chunk_geometry(chunk_dimensions, *version, dataspace, elem_size)?;
|
||||
// For chunked data, only read chunks that intersect the selection
|
||||
let chunk_dims: Vec<u64> = chunk_dimensions.iter().map(|&d| d as u64).collect();
|
||||
let rank = dims.len();
|
||||
|
||||
// Collect chunk info from B-tree
|
||||
let chunks = if *version == 4 {
|
||||
match chunk_index_type {
|
||||
Some(2) => {
|
||||
// Implicit index
|
||||
crate::chunked_read::generate_implicit_chunks(
|
||||
btree_address.unwrap_or(0),
|
||||
dims,
|
||||
chunk_dimensions,
|
||||
elem_size as u32,
|
||||
)
|
||||
}
|
||||
_ => {
|
||||
if let Some(_addr) = btree_address {
|
||||
// Use extensible array or fixed array
|
||||
// Fall back to full read for complex v4 index types
|
||||
let full_data = read_raw_data_full(
|
||||
file_data,
|
||||
layout,
|
||||
dataspace,
|
||||
datatype,
|
||||
pipeline,
|
||||
offset_size,
|
||||
length_size,
|
||||
)?;
|
||||
return extract_selection_from_buffer(
|
||||
&full_data, dims, elem_size, selection,
|
||||
);
|
||||
} else {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// v3: B-tree v1
|
||||
if let Some(addr) = btree_address {
|
||||
crate::chunked_read::collect_chunk_info_checked(
|
||||
file_data,
|
||||
*addr,
|
||||
chunk_dimensions,
|
||||
offset_size,
|
||||
length_size,
|
||||
)?
|
||||
} else {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
};
|
||||
|
||||
// Filter chunks to only those that intersect the selection
|
||||
let intersecting: Vec<_> = chunks
|
||||
.iter()
|
||||
.filter(|ci| {
|
||||
let offsets: Vec<u64> = ci.offsets.iter().take(rank).copied().collect();
|
||||
selection.intersects_chunk(&offsets, &chunk_dims[..rank])
|
||||
})
|
||||
.collect();
|
||||
|
||||
if intersecting.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
// Decompress only the intersecting chunks
|
||||
let _chunk_total_bytes: usize =
|
||||
chunk_dims.iter().map(|&d| d as usize).product::<usize>() * elem_size;
|
||||
let _element_size_u32 = elem_size as u32;
|
||||
|
||||
// First, assemble only the intersecting chunks into a partial buffer,
|
||||
// then extract the selection. For simplicity, we assemble into a full
|
||||
// dataset buffer and extract (same as contiguous path).
|
||||
let full_data = read_raw_data_full(
|
||||
file_data,
|
||||
layout,
|
||||
|
||||
@@ -0,0 +1,236 @@
|
||||
//! 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");
|
||||
}
|
||||
Reference in New Issue
Block a user