use std::collections::{BTreeMap, BTreeSet, HashMap};
use std::path::Path;
use std::sync::Arc;
use crate::chunked_read::ChunkInfo;
use crate::chunked_write::{ChunkMeta, ChunkProvider};
use crate::convert::TryToUsize;
use crate::data_layout::DataLayout;
use crate::datatype::{Datatype, ReferenceType};
use crate::error::{Error, FormatError};
use crate::filter_pipeline::{
FILTER_DEFLATE, FILTER_FLETCHER32, FILTER_SCALEOFFSET, FILTER_SHUFFLE, FilterPipeline,
};
use crate::reader::{Dataset, File, Group};
use crate::scaleoffset::{self, ScaleOffset};
use crate::source::FileSource;
use crate::type_builders::{
AttrValue, DatasetBuilder, FinishedGroup, GroupBuilder, ObjectRefTarget, VlStringElement,
};
use crate::vl_data::{VlByteObject, VlenStringReadOptions, is_vlen_string_datatype};
use crate::writer::FileBuilder;
#[derive(Debug, Default, Clone)]
pub struct RepackOptions {
pub drop: Vec<String>,
}
impl RepackOptions {
pub fn new() -> Self {
Self::default()
}
pub fn drop_path(mut self, path: &str) -> Self {
self.drop.push(path.to_string());
self
}
}
pub fn repack<P: AsRef<Path>, Q: AsRef<Path>>(
src: P,
dst: Q,
options: &RepackOptions,
) -> Result<(), Error> {
let file = Arc::new(File::open_streaming(src)?);
let drop: BTreeSet<String> = options.drop.iter().map(|p| normalize(p)).collect();
let mut matched: BTreeSet<String> = BTreeSet::new();
let mut builder = FileBuilder::new();
if let Some(info) = file.file_space_info() {
builder
.with_file_space_strategy(info.strategy, false, info.threshold)
.with_file_space_page_size(info.page_size);
}
let addr_map = build_object_address_map(&file)?;
let root = file.root();
populate(
&mut builder,
&root,
"",
&drop,
&mut matched,
&file,
&addr_map,
)?;
if let Some(missing) = drop.iter().find(|d| !matched.contains(*d)) {
return Err(Error::RepackUnsupported(format!(
"drop path does not exist in the source: {missing}"
)));
}
builder.write(dst)?;
Ok(())
}
trait GroupSink {
fn sink_dataset(&mut self, name: &str) -> &mut DatasetBuilder;
fn sink_add_group(&mut self, group: FinishedGroup);
fn sink_set_attr(&mut self, name: &str, value: AttrValue);
}
impl GroupSink for FileBuilder {
fn sink_dataset(&mut self, name: &str) -> &mut DatasetBuilder {
self.create_dataset(name)
}
fn sink_add_group(&mut self, group: FinishedGroup) {
self.add_group(group);
}
fn sink_set_attr(&mut self, name: &str, value: AttrValue) {
self.set_attr(name, value);
}
}
impl GroupSink for GroupBuilder {
fn sink_dataset(&mut self, name: &str) -> &mut DatasetBuilder {
self.create_dataset(name)
}
fn sink_add_group(&mut self, group: FinishedGroup) {
self.add_group(group);
}
fn sink_set_attr(&mut self, name: &str, value: AttrValue) {
self.set_attr(name, value);
}
}
fn populate<S: GroupSink>(
sink: &mut S,
src: &Group,
path: &str,
drop: &BTreeSet<String>,
matched: &mut BTreeSet<String>,
file: &Arc<File>,
addr_map: &HashMap<u64, String>,
) -> Result<(), Error> {
let attrs = src.attrs()?;
let owner = if path.is_empty() {
"root group".to_string()
} else {
format!("group {path}")
};
check_attr_completeness(&attrs, &src.attr_names()?, &owner)?;
for (name, value) in sorted(attrs) {
sink.sink_set_attr(&name, value);
}
let mut dataset_names = src.datasets()?;
dataset_names.sort();
for name in dataset_names {
let child_path = join(path, &name);
if drop.contains(&child_path) {
matched.insert(child_path);
continue;
}
let ds = src.dataset(&name)?;
emit_dataset(
sink.sink_dataset(&name),
&ds,
&child_path,
file,
drop,
addr_map,
)?;
}
let mut group_names = src.groups()?;
group_names.sort();
for name in group_names {
let child_path = join(path, &name);
if drop.contains(&child_path) {
matched.insert(child_path);
continue;
}
let child = src.group(&name)?;
let mut gb = GroupBuilder::new(&name);
populate(&mut gb, &child, &child_path, drop, matched, file, addr_map)?;
sink.sink_add_group(gb.finish());
}
Ok(())
}
fn emit_dataset(
db: &mut DatasetBuilder,
ds: &Dataset,
path: &str,
file: &Arc<File>,
drop: &BTreeSet<String>,
addr_map: &HashMap<u64, String>,
) -> Result<(), Error> {
let datatype = ds.datatype()?;
let dataspace = ds.dataspace()?;
let layout = ds.data_layout()?;
let pipeline = ds.filter_pipeline();
check_datatype(&datatype, path)?;
check_layout(&layout, path)?;
let dims = dataspace.dimensions.clone();
let n_elements: u64 = dims.iter().product();
if is_vlen_string_datatype(&datatype) {
emit_vlen_string_dataset(db, ds, path, &datatype, &dims, &layout)?;
let attrs = ds.attrs()?;
check_attr_completeness(&attrs, &ds.attr_names()?, &format!("dataset {path}"))?;
for (name, value) in sorted(attrs) {
db.set_attr(&name, value);
}
return Ok(());
}
if is_nonstring_vlen(&datatype) {
emit_vlen_sequence_dataset(db, ds, path, &datatype, &dims, &layout)?;
let attrs = ds.attrs()?;
check_attr_completeness(&attrs, &ds.attr_names()?, &format!("dataset {path}"))?;
for (name, value) in sorted(attrs) {
db.set_attr(&name, value);
}
return Ok(());
}
if is_object_reference(&datatype) {
emit_object_reference_dataset(db, ds, path, &dims, &layout, file, drop, addr_map)?;
let attrs = ds.attrs()?;
check_attr_completeness(&attrs, &ds.attr_names()?, &format!("dataset {path}"))?;
for (name, value) in sorted(attrs) {
db.set_attr(&name, value);
}
return Ok(());
}
if let DataLayout::Chunked {
chunk_dimensions, ..
} = &layout
&& n_elements > 0
{
let rank = dims.len();
let chunk_dims: Vec<u64> = chunk_dimensions
.iter()
.take(rank)
.map(|&c| c as u64)
.collect();
if let Some(DenseChunkPlan { meta, grid_order }) =
try_plan_dense_chunks(ds, &dims, &chunk_dims)?
{
let maxshape = dataspace
.max_dimensions
.as_ref()
.filter(|ms| *ms != &dims)
.map(|ms| ms.as_slice());
let elem_size = datatype.type_size() as usize;
let provider = DatasetChunkProvider {
file: Arc::clone(file),
grid_order,
};
db.with_raw_chunks_lazy(
datatype,
&dims,
maxshape,
&chunk_dims,
elem_size,
ds.filter_pipeline_message_bytes(),
meta,
Box::new(provider),
);
let attrs = ds.attrs()?;
check_attr_completeness(&attrs, &ds.attr_names()?, &format!("dataset {path}"))?;
for (name, value) in sorted(attrs) {
db.set_attr(&name, value);
}
return Ok(());
}
}
check_pipeline(pipeline.as_ref(), path)?;
if n_elements == 0 {
db.with_dtype(datatype).with_shape(&dims);
} else {
let raw = ds.read_raw()?;
db.with_raw_data(datatype, raw, n_elements)
.with_shape(&dims);
}
if let Some(maxshape) = &dataspace.max_dimensions
&& maxshape != &dims
{
db.with_maxshape(maxshape);
}
if let DataLayout::Chunked {
chunk_dimensions, ..
} = &layout
{
let rank = dims.len();
let logical: Vec<u64> = chunk_dimensions
.iter()
.take(rank)
.map(|&c| c as u64)
.collect();
db.with_chunks(&logical);
}
if let Some(p) = &pipeline {
for f in &p.filters {
match f.filter_id {
FILTER_SHUFFLE => {
db.with_shuffle();
}
FILTER_FLETCHER32 => {
db.with_fletcher32();
}
FILTER_DEFLATE => {
db.with_deflate(f.client_data.first().copied().unwrap_or(6));
}
FILTER_SCALEOFFSET => {
if let Some(mode @ ScaleOffset::Integer(_)) =
scaleoffset::scale_offset_mode(&f.client_data)
{
db.with_scale_offset(mode);
} else {
unreachable!("check_pipeline rejected non-integer scale-offset");
}
}
_ => unreachable!("check_pipeline rejected unsupported filters"),
}
}
}
let attrs = ds.attrs()?;
check_attr_completeness(&attrs, &ds.attr_names()?, &format!("dataset {path}"))?;
for (name, value) in sorted(attrs) {
db.set_attr(&name, value);
}
Ok(())
}
fn emit_vlen_string_dataset(
db: &mut DatasetBuilder,
ds: &Dataset,
path: &str,
datatype: &Datatype,
dims: &[u64],
layout: &DataLayout,
) -> Result<(), Error> {
if matches!(layout, DataLayout::Chunked { .. }) {
return Err(Error::RepackUnsupported(format!(
"dataset {path}: chunked or filtered variable-length string datasets cannot be \
repacked (their element references live inside compressed chunks before the global \
heap addresses are known)"
)));
}
if let Some(maxshape) = &ds.dataspace()?.max_dimensions
&& maxshape != dims
{
return Err(Error::RepackUnsupported(format!(
"dataset {path}: resizable variable-length string datasets cannot be repacked"
)));
}
let objects = ds.read_vlen_string_bytes(VlenStringReadOptions::default())?;
let elements: Vec<VlStringElement> = objects
.into_iter()
.map(|o| match o {
VlByteObject::Null => VlStringElement::Null,
VlByteObject::Bytes(bytes) => VlStringElement::Bytes(bytes),
})
.collect();
db.with_vlen_string_elements(datatype.clone(), &elements)
.map_err(Error::Format)?;
db.with_shape(dims);
Ok(())
}
fn emit_vlen_sequence_dataset(
db: &mut DatasetBuilder,
ds: &Dataset,
path: &str,
datatype: &Datatype,
dims: &[u64],
layout: &DataLayout,
) -> Result<(), Error> {
if matches!(layout, DataLayout::Chunked { .. }) {
return Err(Error::RepackUnsupported(format!(
"dataset {path}: chunked or filtered non-string variable-length datasets cannot be \
repacked (their element references live inside compressed chunks before the global \
heap addresses are known)"
)));
}
if let Some(maxshape) = &ds.dataspace()?.max_dimensions
&& maxshape != dims
{
return Err(Error::RepackUnsupported(format!(
"dataset {path}: resizable non-string variable-length datasets cannot be repacked"
)));
}
let (objects, _element_size) = ds.read_vlen_sequence_bytes(VlenStringReadOptions::default())?;
let elements: Vec<VlStringElement> = objects
.into_iter()
.map(|o| match o {
VlByteObject::Null => VlStringElement::Null,
VlByteObject::Bytes(bytes) => VlStringElement::Bytes(bytes),
})
.collect();
db.with_vlen_sequence_elements(datatype.clone(), &elements)
.map_err(Error::Format)?;
db.with_shape(dims);
Ok(())
}
struct DatasetChunkProvider {
file: Arc<File>,
grid_order: Vec<ChunkInfo>,
}
impl ChunkProvider for DatasetChunkProvider {
fn chunk_bytes(&self, index: usize) -> Result<Vec<u8>, FormatError> {
let info = &self.grid_order[index];
self.file
.source()
.read_exact_at(info.address, info.chunk_size as usize)
}
}
struct DenseChunkPlan {
meta: Vec<ChunkMeta>,
grid_order: Vec<ChunkInfo>,
}
fn try_plan_dense_chunks(
ds: &Dataset,
dims: &[u64],
chunk_dims: &[u64],
) -> Result<Option<DenseChunkPlan>, Error> {
let Some(grid) = crate::chunked_read::plan_dense_grid(ds.raw_chunks()?, dims, chunk_dims)
else {
return Ok(None);
};
let grid_order = grid.grid_order;
let meta = grid_order
.iter()
.map(|info| ChunkMeta {
compressed_size: u64::from(info.chunk_size),
filter_mask: info.filter_mask,
})
.collect();
Ok(Some(DenseChunkPlan { meta, grid_order }))
}
fn check_attr_completeness(
decoded: &std::collections::HashMap<String, AttrValue>,
names: &[String],
owner: &str,
) -> Result<(), Error> {
for name in names {
if !decoded.contains_key(name) {
return Err(Error::RepackUnsupported(format!(
"{owner}: attribute {name:?} has a datatype that cannot be repacked faithfully yet"
)));
}
}
Ok(())
}
fn check_datatype(dt: &Datatype, path: &str) -> Result<(), Error> {
let bad = |what: &str| {
Err(Error::RepackUnsupported(format!(
"dataset {path}: {what} datatype cannot be repacked faithfully yet"
)))
};
match dt {
Datatype::FixedPoint { .. }
| Datatype::FloatingPoint { .. }
| Datatype::Time { .. }
| Datatype::String { .. }
| Datatype::BitField { .. }
| Datatype::Opaque { .. } => Ok(()),
Datatype::VariableLength { .. } if is_vlen_string_datatype(dt) => Ok(()),
Datatype::VariableLength { base_type, .. } => check_vlen_base_type(base_type, path),
Datatype::Reference {
ref_type: ReferenceType::Object,
size: 8,
} => Ok(()),
Datatype::Reference {
ref_type: ReferenceType::Object,
..
} => bad("non-8-byte object reference"),
Datatype::Reference {
ref_type: ReferenceType::DatasetRegion,
..
} => bad("dataset-region reference"),
Datatype::Compound { members, .. } => {
for m in members {
check_datatype(&m.datatype, path)?;
}
Ok(())
}
Datatype::Enumeration { base_type, .. } => check_datatype(base_type, path),
Datatype::Array { base_type, .. } => check_datatype(base_type, path),
}
}
fn is_nonstring_vlen(dt: &Datatype) -> bool {
matches!(dt, Datatype::VariableLength { .. }) && !is_vlen_string_datatype(dt)
}
fn check_vlen_base_type(dt: &Datatype, path: &str) -> Result<(), Error> {
let bad = |what: &str| {
Err(Error::RepackUnsupported(format!(
"dataset {path}: variable-length sequence of {what} cannot be repacked faithfully yet"
)))
};
match dt {
Datatype::FixedPoint { .. }
| Datatype::FloatingPoint { .. }
| Datatype::Time { .. }
| Datatype::String { .. }
| Datatype::BitField { .. }
| Datatype::Opaque { .. } => Ok(()),
Datatype::Reference { .. } => bad("references"),
Datatype::VariableLength { .. } => bad("variable-length elements"),
Datatype::Compound { members, .. } => {
for m in members {
check_vlen_base_type(&m.datatype, path)?;
}
Ok(())
}
Datatype::Enumeration { base_type, .. } => check_vlen_base_type(base_type, path),
Datatype::Array { base_type, .. } => check_vlen_base_type(base_type, path),
}
}
fn is_object_reference(dt: &Datatype) -> bool {
matches!(
dt,
Datatype::Reference {
ref_type: ReferenceType::Object,
..
}
)
}
fn is_dropped(path: &str, drop: &BTreeSet<String>) -> bool {
if drop.contains(path) {
return true;
}
let mut p = path;
while let Some(idx) = p.rfind('/') {
p = &p[..idx];
if drop.contains(p) {
return true;
}
}
false
}
fn build_object_address_map(file: &File) -> Result<HashMap<u64, String>, Error> {
let mut map = HashMap::new();
let root = file.root();
map.insert(root.header_address(), String::new());
collect_addresses(&root, "", &mut map)?;
Ok(map)
}
fn collect_addresses(
group: &Group,
prefix: &str,
map: &mut HashMap<u64, String>,
) -> Result<(), Error> {
for name in group.datasets()? {
let ds = group.dataset(&name)?;
map.insert(ds.header_address(), join(prefix, &name));
}
for name in group.groups()? {
let child = group.group(&name)?;
let child_path = join(prefix, &name);
map.insert(child.header_address(), child_path.clone());
collect_addresses(&child, &child_path, map)?;
}
Ok(())
}
#[allow(clippy::too_many_arguments)]
fn emit_object_reference_dataset(
db: &mut DatasetBuilder,
ds: &Dataset,
path: &str,
dims: &[u64],
layout: &DataLayout,
file: &Arc<File>,
drop: &BTreeSet<String>,
addr_map: &HashMap<u64, String>,
) -> Result<(), Error> {
if matches!(layout, DataLayout::Chunked { .. }) {
return Err(Error::RepackUnsupported(format!(
"dataset {path}: chunked or filtered object-reference datasets cannot be repacked \
(their addresses live inside compressed chunks and would need rewriting in place)"
)));
}
if let Some(maxshape) = &ds.dataspace()?.max_dimensions
&& maxshape != dims
{
return Err(Error::RepackUnsupported(format!(
"dataset {path}: resizable object-reference datasets cannot be repacked"
)));
}
if file.base_address() != 0 {
return Err(Error::RepackUnsupported(format!(
"dataset {path}: object references in a file with a non-zero base address (userblock) \
cannot be repacked yet"
)));
}
let n_elements: usize = dims.iter().product::<u64>().to_usize()?;
let targets = if n_elements == 0 {
Vec::new()
} else {
let raw = ds.read_raw()?;
let needed = n_elements
.checked_mul(8)
.ok_or(FormatError::OffsetOverflow {
offset: n_elements as u64,
length: 8,
})?;
if raw.len() < needed {
return Err(FormatError::UnexpectedEof {
expected: needed,
available: raw.len(),
}
.into());
}
let mut targets = Vec::with_capacity(n_elements);
for chunk in raw[..needed].chunks_exact(8) {
let v = u64::from_le_bytes(chunk.try_into().expect("chunks_exact(8) yields 8 bytes"));
if v == 0 || v == u64::MAX {
targets.push(ObjectRefTarget::Raw(v));
continue;
}
match addr_map.get(&v) {
Some(target_path) if is_dropped(target_path, drop) => {
return Err(Error::RepackUnsupported(format!(
"dataset {path}: object reference to dropped object {target_path:?} \
cannot be repacked"
)));
}
Some(target_path) => targets.push(ObjectRefTarget::Path(target_path.clone())),
None => {
return Err(Error::RepackUnsupported(format!(
"dataset {path}: object reference to address {v:#x} resolves to no \
hard-linked object in the source (dangling, or a named-datatype / \
region target not supported yet)"
)));
}
}
}
targets
};
db.with_object_references(targets);
db.with_shape(dims);
Ok(())
}
fn check_layout(layout: &DataLayout, path: &str) -> Result<(), Error> {
match layout {
DataLayout::Compact { .. } | DataLayout::Contiguous { .. } | DataLayout::Chunked { .. } => {
Ok(())
}
DataLayout::Virtual { .. } => Err(Error::RepackUnsupported(format!(
"dataset {path}: virtual data layout cannot be repacked"
))),
}
}
fn check_pipeline(pipeline: Option<&FilterPipeline>, path: &str) -> Result<(), Error> {
let Some(p) = pipeline else {
return Ok(());
};
for f in &p.filters {
match f.filter_id {
FILTER_DEFLATE | FILTER_SHUFFLE | FILTER_FLETCHER32 => {}
FILTER_SCALEOFFSET => match scaleoffset::scale_offset_mode(&f.client_data) {
Some(ScaleOffset::Integer(_)) => {}
_ => {
return Err(Error::RepackUnsupported(format!(
"dataset {path}: only lossless integer scale-offset with an undefined fill value can be repacked faithfully"
)));
}
},
other => {
return Err(Error::RepackUnsupported(format!(
"dataset {path}: filter id {other} cannot be repacked yet"
)));
}
}
}
Ok(())
}
fn sorted(attrs: std::collections::HashMap<String, AttrValue>) -> Vec<(String, AttrValue)> {
attrs
.into_iter()
.collect::<BTreeMap<_, _>>()
.into_iter()
.collect()
}
fn normalize(path: &str) -> String {
path.split('/')
.filter(|c| !c.is_empty())
.collect::<Vec<_>>()
.join("/")
}
fn join(parent: &str, name: &str) -> String {
if parent.is_empty() {
name.to_string()
} else {
format!("{parent}/{name}")
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn repack_preserves_big_endian_time_dataset() {
use crate::datatype::{Datatype, DatatypeByteOrder};
use crate::reader::File;
use crate::writer::FileBuilder;
let dir = std::env::temp_dir();
let src = dir.join("hdf5_pure_repack_time_src.h5");
let dst = dir.join("hdf5_pure_repack_time_dst.h5");
let dt = Datatype::Time {
size: 4,
byte_order: DatatypeByteOrder::BigEndian,
bit_precision: 32,
};
let raw: Vec<u8> = vec![
0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x02, 0x00, 0x00, 0x00, 0x03,
];
{
let mut b = FileBuilder::new();
b.create_dataset("t")
.with_raw_data(dt.clone(), raw.clone(), 3)
.with_shape(&[3]);
b.write(&src).unwrap();
}
repack(&src, &dst, &RepackOptions::new()).unwrap();
let f = File::open(&dst).unwrap();
let ds = f.dataset("t").unwrap();
assert_eq!(
ds.datatype().unwrap(),
dt,
"time datatype incl. byte order must survive repack"
);
assert_eq!(
ds.read_raw().unwrap(),
raw,
"time element bytes must be preserved"
);
std::fs::remove_file(&src).ok();
std::fs::remove_file(&dst).ok();
}
#[test]
fn is_dropped_matches_self_and_ancestors() {
let drop: BTreeSet<String> = ["g/old", "lone"].iter().map(|s| s.to_string()).collect();
assert!(is_dropped("lone", &drop));
assert!(is_dropped("g/old", &drop));
assert!(is_dropped("g/old/child", &drop));
assert!(is_dropped("g/old/a/b", &drop));
assert!(!is_dropped("g", &drop));
assert!(!is_dropped("g/older", &drop));
assert!(!is_dropped("lonely", &drop));
assert!(!is_dropped("other/old", &drop));
}
}