fix(py): index lists of padded compounds no longer return uninitialised padding

np.concatenate copies structured dtypes field by field into np.empty, so
the padding of ds[[0, 3, 6]] held process memory. The runs' bytes are
joined in Rust, whole elements at a time, before anything becomes numpy:
the padding is the file's bytes (h5py's) and the result is still a view
of the Rust buffer. The h5py comparisons now compare every byte of
structured values; the new test failed on the padding before.

Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>
This commit is contained in:
osobh
2026-09-26 08:54:04 -05:00
co-authored by Claude Opus 5.5
parent 24412a0e59
commit 8c51b05b9c
4 changed files with 116 additions and 52 deletions
+46 -1
View File
@@ -56,6 +56,12 @@ impl Plan {
.collect()
}
/// The shape of the result before the integer-indexed axes are dropped
/// (they have length 1 here): the shape of the joined reads.
pub fn read_shape(&self) -> Vec<usize> {
self.axes.iter().map(|a| a.len() as usize).collect()
}
/// Whether the selection is empty.
pub fn is_empty(&self) -> bool {
self.axes.iter().any(|a| a.len() == 0)
@@ -64,7 +70,7 @@ impl Plan {
/// The hyperslab reads that make up this selection, each with the shape
/// of its block (index axes kept at length 1). More than one only when an
/// axis is indexed by a list: one read per run of consecutive indices,
/// concatenated along `list_axis` afterwards.
/// joined along `list_axis` afterwards (`join_along`).
pub fn reads(&self, dims: &[u64]) -> (Vec<(Selection, Vec<usize>)>, Option<usize>) {
let list_axis = self.axes.iter().position(|a| matches!(a, Axis::List(_)));
let runs: Vec<(u64, u64)> = match list_axis.map(|i| &self.axes[i]) {
@@ -119,6 +125,33 @@ fn consecutive_runs(idx: &[u64]) -> Vec<(u64, u64)> {
runs
}
/// Join row-major blocks of `elem_size`-byte elements whose shapes differ
/// only along `axis` into one buffer, in order along that axis. Whole
/// elements are copied, so compound padding keeps the bytes that were read.
pub(crate) fn join_along(
blocks: &[(Vec<u8>, Vec<usize>)],
axis: usize,
elem_size: usize,
) -> Vec<u8> {
let Some((_, first)) = blocks.first() else {
return Vec::new();
};
let outer: usize = first[..axis].iter().product();
let inner: usize = first[axis + 1..].iter().product::<usize>() * elem_size;
let total: usize = blocks.iter().map(|(_, s)| s[axis]).sum();
let mut out = vec![0u8; outer * total * inner];
let mut at = 0;
for (bytes, shape) in blocks {
let len = shape[axis] * inner;
for o in 0..outer {
let dst = (o * total) * inner + at;
out[dst..dst + len].copy_from_slice(&bytes[o * len..(o + 1) * len]);
}
at += len;
}
out
}
/// Parse `key` for a dataset of shape `dims`.
pub(crate) fn parse(key: &Bound<'_, PyAny>, dims: &[u64]) -> PyResult<Plan> {
let items: Vec<Bound<'_, PyAny>> = match key.cast::<PyTuple>() {
@@ -311,6 +344,18 @@ mod tests {
assert_eq!(consecutive_runs(&[]), vec![]);
}
#[test]
fn blocks_join_along_the_list_axis() {
// Two 2x1 and 2x2 blocks of 1-byte elements, joined along axis 1.
let a = (vec![1, 2], vec![2, 1]);
let b = (vec![3, 4, 5, 6], vec![2, 2]);
assert_eq!(join_along(&[a, b], 1, 1), vec![1, 3, 4, 2, 5, 6]);
// Along axis 0 it is concatenation; 2-byte elements stay whole.
let a = (vec![1, 2, 3, 4], vec![1, 2]);
let b = (vec![5, 6, 7, 8], vec![1, 2]);
assert_eq!(join_along(&[a, b], 0, 2), vec![1, 2, 3, 4, 5, 6, 7, 8]);
}
#[test]
fn full_selection_reads_everything() {
let plan = Plan {