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:
@@ -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 {
|
||||
|
||||
Reference in New Issue
Block a user