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
+103
-45
@@ -1,24 +1,31 @@
|
||||
//! PyAttrs — dict-like access to HDF5 attributes.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use clawhdf5_format::attribute::AttributeMessage;
|
||||
use pyo3::exceptions::{PyKeyError, PyTypeError, PyValueError};
|
||||
use pyo3::prelude::*;
|
||||
use pyo3::types::PyList;
|
||||
use pyo3::types::{PyList, PyTuple};
|
||||
|
||||
use crate::{OwnedAttrValue, attr_value_to_py, py_to_attr_value};
|
||||
use crate::convert::{Converter, Elements, resolve_vl};
|
||||
use crate::{OwnedAttrValue, PyEmpty, attr_value_to_py, node, py_to_attr_value};
|
||||
|
||||
/// Backing storage for attributes.
|
||||
enum AttrsInner {
|
||||
/// Read-only attributes from an existing HDF5 object.
|
||||
Read(HashMap<String, clawhdf5_rs::AttrValue>),
|
||||
/// Attributes of an object in a file opened for reading, sorted by name.
|
||||
Read {
|
||||
file: Arc<clawhdf5_rs::File>,
|
||||
attrs: Vec<AttributeMessage>,
|
||||
},
|
||||
/// Writable attribute list shared with a parent (PyFile or PyGroup).
|
||||
Write(Arc<Mutex<Vec<(String, OwnedAttrValue)>>>),
|
||||
}
|
||||
|
||||
/// Dict-like access to HDF5 attributes.
|
||||
///
|
||||
/// In read mode, provides immutable access to attribute key/value pairs.
|
||||
/// In read mode, values are what h5py returns: numpy scalars for scalar
|
||||
/// attributes, numpy arrays otherwise, `str` for variable-length strings,
|
||||
/// `numpy.bytes_` for fixed-length ones, and `Empty` for a null dataspace.
|
||||
/// In write mode, attributes set here are accumulated and written when
|
||||
/// the parent file is closed.
|
||||
#[pyclass(name = "Attrs")]
|
||||
@@ -27,11 +34,12 @@ pub struct PyAttrs {
|
||||
}
|
||||
|
||||
impl PyAttrs {
|
||||
/// Create a read-only attrs from an existing attribute map.
|
||||
pub(crate) fn from_read(map: HashMap<String, clawhdf5_rs::AttrValue>) -> Self {
|
||||
Self {
|
||||
inner: AttrsInner::Read(map),
|
||||
}
|
||||
/// The attributes of the object at `path` in a file opened for reading.
|
||||
pub(crate) fn read(file: Arc<clawhdf5_rs::File>, path: &str) -> PyResult<Self> {
|
||||
let attrs = node::attributes(&file, path)?;
|
||||
Ok(Self {
|
||||
inner: AttrsInner::Read { file, attrs },
|
||||
})
|
||||
}
|
||||
|
||||
/// Create a writable attrs that shares storage with a parent object.
|
||||
@@ -46,11 +54,11 @@ impl PyAttrs {
|
||||
impl PyAttrs {
|
||||
fn __getitem__(&self, py: Python<'_>, key: &str) -> PyResult<Py<PyAny>> {
|
||||
match &self.inner {
|
||||
AttrsInner::Read(map) => match map.get(key) {
|
||||
Some(val) => Ok(attr_value_to_py(py, val)),
|
||||
None => Err(PyErr::new::<pyo3::exceptions::PyKeyError, _>(
|
||||
key.to_string(),
|
||||
)),
|
||||
AttrsInner::Read { file, attrs } => match attrs.iter().find(|a| a.name == key) {
|
||||
Some(attr) => Ok(attr_to_py(py, file, attr)?.unbind()),
|
||||
None => Err(PyKeyError::new_err(format!(
|
||||
"Can't open attribute (can't locate attribute: '{key}')"
|
||||
))),
|
||||
},
|
||||
AttrsInner::Write(store) => {
|
||||
let guard = store.lock().unwrap();
|
||||
@@ -60,16 +68,14 @@ impl PyAttrs {
|
||||
return Ok(attr_value_to_py(py, &attr_val));
|
||||
}
|
||||
}
|
||||
Err(PyErr::new::<pyo3::exceptions::PyKeyError, _>(
|
||||
key.to_string(),
|
||||
))
|
||||
Err(PyKeyError::new_err(key.to_string()))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn __setitem__(&self, key: &str, value: &Bound<'_, PyAny>) -> PyResult<()> {
|
||||
match &self.inner {
|
||||
AttrsInner::Read(_) => Err(PyErr::new::<pyo3::exceptions::PyIOError, _>(
|
||||
AttrsInner::Read { .. } => Err(PyErr::new::<pyo3::exceptions::PyIOError, _>(
|
||||
"cannot set attributes on a read-only file",
|
||||
)),
|
||||
AttrsInner::Write(store) => {
|
||||
@@ -88,14 +94,14 @@ impl PyAttrs {
|
||||
|
||||
fn __len__(&self) -> usize {
|
||||
match &self.inner {
|
||||
AttrsInner::Read(map) => map.len(),
|
||||
AttrsInner::Read { attrs, .. } => attrs.len(),
|
||||
AttrsInner::Write(store) => store.lock().unwrap().len(),
|
||||
}
|
||||
}
|
||||
|
||||
fn __contains__(&self, key: &str) -> bool {
|
||||
match &self.inner {
|
||||
AttrsInner::Read(map) => map.contains_key(key),
|
||||
AttrsInner::Read { attrs, .. } => attrs.iter().any(|a| a.name == key),
|
||||
AttrsInner::Write(store) => store.lock().unwrap().iter().any(|(k, _)| k == key),
|
||||
}
|
||||
}
|
||||
@@ -111,10 +117,20 @@ impl PyAttrs {
|
||||
format!("<HDF5 Attrs ({n} members)>")
|
||||
}
|
||||
|
||||
/// The value of `key`, or `default` if there is no such attribute.
|
||||
#[pyo3(signature = (key, default=None))]
|
||||
fn get(&self, py: Python<'_>, key: &str, default: Option<Py<PyAny>>) -> PyResult<Py<PyAny>> {
|
||||
if self.__contains__(key) {
|
||||
self.__getitem__(py, key)
|
||||
} else {
|
||||
Ok(default.unwrap_or_else(|| py.None()))
|
||||
}
|
||||
}
|
||||
|
||||
/// Return attribute names as a list.
|
||||
fn keys(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
|
||||
let names: Vec<String> = match &self.inner {
|
||||
AttrsInner::Read(map) => map.keys().cloned().collect(),
|
||||
AttrsInner::Read { attrs, .. } => attrs.iter().map(|a| a.name.clone()).collect(),
|
||||
AttrsInner::Write(store) => store
|
||||
.lock()
|
||||
.unwrap()
|
||||
@@ -129,7 +145,10 @@ impl PyAttrs {
|
||||
/// Return attribute values as a list.
|
||||
fn values(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
|
||||
let vals: Vec<Py<PyAny>> = match &self.inner {
|
||||
AttrsInner::Read(map) => map.values().map(|v| attr_value_to_py(py, v)).collect(),
|
||||
AttrsInner::Read { file, attrs } => attrs
|
||||
.iter()
|
||||
.map(|a| attr_to_py(py, file, a).map(Bound::unbind))
|
||||
.collect::<PyResult<_>>()?,
|
||||
AttrsInner::Write(store) => store
|
||||
.lock()
|
||||
.unwrap()
|
||||
@@ -147,10 +166,10 @@ impl PyAttrs {
|
||||
/// Return attribute (key, value) pairs as a list of tuples.
|
||||
fn items(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
|
||||
let pairs: Vec<(String, Py<PyAny>)> = match &self.inner {
|
||||
AttrsInner::Read(map) => map
|
||||
AttrsInner::Read { file, attrs } => attrs
|
||||
.iter()
|
||||
.map(|(k, v)| (k.clone(), attr_value_to_py(py, v)))
|
||||
.collect(),
|
||||
.map(|a| Ok((a.name.clone(), attr_to_py(py, file, a)?.unbind())))
|
||||
.collect::<PyResult<_>>()?,
|
||||
AttrsInner::Write(store) => store
|
||||
.lock()
|
||||
.unwrap()
|
||||
@@ -166,28 +185,67 @@ impl PyAttrs {
|
||||
}
|
||||
}
|
||||
|
||||
/// An attribute's value as h5py returns it.
|
||||
fn attr_to_py<'py>(
|
||||
py: Python<'py>,
|
||||
file: &clawhdf5_rs::File,
|
||||
attr: &AttributeMessage,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let sb = file.superblock();
|
||||
let conv = Converter::new(py, &attr.datatype, sb.offset_size)
|
||||
.map_err(|e| prefix_err(py, &attr.name, e))?;
|
||||
if node::is_null(&attr.dataspace) {
|
||||
return Ok(PyEmpty::new(conv.dtype).into_pyobject(py)?.into_any());
|
||||
}
|
||||
let shape: Vec<usize> = attr
|
||||
.dataspace
|
||||
.dimensions
|
||||
.iter()
|
||||
.map(|&d| d as usize)
|
||||
.collect();
|
||||
let n: usize = shape.iter().product();
|
||||
let data = if conv.is_vl() {
|
||||
let want = n * conv.elem_size;
|
||||
if attr.raw_data.len() < want {
|
||||
return Err(PyValueError::new_err(format!(
|
||||
"attribute {}: {} bytes of variable-length references, expected {want}",
|
||||
attr.name,
|
||||
attr.raw_data.len(),
|
||||
)));
|
||||
}
|
||||
let raw = &attr.raw_data[..want];
|
||||
let file_data = file.as_bytes();
|
||||
let (osz, lsz, unit) = (sb.offset_size, sb.length_size, conv.vl_unit);
|
||||
Elements::Vl(
|
||||
py.detach(|| resolve_vl(file_data, raw, n, osz, lsz, unit))
|
||||
.map_err(|e| PyValueError::new_err(format!("attribute {}: {e}", attr.name)))?,
|
||||
)
|
||||
} else {
|
||||
Elements::Bytes(attr.raw_data.clone())
|
||||
};
|
||||
let arr = conv
|
||||
.to_array(py, data, &shape, true)
|
||||
.map_err(|e| prefix_err(py, &attr.name, e))?;
|
||||
if shape.is_empty() {
|
||||
// A scalar dataspace: h5py returns the element itself.
|
||||
return arr.get_item(PyTuple::empty(py));
|
||||
}
|
||||
Ok(arr)
|
||||
}
|
||||
|
||||
fn prefix_err(py: Python<'_>, name: &str, e: PyErr) -> PyErr {
|
||||
let msg = format!("attribute {name}: {}", e.value(py));
|
||||
if e.is_instance_of::<PyTypeError>(py) {
|
||||
PyTypeError::new_err(msg)
|
||||
} else {
|
||||
PyValueError::new_err(msg)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn read_attrs_len() {
|
||||
let mut map = HashMap::new();
|
||||
map.insert("a".into(), clawhdf5_rs::AttrValue::I64(1));
|
||||
map.insert("b".into(), clawhdf5_rs::AttrValue::F64(2.0));
|
||||
let attrs = PyAttrs::from_read(map);
|
||||
assert_eq!(attrs.__len__(), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn read_attrs_contains() {
|
||||
let mut map = HashMap::new();
|
||||
map.insert("x".into(), clawhdf5_rs::AttrValue::String("hello".into()));
|
||||
let attrs = PyAttrs::from_read(map);
|
||||
assert!(attrs.__contains__("x"));
|
||||
assert!(!attrs.__contains__("y"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn write_attrs_len() {
|
||||
let store = Arc::new(Mutex::new(Vec::new()));
|
||||
|
||||
Reference in New Issue
Block a user