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
+135
-93
@@ -2,12 +2,12 @@
|
||||
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use pyo3::exceptions::{PyIOError, PyKeyError};
|
||||
use pyo3::prelude::*;
|
||||
use pyo3::types::PyList;
|
||||
|
||||
use crate::attrs::PyAttrs;
|
||||
use crate::dataset::PyDataset;
|
||||
use crate::{DatasetSpec, OwnedAttrValue, apply_dataset_spec, extract_numpy_data, to_py_err};
|
||||
use crate::{DatasetSpec, OwnedAttrValue, apply_dataset_spec, extract_numpy_data, node, to_py_err};
|
||||
|
||||
/// Shared state for a group being written.
|
||||
pub(crate) struct WriteGroupState {
|
||||
@@ -18,15 +18,10 @@ pub(crate) struct WriteGroupState {
|
||||
|
||||
/// An HDF5 group.
|
||||
///
|
||||
/// In read mode, provides `__getitem__` navigation and child listing.
|
||||
/// In write mode, supports `create_dataset` and `create_group` and
|
||||
/// attribute setting.
|
||||
///
|
||||
/// ```python
|
||||
/// grp = f['group_name']
|
||||
/// grp.keys()
|
||||
/// ds = grp['dataset']
|
||||
/// ```
|
||||
/// In read mode it behaves like an h5py group: `grp['name']`,
|
||||
/// `grp['sub/path']` and `grp['/absolute/path']`, `keys()`, `values()`,
|
||||
/// `items()`, iteration, `len()`, `in`, `get()`, `name` and `attrs`.
|
||||
/// In write mode, supports `create_dataset` and attribute setting.
|
||||
#[pyclass(name = "Group")]
|
||||
pub struct PyGroup {
|
||||
inner: GroupInner,
|
||||
@@ -52,46 +47,90 @@ impl PyGroup {
|
||||
inner: GroupInner::Write(state),
|
||||
}
|
||||
}
|
||||
|
||||
fn read_parts(&self, what: &str) -> PyResult<(&Arc<clawhdf5_rs::File>, &str)> {
|
||||
match &self.inner {
|
||||
GroupInner::Read { file, path } => Ok((file, path)),
|
||||
GroupInner::Write(_) => Err(PyIOError::new_err(format!(
|
||||
"cannot {what} a group opened for writing"
|
||||
))),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Read-mode operations shared by `Group` and `File` (a file is its root
|
||||
// group, as in h5py).
|
||||
|
||||
/// `group[key]`.
|
||||
pub(crate) fn get_item(
|
||||
py: Python<'_>,
|
||||
file: &Arc<clawhdf5_rs::File>,
|
||||
path: &str,
|
||||
key: &str,
|
||||
) -> PyResult<Py<PyAny>> {
|
||||
node::open(py, file, node::join(path, key))
|
||||
}
|
||||
|
||||
/// Names of the group's datasets and subgroups, sorted (h5py's order).
|
||||
pub(crate) fn member_names(file: &clawhdf5_rs::File, path: &str) -> PyResult<Vec<String>> {
|
||||
let group = if path.is_empty() {
|
||||
file.root()
|
||||
} else {
|
||||
file.group(path).map_err(to_py_err)?
|
||||
};
|
||||
let mut names = group.datasets().map_err(to_py_err)?;
|
||||
names.extend(group.groups().map_err(to_py_err)?);
|
||||
names.sort_by(|a, b| a.as_bytes().cmp(b.as_bytes()));
|
||||
names.dedup();
|
||||
Ok(names)
|
||||
}
|
||||
|
||||
pub(crate) fn contains(file: &clawhdf5_rs::File, path: &str, key: &str) -> bool {
|
||||
node::exists(file, &node::join(path, key))
|
||||
}
|
||||
|
||||
pub(crate) fn values(
|
||||
py: Python<'_>,
|
||||
file: &Arc<clawhdf5_rs::File>,
|
||||
path: &str,
|
||||
) -> PyResult<Vec<Py<PyAny>>> {
|
||||
member_names(file, path)?
|
||||
.iter()
|
||||
.map(|n| get_item(py, file, path, n))
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub(crate) fn items(
|
||||
py: Python<'_>,
|
||||
file: &Arc<clawhdf5_rs::File>,
|
||||
path: &str,
|
||||
) -> PyResult<Vec<(String, Py<PyAny>)>> {
|
||||
member_names(file, path)?
|
||||
.into_iter()
|
||||
.map(|n| {
|
||||
let v = get_item(py, file, path, &n)?;
|
||||
Ok((n, v))
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl PyGroup {
|
||||
/// Get a child object (dataset or subgroup) by name or path.
|
||||
fn __getitem__(&self, py: Python<'_>, key: &str) -> PyResult<Py<PyAny>> {
|
||||
match &self.inner {
|
||||
GroupInner::Read { file, path } => {
|
||||
let full_path = if path.is_empty() {
|
||||
key.to_string()
|
||||
} else {
|
||||
format!("{path}/{key}")
|
||||
};
|
||||
// Try dataset first
|
||||
match file.dataset(&full_path) {
|
||||
Ok(_) => {
|
||||
let ds = PyDataset::new(Arc::clone(file), full_path)?;
|
||||
Ok(ds.into_pyobject(py)?.into_any().unbind())
|
||||
}
|
||||
Err(clawhdf5_rs::Error::NotADataset(_)) => {
|
||||
let grp = PyGroup::from_read(Arc::clone(file), full_path);
|
||||
Ok(grp.into_pyobject(py)?.into_any().unbind())
|
||||
}
|
||||
Err(e) => {
|
||||
// Could be a group without a DataLayout message
|
||||
match file.group(&full_path) {
|
||||
Ok(_) => {
|
||||
let grp = PyGroup::from_read(Arc::clone(file), full_path);
|
||||
Ok(grp.into_pyobject(py)?.into_any().unbind())
|
||||
}
|
||||
Err(_) => Err(PyErr::new::<pyo3::exceptions::PyKeyError, _>(format!(
|
||||
"{key}: {e}"
|
||||
))),
|
||||
}
|
||||
}
|
||||
}
|
||||
let (file, path) = self.read_parts("read children from")?;
|
||||
get_item(py, file, path, key)
|
||||
}
|
||||
|
||||
/// `group.get(key, default=None)`.
|
||||
#[pyo3(signature = (key, default=None))]
|
||||
fn get(&self, py: Python<'_>, key: &str, default: Option<Py<PyAny>>) -> PyResult<Py<PyAny>> {
|
||||
let (file, path) = self.read_parts("read children from")?;
|
||||
match get_item(py, file, path, key) {
|
||||
Err(e) if e.is_instance_of::<PyKeyError>(py) => {
|
||||
Ok(default.unwrap_or_else(|| py.None()))
|
||||
}
|
||||
GroupInner::Write(_) => Err(PyErr::new::<pyo3::exceptions::PyIOError, _>(
|
||||
"cannot read children from a group opened for writing",
|
||||
)),
|
||||
other => other,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -99,16 +138,7 @@ impl PyGroup {
|
||||
fn keys(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
|
||||
match &self.inner {
|
||||
GroupInner::Read { file, path } => {
|
||||
let group = if path.is_empty() {
|
||||
file.root()
|
||||
} else {
|
||||
file.group(path).map_err(to_py_err)?
|
||||
};
|
||||
let mut names = group.datasets().map_err(to_py_err)?;
|
||||
let groups = group.groups().map_err(to_py_err)?;
|
||||
names.extend(groups);
|
||||
names.sort();
|
||||
let list = PyList::new(py, &names)?;
|
||||
let list = PyList::new(py, member_names(file, path)?)?;
|
||||
Ok(list.into_any().unbind())
|
||||
}
|
||||
GroupInner::Write(state) => {
|
||||
@@ -120,6 +150,38 @@ impl PyGroup {
|
||||
}
|
||||
}
|
||||
|
||||
fn values(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
|
||||
let (file, path) = self.read_parts("read children from")?;
|
||||
Ok(PyList::new(py, values(py, file, path)?)?
|
||||
.into_any()
|
||||
.unbind())
|
||||
}
|
||||
|
||||
fn items(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
|
||||
let (file, path) = self.read_parts("read children from")?;
|
||||
Ok(PyList::new(py, items(py, file, path)?)?.into_any().unbind())
|
||||
}
|
||||
|
||||
fn __iter__(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
|
||||
self.keys(py)?.call_method0(py, "__iter__")
|
||||
}
|
||||
|
||||
fn __len__(&self) -> PyResult<usize> {
|
||||
match &self.inner {
|
||||
GroupInner::Read { file, path } => Ok(member_names(file, path)?.len()),
|
||||
GroupInner::Write(state) => Ok(state.lock().unwrap().datasets.len()),
|
||||
}
|
||||
}
|
||||
|
||||
/// The group's full name, e.g. `/sensors`.
|
||||
#[getter]
|
||||
fn name(&self) -> String {
|
||||
match &self.inner {
|
||||
GroupInner::Read { path, .. } => node::name(path),
|
||||
GroupInner::Write(state) => node::name(&state.lock().unwrap().name),
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a dataset inside this group (write mode only).
|
||||
///
|
||||
/// Parameters:
|
||||
@@ -161,7 +223,7 @@ impl PyGroup {
|
||||
state.lock().unwrap().datasets.push(spec);
|
||||
Ok(())
|
||||
}
|
||||
GroupInner::Read { .. } => Err(PyErr::new::<pyo3::exceptions::PyIOError, _>(
|
||||
GroupInner::Read { .. } => Err(PyIOError::new_err(
|
||||
"cannot create datasets on a read-only group",
|
||||
)),
|
||||
}
|
||||
@@ -171,15 +233,7 @@ impl PyGroup {
|
||||
#[getter]
|
||||
fn attrs(&self) -> PyResult<PyAttrs> {
|
||||
match &self.inner {
|
||||
GroupInner::Read { file, path } => {
|
||||
let group = if path.is_empty() {
|
||||
file.root()
|
||||
} else {
|
||||
file.group(path).map_err(to_py_err)?
|
||||
};
|
||||
let map = group.attrs().map_err(to_py_err)?;
|
||||
Ok(PyAttrs::from_read(map))
|
||||
}
|
||||
GroupInner::Read { file, path } => PyAttrs::read(Arc::clone(file), path),
|
||||
GroupInner::Write(state) => {
|
||||
let store = Arc::clone(&state.lock().unwrap().attrs);
|
||||
Ok(PyAttrs::from_write(store))
|
||||
@@ -189,12 +243,9 @@ impl PyGroup {
|
||||
|
||||
fn __repr__(&self) -> String {
|
||||
match &self.inner {
|
||||
GroupInner::Read { path, .. } => {
|
||||
if path.is_empty() {
|
||||
"<HDF5 Group \"/\" (root)>".to_string()
|
||||
} else {
|
||||
format!("<HDF5 Group \"/{path}\">")
|
||||
}
|
||||
GroupInner::Read { file, path } => {
|
||||
let n = member_names(file, path).map_or(0, |m| m.len());
|
||||
format!("<HDF5 group \"{}\" ({n} members)>", node::name(path))
|
||||
}
|
||||
GroupInner::Write(state) => {
|
||||
let name = &state.lock().unwrap().name;
|
||||
@@ -205,14 +256,7 @@ impl PyGroup {
|
||||
|
||||
fn __contains__(&self, key: &str) -> PyResult<bool> {
|
||||
match &self.inner {
|
||||
GroupInner::Read { file, path } => {
|
||||
let full_path = if path.is_empty() {
|
||||
key.to_string()
|
||||
} else {
|
||||
format!("{path}/{key}")
|
||||
};
|
||||
Ok(file.dataset(&full_path).is_ok() || file.group(&full_path).is_ok())
|
||||
}
|
||||
GroupInner::Read { file, path } => Ok(contains(file, path, key)),
|
||||
GroupInner::Write(state) => {
|
||||
let guard = state.lock().unwrap();
|
||||
Ok(guard.datasets.iter().any(|d| d.name == key))
|
||||
@@ -244,26 +288,24 @@ mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn read_group_construction() {
|
||||
fn member_names_are_sorted() {
|
||||
let mut b = clawhdf5_rs::FileBuilder::new();
|
||||
let mut g = b.create_group("grp");
|
||||
b.create_dataset("zeta").with_f64_data(&[1.0]);
|
||||
b.create_dataset("alpha").with_f64_data(&[1.0]);
|
||||
let mut g = b.create_group("mid");
|
||||
g.create_dataset("x").with_f64_data(&[1.0]);
|
||||
let finished = g.finish();
|
||||
b.add_group(finished);
|
||||
let bytes = b.finish().unwrap();
|
||||
let file = Arc::new(clawhdf5_rs::File::from_bytes(bytes).unwrap());
|
||||
let _grp = PyGroup::from_read(file, "grp".into());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn write_group_state() {
|
||||
let state = WriteGroupState {
|
||||
name: "test".into(),
|
||||
datasets: vec![],
|
||||
attrs: Arc::new(Mutex::new(vec![])),
|
||||
};
|
||||
let arc = Arc::new(Mutex::new(state));
|
||||
let _grp = PyGroup::from_write(arc);
|
||||
let file = clawhdf5_rs::File::from_bytes(bytes).unwrap();
|
||||
assert_eq!(
|
||||
member_names(&file, "").unwrap(),
|
||||
vec!["alpha", "mid", "zeta"]
|
||||
);
|
||||
assert_eq!(member_names(&file, "mid").unwrap(), vec!["x"]);
|
||||
assert!(contains(&file, "", "mid/x"));
|
||||
assert!(contains(&file, "mid", "/alpha"));
|
||||
assert!(!contains(&file, "", "nope"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
Reference in New Issue
Block a user