perf(py): datasets and groups keep their address; groups their links
Every ds[...] and g[k] resolved the path from the root again, two or three times per open, and resolving a name in a large group scans its links: visiting a group was O(n^2). 4000 scalar datasets in one group took 39 s (v1 group) and 131 s (dense) to list, read and re-read; now 0.3 s each. A Dataset keeps its object address, a Group (and the file's root) its address and, after the first lookup, its link table. New facade API File::dataset_at(address), tested in integration_tests. Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>
This commit is contained in:
osobh
+174
-102
@@ -1,13 +1,14 @@
|
||||
//! PyGroup — navigable HDF5 group with read and write support.
|
||||
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::collections::HashMap;
|
||||
use std::sync::{Arc, Mutex, OnceLock};
|
||||
|
||||
use pyo3::exceptions::{PyIOError, PyKeyError};
|
||||
use pyo3::exceptions::{PyIOError, PyKeyError, PyValueError};
|
||||
use pyo3::prelude::*;
|
||||
use pyo3::types::PyList;
|
||||
|
||||
use crate::attrs::PyAttrs;
|
||||
use crate::{DatasetSpec, OwnedAttrValue, apply_dataset_spec, extract_numpy_data, node, to_py_err};
|
||||
use crate::{DatasetSpec, OwnedAttrValue, apply_dataset_spec, extract_numpy_data, node};
|
||||
|
||||
/// Shared state for a group being written.
|
||||
pub(crate) struct WriteGroupState {
|
||||
@@ -28,17 +29,14 @@ pub struct PyGroup {
|
||||
}
|
||||
|
||||
enum GroupInner {
|
||||
Read {
|
||||
file: Arc<clawhdf5_rs::File>,
|
||||
path: String,
|
||||
},
|
||||
Read(ReadGroup),
|
||||
Write(Arc<Mutex<WriteGroupState>>),
|
||||
}
|
||||
|
||||
impl PyGroup {
|
||||
pub(crate) fn from_read(file: Arc<clawhdf5_rs::File>, path: String) -> Self {
|
||||
pub(crate) fn from_read(file: Arc<clawhdf5_rs::File>, path: String, addr: u64) -> Self {
|
||||
Self {
|
||||
inner: GroupInner::Read { file, path },
|
||||
inner: GroupInner::Read(ReadGroup::new(file, path, addr)),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -48,9 +46,9 @@ impl PyGroup {
|
||||
}
|
||||
}
|
||||
|
||||
fn read_parts(&self, what: &str) -> PyResult<(&Arc<clawhdf5_rs::File>, &str)> {
|
||||
fn read_group(&self, what: &str) -> PyResult<&ReadGroup> {
|
||||
match &self.inner {
|
||||
GroupInner::Read { file, path } => Ok((file, path)),
|
||||
GroupInner::Read(g) => Ok(g),
|
||||
GroupInner::Write(_) => Err(PyIOError::new_err(format!(
|
||||
"cannot {what} a group opened for writing"
|
||||
))),
|
||||
@@ -58,77 +56,90 @@ impl PyGroup {
|
||||
}
|
||||
}
|
||||
|
||||
// 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))
|
||||
/// A group in a file opened for reading (a file is its root group, as in
|
||||
/// h5py). It keeps its own address and, once listed, its links, so looking
|
||||
/// up a child neither resolves the path from the root nor scans the group's
|
||||
/// links again: visiting every member of a large group is linear, not
|
||||
/// quadratic.
|
||||
pub(crate) struct ReadGroup {
|
||||
pub file: Arc<clawhdf5_rs::File>,
|
||||
pub path: String,
|
||||
pub addr: u64,
|
||||
/// Link name -> object address (soft links resolved), filled on first use.
|
||||
links: OnceLock<HashMap<String, u64>>,
|
||||
/// Names of the datasets and subgroups, sorted (h5py's order).
|
||||
members: OnceLock<Vec<String>>,
|
||||
}
|
||||
|
||||
/// 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>> {
|
||||
crate::no_panic(|| {
|
||||
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>> {
|
||||
let (file, path) = self.read_parts("read children from")?;
|
||||
get_item(py, file, path, key)
|
||||
impl ReadGroup {
|
||||
pub(crate) fn new(file: Arc<clawhdf5_rs::File>, path: String, addr: u64) -> Self {
|
||||
Self {
|
||||
file,
|
||||
path,
|
||||
addr,
|
||||
links: OnceLock::new(),
|
||||
members: OnceLock::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// `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) {
|
||||
fn links(&self) -> PyResult<&HashMap<String, u64>> {
|
||||
if let Some(links) = self.links.get() {
|
||||
return Ok(links);
|
||||
}
|
||||
let entries = crate::no_panic(|| {
|
||||
clawhdf5_format::group_v2::resolve_group_children(
|
||||
self.file.as_bytes(),
|
||||
self.file.superblock(),
|
||||
self.addr,
|
||||
)
|
||||
.map_err(|e| PyValueError::new_err(format!("{}: {e}", node::name(&self.path))))
|
||||
})?;
|
||||
let map = entries
|
||||
.into_iter()
|
||||
.map(|e| (e.name, e.object_header_address))
|
||||
.collect();
|
||||
Ok(self.links.get_or_init(|| map))
|
||||
}
|
||||
|
||||
/// The path and address of `key` (a name, a relative or an absolute path).
|
||||
fn locate(&self, key: &str) -> PyResult<(String, u64)> {
|
||||
let path = node::join(&self.path, key);
|
||||
let rel = if self.path.is_empty() {
|
||||
Some(path.as_str())
|
||||
} else if path == self.path {
|
||||
Some("")
|
||||
} else {
|
||||
path.strip_prefix(self.path.as_str())
|
||||
.and_then(|r| r.strip_prefix('/'))
|
||||
};
|
||||
let addr = match rel {
|
||||
// A direct child: the link table, when it has the name.
|
||||
Some(name) if !name.is_empty() && !name.contains('/') => {
|
||||
match self.links()?.get(name) {
|
||||
Some(&a) => a,
|
||||
None => node::resolve_from(&self.file, self.addr, name, &path)?,
|
||||
}
|
||||
}
|
||||
Some(rel) => node::resolve_from(&self.file, self.addr, rel, &path)?,
|
||||
None => node::address(&self.file, &path)?,
|
||||
};
|
||||
Ok((path, addr))
|
||||
}
|
||||
|
||||
/// `group[key]`.
|
||||
pub(crate) fn get_item(&self, py: Python<'_>, key: &str) -> PyResult<Py<PyAny>> {
|
||||
let (path, addr) = self.locate(key)?;
|
||||
node::open(py, &self.file, path, addr)
|
||||
}
|
||||
|
||||
/// `group.get(key, default)`.
|
||||
pub(crate) fn get(
|
||||
&self,
|
||||
py: Python<'_>,
|
||||
key: &str,
|
||||
default: Option<Py<PyAny>>,
|
||||
) -> PyResult<Py<PyAny>> {
|
||||
match self.get_item(py, key) {
|
||||
Err(e) if e.is_instance_of::<PyKeyError>(py) => {
|
||||
Ok(default.unwrap_or_else(|| py.None()))
|
||||
}
|
||||
@@ -136,11 +147,70 @@ impl PyGroup {
|
||||
}
|
||||
}
|
||||
|
||||
/// Names of the group's datasets and subgroups, sorted (h5py's order).
|
||||
pub(crate) fn member_names(&self) -> PyResult<&[String]> {
|
||||
if let Some(m) = self.members.get() {
|
||||
return Ok(m);
|
||||
}
|
||||
let mut names = Vec::new();
|
||||
for (name, &addr) in self.links()? {
|
||||
let hdr = node::header_at(&self.file, addr, &node::join(&self.path, name))?;
|
||||
if matches!(
|
||||
node::kind(&hdr),
|
||||
Some(node::Kind::Dataset | node::Kind::Group)
|
||||
) {
|
||||
names.push(name.clone());
|
||||
}
|
||||
}
|
||||
names.sort_by(|a, b| a.as_bytes().cmp(b.as_bytes()));
|
||||
Ok(self.members.get_or_init(|| names))
|
||||
}
|
||||
|
||||
pub(crate) fn contains(&self, key: &str) -> bool {
|
||||
self.locate(key)
|
||||
.and_then(|(path, addr)| node::header_at(&self.file, addr, &path))
|
||||
.ok()
|
||||
.and_then(|h| node::kind(&h))
|
||||
.is_some_and(|k| k != node::Kind::Datatype)
|
||||
}
|
||||
|
||||
pub(crate) fn values(&self, py: Python<'_>) -> PyResult<Vec<Py<PyAny>>> {
|
||||
self.member_names()?
|
||||
.iter()
|
||||
.map(|n| self.get_item(py, n))
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub(crate) fn items(&self, py: Python<'_>) -> PyResult<Vec<(String, Py<PyAny>)>> {
|
||||
self.member_names()?
|
||||
.iter()
|
||||
.map(|n| Ok((n.clone(), self.get_item(py, n)?)))
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub(crate) fn attrs(&self) -> PyResult<PyAttrs> {
|
||||
PyAttrs::read(Arc::clone(&self.file), self.addr, &self.path)
|
||||
}
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl PyGroup {
|
||||
/// Get a child object (dataset or subgroup) by name or path.
|
||||
fn __getitem__(&self, py: Python<'_>, key: &str) -> PyResult<Py<PyAny>> {
|
||||
self.read_group("read children from")?.get_item(py, 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>> {
|
||||
self.read_group("read children from")?.get(py, key, default)
|
||||
}
|
||||
|
||||
/// List the names of all children (datasets and subgroups).
|
||||
fn keys(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
|
||||
match &self.inner {
|
||||
GroupInner::Read { file, path } => {
|
||||
let list = PyList::new(py, member_names(file, path)?)?;
|
||||
GroupInner::Read(g) => {
|
||||
let list = PyList::new(py, g.member_names()?)?;
|
||||
Ok(list.into_any().unbind())
|
||||
}
|
||||
GroupInner::Write(state) => {
|
||||
@@ -153,15 +223,13 @@ 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())
|
||||
let g = self.read_group("read children from")?;
|
||||
Ok(PyList::new(py, g.values(py)?)?.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())
|
||||
let g = self.read_group("read children from")?;
|
||||
Ok(PyList::new(py, g.items(py)?)?.into_any().unbind())
|
||||
}
|
||||
|
||||
fn __iter__(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
|
||||
@@ -170,7 +238,7 @@ impl PyGroup {
|
||||
|
||||
fn __len__(&self) -> PyResult<usize> {
|
||||
match &self.inner {
|
||||
GroupInner::Read { file, path } => Ok(member_names(file, path)?.len()),
|
||||
GroupInner::Read(g) => Ok(g.member_names()?.len()),
|
||||
GroupInner::Write(state) => Ok(state.lock().unwrap().datasets.len()),
|
||||
}
|
||||
}
|
||||
@@ -179,7 +247,7 @@ impl PyGroup {
|
||||
#[getter]
|
||||
fn name(&self) -> String {
|
||||
match &self.inner {
|
||||
GroupInner::Read { path, .. } => node::name(path),
|
||||
GroupInner::Read(g) => node::name(&g.path),
|
||||
GroupInner::Write(state) => node::name(&state.lock().unwrap().name),
|
||||
}
|
||||
}
|
||||
@@ -235,7 +303,7 @@ impl PyGroup {
|
||||
#[getter]
|
||||
fn attrs(&self) -> PyResult<PyAttrs> {
|
||||
match &self.inner {
|
||||
GroupInner::Read { file, path } => PyAttrs::read(Arc::clone(file), path),
|
||||
GroupInner::Read(g) => g.attrs(),
|
||||
GroupInner::Write(state) => {
|
||||
let store = Arc::clone(&state.lock().unwrap().attrs);
|
||||
Ok(PyAttrs::from_write(store))
|
||||
@@ -245,9 +313,9 @@ impl PyGroup {
|
||||
|
||||
fn __repr__(&self) -> String {
|
||||
match &self.inner {
|
||||
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::Read(g) => {
|
||||
let n = g.member_names().map_or(0, |m| m.len());
|
||||
format!("<HDF5 group \"{}\" ({n} members)>", node::name(&g.path))
|
||||
}
|
||||
GroupInner::Write(state) => {
|
||||
let name = &state.lock().unwrap().name;
|
||||
@@ -258,7 +326,7 @@ impl PyGroup {
|
||||
|
||||
fn __contains__(&self, key: &str) -> PyResult<bool> {
|
||||
match &self.inner {
|
||||
GroupInner::Read { file, path } => Ok(contains(file, path, key)),
|
||||
GroupInner::Read(g) => Ok(g.contains(key)),
|
||||
GroupInner::Write(state) => {
|
||||
let guard = state.lock().unwrap();
|
||||
Ok(guard.datasets.iter().any(|d| d.name == key))
|
||||
@@ -299,15 +367,19 @@ mod tests {
|
||||
let finished = g.finish();
|
||||
b.add_group(finished);
|
||||
let bytes = b.finish().unwrap();
|
||||
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"));
|
||||
let file = Arc::new(clawhdf5_rs::File::from_bytes(bytes).unwrap());
|
||||
let root = file.superblock().root_group_address;
|
||||
let top = ReadGroup::new(Arc::clone(&file), String::new(), root);
|
||||
assert_eq!(top.member_names().unwrap(), ["alpha", "mid", "zeta"]);
|
||||
let (path, addr) = top.locate("mid").unwrap();
|
||||
assert_eq!(path, "mid");
|
||||
let mid = ReadGroup::new(Arc::clone(&file), path, addr);
|
||||
assert_eq!(mid.member_names().unwrap(), ["x"]);
|
||||
assert!(top.contains("mid/x"));
|
||||
assert!(mid.contains("/alpha"));
|
||||
assert!(mid.contains("x") && mid.contains("./x"));
|
||||
assert!(!top.contains("nope"));
|
||||
assert!(!mid.contains("alpha"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
Reference in New Issue
Block a user