Compare commits
3
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
55959b4920 | ||
|
|
b70d594c4f | ||
|
|
b9898c2a9c |
@@ -0,0 +1,26 @@
|
|||||||
|
name: CI
|
||||||
|
on:
|
||||||
|
push:
|
||||||
|
branches: [main]
|
||||||
|
pull_request:
|
||||||
|
branches: [main]
|
||||||
|
jobs:
|
||||||
|
test:
|
||||||
|
runs-on: ubuntu-latest
|
||||||
|
container: rust:latest
|
||||||
|
steps:
|
||||||
|
- uses: actions/checkout@v4
|
||||||
|
- name: Cache cargo registry/target
|
||||||
|
uses: actions/cache@v4
|
||||||
|
with:
|
||||||
|
path: |
|
||||||
|
~/.cargo/registry
|
||||||
|
~/.cargo/git
|
||||||
|
target
|
||||||
|
key: ${{ runner.os }}-cargo-${{ hashFiles('**/Cargo.lock') }}
|
||||||
|
- name: Install rustfmt & clippy components
|
||||||
|
run: rustup component add rustfmt clippy
|
||||||
|
- name: Install thumbv7em-none-eabihf target
|
||||||
|
run: rustup target add thumbv7em-none-eabihf
|
||||||
|
- name: Run CI script
|
||||||
|
run: bash scripts/ci-test.sh
|
||||||
@@ -13,42 +13,44 @@ use std::arch::x86_64::*;
|
|||||||
/// Caller must verify is_x86_feature_detected!("avx512f").
|
/// Caller must verify is_x86_feature_detected!("avx512f").
|
||||||
// SAFETY: Caller must have verified avx512f via is_x86_feature_detected!.
|
// SAFETY: Caller must have verified avx512f via is_x86_feature_detected!.
|
||||||
#[target_feature(enable = "avx512f")]
|
#[target_feature(enable = "avx512f")]
|
||||||
pub unsafe fn dot_product(a: &[f32], b: &[f32]) -> f32 { unsafe {
|
pub unsafe fn dot_product(a: &[f32], b: &[f32]) -> f32 {
|
||||||
assert_eq!(a.len(), b.len());
|
unsafe {
|
||||||
let len = a.len();
|
assert_eq!(a.len(), b.len());
|
||||||
let mut i = 0;
|
let len = a.len();
|
||||||
let mut acc0 = _mm512_setzero_ps();
|
let mut i = 0;
|
||||||
let mut acc1 = _mm512_setzero_ps();
|
let mut acc0 = _mm512_setzero_ps();
|
||||||
|
let mut acc1 = _mm512_setzero_ps();
|
||||||
|
|
||||||
// Process 32 elements per iteration (2x16 unrolled)
|
// Process 32 elements per iteration (2x16 unrolled)
|
||||||
while i + 32 <= len {
|
while i + 32 <= len {
|
||||||
let va0 = _mm512_loadu_ps(a.as_ptr().add(i));
|
let va0 = _mm512_loadu_ps(a.as_ptr().add(i));
|
||||||
let vb0 = _mm512_loadu_ps(b.as_ptr().add(i));
|
let vb0 = _mm512_loadu_ps(b.as_ptr().add(i));
|
||||||
acc0 = _mm512_fmadd_ps(va0, vb0, acc0);
|
acc0 = _mm512_fmadd_ps(va0, vb0, acc0);
|
||||||
|
|
||||||
let va1 = _mm512_loadu_ps(a.as_ptr().add(i + 16));
|
let va1 = _mm512_loadu_ps(a.as_ptr().add(i + 16));
|
||||||
let vb1 = _mm512_loadu_ps(b.as_ptr().add(i + 16));
|
let vb1 = _mm512_loadu_ps(b.as_ptr().add(i + 16));
|
||||||
acc1 = _mm512_fmadd_ps(va1, vb1, acc1);
|
acc1 = _mm512_fmadd_ps(va1, vb1, acc1);
|
||||||
|
|
||||||
i += 32;
|
i += 32;
|
||||||
|
}
|
||||||
|
|
||||||
|
if i + 16 <= len {
|
||||||
|
let va = _mm512_loadu_ps(a.as_ptr().add(i));
|
||||||
|
let vb = _mm512_loadu_ps(b.as_ptr().add(i));
|
||||||
|
acc0 = _mm512_fmadd_ps(va, vb, acc0);
|
||||||
|
i += 16;
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut sum = _mm512_reduce_add_ps(_mm512_add_ps(acc0, acc1));
|
||||||
|
|
||||||
|
while i < len {
|
||||||
|
sum += a[i] * b[i];
|
||||||
|
i += 1;
|
||||||
|
}
|
||||||
|
|
||||||
|
sum
|
||||||
}
|
}
|
||||||
|
}
|
||||||
if i + 16 <= len {
|
|
||||||
let va = _mm512_loadu_ps(a.as_ptr().add(i));
|
|
||||||
let vb = _mm512_loadu_ps(b.as_ptr().add(i));
|
|
||||||
acc0 = _mm512_fmadd_ps(va, vb, acc0);
|
|
||||||
i += 16;
|
|
||||||
}
|
|
||||||
|
|
||||||
let mut sum = _mm512_reduce_add_ps(_mm512_add_ps(acc0, acc1));
|
|
||||||
|
|
||||||
while i < len {
|
|
||||||
sum += a[i] * b[i];
|
|
||||||
i += 1;
|
|
||||||
}
|
|
||||||
|
|
||||||
sum
|
|
||||||
}}
|
|
||||||
|
|
||||||
/// AVX-512 cosine similarity — fused single pass.
|
/// AVX-512 cosine similarity — fused single pass.
|
||||||
///
|
///
|
||||||
@@ -56,38 +58,40 @@ pub unsafe fn dot_product(a: &[f32], b: &[f32]) -> f32 { unsafe {
|
|||||||
/// Caller must verify is_x86_feature_detected!("avx512f").
|
/// Caller must verify is_x86_feature_detected!("avx512f").
|
||||||
// SAFETY: Caller must have verified avx512f via is_x86_feature_detected!.
|
// SAFETY: Caller must have verified avx512f via is_x86_feature_detected!.
|
||||||
#[target_feature(enable = "avx512f")]
|
#[target_feature(enable = "avx512f")]
|
||||||
pub unsafe fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 { unsafe {
|
pub unsafe fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
|
||||||
assert_eq!(a.len(), b.len());
|
unsafe {
|
||||||
let len = a.len();
|
assert_eq!(a.len(), b.len());
|
||||||
let mut i = 0;
|
let len = a.len();
|
||||||
|
let mut i = 0;
|
||||||
|
|
||||||
let mut dot_acc = _mm512_setzero_ps();
|
let mut dot_acc = _mm512_setzero_ps();
|
||||||
let mut norm_a_acc = _mm512_setzero_ps();
|
let mut norm_a_acc = _mm512_setzero_ps();
|
||||||
let mut norm_b_acc = _mm512_setzero_ps();
|
let mut norm_b_acc = _mm512_setzero_ps();
|
||||||
|
|
||||||
while i + 16 <= len {
|
while i + 16 <= len {
|
||||||
let va = _mm512_loadu_ps(a.as_ptr().add(i));
|
let va = _mm512_loadu_ps(a.as_ptr().add(i));
|
||||||
let vb = _mm512_loadu_ps(b.as_ptr().add(i));
|
let vb = _mm512_loadu_ps(b.as_ptr().add(i));
|
||||||
dot_acc = _mm512_fmadd_ps(va, vb, dot_acc);
|
dot_acc = _mm512_fmadd_ps(va, vb, dot_acc);
|
||||||
norm_a_acc = _mm512_fmadd_ps(va, va, norm_a_acc);
|
norm_a_acc = _mm512_fmadd_ps(va, va, norm_a_acc);
|
||||||
norm_b_acc = _mm512_fmadd_ps(vb, vb, norm_b_acc);
|
norm_b_acc = _mm512_fmadd_ps(vb, vb, norm_b_acc);
|
||||||
i += 16;
|
i += 16;
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut dot = _mm512_reduce_add_ps(dot_acc);
|
||||||
|
let mut norm_a = _mm512_reduce_add_ps(norm_a_acc);
|
||||||
|
let mut norm_b = _mm512_reduce_add_ps(norm_b_acc);
|
||||||
|
|
||||||
|
while i < len {
|
||||||
|
dot += a[i] * b[i];
|
||||||
|
norm_a += a[i] * a[i];
|
||||||
|
norm_b += b[i] * b[i];
|
||||||
|
i += 1;
|
||||||
|
}
|
||||||
|
|
||||||
|
let denom = (norm_a * norm_b).sqrt();
|
||||||
|
if denom == 0.0 { 0.0 } else { dot / denom }
|
||||||
}
|
}
|
||||||
|
}
|
||||||
let mut dot = _mm512_reduce_add_ps(dot_acc);
|
|
||||||
let mut norm_a = _mm512_reduce_add_ps(norm_a_acc);
|
|
||||||
let mut norm_b = _mm512_reduce_add_ps(norm_b_acc);
|
|
||||||
|
|
||||||
while i < len {
|
|
||||||
dot += a[i] * b[i];
|
|
||||||
norm_a += a[i] * a[i];
|
|
||||||
norm_b += b[i] * b[i];
|
|
||||||
i += 1;
|
|
||||||
}
|
|
||||||
|
|
||||||
let denom = (norm_a * norm_b).sqrt();
|
|
||||||
if denom == 0.0 { 0.0 } else { dot / denom }
|
|
||||||
}}
|
|
||||||
|
|
||||||
/// AVX-512 L2 distance.
|
/// AVX-512 L2 distance.
|
||||||
///
|
///
|
||||||
@@ -95,27 +99,29 @@ pub unsafe fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 { unsafe {
|
|||||||
/// Caller must verify is_x86_feature_detected!("avx512f").
|
/// Caller must verify is_x86_feature_detected!("avx512f").
|
||||||
// SAFETY: Caller must have verified avx512f via is_x86_feature_detected!.
|
// SAFETY: Caller must have verified avx512f via is_x86_feature_detected!.
|
||||||
#[target_feature(enable = "avx512f")]
|
#[target_feature(enable = "avx512f")]
|
||||||
pub unsafe fn l2_distance(a: &[f32], b: &[f32]) -> f32 { unsafe {
|
pub unsafe fn l2_distance(a: &[f32], b: &[f32]) -> f32 {
|
||||||
assert_eq!(a.len(), b.len());
|
unsafe {
|
||||||
let len = a.len();
|
assert_eq!(a.len(), b.len());
|
||||||
let mut i = 0;
|
let len = a.len();
|
||||||
let mut acc = _mm512_setzero_ps();
|
let mut i = 0;
|
||||||
|
let mut acc = _mm512_setzero_ps();
|
||||||
|
|
||||||
while i + 16 <= len {
|
while i + 16 <= len {
|
||||||
let va = _mm512_loadu_ps(a.as_ptr().add(i));
|
let va = _mm512_loadu_ps(a.as_ptr().add(i));
|
||||||
let vb = _mm512_loadu_ps(b.as_ptr().add(i));
|
let vb = _mm512_loadu_ps(b.as_ptr().add(i));
|
||||||
let diff = _mm512_sub_ps(va, vb);
|
let diff = _mm512_sub_ps(va, vb);
|
||||||
acc = _mm512_fmadd_ps(diff, diff, acc);
|
acc = _mm512_fmadd_ps(diff, diff, acc);
|
||||||
i += 16;
|
i += 16;
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut sum = _mm512_reduce_add_ps(acc);
|
||||||
|
|
||||||
|
while i < len {
|
||||||
|
let d = a[i] - b[i];
|
||||||
|
sum += d * d;
|
||||||
|
i += 1;
|
||||||
|
}
|
||||||
|
|
||||||
|
sum.sqrt()
|
||||||
}
|
}
|
||||||
|
}
|
||||||
let mut sum = _mm512_reduce_add_ps(acc);
|
|
||||||
|
|
||||||
while i < len {
|
|
||||||
let d = a[i] - b[i];
|
|
||||||
sum += d * d;
|
|
||||||
i += 1;
|
|
||||||
}
|
|
||||||
|
|
||||||
sum.sqrt()
|
|
||||||
}}
|
|
||||||
|
|||||||
@@ -116,14 +116,15 @@ impl GpuSearchBackend {
|
|||||||
|
|
||||||
// If we don't have an accelerator but now above threshold, try init
|
// If we don't have an accelerator but now above threshold, try init
|
||||||
if vectors.len() >= self.threshold
|
if vectors.len() >= self.threshold
|
||||||
&& let Ok(mut accel) = clawhdf5_gpu::GpuAccelerator::new() {
|
&& let Ok(mut accel) = clawhdf5_gpu::GpuAccelerator::new()
|
||||||
let flat: Vec<f32> = vectors.iter().flat_map(|v| v.iter().copied()).collect();
|
{
|
||||||
if accel.upload_vectors(&flat, self.dim).is_ok()
|
let flat: Vec<f32> = vectors.iter().flat_map(|v| v.iter().copied()).collect();
|
||||||
&& accel.upload_norms(norms).is_ok()
|
if accel.upload_vectors(&flat, self.dim).is_ok()
|
||||||
{
|
&& accel.upload_norms(norms).is_ok()
|
||||||
self.accelerator = Some(accel);
|
{
|
||||||
}
|
self.accelerator = Some(accel);
|
||||||
}
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
#[cfg(not(feature = "gpu"))]
|
#[cfg(not(feature = "gpu"))]
|
||||||
|
|||||||
@@ -26,9 +26,7 @@ impl HDF5Memory {
|
|||||||
) -> Vec<(usize, f32)> {
|
) -> Vec<(usize, f32)> {
|
||||||
self.ensure_hnsw_fresh();
|
self.ensure_hnsw_fresh();
|
||||||
match self.hnsw.as_ref() {
|
match self.hnsw.as_ref() {
|
||||||
Some(index)
|
Some(index) if !index.is_empty() && index.dimension() == query_embedding.len() => {
|
||||||
if !index.is_empty() && index.dimension() == query_embedding.len() =>
|
|
||||||
{
|
|
||||||
// Over-fetch so the merge sees a useful vector pool; cosine
|
// Over-fetch so the merge sees a useful vector pool; cosine
|
||||||
// distance from the index converts back to similarity (1 - d).
|
// distance from the index converts back to similarity (1 - d).
|
||||||
let pool = (k * 8).max(64);
|
let pool = (k * 8).max(64);
|
||||||
@@ -38,7 +36,13 @@ impl HDF5Memory {
|
|||||||
.map(|(id, dist)| (id, 1.0 - dist))
|
.map(|(id, dist)| (id, 1.0 - dist))
|
||||||
.collect();
|
.collect();
|
||||||
let kw_scores = bm25.search(query_text, self.cache.len());
|
let kw_scores = bm25.search(query_text, self.cache.len());
|
||||||
hybrid::merge_vector_keyword(vec_scores, kw_scores, vector_weight, keyword_weight, k)
|
hybrid::merge_vector_keyword(
|
||||||
|
vec_scores,
|
||||||
|
kw_scores,
|
||||||
|
vector_weight,
|
||||||
|
keyword_weight,
|
||||||
|
k,
|
||||||
|
)
|
||||||
}
|
}
|
||||||
_ => hybrid::hybrid_search(
|
_ => hybrid::hybrid_search(
|
||||||
query_embedding,
|
query_embedding,
|
||||||
|
|||||||
@@ -80,8 +80,7 @@ fn hnsw_matches_bruteforce_oracle() {
|
|||||||
oracle.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
|
oracle.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
|
||||||
let oracle_ids: std::collections::HashSet<usize> =
|
let oracle_ids: std::collections::HashSet<usize> =
|
||||||
oracle.iter().take(k).map(|(i, _)| *i).collect();
|
oracle.iter().take(k).map(|(i, _)| *i).collect();
|
||||||
let hnsw_ids: std::collections::HashSet<usize> =
|
let hnsw_ids: std::collections::HashSet<usize> = results.iter().map(|r| r.index).collect();
|
||||||
results.iter().map(|r| r.index).collect();
|
|
||||||
|
|
||||||
let overlap = oracle_ids.intersection(&hnsw_ids).count();
|
let overlap = oracle_ids.intersection(&hnsw_ids).count();
|
||||||
assert!(
|
assert!(
|
||||||
@@ -127,17 +126,19 @@ fn incremental_inserts_after_search_are_found() {
|
|||||||
// First batch, then a search to force the index to build.
|
// First batch, then a search to force the index to build.
|
||||||
for i in 0..40 {
|
for i in 0..40 {
|
||||||
let v = make_vector(&mut seed, dim);
|
let v = make_vector(&mut seed, dim);
|
||||||
mem.save(entry(&format!("a{i}"), v, &format!("a{i}"))).unwrap();
|
mem.save(entry(&format!("a{i}"), v, &format!("a{i}")))
|
||||||
|
.unwrap();
|
||||||
}
|
}
|
||||||
let _ = mem.hybrid_search(&make_vector(&mut seed, dim), "", 1.0, 0.0, 5);
|
let _ = mem.hybrid_search(&make_vector(&mut seed, dim), "", 1.0, 0.0, 5);
|
||||||
|
|
||||||
// Now insert a distinctive vector incrementally and confirm we can find it.
|
// Now insert a distinctive vector incrementally and confirm we can find it.
|
||||||
let needle = vec![10.0f32; dim];
|
let needle = vec![10.0f32; dim];
|
||||||
let idx = mem
|
let idx = mem.save(entry("needle", needle.clone(), "needle")).unwrap();
|
||||||
.save(entry("needle", needle.clone(), "needle"))
|
|
||||||
.unwrap();
|
|
||||||
let hits = mem.hybrid_search(&needle, "", 1.0, 0.0, 1);
|
let hits = mem.hybrid_search(&needle, "", 1.0, 0.0, 1);
|
||||||
assert_eq!(hits[0].index, idx, "incrementally inserted vector must be found");
|
assert_eq!(
|
||||||
|
hits[0].index, idx,
|
||||||
|
"incrementally inserted vector must be found"
|
||||||
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
@@ -158,6 +159,9 @@ fn save_batch_then_search_is_consistent() {
|
|||||||
// Exact-match queries should resolve to themselves after a batch insert.
|
// Exact-match queries should resolve to themselves after a batch insert.
|
||||||
for probe in [0usize, 17, 49] {
|
for probe in [0usize, 17, 49] {
|
||||||
let hits = mem.hybrid_search(&vectors[probe], "", 1.0, 0.0, 1);
|
let hits = mem.hybrid_search(&vectors[probe], "", 1.0, 0.0, 1);
|
||||||
assert_eq!(hits[0].index, probe, "batch-inserted vector {probe} not found");
|
assert_eq!(
|
||||||
|
hits[0].index, probe,
|
||||||
|
"batch-inserted vector {probe} not found"
|
||||||
|
);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -379,7 +379,13 @@ impl HnswIndex {
|
|||||||
|
|
||||||
// Phase 1: greedy descent from the top down to node_level + 1.
|
// Phase 1: greedy descent from the top down to node_level + 1.
|
||||||
for layer in (node_level + 1..=ep_level).rev() {
|
for layer in (node_level + 1..=ep_level).rev() {
|
||||||
ep = greedy_closest(&self.vectors, &self.graph[layer], &self.vectors[id], ep, self.metric);
|
ep = greedy_closest(
|
||||||
|
&self.vectors,
|
||||||
|
&self.graph[layer],
|
||||||
|
&self.vectors[id],
|
||||||
|
ep,
|
||||||
|
self.metric,
|
||||||
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
// Phase 2: search and connect from min(node_level, ep_level) down to 0.
|
// Phase 2: search and connect from min(node_level, ep_level) down to 0.
|
||||||
@@ -1012,11 +1018,14 @@ fn get_attr_i64(attrs: &[(String, AttrValue)], name: &str) -> Result<i64, Format
|
|||||||
/// Like [`get_attr_i64`] but returns `None` when the attribute is absent or not
|
/// Like [`get_attr_i64`] but returns `None` when the attribute is absent or not
|
||||||
/// an integer, instead of erroring. Used for optional/back-compat attributes.
|
/// an integer, instead of erroring. Used for optional/back-compat attributes.
|
||||||
fn get_attr_i64_opt(attrs: &[(String, AttrValue)], name: &str) -> Option<i64> {
|
fn get_attr_i64_opt(attrs: &[(String, AttrValue)], name: &str) -> Option<i64> {
|
||||||
attrs.iter().find(|(n, _)| n == name).and_then(|(_, v)| match v {
|
attrs
|
||||||
AttrValue::I64(val) => Some(*val),
|
.iter()
|
||||||
AttrValue::U64(val) => Some(*val as i64),
|
.find(|(n, _)| n == name)
|
||||||
_ => None,
|
.and_then(|(_, v)| match v {
|
||||||
})
|
AttrValue::I64(val) => Some(*val),
|
||||||
|
AttrValue::U64(val) => Some(*val as i64),
|
||||||
|
_ => None,
|
||||||
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
fn get_attr_string(attrs: &[(String, AttrValue)], name: &str) -> Result<String, FormatError> {
|
fn get_attr_string(attrs: &[(String, AttrValue)], name: &str) -> Result<String, FormatError> {
|
||||||
|
|||||||
@@ -41,11 +41,7 @@ fn bench_metadata_attrs_write(c: &mut Criterion) {
|
|||||||
let path = tmp.path().join("attrs_libhdf5.h5");
|
let path = tmp.path().join("attrs_libhdf5.h5");
|
||||||
b.iter(|| {
|
b.iter(|| {
|
||||||
let file = hdf5::File::create(&path).unwrap();
|
let file = hdf5::File::create(&path).unwrap();
|
||||||
let ds = file
|
let ds = file.new_dataset::<f64>().shape([3]).create("data").unwrap();
|
||||||
.new_dataset::<f64>()
|
|
||||||
.shape([3])
|
|
||||||
.create("data")
|
|
||||||
.unwrap();
|
|
||||||
ds.write(&[1.0f64, 2.0, 3.0]).unwrap();
|
ds.write(&[1.0f64, 2.0, 3.0]).unwrap();
|
||||||
for i in 0..k {
|
for i in 0..k {
|
||||||
ds.new_attr::<i64>()
|
ds.new_attr::<i64>()
|
||||||
@@ -258,11 +254,7 @@ fn bench_metadata_open_from_disk(c: &mut Criterion) {
|
|||||||
let libhdf5_path = tmp.path().join("open_libhdf5.h5");
|
let libhdf5_path = tmp.path().join("open_libhdf5.h5");
|
||||||
{
|
{
|
||||||
let file = hdf5::File::create(&libhdf5_path).unwrap();
|
let file = hdf5::File::create(&libhdf5_path).unwrap();
|
||||||
let ds = file
|
let ds = file.new_dataset::<f64>().shape([3]).create("data").unwrap();
|
||||||
.new_dataset::<f64>()
|
|
||||||
.shape([3])
|
|
||||||
.create("data")
|
|
||||||
.unwrap();
|
|
||||||
ds.write(&[1.0f64, 2.0, 3.0]).unwrap();
|
ds.write(&[1.0f64, 2.0, 3.0]).unwrap();
|
||||||
ds.new_attr::<i64>()
|
ds.new_attr::<i64>()
|
||||||
.create("label")
|
.create("label")
|
||||||
@@ -307,13 +299,17 @@ fn bench_metadata_parse_in_memory(c: &mut Criterion) {
|
|||||||
fb.finish().unwrap()
|
fb.finish().unwrap()
|
||||||
};
|
};
|
||||||
|
|
||||||
group.bench_with_input(BenchmarkId::new("clawhdf5", "in_memory"), &bytes, |b, raw| {
|
group.bench_with_input(
|
||||||
b.iter(|| {
|
BenchmarkId::new("clawhdf5", "in_memory"),
|
||||||
let file = File::from_bytes(raw.clone()).unwrap();
|
&bytes,
|
||||||
let ds = file.dataset("data").unwrap();
|
|b, raw| {
|
||||||
ds.attrs().unwrap()
|
b.iter(|| {
|
||||||
});
|
let file = File::from_bytes(raw.clone()).unwrap();
|
||||||
});
|
let ds = file.dataset("data").unwrap();
|
||||||
|
ds.attrs().unwrap()
|
||||||
|
});
|
||||||
|
},
|
||||||
|
);
|
||||||
|
|
||||||
group.finish();
|
group.finish();
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -74,11 +74,7 @@ fn bench_read_sequential(c: &mut Criterion) {
|
|||||||
let data: Vec<f32> = (0..nn).map(|i| i as f32 * 0.001).collect();
|
let data: Vec<f32> = (0..nn).map(|i| i as f32 * 0.001).collect();
|
||||||
{
|
{
|
||||||
let lf = hdf5::File::create(&path).unwrap();
|
let lf = hdf5::File::create(&path).unwrap();
|
||||||
let lds = lf
|
let lds = lf.new_dataset::<f32>().shape([nn]).create("data").unwrap();
|
||||||
.new_dataset::<f32>()
|
|
||||||
.shape([nn])
|
|
||||||
.create("data")
|
|
||||||
.unwrap();
|
|
||||||
lds.write(data.as_slice()).unwrap();
|
lds.write(data.as_slice()).unwrap();
|
||||||
}
|
}
|
||||||
b.iter(|| {
|
b.iter(|| {
|
||||||
@@ -235,19 +231,23 @@ fn bench_read_zerocopy_mmap(c: &mut Criterion) {
|
|||||||
|
|
||||||
group.throughput(Throughput::Bytes((n * size_of::<f64>()) as u64));
|
group.throughput(Throughput::Bytes((n * size_of::<f64>()) as u64));
|
||||||
|
|
||||||
group.bench_with_input(BenchmarkId::new("clawhdf5_mmap_zerocopy", n), &path, |b, p| {
|
group.bench_with_input(
|
||||||
b.iter(|| {
|
BenchmarkId::new("clawhdf5_mmap_zerocopy", n),
|
||||||
let file = MmapFile::open(p).unwrap();
|
&path,
|
||||||
let ds = file.dataset("data").unwrap();
|
|b, p| {
|
||||||
let slice = ds.read_f64_zerocopy().unwrap();
|
b.iter(|| {
|
||||||
// Sum every element to force the mapped pages to actually be
|
let file = MmapFile::open(p).unwrap();
|
||||||
// faulted in — returning just `.len()` would measure nothing
|
let ds = file.dataset("data").unwrap();
|
||||||
// but the mmap() syscall, repeating the exact "too-fast-to-
|
let slice = ds.read_f64_zerocopy().unwrap();
|
||||||
// be-real" mistake this benchmark exists to fix.
|
// Sum every element to force the mapped pages to actually be
|
||||||
let sum: f64 = slice.map(|s| s.iter().sum()).unwrap_or(0.0);
|
// faulted in — returning just `.len()` would measure nothing
|
||||||
criterion::black_box(sum)
|
// but the mmap() syscall, repeating the exact "too-fast-to-
|
||||||
});
|
// be-real" mistake this benchmark exists to fix.
|
||||||
});
|
let sum: f64 = slice.map(|s| s.iter().sum()).unwrap_or(0.0);
|
||||||
|
criterion::black_box(sum)
|
||||||
|
});
|
||||||
|
},
|
||||||
|
);
|
||||||
|
|
||||||
group.bench_with_input(BenchmarkId::new("clawhdf5_copy", n), &path, |b, p| {
|
group.bench_with_input(BenchmarkId::new("clawhdf5_copy", n), &path, |b, p| {
|
||||||
b.iter(|| {
|
b.iter(|| {
|
||||||
|
|||||||
@@ -63,11 +63,8 @@ fn bench_write_2d_chunked(c: &mut Criterion) {
|
|||||||
let mut group = c.benchmark_group("write_2d_chunked");
|
let mut group = c.benchmark_group("write_2d_chunked");
|
||||||
|
|
||||||
// (rows, cols, chunk_rows, chunk_cols)
|
// (rows, cols, chunk_rows, chunk_cols)
|
||||||
let configs: &[(usize, usize, u64, u64)] = &[
|
let configs: &[(usize, usize, u64, u64)] =
|
||||||
(32, 32, 8, 32),
|
&[(32, 32, 8, 32), (128, 128, 32, 128), (512, 512, 64, 512)];
|
||||||
(128, 128, 32, 128),
|
|
||||||
(512, 512, 64, 512),
|
|
||||||
];
|
|
||||||
|
|
||||||
for &(rows, cols, cr, cc) in configs {
|
for &(rows, cols, cr, cc) in configs {
|
||||||
let n = rows * cols;
|
let n = rows * cols;
|
||||||
@@ -120,11 +117,8 @@ fn bench_write_2d_chunked(c: &mut Criterion) {
|
|||||||
fn bench_write_2d_chunked_zstd(c: &mut Criterion) {
|
fn bench_write_2d_chunked_zstd(c: &mut Criterion) {
|
||||||
let mut group = c.benchmark_group("write_2d_chunked_zstd");
|
let mut group = c.benchmark_group("write_2d_chunked_zstd");
|
||||||
|
|
||||||
let configs: &[(usize, usize, u64, u64)] = &[
|
let configs: &[(usize, usize, u64, u64)] =
|
||||||
(32, 32, 8, 32),
|
&[(32, 32, 8, 32), (128, 128, 32, 128), (512, 512, 64, 512)];
|
||||||
(128, 128, 32, 128),
|
|
||||||
(512, 512, 64, 512),
|
|
||||||
];
|
|
||||||
|
|
||||||
for &(rows, cols, cr, cc) in configs {
|
for &(rows, cols, cr, cc) in configs {
|
||||||
let n = rows * cols;
|
let n = rows * cols;
|
||||||
@@ -132,19 +126,23 @@ fn bench_write_2d_chunked_zstd(c: &mut Criterion) {
|
|||||||
let label = format!("{rows}x{cols}");
|
let label = format!("{rows}x{cols}");
|
||||||
group.throughput(Throughput::Bytes((n * size_of::<f32>()) as u64));
|
group.throughput(Throughput::Bytes((n * size_of::<f32>()) as u64));
|
||||||
|
|
||||||
group.bench_with_input(BenchmarkId::new("clawhdf5/zstd-3", &label), &data, |b, d| {
|
group.bench_with_input(
|
||||||
let tmp = TempDir::new().unwrap();
|
BenchmarkId::new("clawhdf5/zstd-3", &label),
|
||||||
let path = tmp.path().join("write_2d_chunked_zstd.h5");
|
&data,
|
||||||
b.iter(|| {
|
|b, d| {
|
||||||
let mut fb = FileBuilder::new();
|
let tmp = TempDir::new().unwrap();
|
||||||
fb.create_dataset("matrix")
|
let path = tmp.path().join("write_2d_chunked_zstd.h5");
|
||||||
.with_f32_data(d)
|
b.iter(|| {
|
||||||
.with_shape(&[rows as u64, cols as u64])
|
let mut fb = FileBuilder::new();
|
||||||
.with_chunks(&[cr, cc])
|
fb.create_dataset("matrix")
|
||||||
.with_zstd(3);
|
.with_f32_data(d)
|
||||||
fb.write(&path).unwrap();
|
.with_shape(&[rows as u64, cols as u64])
|
||||||
});
|
.with_chunks(&[cr, cc])
|
||||||
});
|
.with_zstd(3);
|
||||||
|
fb.write(&path).unwrap();
|
||||||
|
});
|
||||||
|
},
|
||||||
|
);
|
||||||
|
|
||||||
group.bench_with_input(
|
group.bench_with_input(
|
||||||
BenchmarkId::new("clawhdf5/deflate-6", &label),
|
BenchmarkId::new("clawhdf5/deflate-6", &label),
|
||||||
@@ -178,11 +176,8 @@ fn bench_write_2d_chunked_zstd(c: &mut Criterion) {
|
|||||||
fn bench_write_2d_chunked_pcodec(c: &mut Criterion) {
|
fn bench_write_2d_chunked_pcodec(c: &mut Criterion) {
|
||||||
let mut group = c.benchmark_group("write_2d_chunked_pcodec");
|
let mut group = c.benchmark_group("write_2d_chunked_pcodec");
|
||||||
|
|
||||||
let configs: &[(usize, usize, u64, u64)] = &[
|
let configs: &[(usize, usize, u64, u64)] =
|
||||||
(32, 32, 8, 32),
|
&[(32, 32, 8, 32), (128, 128, 32, 128), (512, 512, 64, 512)];
|
||||||
(128, 128, 32, 128),
|
|
||||||
(512, 512, 64, 512),
|
|
||||||
];
|
|
||||||
|
|
||||||
for &(rows, cols, cr, cc) in configs {
|
for &(rows, cols, cr, cc) in configs {
|
||||||
let n = rows * cols;
|
let n = rows * cols;
|
||||||
|
|||||||
@@ -270,14 +270,29 @@ pub(crate) fn flate2_decompress_preallocated(
|
|||||||
Ok(output)
|
Ok(output)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Absolute ceiling on decompressed output when the caller has no size hint,
|
||||||
|
/// preventing unbounded allocation from a hostile/corrupted zlib stream.
|
||||||
|
const MAX_DECOMPRESS_SIZE: usize = 256 * 1024 * 1024;
|
||||||
|
|
||||||
/// Streaming decompress with dynamic sizing (when output size is unknown).
|
/// Streaming decompress with dynamic sizing (when output size is unknown).
|
||||||
|
///
|
||||||
|
/// Bounded by [`MAX_DECOMPRESS_SIZE`] since there is no chunk-size hint to
|
||||||
|
/// validate against here — an unbounded `read_to_end` would let a hostile
|
||||||
|
/// zlib stream force arbitrarily large allocation (a "zlib bomb").
|
||||||
pub(crate) fn flate2_decompress_streaming(data: &[u8]) -> Result<Vec<u8>, String> {
|
pub(crate) fn flate2_decompress_streaming(data: &[u8]) -> Result<Vec<u8>, String> {
|
||||||
use std::io::Read;
|
use std::io::Read;
|
||||||
let mut decoder = flate2::read::ZlibDecoder::new(data);
|
let decoder = flate2::read::ZlibDecoder::new(data);
|
||||||
let mut result = Vec::new();
|
let mut result = Vec::new();
|
||||||
decoder
|
decoder
|
||||||
|
.take(MAX_DECOMPRESS_SIZE as u64 + 1)
|
||||||
.read_to_end(&mut result)
|
.read_to_end(&mut result)
|
||||||
.map_err(|e| e.to_string())?;
|
.map_err(|e| e.to_string())?;
|
||||||
|
if result.len() > MAX_DECOMPRESS_SIZE {
|
||||||
|
return Err(format!(
|
||||||
|
"decompressed output exceeds {} MiB limit",
|
||||||
|
MAX_DECOMPRESS_SIZE / 1024 / 1024
|
||||||
|
));
|
||||||
|
}
|
||||||
Ok(result)
|
Ok(result)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -11,6 +11,7 @@ categories = ["parser-implementations", "science", "encoding", "no-std"]
|
|||||||
|
|
||||||
[dependencies]
|
[dependencies]
|
||||||
byteorder = { version = "1", default-features = false }
|
byteorder = { version = "1", default-features = false }
|
||||||
|
portable-atomic = { version = "1" }
|
||||||
flate2 = { version = "1", default-features = false, features = ["rust_backend"], optional = true }
|
flate2 = { version = "1", default-features = false, features = ["rust_backend"], optional = true }
|
||||||
sha2 = { version = "0.10", default-features = false, optional = true }
|
sha2 = { version = "0.10", default-features = false, optional = true }
|
||||||
rayon = { version = "1", optional = true }
|
rayon = { version = "1", optional = true }
|
||||||
|
|||||||
@@ -16,6 +16,8 @@ use core::ops::{Deref, DerefMut};
|
|||||||
use alloc::collections::BTreeMap;
|
use alloc::collections::BTreeMap;
|
||||||
#[cfg(feature = "std")]
|
#[cfg(feature = "std")]
|
||||||
use std::collections::HashMap;
|
use std::collections::HashMap;
|
||||||
|
#[cfg(feature = "std")]
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
use crate::chunk_index::{ChunkIndex, ChunkLayout};
|
use crate::chunk_index::{ChunkIndex, ChunkLayout};
|
||||||
use crate::chunked_read::ChunkInfo;
|
use crate::chunked_read::ChunkInfo;
|
||||||
@@ -64,6 +66,11 @@ pub struct CacheAlignedBuffer {
|
|||||||
|
|
||||||
// SAFETY: The raw pointer is exclusively owned — no aliasing.
|
// SAFETY: The raw pointer is exclusively owned — no aliasing.
|
||||||
unsafe impl Send for CacheAlignedBuffer {}
|
unsafe impl Send for CacheAlignedBuffer {}
|
||||||
|
// SAFETY: `CacheAlignedBuffer` exposes its contents only via `&[u8]`/`&mut
|
||||||
|
// [u8]` through the ordinary borrow-checked `Deref`/`DerefMut` impls below —
|
||||||
|
// the same access pattern as `Vec<u8>`, which is `Sync`. Needed so
|
||||||
|
// `Arc<CacheAlignedBuffer>` (used by the chunk cache) is itself `Send`.
|
||||||
|
unsafe impl Sync for CacheAlignedBuffer {}
|
||||||
|
|
||||||
impl CacheAlignedBuffer {
|
impl CacheAlignedBuffer {
|
||||||
/// Allocate a new cache-line-aligned buffer of exactly `len` bytes,
|
/// Allocate a new cache-line-aligned buffer of exactly `len` bytes,
|
||||||
@@ -223,7 +230,9 @@ pub const DEFAULT_MAX_SLOTS: usize = 521;
|
|||||||
#[cfg(feature = "std")]
|
#[cfg(feature = "std")]
|
||||||
struct CachedChunk {
|
struct CachedChunk {
|
||||||
coord: ChunkCoord,
|
coord: ChunkCoord,
|
||||||
data: CacheAlignedBuffer,
|
/// Shared so a cache hit is a refcount bump, not a copy of the whole
|
||||||
|
/// (potentially large) decompressed chunk.
|
||||||
|
data: Arc<CacheAlignedBuffer>,
|
||||||
/// Monotonically increasing access counter for LRU ordering.
|
/// Monotonically increasing access counter for LRU ordering.
|
||||||
last_access: u64,
|
last_access: u64,
|
||||||
}
|
}
|
||||||
@@ -267,6 +276,12 @@ struct CacheInner {
|
|||||||
/// LRU cache of decompressed chunk data.
|
/// LRU cache of decompressed chunk data.
|
||||||
slots: Vec<CachedChunk>,
|
slots: Vec<CachedChunk>,
|
||||||
|
|
||||||
|
/// Coordinate -> index into `slots`, for O(1) lookup instead of a linear
|
||||||
|
/// scan. Kept in sync with `slots` on every insert/evict/clear — in
|
||||||
|
/// particular, `slots.swap_remove(i)` moves the last element into slot
|
||||||
|
/// `i`, so the moved element's index entry must be updated too.
|
||||||
|
slot_index: HashMap<ChunkCoord, usize>,
|
||||||
|
|
||||||
/// Current total bytes of cached decompressed data.
|
/// Current total bytes of cached decompressed data.
|
||||||
current_bytes: usize,
|
current_bytes: usize,
|
||||||
|
|
||||||
@@ -344,6 +359,7 @@ impl ChunkCache {
|
|||||||
index: None,
|
index: None,
|
||||||
index_addr: None,
|
index_addr: None,
|
||||||
slots: Vec::with_capacity(max_slots.min(64)),
|
slots: Vec::with_capacity(max_slots.min(64)),
|
||||||
|
slot_index: HashMap::with_capacity(max_slots.min(64)),
|
||||||
current_bytes: 0,
|
current_bytes: 0,
|
||||||
max_bytes,
|
max_bytes,
|
||||||
max_slots,
|
max_slots,
|
||||||
@@ -375,6 +391,7 @@ impl ChunkCache {
|
|||||||
inner.chunk_index = None;
|
inner.chunk_index = None;
|
||||||
inner.chunk_layout = None;
|
inner.chunk_layout = None;
|
||||||
inner.slots.clear();
|
inner.slots.clear();
|
||||||
|
inner.slot_index.clear();
|
||||||
inner.current_bytes = 0;
|
inner.current_bytes = 0;
|
||||||
inner.last_coord = None;
|
inner.last_coord = None;
|
||||||
inner.index_addr = Some(addr);
|
inner.index_addr = Some(addr);
|
||||||
@@ -477,8 +494,20 @@ impl ChunkCache {
|
|||||||
|
|
||||||
/// Try to get cached decompressed data for a chunk coordinate.
|
/// Try to get cached decompressed data for a chunk coordinate.
|
||||||
///
|
///
|
||||||
/// Returns a clone of the cache-line-aligned buffer.
|
/// O(1) lookup. Returns an owned copy for API compatibility with callers
|
||||||
|
/// that need a `Vec<u8>`; prefer [`Self::get_decompressed_aligned`] when
|
||||||
|
/// an `Arc`-shared buffer works for the caller, since that avoids the
|
||||||
|
/// copy entirely.
|
||||||
pub fn get_decompressed(&self, coord: &[u64]) -> Option<Vec<u8>> {
|
pub fn get_decompressed(&self, coord: &[u64]) -> Option<Vec<u8>> {
|
||||||
|
self.get_decompressed_aligned(coord)
|
||||||
|
.map(|arc| arc.as_slice().to_vec())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Try to get a reference-counted clone of the aligned buffer for a chunk.
|
||||||
|
///
|
||||||
|
/// O(1) index lookup; the clone is an `Arc` refcount bump, not a copy of
|
||||||
|
/// the underlying decompressed data.
|
||||||
|
pub fn get_decompressed_aligned(&self, coord: &[u64]) -> Option<Arc<CacheAlignedBuffer>> {
|
||||||
let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
|
let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
|
||||||
inner.tick += 1;
|
inner.tick += 1;
|
||||||
let tick = inner.tick;
|
let tick = inner.tick;
|
||||||
@@ -500,36 +529,12 @@ impl ChunkCache {
|
|||||||
}
|
}
|
||||||
inner.last_coord = Some(coord.to_vec());
|
inner.last_coord = Some(coord.to_vec());
|
||||||
|
|
||||||
let mut found = None;
|
let found = if let Some(&idx) = inner.slot_index.get(coord) {
|
||||||
for slot in inner.slots.iter_mut() {
|
inner.slots[idx].last_access = tick;
|
||||||
if slot.coord.as_slice() == coord {
|
Some(Arc::clone(&inner.slots[idx].data))
|
||||||
slot.last_access = tick;
|
|
||||||
found = Some(slot.data.to_vec());
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
if let Some(ref data) = found {
|
|
||||||
inner.stats.hits += 1;
|
|
||||||
inner.stats.bytes_read += data.len() as u64;
|
|
||||||
} else {
|
} else {
|
||||||
inner.stats.misses += 1;
|
None
|
||||||
}
|
};
|
||||||
found
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Try to get a reference-counted clone of the aligned buffer for a chunk.
|
|
||||||
pub fn get_decompressed_aligned(&self, coord: &[u64]) -> Option<CacheAlignedBuffer> {
|
|
||||||
let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
|
|
||||||
inner.tick += 1;
|
|
||||||
let tick = inner.tick;
|
|
||||||
let mut found = None;
|
|
||||||
for slot in inner.slots.iter_mut() {
|
|
||||||
if slot.coord.as_slice() == coord {
|
|
||||||
slot.last_access = tick;
|
|
||||||
found = Some(slot.data.clone());
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
if let Some(ref data) = found {
|
if let Some(ref data) = found {
|
||||||
inner.stats.hits += 1;
|
inner.stats.hits += 1;
|
||||||
inner.stats.bytes_read += data.len() as u64;
|
inner.stats.bytes_read += data.len() as u64;
|
||||||
@@ -542,30 +547,39 @@ impl ChunkCache {
|
|||||||
/// Insert decompressed chunk data into the LRU cache.
|
/// Insert decompressed chunk data into the LRU cache.
|
||||||
///
|
///
|
||||||
/// The data is stored in a [`CacheAlignedBuffer`] so subsequent reads
|
/// The data is stored in a [`CacheAlignedBuffer`] so subsequent reads
|
||||||
/// return cache-line-aligned memory.
|
/// return cache-line-aligned memory. Returns the `Arc`-shared buffer that
|
||||||
pub fn put_decompressed(&self, coord: ChunkCoord, data: Vec<u8>) {
|
/// is now cached (or already was), so the caller can reuse it directly
|
||||||
let aligned = CacheAlignedBuffer::from_slice(&data);
|
/// instead of holding a separate copy of the same data.
|
||||||
self.put_decompressed_aligned(coord, aligned);
|
pub fn put_decompressed(&self, coord: ChunkCoord, data: Vec<u8>) -> Arc<CacheAlignedBuffer> {
|
||||||
|
let aligned = CacheAlignedBuffer::from_vec(data);
|
||||||
|
self.put_decompressed_aligned(coord, aligned)
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Insert an already-aligned buffer into the LRU cache.
|
/// Insert an already-aligned buffer into the LRU cache.
|
||||||
pub fn put_decompressed_aligned(&self, coord: ChunkCoord, data: CacheAlignedBuffer) {
|
///
|
||||||
|
/// Returns the `Arc`-shared buffer now held by the cache (the one just
|
||||||
|
/// inserted, or the existing cached copy if `coord` was already present).
|
||||||
|
pub fn put_decompressed_aligned(
|
||||||
|
&self,
|
||||||
|
coord: ChunkCoord,
|
||||||
|
data: CacheAlignedBuffer,
|
||||||
|
) -> Arc<CacheAlignedBuffer> {
|
||||||
|
let data = Arc::new(data);
|
||||||
let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
|
let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
|
||||||
let data_len = data.len();
|
let data_len = data.len();
|
||||||
|
|
||||||
// Don't cache if single chunk exceeds budget
|
// Don't cache if single chunk exceeds budget — still return the data
|
||||||
|
// to the caller, just don't retain it.
|
||||||
if data_len > inner.max_bytes {
|
if data_len > inner.max_bytes {
|
||||||
return;
|
return data;
|
||||||
}
|
}
|
||||||
|
|
||||||
// Check if already present
|
// Check if already present
|
||||||
inner.tick += 1;
|
inner.tick += 1;
|
||||||
let tick = inner.tick;
|
let tick = inner.tick;
|
||||||
for slot in inner.slots.iter_mut() {
|
if let Some(&idx) = inner.slot_index.get(&coord) {
|
||||||
if slot.coord == coord {
|
inner.slots[idx].last_access = tick;
|
||||||
slot.last_access = tick;
|
return Arc::clone(&inner.slots[idx].data); // already cached
|
||||||
return; // already cached
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
// Evict until we have room
|
// Evict until we have room
|
||||||
@@ -581,16 +595,26 @@ impl ChunkCache {
|
|||||||
.map(|(i, _)| i)
|
.map(|(i, _)| i)
|
||||||
.unwrap();
|
.unwrap();
|
||||||
let removed = inner.slots.swap_remove(lru_idx);
|
let removed = inner.slots.swap_remove(lru_idx);
|
||||||
|
inner.slot_index.remove(&removed.coord);
|
||||||
|
// swap_remove moved the former last element into `lru_idx` (unless
|
||||||
|
// it *was* the last element) — fix up that element's index entry.
|
||||||
|
if lru_idx < inner.slots.len() {
|
||||||
|
let moved_coord = inner.slots[lru_idx].coord.clone();
|
||||||
|
inner.slot_index.insert(moved_coord, lru_idx);
|
||||||
|
}
|
||||||
inner.current_bytes -= removed.data.len();
|
inner.current_bytes -= removed.data.len();
|
||||||
inner.stats.evictions += 1;
|
inner.stats.evictions += 1;
|
||||||
}
|
}
|
||||||
|
|
||||||
inner.current_bytes += data_len;
|
inner.current_bytes += data_len;
|
||||||
|
let new_idx = inner.slots.len();
|
||||||
|
inner.slot_index.insert(coord.clone(), new_idx);
|
||||||
inner.slots.push(CachedChunk {
|
inner.slots.push(CachedChunk {
|
||||||
coord,
|
coord,
|
||||||
data,
|
data: Arc::clone(&data),
|
||||||
last_access: tick,
|
last_access: tick,
|
||||||
});
|
});
|
||||||
|
data
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Clear the entire cache (index + decompressed data).
|
/// Clear the entire cache (index + decompressed data).
|
||||||
@@ -599,6 +623,7 @@ impl ChunkCache {
|
|||||||
inner.index = None;
|
inner.index = None;
|
||||||
inner.index_addr = None;
|
inner.index_addr = None;
|
||||||
inner.slots.clear();
|
inner.slots.clear();
|
||||||
|
inner.slot_index.clear();
|
||||||
inner.current_bytes = 0;
|
inner.current_bytes = 0;
|
||||||
inner.tick = 0;
|
inner.tick = 0;
|
||||||
inner.last_coord = None;
|
inner.last_coord = None;
|
||||||
@@ -607,11 +632,13 @@ impl ChunkCache {
|
|||||||
inner.chunk_layout = None;
|
inner.chunk_layout = None;
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Hint that the given chunk coordinates will be accessed soon.
|
/// Record that the given chunk coordinates are predicted to be accessed
|
||||||
|
/// soon (bookkeeping only).
|
||||||
///
|
///
|
||||||
/// Pre-populates the chunk index for these coordinates so that
|
/// This does **not** prefetch or pre-decompress anything — it only
|
||||||
/// subsequent lookups are O(1). This does NOT pre-decompress the
|
/// checks whether each coordinate is already in the chunk index and
|
||||||
/// chunks — it only ensures the index entries exist.
|
/// updates access-pattern stats accordingly. Real prefetching (e.g.
|
||||||
|
/// background pre-decompression) is not implemented.
|
||||||
pub fn prefetch_hint(&self, next_coords: &[ChunkCoord]) {
|
pub fn prefetch_hint(&self, next_coords: &[ChunkCoord]) {
|
||||||
let inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
|
let inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
|
||||||
if inner.index.is_none() {
|
if inner.index.is_none() {
|
||||||
@@ -785,6 +812,50 @@ mod tests {
|
|||||||
assert_eq!(cache.cached_bytes(), 3);
|
assert_eq!(cache.cached_bytes(), 3);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn slot_index_consistent_after_many_evictions() {
|
||||||
|
// Force repeated swap_remove evictions (small slot budget, many
|
||||||
|
// inserts) and confirm the coord -> slot index stays correct: every
|
||||||
|
// remaining coord must still resolve to its own data, not another
|
||||||
|
// slot's (which would happen if swap_remove's index fixup were wrong).
|
||||||
|
let cache = ChunkCache::with_capacity(1024 * 1024, 4); // max 4 slots
|
||||||
|
|
||||||
|
for i in 0..50u64 {
|
||||||
|
cache.put_decompressed(vec![i], vec![(i % 256) as u8; 8]);
|
||||||
|
// Interleave reads of a couple of earlier coords to churn LRU
|
||||||
|
// order (and thus which slot gets swap_remove'd) beyond simple
|
||||||
|
// FIFO eviction.
|
||||||
|
if i >= 2 {
|
||||||
|
let _ = cache.get_decompressed(&[i - 2]);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Whatever remains in the cache (at most 4 slots) must return its
|
||||||
|
// own correct data.
|
||||||
|
for i in 0..50u64 {
|
||||||
|
if let Some(data) = cache.get_decompressed(&[i]) {
|
||||||
|
assert_eq!(
|
||||||
|
data,
|
||||||
|
vec![(i % 256) as u8; 8],
|
||||||
|
"coord {i} returned wrong data after eviction churn"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
assert!(cache.cached_chunk_count() <= 4);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn get_decompressed_aligned_shares_arc_on_hit() {
|
||||||
|
let cache = ChunkCache::new();
|
||||||
|
cache.put_decompressed(vec![0, 0], vec![9, 9, 9, 9]);
|
||||||
|
let a = cache.get_decompressed_aligned(&[0, 0]).unwrap();
|
||||||
|
let b = cache.get_decompressed_aligned(&[0, 0]).unwrap();
|
||||||
|
// A cache hit clones the Arc (refcount bump), not the underlying
|
||||||
|
// buffer — both handles point at the same allocation.
|
||||||
|
assert!(Arc::ptr_eq(&a, &b));
|
||||||
|
assert_eq!(a.as_slice(), &[9, 9, 9, 9]);
|
||||||
|
}
|
||||||
|
|
||||||
// --- CacheAlignedBuffer tests ---
|
// --- CacheAlignedBuffer tests ---
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
|
|||||||
@@ -17,6 +17,8 @@ use crate::extensible_array::{ExtensibleArrayHeader, read_extensible_array_chunk
|
|||||||
use crate::filter_pipeline::FilterPipeline;
|
use crate::filter_pipeline::FilterPipeline;
|
||||||
use crate::filters::decompress_chunk;
|
use crate::filters::decompress_chunk;
|
||||||
use crate::fixed_array::{FixedArrayHeader, read_fixed_array_chunks};
|
use crate::fixed_array::{FixedArrayHeader, read_fixed_array_chunks};
|
||||||
|
#[cfg(feature = "std")]
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
#[cfg(feature = "parallel")]
|
#[cfg(feature = "parallel")]
|
||||||
use crate::parallel_read;
|
use crate::parallel_read;
|
||||||
@@ -689,7 +691,7 @@ pub fn read_chunked_data_cached(
|
|||||||
let coord: Vec<u64> = chunk_info.offsets.iter().take(rank).copied().collect();
|
let coord: Vec<u64> = chunk_info.offsets.iter().take(rank).copied().collect();
|
||||||
|
|
||||||
// Try decompressed cache first
|
// Try decompressed cache first
|
||||||
let decompressed = if let Some(cached) = cache.get_decompressed(&coord) {
|
let decompressed = if let Some(cached) = cache.get_decompressed_aligned(&coord) {
|
||||||
cached
|
cached
|
||||||
} else {
|
} else {
|
||||||
// Decompress from file
|
// Decompress from file
|
||||||
@@ -711,8 +713,7 @@ pub fn read_chunked_data_cached(
|
|||||||
} else {
|
} else {
|
||||||
raw_chunk.to_vec()
|
raw_chunk.to_vec()
|
||||||
};
|
};
|
||||||
cache.put_decompressed(coord, dec.clone());
|
cache.put_decompressed(coord, dec)
|
||||||
dec
|
|
||||||
};
|
};
|
||||||
|
|
||||||
let chunk_offsets: Vec<usize> = chunk_info
|
let chunk_offsets: Vec<usize> = chunk_info
|
||||||
@@ -1055,7 +1056,7 @@ pub fn read_chunked_data_sweep(
|
|||||||
}
|
}
|
||||||
|
|
||||||
// Try decompressed cache first
|
// Try decompressed cache first
|
||||||
let decompressed = if let Some(cached) = cache.get_decompressed(&coord) {
|
let decompressed = if let Some(cached) = cache.get_decompressed_aligned(&coord) {
|
||||||
cached
|
cached
|
||||||
} else {
|
} else {
|
||||||
// Decompress from file
|
// Decompress from file
|
||||||
@@ -1077,8 +1078,7 @@ pub fn read_chunked_data_sweep(
|
|||||||
} else {
|
} else {
|
||||||
raw_chunk.to_vec()
|
raw_chunk.to_vec()
|
||||||
};
|
};
|
||||||
cache.put_decompressed(coord, dec.clone());
|
cache.put_decompressed(coord, dec)
|
||||||
dec
|
|
||||||
};
|
};
|
||||||
|
|
||||||
let chunk_offsets: Vec<usize> = chunk_info
|
let chunk_offsets: Vec<usize> = chunk_info
|
||||||
@@ -1271,7 +1271,7 @@ pub fn read_chunked_data_indexed(
|
|||||||
.ok_or_else(|| FormatError::ChunkedReadError("chunk layout not available".into()))?;
|
.ok_or_else(|| FormatError::ChunkedReadError("chunk layout not available".into()))?;
|
||||||
|
|
||||||
// Decompress chunks (using LRU cache where possible)
|
// Decompress chunks (using LRU cache where possible)
|
||||||
let mut chunk_buffers: Vec<CacheAlignedBuffer> = Vec::with_capacity(mappings_info.len());
|
let mut chunk_buffers: Vec<Arc<CacheAlignedBuffer>> = Vec::with_capacity(mappings_info.len());
|
||||||
for (coord, file_offset, file_size, filter_mask) in &mappings_info {
|
for (coord, file_offset, file_size, filter_mask) in &mappings_info {
|
||||||
if let Some(cached) = cache.get_decompressed_aligned(coord) {
|
if let Some(cached) = cache.get_decompressed_aligned(coord) {
|
||||||
chunk_buffers.push(cached);
|
chunk_buffers.push(cached);
|
||||||
@@ -1295,8 +1295,8 @@ pub fn read_chunked_data_indexed(
|
|||||||
raw_chunk.to_vec()
|
raw_chunk.to_vec()
|
||||||
};
|
};
|
||||||
let aligned = CacheAlignedBuffer::from_vec(decompressed);
|
let aligned = CacheAlignedBuffer::from_vec(decompressed);
|
||||||
cache.put_decompressed_aligned(coord.clone(), aligned.clone());
|
let arc = cache.put_decompressed_aligned(coord.clone(), aligned);
|
||||||
chunk_buffers.push(aligned);
|
chunk_buffers.push(arc);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -65,10 +65,8 @@ impl ChunkOptions {
|
|||||||
pub fn build_pipeline(&self, element_size: u32) -> Option<FilterPipeline> {
|
pub fn build_pipeline(&self, element_size: u32) -> Option<FilterPipeline> {
|
||||||
let mut filters = Vec::new();
|
let mut filters = Vec::new();
|
||||||
|
|
||||||
let has_compression = self.deflate_level.is_some()
|
let has_compression =
|
||||||
|| self.zstd_level.is_some()
|
self.deflate_level.is_some() || self.zstd_level.is_some() || self.lz4 || self.pcodec;
|
||||||
|| self.lz4
|
|
||||||
|| self.pcodec;
|
|
||||||
|
|
||||||
// Shuffle before compression. Applied if explicitly requested OR if compression
|
// Shuffle before compression. Applied if explicitly requested OR if compression
|
||||||
// is active and the caller hasn't disabled it — matches h5py default behavior
|
// is active and the caller hasn't disabled it — matches h5py default behavior
|
||||||
@@ -608,7 +606,7 @@ pub fn precompress_chunks(
|
|||||||
|
|
||||||
let chunks = raw_chunks
|
let chunks = raw_chunks
|
||||||
.into_iter()
|
.into_iter()
|
||||||
.zip(compressed.into_iter())
|
.zip(compressed)
|
||||||
.map(|((_offsets, raw_bytes), c)| (raw_bytes.len() as u64, c))
|
.map(|((_offsets, raw_bytes), c)| (raw_bytes.len() as u64, c))
|
||||||
.collect();
|
.collect();
|
||||||
|
|
||||||
@@ -755,7 +753,11 @@ pub fn build_chunked_data_at_ext(
|
|||||||
maxshape: Option<&[u64]>,
|
maxshape: Option<&[u64]>,
|
||||||
) -> Result<ChunkedDataResult, FormatError> {
|
) -> Result<ChunkedDataResult, FormatError> {
|
||||||
let pre = precompress_chunks(raw_data, shape, chunk_dims, element_size, options)?;
|
let pre = precompress_chunks(raw_data, shape, chunk_dims, element_size, options)?;
|
||||||
Ok(build_chunked_data_from_precompressed(&pre, base_address, maxshape))
|
Ok(build_chunked_data_from_precompressed(
|
||||||
|
&pre,
|
||||||
|
base_address,
|
||||||
|
maxshape,
|
||||||
|
))
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Write selected elements into an existing in-memory dataset buffer.
|
/// Write selected elements into an existing in-memory dataset buffer.
|
||||||
|
|||||||
@@ -821,7 +821,8 @@ mod tests {
|
|||||||
let blob = [
|
let blob = [
|
||||||
0x00u8, // block version 0
|
0x00u8, // block version 0
|
||||||
0x01, 0, 0, 0, 0, 0, 0, 0, // nused = 1
|
0x01, 0, 0, 0, 0, 0, 0, 0, // nused = 1
|
||||||
0x73, 0x72, 0x63, 0x5f, 0x65, 0x78, 0x74, 0x2e, 0x68, 0x35, 0x00, // "src_ext.h5\0"
|
0x73, 0x72, 0x63, 0x5f, 0x65, 0x78, 0x74, 0x2e, 0x68, 0x35,
|
||||||
|
0x00, // "src_ext.h5\0"
|
||||||
0x64, 0x61, 0x74, 0x61, 0x00, // "data\0"
|
0x64, 0x61, 0x74, 0x61, 0x00, // "data\0"
|
||||||
0x03, 0, 0, 0, 0x01, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, // source sel = ALL
|
0x03, 0, 0, 0, 0x01, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, // source sel = ALL
|
||||||
0x03, 0, 0, 0, 0x01, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, // virtual sel = ALL
|
0x03, 0, 0, 0, 0x01, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, // virtual sel = ALL
|
||||||
@@ -846,7 +847,7 @@ mod tests {
|
|||||||
// One entry whose source selection (ALL) is truncated to 8 of 16 bytes.
|
// One entry whose source selection (ALL) is truncated to 8 of 16 bytes.
|
||||||
let blob = [
|
let blob = [
|
||||||
0x01u8, // version 1
|
0x01u8, // version 1
|
||||||
0x01, 0, 0, 0, 0, 0, 0, 0, // nused = 1
|
0x01, 0, 0, 0, 0, 0, 0, 0, // nused = 1
|
||||||
0x04, // same-file marker
|
0x04, // same-file marker
|
||||||
0x78, 0x00, // "x\0"
|
0x78, 0x00, // "x\0"
|
||||||
0x03, 0, 0, 0, 0x01, 0, 0, 0, // ALL header, truncated (8 of 16 bytes)
|
0x03, 0, 0, 0, 0x01, 0, 0, 0, // ALL header, truncated (8 of 16 bytes)
|
||||||
|
|||||||
@@ -94,7 +94,14 @@ pub fn read_raw_data_full(
|
|||||||
length_size: u8,
|
length_size: u8,
|
||||||
) -> Result<Vec<u8>, FormatError> {
|
) -> Result<Vec<u8>, FormatError> {
|
||||||
read_raw_data_full_impl(
|
read_raw_data_full_impl(
|
||||||
file_data, layout, dataspace, datatype, pipeline, offset_size, length_size, None,
|
file_data,
|
||||||
|
layout,
|
||||||
|
dataspace,
|
||||||
|
datatype,
|
||||||
|
pipeline,
|
||||||
|
offset_size,
|
||||||
|
length_size,
|
||||||
|
None,
|
||||||
)
|
)
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -112,7 +119,14 @@ pub fn read_raw_data_full_with_resolver(
|
|||||||
resolver: Option<&VdsSourceResolver>,
|
resolver: Option<&VdsSourceResolver>,
|
||||||
) -> Result<Vec<u8>, FormatError> {
|
) -> Result<Vec<u8>, FormatError> {
|
||||||
read_raw_data_full_impl(
|
read_raw_data_full_impl(
|
||||||
file_data, layout, dataspace, datatype, pipeline, offset_size, length_size, resolver,
|
file_data,
|
||||||
|
layout,
|
||||||
|
dataspace,
|
||||||
|
datatype,
|
||||||
|
pipeline,
|
||||||
|
offset_size,
|
||||||
|
length_size,
|
||||||
|
resolver,
|
||||||
)
|
)
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -465,12 +479,12 @@ fn read_virtual_data(
|
|||||||
FormatError::ChunkedReadError("virtual dataset has no mapping global heap".into())
|
FormatError::ChunkedReadError("virtual dataset has no mapping global heap".into())
|
||||||
})?;
|
})?;
|
||||||
let coll = GlobalHeapCollection::parse(file_data, addr as usize, length_size)?;
|
let coll = GlobalHeapCollection::parse(file_data, addr as usize, length_size)?;
|
||||||
let obj = coll
|
let obj =
|
||||||
.get_object(global_heap_index as u16)
|
coll.get_object(global_heap_index as u16)
|
||||||
.ok_or(FormatError::GlobalHeapObjectNotFound {
|
.ok_or(FormatError::GlobalHeapObjectNotFound {
|
||||||
collection_address: addr,
|
collection_address: addr,
|
||||||
index: global_heap_index as u16,
|
index: global_heap_index as u16,
|
||||||
})?;
|
})?;
|
||||||
let mappings = parse_vds_mappings(&obj.data, length_size)?;
|
let mappings = parse_vds_mappings(&obj.data, length_size)?;
|
||||||
|
|
||||||
for m in &mappings {
|
for m in &mappings {
|
||||||
@@ -1815,13 +1829,19 @@ mod tests {
|
|||||||
0xff, 0xff, 0x00, 0x00, 0x64, 0x00, 0x00, 0x00, 0xe8, 0x03, 0x00, 0x00, 0x00, 0x80,
|
0xff, 0xff, 0x00, 0x00, 0x64, 0x00, 0x00, 0x00, 0xe8, 0x03, 0x00, 0x00, 0x00, 0x80,
|
||||||
0x00, 0x00,
|
0x00, 0x00,
|
||||||
];
|
];
|
||||||
assert_eq!(read_as_i32(&raw, &arr).unwrap(), vec![-1, 100, 1000, -32768]);
|
assert_eq!(
|
||||||
|
read_as_i32(&raw, &arr).unwrap(),
|
||||||
|
vec![-1, 100, 1000, -32768]
|
||||||
|
);
|
||||||
// Nested array-of-array unwraps recursively.
|
// Nested array-of-array unwraps recursively.
|
||||||
let nested = Datatype::Array {
|
let nested = Datatype::Array {
|
||||||
base_type: Box::new(arr),
|
base_type: Box::new(arr),
|
||||||
dimensions: vec![2],
|
dimensions: vec![2],
|
||||||
};
|
};
|
||||||
assert_eq!(read_as_i32(&raw, &nested).unwrap(), vec![-1, 100, 1000, -32768]);
|
assert_eq!(
|
||||||
|
read_as_i32(&raw, &nested).unwrap(),
|
||||||
|
vec![-1, 100, 1000, -32768]
|
||||||
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
fn make_f64_le_type() -> Datatype {
|
fn make_f64_le_type() -> Datatype {
|
||||||
|
|||||||
@@ -1030,14 +1030,8 @@ mod tests {
|
|||||||
Datatype::Compound { size, members } => {
|
Datatype::Compound { size, members } => {
|
||||||
assert_eq!(size, 20);
|
assert_eq!(size, 20);
|
||||||
assert_eq!(members.len(), 3);
|
assert_eq!(members.len(), 3);
|
||||||
assert_eq!(
|
assert_eq!((members[0].name.as_str(), members[0].byte_offset), ("x", 0));
|
||||||
(members[0].name.as_str(), members[0].byte_offset),
|
assert_eq!((members[1].name.as_str(), members[1].byte_offset), ("y", 8));
|
||||||
("x", 0)
|
|
||||||
);
|
|
||||||
assert_eq!(
|
|
||||||
(members[1].name.as_str(), members[1].byte_offset),
|
|
||||||
("y", 8)
|
|
||||||
);
|
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
(members[2].name.as_str(), members[2].byte_offset),
|
(members[2].name.as_str(), members[2].byte_offset),
|
||||||
("id", 16)
|
("id", 16)
|
||||||
@@ -1076,7 +1070,10 @@ mod tests {
|
|||||||
dimensions,
|
dimensions,
|
||||||
} => {
|
} => {
|
||||||
assert_eq!(dimensions, vec![3]);
|
assert_eq!(dimensions, vec![3]);
|
||||||
assert!(matches!(*base_type, Datatype::FloatingPoint { size: 8, .. }));
|
assert!(matches!(
|
||||||
|
*base_type,
|
||||||
|
Datatype::FloatingPoint { size: 8, .. }
|
||||||
|
));
|
||||||
}
|
}
|
||||||
other => panic!("expected Array, got {other:?}"),
|
other => panic!("expected Array, got {other:?}"),
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -20,7 +20,7 @@
|
|||||||
//! ```
|
//! ```
|
||||||
|
|
||||||
#[cfg(not(feature = "std"))]
|
#[cfg(not(feature = "std"))]
|
||||||
use alloc::{string::String, vec, vec::Vec};
|
use alloc::{format, string::String, vec, vec::Vec};
|
||||||
|
|
||||||
#[cfg(not(feature = "std"))]
|
#[cfg(not(feature = "std"))]
|
||||||
use alloc::collections::BTreeMap;
|
use alloc::collections::BTreeMap;
|
||||||
|
|||||||
@@ -375,7 +375,8 @@ fn build_multiblock_fractal_heap(
|
|||||||
let table_width: u16 = 4;
|
let table_width: u16 = 4;
|
||||||
let starting_block_size: u64 = 512;
|
let starting_block_size: u64 = 512;
|
||||||
let dblock_header_size = 4 + 1 + os + block_offset_bytes + 4;
|
let dblock_header_size = 4 + 1 + os + block_offset_bytes + 4;
|
||||||
let block_capacity = |row: usize| block_size_for_row(starting_block_size, row) - dblock_header_size as u64;
|
let block_capacity =
|
||||||
|
|row: usize| block_size_for_row(starting_block_size, row) - dblock_header_size as u64;
|
||||||
|
|
||||||
// ---- Pack objects into direct blocks (row-major over the doubling table) ----
|
// ---- Pack objects into direct blocks (row-major over the doubling table) ----
|
||||||
struct Blk {
|
struct Blk {
|
||||||
@@ -542,7 +543,26 @@ fn block_size_for_row(starting_block_size: u64, row: usize) -> u64 {
|
|||||||
|
|
||||||
/// Size in bytes of the FRHP header for the given offset/length sizes.
|
/// Size in bytes of the FRHP header for the given offset/length sizes.
|
||||||
fn frhp_header_size(os: usize, ls: usize) -> usize {
|
fn frhp_header_size(os: usize, ls: usize) -> usize {
|
||||||
4 + 1 + 2 + 2 + 1 + 4 + ls + os + ls + os + ls + ls + ls + ls + ls + ls + ls + ls + 2 + ls + ls
|
4 + 1
|
||||||
|
+ 2
|
||||||
|
+ 2
|
||||||
|
+ 1
|
||||||
|
+ 4
|
||||||
|
+ ls
|
||||||
|
+ os
|
||||||
|
+ ls
|
||||||
|
+ os
|
||||||
|
+ ls
|
||||||
|
+ ls
|
||||||
|
+ ls
|
||||||
|
+ ls
|
||||||
|
+ ls
|
||||||
|
+ ls
|
||||||
|
+ ls
|
||||||
|
+ ls
|
||||||
|
+ 2
|
||||||
|
+ ls
|
||||||
|
+ ls
|
||||||
+ 2
|
+ 2
|
||||||
+ 2
|
+ 2
|
||||||
+ os
|
+ os
|
||||||
@@ -670,8 +690,7 @@ pub(crate) fn build_dense_attrs(attrs: &[AttributeMessage], base_address: u64) -
|
|||||||
// Pad to node_size
|
// Pad to node_size
|
||||||
btlf.resize(node_size as usize, 0);
|
btlf.resize(node_size as usize, 0);
|
||||||
|
|
||||||
let mut blob =
|
let mut blob = Vec::with_capacity(heap.blob.len() + bthd.len() + btlf.len());
|
||||||
Vec::with_capacity(heap.blob.len() + bthd.len() + btlf.len());
|
|
||||||
blob.extend_from_slice(&heap.blob);
|
blob.extend_from_slice(&heap.blob);
|
||||||
blob.extend_from_slice(&bthd);
|
blob.extend_from_slice(&bthd);
|
||||||
blob.extend_from_slice(&btlf);
|
blob.extend_from_slice(&btlf);
|
||||||
@@ -762,8 +781,7 @@ pub(crate) fn build_dense_links(links: &[LinkMessage], base_address: u64) -> Den
|
|||||||
btlf.extend_from_slice(&btlf_checksum.to_le_bytes());
|
btlf.extend_from_slice(&btlf_checksum.to_le_bytes());
|
||||||
btlf.resize(node_size as usize, 0);
|
btlf.resize(node_size as usize, 0);
|
||||||
|
|
||||||
let mut blob =
|
let mut blob = Vec::with_capacity(heap.blob.len() + bthd.len() + btlf.len());
|
||||||
Vec::with_capacity(heap.blob.len() + bthd.len() + btlf.len());
|
|
||||||
blob.extend_from_slice(&heap.blob);
|
blob.extend_from_slice(&heap.blob);
|
||||||
blob.extend_from_slice(&bthd);
|
blob.extend_from_slice(&bthd);
|
||||||
blob.extend_from_slice(&btlf);
|
blob.extend_from_slice(&btlf);
|
||||||
@@ -1096,10 +1114,7 @@ impl FileWriter {
|
|||||||
root_attrs.push(build_attr_message(n, v));
|
root_attrs.push(build_attr_message(n, v));
|
||||||
}
|
}
|
||||||
|
|
||||||
let is_vds: Vec<bool> = all_ds
|
let is_vds: Vec<bool> = all_ds.iter().map(|d| d.virtual_sources.is_some()).collect();
|
||||||
.iter()
|
|
||||||
.map(|d| d.virtual_sources.is_some())
|
|
||||||
.collect();
|
|
||||||
let is_chunked: Vec<bool> = all_ds
|
let is_chunked: Vec<bool> = all_ds
|
||||||
.iter()
|
.iter()
|
||||||
.enumerate()
|
.enumerate()
|
||||||
@@ -1198,7 +1213,8 @@ impl FileWriter {
|
|||||||
// Global heap blob size is address-independent; compute it now
|
// Global heap blob size is address-independent; compute it now
|
||||||
// so pass 2 can place it correctly.
|
// so pass 2 can place it correctly.
|
||||||
let vds_mappings = d.virtual_sources.as_deref().unwrap_or(&[]);
|
let vds_mappings = d.virtual_sources.as_deref().unwrap_or(&[]);
|
||||||
let gcol_bytes = build_global_heap_collection(&serialize_vds_mappings(vds_mappings));
|
let gcol_bytes =
|
||||||
|
build_global_heap_collection(&serialize_vds_mappings(vds_mappings));
|
||||||
dummy_blobs.push(DataBlob {
|
dummy_blobs.push(DataBlob {
|
||||||
data: gcol_bytes, // store heap blob here temporarily
|
data: gcol_bytes, // store heap blob here temporarily
|
||||||
oh_bytes: oh,
|
oh_bytes: oh,
|
||||||
@@ -1216,8 +1232,11 @@ impl FileWriter {
|
|||||||
elem_size,
|
elem_size,
|
||||||
&d.chunk_options,
|
&d.chunk_options,
|
||||||
)?;
|
)?;
|
||||||
let result =
|
let result = build_chunked_data_from_precompressed(
|
||||||
build_chunked_data_from_precompressed(&pre, dummy_cursor, d.maxshape.as_deref());
|
&pre,
|
||||||
|
dummy_cursor,
|
||||||
|
d.maxshape.as_deref(),
|
||||||
|
);
|
||||||
dummy_cursor += result.data_bytes.len() as u64;
|
dummy_cursor += result.data_bytes.len() as u64;
|
||||||
let dense_blob = if ds_dense[i] {
|
let dense_blob = if ds_dense[i] {
|
||||||
Some(build_dense_attrs(&d.attrs, 0))
|
Some(build_dense_attrs(&d.attrs, 0))
|
||||||
@@ -1391,7 +1410,10 @@ impl FileWriter {
|
|||||||
// Reuse precompressed chunks from Pass 1 — avoids re-compressing
|
// Reuse precompressed chunks from Pass 1 — avoids re-compressing
|
||||||
// the same data a second time.
|
// the same data a second time.
|
||||||
let result = build_chunked_data_from_precompressed(
|
let result = build_chunked_data_from_precompressed(
|
||||||
dummy_blobs[i].precompressed.as_ref().expect("chunked dataset missing precompressed cache"),
|
dummy_blobs[i]
|
||||||
|
.precompressed
|
||||||
|
.as_ref()
|
||||||
|
.expect("chunked dataset missing precompressed cache"),
|
||||||
base_address,
|
base_address,
|
||||||
d.maxshape.as_deref(),
|
d.maxshape.as_deref(),
|
||||||
);
|
);
|
||||||
@@ -1492,7 +1514,9 @@ impl FileWriter {
|
|||||||
// Rebuild the root link blob with real target addresses (same size as
|
// Rebuild the root link blob with real target addresses (same size as
|
||||||
// the dummy used for layout); its LinkInfo goes in the OH.
|
// the dummy used for layout); its LinkInfo goes in the OH.
|
||||||
let root_link_blob = root_link_blob_addr.map(|addr| build_dense_links(&root_links, addr));
|
let root_link_blob = root_link_blob_addr.map(|addr| build_dense_links(&root_links, addr));
|
||||||
let root_dl = root_link_blob.as_ref().map(|b| b.link_info_message.as_slice());
|
let root_dl = root_link_blob
|
||||||
|
.as_ref()
|
||||||
|
.map(|b| b.link_info_message.as_slice());
|
||||||
buf.extend_from_slice(&build_group_oh(
|
buf.extend_from_slice(&build_group_oh(
|
||||||
&root_links,
|
&root_links,
|
||||||
root_dl,
|
root_dl,
|
||||||
@@ -1903,14 +1927,14 @@ mod tests {
|
|||||||
fn sel_hyper_1d(start: u16, block: u16) -> Vec<u8> {
|
fn sel_hyper_1d(start: u16, block: u16) -> Vec<u8> {
|
||||||
let mut v = vec![
|
let mut v = vec![
|
||||||
2, 0, 0, 0, // type = HYPER
|
2, 0, 0, 0, // type = HYPER
|
||||||
3, 0, 0, 0, // version 3
|
3, 0, 0, 0, // version 3
|
||||||
0x01, // flags = regular
|
0x01, // flags = regular
|
||||||
0x02, // enc_size = 2 (u16 per coordinate)
|
0x02, // enc_size = 2 (u16 per coordinate)
|
||||||
1, 0, 0, 0, // rank = 1
|
1, 0, 0, 0, // rank = 1
|
||||||
];
|
];
|
||||||
v.extend_from_slice(&start.to_le_bytes()); // start
|
v.extend_from_slice(&start.to_le_bytes()); // start
|
||||||
v.extend_from_slice(&1u16.to_le_bytes()); // stride
|
v.extend_from_slice(&1u16.to_le_bytes()); // stride
|
||||||
v.extend_from_slice(&1u16.to_le_bytes()); // count
|
v.extend_from_slice(&1u16.to_le_bytes()); // count
|
||||||
v.extend_from_slice(&block.to_le_bytes()); // block
|
v.extend_from_slice(&block.to_le_bytes()); // block
|
||||||
v
|
v
|
||||||
}
|
}
|
||||||
@@ -1936,8 +1960,10 @@ mod tests {
|
|||||||
|
|
||||||
let mut fw = FileWriter::new();
|
let mut fw = FileWriter::new();
|
||||||
// Source datasets (real data in this file)
|
// Source datasets (real data in this file)
|
||||||
fw.create_dataset("src_a").with_f64_data(&[1.0, 2.0, 3.0, 4.0]);
|
fw.create_dataset("src_a")
|
||||||
fw.create_dataset("src_b").with_f64_data(&[5.0, 6.0, 7.0, 8.0]);
|
.with_f64_data(&[1.0, 2.0, 3.0, 4.0]);
|
||||||
|
fw.create_dataset("src_b")
|
||||||
|
.with_f64_data(&[5.0, 6.0, 7.0, 8.0]);
|
||||||
// Virtual dataset
|
// Virtual dataset
|
||||||
fw.create_dataset("vds")
|
fw.create_dataset("vds")
|
||||||
.with_shape(&[8])
|
.with_shape(&[8])
|
||||||
@@ -1950,13 +1976,8 @@ mod tests {
|
|||||||
let sig = signature::find_signature(&bytes).unwrap();
|
let sig = signature::find_signature(&bytes).unwrap();
|
||||||
let sb = Superblock::parse(&bytes, sig).unwrap();
|
let sb = Superblock::parse(&bytes, sig).unwrap();
|
||||||
let vds_addr = resolve_path_any(&bytes, &sb, "vds").unwrap();
|
let vds_addr = resolve_path_any(&bytes, &sb, "vds").unwrap();
|
||||||
let hdr = ObjectHeader::parse(
|
let hdr =
|
||||||
&bytes,
|
ObjectHeader::parse(&bytes, vds_addr as usize, sb.offset_size, sb.length_size).unwrap();
|
||||||
vds_addr as usize,
|
|
||||||
sb.offset_size,
|
|
||||||
sb.length_size,
|
|
||||||
)
|
|
||||||
.unwrap();
|
|
||||||
|
|
||||||
let dl_data = &hdr
|
let dl_data = &hdr
|
||||||
.messages
|
.messages
|
||||||
@@ -1965,8 +1986,7 @@ mod tests {
|
|||||||
.unwrap()
|
.unwrap()
|
||||||
.data;
|
.data;
|
||||||
|
|
||||||
let mut layout =
|
let mut layout = DataLayout::parse(dl_data, sb.offset_size, sb.length_size).unwrap();
|
||||||
DataLayout::parse(dl_data, sb.offset_size, sb.length_size).unwrap();
|
|
||||||
|
|
||||||
// Before resolution, mappings field is empty.
|
// Before resolution, mappings field is empty.
|
||||||
assert!(
|
assert!(
|
||||||
@@ -2033,8 +2053,7 @@ mod tests {
|
|||||||
.find(|m| m.msg_type == MessageType::DataLayout)
|
.find(|m| m.msg_type == MessageType::DataLayout)
|
||||||
.unwrap()
|
.unwrap()
|
||||||
.data;
|
.data;
|
||||||
let mut layout =
|
let mut layout = DataLayout::parse(dl_data, sb.offset_size, sb.length_size).unwrap();
|
||||||
DataLayout::parse(dl_data, sb.offset_size, sb.length_size).unwrap();
|
|
||||||
layout.resolve_vds_mappings(&bytes, sb.length_size).unwrap();
|
layout.resolve_vds_mappings(&bytes, sb.length_size).unwrap();
|
||||||
|
|
||||||
match &layout {
|
match &layout {
|
||||||
@@ -2112,7 +2131,10 @@ mod tests {
|
|||||||
.expect("external link 'remote_temp' not found in group OH");
|
.expect("external link 'remote_temp' not found in group OH");
|
||||||
|
|
||||||
match &ext_link.link_target {
|
match &ext_link.link_target {
|
||||||
crate::link_message::LinkTarget::External { filename, object_path } => {
|
crate::link_message::LinkTarget::External {
|
||||||
|
filename,
|
||||||
|
object_path,
|
||||||
|
} => {
|
||||||
assert_eq!(filename, "other_file.h5");
|
assert_eq!(filename, "other_file.h5");
|
||||||
assert_eq!(object_path, "/temperature");
|
assert_eq!(object_path, "/temperature");
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -4,14 +4,19 @@
|
|||||||
extern crate alloc;
|
extern crate alloc;
|
||||||
|
|
||||||
#[cfg(not(feature = "std"))]
|
#[cfg(not(feature = "std"))]
|
||||||
use alloc::{vec, vec::Vec};
|
use alloc::{boxed::Box, vec, vec::Vec};
|
||||||
|
|
||||||
use crate::error::FormatError;
|
use crate::error::FormatError;
|
||||||
use crate::filter_pipeline::{
|
use crate::filter_pipeline::{
|
||||||
FILTER_DEFLATE, FILTER_FLETCHER32, FILTER_LZ4, FILTER_NBIT, FILTER_PCODEC,
|
FILTER_DEFLATE, FILTER_FLETCHER32, FILTER_LZ4, FILTER_NBIT, FILTER_PCODEC, FILTER_SCALEOFFSET,
|
||||||
FILTER_SCALEOFFSET, FILTER_SHUFFLE, FILTER_SZIP, FILTER_ZSTD, FilterPipeline,
|
FILTER_SHUFFLE, FILTER_SZIP, FILTER_ZSTD, FilterPipeline,
|
||||||
};
|
};
|
||||||
|
|
||||||
|
/// Absolute ceiling on a single decompressed chunk's output size, used only
|
||||||
|
/// when the pipeline's declared `chunk_size` is unavailable (0). Prevents
|
||||||
|
/// unbounded-allocation DoS from a malicious/corrupted compressed chunk.
|
||||||
|
pub(crate) const MAX_DECOMPRESS_SIZE: usize = 256 * 1024 * 1024;
|
||||||
|
|
||||||
/// Apply a filter pipeline to decompress a chunk.
|
/// Apply a filter pipeline to decompress a chunk.
|
||||||
/// Filters are applied in REVERSE order for decompression.
|
/// Filters are applied in REVERSE order for decompression.
|
||||||
pub fn decompress_chunk(
|
pub fn decompress_chunk(
|
||||||
@@ -25,16 +30,22 @@ pub fn decompress_chunk(
|
|||||||
for filter in pipeline.filters.iter().rev() {
|
for filter in pipeline.filters.iter().rev() {
|
||||||
data = match filter.filter_id {
|
data = match filter.filter_id {
|
||||||
FILTER_SHUFFLE => shuffle_decompress(&data, element_size as usize)?,
|
FILTER_SHUFFLE => shuffle_decompress(&data, element_size as usize)?,
|
||||||
FILTER_DEFLATE => deflate_decompress(&data)?,
|
// `chunk_size` is the expected decompressed size (shuffle/fletcher32
|
||||||
FILTER_LZ4 => lz4_decompress(&data)?,
|
// are size-preserving, so it bounds these too); pass it so these
|
||||||
FILTER_ZSTD => zstd_decompress(&data)?,
|
// decoders can't be forced into unbounded allocation by a hostile
|
||||||
|
// or corrupted compressed payload.
|
||||||
|
FILTER_DEFLATE => deflate_decompress(&data, chunk_size)?,
|
||||||
|
FILTER_LZ4 => lz4_decompress(&data, chunk_size)?,
|
||||||
|
FILTER_ZSTD => zstd_decompress(&data, chunk_size)?,
|
||||||
FILTER_FLETCHER32 => fletcher32_verify(&data)?,
|
FILTER_FLETCHER32 => fletcher32_verify(&data)?,
|
||||||
FILTER_PCODEC => pcodec_decompress(&data, element_size as usize)?,
|
FILTER_PCODEC => pcodec_decompress(&data, element_size as usize, chunk_size)?,
|
||||||
// `chunk_size` is the expected decompressed size; pass it so these
|
// `chunk_size` is the expected decompressed size; pass it so these
|
||||||
// decoders can reject an element count that would over-allocate.
|
// decoders can reject an element count that would over-allocate.
|
||||||
FILTER_SCALEOFFSET => scaleoffset_decompress(&data, &filter.client_data, chunk_size)?,
|
FILTER_SCALEOFFSET => scaleoffset_decompress(&data, &filter.client_data, chunk_size)?,
|
||||||
FILTER_NBIT => nbit_decompress(&data, &filter.client_data, chunk_size)?,
|
FILTER_NBIT => nbit_decompress(&data, &filter.client_data, chunk_size)?,
|
||||||
FILTER_SZIP => crate::filters_szip::szip_decompress(&data, &filter.client_data, chunk_size)?,
|
FILTER_SZIP => {
|
||||||
|
crate::filters_szip::szip_decompress(&data, &filter.client_data, chunk_size)?
|
||||||
|
}
|
||||||
other => return Err(FormatError::UnsupportedFilter(other)),
|
other => return Err(FormatError::UnsupportedFilter(other)),
|
||||||
};
|
};
|
||||||
}
|
}
|
||||||
@@ -89,6 +100,27 @@ pub fn compress_chunk(
|
|||||||
/// E-scale, interpreted as i32 for negative exponents), `[2]`=element count,
|
/// E-scale, interpreted as i32 for negative exponents), `[2]`=element count,
|
||||||
/// `[4]`=element size, `[5]`=signed flag, `[6]`=byte order (1 = big-endian),
|
/// `[4]`=element size, `[5]`=signed flag, `[6]`=byte order (1 = big-endian),
|
||||||
/// `[7]`=fill defined, `[8..]`=fill value bits.
|
/// `[7]`=fill defined, `[8..]`=fill value bits.
|
||||||
|
/// `f64::powi` equivalent that works under `no_std` (no libm/std available).
|
||||||
|
/// Exponentiation by squaring, matching `powi`'s semantics for negative
|
||||||
|
/// exponents via reciprocal.
|
||||||
|
fn powi_f64(base: f64, mut exp: i32) -> f64 {
|
||||||
|
let neg = exp < 0;
|
||||||
|
if neg {
|
||||||
|
exp = -exp;
|
||||||
|
}
|
||||||
|
let mut result = 1.0f64;
|
||||||
|
let mut b = base;
|
||||||
|
let mut e = exp as u32;
|
||||||
|
while e > 0 {
|
||||||
|
if e & 1 == 1 {
|
||||||
|
result *= b;
|
||||||
|
}
|
||||||
|
b *= b;
|
||||||
|
e >>= 1;
|
||||||
|
}
|
||||||
|
if neg { 1.0 / result } else { result }
|
||||||
|
}
|
||||||
|
|
||||||
fn scaleoffset_decompress(
|
fn scaleoffset_decompress(
|
||||||
data: &[u8],
|
data: &[u8],
|
||||||
cd: &[u32],
|
cd: &[u32],
|
||||||
@@ -207,9 +239,9 @@ fn scaleoffset_decompress(
|
|||||||
if has_fill_code && code == fill_code {
|
if has_fill_code && code == fill_code {
|
||||||
fill_value
|
fill_value
|
||||||
} else if is_escale {
|
} else if is_escale {
|
||||||
minval + code as f64 * 2f64.powi(scale_factor)
|
minval + code as f64 * powi_f64(2.0, scale_factor)
|
||||||
} else {
|
} else {
|
||||||
minval + code as f64 / 10f64.powi(scale_factor)
|
minval + code as f64 / powi_f64(10.0, scale_factor)
|
||||||
}
|
}
|
||||||
})
|
})
|
||||||
.collect();
|
.collect();
|
||||||
@@ -569,24 +601,48 @@ fn nbit_decompress(data: &[u8], cd: &[u32], expected_bytes: usize) -> Result<Vec
|
|||||||
}
|
}
|
||||||
|
|
||||||
/// Decompress zlib-compressed data.
|
/// Decompress zlib-compressed data.
|
||||||
|
///
|
||||||
|
/// `expected_bytes` is the pipeline's declared decompressed chunk size (0 if
|
||||||
|
/// unavailable); output is rejected if it exceeds this bound (or, when
|
||||||
|
/// unavailable, [`MAX_DECOMPRESS_SIZE`]), preventing a hostile/corrupted
|
||||||
|
/// compressed payload from forcing unbounded allocation (a "zlib bomb").
|
||||||
#[cfg(feature = "deflate")]
|
#[cfg(feature = "deflate")]
|
||||||
fn deflate_decompress(data: &[u8]) -> Result<Vec<u8>, FormatError> {
|
fn deflate_decompress(data: &[u8], expected_bytes: usize) -> Result<Vec<u8>, FormatError> {
|
||||||
|
let limit = if expected_bytes != 0 {
|
||||||
|
expected_bytes
|
||||||
|
} else {
|
||||||
|
MAX_DECOMPRESS_SIZE
|
||||||
|
};
|
||||||
|
|
||||||
// Try system zlib first on macOS (Apple's ARM64-optimized libz is ~1.4x
|
// Try system zlib first on macOS (Apple's ARM64-optimized libz is ~1.4x
|
||||||
// faster at decompression than zlib-ng on Apple Silicon).
|
// faster at decompression than zlib-ng on Apple Silicon).
|
||||||
#[cfg(all(target_os = "macos", feature = "system-zlib-decompress"))]
|
#[cfg(all(target_os = "macos", feature = "system-zlib-decompress"))]
|
||||||
{
|
{
|
||||||
if let Ok(result) = sysz::decompress(data) {
|
if let Ok(result) = sysz::decompress(data) {
|
||||||
|
if result.len() > limit {
|
||||||
|
return Err(FormatError::DecompressionError(
|
||||||
|
"deflate: output exceeds expected chunk size".into(),
|
||||||
|
));
|
||||||
|
}
|
||||||
return Ok(result);
|
return Ok(result);
|
||||||
}
|
}
|
||||||
// Fall through to flate2 on error
|
// Fall through to flate2 on error
|
||||||
}
|
}
|
||||||
|
|
||||||
use std::io::Read;
|
use std::io::Read;
|
||||||
let mut decoder = flate2::read::ZlibDecoder::new(data);
|
let decoder = flate2::read::ZlibDecoder::new(data);
|
||||||
let mut result = Vec::new();
|
let mut result = Vec::with_capacity(limit.min(1 << 20));
|
||||||
|
// Read one byte past the limit so an over-size stream is distinguishable
|
||||||
|
// from one that legitimately ends exactly at the limit.
|
||||||
decoder
|
decoder
|
||||||
|
.take(limit as u64 + 1)
|
||||||
.read_to_end(&mut result)
|
.read_to_end(&mut result)
|
||||||
.map_err(|e| FormatError::DecompressionError(e.to_string()))?;
|
.map_err(|e| FormatError::DecompressionError(e.to_string()))?;
|
||||||
|
if result.len() > limit {
|
||||||
|
return Err(FormatError::DecompressionError(
|
||||||
|
"deflate: output exceeds size limit".into(),
|
||||||
|
));
|
||||||
|
}
|
||||||
Ok(result)
|
Ok(result)
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -659,7 +715,7 @@ mod sysz {
|
|||||||
}
|
}
|
||||||
|
|
||||||
#[cfg(not(feature = "deflate"))]
|
#[cfg(not(feature = "deflate"))]
|
||||||
fn deflate_decompress(_data: &[u8]) -> Result<Vec<u8>, FormatError> {
|
fn deflate_decompress(_data: &[u8], _expected_bytes: usize) -> Result<Vec<u8>, FormatError> {
|
||||||
Err(FormatError::UnsupportedFilter(FILTER_DEFLATE))
|
Err(FormatError::UnsupportedFilter(FILTER_DEFLATE))
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -682,20 +738,35 @@ fn deflate_compress(_data: &[u8], _level: u32) -> Result<Vec<u8>, FormatError> {
|
|||||||
}
|
}
|
||||||
|
|
||||||
/// Decompress LZ4 data. Format: 4 bytes LE original size + LZ4 block data.
|
/// Decompress LZ4 data. Format: 4 bytes LE original size + LZ4 block data.
|
||||||
|
///
|
||||||
|
/// The 4-byte "original size" header is part of the attacker-controlled
|
||||||
|
/// compressed payload itself, so it is bounded against `expected_bytes` (the
|
||||||
|
/// pipeline's declared chunk size) before being used to size the output
|
||||||
|
/// allocation — otherwise a crafted 4-byte value can request up to ~4 GiB.
|
||||||
#[cfg(feature = "lz4")]
|
#[cfg(feature = "lz4")]
|
||||||
fn lz4_decompress(data: &[u8]) -> Result<Vec<u8>, FormatError> {
|
fn lz4_decompress(data: &[u8], expected_bytes: usize) -> Result<Vec<u8>, FormatError> {
|
||||||
if data.len() < 4 {
|
if data.len() < 4 {
|
||||||
return Err(FormatError::DecompressionError(
|
return Err(FormatError::DecompressionError(
|
||||||
"lz4: data too short".into(),
|
"lz4: data too short".into(),
|
||||||
));
|
));
|
||||||
}
|
}
|
||||||
let orig_size = u32::from_le_bytes([data[0], data[1], data[2], data[3]]) as usize;
|
let orig_size = u32::from_le_bytes([data[0], data[1], data[2], data[3]]) as usize;
|
||||||
|
if expected_bytes != 0 && orig_size > expected_bytes {
|
||||||
|
return Err(FormatError::DecompressionError(
|
||||||
|
"lz4: declared size exceeds chunk size".into(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
if orig_size > MAX_DECOMPRESS_SIZE {
|
||||||
|
return Err(FormatError::DecompressionError(
|
||||||
|
"lz4: declared size exceeds limit".into(),
|
||||||
|
));
|
||||||
|
}
|
||||||
lz4_flex::block::decompress(&data[4..], orig_size)
|
lz4_flex::block::decompress(&data[4..], orig_size)
|
||||||
.map_err(|e| FormatError::DecompressionError(format!("lz4: {e}")))
|
.map_err(|e| FormatError::DecompressionError(format!("lz4: {e}")))
|
||||||
}
|
}
|
||||||
|
|
||||||
#[cfg(not(feature = "lz4"))]
|
#[cfg(not(feature = "lz4"))]
|
||||||
fn lz4_decompress(_data: &[u8]) -> Result<Vec<u8>, FormatError> {
|
fn lz4_decompress(_data: &[u8], _expected_bytes: usize) -> Result<Vec<u8>, FormatError> {
|
||||||
Err(FormatError::UnsupportedFilter(FILTER_LZ4))
|
Err(FormatError::UnsupportedFilter(FILTER_LZ4))
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -715,13 +786,35 @@ fn lz4_compress(_data: &[u8]) -> Result<Vec<u8>, FormatError> {
|
|||||||
}
|
}
|
||||||
|
|
||||||
/// Decompress zstd data.
|
/// Decompress zstd data.
|
||||||
|
///
|
||||||
|
/// `expected_bytes` bounds the output (or [`MAX_DECOMPRESS_SIZE`] when
|
||||||
|
/// unavailable) to guard against a zstd decompression bomb, since zstd's
|
||||||
|
/// compression ratio can exceed 1000:1.
|
||||||
#[cfg(feature = "zstd")]
|
#[cfg(feature = "zstd")]
|
||||||
fn zstd_decompress(data: &[u8]) -> Result<Vec<u8>, FormatError> {
|
fn zstd_decompress(data: &[u8], expected_bytes: usize) -> Result<Vec<u8>, FormatError> {
|
||||||
zstd::decode_all(data).map_err(|e| FormatError::DecompressionError(format!("zstd: {e}")))
|
use std::io::Read;
|
||||||
|
let limit = if expected_bytes != 0 {
|
||||||
|
expected_bytes
|
||||||
|
} else {
|
||||||
|
MAX_DECOMPRESS_SIZE
|
||||||
|
};
|
||||||
|
let decoder = zstd::stream::Decoder::new(data)
|
||||||
|
.map_err(|e| FormatError::DecompressionError(format!("zstd: {e}")))?;
|
||||||
|
let mut out = Vec::with_capacity(limit.min(1 << 20));
|
||||||
|
decoder
|
||||||
|
.take(limit as u64 + 1)
|
||||||
|
.read_to_end(&mut out)
|
||||||
|
.map_err(|e| FormatError::DecompressionError(format!("zstd: {e}")))?;
|
||||||
|
if out.len() > limit {
|
||||||
|
return Err(FormatError::DecompressionError(
|
||||||
|
"zstd: output exceeds chunk size".into(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
Ok(out)
|
||||||
}
|
}
|
||||||
|
|
||||||
#[cfg(not(feature = "zstd"))]
|
#[cfg(not(feature = "zstd"))]
|
||||||
fn zstd_decompress(_data: &[u8]) -> Result<Vec<u8>, FormatError> {
|
fn zstd_decompress(_data: &[u8], _expected_bytes: usize) -> Result<Vec<u8>, FormatError> {
|
||||||
Err(FormatError::UnsupportedFilter(FILTER_ZSTD))
|
Err(FormatError::UnsupportedFilter(FILTER_ZSTD))
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -982,30 +1075,74 @@ fn pcodec_compress(_data: &[u8], _element_size: usize) -> Result<Vec<u8>, Format
|
|||||||
Err(FormatError::UnsupportedFilter(FILTER_PCODEC))
|
Err(FormatError::UnsupportedFilter(FILTER_PCODEC))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// `expected_bytes` bounds the number of elements decoded: the output buffer
|
||||||
|
/// is pre-sized to exactly `expected_bytes / element_size` elements and
|
||||||
|
/// `simple_decompress_into` never writes past it, so a corrupted/hostile pco
|
||||||
|
/// stream cannot force over-allocation the way an unbounded `simple_decompress`
|
||||||
|
/// (which allocates however many elements the stream claims) could.
|
||||||
#[cfg(feature = "pcodec")]
|
#[cfg(feature = "pcodec")]
|
||||||
fn pcodec_decompress(data: &[u8], element_size: usize) -> Result<Vec<u8>, FormatError> {
|
fn pcodec_decompress(
|
||||||
use pco::standalone::simple_decompress;
|
data: &[u8],
|
||||||
|
element_size: usize,
|
||||||
|
expected_bytes: usize,
|
||||||
|
) -> Result<Vec<u8>, FormatError> {
|
||||||
|
use pco::standalone::simple_decompress_into;
|
||||||
|
let limit_bytes = if expected_bytes != 0 {
|
||||||
|
expected_bytes
|
||||||
|
} else {
|
||||||
|
MAX_DECOMPRESS_SIZE
|
||||||
|
};
|
||||||
|
let n = if element_size != 0 {
|
||||||
|
limit_bytes / element_size
|
||||||
|
} else {
|
||||||
|
0
|
||||||
|
};
|
||||||
match element_size {
|
match element_size {
|
||||||
4 => {
|
4 => {
|
||||||
let nums = simple_decompress::<f32>(data)
|
let mut buf = vec![0f32; n];
|
||||||
|
let progress = simple_decompress_into(data, &mut buf)
|
||||||
.map_err(|e| FormatError::DecompressionError(format!("pco: {e}")))?;
|
.map_err(|e| FormatError::DecompressionError(format!("pco: {e}")))?;
|
||||||
Ok(nums.iter().flat_map(|x| x.to_le_bytes()).collect())
|
if !progress.finished {
|
||||||
|
return Err(FormatError::DecompressionError(
|
||||||
|
"pco: stream contains more data than expected chunk size allows".into(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
buf.truncate(progress.n_processed);
|
||||||
|
Ok(buf.iter().flat_map(|x| x.to_le_bytes()).collect())
|
||||||
}
|
}
|
||||||
8 => {
|
8 => {
|
||||||
let nums = simple_decompress::<f64>(data)
|
let mut buf = vec![0f64; n];
|
||||||
|
let progress = simple_decompress_into(data, &mut buf)
|
||||||
.map_err(|e| FormatError::DecompressionError(format!("pco: {e}")))?;
|
.map_err(|e| FormatError::DecompressionError(format!("pco: {e}")))?;
|
||||||
Ok(nums.iter().flat_map(|x| x.to_le_bytes()).collect())
|
if !progress.finished {
|
||||||
|
return Err(FormatError::DecompressionError(
|
||||||
|
"pco: stream contains more data than expected chunk size allows".into(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
buf.truncate(progress.n_processed);
|
||||||
|
Ok(buf.iter().flat_map(|x| x.to_le_bytes()).collect())
|
||||||
}
|
}
|
||||||
_ => {
|
_ => {
|
||||||
let nums = simple_decompress::<u32>(data)
|
let mut buf = vec![0u32; n];
|
||||||
|
let progress = simple_decompress_into(data, &mut buf)
|
||||||
.map_err(|e| FormatError::DecompressionError(format!("pco: {e}")))?;
|
.map_err(|e| FormatError::DecompressionError(format!("pco: {e}")))?;
|
||||||
Ok(nums.iter().flat_map(|x| x.to_le_bytes()).collect())
|
if !progress.finished {
|
||||||
|
return Err(FormatError::DecompressionError(
|
||||||
|
"pco: stream contains more data than expected chunk size allows".into(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
buf.truncate(progress.n_processed);
|
||||||
|
Ok(buf.iter().flat_map(|x| x.to_le_bytes()).collect())
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
#[cfg(not(feature = "pcodec"))]
|
#[cfg(not(feature = "pcodec"))]
|
||||||
fn pcodec_decompress(_data: &[u8], _element_size: usize) -> Result<Vec<u8>, FormatError> {
|
fn pcodec_decompress(
|
||||||
|
_data: &[u8],
|
||||||
|
_element_size: usize,
|
||||||
|
_expected_bytes: usize,
|
||||||
|
) -> Result<Vec<u8>, FormatError> {
|
||||||
Err(FormatError::UnsupportedFilter(FILTER_PCODEC))
|
Err(FormatError::UnsupportedFilter(FILTER_PCODEC))
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -1021,7 +1158,7 @@ mod tests {
|
|||||||
fn deflate_compress_decompress_roundtrip() {
|
fn deflate_compress_decompress_roundtrip() {
|
||||||
let data: Vec<u8> = (0..256).map(|i| (i % 256) as u8).collect();
|
let data: Vec<u8> = (0..256).map(|i| (i % 256) as u8).collect();
|
||||||
let compressed = deflate_compress(&data, 6).unwrap();
|
let compressed = deflate_compress(&data, 6).unwrap();
|
||||||
let decompressed = deflate_decompress(&compressed).unwrap();
|
let decompressed = deflate_decompress(&compressed, data.len()).unwrap();
|
||||||
assert_eq!(decompressed, data);
|
assert_eq!(decompressed, data);
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -1034,7 +1171,7 @@ mod tests {
|
|||||||
let compressed: Vec<u8> = vec![
|
let compressed: Vec<u8> = vec![
|
||||||
120, 156, 99, 96, 100, 98, 102, 97, 101, 99, 231, 224, 4, 0, 0, 175, 0, 46,
|
120, 156, 99, 96, 100, 98, 102, 97, 101, 99, 231, 224, 4, 0, 0, 175, 0, 46,
|
||||||
];
|
];
|
||||||
let decompressed = deflate_decompress(&compressed).unwrap();
|
let decompressed = deflate_decompress(&compressed, 10).unwrap();
|
||||||
assert_eq!(decompressed, vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9]);
|
assert_eq!(decompressed, vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9]);
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -1045,7 +1182,7 @@ mod tests {
|
|||||||
let data = vec![0u8, 1, 2, 3, 4, 5, 6, 7, 8, 9];
|
let data = vec![0u8, 1, 2, 3, 4, 5, 6, 7, 8, 9];
|
||||||
let compressed = deflate_compress(&data, 6).unwrap();
|
let compressed = deflate_compress(&data, 6).unwrap();
|
||||||
assert!(!compressed.is_empty());
|
assert!(!compressed.is_empty());
|
||||||
let decompressed = deflate_decompress(&compressed).unwrap();
|
let decompressed = deflate_decompress(&compressed, data.len()).unwrap();
|
||||||
assert_eq!(decompressed, data);
|
assert_eq!(decompressed, data);
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -1246,7 +1383,7 @@ mod tests {
|
|||||||
fn lz4_compress_decompress_roundtrip() {
|
fn lz4_compress_decompress_roundtrip() {
|
||||||
let data: Vec<u8> = (0..256).map(|i| (i % 256) as u8).collect();
|
let data: Vec<u8> = (0..256).map(|i| (i % 256) as u8).collect();
|
||||||
let compressed = lz4_compress(&data).unwrap();
|
let compressed = lz4_compress(&data).unwrap();
|
||||||
let decompressed = lz4_decompress(&compressed).unwrap();
|
let decompressed = lz4_decompress(&compressed, data.len()).unwrap();
|
||||||
assert_eq!(decompressed, data);
|
assert_eq!(decompressed, data);
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -1301,7 +1438,7 @@ mod tests {
|
|||||||
fn zstd_compress_decompress_roundtrip() {
|
fn zstd_compress_decompress_roundtrip() {
|
||||||
let data: Vec<u8> = (0..256).map(|i| (i % 256) as u8).collect();
|
let data: Vec<u8> = (0..256).map(|i| (i % 256) as u8).collect();
|
||||||
let compressed = zstd_compress(&data, 3).unwrap();
|
let compressed = zstd_compress(&data, 3).unwrap();
|
||||||
let decompressed = zstd_decompress(&compressed).unwrap();
|
let decompressed = zstd_decompress(&compressed, data.len()).unwrap();
|
||||||
assert_eq!(decompressed, data);
|
assert_eq!(decompressed, data);
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -1372,7 +1509,10 @@ mod tests {
|
|||||||
0x02, 0x00, 0x00, 0x00, 0x08, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
|
0x02, 0x00, 0x00, 0x00, 0x08, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
|
||||||
0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0xc6, 0x00,
|
0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0xc6, 0x00,
|
||||||
];
|
];
|
||||||
assert_eq!(scaleoffset_decompress(&raw, &cd, 0).unwrap(), i32_le(&[0, 1, 2, 3]));
|
assert_eq!(
|
||||||
|
scaleoffset_decompress(&raw, &cd, 0).unwrap(),
|
||||||
|
i32_le(&[0, 1, 2, 3])
|
||||||
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
@@ -1450,9 +1590,9 @@ mod tests {
|
|||||||
let cd = [1u32, 1, 4, 0, 8, 0, 0, 0];
|
let cd = [1u32, 1, 4, 0, 8, 0, 0, 0];
|
||||||
let raw: &[u8] = &[
|
let raw: &[u8] = &[
|
||||||
2, 0, 0, 0, // minbits=2
|
2, 0, 0, 0, // minbits=2
|
||||||
8, // minval_width=8
|
8, // minval_width=8
|
||||||
0, 0, 0, 0, 0, 0, 0, 0, // minval=0.0f64
|
0, 0, 0, 0, 0, 0, 0, 0, // minval=0.0f64
|
||||||
0, 0, 0, 0, 0, 0, 0, 0, // 8 reserved bytes
|
0, 0, 0, 0, 0, 0, 0, 0, // 8 reserved bytes
|
||||||
0x1B, // packed codes: 00 01 10 11 MSB-first
|
0x1B, // packed codes: 00 01 10 11 MSB-first
|
||||||
];
|
];
|
||||||
let got = as_f64(&scaleoffset_decompress(raw, &cd, 0).unwrap());
|
let got = as_f64(&scaleoffset_decompress(raw, &cd, 0).unwrap());
|
||||||
@@ -1466,9 +1606,9 @@ mod tests {
|
|||||||
let cd = [1u32, 0xFFFF_FFFF, 4, 0, 8, 0, 0, 0];
|
let cd = [1u32, 0xFFFF_FFFF, 4, 0, 8, 0, 0, 0];
|
||||||
let raw: &[u8] = &[
|
let raw: &[u8] = &[
|
||||||
2, 0, 0, 0, // minbits=2
|
2, 0, 0, 0, // minbits=2
|
||||||
8, // minval_width=8
|
8, // minval_width=8
|
||||||
0, 0, 0, 0, 0, 0, 0, 0, // minval=0.0f64
|
0, 0, 0, 0, 0, 0, 0, 0, // minval=0.0f64
|
||||||
0, 0, 0, 0, 0, 0, 0, 0, // 8 reserved bytes
|
0, 0, 0, 0, 0, 0, 0, 0, // 8 reserved bytes
|
||||||
0x1B, // packed codes: 00 01 10 11 MSB-first
|
0x1B, // packed codes: 00 01 10 11 MSB-first
|
||||||
];
|
];
|
||||||
let got = as_f64(&scaleoffset_decompress(raw, &cd, 0).unwrap());
|
let got = as_f64(&scaleoffset_decompress(raw, &cd, 0).unwrap());
|
||||||
@@ -1531,8 +1671,12 @@ mod tests {
|
|||||||
fn nbit_compound_with_array_member() {
|
fn nbit_compound_with_array_member() {
|
||||||
// Compound { a: array(2,) of i32 prec 16 @0; b: u32@8 prec 8 }, 2 elements.
|
// Compound { a: array(2,) of i32 prec 16 @0; b: u32@8 prec 8 }, 2 elements.
|
||||||
// data = [([-1,100],200), ([1000,-32768],7)].
|
// data = [([-1,100],200), ([1000,-32768],7)].
|
||||||
let cd = [20u32, 0, 2, 3, 12, 2, 0, 2, 8, 1, 4, 0, 16, 0, 8, 1, 4, 0, 8, 0];
|
let cd = [
|
||||||
let raw = [0xff, 0xff, 0x00, 0x64, 0xc8, 0x03, 0xe8, 0x80, 0x00, 0x07, 0x00];
|
20u32, 0, 2, 3, 12, 2, 0, 2, 8, 1, 4, 0, 16, 0, 8, 1, 4, 0, 8, 0,
|
||||||
|
];
|
||||||
|
let raw = [
|
||||||
|
0xff, 0xff, 0x00, 0x64, 0xc8, 0x03, 0xe8, 0x80, 0x00, 0x07, 0x00,
|
||||||
|
];
|
||||||
#[rustfmt::skip]
|
#[rustfmt::skip]
|
||||||
let expected: Vec<u8> = vec![
|
let expected: Vec<u8> = vec![
|
||||||
0xff,0xff,0x00,0x00, 0x64,0x00,0x00,0x00, 0xc8,0x00,0x00,0x00, // ([-1,100], 200)
|
0xff,0xff,0x00,0x00, 0x64,0x00,0x00,0x00, 0xc8,0x00,0x00,0x00, // ([-1,100], 200)
|
||||||
@@ -1623,4 +1767,79 @@ mod tests {
|
|||||||
// Missing client data entirely.
|
// Missing client data entirely.
|
||||||
assert!(scaleoffset_decompress(&[0u8; 32], &[2, 0], 4).is_err());
|
assert!(scaleoffset_decompress(&[0u8; 32], &[2, 0], 4).is_err());
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// ----- Decompression-bomb hardening: hostile compressed data must not -----
|
||||||
|
// ----- force unbounded allocation. -----
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
#[cfg(feature = "lz4")]
|
||||||
|
fn lz4_decompress_rejects_oversized_orig_size() {
|
||||||
|
// 4-byte LE header claiming ~4 GiB, followed by a few garbage bytes.
|
||||||
|
let mut data = u32::MAX.to_le_bytes().to_vec();
|
||||||
|
data.extend_from_slice(&[0u8; 8]);
|
||||||
|
assert!(lz4_decompress(&data, 64).is_err());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
#[cfg(feature = "lz4")]
|
||||||
|
fn lz4_decompress_rejects_size_exceeding_chunk_size() {
|
||||||
|
// orig_size (1000) is well under MAX_DECOMPRESS_SIZE but exceeds the
|
||||||
|
// pipeline's declared chunk size (64) — must be rejected by the
|
||||||
|
// chunk-size check specifically, not just the absolute cap.
|
||||||
|
let mut data = 1000u32.to_le_bytes().to_vec();
|
||||||
|
data.extend_from_slice(&[0u8; 8]);
|
||||||
|
assert!(lz4_decompress(&data, 64).is_err());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
#[cfg(feature = "deflate")]
|
||||||
|
fn deflate_decompress_rejects_output_exceeding_chunk_size() {
|
||||||
|
// A highly-compressible deflate bomb (1 MiB of zeros compresses to a
|
||||||
|
// tiny payload); declared chunk size is far smaller than the real
|
||||||
|
// decompressed size, so this must be rejected rather than allocating
|
||||||
|
// the full 1 MiB.
|
||||||
|
let data = vec![0u8; 1024 * 1024];
|
||||||
|
let compressed = deflate_compress(&data, 6).unwrap();
|
||||||
|
assert!(deflate_decompress(&compressed, 64).is_err());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
#[cfg(feature = "zstd")]
|
||||||
|
fn zstd_decompress_rejects_output_exceeding_chunk_size() {
|
||||||
|
let data = vec![0u8; 1024 * 1024];
|
||||||
|
let compressed = zstd_compress(&data, 3).unwrap();
|
||||||
|
assert!(zstd_decompress(&compressed, 64).is_err());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
#[cfg(feature = "pcodec")]
|
||||||
|
fn pcodec_decompress_rejects_element_count_exceeding_chunk_size() {
|
||||||
|
let data: Vec<f32> = (0..1000).map(|i| i as f32).collect();
|
||||||
|
let raw: Vec<u8> = data.iter().flat_map(|x| x.to_le_bytes()).collect();
|
||||||
|
let compressed = pcodec_compress(&raw, 4).unwrap();
|
||||||
|
// Declared chunk size only fits 4 f32 elements, far fewer than the
|
||||||
|
// 1000 the stream actually contains.
|
||||||
|
assert!(pcodec_decompress(&compressed, 4, 16).is_err());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
#[cfg(feature = "lz4")]
|
||||||
|
fn decompress_chunk_rejects_hostile_lz4_size_via_public_entrypoint() {
|
||||||
|
// The actually-exploited path: a FilterPipeline claiming a small
|
||||||
|
// chunk_size, but whose LZ4-compressed data header claims a huge
|
||||||
|
// decompressed size.
|
||||||
|
use crate::filter_pipeline::{FilterDescription, FilterPipeline};
|
||||||
|
let mut data = u32::MAX.to_le_bytes().to_vec();
|
||||||
|
data.extend_from_slice(&[0u8; 8]);
|
||||||
|
let pipeline = FilterPipeline {
|
||||||
|
version: 2,
|
||||||
|
filters: vec![FilterDescription {
|
||||||
|
filter_id: FILTER_LZ4,
|
||||||
|
name: None,
|
||||||
|
flags: 0,
|
||||||
|
client_data: vec![],
|
||||||
|
}],
|
||||||
|
};
|
||||||
|
assert!(decompress_chunk(&data, &pipeline, 16, 1).is_err());
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -2,6 +2,9 @@
|
|||||||
//!
|
//!
|
||||||
//! Gated by the `szip` feature which links against the system libaec library.
|
//! Gated by the `szip` feature which links against the system libaec library.
|
||||||
|
|
||||||
|
#[cfg(not(feature = "std"))]
|
||||||
|
use alloc::vec::Vec;
|
||||||
|
|
||||||
use crate::error::FormatError;
|
use crate::error::FormatError;
|
||||||
|
|
||||||
/// Decompress SZIP-compressed data using libaec.
|
/// Decompress SZIP-compressed data using libaec.
|
||||||
@@ -49,9 +52,7 @@ fn szip_decode_impl(data: &[u8], cd: &[u32], chunk_size: usize) -> Result<Vec<u8
|
|||||||
));
|
));
|
||||||
}
|
}
|
||||||
if data.is_empty() {
|
if data.is_empty() {
|
||||||
return Err(FormatError::ChunkedReadError(
|
return Err(FormatError::ChunkedReadError("szip: empty input".into()));
|
||||||
"szip: empty input".into(),
|
|
||||||
));
|
|
||||||
}
|
}
|
||||||
|
|
||||||
// Map HDF5 option mask to libaec flags.
|
// Map HDF5 option mask to libaec flags.
|
||||||
@@ -112,7 +113,7 @@ mod tests {
|
|||||||
#[cfg(feature = "szip")]
|
#[cfg(feature = "szip")]
|
||||||
#[test]
|
#[test]
|
||||||
fn roundtrip_u8_msb_no_nn() {
|
fn roundtrip_u8_msb_no_nn() {
|
||||||
use libaec_sys::{AecStream, AEC_DATA_MSB};
|
use libaec_sys::{AEC_DATA_MSB, AecStream};
|
||||||
|
|
||||||
let original: Vec<u8> = (0..1024u32).map(|i| (i % 256) as u8).collect();
|
let original: Vec<u8> = (0..1024u32).map(|i| (i % 256) as u8).collect();
|
||||||
|
|
||||||
@@ -144,7 +145,7 @@ mod tests {
|
|||||||
#[cfg(feature = "szip")]
|
#[cfg(feature = "szip")]
|
||||||
#[test]
|
#[test]
|
||||||
fn roundtrip_u8_msb_with_nn() {
|
fn roundtrip_u8_msb_with_nn() {
|
||||||
use libaec_sys::{AecStream, AEC_DATA_MSB, AEC_DATA_PREPROCESS};
|
use libaec_sys::{AEC_DATA_MSB, AEC_DATA_PREPROCESS, AecStream};
|
||||||
|
|
||||||
let original: Vec<u8> = (0..1024u32).map(|i| (i % 256) as u8).collect();
|
let original: Vec<u8> = (0..1024u32).map(|i| (i % 256) as u8).collect();
|
||||||
|
|
||||||
@@ -167,6 +168,9 @@ mod tests {
|
|||||||
let cd = [0x20u32, 8, 8, 1024];
|
let cd = [0x20u32, 8, 8, 1024];
|
||||||
let decoded = szip_decompress(&encoded, &cd, original.len())
|
let decoded = szip_decompress(&encoded, &cd, original.len())
|
||||||
.expect("szip_decompress with NN must succeed");
|
.expect("szip_decompress with NN must succeed");
|
||||||
assert_eq!(decoded, original, "NN round-trip must reproduce original data");
|
assert_eq!(
|
||||||
|
decoded, original,
|
||||||
|
"NN round-trip must reproduce original data"
|
||||||
|
);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -186,25 +186,26 @@ pub fn read_fixed_array_chunks(
|
|||||||
chunk_dimensions.iter().map(|&d| d as u64).product::<u64>() * element_size as u64;
|
chunk_dimensions.iter().map(|&d| d as u64).product::<u64>() * element_size as u64;
|
||||||
|
|
||||||
let mut chunks = Vec::new();
|
let mut chunks = Vec::new();
|
||||||
let push_element = |i: usize, abs: usize, chunks: &mut Vec<ChunkInfo>| -> Result<(), FormatError> {
|
let push_element =
|
||||||
if let Some((address, chunk_size, filter_mask)) = parse_fa_element(
|
|i: usize, abs: usize, chunks: &mut Vec<ChunkInfo>| -> Result<(), FormatError> {
|
||||||
file_data,
|
if let Some((address, chunk_size, filter_mask)) = parse_fa_element(
|
||||||
abs,
|
file_data,
|
||||||
header.client_id,
|
abs,
|
||||||
offset_size,
|
header.client_id,
|
||||||
header.element_size,
|
offset_size,
|
||||||
chunk_byte_size,
|
header.element_size,
|
||||||
)? {
|
chunk_byte_size,
|
||||||
let offsets = index_to_chunk_offsets(i, &num_chunks_per_dim, chunk_dimensions);
|
)? {
|
||||||
chunks.push(ChunkInfo {
|
let offsets = index_to_chunk_offsets(i, &num_chunks_per_dim, chunk_dimensions);
|
||||||
chunk_size,
|
chunks.push(ChunkInfo {
|
||||||
filter_mask,
|
chunk_size,
|
||||||
offsets,
|
filter_mask,
|
||||||
address,
|
offsets,
|
||||||
});
|
address,
|
||||||
}
|
});
|
||||||
Ok(())
|
}
|
||||||
};
|
Ok(())
|
||||||
|
};
|
||||||
|
|
||||||
// A data block is paged when it holds more elements than fit in one page.
|
// A data block is paged when it holds more elements than fit in one page.
|
||||||
// `max_nelmts_bits` is an untrusted u8; a shift >= the pointer width would
|
// `max_nelmts_bits` is an untrusted u8; a shift >= the pointer width would
|
||||||
@@ -232,9 +233,8 @@ pub fn read_fixed_array_chunks(
|
|||||||
// only the final page holds fewer elements. Uninitialized pages (bit clear)
|
// only the final page holds fewer elements. Uninitialized pages (bit clear)
|
||||||
// still occupy their slot on disk but are zero-filled, so the bitmap — not a
|
// still occupy their slot on disk but are zero-filled, so the bitmap — not a
|
||||||
// 0xFF sentinel — is what marks a whole page as unallocated.
|
// 0xFF sentinel — is what marks a whole page as unallocated.
|
||||||
let stride_overflow = || {
|
let stride_overflow =
|
||||||
FormatError::ChunkedReadError("Fixed Array page offset overflow".into())
|
|| FormatError::ChunkedReadError("Fixed Array page offset overflow".into());
|
||||||
};
|
|
||||||
let npages = num_elements.div_ceil(page_nelmts);
|
let npages = num_elements.div_ceil(page_nelmts);
|
||||||
let bitmap_size = npages.div_ceil(8);
|
let bitmap_size = npages.div_ceil(8);
|
||||||
let bitmap_start = elements_start;
|
let bitmap_start = elements_start;
|
||||||
@@ -720,10 +720,7 @@ mod tests {
|
|||||||
// Page 1 (elements 4,5,6,7) is uninitialized => skipped. The remaining
|
// Page 1 (elements 4,5,6,7) is uninitialized => skipped. The remaining
|
||||||
// 7 chunks (0..4 and 8..11) come back with their original linear index.
|
// 7 chunks (0..4 and 8..11) come back with their original linear index.
|
||||||
assert_eq!(chunks.len(), 7);
|
assert_eq!(chunks.len(), 7);
|
||||||
let mut got: Vec<(u64, u64)> = chunks
|
let mut got: Vec<(u64, u64)> = chunks.iter().map(|c| (c.offsets[0], c.address)).collect();
|
||||||
.iter()
|
|
||||||
.map(|c| (c.offsets[0], c.address))
|
|
||||||
.collect();
|
|
||||||
got.sort();
|
got.sort();
|
||||||
let expect: Vec<(u64, u64)> = [0usize, 1, 2, 3, 8, 9, 10]
|
let expect: Vec<(u64, u64)> = [0usize, 1, 2, 3, 8, 9, 10]
|
||||||
.iter()
|
.iter()
|
||||||
|
|||||||
@@ -6,7 +6,7 @@ use alloc::vec::Vec;
|
|||||||
use crate::error::FormatError;
|
use crate::error::FormatError;
|
||||||
|
|
||||||
/// Magic signature for global heap collections.
|
/// Magic signature for global heap collections.
|
||||||
const GCOL_SIGNATURE: [u8; 4] = [b'G', b'C', b'O', b'L'];
|
const GCOL_SIGNATURE: [u8; 4] = *b"GCOL";
|
||||||
|
|
||||||
/// A parsed global heap collection.
|
/// A parsed global heap collection.
|
||||||
#[derive(Debug, Clone)]
|
#[derive(Debug, Clone)]
|
||||||
|
|||||||
@@ -9,10 +9,10 @@ use crate::error::FormatError;
|
|||||||
use crate::message_type::MessageType;
|
use crate::message_type::MessageType;
|
||||||
|
|
||||||
/// OHDR signature for v2 object headers.
|
/// OHDR signature for v2 object headers.
|
||||||
const OHDR_SIGNATURE: [u8; 4] = [b'O', b'H', b'D', b'R'];
|
const OHDR_SIGNATURE: [u8; 4] = *b"OHDR";
|
||||||
|
|
||||||
/// OCHK signature for v2 continuation chunks.
|
/// OCHK signature for v2 continuation chunks.
|
||||||
const OCHK_SIGNATURE: [u8; 4] = [b'O', b'C', b'H', b'K'];
|
const OCHK_SIGNATURE: [u8; 4] = *b"OCHK";
|
||||||
|
|
||||||
/// A single parsed header message.
|
/// A single parsed header message.
|
||||||
#[derive(Debug, Clone)]
|
#[derive(Debug, Clone)]
|
||||||
@@ -555,13 +555,12 @@ mod tests {
|
|||||||
buf.push(2); // version
|
buf.push(2); // version
|
||||||
buf.push(flags);
|
buf.push(flags);
|
||||||
|
|
||||||
if has_timestamps
|
if has_timestamps && let Some((at, mt, ct, bt)) = timestamps {
|
||||||
&& let Some((at, mt, ct, bt)) = timestamps {
|
buf.extend_from_slice(&at.to_le_bytes());
|
||||||
buf.extend_from_slice(&at.to_le_bytes());
|
buf.extend_from_slice(&mt.to_le_bytes());
|
||||||
buf.extend_from_slice(&mt.to_le_bytes());
|
buf.extend_from_slice(&ct.to_le_bytes());
|
||||||
buf.extend_from_slice(&ct.to_le_bytes());
|
buf.extend_from_slice(&bt.to_le_bytes());
|
||||||
buf.extend_from_slice(&bt.to_le_bytes());
|
}
|
||||||
}
|
|
||||||
|
|
||||||
if flags & 0x10 != 0 {
|
if flags & 0x10 != 0 {
|
||||||
buf.extend_from_slice(&8u16.to_le_bytes()); // max_compact
|
buf.extend_from_slice(&8u16.to_le_bytes()); // max_compact
|
||||||
|
|||||||
@@ -4,7 +4,7 @@
|
|||||||
//! events. The [`DefaultProfiler`] implementation uses atomic counters for
|
//! events. The [`DefaultProfiler`] implementation uses atomic counters for
|
||||||
//! thread-safe, low-overhead profiling.
|
//! thread-safe, low-overhead profiling.
|
||||||
|
|
||||||
use core::sync::atomic::{AtomicU64, Ordering};
|
use portable_atomic::{AtomicU64, Ordering};
|
||||||
|
|
||||||
/// Trait for profiling I/O operations.
|
/// Trait for profiling I/O operations.
|
||||||
///
|
///
|
||||||
|
|||||||
@@ -587,7 +587,7 @@ mod tests {
|
|||||||
// start=0 stride=1 count=1 block=4, version 3, enc_size 2, rank 1.
|
// start=0 stride=1 count=1 block=4, version 3, enc_size 2, rank 1.
|
||||||
let bytes = [
|
let bytes = [
|
||||||
0x02, 0, 0, 0, // type = HYPER
|
0x02, 0, 0, 0, // type = HYPER
|
||||||
0x03, 0, 0, 0, // version 3
|
0x03, 0, 0, 0, // version 3
|
||||||
0x01, // flags = regular
|
0x01, // flags = regular
|
||||||
0x02, // enc_size = 2
|
0x02, // enc_size = 2
|
||||||
0x01, 0, 0, 0, // rank = 1
|
0x01, 0, 0, 0, // rank = 1
|
||||||
|
|||||||
@@ -321,9 +321,9 @@ fn roundtrip_through_file_writer() {
|
|||||||
&& let clawhdf5_format::link_message::LinkTarget::Hard {
|
&& let clawhdf5_format::link_message::LinkTarget::Hard {
|
||||||
object_header_address,
|
object_header_address,
|
||||||
} = link.link_target
|
} = link.link_target
|
||||||
{
|
{
|
||||||
ds_addr = Some(object_header_address);
|
ds_addr = Some(object_header_address);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
let ds_addr = ds_addr.expect("compound_ds link not found");
|
let ds_addr = ds_addr.expect("compound_ds link not found");
|
||||||
|
|||||||
@@ -706,7 +706,10 @@ fn scaleoffset_float_escale_reads_as_raw() {
|
|||||||
let (raw, datatype, _) = read_chunked_dataset(file_data, "x");
|
let (raw, datatype, _) = read_chunked_dataset(file_data, "x");
|
||||||
let values = read_as_f64(&raw, &datatype).unwrap();
|
let values = read_as_f64(&raw, &datatype).unwrap();
|
||||||
let expect: Vec<f64> = (0..20).map(|i| i as f64 * 0.25).collect();
|
let expect: Vec<f64> = (0..20).map(|i| i as f64 * 0.25).collect();
|
||||||
assert_eq!(values, expect, "E-scale (raw + masked filter) must read verbatim");
|
assert_eq!(
|
||||||
|
values, expect,
|
||||||
|
"E-scale (raw + masked filter) must read verbatim"
|
||||||
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
@@ -717,26 +720,48 @@ fn v4_virtual_dataset_cycle_errors_not_overflow() {
|
|||||||
let offset = find_signature(file_data).unwrap();
|
let offset = find_signature(file_data).unwrap();
|
||||||
let sb = Superblock::parse(file_data, offset).unwrap();
|
let sb = Superblock::parse(file_data, offset).unwrap();
|
||||||
let addr = resolve_path_any(file_data, &sb, "virt").unwrap();
|
let addr = resolve_path_any(file_data, &sb, "virt").unwrap();
|
||||||
let hdr = ObjectHeader::parse(file_data, addr as usize, sb.offset_size, sb.length_size).unwrap();
|
let hdr =
|
||||||
|
ObjectHeader::parse(file_data, addr as usize, sb.offset_size, sb.length_size).unwrap();
|
||||||
let ds = Dataspace::parse(
|
let ds = Dataspace::parse(
|
||||||
&hdr.messages.iter().find(|m| m.msg_type == MessageType::Dataspace).unwrap().data,
|
&hdr.messages
|
||||||
|
.iter()
|
||||||
|
.find(|m| m.msg_type == MessageType::Dataspace)
|
||||||
|
.unwrap()
|
||||||
|
.data,
|
||||||
sb.length_size,
|
sb.length_size,
|
||||||
)
|
)
|
||||||
.unwrap();
|
.unwrap();
|
||||||
let (dt, _) = Datatype::parse(
|
let (dt, _) = Datatype::parse(
|
||||||
&hdr.messages.iter().find(|m| m.msg_type == MessageType::Datatype).unwrap().data,
|
&hdr.messages
|
||||||
|
.iter()
|
||||||
|
.find(|m| m.msg_type == MessageType::Datatype)
|
||||||
|
.unwrap()
|
||||||
|
.data,
|
||||||
)
|
)
|
||||||
.unwrap();
|
.unwrap();
|
||||||
let layout = DataLayout::parse(
|
let layout = DataLayout::parse(
|
||||||
&hdr.messages.iter().find(|m| m.msg_type == MessageType::DataLayout).unwrap().data,
|
&hdr.messages
|
||||||
|
.iter()
|
||||||
|
.find(|m| m.msg_type == MessageType::DataLayout)
|
||||||
|
.unwrap()
|
||||||
|
.data,
|
||||||
sb.offset_size,
|
sb.offset_size,
|
||||||
sb.length_size,
|
sb.length_size,
|
||||||
)
|
)
|
||||||
.unwrap();
|
.unwrap();
|
||||||
let r = read_raw_data_full(
|
let r = read_raw_data_full(
|
||||||
file_data, &layout, &ds, &dt, None, sb.offset_size, sb.length_size,
|
file_data,
|
||||||
|
&layout,
|
||||||
|
&ds,
|
||||||
|
&dt,
|
||||||
|
None,
|
||||||
|
sb.offset_size,
|
||||||
|
sb.length_size,
|
||||||
|
);
|
||||||
|
assert!(
|
||||||
|
r.is_err(),
|
||||||
|
"cyclic virtual dataset must error, not overflow"
|
||||||
);
|
);
|
||||||
assert!(r.is_err(), "cyclic virtual dataset must error, not overflow");
|
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
@@ -751,16 +776,28 @@ fn v4_virtual_dataset_external_file_read() {
|
|||||||
let addr = resolve_path_any(virt, &sb, "virt").unwrap();
|
let addr = resolve_path_any(virt, &sb, "virt").unwrap();
|
||||||
let hdr = ObjectHeader::parse(virt, addr as usize, sb.offset_size, sb.length_size).unwrap();
|
let hdr = ObjectHeader::parse(virt, addr as usize, sb.offset_size, sb.length_size).unwrap();
|
||||||
let ds = Dataspace::parse(
|
let ds = Dataspace::parse(
|
||||||
&hdr.messages.iter().find(|m| m.msg_type == MessageType::Dataspace).unwrap().data,
|
&hdr.messages
|
||||||
|
.iter()
|
||||||
|
.find(|m| m.msg_type == MessageType::Dataspace)
|
||||||
|
.unwrap()
|
||||||
|
.data,
|
||||||
sb.length_size,
|
sb.length_size,
|
||||||
)
|
)
|
||||||
.unwrap();
|
.unwrap();
|
||||||
let (dt, _) = Datatype::parse(
|
let (dt, _) = Datatype::parse(
|
||||||
&hdr.messages.iter().find(|m| m.msg_type == MessageType::Datatype).unwrap().data,
|
&hdr.messages
|
||||||
|
.iter()
|
||||||
|
.find(|m| m.msg_type == MessageType::Datatype)
|
||||||
|
.unwrap()
|
||||||
|
.data,
|
||||||
)
|
)
|
||||||
.unwrap();
|
.unwrap();
|
||||||
let layout = DataLayout::parse(
|
let layout = DataLayout::parse(
|
||||||
&hdr.messages.iter().find(|m| m.msg_type == MessageType::DataLayout).unwrap().data,
|
&hdr.messages
|
||||||
|
.iter()
|
||||||
|
.find(|m| m.msg_type == MessageType::DataLayout)
|
||||||
|
.unwrap()
|
||||||
|
.data,
|
||||||
sb.offset_size,
|
sb.offset_size,
|
||||||
sb.length_size,
|
sb.length_size,
|
||||||
)
|
)
|
||||||
@@ -790,7 +827,14 @@ fn v4_virtual_dataset_external_file_read() {
|
|||||||
|
|
||||||
// With no resolver, an external source is a clean error (not wrong data).
|
// With no resolver, an external source is a clean error (not wrong data).
|
||||||
let no_resolver = read_raw_data_full_with_resolver(
|
let no_resolver = read_raw_data_full_with_resolver(
|
||||||
virt, &layout, &ds, &dt, None, sb.offset_size, sb.length_size, None,
|
virt,
|
||||||
|
&layout,
|
||||||
|
&ds,
|
||||||
|
&dt,
|
||||||
|
None,
|
||||||
|
sb.offset_size,
|
||||||
|
sb.length_size,
|
||||||
|
None,
|
||||||
);
|
);
|
||||||
assert!(no_resolver.is_err());
|
assert!(no_resolver.is_err());
|
||||||
}
|
}
|
||||||
@@ -805,9 +849,18 @@ fn v4_paged_fixed_array_read() {
|
|||||||
let values = read_as_i32(&raw, &datatype).unwrap();
|
let values = read_as_i32(&raw, &datatype).unwrap();
|
||||||
assert_eq!(values.len(), 1025 * 16);
|
assert_eq!(values.len(), 1025 * 16);
|
||||||
for k in 0..1025usize {
|
for k in 0..1025usize {
|
||||||
assert_eq!(values[k * 16], k as i32, "chunk-start mismatch at chunk {k}");
|
assert_eq!(
|
||||||
|
values[k * 16],
|
||||||
|
k as i32,
|
||||||
|
"chunk-start mismatch at chunk {k}"
|
||||||
|
);
|
||||||
for j in 1..16 {
|
for j in 1..16 {
|
||||||
assert_eq!(values[k * 16 + j], 0, "non-start element nonzero at {}", k * 16 + j);
|
assert_eq!(
|
||||||
|
values[k * 16 + j],
|
||||||
|
0,
|
||||||
|
"non-start element nonzero at {}",
|
||||||
|
k * 16 + j
|
||||||
|
);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -214,9 +214,9 @@ print('ok')
|
|||||||
&& let clawhdf5_format::link_message::LinkTarget::Hard {
|
&& let clawhdf5_format::link_message::LinkTarget::Hard {
|
||||||
object_header_address,
|
object_header_address,
|
||||||
} = link.link_target
|
} = link.link_target
|
||||||
{
|
{
|
||||||
refs_addr = Some(object_header_address);
|
refs_addr = Some(object_header_address);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -309,7 +309,8 @@ mod tests {
|
|||||||
insert_relation(&conn, 1, 1, "self");
|
insert_relation(&conn, 1, 1, "self");
|
||||||
drop(conn);
|
drop(conn);
|
||||||
|
|
||||||
let data = sqlite_reader::read_sqlite(&db_path, false, None, &SchemaConfig::default()).unwrap();
|
let data =
|
||||||
|
sqlite_reader::read_sqlite(&db_path, false, None, &SchemaConfig::default()).unwrap();
|
||||||
let opts = hdf5_writer::WriteOptions {
|
let opts = hdf5_writer::WriteOptions {
|
||||||
agent_id: "test-agent".into(),
|
agent_id: "test-agent".into(),
|
||||||
embedder: "test-embed".into(),
|
embedder: "test-embed".into(),
|
||||||
@@ -319,7 +320,8 @@ mod tests {
|
|||||||
};
|
};
|
||||||
hdf5_writer::write_hdf5(h5_path.to_str().unwrap(), &data, &opts).unwrap();
|
hdf5_writer::write_hdf5(h5_path.to_str().unwrap(), &data, &opts).unwrap();
|
||||||
|
|
||||||
let summary = validate::validate_hdf5(h5_path.to_str().unwrap(), &data, false, false).unwrap();
|
let summary =
|
||||||
|
validate::validate_hdf5(h5_path.to_str().unwrap(), &data, false, false).unwrap();
|
||||||
assert_eq!(summary.chunks, 2);
|
assert_eq!(summary.chunks, 2);
|
||||||
assert_eq!(summary.sessions, 1);
|
assert_eq!(summary.sessions, 1);
|
||||||
assert_eq!(summary.entities, 1);
|
assert_eq!(summary.entities, 1);
|
||||||
@@ -340,7 +342,8 @@ mod tests {
|
|||||||
insert_chunk(&conn, 3, "also active", &make_embedding(4, 3.0), 0);
|
insert_chunk(&conn, 3, "also active", &make_embedding(4, 3.0), 0);
|
||||||
drop(conn);
|
drop(conn);
|
||||||
|
|
||||||
let data = sqlite_reader::read_sqlite(&db_path, true, None, &SchemaConfig::default()).unwrap();
|
let data =
|
||||||
|
sqlite_reader::read_sqlite(&db_path, true, None, &SchemaConfig::default()).unwrap();
|
||||||
assert_eq!(data.chunks.len(), 2);
|
assert_eq!(data.chunks.len(), 2);
|
||||||
|
|
||||||
let opts = hdf5_writer::WriteOptions {
|
let opts = hdf5_writer::WriteOptions {
|
||||||
@@ -352,7 +355,8 @@ mod tests {
|
|||||||
};
|
};
|
||||||
hdf5_writer::write_hdf5(h5_path.to_str().unwrap(), &data, &opts).unwrap();
|
hdf5_writer::write_hdf5(h5_path.to_str().unwrap(), &data, &opts).unwrap();
|
||||||
|
|
||||||
let summary = validate::validate_hdf5(h5_path.to_str().unwrap(), &data, false, false).unwrap();
|
let summary =
|
||||||
|
validate::validate_hdf5(h5_path.to_str().unwrap(), &data, false, false).unwrap();
|
||||||
assert_eq!(summary.chunks, 2);
|
assert_eq!(summary.chunks, 2);
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -367,7 +371,8 @@ mod tests {
|
|||||||
insert_chunk(&conn, 2, "deleted", &make_embedding(4, 2.0), 1);
|
insert_chunk(&conn, 2, "deleted", &make_embedding(4, 2.0), 1);
|
||||||
drop(conn);
|
drop(conn);
|
||||||
|
|
||||||
let data = sqlite_reader::read_sqlite(&db_path, false, None, &SchemaConfig::default()).unwrap();
|
let data =
|
||||||
|
sqlite_reader::read_sqlite(&db_path, false, None, &SchemaConfig::default()).unwrap();
|
||||||
assert_eq!(data.chunks.len(), 2);
|
assert_eq!(data.chunks.len(), 2);
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -381,7 +386,8 @@ mod tests {
|
|||||||
insert_chunk(&conn, 1, "test", &make_embedding(16, 0.5), 0);
|
insert_chunk(&conn, 1, "test", &make_embedding(16, 0.5), 0);
|
||||||
drop(conn);
|
drop(conn);
|
||||||
|
|
||||||
let data = sqlite_reader::read_sqlite(&db_path, false, None, &SchemaConfig::default()).unwrap();
|
let data =
|
||||||
|
sqlite_reader::read_sqlite(&db_path, false, None, &SchemaConfig::default()).unwrap();
|
||||||
assert_eq!(data.embedding_dim, 16);
|
assert_eq!(data.embedding_dim, 16);
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -395,7 +401,8 @@ mod tests {
|
|||||||
insert_chunk(&conn, 1, "test", &make_embedding(16, 0.5), 0);
|
insert_chunk(&conn, 1, "test", &make_embedding(16, 0.5), 0);
|
||||||
drop(conn);
|
drop(conn);
|
||||||
|
|
||||||
let data = sqlite_reader::read_sqlite(&db_path, false, Some(8), &SchemaConfig::default()).unwrap();
|
let data =
|
||||||
|
sqlite_reader::read_sqlite(&db_path, false, Some(8), &SchemaConfig::default()).unwrap();
|
||||||
assert_eq!(data.embedding_dim, 8);
|
assert_eq!(data.embedding_dim, 8);
|
||||||
// Embedding truncated to dim 8
|
// Embedding truncated to dim 8
|
||||||
assert_eq!(data.chunks[0].embedding.len(), 8);
|
assert_eq!(data.chunks[0].embedding.len(), 8);
|
||||||
@@ -413,7 +420,8 @@ mod tests {
|
|||||||
insert_chunk(&conn, 1, "test", &emb, 0);
|
insert_chunk(&conn, 1, "test", &emb, 0);
|
||||||
drop(conn);
|
drop(conn);
|
||||||
|
|
||||||
let data = sqlite_reader::read_sqlite(&db_path, false, None, &SchemaConfig::default()).unwrap();
|
let data =
|
||||||
|
sqlite_reader::read_sqlite(&db_path, false, None, &SchemaConfig::default()).unwrap();
|
||||||
let opts = hdf5_writer::WriteOptions {
|
let opts = hdf5_writer::WriteOptions {
|
||||||
agent_id: "t".into(),
|
agent_id: "t".into(),
|
||||||
embedder: "t".into(),
|
embedder: "t".into(),
|
||||||
@@ -424,7 +432,8 @@ mod tests {
|
|||||||
hdf5_writer::write_hdf5(h5_path.to_str().unwrap(), &data, &opts).unwrap();
|
hdf5_writer::write_hdf5(h5_path.to_str().unwrap(), &data, &opts).unwrap();
|
||||||
|
|
||||||
// Content-validate with the float16 tolerance enabled.
|
// Content-validate with the float16 tolerance enabled.
|
||||||
let summary = validate::validate_hdf5(h5_path.to_str().unwrap(), &data, true, true).unwrap();
|
let summary =
|
||||||
|
validate::validate_hdf5(h5_path.to_str().unwrap(), &data, true, true).unwrap();
|
||||||
assert_eq!(summary.chunks, 1);
|
assert_eq!(summary.chunks, 1);
|
||||||
|
|
||||||
// Verify float16 values are within tolerance
|
// Verify float16 values are within tolerance
|
||||||
@@ -453,7 +462,8 @@ mod tests {
|
|||||||
}
|
}
|
||||||
drop(conn);
|
drop(conn);
|
||||||
|
|
||||||
let data = sqlite_reader::read_sqlite(&db_path, false, None, &SchemaConfig::default()).unwrap();
|
let data =
|
||||||
|
sqlite_reader::read_sqlite(&db_path, false, None, &SchemaConfig::default()).unwrap();
|
||||||
|
|
||||||
let opts_compressed = hdf5_writer::WriteOptions {
|
let opts_compressed = hdf5_writer::WriteOptions {
|
||||||
agent_id: "t".into(),
|
agent_id: "t".into(),
|
||||||
@@ -493,7 +503,8 @@ mod tests {
|
|||||||
drop(conn);
|
drop(conn);
|
||||||
|
|
||||||
// Simulate dry-run: read data but don't write
|
// Simulate dry-run: read data but don't write
|
||||||
let data = sqlite_reader::read_sqlite(&db_path, false, None, &SchemaConfig::default()).unwrap();
|
let data =
|
||||||
|
sqlite_reader::read_sqlite(&db_path, false, None, &SchemaConfig::default()).unwrap();
|
||||||
assert_eq!(data.chunks.len(), 1);
|
assert_eq!(data.chunks.len(), 1);
|
||||||
assert!(!h5_path.exists());
|
assert!(!h5_path.exists());
|
||||||
}
|
}
|
||||||
@@ -505,7 +516,8 @@ mod tests {
|
|||||||
let db_path = create_test_db(&dir);
|
let db_path = create_test_db(&dir);
|
||||||
let h5_path = dir.path().join("out.h5");
|
let h5_path = dir.path().join("out.h5");
|
||||||
|
|
||||||
let data = sqlite_reader::read_sqlite(&db_path, false, None, &SchemaConfig::default()).unwrap();
|
let data =
|
||||||
|
sqlite_reader::read_sqlite(&db_path, false, None, &SchemaConfig::default()).unwrap();
|
||||||
assert_eq!(data.chunks.len(), 0);
|
assert_eq!(data.chunks.len(), 0);
|
||||||
assert_eq!(data.sessions.len(), 0);
|
assert_eq!(data.sessions.len(), 0);
|
||||||
assert_eq!(data.entities.len(), 0);
|
assert_eq!(data.entities.len(), 0);
|
||||||
@@ -520,7 +532,8 @@ mod tests {
|
|||||||
};
|
};
|
||||||
hdf5_writer::write_hdf5(h5_path.to_str().unwrap(), &data, &opts).unwrap();
|
hdf5_writer::write_hdf5(h5_path.to_str().unwrap(), &data, &opts).unwrap();
|
||||||
|
|
||||||
let summary = validate::validate_hdf5(h5_path.to_str().unwrap(), &data, false, false).unwrap();
|
let summary =
|
||||||
|
validate::validate_hdf5(h5_path.to_str().unwrap(), &data, false, false).unwrap();
|
||||||
assert_eq!(summary.chunks, 0);
|
assert_eq!(summary.chunks, 0);
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -543,7 +556,8 @@ mod tests {
|
|||||||
}
|
}
|
||||||
drop(conn);
|
drop(conn);
|
||||||
|
|
||||||
let data = sqlite_reader::read_sqlite(&db_path, false, None, &SchemaConfig::default()).unwrap();
|
let data =
|
||||||
|
sqlite_reader::read_sqlite(&db_path, false, None, &SchemaConfig::default()).unwrap();
|
||||||
assert_eq!(data.chunks.len(), 1000);
|
assert_eq!(data.chunks.len(), 1000);
|
||||||
|
|
||||||
let opts = hdf5_writer::WriteOptions {
|
let opts = hdf5_writer::WriteOptions {
|
||||||
@@ -573,7 +587,8 @@ mod tests {
|
|||||||
insert_session(&conn, "session-gamma", 21, 30);
|
insert_session(&conn, "session-gamma", 21, 30);
|
||||||
drop(conn);
|
drop(conn);
|
||||||
|
|
||||||
let data = sqlite_reader::read_sqlite(&db_path, false, None, &SchemaConfig::default()).unwrap();
|
let data =
|
||||||
|
sqlite_reader::read_sqlite(&db_path, false, None, &SchemaConfig::default()).unwrap();
|
||||||
assert_eq!(data.sessions.len(), 3);
|
assert_eq!(data.sessions.len(), 3);
|
||||||
|
|
||||||
let opts = hdf5_writer::WriteOptions {
|
let opts = hdf5_writer::WriteOptions {
|
||||||
@@ -585,7 +600,8 @@ mod tests {
|
|||||||
};
|
};
|
||||||
hdf5_writer::write_hdf5(h5_path.to_str().unwrap(), &data, &opts).unwrap();
|
hdf5_writer::write_hdf5(h5_path.to_str().unwrap(), &data, &opts).unwrap();
|
||||||
|
|
||||||
let summary = validate::validate_hdf5(h5_path.to_str().unwrap(), &data, false, false).unwrap();
|
let summary =
|
||||||
|
validate::validate_hdf5(h5_path.to_str().unwrap(), &data, false, false).unwrap();
|
||||||
assert_eq!(summary.sessions, 3);
|
assert_eq!(summary.sessions, 3);
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -605,7 +621,8 @@ mod tests {
|
|||||||
insert_relation(&conn, 2, 3, "uses");
|
insert_relation(&conn, 2, 3, "uses");
|
||||||
drop(conn);
|
drop(conn);
|
||||||
|
|
||||||
let data = sqlite_reader::read_sqlite(&db_path, false, None, &SchemaConfig::default()).unwrap();
|
let data =
|
||||||
|
sqlite_reader::read_sqlite(&db_path, false, None, &SchemaConfig::default()).unwrap();
|
||||||
assert_eq!(data.entities.len(), 3);
|
assert_eq!(data.entities.len(), 3);
|
||||||
assert_eq!(data.relations.len(), 3);
|
assert_eq!(data.relations.len(), 3);
|
||||||
|
|
||||||
@@ -618,7 +635,8 @@ mod tests {
|
|||||||
};
|
};
|
||||||
hdf5_writer::write_hdf5(h5_path.to_str().unwrap(), &data, &opts).unwrap();
|
hdf5_writer::write_hdf5(h5_path.to_str().unwrap(), &data, &opts).unwrap();
|
||||||
|
|
||||||
let summary = validate::validate_hdf5(h5_path.to_str().unwrap(), &data, false, false).unwrap();
|
let summary =
|
||||||
|
validate::validate_hdf5(h5_path.to_str().unwrap(), &data, false, false).unwrap();
|
||||||
assert_eq!(summary.entities, 3);
|
assert_eq!(summary.entities, 3);
|
||||||
assert_eq!(summary.relations, 3);
|
assert_eq!(summary.relations, 3);
|
||||||
}
|
}
|
||||||
@@ -634,7 +652,8 @@ mod tests {
|
|||||||
insert_chunk(&conn, 1, "test", &make_embedding(4, 1.0), 0);
|
insert_chunk(&conn, 1, "test", &make_embedding(4, 1.0), 0);
|
||||||
drop(conn);
|
drop(conn);
|
||||||
|
|
||||||
let data = sqlite_reader::read_sqlite(&db_path, false, None, &SchemaConfig::default()).unwrap();
|
let data =
|
||||||
|
sqlite_reader::read_sqlite(&db_path, false, None, &SchemaConfig::default()).unwrap();
|
||||||
let opts = hdf5_writer::WriteOptions {
|
let opts = hdf5_writer::WriteOptions {
|
||||||
agent_id: "t".into(),
|
agent_id: "t".into(),
|
||||||
embedder: "t".into(),
|
embedder: "t".into(),
|
||||||
@@ -645,8 +664,8 @@ mod tests {
|
|||||||
hdf5_writer::write_hdf5(h5_path.to_str().unwrap(), &data, &opts).unwrap();
|
hdf5_writer::write_hdf5(h5_path.to_str().unwrap(), &data, &opts).unwrap();
|
||||||
|
|
||||||
// Validating against a source with an extra (unwritten) chunk must fail.
|
// Validating against a source with an extra (unwritten) chunk must fail.
|
||||||
let mut bigger = sqlite_reader::read_sqlite(&db_path, false, None, &SchemaConfig::default())
|
let mut bigger =
|
||||||
.unwrap();
|
sqlite_reader::read_sqlite(&db_path, false, None, &SchemaConfig::default()).unwrap();
|
||||||
let mut extra = bigger.chunks[0].clone();
|
let mut extra = bigger.chunks[0].clone();
|
||||||
extra.id = 999;
|
extra.id = 999;
|
||||||
bigger.chunks.push(extra);
|
bigger.chunks.push(extra);
|
||||||
@@ -666,7 +685,8 @@ mod tests {
|
|||||||
insert_chunk(&conn, 1, "test", &make_embedding(8, 1.0), 0);
|
insert_chunk(&conn, 1, "test", &make_embedding(8, 1.0), 0);
|
||||||
drop(conn);
|
drop(conn);
|
||||||
|
|
||||||
let data = sqlite_reader::read_sqlite(&db_path, false, None, &SchemaConfig::default()).unwrap();
|
let data =
|
||||||
|
sqlite_reader::read_sqlite(&db_path, false, None, &SchemaConfig::default()).unwrap();
|
||||||
let opts = hdf5_writer::WriteOptions {
|
let opts = hdf5_writer::WriteOptions {
|
||||||
agent_id: "my-agent-42".into(),
|
agent_id: "my-agent-42".into(),
|
||||||
embedder: "openai-ada".into(),
|
embedder: "openai-ada".into(),
|
||||||
@@ -712,7 +732,8 @@ mod tests {
|
|||||||
insert_chunk(&conn, 1, "test", &emb, 0);
|
insert_chunk(&conn, 1, "test", &emb, 0);
|
||||||
drop(conn);
|
drop(conn);
|
||||||
|
|
||||||
let data = sqlite_reader::read_sqlite(&db_path, false, None, &SchemaConfig::default()).unwrap();
|
let data =
|
||||||
|
sqlite_reader::read_sqlite(&db_path, false, None, &SchemaConfig::default()).unwrap();
|
||||||
let opts = hdf5_writer::WriteOptions {
|
let opts = hdf5_writer::WriteOptions {
|
||||||
agent_id: "t".into(),
|
agent_id: "t".into(),
|
||||||
embedder: "t".into(),
|
embedder: "t".into(),
|
||||||
@@ -758,7 +779,8 @@ mod tests {
|
|||||||
drop(conn);
|
drop(conn);
|
||||||
|
|
||||||
// Skip deleted
|
// Skip deleted
|
||||||
let data = sqlite_reader::read_sqlite(&db_path, true, None, &SchemaConfig::default()).unwrap();
|
let data =
|
||||||
|
sqlite_reader::read_sqlite(&db_path, true, None, &SchemaConfig::default()).unwrap();
|
||||||
assert_eq!(data.chunks.len(), 4); // chunk 3 is deleted
|
assert_eq!(data.chunks.len(), 4); // chunk 3 is deleted
|
||||||
|
|
||||||
let opts = hdf5_writer::WriteOptions {
|
let opts = hdf5_writer::WriteOptions {
|
||||||
@@ -770,7 +792,8 @@ mod tests {
|
|||||||
};
|
};
|
||||||
hdf5_writer::write_hdf5(h5_path.to_str().unwrap(), &data, &opts).unwrap();
|
hdf5_writer::write_hdf5(h5_path.to_str().unwrap(), &data, &opts).unwrap();
|
||||||
|
|
||||||
let summary = validate::validate_hdf5(h5_path.to_str().unwrap(), &data, false, false).unwrap();
|
let summary =
|
||||||
|
validate::validate_hdf5(h5_path.to_str().unwrap(), &data, false, false).unwrap();
|
||||||
assert_eq!(summary.chunks, 4);
|
assert_eq!(summary.chunks, 4);
|
||||||
assert_eq!(summary.sessions, 2);
|
assert_eq!(summary.sessions, 2);
|
||||||
assert_eq!(summary.entities, 2);
|
assert_eq!(summary.entities, 2);
|
||||||
@@ -789,7 +812,8 @@ mod tests {
|
|||||||
insert_session(&conn, "s1", 0, 10);
|
insert_session(&conn, "s1", 0, 10);
|
||||||
drop(conn);
|
drop(conn);
|
||||||
|
|
||||||
let data = sqlite_reader::read_sqlite(&db_path, false, None, &SchemaConfig::default()).unwrap();
|
let data =
|
||||||
|
sqlite_reader::read_sqlite(&db_path, false, None, &SchemaConfig::default()).unwrap();
|
||||||
let opts = hdf5_writer::WriteOptions {
|
let opts = hdf5_writer::WriteOptions {
|
||||||
agent_id: "t".into(),
|
agent_id: "t".into(),
|
||||||
embedder: "t".into(),
|
embedder: "t".into(),
|
||||||
@@ -800,8 +824,8 @@ mod tests {
|
|||||||
hdf5_writer::write_hdf5(h5_path.to_str().unwrap(), &data, &opts).unwrap();
|
hdf5_writer::write_hdf5(h5_path.to_str().unwrap(), &data, &opts).unwrap();
|
||||||
|
|
||||||
// Validating against a source whose session content differs must fail.
|
// Validating against a source whose session content differs must fail.
|
||||||
let mut tampered = sqlite_reader::read_sqlite(&db_path, false, None, &SchemaConfig::default())
|
let mut tampered =
|
||||||
.unwrap();
|
sqlite_reader::read_sqlite(&db_path, false, None, &SchemaConfig::default()).unwrap();
|
||||||
tampered.sessions[0].summary = "DIFFERENT".into();
|
tampered.sessions[0].summary = "DIFFERENT".into();
|
||||||
let result = validate::validate_hdf5(h5_path.to_str().unwrap(), &tampered, false, false);
|
let result = validate::validate_hdf5(h5_path.to_str().unwrap(), &tampered, false, false);
|
||||||
assert!(result.is_err());
|
assert!(result.is_err());
|
||||||
@@ -819,7 +843,8 @@ mod tests {
|
|||||||
insert_chunk(&conn, 1, "hello", &make_embedding(8, 1.0), 0);
|
insert_chunk(&conn, 1, "hello", &make_embedding(8, 1.0), 0);
|
||||||
drop(conn);
|
drop(conn);
|
||||||
|
|
||||||
let data = sqlite_reader::read_sqlite(&db_path, false, None, &SchemaConfig::default()).unwrap();
|
let data =
|
||||||
|
sqlite_reader::read_sqlite(&db_path, false, None, &SchemaConfig::default()).unwrap();
|
||||||
let opts = hdf5_writer::WriteOptions {
|
let opts = hdf5_writer::WriteOptions {
|
||||||
agent_id: "t".into(),
|
agent_id: "t".into(),
|
||||||
embedder: "t".into(),
|
embedder: "t".into(),
|
||||||
@@ -830,8 +855,8 @@ mod tests {
|
|||||||
hdf5_writer::write_hdf5(h5_path.to_str().unwrap(), &data, &opts).unwrap();
|
hdf5_writer::write_hdf5(h5_path.to_str().unwrap(), &data, &opts).unwrap();
|
||||||
|
|
||||||
// A source whose embedding differs (but counts match) must fail validation.
|
// A source whose embedding differs (but counts match) must fail validation.
|
||||||
let mut tampered = sqlite_reader::read_sqlite(&db_path, false, None, &SchemaConfig::default())
|
let mut tampered =
|
||||||
.unwrap();
|
sqlite_reader::read_sqlite(&db_path, false, None, &SchemaConfig::default()).unwrap();
|
||||||
tampered.chunks[0].embedding[3] += 9.0;
|
tampered.chunks[0].embedding[3] += 9.0;
|
||||||
let result = validate::validate_hdf5(h5_path.to_str().unwrap(), &tampered, true, false);
|
let result = validate::validate_hdf5(h5_path.to_str().unwrap(), &tampered, true, false);
|
||||||
assert!(result.is_err());
|
assert!(result.is_err());
|
||||||
@@ -856,7 +881,10 @@ mod tests {
|
|||||||
CREATE TABLE relations (src INTEGER, tgt INTEGER, relation TEXT, weight REAL, timestamp REAL);",
|
CREATE TABLE relations (src INTEGER, tgt INTEGER, relation TEXT, weight REAL, timestamp REAL);",
|
||||||
)
|
)
|
||||||
.unwrap();
|
.unwrap();
|
||||||
let blob: Vec<u8> = make_embedding(4, 1.0).iter().flat_map(|v| v.to_le_bytes()).collect();
|
let blob: Vec<u8> = make_embedding(4, 1.0)
|
||||||
|
.iter()
|
||||||
|
.flat_map(|v| v.to_le_bytes())
|
||||||
|
.collect();
|
||||||
conn.execute(
|
conn.execute(
|
||||||
"INSERT INTO my_chunks VALUES (1, 'hi', ?1, 'api', 1.0, 's', '', 0)",
|
"INSERT INTO my_chunks VALUES (1, 'hi', ?1, 'api', 1.0, 's', '', 0)",
|
||||||
rusqlite::params![blob],
|
rusqlite::params![blob],
|
||||||
@@ -908,7 +936,8 @@ mod tests {
|
|||||||
let base = hdf5_reader::read_hdf5(h5_path.to_str().unwrap()).unwrap();
|
let base = hdf5_reader::read_hdf5(h5_path.to_str().unwrap()).unwrap();
|
||||||
let max_id = base.chunks.iter().map(|c| c.id).max().unwrap_or(0);
|
let max_id = base.chunks.iter().map(|c| c.id).max().unwrap_or(0);
|
||||||
assert_eq!(max_id, 2);
|
assert_eq!(max_id, 2);
|
||||||
let new = sqlite_reader::read_sqlite_filtered(&db_path, false, Some(4), &cfg, max_id).unwrap();
|
let new =
|
||||||
|
sqlite_reader::read_sqlite_filtered(&db_path, false, Some(4), &cfg, max_id).unwrap();
|
||||||
assert_eq!(new.chunks.len(), 2); // only id 3 and 4
|
assert_eq!(new.chunks.len(), 2); // only id 3 and 4
|
||||||
|
|
||||||
let mut merged = base;
|
let mut merged = base;
|
||||||
|
|||||||
@@ -91,7 +91,14 @@ impl Default for SchemaConfig {
|
|||||||
},
|
},
|
||||||
sessions: TableSchema {
|
sessions: TableSchema {
|
||||||
table: "sessions".into(),
|
table: "sessions".into(),
|
||||||
columns: vec!["id", "start_idx", "end_idx", "channel", "timestamp", "summary"],
|
columns: vec![
|
||||||
|
"id",
|
||||||
|
"start_idx",
|
||||||
|
"end_idx",
|
||||||
|
"channel",
|
||||||
|
"timestamp",
|
||||||
|
"summary",
|
||||||
|
],
|
||||||
},
|
},
|
||||||
entities: TableSchema {
|
entities: TableSchema {
|
||||||
table: "entities".into(),
|
table: "entities".into(),
|
||||||
|
|||||||
@@ -77,10 +77,9 @@ pub fn validate_hdf5(
|
|||||||
}
|
}
|
||||||
for (k, (&a, &b)) in s.embedding.iter().zip(g.embedding.iter()).enumerate() {
|
for (k, (&a, &b)) in s.embedding.iter().zip(g.embedding.iter()).enumerate() {
|
||||||
if (a - b).abs() > emb_abs + emb_rel * a.abs() {
|
if (a - b).abs() > emb_abs + emb_rel * a.abs() {
|
||||||
return Err(format!(
|
return Err(
|
||||||
"chunk[{i}].embedding[{k}] mismatch: source {a}, HDF5 {b}"
|
format!("chunk[{i}].embedding[{k}] mismatch: source {a}, HDF5 {b}").into(),
|
||||||
)
|
);
|
||||||
.into());
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
rows_checked += 1;
|
rows_checked += 1;
|
||||||
@@ -108,7 +107,12 @@ pub fn validate_hdf5(
|
|||||||
}
|
}
|
||||||
rows_checked += 1;
|
rows_checked += 1;
|
||||||
}
|
}
|
||||||
for (i, (s, g)) in source.relations.iter().zip(got.relations.iter()).enumerate() {
|
for (i, (s, g)) in source
|
||||||
|
.relations
|
||||||
|
.iter()
|
||||||
|
.zip(got.relations.iter())
|
||||||
|
.enumerate()
|
||||||
|
{
|
||||||
if s.src != g.src || s.tgt != g.tgt || s.relation != g.relation {
|
if s.src != g.src || s.tgt != g.tgt || s.relation != g.relation {
|
||||||
return Err(format!("relation[{i}] mismatch").into());
|
return Err(format!("relation[{i}] mismatch").into());
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -494,7 +494,9 @@ mod tests {
|
|||||||
let file = File::open(&path).unwrap();
|
let file = File::open(&path).unwrap();
|
||||||
let ds = file.dataset("data").unwrap();
|
let ds = file.dataset("data").unwrap();
|
||||||
|
|
||||||
if let Ok(slice) = ds.read_f32_zerocopy() { assert_eq!(slice, &original[..]) }
|
if let Ok(slice) = ds.read_f32_zerocopy() {
|
||||||
|
assert_eq!(slice, &original[..])
|
||||||
|
}
|
||||||
assert_eq!(ds.read_f32().unwrap(), original);
|
assert_eq!(ds.read_f32().unwrap(), original);
|
||||||
|
|
||||||
std::fs::remove_file(&path).ok();
|
std::fs::remove_file(&path).ok();
|
||||||
|
|||||||
@@ -715,7 +715,9 @@ fn multiple_chunked_datasets_share_file_cache() {
|
|||||||
use clawhdf5_format::datatype::{CharacterSet, Datatype, StringPadding};
|
use clawhdf5_format::datatype::{CharacterSet, Datatype, StringPadding};
|
||||||
|
|
||||||
// 1-D chunked + compressed fixed-length strings (payload > compress threshold).
|
// 1-D chunked + compressed fixed-length strings (payload > compress threshold).
|
||||||
let strings: Vec<String> = (0..64).map(|i| format!("entry-{i:06}-{}", "x".repeat(80))).collect();
|
let strings: Vec<String> = (0..64)
|
||||||
|
.map(|i| format!("entry-{i:06}-{}", "x".repeat(80)))
|
||||||
|
.collect();
|
||||||
let max_len = strings.iter().map(|s| s.len()).max().unwrap();
|
let max_len = strings.iter().map(|s| s.len()).max().unwrap();
|
||||||
let mut sraw = Vec::new();
|
let mut sraw = Vec::new();
|
||||||
for s in &strings {
|
for s in &strings {
|
||||||
@@ -743,7 +745,9 @@ fn multiple_chunked_datasets_share_file_cache() {
|
|||||||
{
|
{
|
||||||
let ds = b.create_dataset("mat");
|
let ds = b.create_dataset("mat");
|
||||||
ds.with_f32_data(&mat).with_shape(&[n as u64, d as u64]);
|
ds.with_f32_data(&mat).with_shape(&[n as u64, d as u64]);
|
||||||
ds.with_chunks(&[10, d as u64]).with_shuffle().with_deflate(6);
|
ds.with_chunks(&[10, d as u64])
|
||||||
|
.with_shuffle()
|
||||||
|
.with_deflate(6);
|
||||||
}
|
}
|
||||||
let bytes = b.finish().unwrap();
|
let bytes = b.finish().unwrap();
|
||||||
let file = File::from_bytes(bytes).unwrap();
|
let file = File::from_bytes(bytes).unwrap();
|
||||||
@@ -755,7 +759,10 @@ fn multiple_chunked_datasets_share_file_cache() {
|
|||||||
let got_mat = file.dataset("mat").unwrap().read_f32().unwrap();
|
let got_mat = file.dataset("mat").unwrap().read_f32().unwrap();
|
||||||
assert_eq!(got_mat, mat);
|
assert_eq!(got_mat, mat);
|
||||||
// Read the 1-D one again to confirm the cache rebinds back correctly.
|
// Read the 1-D one again to confirm the cache rebinds back correctly.
|
||||||
assert_eq!(file.dataset("strs").unwrap().read_string().unwrap(), strings);
|
assert_eq!(
|
||||||
|
file.dataset("strs").unwrap().read_string().unwrap(),
|
||||||
|
strings
|
||||||
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
@@ -850,8 +857,14 @@ fn dense_group_links_roundtrip() {
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
// The small (compact) group still works.
|
// The small (compact) group still works.
|
||||||
assert_eq!(file.dataset("small/a").unwrap().read_f64().unwrap(), vec![1.0]);
|
assert_eq!(
|
||||||
assert_eq!(file.dataset("small/b").unwrap().read_f64().unwrap(), vec![2.0]);
|
file.dataset("small/a").unwrap().read_f64().unwrap(),
|
||||||
|
vec![1.0]
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
file.dataset("small/b").unwrap().read_f64().unwrap(),
|
||||||
|
vec![2.0]
|
||||||
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
@@ -908,7 +921,10 @@ fn dense_links_multiblock_fractal_heap_roundtrip() {
|
|||||||
b.add_group(g.finish());
|
b.add_group(g.finish());
|
||||||
let file = File::from_bytes(b.finish().unwrap()).unwrap();
|
let file = File::from_bytes(b.finish().unwrap()).unwrap();
|
||||||
|
|
||||||
assert_eq!(file.group("big").unwrap().datasets().unwrap().len(), n as usize);
|
assert_eq!(
|
||||||
|
file.group("big").unwrap().datasets().unwrap().len(),
|
||||||
|
n as usize
|
||||||
|
);
|
||||||
for i in [0, 1, 1234, n - 1] {
|
for i in [0, 1, 1234, n - 1] {
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
file.dataset(&format!("big/dataset_number_{i:05}"))
|
file.dataset(&format!("big/dataset_number_{i:05}"))
|
||||||
|
|||||||
@@ -1,5 +1,5 @@
|
|||||||
#!/usr/bin/env bash
|
#!/usr/bin/env bash
|
||||||
# CI check: verify rustyhdf5-format compiles under no_std (thumbv7em-none-eabihf).
|
# CI check: verify clawhdf5-format compiles under no_std (thumbv7em-none-eabihf).
|
||||||
#
|
#
|
||||||
# Usage:
|
# Usage:
|
||||||
# ./scripts/check-nostd.sh
|
# ./scripts/check-nostd.sh
|
||||||
@@ -11,7 +11,7 @@ set -euo pipefail
|
|||||||
|
|
||||||
TARGET="thumbv7em-none-eabihf"
|
TARGET="thumbv7em-none-eabihf"
|
||||||
|
|
||||||
echo "==> Checking no_std build for rustyhdf5-format (target: $TARGET)"
|
echo "==> Checking no_std build for clawhdf5-format (target: $TARGET)"
|
||||||
|
|
||||||
# Ensure the target is installed
|
# Ensure the target is installed
|
||||||
if ! rustup target list --installed | grep -q "$TARGET"; then
|
if ! rustup target list --installed | grep -q "$TARGET"; then
|
||||||
@@ -20,13 +20,13 @@ if ! rustup target list --installed | grep -q "$TARGET"; then
|
|||||||
fi
|
fi
|
||||||
|
|
||||||
# Build with no default features (no std, no flate2, no sha2)
|
# Build with no default features (no std, no flate2, no sha2)
|
||||||
cargo build --target "$TARGET" -p rustyhdf5-format --no-default-features
|
cargo build --target "$TARGET" -p clawhdf5-format --no-default-features
|
||||||
|
|
||||||
echo "==> no_std build succeeded"
|
echo "==> no_std build succeeded"
|
||||||
|
|
||||||
# Also verify the default-features (std) build still works
|
# Also verify the default-features (std) build still works
|
||||||
echo "==> Checking default-features build"
|
echo "==> Checking default-features build"
|
||||||
cargo build -p rustyhdf5-format
|
cargo build -p clawhdf5-format
|
||||||
echo "==> default-features build succeeded"
|
echo "==> default-features build succeeded"
|
||||||
|
|
||||||
echo "==> All no_std checks passed"
|
echo "==> All no_std checks passed"
|
||||||
|
|||||||
+4
-4
@@ -34,16 +34,16 @@ run_step() {
|
|||||||
# 1. Format check
|
# 1. Format check
|
||||||
run_step "cargo fmt --check" cargo fmt --check
|
run_step "cargo fmt --check" cargo fmt --check
|
||||||
|
|
||||||
# 2. Clippy (exclude rustyhdf5-py which needs PyO3/Python)
|
# 2. Clippy (exclude clawhdf5-py which needs PyO3/Python)
|
||||||
run_step "cargo clippy" cargo clippy \
|
run_step "cargo clippy" cargo clippy \
|
||||||
--workspace \
|
--workspace \
|
||||||
--exclude rustyhdf5-py \
|
--exclude clawhdf5-py \
|
||||||
-- -D warnings
|
-- -D warnings
|
||||||
|
|
||||||
# 3. Tests (exclude rustyhdf5-py)
|
# 3. Tests (exclude clawhdf5-py)
|
||||||
run_step "cargo test" cargo test \
|
run_step "cargo test" cargo test \
|
||||||
--workspace \
|
--workspace \
|
||||||
--exclude rustyhdf5-py
|
--exclude clawhdf5-py
|
||||||
|
|
||||||
# 4. no_std check
|
# 4. no_std check
|
||||||
run_step "check-nostd.sh" "$SCRIPT_DIR/check-nostd.sh"
|
run_step "check-nostd.sh" "$SCRIPT_DIR/check-nostd.sh"
|
||||||
|
|||||||
Reference in New Issue
Block a user