feat(py): h5py-style reads of only the selected elements, GIL released
ds[key] read the whole dataset and sliced it in numpy, and knew six dtypes. Keys (ints, positive-step slices, Ellipsis, one increasing index list, compound field names) now map onto hyperslab selections, and the facade's read_selection bytes become the numpy buffer without a copy (PyArray::from_vec viewed as the dtype). dtype mapping follows h5py for all integer/IEEE float widths and byte orders, bool, enum, complex, fixed and variable-length strings, vlen sequences, opaque, array types and (nested, padded) compounds; anything it cannot describe exactly is a TypeError. Attributes return what h5py returns; groups and files gain the rest of the h5py mapping interface. Reads run under py.detach. tests/test_read_vs_h5py.py compares >500 reads with h5py 3.16 on an h5py-written file, checks errors match, that a damaged chunk outside the selection is never touched, and 8 threads reading at once. Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>
This commit is contained in:
osobh
@@ -0,0 +1,366 @@
|
||||
//! h5py-style indexing (`ds[1, 2:10:3, ...]`) mapped onto hyperslab
|
||||
//! selections, so only the selected elements are read.
|
||||
//!
|
||||
//! The rules and error messages follow h5py's `selections.py`: integers
|
||||
//! (negative from the end) drop their axis, slices must have a positive
|
||||
//! step, one `Ellipsis` fills the unmentioned axes, a single increasing list
|
||||
//! of integers may index one axis, and strings name compound fields.
|
||||
//! Everything else (`None`/`np.newaxis`, boolean masks, several index lists)
|
||||
//! is refused with the error h5py gives.
|
||||
|
||||
use clawhdf5_format::selection::Selection;
|
||||
use pyo3::exceptions::{PyIndexError, PyTypeError, PyValueError};
|
||||
use pyo3::prelude::*;
|
||||
use pyo3::types::{PyEllipsis, PySlice, PyString, PyTuple};
|
||||
|
||||
/// The selection along one axis.
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub(crate) enum Axis {
|
||||
/// A single index: the axis is dropped from the result.
|
||||
Index(u64),
|
||||
/// `start, start + step, ...`, `count` of them.
|
||||
Slice { start: u64, step: u64, count: u64 },
|
||||
/// Increasing, distinct indices.
|
||||
List(Vec<u64>),
|
||||
}
|
||||
|
||||
impl Axis {
|
||||
fn len(&self) -> u64 {
|
||||
match self {
|
||||
Axis::Index(_) => 1,
|
||||
Axis::Slice { count, .. } => *count,
|
||||
Axis::List(v) => v.len() as u64,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// A parsed index expression.
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub(crate) struct Plan {
|
||||
/// One entry per dataset axis.
|
||||
pub axes: Vec<Axis>,
|
||||
/// Compound field names to keep (empty: all).
|
||||
pub fields: Vec<String>,
|
||||
/// For a scalar dataset: `ds[()]` gives a scalar, `ds[...]` a 0-d array.
|
||||
/// For other datasets: every axis was an integer, so h5py gives a scalar.
|
||||
pub scalar: bool,
|
||||
}
|
||||
|
||||
impl Plan {
|
||||
/// The shape of the result.
|
||||
pub fn out_shape(&self) -> Vec<usize> {
|
||||
self.axes
|
||||
.iter()
|
||||
.filter(|a| !matches!(a, Axis::Index(_)))
|
||||
.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)
|
||||
}
|
||||
|
||||
/// 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.
|
||||
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]) {
|
||||
Some(Axis::List(idx)) => consecutive_runs(idx),
|
||||
_ => vec![(0, 0)],
|
||||
};
|
||||
let mut out = Vec::with_capacity(runs.len());
|
||||
for (run_start, run_len) in runs {
|
||||
let mut start = Vec::with_capacity(dims.len());
|
||||
let mut stride = Vec::with_capacity(dims.len());
|
||||
let mut count = Vec::with_capacity(dims.len());
|
||||
for axis in &self.axes {
|
||||
let (s, st, c) = match axis {
|
||||
Axis::Index(i) => (*i, 1, 1),
|
||||
Axis::Slice { start, step, count } => (*start, *step, *count),
|
||||
Axis::List(_) => (run_start, 1, run_len),
|
||||
};
|
||||
start.push(s);
|
||||
// A stride only matters between blocks; keep it >= 1.
|
||||
stride.push(if c <= 1 { 1 } else { st });
|
||||
count.push(c);
|
||||
}
|
||||
let block_shape: Vec<usize> = count.iter().map(|&c| c as usize).collect();
|
||||
let whole = start.iter().all(|&s| s == 0)
|
||||
&& stride.iter().all(|&s| s == 1)
|
||||
&& count.as_slice() == dims;
|
||||
let sel = if whole {
|
||||
Selection::All
|
||||
} else {
|
||||
let block = vec![1; dims.len()];
|
||||
Selection::Hyperslab {
|
||||
start,
|
||||
stride,
|
||||
count,
|
||||
block,
|
||||
}
|
||||
};
|
||||
out.push((sel, block_shape));
|
||||
}
|
||||
(out, list_axis)
|
||||
}
|
||||
}
|
||||
|
||||
fn consecutive_runs(idx: &[u64]) -> Vec<(u64, u64)> {
|
||||
let mut runs: Vec<(u64, u64)> = Vec::new();
|
||||
for &i in idx {
|
||||
match runs.last_mut() {
|
||||
Some((s, n)) if *s + *n == i => *n += 1,
|
||||
_ => runs.push((i, 1)),
|
||||
}
|
||||
}
|
||||
runs
|
||||
}
|
||||
|
||||
/// 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>() {
|
||||
Ok(t) => t.iter().collect(),
|
||||
Err(_) => vec![key.clone()],
|
||||
};
|
||||
let mut fields = Vec::new();
|
||||
let mut args = Vec::new();
|
||||
for item in items {
|
||||
if let Ok(s) = item.cast::<PyString>() {
|
||||
fields.push(s.to_str()?.to_owned());
|
||||
} else {
|
||||
args.push(item);
|
||||
}
|
||||
}
|
||||
|
||||
if args.iter().any(|a| a.is_none()) {
|
||||
return Err(PyTypeError::new_err(
|
||||
"Indexing with None (or np.newaxis) is not supported",
|
||||
));
|
||||
}
|
||||
|
||||
let rank = dims.len();
|
||||
if rank == 0 {
|
||||
return match args.as_slice() {
|
||||
[] => Ok(Plan {
|
||||
axes: vec![],
|
||||
fields,
|
||||
scalar: true,
|
||||
}),
|
||||
[a] if a.is_instance_of::<PyEllipsis>() => Ok(Plan {
|
||||
axes: vec![],
|
||||
fields,
|
||||
scalar: false,
|
||||
}),
|
||||
_ => Err(PyValueError::new_err(
|
||||
"Illegal slicing argument for scalar dataspace",
|
||||
)),
|
||||
};
|
||||
}
|
||||
|
||||
// Expand the ellipsis (at most one) to full slices.
|
||||
let n_ellipsis = args
|
||||
.iter()
|
||||
.filter(|a| a.is_instance_of::<PyEllipsis>())
|
||||
.count();
|
||||
if n_ellipsis > 1 {
|
||||
return Err(PyValueError::new_err("Only one ellipsis may be used."));
|
||||
}
|
||||
let explicit = args.len() - n_ellipsis;
|
||||
if explicit > rank {
|
||||
return Err(PyValueError::new_err(format!(
|
||||
"{explicit} indexing arguments for {rank} dimensions"
|
||||
)));
|
||||
}
|
||||
let py = key.py();
|
||||
let mut expanded: Vec<Option<Bound<'_, PyAny>>> = Vec::with_capacity(rank);
|
||||
for a in args {
|
||||
if a.is_instance_of::<PyEllipsis>() {
|
||||
for _ in 0..(rank - explicit) {
|
||||
expanded.push(None);
|
||||
}
|
||||
} else {
|
||||
expanded.push(Some(a));
|
||||
}
|
||||
}
|
||||
while expanded.len() < rank {
|
||||
expanded.push(None);
|
||||
}
|
||||
|
||||
let mut axes = Vec::with_capacity(rank);
|
||||
for (arg, &n) in expanded.iter().zip(dims) {
|
||||
axes.push(match arg {
|
||||
None => Axis::Slice {
|
||||
start: 0,
|
||||
step: 1,
|
||||
count: n,
|
||||
},
|
||||
Some(a) => parse_axis(py, a, n)?,
|
||||
});
|
||||
}
|
||||
if axes.iter().filter(|a| matches!(a, Axis::List(_))).count() > 1 {
|
||||
return Err(PyTypeError::new_err(
|
||||
"Only one indexing vector or array is currently allowed for fancy indexing",
|
||||
));
|
||||
}
|
||||
let scalar = axes.iter().all(|a| matches!(a, Axis::Index(_)));
|
||||
Ok(Plan {
|
||||
axes,
|
||||
fields,
|
||||
scalar,
|
||||
})
|
||||
}
|
||||
|
||||
fn parse_axis(py: Python<'_>, a: &Bound<'_, PyAny>, n: u64) -> PyResult<Axis> {
|
||||
if a.is_none() {
|
||||
return Err(PyTypeError::new_err(
|
||||
"Indexing with None (or np.newaxis) is not supported",
|
||||
));
|
||||
}
|
||||
if let Ok(s) = a.cast::<PySlice>() {
|
||||
let n_isize = isize::try_from(n)
|
||||
.map_err(|_| PyValueError::new_err("dimension too large to slice"))?;
|
||||
let ind = s.indices(n_isize)?;
|
||||
if ind.step < 1 {
|
||||
return Err(PyValueError::new_err(format!(
|
||||
"Step must be >= 1 (got {})",
|
||||
ind.step
|
||||
)));
|
||||
}
|
||||
// `slicelength` is the number of elements selected, >= 0.
|
||||
let count = ind.slicelength as u64;
|
||||
let start = if count == 0 { 0 } else { ind.start as u64 };
|
||||
return Ok(Axis::Slice {
|
||||
start,
|
||||
step: ind.step as u64,
|
||||
count,
|
||||
});
|
||||
}
|
||||
let np = py.import("numpy")?;
|
||||
let is_bool =
|
||||
a.is_instance_of::<pyo3::types::PyBool>() || a.is_instance(&np.getattr("bool_")?)?;
|
||||
let is_array_like = a.is_instance(&np.getattr("ndarray")?)?
|
||||
|| a.is_instance_of::<pyo3::types::PyList>()
|
||||
|| a.is_instance_of::<PyTuple>();
|
||||
if !is_bool && !is_array_like && a.hasattr("__index__")? {
|
||||
let i: i128 = a.call_method0("__index__")?.extract()?;
|
||||
return Ok(Axis::Index(normalize(i, n)?));
|
||||
}
|
||||
if is_array_like {
|
||||
let arr = np.call_method1("asarray", (a,))?;
|
||||
let kind: String = arr.getattr("dtype")?.getattr("kind")?.extract()?;
|
||||
if kind == "b" {
|
||||
return Err(PyTypeError::new_err(
|
||||
"Boolean mask indexing is not supported by clawhdf5",
|
||||
));
|
||||
}
|
||||
let ndim: usize = arr.getattr("ndim")?.extract()?;
|
||||
let size: usize = arr.getattr("size")?.extract()?;
|
||||
if size > 0 && kind != "i" && kind != "u" {
|
||||
return Err(PyTypeError::new_err(
|
||||
"Indexing arrays must have integer dtypes",
|
||||
));
|
||||
}
|
||||
if ndim > 1 {
|
||||
return Err(PyTypeError::new_err(
|
||||
"Only 1-D integer lists or arrays can be used for fancy indexing",
|
||||
));
|
||||
}
|
||||
let vals: Vec<i128> = arr.call_method0("tolist")?.extract()?;
|
||||
let mut idx = Vec::with_capacity(vals.len());
|
||||
for v in vals {
|
||||
idx.push(normalize(v, n)?);
|
||||
}
|
||||
if idx.windows(2).any(|w| w[0] >= w[1]) {
|
||||
return Err(PyTypeError::new_err(
|
||||
"Indexing elements must be in increasing order",
|
||||
));
|
||||
}
|
||||
return Ok(Axis::List(idx));
|
||||
}
|
||||
Err(PyTypeError::new_err(format!(
|
||||
"Illegal index type for clawhdf5 datasets: {}",
|
||||
a.get_type().name()?
|
||||
)))
|
||||
}
|
||||
|
||||
fn normalize(i: i128, n: u64) -> PyResult<u64> {
|
||||
let n_i = i128::from(n);
|
||||
let j = if i < 0 { i + n_i } else { i };
|
||||
if j < 0 || j >= n_i {
|
||||
let hi = n_i - 1;
|
||||
return Err(PyIndexError::new_err(format!(
|
||||
"Index ({i}) out of range for (0-{hi})"
|
||||
)));
|
||||
}
|
||||
Ok(j as u64)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn runs_group_consecutive_indices() {
|
||||
assert_eq!(
|
||||
consecutive_runs(&[1, 2, 3, 7, 9, 10]),
|
||||
vec![(1, 3), (7, 1), (9, 2)]
|
||||
);
|
||||
assert_eq!(consecutive_runs(&[]), vec![]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn full_selection_reads_everything() {
|
||||
let plan = Plan {
|
||||
axes: vec![
|
||||
Axis::Slice {
|
||||
start: 0,
|
||||
step: 1,
|
||||
count: 4,
|
||||
},
|
||||
Axis::Slice {
|
||||
start: 0,
|
||||
step: 1,
|
||||
count: 3,
|
||||
},
|
||||
],
|
||||
fields: vec![],
|
||||
scalar: false,
|
||||
};
|
||||
let (reads, list) = plan.reads(&[4, 3]);
|
||||
assert_eq!(list, None);
|
||||
assert_eq!(reads, vec![(Selection::All, vec![4, 3])]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn index_and_step_map_to_a_hyperslab() {
|
||||
let plan = Plan {
|
||||
axes: vec![
|
||||
Axis::Index(2),
|
||||
Axis::Slice {
|
||||
start: 1,
|
||||
step: 3,
|
||||
count: 2,
|
||||
},
|
||||
],
|
||||
fields: vec![],
|
||||
scalar: false,
|
||||
};
|
||||
let (reads, _) = plan.reads(&[4, 8]);
|
||||
assert_eq!(
|
||||
reads,
|
||||
vec![(
|
||||
Selection::Hyperslab {
|
||||
start: vec![2, 1],
|
||||
stride: vec![1, 3],
|
||||
count: vec![1, 2],
|
||||
block: vec![1, 1],
|
||||
},
|
||||
vec![1, 2]
|
||||
)]
|
||||
);
|
||||
assert_eq!(plan.out_shape(), vec![2]);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user