quantumclaw/clawhdf5
3
1
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 1 Wiki Activity

feat(py): h5py-style reads of only the selected elements, GIL released

Browse Source 
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
2026-09-26 08:19:48 -05:00
co-authored by Claude Opus 5.5
parent 006bf3b131
commit 2d4b211523
12 changed files with 2305 additions and 385 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/clawhdf5-py/src/attrs.rs
+103 -45
View File
@@ -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()));
crates/clawhdf5-py/src/convert.rs
+600
View File
@@ -0,0 +1,600 @@
//! HDF5 datatypes as numpy dtypes, and element bytes as numpy arrays.
//!
//! The dtype a file's datatype maps to is the one h5py reports for it
//! (byte order kept, compound offsets and padding kept, `r`/`i` compounds as
//! complex, the `FALSE`/`TRUE` enum as `bool`, fixed strings as `S<n>`,
//! variable-length data as `object`). For every fixed-size type that dtype
//! describes the file's element bytes exactly, so the bytes the library
//! returns become the array's buffer as they are: the `Vec<u8>` is handed to
//! numpy without a copy and viewed as the dtype.
//!
//! Anything this mapping cannot describe exactly — non-IEEE floats, integers
//! with padding bits, VAX byte order, references, bitfields, time, and
//! variable-length members inside compounds or arrays — is a `TypeError`,
//! never a best-effort guess.
use std::collections::HashMap;
use clawhdf5_format::datatype::{CharacterSet, Datatype, DatatypeByteOrder};
use clawhdf5_format::global_heap::GlobalHeapCollection;
use numpy::PyArray1;
use pyo3::exceptions::{PyTypeError, PyValueError};
use pyo3::prelude::*;
use pyo3::types::{PyBytes, PyDict, PyList, PyString, PyTuple};
/// How the elements of a datatype become Python values.
#[derive(Clone, Debug, PartialEq)]
pub(crate) enum Layout {
/// Fixed-size elements numpy reads as they are.
Fixed,
/// A top-level HDF5 array type: elements are viewed as the base dtype and
/// the array's dimensions are appended to the shape (as h5py does).
Subarray(Vec<usize>),
/// Variable-length string: a global heap reference per element.
VlString { utf8: bool },
/// Variable-length sequence of a fixed-size base type.
VlSequence,
}
/// Everything needed to turn a dataset's or attribute's bytes into numpy.
pub(crate) struct Converter {
/// The dtype reported to Python (`Dataset.dtype`).
pub dtype: Py<PyAny>,
/// The dtype the element bytes are viewed as: `dtype` itself, the base
/// of a subarray, or the base of a variable-length sequence.
pub view: Py<PyAny>,
pub layout: Layout,
/// Bytes per element in the raw buffer the library returns.
pub elem_size: usize,
/// For variable-length data, bytes per unit of an element's stored
/// length: 1 for strings, the base type's size for sequences.
pub vl_unit: usize,
}
fn unsupported(what: impl std::fmt::Display) -> PyErr {
PyTypeError::new_err(format!(
"clawhdf5 cannot read this datatype into numpy: {what}"
))
}
fn byte_order_char(order: &DatatypeByteOrder, size: u32) -> PyResult<&'static str> {
if size == 1 {
return Ok("|");
}
match order {
DatatypeByteOrder::LittleEndian => Ok("<"),
DatatypeByteOrder::BigEndian => Ok(">"),
DatatypeByteOrder::Vax => Err(unsupported("VAX byte order")),
}
}
/// The numpy format string of an integer type, if it is a plain one.
fn int_format(dt: &Datatype) -> PyResult<String> {
match dt {
Datatype::FixedPoint {
size,
byte_order,
signed,
bit_offset,
bit_precision,
} => {
if !matches!(size, 1 | 2 | 4 | 8) {
return Err(unsupported(format!("{size}-byte integer")));
}
if *bit_offset != 0 || u32::from(*bit_precision) != size * 8 {
return Err(unsupported(format!(
"integer with {bit_precision} significant bits at offset {bit_offset} in {size} bytes"
)));
}
let kind = if *signed { 'i' } else { 'u' };
Ok(format!(
"{}{kind}{size}",
byte_order_char(byte_order, *size)?
))
}
other => Err(unsupported(format!("{other:?} is not an integer"))),
}
}
/// The numpy format string of an IEEE 754 binary16/32/64 type.
fn float_format(dt: &Datatype) -> PyResult<String> {
let Datatype::FloatingPoint {
size,
byte_order,
bit_offset,
bit_precision,
exponent_location,
exponent_size,
mantissa_location,
mantissa_size,
exponent_bias,
} = dt
else {
return Err(unsupported(format!("{dt:?} is not a float")));
};
// (exponent location, exponent size, mantissa size, bias) of IEEE 754.
let ieee = match size {
2 => (10, 5, 10, 15),
4 => (23, 8, 23, 127),
8 => (52, 11, 52, 1023),
_ => return Err(unsupported(format!("{size}-byte float"))),
};
let layout = (
*exponent_location,
*exponent_size,
*mantissa_size,
*exponent_bias,
);
if *bit_offset != 0
|| u32::from(*bit_precision) != size * 8
|| *mantissa_location != 0
|| layout != ieee
{
return Err(unsupported(format!(
"non-IEEE {size}-byte float (exponent {exponent_size} bits at {exponent_location}, \
mantissa {mantissa_size} bits at {mantissa_location}, bias {exponent_bias})"
)));
}
Ok(format!("{}f{size}", byte_order_char(byte_order, *size)?))
}
/// `r`/`i` compounds of two identical IEEE floats are complex numbers in h5py.
fn complex_format(
size: u32,
members: &[clawhdf5_format::datatype::CompoundMember],
) -> Option<String> {
let [re, im] = members else { return None };
if re.name != "r" || im.name != "i" || re.datatype != im.datatype {
return None;
}
let Datatype::FloatingPoint {
size: fsize,
byte_order,
..
} = &re.datatype
else {
return None;
};
if !matches!(fsize, 4 | 8)
|| re.byte_offset != 0
|| im.byte_offset != u64::from(*fsize)
|| size != 2 * fsize
{
return None;
}
float_format(&re.datatype).ok()?;
let order = byte_order_char(byte_order, *fsize).ok()?;
Some(format!("{order}c{}", 2 * fsize))
}
/// The members of an enum as `{name: value}`.
fn enum_members<'py>(
py: Python<'py>,
base: &Datatype,
members: &[clawhdf5_format::datatype::EnumMember],
) -> PyResult<Bound<'py, PyDict>> {
let signed = matches!(base, Datatype::FixedPoint { signed: true, .. });
let dict = PyDict::new(py);
for m in members {
let value: Py<PyAny> = if signed {
let v = clawhdf5_format::data_read::read_as_i64(&m.value, base)
.map_err(|e| PyValueError::new_err(format!("enum member {}: {e}", m.name)))?;
let v = *v.first().ok_or_else(|| {
PyValueError::new_err(format!("enum member {} has no value", m.name))
})?;
v.into_pyobject(py)?.into_any().unbind()
} else {
let v = clawhdf5_format::data_read::read_as_u64(&m.value, base)
.map_err(|e| PyValueError::new_err(format!("enum member {}: {e}", m.name)))?;
let v = *v.first().ok_or_else(|| {
PyValueError::new_err(format!("enum member {} has no value", m.name))
})?;
v.into_pyobject(py)?.into_any().unbind()
};
dict.set_item(&m.name, value)?;
}
Ok(dict)
}
/// Whether an enum is h5py's boolean: a one-byte integer with exactly the
/// members `FALSE` = 0 and `TRUE` = 1.
fn is_h5py_bool(base: &Datatype, members: &[clawhdf5_format::datatype::EnumMember]) -> bool {
if base.type_size() != 1 || members.len() != 2 {
return false;
}
let value = |name: &str| {
members
.iter()
.find(|m| m.name == name)
.and_then(|m| m.value.first().copied())
};
value("FALSE") == Some(0) && value("TRUE") == Some(1)
}
fn np_dtype<'py>(py: Python<'py>, spec: impl IntoPyObject<'py>) -> PyResult<Bound<'py, PyAny>> {
py.import("numpy")?.getattr("dtype")?.call1((spec,))
}
fn np_dtype_with_metadata<'py>(
py: Python<'py>,
spec: impl IntoPyObject<'py>,
metadata: Bound<'py, PyDict>,
) -> PyResult<Bound<'py, PyAny>> {
let kwargs = PyDict::new(py);
kwargs.set_item("metadata", metadata)?;
py.import("numpy")?
.getattr("dtype")?
.call((spec,), Some(&kwargs))
}
/// The numpy dtype of a fixed-size datatype, whose element bytes numpy can
/// read as they are.
pub(crate) fn fixed_dtype<'py>(py: Python<'py>, dt: &Datatype) -> PyResult<Bound<'py, PyAny>> {
match dt {
Datatype::FixedPoint { .. } => np_dtype(py, int_format(dt)?),
Datatype::FloatingPoint { .. } => np_dtype(py, float_format(dt)?),
Datatype::String { size, charset, .. } => {
if *size == 0 {
return Err(unsupported("zero-length fixed string"));
}
let meta = PyDict::new(py);
let enc = match charset {
CharacterSet::Ascii => "ascii",
CharacterSet::Utf8 => "utf-8",
};
meta.set_item("h5py_encoding", enc)?;
np_dtype_with_metadata(py, format!("S{size}"), meta)
}
Datatype::Opaque { size, .. } => {
if *size == 0 {
return Err(unsupported("zero-length opaque type"));
}
np_dtype(py, format!("V{size}"))
}
Datatype::Enumeration {
base_type, members, ..
} => {
let base = int_format(base_type)?;
if is_h5py_bool(base_type, members) {
return np_dtype(py, "?");
}
let meta = PyDict::new(py);
meta.set_item("enum", enum_members(py, base_type, members)?)?;
np_dtype_with_metadata(py, base, meta)
}
Datatype::Compound { size, members } => {
if let Some(c) = complex_format(*size, members) {
return np_dtype(py, c);
}
let names = PyList::empty(py);
let formats = PyList::empty(py);
let offsets = PyList::empty(py);
for m in members {
let end = m.byte_offset.checked_add(u64::from(m.datatype.type_size()));
if end.is_none_or(|end| end > u64::from(*size)) {
return Err(PyValueError::new_err(format!(
"compound member {} lies outside the {size}-byte compound",
m.name
)));
}
names.append(&m.name)?;
formats.append(fixed_dtype(py, &m.datatype).map_err(|e| {
unsupported(format!("compound member {}: {}", m.name, e.value(py)))
})?)?;
offsets.append(m.byte_offset)?;
}
let spec = PyDict::new(py);
spec.set_item("names", names)?;
spec.set_item("formats", formats)?;
spec.set_item("offsets", offsets)?;
spec.set_item("itemsize", size)?;
np_dtype(py, spec)
}
Datatype::Array {
base_type,
dimensions,
} => {
let base = fixed_dtype(py, base_type)?;
let dims = PyTuple::new(py, dimensions)?;
np_dtype(py, (base, dims))
}
Datatype::VariableLength { is_string, .. } => Err(unsupported(if *is_string {
"variable-length string inside a compound or array type"
} else {
"variable-length sequence inside a compound or array type"
})),
Datatype::Reference { .. } => Err(unsupported("object/region references")),
Datatype::BitField { .. } => Err(unsupported("bitfield")),
Datatype::Time { .. } => Err(unsupported("time")),
}
}
impl Converter {
/// The converter for a dataset's or attribute's datatype.
pub(crate) fn new(py: Python<'_>, dt: &Datatype, offset_size: u8) -> PyResult<Self> {
match dt {
Datatype::VariableLength {
is_string: true,
charset,
..
} => {
let utf8 = matches!(charset, Some(CharacterSet::Utf8));
let meta = PyDict::new(py);
if utf8 {
meta.set_item("vlen", py.get_type::<PyString>())?;
} else {
meta.set_item("vlen", py.get_type::<PyBytes>())?;
}
let dtype = np_dtype_with_metadata(py, "O", meta)?;
Ok(Self {
view: dtype.clone().unbind(),
dtype: dtype.unbind(),
layout: Layout::VlString { utf8 },
elem_size: vl_ref_size(offset_size)?,
vl_unit: 1,
})
}
Datatype::VariableLength {
is_string: false,
base_type,
..
} => {
let base = fixed_dtype(py, base_type)?;
let meta = PyDict::new(py);
meta.set_item("vlen", &base)?;
let dtype = np_dtype_with_metadata(py, "O", meta)?;
Ok(Self {
dtype: dtype.unbind(),
view: base.unbind(),
layout: Layout::VlSequence,
elem_size: vl_ref_size(offset_size)?,
vl_unit: base_type.type_size() as usize,
})
}
Datatype::Array {
base_type,
dimensions,
} => {
let dtype = fixed_dtype(py, dt)?;
let base = fixed_dtype(py, base_type)?;
Ok(Self {
dtype: dtype.unbind(),
view: base.unbind(),
layout: Layout::Subarray(dimensions.iter().map(|&d| d as usize).collect()),
elem_size: dt.type_size() as usize,
vl_unit: 0,
})
}
_ => {
let dtype = fixed_dtype(py, dt)?;
Ok(Self {
view: dtype.clone().unbind(),
dtype: dtype.unbind(),
layout: Layout::Fixed,
elem_size: dt.type_size() as usize,
vl_unit: 0,
})
}
}
}
pub(crate) fn is_vl(&self) -> bool {
matches!(self.layout, Layout::VlString { .. } | Layout::VlSequence)
}
/// An empty array of `shape` (some dimension is zero).
pub(crate) fn empty<'py>(
&self,
py: Python<'py>,
shape: &[usize],
) -> PyResult<Bound<'py, PyAny>> {
let np = py.import("numpy")?;
match &self.layout {
Layout::Subarray(dims) => {
let mut full = shape.to_vec();
full.extend_from_slice(dims);
np.call_method1("empty", (PyTuple::new(py, full)?, self.view.bind(py)))
}
_ => np.call_method1("empty", (PyTuple::new(py, shape)?, self.dtype.bind(py))),
}
}
/// Turn decoded element data into a numpy array of `shape`.
///
/// `str_values` decodes variable-length strings to `str` (what h5py
/// does for attributes) instead of `bytes` (what it does for datasets).
pub(crate) fn to_array<'py>(
&self,
py: Python<'py>,
data: Elements,
shape: &[usize],
str_values: bool,
) -> PyResult<Bound<'py, PyAny>> {
let n: usize = shape.iter().product();
match (data, &self.layout) {
(Elements::Bytes(bytes), Layout::Fixed) => {
bytes_as_array(py, bytes, self.view.bind(py), shape)
}
(Elements::Bytes(bytes), Layout::Subarray(dims)) => {
let mut full = shape.to_vec();
full.extend_from_slice(dims);
bytes_as_array(py, bytes, self.view.bind(py), &full)
}
(Elements::Vl(items), Layout::VlString { .. }) => {
check_count(items.len(), n)?;
let mut objs: Vec<Py<PyAny>> = Vec::with_capacity(items.len());
for item in items {
let obj = if str_values {
PyBytes::new(py, &item)
.call_method1("decode", ("utf-8", "surrogateescape"))?
.unbind()
} else {
PyBytes::new(py, &item).into_any().unbind()
};
objs.push(obj);
}
object_array(py, objs, shape)
}
(Elements::Vl(items), Layout::VlSequence) => {
check_count(items.len(), n)?;
let base = self.view.bind(py);
let itemsize: usize = base.getattr("itemsize")?.extract()?;
let mut objs: Vec<Py<PyAny>> = Vec::with_capacity(items.len());
for item in items {
if item.len() % itemsize != 0 {
return Err(PyValueError::new_err(format!(
"variable-length element of {} bytes is not a whole number of {itemsize}-byte values",
item.len()
)));
}
let len = item.len() / itemsize;
objs.push(bytes_as_array(py, item, base, &[len])?.unbind());
}
object_array(py, objs, shape)
}
_ => Err(PyValueError::new_err(
"internal error: element data does not match the datatype",
)),
}
}
}
/// Element data as read, before it becomes numpy.
pub(crate) enum Elements {
/// The elements' bytes, back to back.
Bytes(Vec<u8>),
/// Each variable-length element's bytes, resolved from the global heap.
Vl(Vec<Vec<u8>>),
}
fn vl_ref_size(offset_size: u8) -> PyResult<usize> {
// The library sizes a variable-length element as 16 bytes (a length, an
// 8-byte heap address and an index) whatever the file's offset size.
// Refuse the other sizes rather than read misaligned references.
if offset_size != 8 {
return Err(unsupported(format!(
"variable-length data in a file with {offset_size}-byte offsets"
)));
}
Ok(4 + usize::from(offset_size) + 4)
}
fn check_count(got: usize, want: usize) -> PyResult<()> {
if got != want {
return Err(PyValueError::new_err(format!(
"read {got} elements, expected {want}"
)));
}
Ok(())
}
/// A numpy array over `bytes` without copying them: the `Vec` becomes the
/// array's buffer and is viewed as `dtype` with `shape`.
pub(crate) fn bytes_as_array<'py>(
py: Python<'py>,
bytes: Vec<u8>,
dtype: &Bound<'py, PyAny>,
shape: &[usize],
) -> PyResult<Bound<'py, PyAny>> {
let itemsize: usize = dtype.getattr("itemsize")?.extract()?;
let n: usize = shape.iter().product();
if n.checked_mul(itemsize) != Some(bytes.len()) {
return Err(PyValueError::new_err(format!(
"read {} bytes, expected {n} elements of {itemsize} bytes",
bytes.len()
)));
}
let shape = PyTuple::new(py, shape)?;
if n == 0 {
return py.import("numpy")?.call_method1("empty", (shape, dtype));
}
let raw = PyArray1::from_vec(py, bytes);
let arr = raw
.call_method1("view", (dtype,))?
.call_method1("reshape", (shape,))?;
// A `Vec<u8>` carries no alignment promise. numpy copes with unaligned
// arrays, but slowly and not in every routine, so hand out an aligned
// copy in the (allocator-dependent, rare) case the buffer is not.
if !arr
.getattr("flags")?
.getattr("aligned")?
.extract::<bool>()?
{
return arr.call_method0("copy");
}
Ok(arr)
}
fn object_array<'py>(
py: Python<'py>,
objs: Vec<Py<PyAny>>,
shape: &[usize],
) -> PyResult<Bound<'py, PyAny>> {
let arr = PyArray1::from_vec(py, objs);
arr.call_method1("reshape", (PyTuple::new(py, shape)?,))
}
/// Resolve variable-length elements (global heap references in `raw`) to
/// their bytes: each element's stored length times `unit` (1 for strings,
/// the base type's size for sequences). Pure Rust, so it runs without the
/// GIL.
pub(crate) fn resolve_vl(
file_data: &[u8],
raw: &[u8],
count: usize,
offset_size: u8,
length_size: u8,
unit: usize,
) -> Result<Vec<Vec<u8>>, String> {
let refs = clawhdf5_format::vl_data::parse_vl_references(raw, count as u64, offset_size)
.map_err(|e| e.to_string())?;
let undefined = match offset_size {
2 => 0xFFFF,
4 => 0xFFFF_FFFF,
_ => u64::MAX,
};
let mut collections: HashMap<u64, GlobalHeapCollection> = HashMap::new();
let mut out = Vec::with_capacity(refs.len());
for vl in &refs {
if vl.collection_address == 0 || vl.collection_address == undefined {
if vl.length != 0 {
return Err(format!(
"variable-length element of length {} has no heap address",
vl.length
));
}
out.push(Vec::new());
continue;
}
let coll = match collections.entry(vl.collection_address) {
std::collections::hash_map::Entry::Occupied(e) => e.into_mut(),
std::collections::hash_map::Entry::Vacant(e) => {
let addr = usize::try_from(vl.collection_address)
.map_err(|_| "global heap address out of range".to_string())?;
e.insert(
GlobalHeapCollection::parse(file_data, addr, length_size)
.map_err(|e| e.to_string())?,
)
}
};
let index = u16::try_from(vl.object_index)
.map_err(|_| format!("global heap object index {} out of range", vl.object_index))?;
let obj = coll.get_object(index).ok_or_else(|| {
format!(
"global heap object {index} not found in the collection at {}",
vl.collection_address
)
})?;
let need = (vl.length as usize)
.checked_mul(unit)
.ok_or("variable-length element too long")?;
if need > obj.data.len() {
return Err(format!(
"variable-length element of {need} bytes in a {}-byte heap object",
obj.data.len()
));
}
out.push(obj.data[..need].to_vec());
}
Ok(out)
}
crates/clawhdf5-py/src/dataset.rs
+301 -201
View File
@@ -1,241 +1,341 @@
//! PyDataset — read access to HDF5 datasets with numpy integration.
//! PyDataset — h5py-style read access to HDF5 datasets.
//!
//! `ds[key]` parses the key into hyperslab selections (see `select`) and
//! reads only those elements through the facade's `read_selection`; the
//! bytes it returns become the numpy array's buffer without a copy (see
//! `convert`). All file access and decoding runs with the GIL released, so
//! Python threads reading the same or different datasets run in parallel.
use std::sync::Arc;
use numpy::PyArrayDyn;
use numpy::ndarray::{ArrayD, IxDyn};
use clawhdf5_format::datatype::Datatype;
use pyo3::exceptions::{PyTypeError, PyValueError};
use pyo3::prelude::*;
use pyo3::types::PyList;
use clawhdf5_rs::DType;
use pyo3::types::{PyList, PyTuple};
use crate::attrs::PyAttrs;
use crate::to_py_err;
use crate::convert::{Converter, Elements, resolve_vl};
use crate::select::{self, Plan};
use crate::{PyEmpty, node, to_py_err};
/// A handle to an HDF5 dataset (read mode).
/// A dataset in a file opened for reading.
///
/// Supports numpy-style indexing via `__getitem__`:
/// ```python
/// ds = f['dataset_name']
/// data = ds[:] # read all data as numpy array
/// shape = ds.shape
/// dtype = ds.dtype
/// ds = f['group/dataset']
/// ds.shape, ds.dtype, ds.attrs['units']
/// block = ds[10:20, ::2] # reads only the selected elements
/// ```
#[pyclass(name = "Dataset")]
pub struct PyDataset {
file: Arc<clawhdf5_rs::File>,
path: String,
cached_shape: Vec<u64>,
cached_dtype: DType,
/// `None` for a dataset with a null dataspace (h5py's `Empty`).
shape: Option<Vec<u64>>,
datatype: Datatype,
/// Why the datatype cannot be read into numpy, if it cannot.
conv: Result<Converter, String>,
}
impl PyDataset {
pub fn new(file: Arc<clawhdf5_rs::File>, path: String) -> PyResult<Self> {
let ds = file.dataset(&path).map_err(to_py_err)?;
let cached_shape = ds.shape().map_err(to_py_err)?;
let cached_dtype = ds.dtype().map_err(to_py_err)?;
pub(crate) fn open(
py: Python<'_>,
file: Arc<clawhdf5_rs::File>,
path: String,
) -> PyResult<Self> {
let hdr = node::header(&file, &path)?;
let null = node::is_null(&node::dataspace(&file, &hdr)?);
let (shape, datatype) = {
let ds = file.dataset(&path).map_err(to_py_err)?;
let shape = if null {
None
} else {
Some(ds.shape().map_err(to_py_err)?)
};
(shape, ds.raw_datatype().map_err(to_py_err)?)
};
let conv = Converter::new(py, &datatype, file.superblock().offset_size)
.map_err(|e| e.value(py).to_string());
Ok(Self {
file,
path,
cached_shape,
cached_dtype,
shape,
datatype,
conv,
})
}
}
/// Map a `DType` to a numpy dtype string.
fn dtype_to_numpy_str(dt: &DType) -> &'static str {
match dt {
DType::F64 => "float64",
DType::F32 => "float32",
DType::I64 => "int64",
DType::I32 => "int32",
DType::I16 => "int16",
DType::I8 => "int8",
DType::U64 => "uint64",
DType::U32 => "uint32",
DType::U16 => "uint16",
DType::U8 => "uint8",
DType::String | DType::VariableLengthString => "object",
_ => "object",
}
}
#[pymethods]
impl PyDataset {
/// The shape of the dataset as a tuple.
#[getter]
fn shape(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
let tuple = pyo3::types::PyTuple::new(py, self.cached_shape.iter().map(|&d| d as usize))?;
Ok(tuple.into_any().unbind())
fn converter(&self) -> PyResult<&Converter> {
self.conv
.as_ref()
.map_err(|msg| PyTypeError::new_err(format!("{}: {msg}", node::name(&self.path))))
}
/// The numpy dtype string of the dataset.
#[getter]
fn dtype(&self) -> &'static str {
dtype_to_numpy_str(&self.cached_dtype)
}
/// Read the selection described by `plan` into a numpy array.
fn read_plan<'py>(&self, py: Python<'py>, plan: &Plan) -> PyResult<Bound<'py, PyAny>> {
let conv = self.converter()?;
let dims = self.shape.as_deref().unwrap_or(&[]);
let out_shape = plan.out_shape();
/// Attribute access (read-only).
#[getter]
fn attrs(&self) -> PyResult<PyAttrs> {
let ds = self.file.dataset(&self.path).map_err(to_py_err)?;
let map = ds.attrs().map_err(to_py_err)?;
Ok(PyAttrs::from_read(map))
}
/// Read data via indexing. Supports `ds[:]`, `ds[0]`, `ds[0:5]`, etc.
///
/// The full dataset is always read from the underlying file; the index
/// is then applied on the resulting numpy array.
fn __getitem__<'py>(&self, py: Python<'py>, key: &Bound<'py, PyAny>) -> PyResult<Py<PyAny>> {
let arr = self.read_as_numpy(py)?;
let indexed = arr.get_item(key)?;
Ok(indexed.unbind())
}
fn __repr__(&self) -> String {
format!(
"<HDF5 Dataset \"{}\": shape {:?}, dtype {}>",
self.path,
self.cached_shape,
dtype_to_numpy_str(&self.cached_dtype),
)
}
fn __len__(&self) -> usize {
self.cached_shape.first().copied().unwrap_or(0) as usize
}
}
impl PyDataset {
/// Read the full dataset and return it as a numpy array (or list for strings).
///
/// For numeric types, the Rust I/O (file reading + decompression) is
/// performed inside `py.detach()` so that the GIL is released
/// during the potentially expensive operation. The numpy array
/// construction still happens with the GIL held.
fn read_as_numpy<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
let file = &self.file;
let path = &self.path;
let shape: Vec<usize> = self.cached_shape.iter().map(|&d| d as usize).collect();
match &self.cached_dtype {
DType::F64 => {
let data = py
.detach(|| file.dataset(path).and_then(|ds| ds.read_f64()))
.map_err(to_py_err)?;
let nd = ArrayD::from_shape_vec(IxDyn(&shape), data)
.map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))?;
let arr = PyArrayDyn::from_owned_array(py, nd);
Ok(arr.into_any())
}
DType::F32 => {
let data = py
.detach(|| file.dataset(path).and_then(|ds| ds.read_f32()))
.map_err(to_py_err)?;
let nd = ArrayD::from_shape_vec(IxDyn(&shape), data)
.map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))?;
let arr = PyArrayDyn::from_owned_array(py, nd);
Ok(arr.into_any())
}
DType::I32 => {
let data = py
.detach(|| file.dataset(path).and_then(|ds| ds.read_i32()))
.map_err(to_py_err)?;
let nd = ArrayD::from_shape_vec(IxDyn(&shape), data)
.map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))?;
let arr = PyArrayDyn::from_owned_array(py, nd);
Ok(arr.into_any())
}
DType::I64 => {
let data = py
.detach(|| file.dataset(path).and_then(|ds| ds.read_i64()))
.map_err(to_py_err)?;
let nd = ArrayD::from_shape_vec(IxDyn(&shape), data)
.map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))?;
let arr = PyArrayDyn::from_owned_array(py, nd);
Ok(arr.into_any())
}
DType::U8 => {
// Try zero-copy first (contiguous layout), fall back to
// read_u64 + cast for chunked/compact datasets.
let data: Vec<u8> = py
.detach(|| {
let ds = file.dataset(path)?;
match ds.read_u8_zerocopy() {
Ok(slice) => Ok(slice.to_vec()),
Err(_) => {
let raw = ds.read_u64()?;
Ok(raw.iter().map(|&v| v as u8).collect())
let arr = if plan.is_empty() {
conv.empty(py, &out_shape)?
} else {
let (reads, list_axis) = plan.reads(dims);
let file = &*self.file;
let path = self.path.as_str();
let (vl, elem_size, unit) = (conv.is_vl(), conv.elem_size, conv.vl_unit);
// Everything below touches only Rust data: release the GIL.
let blocks: Vec<(Elements, Vec<usize>)> = py
.detach(|| -> Result<_, ReadError> {
let ds = file.dataset(path)?;
let sb = file.superblock();
let mut blocks = Vec::with_capacity(reads.len());
for (sel, shape) in reads {
let raw = ds.read_selection(&sel)?;
let n: usize = shape.iter().product();
let data = if vl {
if raw.len() != n * elem_size {
return Err(ReadError::Other(format!(
"read {} bytes of variable-length references, expected {}",
raw.len(),
n * elem_size
)));
}
}
})
.map_err(to_py_err)?;
let nd = ArrayD::from_shape_vec(IxDyn(&shape), data)
.map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))?;
let arr = PyArrayDyn::from_owned_array(py, nd);
Ok(arr.into_any())
Elements::Vl(
resolve_vl(
file.as_bytes(),
&raw,
n,
sb.offset_size,
sb.length_size,
unit,
)
.map_err(ReadError::Other)?,
)
} else {
Elements::Bytes(raw)
};
blocks.push((data, shape));
}
Ok(blocks)
})
.map_err(|e| e.into_py(&self.path))?;
let mut arrays = Vec::with_capacity(blocks.len());
for (data, shape) in blocks {
arrays.push(conv.to_array(py, data, &shape, false)?);
}
DType::U64 => {
let data = py
.detach(|| file.dataset(path).and_then(|ds| ds.read_u64()))
.map_err(to_py_err)?;
let nd = ArrayD::from_shape_vec(IxDyn(&shape), data)
.map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))?;
let arr = PyArrayDyn::from_owned_array(py, nd);
Ok(arr.into_any())
let joined = if arrays.len() == 1 {
arrays.pop().expect("one block")
} else {
let axis = list_axis.expect("several reads only for a list index");
// Name the dtype: left to itself numpy canonicalises a
// structured dtype here (drops padding, native byte order).
let kwargs = pyo3::types::PyDict::new(py);
kwargs.set_item("axis", axis)?;
kwargs.set_item("dtype", arrays[0].getattr("dtype")?)?;
kwargs.set_item("casting", "no")?;
py.import("numpy")?.call_method(
"concatenate",
(PyList::new(py, arrays)?,),
Some(&kwargs),
)?
};
// Drop the axes indexed by an integer (length 1 in the blocks).
let mut shape = out_shape.clone();
if let crate::convert::Layout::Subarray(sub) = &conv.layout {
shape.extend_from_slice(sub);
}
DType::String | DType::VariableLengthString => {
// String reads need the GIL for PyList construction, but we
// release it during the Rust I/O portion.
let data = py
.detach(|| file.dataset(path).and_then(|ds| ds.read_string()))
.map_err(to_py_err)?;
let list = PyList::new(py, &data)?;
Ok(list.into_any())
}
other => Err(PyErr::new::<pyo3::exceptions::PyTypeError, _>(format!(
"unsupported dataset dtype for reading: {other}"
))),
joined.call_method1("reshape", (PyTuple::new(py, shape)?,))?
};
let arr = select_fields(py, arr, &plan.fields)?;
if plan.scalar {
return arr.get_item(PyTuple::empty(py));
}
Ok(arr)
}
}
/// An error from the read closure, turned into a Python error with the GIL.
enum ReadError {
Lib(clawhdf5_rs::Error),
Other(String),
}
impl From<clawhdf5_rs::Error> for ReadError {
fn from(e: clawhdf5_rs::Error) -> Self {
ReadError::Lib(e)
}
}
impl ReadError {
fn into_py(self, path: &str) -> PyErr {
match self {
ReadError::Lib(e) => to_py_err(e),
ReadError::Other(msg) => PyValueError::new_err(format!("{}: {msg}", node::name(path))),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
/// Keep only the named compound fields, as h5py's `ds['x']` / `ds['x', 'y']`.
fn select_fields<'py>(
py: Python<'py>,
arr: Bound<'py, PyAny>,
fields: &[String],
) -> PyResult<Bound<'py, PyAny>> {
if fields.is_empty() {
return Ok(arr);
}
let names = arr.getattr("dtype")?.getattr("names")?;
if names.is_none() {
return Err(PyValueError::new_err(
"Field names only allowed for compound types",
));
}
let names: Vec<String> = names.extract()?;
for f in fields {
if !names.contains(f) {
return Err(PyValueError::new_err(format!(
"Field {f} does not appear in this type."
)));
}
}
let np = py.import("numpy")?;
if let [one] = fields {
return np.call_method1("ascontiguousarray", (arr.get_item(one)?,));
}
let picked = arr.get_item(PyList::new(py, fields)?)?;
py.import("numpy.lib.recfunctions")?
.call_method1("repack_fields", (picked,))
}
#[test]
fn dtype_mapping() {
assert_eq!(dtype_to_numpy_str(&DType::F64), "float64");
assert_eq!(dtype_to_numpy_str(&DType::F32), "float32");
assert_eq!(dtype_to_numpy_str(&DType::I32), "int32");
assert_eq!(dtype_to_numpy_str(&DType::I64), "int64");
assert_eq!(dtype_to_numpy_str(&DType::U8), "uint8");
assert_eq!(dtype_to_numpy_str(&DType::String), "object");
#[pymethods]
impl PyDataset {
/// The shape of the dataset (`None` for an empty/null dataspace).
#[getter]
fn shape<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
match &self.shape {
Some(s) => Ok(PyTuple::new(py, s)?.into_any()),
None => Ok(py.None().into_bound(py)),
}
}
#[test]
fn dataset_from_file() {
let mut b = clawhdf5_rs::FileBuilder::new();
b.create_dataset("vals").with_f64_data(&[1.0, 2.0, 3.0]);
let bytes = b.finish().unwrap();
let file = Arc::new(clawhdf5_rs::File::from_bytes(bytes).unwrap());
let ds = PyDataset::new(file, "vals".into()).unwrap();
assert_eq!(ds.cached_shape, vec![3]);
assert_eq!(ds.cached_dtype, DType::F64);
/// The maximum shape (`None` per unlimited dimension), like h5py.
#[getter]
fn maxshape<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
let Some(shape) = &self.shape else {
return Ok(py.None().into_bound(py));
};
let max = self
.file
.dataset(&self.path)
.and_then(|ds| ds.max_dimensions())
.map_err(to_py_err)?
.unwrap_or_else(|| shape.clone());
let items: Vec<Option<u64>> = max
.into_iter()
.map(|d| (d != u64::MAX).then_some(d))
.collect();
Ok(PyTuple::new(py, items)?.into_any())
}
#[test]
fn dataset_len() {
let mut b = clawhdf5_rs::FileBuilder::new();
b.create_dataset("data")
.with_i32_data(&[10, 20, 30, 40])
.with_shape(&[2, 2]);
let bytes = b.finish().unwrap();
let file = Arc::new(clawhdf5_rs::File::from_bytes(bytes).unwrap());
let ds = PyDataset::new(file, "data".into()).unwrap();
assert_eq!(ds.__len__(), 2);
/// The dataset's numpy dtype, as h5py reports it.
#[getter]
fn dtype<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
Ok(self.converter()?.dtype.bind(py).clone())
}
#[getter]
fn ndim(&self) -> usize {
self.shape.as_ref().map_or(0, Vec::len)
}
/// Number of elements (`None` for an empty/null dataspace, as h5py).
#[getter]
fn size(&self) -> Option<u64> {
self.shape.as_ref().map(|s| s.iter().product())
}
/// The dataset's full name, e.g. `/group/data`.
#[getter]
fn name(&self) -> String {
node::name(&self.path)
}
/// The dataset's attributes (read-only, dict-like).
#[getter]
fn attrs(&self) -> PyResult<PyAttrs> {
PyAttrs::read(Arc::clone(&self.file), &self.path)
}
/// Read with h5py indexing: integers, slices with positive steps,
/// `...`, one increasing list of integers, and compound field names.
/// Only the selected elements are read from the file.
fn __getitem__<'py>(
&self,
py: Python<'py>,
key: &Bound<'py, PyAny>,
) -> PyResult<Bound<'py, PyAny>> {
let Some(dims) = &self.shape else {
let is_empty_tuple = key.cast::<PyTuple>().is_ok_and(|t| t.is_empty());
let is_ellipsis = key.is_instance_of::<pyo3::types::PyEllipsis>();
if is_empty_tuple || is_ellipsis {
let empty = PyEmpty::new(self.converter()?.dtype.clone_ref(py));
return Ok(empty.into_pyobject(py)?.into_any());
}
return Err(PyValueError::new_err("Empty datasets cannot be sliced"));
};
let plan = select::parse(key, dims)?;
self.read_plan(py, &plan)
}
/// `numpy.asarray(ds)` reads the whole dataset.
#[pyo3(signature = (dtype=None, copy=None))]
fn __array__<'py>(
&self,
py: Python<'py>,
dtype: Option<&Bound<'py, PyAny>>,
copy: Option<bool>,
) -> PyResult<Bound<'py, PyAny>> {
let _ = copy; // every read is a fresh array
let Some(dims) = &self.shape else {
return Err(PyValueError::new_err("an empty dataset has no array value"));
};
let ellipsis = pyo3::types::PyEllipsis::get(py).to_owned().into_any();
let plan = select::parse(&ellipsis, dims)?;
let arr = self.read_plan(py, &plan)?;
match dtype {
Some(dt) => arr.call_method1("astype", (dt,)),
None => Ok(arr),
}
}
fn __len__(&self) -> PyResult<usize> {
match self.shape.as_deref() {
Some([first, ..]) => Ok(*first as usize),
_ => Err(PyTypeError::new_err(
"Attempt to take len() of scalar dataset",
)),
}
}
fn __repr__(&self, py: Python<'_>) -> String {
let dtype = match &self.conv {
Ok(c) => c
.dtype
.bind(py)
.str()
.map(|s| s.to_string())
.unwrap_or_default(),
Err(_) => format!("{:?}", self.datatype),
};
let shape = match &self.shape {
Some(s) => format!("{s:?}"),
None => "None".to_string(),
};
format!(
"<HDF5 dataset \"{}\": shape {shape}, type \"{dtype}\">",
node::name(&self.path)
)
}
}
crates/clawhdf5-py/src/file.rs
+56 -41
View File
@@ -3,11 +3,12 @@
use std::path::PathBuf;
use std::sync::{Arc, Mutex};
use pyo3::exceptions::PyKeyError;
use pyo3::prelude::*;
use pyo3::types::PyList;
use crate::attrs::PyAttrs;
use crate::dataset::PyDataset;
use crate::group::{PyGroup, WriteGroupState, finalize_write_group};
use crate::group::{self, PyGroup, WriteGroupState, finalize_write_group};
use crate::{DatasetSpec, OwnedAttrValue, apply_dataset_spec, extract_numpy_data, to_py_err};
/// Internal state for write mode.
@@ -35,6 +36,7 @@ struct WriteState {
#[pyclass(name = "File")]
pub struct PyFile {
inner: Option<FileInner>,
filename: String,
}
enum FileInner {
@@ -51,15 +53,20 @@ impl PyFile {
/// mode: 'r' for read (default), 'w' for write
#[new]
#[pyo3(signature = (path, mode="r"))]
fn new(path: &str, mode: &str) -> PyResult<Self> {
fn new(py: Python<'_>, path: &str, mode: &str) -> PyResult<Self> {
let filename = path.to_string();
match mode {
"r" => {
let file = clawhdf5_rs::File::open(path).map_err(to_py_err)?;
let file = py
.detach(|| clawhdf5_rs::File::open(path))
.map_err(to_py_err)?;
Ok(Self {
inner: Some(FileInner::Read(Arc::new(file))),
filename,
})
}
"w" => Ok(Self {
filename,
inner: Some(FileInner::Write(WriteState {
path: PathBuf::from(path),
root_datasets: Vec::new(),
@@ -101,44 +108,56 @@ impl PyFile {
Ok(false) // don't suppress exceptions
}
/// Get a child object (dataset or group) by path.
/// Get a child object (dataset or group) by path; `f['/']` is the root.
fn __getitem__(&self, py: Python<'_>, key: &str) -> PyResult<Py<PyAny>> {
let file = self.read_file()?;
// Try dataset first
match file.dataset(key) {
Ok(_) => {
let ds = PyDataset::new(Arc::clone(file), key.to_string())?;
Ok(ds.into_pyobject(py)?.into_any().unbind())
}
Err(clawhdf5_rs::Error::NotADataset(_)) => {
let grp = PyGroup::from_read(Arc::clone(file), key.to_string());
Ok(grp.into_pyobject(py)?.into_any().unbind())
}
Err(_) => {
// Could be a group (no DataLayout message, no error)
match file.group(key) {
Ok(_) => {
let grp = PyGroup::from_read(Arc::clone(file), key.to_string());
Ok(grp.into_pyobject(py)?.into_any().unbind())
}
Err(e) => Err(PyErr::new::<pyo3::exceptions::PyKeyError, _>(format!(
"{key}: {e}"
))),
}
group::get_item(py, self.read_file()?, "", key)
}
/// `f.get(key, default=None)`.
#[pyo3(signature = (key, default=None))]
fn get(&self, py: Python<'_>, key: &str, default: Option<Py<PyAny>>) -> PyResult<Py<PyAny>> {
match group::get_item(py, self.read_file()?, "", key) {
Err(e) if e.is_instance_of::<PyKeyError>(py) => {
Ok(default.unwrap_or_else(|| py.None()))
}
other => other,
}
}
/// List the names of all children in the root group.
fn keys(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
let file = self.read_file()?;
let root = file.root();
let mut names = root.datasets().map_err(to_py_err)?;
let groups = root.groups().map_err(to_py_err)?;
names.extend(groups);
names.sort();
let list = pyo3::types::PyList::new(py, &names)?;
Ok(list.into_any().unbind())
let names = group::member_names(self.read_file()?, "")?;
Ok(PyList::new(py, names)?.into_any().unbind())
}
fn values(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
let vals = group::values(py, self.read_file()?, "")?;
Ok(PyList::new(py, vals)?.into_any().unbind())
}
fn items(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
let items = group::items(py, self.read_file()?, "")?;
Ok(PyList::new(py, items)?.into_any().unbind())
}
fn __iter__(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
self.keys(py)?.call_method0(py, "__iter__")
}
fn __len__(&self) -> PyResult<usize> {
Ok(group::member_names(self.read_file()?, "")?.len())
}
/// The root group's name, `/`.
#[getter]
fn name(&self) -> &'static str {
"/"
}
/// The path the file was opened with.
#[getter]
fn filename(&self) -> &str {
&self.filename
}
/// Create a dataset in the root group (write mode only).
@@ -192,10 +211,7 @@ impl PyFile {
#[getter]
fn attrs(&self) -> PyResult<PyAttrs> {
match self.inner.as_ref() {
Some(FileInner::Read(file)) => {
let map = file.root().attrs().map_err(to_py_err)?;
Ok(PyAttrs::from_read(map))
}
Some(FileInner::Read(file)) => PyAttrs::read(Arc::clone(file), ""),
Some(FileInner::Write(state)) => Ok(PyAttrs::from_write(Arc::clone(&state.root_attrs))),
None => Err(PyErr::new::<pyo3::exceptions::PyIOError, _>(
"file is closed",
@@ -216,8 +232,7 @@ impl PyFile {
}
fn __contains__(&self, key: &str) -> PyResult<bool> {
let file = self.read_file()?;
Ok(file.dataset(key).is_ok() || file.group(key).is_ok())
Ok(group::contains(self.read_file()?, "", key))
}
}
crates/clawhdf5-py/src/group.rs
+135 -93
View File
@@ -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]
crates/clawhdf5-py/src/lib.rs
+54 -2
View File
@@ -10,9 +10,12 @@
//! ```
mod attrs;
mod convert;
mod dataset;
mod file;
mod group;
mod node;
mod select;
use pyo3::prelude::*;
@@ -46,6 +49,54 @@ pub(crate) fn to_py_err(e: clawhdf5_rs::Error) -> PyErr {
}
}
/// The value of a dataset or attribute with a null dataspace: a type but no
/// data. Mirrors `h5py.Empty`.
#[pyclass(name = "Empty", frozen)]
pub struct PyEmpty {
dtype: Py<PyAny>,
}
impl PyEmpty {
pub(crate) fn new(dtype: Py<PyAny>) -> Self {
Self { dtype }
}
}
#[pymethods]
impl PyEmpty {
#[new]
fn py_new(py: Python<'_>, dtype: &Bound<'_, PyAny>) -> PyResult<Self> {
let dtype = py.import("numpy")?.getattr("dtype")?.call1((dtype,))?;
Ok(Self::new(dtype.unbind()))
}
#[getter]
fn dtype(&self, py: Python<'_>) -> Py<PyAny> {
self.dtype.clone_ref(py)
}
#[getter]
fn shape(&self, py: Python<'_>) -> Py<PyAny> {
py.None()
}
#[getter]
fn size(&self, py: Python<'_>) -> Py<PyAny> {
py.None()
}
fn __eq__(&self, py: Python<'_>, other: &Bound<'_, PyAny>) -> PyResult<bool> {
match other.cast::<PyEmpty>() {
Ok(o) => self.dtype.bind(py).eq(o.get().dtype.bind(py)),
Err(_) => Ok(false),
}
}
fn __repr__(&self, py: Python<'_>) -> PyResult<String> {
Ok(format!("Empty(dtype={})", self.dtype.bind(py).repr()?))
}
}
/// The data payload for a dataset being written.
#[derive(Clone)]
pub(crate) enum DatasetData {
@@ -224,6 +275,7 @@ fn clawhdf5(m: &Bound<'_, PyModule>) -> PyResult<()> {
m.add_class::<PyDataset>()?;
m.add_class::<PyGroup>()?;
m.add_class::<PyAttrs>()?;
m.add_class::<PyEmpty>()?;
Ok(())
}
@@ -233,9 +285,9 @@ mod tests {
#[test]
fn owned_attr_value_roundtrip() {
let val = OwnedAttrValue::F64(3.14);
let val = OwnedAttrValue::F64(2.5);
let attr: clawhdf5_rs::AttrValue = val.into();
assert!(matches!(attr, clawhdf5_rs::AttrValue::F64(v) if (v - 3.14).abs() < 1e-10));
assert!(matches!(attr, clawhdf5_rs::AttrValue::F64(v) if (v - 2.5).abs() < 1e-10));
}
#[test]
crates/clawhdf5-py/src/node.rs
+168
View File
@@ -0,0 +1,168 @@
//! Resolving paths to objects in a file opened for reading.
use std::sync::Arc;
use clawhdf5_format::attribute::AttributeMessage;
use clawhdf5_format::dataspace::{Dataspace, DataspaceType};
use clawhdf5_format::message_type::MessageType;
use clawhdf5_format::object_header::ObjectHeader;
use pyo3::exceptions::{PyKeyError, PyTypeError, PyValueError};
use pyo3::prelude::*;
use crate::dataset::PyDataset;
use crate::group::PyGroup;
/// Join `key` onto the group path `base` the way h5py does: an absolute key
/// starts from the root, a relative one from `base`. Paths are kept without
/// a leading `/`; the root is `""`.
pub(crate) fn join(base: &str, key: &str) -> String {
let parts = if key.starts_with('/') {
key.split('/').collect::<Vec<_>>()
} else {
base.split('/').chain(key.split('/')).collect()
};
parts
.into_iter()
.filter(|p| !p.is_empty() && *p != ".")
.collect::<Vec<_>>()
.join("/")
}
/// The HDF5 name (`/a/b`) of a path.
pub(crate) fn name(path: &str) -> String {
format!("/{path}")
}
/// The object header of the object at `path`.
pub(crate) fn header(file: &clawhdf5_rs::File, path: &str) -> PyResult<ObjectHeader> {
let sb = file.superblock();
let data = file.as_bytes();
let addr = if path.is_empty() {
sb.root_group_address
} else {
clawhdf5_format::group_v2::resolve_path_any(data, sb, path).map_err(|e| {
PyKeyError::new_err(format!(
"Unable to open object (object '{}' doesn't exist): {e}",
name(path)
))
})?
};
let addr = usize::try_from(addr)
.map_err(|_| PyValueError::new_err(format!("{}: address out of range", name(path))))?;
ObjectHeader::parse(data, addr, sb.offset_size, sb.length_size)
.map_err(|e| PyValueError::new_err(format!("{}: {e}", name(path))))
}
/// What an object header describes.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum Kind {
Dataset,
Group,
Datatype,
}
pub(crate) fn kind(hdr: &ObjectHeader) -> Option<Kind> {
let has = |t: MessageType| hdr.messages.iter().any(|m| m.msg_type == t);
if has(MessageType::DataLayout) {
Some(Kind::Dataset)
} else if has(MessageType::LinkInfo)
|| has(MessageType::Link)
|| has(MessageType::SymbolTable)
|| has(MessageType::GroupInfo)
{
Some(Kind::Group)
} else if has(MessageType::Datatype) {
Some(Kind::Datatype)
} else {
None
}
}
/// Open the object at `path` as a `Dataset` or `Group`.
pub(crate) fn open(
py: Python<'_>,
file: &Arc<clawhdf5_rs::File>,
path: String,
) -> PyResult<Py<PyAny>> {
let hdr = header(file, &path)?;
match kind(&hdr) {
Some(Kind::Dataset) => Ok(PyDataset::open(py, Arc::clone(file), path)?
.into_pyobject(py)?
.into_any()
.unbind()),
Some(Kind::Group) => Ok(PyGroup::from_read(Arc::clone(file), path)
.into_pyobject(py)?
.into_any()
.unbind()),
Some(Kind::Datatype) => Err(PyTypeError::new_err(format!(
"{}: committed (named) datatypes are not supported by clawhdf5",
name(&path)
))),
None => Err(PyValueError::new_err(format!(
"{}: not a dataset, group or datatype",
name(&path)
))),
}
}
/// Whether `path` names a dataset or group.
pub(crate) fn exists(file: &clawhdf5_rs::File, path: &str) -> bool {
header(file, path)
.ok()
.and_then(|h| kind(&h))
.is_some_and(|k| k != Kind::Datatype)
}
/// The dataspace message of an object header.
pub(crate) fn dataspace(file: &clawhdf5_rs::File, hdr: &ObjectHeader) -> PyResult<Dataspace> {
let sb = file.superblock();
let msg = hdr
.messages
.iter()
.find(|m| m.msg_type == MessageType::Dataspace)
.ok_or_else(|| PyValueError::new_err("object has no dataspace message"))?;
let data = clawhdf5_format::shared_message::message_data(
file.as_bytes(),
msg,
sb.offset_size,
sb.length_size,
)
.map_err(|e| PyValueError::new_err(e.to_string()))?;
Dataspace::parse(&data, sb.length_size).map_err(|e| PyValueError::new_err(e.to_string()))
}
pub(crate) fn is_null(space: &Dataspace) -> bool {
space.space_type == DataspaceType::Null
}
/// The attributes of the object at `path`, sorted by name (h5py's order).
/// Attributes whose messages cannot be parsed are left out, as the facade's
/// `attrs()` does.
pub(crate) fn attributes(file: &clawhdf5_rs::File, path: &str) -> PyResult<Vec<AttributeMessage>> {
let hdr = header(file, path)?;
let sb = file.superblock();
let (mut attrs, _errors) = clawhdf5_format::attribute::extract_attributes_tolerant(
file.as_bytes(),
&hdr,
sb.offset_size,
sb.length_size,
)
.map_err(|e| PyValueError::new_err(format!("{}: {e}", name(path))))?;
attrs.sort_by(|a, b| a.name.as_bytes().cmp(b.name.as_bytes()));
Ok(attrs)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn join_paths() {
assert_eq!(join("", "a"), "a");
assert_eq!(join("a", "b/c"), "a/b/c");
assert_eq!(join("a/b", "/x"), "x");
assert_eq!(join("a", "/"), "");
assert_eq!(join("", "/a//b/"), "a/b");
assert_eq!(join("a", "./b"), "a/b");
}
}
crates/clawhdf5-py/src/select.rs
+366
View File
@@ -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]);
}
}