use std::collections::{BTreeMap, BTreeSet};
use std::path::Path;
use crate::data_layout::DataLayout;
use crate::datatype::Datatype;
use crate::error::Error;
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::type_builders::{AttrValue, DatasetBuilder, FinishedGroup, GroupBuilder};
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 = File::open(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 root = file.root();
populate(&mut builder, &root, "", &drop, &mut matched)?;
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>,
) -> 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)?;
}
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)?;
sink.sink_add_group(gb.finish());
}
Ok(())
}
fn emit_dataset(db: &mut DatasetBuilder, ds: &Dataset, path: &str) -> 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)?;
check_pipeline(pipeline.as_ref(), path)?;
let dims = dataspace.dimensions.clone();
let n_elements: u64 = dims.iter().product();
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 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::String { .. }
| Datatype::BitField { .. }
| Datatype::Opaque { .. } => Ok(()),
Datatype::Time { .. } => bad("time"),
Datatype::VariableLength { .. } => bad("variable-length"),
Datatype::Reference { .. } => bad("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 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}")
}
}