use std::path::{Path, PathBuf};
use std::sync::{Mutex, OnceLock};
use std::time::SystemTime;
use hdf5::types::{FloatSize, IntSize, TypeDescriptor};
use crate::array::ByteOrder;
use crate::array::{ArrayBuffer, ArrayData, ArrayMetadata, ArrayValues};
use crate::backend::ArrayBackend;
use crate::error::{Error, ExternalDataReason, Result};
pub struct Hdf5Backend;
const HDF5_SIGNATURE: [u8; 8] = [0x89, b'H', b'D', b'F', b'\r', b'\n', 0x1a, b'\n'];
fn has_hdf5_signature(file: &Path) -> bool {
use std::io::Read;
let mut head = [0u8; 8];
std::fs::File::open(file)
.and_then(|mut handle| handle.read_exact(&mut head))
.is_ok()
&& head == HDF5_SIGNATURE
}
impl ArrayBackend for Hdf5Backend {
fn name(&self) -> &'static str {
"hdf5"
}
fn handles(&self, mime: Option<&str>, file: &Path) -> bool {
if let Some(mime) = mime {
return mime.contains("hdf5") || mime.contains("hdf");
}
if file
.extension()
.is_some_and(|e| e.eq_ignore_ascii_case("h5") || e.eq_ignore_ascii_case("hdf5"))
{
return true;
}
has_hdf5_signature(file)
}
fn metadata(&self, file: &Path, path_in_file: &str) -> Result<ArrayMetadata> {
let handle = open(file)?;
let dataset = dataset(&handle, file, path_in_file)?;
Ok(ArrayMetadata::new(
dataset.shape().iter().map(|d| *d as i64).collect(),
element_type(&dataset, file, path_in_file)?,
))
}
fn byte_order(&self, file: &Path, path_in_file: &str) -> Result<Option<ByteOrder>> {
let handle = open(file)?;
let dataset = dataset(&handle, file, path_in_file)?;
let dtype = dataset
.dtype()
.map_err(|e| external(file, path_in_file, ExternalDataReason::Backend(e.to_string())))?;
if dtype.size() <= 1 {
return Ok(None);
}
Ok(match dtype.byte_order() {
hdf5::datatype::ByteOrder::LittleEndian => Some(ByteOrder::LittleEndian),
hdf5::datatype::ByteOrder::BigEndian => Some(ByteOrder::BigEndian),
_ => None,
})
}
fn read(&self, file: &Path, path_in_file: &str, start: &[i64], count: &[i64]) -> Result<ArrayData> {
let handle = open(file)?;
let dataset = dataset(&handle, file, path_in_file)?;
let shape: Vec<i64> = dataset.shape().iter().map(|d| *d as i64).collect();
let selection = Selection::of(&shape, start, count, file, path_in_file)?;
let descriptor = descriptor(&dataset, file, path_in_file)?;
let wanted: usize = selection.count.iter().product::<i64>().max(0) as usize;
let values = match &descriptor {
TypeDescriptor::Float(FloatSize::U8) => {
ArrayValues::F64(fill(&dataset, &selection, wanted, file, path_in_file)?)
}
TypeDescriptor::Float(FloatSize::U4) => {
ArrayValues::F32(fill(&dataset, &selection, wanted, file, path_in_file)?)
}
TypeDescriptor::Integer(IntSize::U8) => {
ArrayValues::I64(fill(&dataset, &selection, wanted, file, path_in_file)?)
}
TypeDescriptor::Integer(IntSize::U4) => {
ArrayValues::I32(fill(&dataset, &selection, wanted, file, path_in_file)?)
}
TypeDescriptor::Integer(IntSize::U2) => {
ArrayValues::I16(fill(&dataset, &selection, wanted, file, path_in_file)?)
}
TypeDescriptor::Integer(IntSize::U1) => {
ArrayValues::I8(fill(&dataset, &selection, wanted, file, path_in_file)?)
}
TypeDescriptor::Unsigned(IntSize::U8) => {
ArrayValues::U64(fill(&dataset, &selection, wanted, file, path_in_file)?)
}
TypeDescriptor::Unsigned(IntSize::U4) => {
ArrayValues::U32(fill(&dataset, &selection, wanted, file, path_in_file)?)
}
TypeDescriptor::Unsigned(IntSize::U2) => {
ArrayValues::U16(fill(&dataset, &selection, wanted, file, path_in_file)?)
}
TypeDescriptor::Unsigned(IntSize::U1) | TypeDescriptor::Boolean => {
ArrayValues::U8(fill(&dataset, &selection, wanted, file, path_in_file)?)
}
other => {
return Err(external(
file,
path_in_file,
ExternalDataReason::Backend(format!(
"the dataset holds `{other:?}`, which this backend does not decode"
)),
));
}
};
Ok(ArrayData {
metadata: ArrayMetadata::new(selection.count.clone(), name_of(&descriptor)),
values,
})
}
fn read_into(
&self,
file: &Path,
path_in_file: &str,
start: &[i64],
count: &[i64],
out: &mut ArrayBuffer<'_>,
) -> Result<ArrayMetadata> {
let handle = open(file)?;
let dataset = dataset(&handle, file, path_in_file)?;
let shape: Vec<i64> = dataset.shape().iter().map(|d| *d as i64).collect();
let selection = Selection::of(&shape, start, count, file, path_in_file)?;
let descriptor = descriptor(&dataset, file, path_in_file)?;
let element_type = name_of(&descriptor);
let wanted: usize = selection.count.iter().product::<i64>().max(0) as usize;
if out.len() != wanted {
return Err(external(
file,
path_in_file,
ExternalDataReason::Backend(format!(
"the buffer holds {} element(s), the selection is {wanted}",
out.len()
)),
));
}
if out.element_type() != element_type {
return Err(external(
file,
path_in_file,
ExternalDataReason::Backend(format!(
"the buffer is `{}`, the dataset is `{element_type}`",
out.element_type()
)),
));
}
match out {
ArrayBuffer::F64(b) => into(&dataset, &selection, b, file, path_in_file)?,
ArrayBuffer::F32(b) => into(&dataset, &selection, b, file, path_in_file)?,
ArrayBuffer::I64(b) => into(&dataset, &selection, b, file, path_in_file)?,
ArrayBuffer::I32(b) => into(&dataset, &selection, b, file, path_in_file)?,
ArrayBuffer::U64(b) => into(&dataset, &selection, b, file, path_in_file)?,
ArrayBuffer::U32(b) => into(&dataset, &selection, b, file, path_in_file)?,
ArrayBuffer::U8(b) => into(&dataset, &selection, b, file, path_in_file)?,
}
Ok(ArrayMetadata::new(selection.count.clone(), element_type))
}
fn file_uuid(&self, file: &Path) -> Result<Option<String>> {
let handle = open(file)?;
let Ok(attribute) = handle.attr("uuid") else {
return Ok(None);
};
let value = attribute
.read_scalar::<hdf5::types::FixedAscii<36>>()
.map(|v| v.to_string())
.or_else(|_| {
attribute
.read_scalar::<hdf5::types::VarLenUnicode>()
.map(|v| v.to_string())
});
Ok(value.ok().map(|v| v.trim().to_string()).filter(|v| !v.is_empty()))
}
fn writes(&self) -> bool {
true
}
fn write(&self, file: &Path, path_in_file: &str, data: &ArrayData) -> Result<()> {
let handle = open_for_writing(file)?;
write_into(&handle, file, path_in_file, data)?;
handle.flush().map_err(|e| {
external(
file,
path_in_file,
ExternalDataReason::Backend(format!("flushing the file: {e}")),
)
})
}
fn write_many(&self, file: &Path, items: &[(&str, &ArrayData)]) -> Vec<Result<()>> {
let every = |message: String| -> Vec<Result<()>> {
items
.iter()
.map(|(path, _)| Err(external(file, path, ExternalDataReason::Backend(message.clone()))))
.collect()
};
let handle = match open_for_writing(file) {
Ok(handle) => handle,
Err(why) => return every(why.to_string()),
};
let written: Vec<Result<()>> = items
.iter()
.map(|(path, data)| write_into(&handle, file, path, data))
.collect();
if let Err(e) = handle.flush() {
return every(format!("flushing the file: {e}"));
}
written
}
fn set_file_uuid(&self, file: &Path, uuid: &str) -> Result<()> {
let handle = open_for_writing(file)?;
let fail = |message: String| Error::ExternalData {
location: file.display().to_string(),
reason: ExternalDataReason::Backend(message),
};
let text = uuid.trim().to_ascii_lowercase();
if !energyml_core::identifiers::is_uuid(&text) {
return Err(fail(format!(
"`{uuid}` is not a uuid; §3.2.3 asks for 36 canonical lower-case characters"
)));
}
let value = hdf5::types::FixedAscii::<36>::from_ascii(text.as_bytes())
.map_err(|_| fail(format!("`{uuid}` is not 36 ASCII characters")))?;
if handle.attr("uuid").is_ok() {
handle
.delete_attr("uuid")
.map_err(|e| fail(format!("replacing the existing uuid attribute: {e}")))?;
}
handle
.new_attr::<hdf5::types::FixedAscii<36>>()
.create("uuid")
.map_err(|e| fail(format!("creating the uuid attribute: {e}")))?
.write_scalar(&value)
.map_err(|e| fail(format!("writing the uuid attribute: {e}")))?;
handle
.flush()
.map_err(|e| fail(format!("flushing the file: {e}")))?;
Ok(())
}
}
struct Selection {
start: Vec<i64>,
count: Vec<i64>,
whole: bool,
}
impl Selection {
fn of(shape: &[i64], start: &[i64], count: &[i64], file: &Path, path_in_file: &str) -> Result<Self> {
let start: Vec<i64> = if start.is_empty() {
vec![0; shape.len()]
} else {
start.to_vec()
};
let count: Vec<i64> = if count.is_empty() {
shape.to_vec()
} else {
count.to_vec()
};
let bad = |message: String| external(file, path_in_file, ExternalDataReason::Backend(message));
if start.len() != shape.len() || count.len() != shape.len() {
if start.iter().all(|s| *s == 0)
&& (count.iter().sum::<i64>() >= shape.iter().product::<i64>()
|| (count.len() == 1 && count[0] == shape[0]))
{
return Ok(Self {
start: vec![0; shape.len()],
count: shape.to_vec(),
whole: true,
});
}
return Err(bad(format!(
"the object asks for {} dimension(s) but the dataset has {}",
count.len(),
shape.len()
)));
}
for (axis, ((s, c), extent)) in start.iter().zip(&count).zip(shape).enumerate() {
if *s < 0 || *c < 0 || s + c > *extent {
return Err(bad(format!(
"dimension {axis}: the object asks for {c} element(s) from {s}, \
the dataset holds {extent}"
)));
}
}
let whole = start.iter().all(|s| *s == 0) && count == shape;
Ok(Self { start, count, whole })
}
}
const DEFAULT_OPEN_LIMIT: usize = 8;
struct Cached {
path: PathBuf,
handle: hdf5::File,
stamp: (Option<SystemTime>, u64),
}
struct OpenFiles {
entries: Vec<Cached>,
limit: usize,
opens: u64,
}
fn cache() -> &'static Mutex<OpenFiles> {
static FILES: OnceLock<Mutex<OpenFiles>> = OnceLock::new();
FILES.get_or_init(|| {
Mutex::new(OpenFiles {
entries: Vec::new(),
limit: DEFAULT_OPEN_LIMIT,
opens: 0,
})
})
}
fn stamp_of(file: &Path) -> (Option<SystemTime>, u64) {
match std::fs::metadata(file) {
Ok(data) => (data.modified().ok(), data.len()),
Err(_) => (None, 0),
}
}
pub fn open_file_limit() -> usize {
cache()
.lock()
.map(|files| files.limit)
.unwrap_or(DEFAULT_OPEN_LIMIT)
}
pub fn set_open_file_limit(limit: usize) {
if let Ok(mut files) = cache().lock() {
files.limit = limit;
while files.entries.len() > limit {
files.entries.remove(0);
}
}
}
pub fn real_opens() -> u64 {
cache().lock().map(|files| files.opens).unwrap_or(0)
}
pub fn close_cached_files() {
if let Ok(mut files) = cache().lock() {
files.entries.clear();
}
}
fn evict(file: &Path) {
if let Ok(mut files) = cache().lock() {
files.entries.retain(|entry| entry.path != file);
}
}
fn open(file: &Path) -> Result<hdf5::File> {
let stamp = stamp_of(file);
if let Ok(mut files) = cache().lock()
&& let Some(at) = files.entries.iter().position(|entry| entry.path == file)
{
if files.entries[at].stamp == stamp {
let entry = files.entries.remove(at);
let handle = entry.handle.clone();
files.entries.push(entry);
return Ok(handle);
}
files.entries.remove(at);
}
if let Ok(mut files) = cache().lock() {
files.opens += 1;
}
let handle = hdf5::File::open(file).map_err(|e| Error::ExternalData {
location: file.display().to_string(),
reason: ExternalDataReason::Backend(e.to_string()),
})?;
if let Ok(mut files) = cache().lock()
&& files.limit > 0
{
files.entries.push(Cached {
path: file.to_path_buf(),
handle: handle.clone(),
stamp,
});
while files.entries.len() > files.limit {
files.entries.remove(0);
}
}
Ok(handle)
}
fn write_into(handle: &hdf5::File, file: &Path, path_in_file: &str, data: &ArrayData) -> Result<()> {
let bad = |message: String| external(file, path_in_file, ExternalDataReason::Backend(message));
let shape: Vec<usize> = data
.metadata
.dimensions
.iter()
.map(|d| (*d).max(0) as usize)
.collect();
let declared: usize = shape.iter().product();
if shape.is_empty() {
return Err(bad("the array declares no dimensions".to_string()));
}
if declared != data.values.len() {
return Err(bad(format!(
"the shape {shape:?} describes {declared} element(s), the values hold {}",
data.values.len()
)));
}
if let Some((groups, _)) = path_in_file.rsplit_once('/') {
let mut current = String::new();
for name in groups.split('/').filter(|s| !s.is_empty()) {
current.push('/');
current.push_str(name);
if handle.group(¤t).is_err() {
handle
.create_group_builder()
.obj_track_times(false)
.create(current.as_str())
.map_err(|e| bad(format!("creating the group `{current}`: {e}")))?;
}
}
}
if handle.dataset(path_in_file).is_ok() {
handle
.unlink(path_in_file)
.map_err(|e| bad(format!("replacing the existing dataset: {e}")))?;
}
macro_rules! put {
($values:expr, $ty:ty) => {{
let builder = handle.new_dataset::<$ty>().shape(shape.clone());
let dataset = builder
.create(path_in_file)
.map_err(|e| bad(format!("creating the dataset: {e}")))?;
dataset
.write_raw($values)
.map_err(|e| bad(format!("writing the values: {e}")))?;
}};
}
match &data.values {
ArrayValues::F64(v) => put!(v, f64),
ArrayValues::F32(v) => put!(v, f32),
ArrayValues::I64(v) => put!(v, i64),
ArrayValues::I32(v) => put!(v, i32),
ArrayValues::I16(v) => put!(v, i16),
ArrayValues::I8(v) => put!(v, i8),
ArrayValues::U64(v) => put!(v, u64),
ArrayValues::U32(v) => put!(v, u32),
ArrayValues::U16(v) => put!(v, u16),
ArrayValues::U8(v) => put!(v, u8),
}
Ok(())
}
fn open_for_writing(file: &Path) -> Result<hdf5::File> {
evict(file);
let fail = |e: hdf5::Error| Error::ExternalData {
location: file.display().to_string(),
reason: ExternalDataReason::Backend(e.to_string()),
};
if file.exists() {
return hdf5::File::open_rw(file).map_err(fail);
}
if let Some(folder) = file.parent()
&& !folder.as_os_str().is_empty()
&& !folder.exists()
{
std::fs::create_dir_all(folder).map_err(|e| Error::ExternalData {
location: folder.display().to_string(),
reason: ExternalDataReason::Backend(e.to_string()),
})?;
}
hdf5::File::with_options()
.with_fcpl(|plist| plist.obj_track_times(false))
.create(file)
.map_err(fail)
}
fn dataset(handle: &hdf5::File, file: &Path, path_in_file: &str) -> Result<hdf5::Dataset> {
handle.dataset(path_in_file).map_err(|_| Error::ExternalData {
location: file.display().to_string(),
reason: ExternalDataReason::PathMissing(path_in_file.to_string()),
})
}
fn descriptor(dataset: &hdf5::Dataset, file: &Path, path_in_file: &str) -> Result<TypeDescriptor> {
dataset
.dtype()
.and_then(|dtype| dtype.to_descriptor())
.map_err(|e| external(file, path_in_file, ExternalDataReason::Backend(e.to_string())))
}
fn element_type(dataset: &hdf5::Dataset, file: &Path, path_in_file: &str) -> Result<String> {
Ok(name_of(&descriptor(dataset, file, path_in_file)?))
}
fn fill<T: hdf5::H5Type + Copy + Default>(
dataset: &hdf5::Dataset,
selection: &Selection,
wanted: usize,
file: &Path,
path_in_file: &str,
) -> Result<Vec<T>> {
let mut buffer = vec![T::default(); wanted];
into(dataset, selection, &mut buffer, file, path_in_file)?;
Ok(buffer)
}
fn into<T: hdf5::H5Type + Copy>(
dataset: &hdf5::Dataset,
selection: &Selection,
buffer: &mut [T],
file: &Path,
path_in_file: &str,
) -> Result<()> {
let fail = |e: hdf5::Error| external(file, path_in_file, ExternalDataReason::Backend(e.to_string()));
if selection.whole {
let read = dataset.read_into_raw(buffer).map_err(fail)?;
if read != buffer.len() {
return Err(external(
file,
path_in_file,
ExternalDataReason::Backend(format!(
"read {read} element(s) into a buffer of {}",
buffer.len()
)),
));
}
return Ok(());
}
let hyperslab: Vec<hdf5::SliceOrIndex> = selection
.start
.iter()
.zip(&selection.count)
.map(|(start, count)| hdf5::SliceOrIndex::SliceCount {
start: *start as usize,
step: 1,
block: 1,
count: *count as usize,
})
.collect();
let array = dataset
.read_slice::<T, _, ndarray::IxDyn>(hdf5::Hyperslab::from(hyperslab))
.map_err(fail)?;
for (slot, value) in buffer.iter_mut().zip(array.iter()) {
*slot = *value;
}
Ok(())
}
fn name_of(descriptor: &TypeDescriptor) -> String {
match descriptor {
TypeDescriptor::Float(FloatSize::U8) => "double".to_string(),
TypeDescriptor::Float(FloatSize::U4) => "float".to_string(),
TypeDescriptor::Integer(size) => format!("int{}", *size as usize * 8),
TypeDescriptor::Unsigned(size) => format!("uint{}", *size as usize * 8),
TypeDescriptor::Boolean => "bool".to_string(),
other => format!("{other:?}"),
}
}
fn external(file: &Path, path_in_file: &str, reason: ExternalDataReason) -> Error {
Error::ExternalData {
location: format!("{}:{path_in_file}", file.display()),
reason,
}
}