Compare commits
130
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
72b9cfb1e1 | ||
|
|
650f355219 | ||
|
|
7f5cfee281 | ||
|
|
c5302e587e | ||
|
|
e1115bc92a | ||
|
|
36d7a6f234 | ||
|
|
4b23ad697c | ||
|
|
e7f2d8575d | ||
|
|
7c1968a34a | ||
|
|
6db13c60b8 | ||
|
|
d99426be94 | ||
|
|
f5505fb03d | ||
|
|
95dcb04454 | ||
|
|
1dba7b465a | ||
|
|
57e938c438 | ||
|
|
3000b40cf3 | ||
|
|
bc820fbd8c | ||
|
|
9066d34eaa | ||
|
|
540fa08907 | ||
|
|
14876b8ae5 | ||
|
|
5935e13866 | ||
|
|
8c3ef996ea | ||
|
|
74fdf0582b | ||
|
|
44f5f8b5c5 | ||
|
|
2f252df084 | ||
|
|
c8c2930fc0 | ||
|
|
417c9516ca | ||
|
|
53dbddb07b | ||
|
|
081341b433 | ||
|
|
d074385944 | ||
|
|
4a1876faf2 | ||
|
|
be88e3fec7 | ||
|
|
e162c013fd | ||
|
|
06dda26d85 | ||
|
|
585e14d5e2 | ||
|
|
183d96ee26 | ||
|
|
bba1560416 | ||
|
|
b36998ef01 | ||
|
|
aef8e766ae | ||
|
|
9ea44d473d | ||
|
|
46203ea761 | ||
|
|
75bdb53342 | ||
|
|
dd5b3f6633 | ||
|
|
79dfa78e8f | ||
|
|
87d64588e5 | ||
|
|
bdadf3447c | ||
|
|
0c65a27b00 | ||
|
|
db9af7972c | ||
|
|
7706697feb | ||
|
|
4ecac65f22 | ||
|
|
c0f704c381 | ||
|
|
00b0cb0035 | ||
|
|
a7920bd4b3 | ||
|
|
7e43b5366c | ||
|
|
dce5559ff2 | ||
|
|
1b3bbb054a | ||
|
|
a8fb758489 | ||
|
|
73bb068264 | ||
|
|
5c8323cb1e | ||
|
|
dbaf3f505d | ||
|
|
c470244a6f | ||
|
|
d0db83812b | ||
|
|
5e4aa1c6bf | ||
|
|
1cceb930b2 | ||
|
|
fc7ae6549a | ||
|
|
735db117a7 | ||
|
|
e9b37a9602 | ||
|
|
4bed8b3765 | ||
|
|
36d689bc2c | ||
|
|
e7c08e06b4 | ||
|
|
c5049eb734 | ||
|
|
6f6bc97850 | ||
|
|
0cb72e8a60 | ||
|
|
e338d58ad5 | ||
|
|
8b85d9364b | ||
|
|
6598a7d02f | ||
|
|
114a2dfcba | ||
|
|
56a8c2f3d0 | ||
|
|
7b16dc90d6 | ||
|
|
4a5544da1d | ||
|
|
f4c6d43a3f | ||
|
|
e5e087f9ab | ||
|
|
16c9ee0554 | ||
|
|
1d767e3b93 | ||
|
|
a91df3f1c3 | ||
|
|
b41272487a | ||
|
|
0bc7a293ae | ||
|
|
dea02f5214 | ||
|
|
97e65f2adf | ||
|
|
fb58300b3f | ||
|
|
eb196e824f | ||
|
|
367faad7f7 | ||
|
|
0901fb1499 | ||
|
|
e9aeb110b7 | ||
|
|
5889b378e9 | ||
|
|
18dc35f7e5 | ||
|
|
e17ab0ceef | ||
|
|
105cf13347 | ||
|
|
0529f72a2c | ||
|
|
a29c1b224b | ||
|
|
c0a9206703 | ||
|
|
57756e69ec | ||
|
|
6ad8ceb426 | ||
|
|
2e7e0456c1 | ||
|
|
dc0113d015 | ||
|
|
8ea455bbcb | ||
|
|
1e18ff5a86 | ||
|
|
306a35347c | ||
|
|
4c60398b30 | ||
|
|
7155409202 | ||
|
|
64d9c5f171 | ||
|
|
84a39ef3c5 | ||
|
|
55ed87d2e8 | ||
|
|
6531158d9f | ||
|
|
aa92fef7bb | ||
|
|
29baabbed2 | ||
|
|
b23946e62d | ||
|
|
52cfcf20b2 | ||
|
|
05c665a898 | ||
|
|
b36c6ec2af | ||
|
|
f3d63dbdcd | ||
|
|
0addf328bc | ||
|
|
d668e45ab5 | ||
|
|
c6a7bbfc67 | ||
|
|
3027380979 | ||
|
|
f507803ec1 | ||
|
|
8803d0754b | ||
|
|
42c3872ec9 | ||
|
|
c19199f3eb | ||
|
|
41db450c92 |
+58
-12
@@ -9,15 +9,15 @@ jobs:
|
|||||||
runs-on: ubuntu-latest
|
runs-on: ubuntu-latest
|
||||||
container: rust:latest
|
container: rust:latest
|
||||||
steps:
|
steps:
|
||||||
- uses: actions/checkout@v4
|
# Plain git rather than actions/checkout: that is a JavaScript action,
|
||||||
- name: Cache cargo registry/target
|
# and rust:latest has no `node`, so it failed with exit 127 before any
|
||||||
uses: actions/cache@v4
|
# code was built — on every push. actions/cache went for the same reason.
|
||||||
with:
|
- name: Check out
|
||||||
path: |
|
run: |
|
||||||
~/.cargo/registry
|
git init -q .
|
||||||
~/.cargo/git
|
git remote add origin "${GITHUB_SERVER_URL}/${GITHUB_REPOSITORY}.git"
|
||||||
target
|
for i in 1 2 3; do git fetch -q --depth 1 origin "${GITHUB_SHA}" && break; sleep 5; done
|
||||||
key: ${{ runner.os }}-cargo-${{ hashFiles('**/Cargo.lock') }}
|
git checkout -q FETCH_HEAD
|
||||||
- name: Install rustfmt & clippy components
|
- name: Install rustfmt & clippy components
|
||||||
run: rustup component add rustfmt clippy
|
run: rustup component add rustfmt clippy
|
||||||
- name: Install thumbv7em-none-eabihf target
|
- name: Install thumbv7em-none-eabihf target
|
||||||
@@ -28,13 +28,59 @@ jobs:
|
|||||||
# dependency a failure (CLAWHDF5_REQUIRE_INTEROP below).
|
# dependency a failure (CLAWHDF5_REQUIRE_INTEROP below).
|
||||||
run: |
|
run: |
|
||||||
apt-get update
|
apt-get update
|
||||||
apt-get install -y --no-install-recommends python3 python3-venv
|
# cmake builds libz-ng-sys for the opt-in `fast-deflate` (zlib-ng)
|
||||||
|
# steps in ci-test.sh; rust:latest does not ship it. The default
|
||||||
|
# build (pure-Rust zlib-rs) does not need it.
|
||||||
|
apt-get install -y --no-install-recommends python3 python3-venv cmake
|
||||||
python3 -m venv /opt/interop
|
python3 -m venv /opt/interop
|
||||||
/opt/interop/bin/pip install --no-cache-dir h5py numpy netCDF4 xarray
|
/opt/interop/bin/pip install --no-cache-dir h5py numpy netCDF4 xarray hdf5plugin
|
||||||
echo "/opt/interop/bin" >> "$GITHUB_PATH"
|
echo "/opt/interop/bin" >> "$GITHUB_PATH"
|
||||||
- name: Show interop library versions
|
- name: Show interop library versions
|
||||||
run: python3 -c "import h5py, netCDF4; print('h5py', h5py.__version__, 'HDF5', h5py.version.hdf5_version, 'netCDF4', netCDF4.__version__)"
|
run: /opt/interop/bin/python -c "import h5py, netCDF4, hdf5plugin; print('h5py', h5py.__version__, 'HDF5', h5py.version.hdf5_version, 'netCDF4', netCDF4.__version__, 'hdf5plugin', hdf5plugin.version)"
|
||||||
- name: Run CI script
|
- name: Run CI script
|
||||||
env:
|
env:
|
||||||
|
# Name the interpreter outright rather than relying on $GITHUB_PATH
|
||||||
|
# reaching the test processes: if `python3` resolved to the system
|
||||||
|
# one instead of the venv, every interop suite would skip.
|
||||||
|
# CLAWHDF5_REQUIRE_INTEROP turns that skip into a failure, so the
|
||||||
|
# two together mean the suites either run or the build goes red.
|
||||||
|
CLAWHDF5_PYTHON: /opt/interop/bin/python
|
||||||
CLAWHDF5_REQUIRE_INTEROP: "1"
|
CLAWHDF5_REQUIRE_INTEROP: "1"
|
||||||
run: bash scripts/ci-test.sh
|
run: bash scripts/ci-test.sh
|
||||||
|
|
||||||
|
test-arm64:
|
||||||
|
# The aarch64 kernels in clawhdf5-accel — NEON `dot_i8`, including the
|
||||||
|
# SDOT path, and the f32 NEON kernels — are cfg'd out on x86, so the job
|
||||||
|
# above never compiles, lints or tests them.
|
||||||
|
#
|
||||||
|
# `linux_arm64` is served by two runners that execute differently:
|
||||||
|
# vision-01 runs steps on the host (Rust already installed) and vision-02
|
||||||
|
# runs them in docker.gitea.com/runner-images. So the steps work in both:
|
||||||
|
# no `container:`, no JavaScript actions (they are fetched from GitHub,
|
||||||
|
# which not every runner reliably reaches), and an explicit `+stable`
|
||||||
|
# toolchain rather than whatever a host happens to default to.
|
||||||
|
runs-on: linux_arm64
|
||||||
|
env:
|
||||||
|
CARGO_NET_RETRY: "10"
|
||||||
|
CARGO_TERM_COLOR: always
|
||||||
|
steps:
|
||||||
|
- name: Check out
|
||||||
|
run: |
|
||||||
|
git init -q .
|
||||||
|
git remote add origin "${GITHUB_SERVER_URL}/${GITHUB_REPOSITORY}.git"
|
||||||
|
for i in 1 2 3; do git fetch -q --depth 1 origin "${GITHUB_SHA}" && break; sleep 5; done
|
||||||
|
git checkout -q FETCH_HEAD
|
||||||
|
- name: Rust stable
|
||||||
|
run: |
|
||||||
|
export PATH="$HOME/.cargo/bin:$PATH"
|
||||||
|
command -v rustup >/dev/null || curl -sSf --retry 5 https://sh.rustup.rs | sh -s -- -y --profile minimal --default-toolchain none
|
||||||
|
rustup toolchain install stable --profile minimal --component clippy
|
||||||
|
echo "$HOME/.cargo/bin" >> "$GITHUB_PATH"
|
||||||
|
- name: Confirm aarch64
|
||||||
|
run: |
|
||||||
|
test "$(uname -m)" = aarch64
|
||||||
|
if grep -q asimddp /proc/cpuinfo; then echo "dot-product extension present: SDOT kernel runs"; else echo "no dot-product extension: plain NEON kernel runs"; fi
|
||||||
|
- name: Clippy (aarch64 kernels)
|
||||||
|
run: cargo +stable clippy -p clawhdf5-accel --all-targets -- -D warnings
|
||||||
|
- name: Test
|
||||||
|
run: cargo +stable test -p clawhdf5-accel -p clawhdf5-ann -p clawhdf5-format
|
||||||
|
|||||||
@@ -4,3 +4,4 @@ benchmarks/longmemeval/*.json
|
|||||||
|
|
||||||
# Local model weights (MiniLM etc.) — large, not committed
|
# Local model weights (MiniLM etc.) — large, not committed
|
||||||
weights/
|
weights/
|
||||||
|
.venv
|
||||||
|
|||||||
+1183
-140
File diff suppressed because it is too large
Load Diff
+703
@@ -1,5 +1,708 @@
|
|||||||
# Changelog
|
# Changelog
|
||||||
|
|
||||||
|
## Unreleased
|
||||||
|
|
||||||
|
### Upgrade Notes
|
||||||
|
- **HDF5 correctness audit (2026-09-25).** A sweep of 686 public files (the
|
||||||
|
libhdf5 test files, the HDF Group's CVE reproducers, pyfive, netcdf-c,
|
||||||
|
netcdf4-python, h5wasm, h5py and xarray corpora), a 567-case read matrix and
|
||||||
|
a 96-case write matrix against HDF5 1.10–2.0 found bugs that returned wrong
|
||||||
|
values with no error, and files we wrote that libhdf5 rejects. The fixes are
|
||||||
|
listed under Correctness and Interop. What changes for callers:
|
||||||
|
- **Chunked datasets whose max shape is larger than their current shape**,
|
||||||
|
or whose unlimited dimension is not the first, were indexed by the current
|
||||||
|
shape instead of the max shape, both when read and when written. Files from
|
||||||
|
libhdf5 now read correctly. Files clawhdf5 wrote with such a max shape were
|
||||||
|
laid out wrongly and now read the way libhdf5 always read them — rewrite
|
||||||
|
them. Agent stores and ClawBrainHub files have no max shape and are
|
||||||
|
unaffected.
|
||||||
|
- Integer reads (`read_i32`/`read_i64`/`read_u64`/...) of float data now
|
||||||
|
convert (truncate toward zero, saturate at the type's range, NaN reads as
|
||||||
|
0) instead of returning the IEEE bit pattern, and out-of-range integers
|
||||||
|
saturate instead of keeping the low bits.
|
||||||
|
- `FileWriter::finish()` now returns an error instead of writing a corrupt
|
||||||
|
file for: a header message over 64 KiB (e.g. an attribute larger than
|
||||||
|
~64 KiB), a group/dataset/link name that is empty, `.` or contains `/`
|
||||||
|
(nested paths were written as one literal link), a max shape smaller than
|
||||||
|
the shape, a page size outside 512 B–1 GiB, and more than 65 535 chunks in
|
||||||
|
a dataset with several unlimited dimensions.
|
||||||
|
- **Breaking (format crate):** `ObjectHeaderWriter::serialize`,
|
||||||
|
`BatchObjectHeaderWriter::compute_sizes`/`serialize_all` and
|
||||||
|
`build_chunked_data_from_precompressed` return `Result`;
|
||||||
|
`read_fixed_array_chunks`/`read_extensible_array_chunks` take `max_dims`;
|
||||||
|
`build_fixed_array_at`/`ea_writer::build_extensible_array_at` take one
|
||||||
|
`Option<WrittenChunk>` per index slot; `fill_value::dataset_fill_value`
|
||||||
|
returns `UnresolvedSharedMessage` for a shared message it cannot resolve
|
||||||
|
instead of `None`. `FillTime::default()` is `IfSet` (libhdf5's default;
|
||||||
|
default files are byte-identical).
|
||||||
|
- **ZeroClaw does not use clawhdf5.** The project described itself as
|
||||||
|
ZeroClaw's memory backend ("imported as a `clawhdf5` Cargo feature"). Checked
|
||||||
|
against ZeroClaw v0.8.5 (the latest release), the `osobh/zeroclaw` fork and
|
||||||
|
their full history: no such feature or backend has ever existed. And
|
||||||
|
`clawhdf5-migrate`'s "ZeroClaw layout" (`memory_chunks`, `sessions`,
|
||||||
|
`entities`, `relations`) is not ZeroClaw's schema — ZeroClaw uses a single
|
||||||
|
`memories` table — so the migrator cannot read a ZeroClaw database. The
|
||||||
|
claims are withdrawn; the migrator's layout is documented as its own.
|
||||||
|
- **OpenClaw is not supported, and never was.** The docs described a
|
||||||
|
"drop-in" OpenClaw memory backend enabled with `memory.backend = "clawhdf5"`.
|
||||||
|
That config was never valid in any OpenClaw release (v2026.2–v2026.7
|
||||||
|
accepted only `builtin`/`qmd` and rejected unknown keys, so a Gateway given
|
||||||
|
it refuses to start; OpenClaw 2.0 removed the key), no plugin was ever built,
|
||||||
|
and `@redclaw/clawhdf5` was never published. The integration docs
|
||||||
|
(`openclaw-integration.md`, `openclaw-config.md`, `migration-guide.md`) are
|
||||||
|
removed; `docs/openclaw.md` explains the status and what a real plugin would
|
||||||
|
need against OpenClaw v2026.9.6. `ClawhdfBackend` stays as a library API.
|
||||||
|
- **Breaking:** `MemoryError` is now `#[non_exhaustive]` and gained
|
||||||
|
`SigningKeyRequired`; a `match` on it needs a wildcard arm. Future variants
|
||||||
|
will no longer be breaking.
|
||||||
|
- **Breaking:** `clawhdf5-agent`'s `agent` feature is removed. It enabled
|
||||||
|
nothing — the agent layer is always built — but the README and guides told
|
||||||
|
people to pass it; drop `agent` from `features = [...]`.
|
||||||
|
- **`clawhdf5-migrate` now writes a real agent store.** Its output used to be
|
||||||
|
a layout of its own (`/chunks`, `/sessions`, `/entities`, `/relations`, no
|
||||||
|
`/meta`) that `HDF5Memory::open` rejected, so a migrated file could not be
|
||||||
|
used as agent memory. Files it wrote before this release are not agent
|
||||||
|
stores; re-run the migration. Also: embeddings default to `float16` like
|
||||||
|
any new store (`--f32` opts out; `--float16` is a hidden no-op); a row with
|
||||||
|
the wrong embedding length is an error instead of being truncated or
|
||||||
|
padded; `--incremental` now matches rows by content against an existing
|
||||||
|
store and follows the source's deleted flags; a source with no memory rows
|
||||||
|
needs `--embedding-dim`. The per-dataset SHA-256 provenance attributes of
|
||||||
|
the old layout are gone (the agent schema has no place for them).
|
||||||
|
- **Files written by clawhdf5 now open in h5py and libhdf5.** Every `f32`
|
||||||
|
dataset we wrote — including every agent store's embeddings — was refused
|
||||||
|
with "sign bit position out of bounds", and every empty dataset with
|
||||||
|
"invalid dataset size". Both were write-side bugs present in every release;
|
||||||
|
clawhdf5's own reader was unaffected. An agent store is rewritten in full at
|
||||||
|
each checkpoint, so it becomes readable at its next checkpoint on this
|
||||||
|
version; other files with `f32` or empty datasets need rewriting. Details in
|
||||||
|
`docs/known-issues.md`.
|
||||||
|
- **New stores store embeddings as half precision by default.**
|
||||||
|
`MemoryConfig::float16` was persisted and otherwise ignored; it now writes
|
||||||
|
`float16` embeddings (48% smaller files at 100K) and rounds each embedding
|
||||||
|
to half precision as it is saved — and it defaults to `true` for new
|
||||||
|
stores. On the full LongMemEval haystack with real MiniLM embeddings every
|
||||||
|
retrieval metric matched `f32`. **Existing stores are unaffected**: every
|
||||||
|
agent store has recorded `float16 = false`, and keeps it (a v2.5.0 fixture
|
||||||
|
guards this). A store that already had `float16 = true` rounds its
|
||||||
|
embeddings when next opened and writes them as `float16` at its next
|
||||||
|
checkpoint. Opt out with `float16 = false` or `create --f32`; the CLI's
|
||||||
|
`--float16` is still accepted and now a no-op. Values beyond ±65504 are
|
||||||
|
refused, so keep `f32` for unnormalised vectors.
|
||||||
|
- **Breaking:** `MemoryError` gained `InvalidEntry`, returned when a
|
||||||
|
`float16` store is given an embedding value beyond ±65504. Exhaustive
|
||||||
|
matches need the new arm.
|
||||||
|
- **The default build no longer compiles any C.** Deflate now defaults to the
|
||||||
|
pure-Rust zlib-rs instead of zlib-ng, so building the core crates needs
|
||||||
|
neither cmake nor a C compiler. Speed on HDF5 reads and writes is within 6%
|
||||||
|
of zlib-ng, and compressed output is byte-identical. To keep zlib-ng, enable
|
||||||
|
`fast-deflate` (on `clawhdf5`, `clawhdf5-format` or `clawhdf5-filters`); it
|
||||||
|
overrides zlib-rs wherever it is on.
|
||||||
|
- **A truncated deflate chunk is now an error.** It used to read back short,
|
||||||
|
with no error.
|
||||||
|
- **Minimum supported Rust is 1.92**, now declared in every crate's
|
||||||
|
`rust-version` and checked in CI.
|
||||||
|
- **New stores use the int8 vector index by default.**
|
||||||
|
`MemoryConfig::quantized_index` now defaults to `true`: a quarter of the
|
||||||
|
index memory, builds 1.8x (x86-64) and 2.3x (Raspberry Pi 5) faster, and
|
||||||
|
searches 1.63x and 1.18x faster at equal recall, measured on every
|
||||||
|
configuration tested. **Existing stores are unaffected** — a store written
|
||||||
|
with v2.6.0 or later keeps its persisted setting, and one written before the
|
||||||
|
setting existed opens as `false` and keeps its f32 index. Set
|
||||||
|
`quantized_index = false`, or pass `create --f32-index` to the CLI, to opt
|
||||||
|
out. The CLI's `--quantized-index` is still accepted but is now a no-op.
|
||||||
|
|
||||||
|
### Signing
|
||||||
|
- `clawhdf5-agent`: **Ed25519-signed checkpoints** — the README's
|
||||||
|
"cryptographically verifiable memory", now true. With
|
||||||
|
`HDF5Memory::set_signing_key(key)`, every checkpoint stores a signed
|
||||||
|
manifest: a SHA-256 per record (text, embedding as stored, channel,
|
||||||
|
timestamp, session, tags, deleted flag, activation) in a Merkle tree, plus
|
||||||
|
hashes of the settings (and WAL mark), sessions and knowledge graph, with
|
||||||
|
the per-record hashes in `/integrity/record_hashes`.
|
||||||
|
`HDF5Memory::verify(path, &public_key)` recomputes everything from the file
|
||||||
|
and reports which part changed and which records (`changed_records`); a
|
||||||
|
forged manifest fails the signature. The key is never persisted; a signed
|
||||||
|
store refuses to checkpoint without it (`MemoryError::SigningKeyRequired`),
|
||||||
|
and `remove_signature()` is the deliberate way back to unsigned. Saves still
|
||||||
|
in the WAL are not covered (`wal_entries_unsigned`). Tests include every
|
||||||
|
kind of edit, and an edit made with h5py in place, which verify pinpoints.
|
||||||
|
Cost: ~20% of a checkpoint, 32 bytes per record (`BENCHMARKS.md`, "Signed
|
||||||
|
checkpoints"). New dependencies `ed25519-dalek`, `sha2`, `rand_core` — pure
|
||||||
|
Rust; the no-C check still passes.
|
||||||
|
- `clawhdf5-cli`: `keygen --out <file>` (owner-only key file),
|
||||||
|
`--signing-key <file>` / `CLAWHDF5_SIGNING_KEY` on writing commands
|
||||||
|
(`create` signs immediately), `verify --public-key <hex|file>` (JSON report;
|
||||||
|
exit status 2 if not valid), and `signed` in `create`/`stats` output.
|
||||||
|
|
||||||
|
### Migration
|
||||||
|
- `clawhdf5-migrate`: writes through the agent's own API (`HDF5Memory::create`
|
||||||
|
/ `open`, `save_batch`, the session cache and knowledge graph), so there is
|
||||||
|
no second copy of the schema. Sessions and entities/relations carry over;
|
||||||
|
deleted rows become deleted records (or are left out with
|
||||||
|
`--skip-deleted`). Every source row is checked before the output is created,
|
||||||
|
so a source that cannot be migrated leaves an existing store untouched.
|
||||||
|
Validation reads the result back with `HDF5Memory::open_read_only`, compares
|
||||||
|
every field (embeddings bit for bit — `round_to_f16` of the source for a
|
||||||
|
`float16` store) and checks that a migrated record is found by search. The
|
||||||
|
`half`-based conversion is gone; `clawhdf5_format::float16` is the only one.
|
||||||
|
42 tests, including h5py opening a migrated store; an adversarial review's
|
||||||
|
two blocker and four major findings are fixed with regression tests.
|
||||||
|
- `clawhdf5-agent`: `HDF5Memory::sessions()` / `sessions_mut()`,
|
||||||
|
`HDF5Memory::delete_batch(&[usize])` (one save, all-or-nothing, never
|
||||||
|
auto-compacts), `SessionCache::add_at`, and `SessionCache` / `SessionEntry`
|
||||||
|
re-exported from the crate root.
|
||||||
|
|
||||||
|
### Search
|
||||||
|
- `clawhdf5-agent`: **`HDF5Memory::search` with `SearchOptions`** — source
|
||||||
|
filtering, re-ranking and confidence rejection in the store's own search
|
||||||
|
path. Re-ranking and confidence rejection used to be reachable only
|
||||||
|
through the OpenClaw backend, which now calls `search` with both on.
|
||||||
|
- `with_sources([..])` restricts a search to records from those source
|
||||||
|
channels. It applies before ranking, so a filtered search still returns up
|
||||||
|
to `k` results, normalised over what it can return. Measured at 100K: the
|
||||||
|
exact filtered top 10 for filters keeping 50%, 10% and 1% of the store and
|
||||||
|
for records far from the query, and never slower than an unfiltered search
|
||||||
|
(2.3 ms for a 1% filter vs 4.6 ms unfiltered). See `BENCHMARKS.md`,
|
||||||
|
"Search options".
|
||||||
|
- `with_rerank(ReRankConfig)` re-ranks a pool of `max(3k, 10)` candidates
|
||||||
|
(`rerank_pool` to change it) by relevance, recency, source authority and
|
||||||
|
activation; `with_confidence(ConfidenceConfig)` drops low-confidence
|
||||||
|
results; `at_time(now)` pins the clock for recency. About 3% on latency.
|
||||||
|
- `hybrid_search` and `hybrid_search_with` are unchanged (tested bit for
|
||||||
|
bit against `search` with default options).
|
||||||
|
- `clawhdf5-agent`: the OpenClaw backend's search now boosts the Hebbian
|
||||||
|
activation of the `k` results it returns, not of the whole `3k` candidate
|
||||||
|
pool it re-ranks.
|
||||||
|
|
||||||
|
### Documentation
|
||||||
|
- OpenClaw claims withdrawn across the README, QUICKSTART, USE_CASES, ROADMAP
|
||||||
|
(Track 7 marked withdrawn) and the `openclaw` module docs; the dead
|
||||||
|
`github.com/redclawsystems/openclaw` link is gone. The Node package is
|
||||||
|
marked unpublished and broken (now `"private": true` so it cannot be
|
||||||
|
published by accident), with its bugs recorded in `docs/known-issues.md`.
|
||||||
|
|
||||||
|
### Benchmarks
|
||||||
|
- Every undated or pre-September section of `BENCHMARKS.md` re-run on one
|
||||||
|
machine on one day (tank, 2026-09-24, commit 5c8323c), with the command for
|
||||||
|
each and every number traced back to the raw output by a separate check.
|
||||||
|
Where a figure moved, the section says so. Two apparent regressions were
|
||||||
|
isolated rather than published: knowledge-graph traversal (a real bug,
|
||||||
|
fixed above) and the write path, which measures the same at v2.3.0 on this
|
||||||
|
machine — the old 18 µs / 6.17 ms figures came from an undated run on other
|
||||||
|
hardware; `float16` adds ~2 µs per save and the int8 index nothing.
|
||||||
|
- New `multimodal_bench`: cross-modal search at 1K and 10K records, which the
|
||||||
|
README claimed but nothing measured.
|
||||||
|
- `footprint_bench` reports whether it built `float16` or `f32` stores and
|
||||||
|
takes `--f32`; it had kept printing "f32" after the default changed.
|
||||||
|
|
||||||
|
### Interop
|
||||||
|
- `clawhdf5-format`: **every `f32` dataset was unreadable by h5py and
|
||||||
|
libhdf5.** The float datatype encoder hard-coded the sign bit's position to
|
||||||
|
63, correct only for `f64`; libhdf5 validates it and refused the dataset. It
|
||||||
|
is now derived from the type (15 / 31 / 63). Our reader ignores the field,
|
||||||
|
and the interop suites only wrote `f64`, which is how it went unnoticed.
|
||||||
|
- `clawhdf5-format`: **every empty dataset was unreadable by h5py and
|
||||||
|
libhdf5.** It was written with a real address and zero bytes, which trips
|
||||||
|
libhdf5's `addr + size <= addr` overflow check. An empty contiguous dataset
|
||||||
|
now gets the undefined address, as libhdf5 writes it. This affected every
|
||||||
|
agent store without sessions or a knowledge graph.
|
||||||
|
- New interop tests: `f32` and `float16` datasets in both directions (our
|
||||||
|
`float16` rounding matches numpy's bit for bit on 4 020 probe values,
|
||||||
|
including ties, subnormals and the overflow boundary), and an agent store —
|
||||||
|
`f32` and `float16` — opened by h5py with every dataset decoded.
|
||||||
|
- `clawhdf5-format` filters, checked against libhdf5 + hdf5plugin:
|
||||||
|
- **LZ4 (32004) now uses the registered HDF5 LZ4 format** (8-byte BE size,
|
||||||
|
4-byte BE block size, BE-length-prefixed blocks). Our old framing (4-byte
|
||||||
|
LE size + one block) was readable only by clawhdf5, and we could not read
|
||||||
|
libhdf5's (`h5ex_d_lz4.h5`). Old clawhdf5 LZ4 chunks still read; they are
|
||||||
|
told apart unambiguously (a registered chunk starts with four zero bytes).
|
||||||
|
- **Zstd (32015) frames now record the content size**, which libhdf5's zstd
|
||||||
|
plugin needs; h5py could not read our zstd datasets.
|
||||||
|
- **Pcodec moved from filter ID 32023 to 480.** 32023 is registered to
|
||||||
|
Granular BitRound, whose decode is a pass-through — libhdf5 with that
|
||||||
|
plugin would have returned compressed bytes as data. Pcodec has no
|
||||||
|
registered ID; 480 is in the registry's private range (256–511) and only
|
||||||
|
clawhdf5 can read it. Chunks written under 32023 with the filter name
|
||||||
|
`pcodec` (clawhdf5 ≤ 2.7.0) still read.
|
||||||
|
- **SZIP decode matches libhdf5.** It returned garbage or zeros with no
|
||||||
|
error for libhdf5-written files (the 4-byte size prefix, 32/64-bit
|
||||||
|
byte-plane interleaving, reference interval, scanline padding and byte
|
||||||
|
order were all handled wrongly) and rejected 64-bit data.
|
||||||
|
- N-Bit honours libhdf5's "need not compress" flag (multi-filter pipelines
|
||||||
|
such as `tfilters.h5` failed) and reads enum/no-op members.
|
||||||
|
- Scale-offset `float` decode uses libhdf5's single-precision arithmetic
|
||||||
|
(was 1 ULP off for some values).
|
||||||
|
- A pipeline with Fletcher32 ahead of the compressor (h5py
|
||||||
|
`set_fletcher32()` then `set_deflate()`) no longer fails with "deflate:
|
||||||
|
output exceeds size limit".
|
||||||
|
|
||||||
|
### Storage
|
||||||
|
- `clawhdf5-format`: **half-precision datasets.**
|
||||||
|
`DatasetBuilder::with_f16_data` writes IEEE binary16 (numpy `float16`),
|
||||||
|
rounding to nearest-even; `make_f16_type`, and `clawhdf5_format::float16`
|
||||||
|
with the conversions, which are checked against the `half` crate on 16.7M
|
||||||
|
values and round-trip all 65 536 half values. Reading `float16` as `f32`
|
||||||
|
gained a little-endian fast path.
|
||||||
|
- `clawhdf5-agent`: **`MemoryConfig::float16` stores embeddings as half
|
||||||
|
precision.** At 100K x 384 the file goes from 154.0 to 80.8 MiB (−48%), a
|
||||||
|
checkpoint from 752 to 512 ms and open from 300 to 252 ms, with the same
|
||||||
|
vector recall@10 against an exact scan (0.999 vs 0.994) and the same
|
||||||
|
`hybrid_search` latency; at 10K open is 3 ms slower. On the full
|
||||||
|
LongMemEval haystack with real MiniLM embeddings every retrieval metric is
|
||||||
|
identical to `f32` (`longmemeval_bench --float16`). The cache rounds each
|
||||||
|
embedding as it is saved, so memory and file agree bit for bit and a store
|
||||||
|
returns the same results before and after a reopen (tested). Out-of-range
|
||||||
|
values are refused with `MemoryError::InvalidEntry` rather than stored as
|
||||||
|
infinity; batches are all or nothing. CLI: `create --float16`. See
|
||||||
|
`BENCHMARKS.md`, "float16 embedding storage".
|
||||||
|
|
||||||
|
### Build
|
||||||
|
- **Pure-Rust default.** `clawhdf5-format`, `clawhdf5-filters` and the
|
||||||
|
`clawhdf5` facade default to the `zlib-rs` deflate backend; `fast-deflate`
|
||||||
|
(zlib-ng) is opt-in. No crate in the default dependency tree of the core
|
||||||
|
crates compiles C, and `ci-test.sh` now fails if one appears. The facade's
|
||||||
|
`fast-deflate` was on by default and is now off. See `BENCHMARKS.md`,
|
||||||
|
"Deflate backend".
|
||||||
|
- `zlib-rs` also enables flate2's `runtime_detection`. Without it zlib-rs has
|
||||||
|
no `std`, cannot detect SIMD at runtime, and inflates 3.5x slower; the
|
||||||
|
workspace builds flate2 with `default-features = false`, which had been
|
||||||
|
switching it off.
|
||||||
|
- `rust-version = "1.92"` for the whole workspace (the floor: `wgpu` requires
|
||||||
|
it), and CI checks the workspace on exactly that toolchain.
|
||||||
|
- CI keeps zlib-ng building and tested; the arm64 job no longer needs cmake.
|
||||||
|
|
||||||
|
### Correctness
|
||||||
|
- `clawhdf5-format` reader — **values returned wrong with no error:**
|
||||||
|
- Fixed Array and Extensible Array chunk indexes were laid out by the
|
||||||
|
dataset's current shape instead of its max shape (23 libhdf5 test files,
|
||||||
|
and any h5py file with e.g. `maxshape=(10, None)` or `(20, 10)` under
|
||||||
|
`libver='latest'`).
|
||||||
|
- Files with 4-byte offsets: unfiltered chunked datasets read as zeros.
|
||||||
|
Chunk B-tree keys store offsets in 8 bytes whatever the file's offset
|
||||||
|
size.
|
||||||
|
- A chunk's filter mask skipped the whole pipeline when any bit was set;
|
||||||
|
only the flagged filters are skipped now.
|
||||||
|
- Float data read as an integer returned the bit pattern; narrowing integer
|
||||||
|
reads kept the low bits; bfloat16 was decoded as IEEE half. Floats are now
|
||||||
|
decoded from their datatype fields (bf16, FP8 E4M3/E5M2, IEEE half, single
|
||||||
|
and double).
|
||||||
|
- `vl_data::read_vl_bytes` truncated sequences of non-byte base types.
|
||||||
|
- A shared fill-value message read as zero fill; it is resolved now,
|
||||||
|
including from the file's shared-message (SOHM) table, which could never
|
||||||
|
resolve because its index version byte was skipped.
|
||||||
|
- Two threads reading two chunked datasets through one `File` could get each
|
||||||
|
other's chunks (the shared chunk cache was switched between datasets
|
||||||
|
across separate lock acquisitions). The cache is now keyed by dataset.
|
||||||
|
- `clawhdf5-format` reader — errors on valid files: enum and bool datasets
|
||||||
|
through the numeric readers; the "don't filter partial edge chunks" layout
|
||||||
|
flag; Fletcher32 ahead of deflate (NetCDF-4's order). Unknown-message flags
|
||||||
|
follow libhdf5 (`tbogus.h5`): "fail if unknown" is refused, "fail if unknown
|
||||||
|
and writing" is ignored by a reader.
|
||||||
|
- `clawhdf5-format` writer — **files libhdf5 rejects or reads wrong:**
|
||||||
|
- Extensible Array (one unlimited dimension): chunks from index 244 on were
|
||||||
|
written but never indexed and read as 0, by libhdf5 and by us.
|
||||||
|
- Fixed Array: more than 1 024 chunks gave checksum errors (data blocks
|
||||||
|
were never paged).
|
||||||
|
- A finite max shape larger than the shape gave libhdf5 "addr overflow"; an
|
||||||
|
unlimited dimension that is not the first scrambled the data; several
|
||||||
|
unlimited dimensions (`(None, None)`) broke the whole file. These now
|
||||||
|
write the index libhdf5 writes (swizzled Extensible Array, or a B-tree v2
|
||||||
|
index for several unlimited dimensions).
|
||||||
|
- Header messages over 64 KiB (the size field is 16 bits) and compact
|
||||||
|
datasets at 65 534–65 535 bytes produced corrupt files.
|
||||||
|
- Reference, Opaque, BitField and Time datatypes were written as empty
|
||||||
|
messages; they now encode as HDF5 2.0 does.
|
||||||
|
- `with_page_size` wrote a nonexistent superblock version 4; it now writes
|
||||||
|
the v3 superblock and File Space Info message libhdf5 writes.
|
||||||
|
- `FillTime` values were rotated on disk (NEVER was written as ALLOC, and so
|
||||||
|
on). New `DatasetBuilder::with_fill_value`.
|
||||||
|
- An empty-string attribute got a zero-size datatype, which made every
|
||||||
|
attribute on the object unreadable in libhdf5.
|
||||||
|
- `maxshape` equal to the shape no longer forces chunked layout.
|
||||||
|
- `clawhdf5-format`: **a truncated deflate chunk read back short, with no
|
||||||
|
error.** The deflate filter used flate2's streaming reader, which returns the
|
||||||
|
bytes it has when the input runs out before the end-of-stream marker. It now
|
||||||
|
decodes in one pass into a buffer sized to the chunk and reports a
|
||||||
|
truncated stream as `DecompressionError`. Same fix in `clawhdf5-filters`,
|
||||||
|
where output longer than the stated size was also silently cut off; it is
|
||||||
|
now an error.
|
||||||
|
|
||||||
|
### Defaults
|
||||||
|
- `clawhdf5-agent`: `MemoryConfig::float16` defaults to `true` for new stores,
|
||||||
|
measured rather than assumed: identical LongMemEval retrieval on real
|
||||||
|
embeddings, 48% smaller files and faster checkpoints and opens at 100K.
|
||||||
|
`clawhdf5-cli create --f32` opts out; like `--f32-index`, it only ever
|
||||||
|
switches the default off.
|
||||||
|
- `clawhdf5-agent`: `MemoryConfig::quantized_index` defaults to `true` for new
|
||||||
|
stores. The reason it had been off — that int8 search was slower on ARM —
|
||||||
|
did not survive measurement (see Corrections). Stores that predate the
|
||||||
|
setting still load it as `false`, so reopening one never changes how its
|
||||||
|
index is held; a store written by the v2.5.0 CLI is now a test fixture that
|
||||||
|
guards exactly that, and the test fails if the load default is changed.
|
||||||
|
- `clawhdf5-cli`: `create --f32-index` opts out. `create` used to assign
|
||||||
|
`--quantized-index` straight into the config, which under the new default
|
||||||
|
would have forced every CLI-created store back to f32 unless the caller
|
||||||
|
knew to ask; it now only ever switches the default off.
|
||||||
|
|
||||||
|
### Performance
|
||||||
|
- `clawhdf5-agent`: consolidation's novelty scoring (each `add_memory` against
|
||||||
|
the whole working tier) computes the new record's norm once, takes each
|
||||||
|
comparison in one vectorised pass instead of three, and splits a working
|
||||||
|
tier of 4 096+ records across threads — same results, tested against the
|
||||||
|
old formula. It had made `consolidation_efficiency` stall at 100K; the
|
||||||
|
complete run now takes 8 min and fills in the 100K cycle row (46.66 ms) and
|
||||||
|
the memory-reduction table.
|
||||||
|
- `clawhdf5-bench`: `consolidation_efficiency` no longer prints a record-count
|
||||||
|
ratio as a "BM25 Speedup" (it was never measured), nor claims cycle time
|
||||||
|
grows sub-linearly (its own numbers grow slightly faster than linearly).
|
||||||
|
- `clawhdf5-agent`: **knowledge-graph traversal was 6.5x slower than it
|
||||||
|
should be.** `bfs_neighbors` and `spreading_activation` built an adjacency
|
||||||
|
index over the whole graph on every call (1efd82c), so a 2-hop BFS over 1K
|
||||||
|
entities took 155 µs. The index is now cached on `KnowledgeCache` and
|
||||||
|
checked against a fingerprint of the graph on each use — one pass over
|
||||||
|
entity ids and relation endpoints, no allocation — so any change, including
|
||||||
|
direct edits of its public `Vec`s, still rebuilds it (tested). BFS over 1K
|
||||||
|
entities: 155.1 -> 23.1 µs; spreading activation over 100: 22.8 -> 10.1 µs.
|
||||||
|
- `clawhdf5-format`, `clawhdf5-filters`: both deflate paths hand the codec the
|
||||||
|
whole chunk in one call, into a buffer allocated once, instead of streaming
|
||||||
|
it through a 32 KiB buffer: about 5% on chunked writes and 10% on zlib-ng's
|
||||||
|
1 MB inflate.
|
||||||
|
- `clawhdf5-accel`: **`dot_i8` has aarch64 kernels** — `SDOT` for CPUs with
|
||||||
|
the ARMv8.2 dot-product extension (Cortex-A76 and later, Neoverse-N1, every
|
||||||
|
Apple Silicon generation) and plain NEON (`vmull_s8` + `vpadalq_s16`) for
|
||||||
|
the rest, selected at runtime. `SDOT` is issued through inline assembly,
|
||||||
|
because the `vdotq_s32` intrinsic is still behind the unstable
|
||||||
|
`stdarch_neon_dotprod` feature. On a Raspberry Pi 5 at N = 100 000 and
|
||||||
|
equal recall, the quantised index answers **1.18x the queries per second**
|
||||||
|
of f32 (7 267 vs 6 164) and builds **2.3x faster** (14 464 vs 33 413 ms).
|
||||||
|
Both kernels are tested bit-for-bit against scalar on real hardware, each
|
||||||
|
explicitly — dispatch only ever takes one path on a given CPU, so testing
|
||||||
|
through it alone would have left the plain-NEON fallback unexercised on any
|
||||||
|
machine with `SDOT`.
|
||||||
|
|
||||||
|
### Corrections
|
||||||
|
- The v2.7.0 entry for `dot_i8` said `quantized_index` stayed off by default
|
||||||
|
because "aarch64 falls back to the scalar loop", implying the ~13% search
|
||||||
|
penalty measured on x86 applied on ARM too. It did not. That figure came
|
||||||
|
from scalar int8 against hand-written AVX2 f32 kernels on x86, whose
|
||||||
|
portable baseline is SSE2; on aarch64 NEON is the baseline, and measured on
|
||||||
|
a Pi 5 the scalar int8 loop already matched f32 for search while building
|
||||||
|
1.76x faster. The claim was extrapolated rather than measured.
|
||||||
|
|
||||||
|
## v2.7.0 (2026-09-20)
|
||||||
|
|
||||||
|
### Upgrade Notes
|
||||||
|
- **Two read-path bugs fixed, one of them silent.** Datasets indexed by an
|
||||||
|
Extensible Array (any dataset with one unlimited dimension) returned data
|
||||||
|
from the wrong chunks past their first few dozen. If you have readings taken
|
||||||
|
from such a dataset with an earlier release, they may be wrong; re-read them.
|
||||||
|
- **A corrupt chunk index is now an error.** Fixed and Extensible Array
|
||||||
|
structures carry checksums that were previously ignored, so damage surfaced
|
||||||
|
as plausible data from the wrong offset. Code that read a damaged file and
|
||||||
|
got numbers will now get `ChecksumMismatch` instead. That is the point.
|
||||||
|
- **Breaking:** `MemoryConfig` gained `hnsw_m`, `hnsw_ef_construction` and
|
||||||
|
`hnsw_ef_search`, so literal constructions need updating;
|
||||||
|
`..Default::default()` does not. All three default to the previous
|
||||||
|
behaviour.
|
||||||
|
|
||||||
|
### Correctness
|
||||||
|
- `clawhdf5-format`: **datasets indexed by an Extensible Array returned wrong
|
||||||
|
data beyond their first few dozen chunks.** One unlimited dimension gives a
|
||||||
|
dataset an Extensible Array chunk index, whose first elements (4 by default)
|
||||||
|
sit inline in the index block and whose rest live in data blocks sized by a
|
||||||
|
formula the reader got wrong. In the default layout everything through the
|
||||||
|
36th chunk happened to line up and the 37th onwards did not: a 400-chunk
|
||||||
|
dataset silently returned wrong values from chunk 37, and datasets past
|
||||||
|
about a thousand chunks failed outright with "invalid Extensible Array data
|
||||||
|
block signature". **Reads were wrong, not
|
||||||
|
merely refused** — the caller got plausible numbers from the wrong chunks.
|
||||||
|
Four separate layout errors, each checked against files written by HDF5 2.0
|
||||||
|
and against the library source:
|
||||||
|
- the number of data blocks in super block `u` is `2^(u/2)`, not `2^u`;
|
||||||
|
- each holds `2^((u+1)/2) * data_blk_min_elmts` elements, which doubles
|
||||||
|
every *other* level rather than every level;
|
||||||
|
- a super block carries a block-offset field before its data block
|
||||||
|
addresses, which was not skipped;
|
||||||
|
- the page-init bitmap belongs to the super block, one bit per page packed
|
||||||
|
across all its data blocks (MSB first), and was being read from inside the
|
||||||
|
data block instead; a paged data block also ends its prefix with a
|
||||||
|
checksum before the first page.
|
||||||
|
Covered now by interop tests at 4, 37, 400, 5 000 and 200 000 chunks (the
|
||||||
|
last large enough for paged data blocks), plus sparse, gzip-filtered and
|
||||||
|
2-D cases. Writing is unaffected; this is a read-path bug.
|
||||||
|
- `clawhdf5-format`: the sibling Fixed Array index (fixed dimensions written
|
||||||
|
with `libver='latest'`) was checked against the same range and is correct,
|
||||||
|
including paged data blocks and sparse datasets — it really does keep its
|
||||||
|
page-init bitmap in the data block, where the Extensible Array does not.
|
||||||
|
It had no real-file coverage above the inline sizes either, so it now has
|
||||||
|
the same tests.
|
||||||
|
|
||||||
|
### Security
|
||||||
|
- `clawhdf5-format`: **a crafted file could crash any reader through B-tree v2
|
||||||
|
traversal.** Recursion was bounded only by the depth the file claimed (a
|
||||||
|
`u16`), and child addresses were never checked for sharing. A node listing
|
||||||
|
itself as its own child under a header claiming 65 535 levels — under 100
|
||||||
|
bytes — overflowed the stack and **aborted the process** (SIGABRT, not a
|
||||||
|
catchable error). Levels whose children all point at one shared node below
|
||||||
|
reached it fan-out^depth times: 29.5 million records from ~5 KB, and one
|
||||||
|
more level would exhaust memory. Both are now errors, returned in under a
|
||||||
|
millisecond: depth is capped at 64 (as the fractal heap already was), and
|
||||||
|
traversal stops once it has produced more records than the file has bytes
|
||||||
|
to hold. Every B-tree v2 user goes through this path — dense attributes,
|
||||||
|
v2 groups, shared messages and chunk indexes. Valid files are unaffected,
|
||||||
|
including a depth-2 HDF5 2.0 chunk index with 40 000 records, now covered by
|
||||||
|
an interop test.
|
||||||
|
|
||||||
|
### Integrity
|
||||||
|
- `clawhdf5-format`: **Fixed and Extensible Array chunk indexes now verify
|
||||||
|
their checksums** (the `checksum` feature, on by default). Every structure
|
||||||
|
in both — header, index block, super block, data block and each data block
|
||||||
|
page — carries a Jenkins lookup3 checksum that was parsed past and ignored.
|
||||||
|
The consequence of skipping it is not a missing warning but wrong data: a
|
||||||
|
single flipped bit in a chunk address still parses, still points inside the
|
||||||
|
file, and the reader hands back whatever bytes now sit there as the chunk's
|
||||||
|
contents. Verified in both directions — the checksums accept files written
|
||||||
|
by HDF5 2.0 at 100 to 200 000 chunks, dense, sparse, filtered and paged,
|
||||||
|
and an interop test corrupts an address to confirm the read now fails
|
||||||
|
instead of returning data (it does return data when the check is removed).
|
||||||
|
|
||||||
|
### Performance
|
||||||
|
- `clawhdf5-agent`: **opening a store is ~28% faster** (455 ms -> 327 ms at
|
||||||
|
100k x 384). `read_from_disk` memory-mapped the file and then copied the
|
||||||
|
entire mapping into a `Vec` for `File::from_bytes`, when `File::open`
|
||||||
|
memory-maps it directly — so every open paid a full-file memcpy for nothing.
|
||||||
|
Process peak memory is unchanged: the peak falls after the parse, during the
|
||||||
|
index build, so the transient never reached the high-water mark. The
|
||||||
|
footprint harness now reports that peak next to the retained figure, which
|
||||||
|
is how this was checked rather than assumed.
|
||||||
|
- `clawhdf5-accel`: **`dot_i8`, a runtime-dispatched int8 dot product** (AVX2:
|
||||||
|
sign-extend each half to `i16`, then `madd_epi16`; scalar fallback
|
||||||
|
elsewhere). The quantised HNSW index used a scalar loop while the `f32` path
|
||||||
|
it was measured against ran AVX2, so the ~13% throughput cost recorded for
|
||||||
|
`MemoryConfig::quantized_index` was a missing kernel rather than a property
|
||||||
|
of int8. With the kernel, at N = 100 000 x 384 and equal recall, the
|
||||||
|
quantised index answers **1.63x as many queries per second** (21 848 vs
|
||||||
|
13 399 at ef=64, recall 0.9940 vs 0.9945) and builds **1.8x faster** (1778
|
||||||
|
vs 3197 ms) — on top of holding a quarter of the vectors. Medians of three
|
||||||
|
alternating runs. It remains off by default only because the kernel is
|
||||||
|
AVX2-only and aarch64 falls back to the scalar loop. Integer arithmetic, so
|
||||||
|
the SIMD path is tested to agree with scalar bit for bit.
|
||||||
|
|
||||||
|
### Tuning
|
||||||
|
- `clawhdf5-agent`: **the HNSW parameters are configurable** —
|
||||||
|
`MemoryConfig::hnsw_m`, `hnsw_ef_construction` and `hnsw_ef_search`
|
||||||
|
(defaults 16, 64, and 0 meaning "scale with `k`", i.e. today's behaviour).
|
||||||
|
They were constants, so a deployment could not trade recall against memory
|
||||||
|
or query speed at all. All three are persisted with the store. Values are
|
||||||
|
clamped where the index requires it: `clawhdf5-ann` asserts a graph degree
|
||||||
|
of at least 2, so a configured 0 — from a file, or from a caller who took 0
|
||||||
|
to mean "default" — used to abort the process inside the builder. Lowering
|
||||||
|
`ef_search` also no longer narrows the candidate pool that fusion sees.
|
||||||
|
**Breaking:** `MemoryConfig` gained fields, so literal constructions need
|
||||||
|
updating; `..Default::default()` does not.
|
||||||
|
|
||||||
|
### Documentation
|
||||||
|
- `clawhdf5-agent`: `BM25Index::search` claimed to use Block-Max WAND for early
|
||||||
|
termination. It never did; it scores every match exhaustively. It now says
|
||||||
|
so, and why no pruning would help the store: `hybrid_search` uses `scores()`,
|
||||||
|
since fusion normalises over every match.
|
||||||
|
|
||||||
|
## v2.6.0 (2026-09-20)
|
||||||
|
|
||||||
|
### Upgrade Notes
|
||||||
|
- **Re-ranked results change, substantially for the better.** `RerankInput`
|
||||||
|
and `ReRankConfig` gained fields (`relevance`, `relevance_weight`), so
|
||||||
|
literal constructions need updating; `..Default::default()` does not. Any
|
||||||
|
caller that re-ranked was previously getting results ordered by age with the
|
||||||
|
retrieval score discarded — see below.
|
||||||
|
- **Breaking:** `MemoryCache::embeddings` is a `cache::Embeddings` rather than
|
||||||
|
a `Vec<Vec<f32>>` (indexing still yields a `&[f32]` row); `embeddings_flat`
|
||||||
|
is gone, replaced by `flat_embeddings()`; `rebuild_flat()` is a deprecated
|
||||||
|
no-op.
|
||||||
|
- `MemoryConfig` gained `quantized_index` (default `false`, so behaviour is
|
||||||
|
unchanged unless you opt in); literal constructions need the field.
|
||||||
|
|
||||||
|
### Retrieval quality
|
||||||
|
- `clawhdf5-agent`: **re-ranking discarded the retrieval score.**
|
||||||
|
`reranker::rerank` built its combined score from temporal decay, source
|
||||||
|
authority and Hebbian activation only — `RerankInput` had no relevance field
|
||||||
|
— so re-ranking a candidate pool reordered it by age and threw the
|
||||||
|
retriever's ordering away. The OpenClaw backend re-ranked every search, so
|
||||||
|
this was its shipping behaviour: measured over the full LongMemEval haystack
|
||||||
|
it cost **40.6pp of Hit@1** (11.0% vs 51.6%) and two thirds of MRR (0.183 vs
|
||||||
|
0.643). `RerankInput::relevance` and `ReRankConfig::relevance_weight` (1.0 by
|
||||||
|
default) fix it: relevance leads and the metadata signals break near-ties,
|
||||||
|
which restores retrieval (Hit@1 +0.4pp vs no re-ranking) and improves
|
||||||
|
recency discrimination by 6–7pp. **Breaking:** `RerankInput` and
|
||||||
|
`ReRankConfig` gained fields, so literal constructions need updating;
|
||||||
|
`..Default::default()` does not.
|
||||||
|
- `clawhdf5-bench`: the LongMemEval harness feeds the dataset's real session
|
||||||
|
dates to the store instead of a synthetic counter (decay needs true
|
||||||
|
intervals, not just the right order), and reports `newest_gold_first` — on a
|
||||||
|
`knowledge-update` question, did the newest gold session outrank the stale
|
||||||
|
one it supersedes? Plain recall cannot see this, because both are labelled
|
||||||
|
gold. New `--rerank-sweep`.
|
||||||
|
|
||||||
|
### Memory
|
||||||
|
- `clawhdf5-agent`: **`MemoryConfig::quantized_index`** stores the vector
|
||||||
|
index's own copy of the embeddings as `i8` rather than `f32`, which at 100k
|
||||||
|
384-dim entries takes the index from 266 to 123 MiB and the whole reopened
|
||||||
|
store from 399 to 256 MiB (2.72x -> **1.74x** the raw vectors). Quantised
|
||||||
|
distances are approximate and `ef` cannot compensate — recall@10 tops out at
|
||||||
|
0.967 against f32's 0.9995 — so the query path re-scores the candidate pool
|
||||||
|
against the exact embeddings the store already holds, which restores recall
|
||||||
|
(0.9940 vs 0.9945 at ef=64) for about 13% of QPS. **Off by default**: it
|
||||||
|
trades query speed for memory, and which side is worth more depends on the
|
||||||
|
deployment. The setting is persisted, so a reopened store does not silently
|
||||||
|
revert to four times the index memory.
|
||||||
|
- `clawhdf5-ann`: `Storage::Int8` and the `build_with` / `new_with` /
|
||||||
|
`from_graph_bytes_with` constructors that select it. The scale is per row,
|
||||||
|
not global — a fixed `[-1, 1]` scale spends fewer than 12 of the 255 levels
|
||||||
|
on a unit-length 128-dim vector and is unusable (0.35 top-10 overlap against
|
||||||
|
an exact ranking, versus 0.99 per row). `compact()` keeps the storage it was
|
||||||
|
given; serialized indexes still carry f32 vectors, so a quantised index is
|
||||||
|
rebuilt rather than loaded.
|
||||||
|
- `clawhdf5-agent`: **a loaded store holds ~30% less memory** (100k 384-dim
|
||||||
|
entries: 505 -> 357 MiB, 3.44x -> 2.43x the raw vectors). The cache kept
|
||||||
|
every embedding twice — a `Vec<Vec<f32>>` and a flattened copy for the
|
||||||
|
batched kernels, maintained in lock-step — so it now stores only the flat
|
||||||
|
buffer and indexes into it. Recall and query latency are unchanged.
|
||||||
|
**Breaking:** `MemoryCache::embeddings` is a `cache::Embeddings` rather than
|
||||||
|
a `Vec<Vec<f32>>` (indexing still yields a `&[f32]` row); `embeddings_flat`
|
||||||
|
is gone, replaced by `flat_embeddings()`; `rebuild_flat()` is a deprecated
|
||||||
|
no-op. Rows are now always exactly `dim` long — shorter ones are
|
||||||
|
zero-padded — which makes the ragged-row case that used to silently
|
||||||
|
misalign the flattened copy unrepresentable.
|
||||||
|
- `clawhdf5-bench`: `search_harness --footprint` reports live heap use per
|
||||||
|
stage, measured with a counting allocator (RSS cannot see a structure freed
|
||||||
|
into the allocator's own pool).
|
||||||
|
|
||||||
|
### Testing
|
||||||
|
- The Python interop suites honour **`CLAWHDF5_PYTHON`**, and `ci-test.sh`
|
||||||
|
picks up a `.venv/bin/python` automatically. On a PEP 668 "externally
|
||||||
|
managed" system h5py cannot be installed into the system interpreter at all,
|
||||||
|
so every interop suite — the h5py writer round-trips, the facade, netCDF4
|
||||||
|
and the reference files — was skipping silently. A silent skip here is
|
||||||
|
exactly how the v5 compound-datatype bug reached a release.
|
||||||
|
`CLAWHDF5_REQUIRE_INTEROP=1` still turns a skip into a failure.
|
||||||
|
|
||||||
|
## v2.5.0 (2026-09-19)
|
||||||
|
|
||||||
|
### Upgrade Notes
|
||||||
|
- **Retrieval rankings change, for the better.** The default fusion weights
|
||||||
|
move from `0.7/0.3` to `0.4/0.6` (`hybrid::DEFAULT_FUSION`), measured over the
|
||||||
|
full LongMemEval haystack: turn-level Hit@1 51.6% vs 44.2%, MRR 0.643 vs
|
||||||
|
0.586. `unified_search` and the OpenClaw backend pick this up automatically;
|
||||||
|
callers passing weights to `hybrid_search` explicitly are unaffected.
|
||||||
|
- **Out-of-range selections are now errors.** `read_*_selection` used to return
|
||||||
|
data for a selection that ran past a dataset edge — a hyperslab came back
|
||||||
|
zero-padded, and a point with an out-of-range coordinate wrapped into the
|
||||||
|
next row. Both are now `FormatError::SelectionOutOfBounds`. Code relying on
|
||||||
|
the old (wrong) values will start seeing errors.
|
||||||
|
- **Large compressed datasets written without explicit chunk dimensions get a
|
||||||
|
different layout.** They used to be stored as one chunk; they are now split
|
||||||
|
to ~1 MiB chunks. The files stay standard and h5py-readable, and explicit
|
||||||
|
`with_chunks` is unaffected.
|
||||||
|
- `rayon` is now a default dependency of `clawhdf5-agent` (the parallel index
|
||||||
|
build). Opt out with `--no-default-features --features float16,hnsw`.
|
||||||
|
- `clawhdf5-ann` search results no longer shrink when records near the query
|
||||||
|
have been deleted, so a search that previously returned fewer than `k`
|
||||||
|
results now returns `k`.
|
||||||
|
|
||||||
|
### Retrieval quality
|
||||||
|
- `clawhdf5-agent`: optional keyword stemming — `bm25::TokenFilter::Stemmed`
|
||||||
|
and `HDF5Memory::set_token_filter`, so "training" and "trains" match. **Off
|
||||||
|
by default**, on measurement rather than principle: over the full LongMemEval
|
||||||
|
haystack it buys depth and costs the top rank (BM25 alone: Hit@5 +2.8pp,
|
||||||
|
Hit@10 +2.4pp, Hit@1 −1.8pp, MRR unchanged), and on the shipping hybrid
|
||||||
|
configuration the trade is narrower still. See `BENCHMARKS.md`.
|
||||||
|
- `clawhdf5-agent`: **`QueryExpander::expand` panicked on ordinary non-ASCII
|
||||||
|
input** — `"İ AI"` was enough. It searched a lowercased copy of the query and
|
||||||
|
then sliced the *original* with those offsets, which only works while
|
||||||
|
lowercasing preserves byte length (Turkish `İ` is 2 bytes and lowercases to
|
||||||
|
3). Depending on where the offsets drifted it either corrupted the output
|
||||||
|
("İstanbul AI trip" lost a character) or panicked. Matching now walks the
|
||||||
|
original string.
|
||||||
|
- `clawhdf5-agent`: query expansion no longer rewrites text inside words.
|
||||||
|
`replace_word_case_insensitive` did a plain substring replace despite its
|
||||||
|
name, so "training" became "trArtificial Intelligencening" and "programming"
|
||||||
|
became "Pull Requestogramming" — every acronym expansion of ordinary prose
|
||||||
|
was corrupt. Matches now require word boundaries; genuine acronyms
|
||||||
|
(`API`, `database`) still expand.
|
||||||
|
- `clawhdf5-agent`: **the default fusion weights are now the measured ones.**
|
||||||
|
A sweep of every 0.1 step over the full LongMemEval haystack (500 questions,
|
||||||
|
real MiniLM embeddings) shows the long-standing `0.7/0.3` default is
|
||||||
|
*strictly dominated* by `0.4/0.6` — turn-level Hit@1 51.6% vs 44.2%, Hit@5
|
||||||
|
81.4% vs 79.2%, Hit@10 87.8% vs 85.8%, MRR 0.643 vs 0.586, and better at
|
||||||
|
session level too. The finding was recorded in `BENCHMARKS.md` but had never
|
||||||
|
been applied: `unified_search` and the OpenClaw backend both hardcoded
|
||||||
|
`0.7/0.3`. They now use `hybrid::DEFAULT_FUSION`. **Callers passing weights
|
||||||
|
to `hybrid_search` explicitly are unaffected** — pass `0.4`/`0.6` (or use
|
||||||
|
`hybrid_search_with`) to get the tuned behaviour.
|
||||||
|
- `clawhdf5-agent`: fusion is now selectable. New `hybrid::Fusion`
|
||||||
|
(`Weighted { vector, keyword }` or `Rrf { k }`), `hybrid::fuse`,
|
||||||
|
`hybrid::hybrid_search_fused` and `HDF5Memory::hybrid_search_with`.
|
||||||
|
Reciprocal rank fusion existed but was unreachable from the store, so it had
|
||||||
|
never been measured against the weighted sum; the LongMemEval bench now has
|
||||||
|
an `RRF` mode.
|
||||||
|
|
||||||
|
### HDF5 Read Path
|
||||||
|
- **Selection reads cost what the selection costs.** `read_*_selection` decoded
|
||||||
|
the *entire* dataset and then picked elements out, so a 64 x 64 window of a
|
||||||
|
64 MB compressed dataset took 105 ms - about as long as reading all of it.
|
||||||
|
Now only the rows (contiguous) or chunks that overlap the selection's
|
||||||
|
bounding box are read and decompressed: that window takes 0.39 ms, one row
|
||||||
|
2.7 ms, one column 5.2 ms. Results are identical to the full-read path
|
||||||
|
(equivalence-tested over random hyperslabs and point lists, ranks 1-3,
|
||||||
|
contiguous / chunked / deflate). New `read_harness` bench binary.
|
||||||
|
- **Faster full reads** (same-moment A/B, 64 MB `f64`): chunked + deflate
|
||||||
|
110 -> 69 ms, chunked 72 -> 60 ms, contiguous 56 -> 30 ms. The facade's
|
||||||
|
cached read path now decompresses cache misses in parallel batches (it was
|
||||||
|
sequential; only the uncached reader was parallel) and caches only datasets
|
||||||
|
that fit the chunk cache; unfiltered chunks are copied straight from the file
|
||||||
|
bytes; a contiguous dataset is converted straight from the file bytes; and
|
||||||
|
the native-endian conversions no longer zero a buffer before overwriting it.
|
||||||
|
- **Datasets indexed by a version-2 B-tree now read** (layout v4, chunk index
|
||||||
|
type 5 — what `libver='latest'` uses for two or more unlimited dimensions;
|
||||||
|
previously "unsupported chunked layout"). The four copies of the chunk-index
|
||||||
|
dispatch are now one shared function, so every read path gets it.
|
||||||
|
- **`H5T_STD_REF` references** (HDF5 1.12+, datatype message version 4) parse:
|
||||||
|
`ReferenceType` gains `Object2`, `DatasetRegion2` and `Attribute`, and
|
||||||
|
`read_object_references` decodes the new object references. Previously any
|
||||||
|
dataset of this type failed with `InvalidReferenceType(2)`. Tested against a
|
||||||
|
file written by HDF5 2.0 itself (fixture + generator script committed).
|
||||||
|
- **Automatic chunk sizes.** Asking for compression (or any filter) without
|
||||||
|
`with_chunks` used to store the whole dataset as one chunk, so any read had
|
||||||
|
to decompress everything and nothing could be decoded in parallel. Datasets up
|
||||||
|
to 1 MiB stay a single chunk, as before; larger ones are split by halving the
|
||||||
|
dimensions in turn until a chunk is at most 1 MiB (the approach h5py takes).
|
||||||
|
**Behaviour change:** large compressed datasets written without explicit
|
||||||
|
chunk dimensions get a different (standard, h5py-readable) layout. Explicit
|
||||||
|
`with_chunks` is unaffected.
|
||||||
|
- **Out-of-range selections are errors.** They used to return data: a hyperslab
|
||||||
|
past an edge came back padded with zeros, and a point whose column was out of
|
||||||
|
range wrapped into the next row and returned that element. Now
|
||||||
|
`FormatError::SelectionOutOfBounds` (also for a rank mismatch or overlapping
|
||||||
|
blocks).
|
||||||
|
|
||||||
|
### Search
|
||||||
|
- `clawhdf5-ann`: **faster index builds.** Back-link pruning is 90% of a
|
||||||
|
build's distance evaluations; the bulk build now inserts in batches and
|
||||||
|
prunes each overflowing neighbour list once per batch (10K: 1676 -> 1074 ms).
|
||||||
|
With the `parallel` feature, planning and pruning run on a thread pool (10K:
|
||||||
|
388 ms, 100K: ~21 s -> 5.9 s on 16 cores). The graph is deterministic and
|
||||||
|
identical with or without the feature. `clawhdf5-agent`'s `parallel` feature
|
||||||
|
enables it for the agent's index and is now **on by default** (adds `rayon`
|
||||||
|
to the default dependency set; build with `--no-default-features --features
|
||||||
|
float16,hnsw` to opt out).
|
||||||
|
- `clawhdf5-ann`: `HnswIndex::search` returned fewer than `k` results — often
|
||||||
|
none — when the records nearest the query had been deleted: it collected `ef`
|
||||||
|
candidates, *then* dropped the deleted ones, *then* took `k`. Deleted nodes
|
||||||
|
are now traversed as waypoints but never occupy a result slot, so a search
|
||||||
|
returns the `k` nearest live records. Matters for any store that deletes or
|
||||||
|
supersedes memories without compacting straight away.
|
||||||
|
|
||||||
## v2.4.0 (2026-09-19)
|
## v2.4.0 (2026-09-19)
|
||||||
|
|
||||||
### Upgrade Notes
|
### Upgrade Notes
|
||||||
|
|||||||
@@ -1,7 +1,7 @@
|
|||||||
# clawhdf5
|
# clawhdf5
|
||||||
|
|
||||||
## Purpose
|
## Purpose
|
||||||
Pure-Rust HDF5 format implementation with HNSW vector search, WAL-backed persistence, agent memory storage, and GPU-accelerated I/O. Used by ZeroClaw as its persistent memory and knowledge graph backend.
|
Pure-Rust HDF5 format implementation with HNSW vector search, WAL-backed persistence, agent memory storage, and GPU-accelerated vector search. A standalone library. Its one verified consumer is ClawBrainHub (`.brain` files); no agent framework integrates it (OpenClaw and ZeroClaw claims were withdrawn on 2026-09-25 — neither was ever true).
|
||||||
|
|
||||||
## Architecture
|
## Architecture
|
||||||
|
|
||||||
@@ -11,15 +11,15 @@ Cargo workspace with 16 crates under `crates/` (plus `libaec-sys`, an internal F
|
|||||||
|-------|------|
|
|-------|------|
|
||||||
| `clawhdf5-format` | HDF5 binary spec parser (superblock, B-tree, heap) — also holds shared type definitions and physical constants |
|
| `clawhdf5-format` | HDF5 binary spec parser (superblock, B-tree, heap) — also holds shared type definitions and physical constants |
|
||||||
| `clawhdf5-io` | Read/write implementation |
|
| `clawhdf5-io` | Read/write implementation |
|
||||||
| `clawhdf5-filters` | Compression filters (gzip, LZ4, Zstd, Blosc) |
|
| `clawhdf5-filters` | Deflate backends (zlib-rs, zlib-ng, Apple Compression); the HDF5 filter pipeline and the other codecs (LZ4, Zstd, SZIP, N-Bit, scale-offset, pcodec) live in `clawhdf5-format`. No Blosc. |
|
||||||
| `clawhdf5-derive` | Proc-macro derive for HDF5-serializable structs |
|
| `clawhdf5-derive` | Proc-macro derive for HDF5-serializable structs |
|
||||||
| `clawhdf5` | Main facade crate |
|
| `clawhdf5` | Main facade crate |
|
||||||
| `clawhdf5-netcdf4` | NetCDF-4 compatibility layer |
|
| `clawhdf5-netcdf4` | NetCDF-4 compatibility layer |
|
||||||
| `clawhdf5-ann` | HNSW approximate nearest-neighbor vector index |
|
| `clawhdf5-ann` | HNSW approximate nearest-neighbor vector index |
|
||||||
| `clawhdf5-agent` | Agent memory, session history, knowledge graph storage |
|
| `clawhdf5-agent` | Agent memory, session history, knowledge graph storage |
|
||||||
| `clawhdf5-gpu` | GPU-accelerated I/O via wgpu (hand-written WGSL compute shaders) |
|
| `clawhdf5-gpu` | GPU vector distance computation via wgpu (hand-written WGSL compute shaders) — not dataset I/O |
|
||||||
| `clawhdf5-accel` | CPU SIMD acceleration path |
|
| `clawhdf5-accel` | CPU SIMD acceleration path |
|
||||||
| `clawhdf5-migrate` | Schema migration engine |
|
| `clawhdf5-migrate` | SQLite → HDF5 agent-memory migration |
|
||||||
| `clawhdf5-android` | Android JNI bindings |
|
| `clawhdf5-android` | Android JNI bindings |
|
||||||
| `clawhdf5-cli` | Command-line interface |
|
| `clawhdf5-cli` | Command-line interface |
|
||||||
| `clawhdf5-napi` | Node.js native addon bindings |
|
| `clawhdf5-napi` | Node.js native addon bindings |
|
||||||
@@ -27,17 +27,37 @@ Cargo workspace with 16 crates under `crates/` (plus `libaec-sys`, an internal F
|
|||||||
| `clawhdf5-bench` | Benchmark suite |
|
| `clawhdf5-bench` | Benchmark suite |
|
||||||
|
|
||||||
## Key Features
|
## Key Features
|
||||||
- Zero-dependency HDF5 read/write (no libhdf5 C library required)
|
- Zero-C-dependency HDF5 read/write: no libhdf5, and deflate defaults to
|
||||||
|
pure-Rust zlib-rs (`fast-deflate` opts into zlib-ng, which needs cmake).
|
||||||
|
`ci-test.sh` fails if a C-building crate enters the core crates' default
|
||||||
|
tree. flate2 must keep `runtime_detection` with zlib-rs — without it zlib-rs
|
||||||
|
loses SIMD and inflates 3.5x slower. MSRV is 1.92 (`rust-version`, checked
|
||||||
|
in CI).
|
||||||
- HNSW vector index for semantic similarity search over agent memories — the
|
- HNSW vector index for semantic similarity search over agent memories — the
|
||||||
`clawhdf5-agent` `hnsw` feature is **on by default**, so `hybrid_search` uses
|
`clawhdf5-agent` `hnsw` feature is **on by default**, so `hybrid_search` uses
|
||||||
the approximate `clawhdf5-ann` index for the vector stage (the index mirrors
|
the approximate `clawhdf5-ann` index for the vector stage (the index mirrors
|
||||||
the cache and self-heals on drift). Build the agent with
|
the cache and self-heals on drift). Build the agent with
|
||||||
`--no-default-features --features float16` to force the exact linear cosine scan.
|
`--no-default-features --features float16` to force the exact linear cosine scan.
|
||||||
|
The agent's `parallel` feature (also default) builds the index on a thread
|
||||||
|
pool; the graph is identical with or without it.
|
||||||
The index uses the HNSW paper's diversity heuristic for neighbour selection
|
The index uses the HNSW paper's diversity heuristic for neighbour selection
|
||||||
(plain closest-M capped recall on clustered data: 0.31 recall@10 at 100K). Its
|
(plain closest-M capped recall on clustered data: 0.31 recall@10 at 100K). Its
|
||||||
graph is saved to `<store>.h5.ann` at each checkpoint and reloaded by `open()`
|
graph is saved to `<store>.h5.ann` at each checkpoint and reloaded by `open()`
|
||||||
(tied to the checkpoint by a generation id; stale/damaged sidecars are
|
(tied to the checkpoint by a generation id; stale/damaged sidecars are
|
||||||
ignored and the index rebuilt). `hybrid_search` keeps one incremental BM25
|
ignored and the index rebuilt). `MemoryConfig::quantized_index` (**on by
|
||||||
|
default** for new stores, persisted; stores predating the setting load as
|
||||||
|
`false` and keep their f32 index — guarded by
|
||||||
|
`tests/fixtures/store_v2_5_0.h5`; CLI opt-out is `create --f32-index`)
|
||||||
|
stores the index's own copy of the embeddings as `i8`,
|
||||||
|
which roughly halves a loaded store's memory (2.72x -> 1.74x the raw vectors
|
||||||
|
at 100K); because quantised distances are approximate and `ef` cannot
|
||||||
|
compensate, the query path then re-scores the candidate pool against the
|
||||||
|
exact embeddings, which holds recall at the f32 index's level. It is also
|
||||||
|
faster at equal recall: 1.63x the QPS on x86-64 (AVX2) and 1.18x on a
|
||||||
|
Raspberry Pi 5 (`clawhdf5_accel::dot_i8`, NEON `SDOT` via inline asm since
|
||||||
|
the intrinsic is unstable; plain NEON on pre-dotprod cores). The aarch64
|
||||||
|
code is `cfg`'d out on x86, so x86 CI never compiles or lints it — test it
|
||||||
|
on real ARM (`rpivision02`, 10.0.2.3, is a Pi 5). `hybrid_search` keeps one incremental BM25
|
||||||
index for the life of the store and never writes the store: Hebbian
|
index for the life of the store and never writes the store: Hebbian
|
||||||
activation boosts are persisted by the next checkpoint (or on drop), not per
|
activation boosts are persisted by the next checkpoint (or on drop), not per
|
||||||
query. Measure any search-path change with
|
query. Measure any search-path change with
|
||||||
@@ -67,8 +87,51 @@ Cargo workspace with 16 crates under `crates/` (plus `libaec-sys`, an internal F
|
|||||||
`export` do). An unreadable WAL (torn header, bad magic) is quarantined to
|
`export` do). An unreadable WAL (torn header, bad magic) is quarantined to
|
||||||
`<store>.h5.wal.corrupt-<ts>` rather than blocking `open()`; a WAL with an
|
`<store>.h5.wal.corrupt-<ts>` rather than blocking `open()`; a WAL with an
|
||||||
unknown *newer* version still fails and is left untouched.
|
unknown *newer* version still fails and is left untouched.
|
||||||
- `MemoryConfig::compression` uses deflate by default; enable the agent's
|
- `MemoryConfig::float16` (**on by default** for new stores, persisted;
|
||||||
`zstd` feature to compress embeddings with Zstd instead (links libzstd).
|
existing stores keep their recorded `false` — guarded by the v2.5.0
|
||||||
|
fixture in `tests/float16_store.rs`; CLI opt-out is `create --f32`) writes
|
||||||
|
`/memory/embeddings` as IEEE half precision (48% smaller file at 100K;
|
||||||
|
LongMemEval with real MiniLM embeddings identical to f32).
|
||||||
|
`MemoryCache::half_precision` rounds each embedding as it enters the cache (push, update, WAL replay, and on load of a store still
|
||||||
|
`f32` on disk), so memory and file agree bit for bit; the conversions live
|
||||||
|
in `clawhdf5_format::float16` and must stay the single implementation.
|
||||||
|
Values beyond ±65504 are `MemoryError::InvalidEntry`. Interop: every file
|
||||||
|
must open in h5py — `f32` datasets and empty datasets did not until
|
||||||
|
2026-09-23 (see `docs/known-issues.md`); the agent's `h5py_interop` test
|
||||||
|
guards a whole store.
|
||||||
|
- `HDF5Memory::search(query_emb, text, &SearchOptions)` is the full search
|
||||||
|
path: optional source-channel filter (applied before ranking; exact scan of
|
||||||
|
the allowed records whenever cheaper than `pool × M` index distance
|
||||||
|
evaluations, and as the fallback when the pool comes back short), fusion,
|
||||||
|
activation scaling, optional re-ranking and confidence rejection.
|
||||||
|
`hybrid_search`/`hybrid_search_with` are thin wrappers; `ClawhdfBackend`
|
||||||
|
(the `openclaw` module) is `search` with re-rank + confidence on.
|
||||||
|
- **OpenClaw is not supported** (decided 2026-09-25): clawhdf5 is not an
|
||||||
|
OpenClaw memory plugin and never was — the old `memory.backend = "clawhdf5"`
|
||||||
|
config was never valid. Don't reintroduce OpenClaw claims; `docs/openclaw.md`
|
||||||
|
records what a real plugin would need.
|
||||||
|
- **ZeroClaw does not use clawhdf5** (checked 2026-09-25 against upstream
|
||||||
|
v0.8.5 and the `osobh/zeroclaw` fork, and their full history): no
|
||||||
|
`clawhdf5` feature or backend exists; ZeroClaw's memory backends are
|
||||||
|
sqlite/lucid/postgres/qdrant/markdown/none behind its own `Memory` trait.
|
||||||
|
`clawhdf5-migrate`'s default SQLite layout (`memory_chunks`, `sessions`,
|
||||||
|
`entities`, `relations`) is not ZeroClaw's schema either (ZeroClaw's is a
|
||||||
|
`memories` table). Don't reintroduce integration claims without an
|
||||||
|
integration and a test against the real consumer. Measure changes with
|
||||||
|
`search_harness --options-study`.
|
||||||
|
- `MemoryConfig::compression` is off by default; when on, embeddings are
|
||||||
|
deflate-compressed, or Zstd with the agent's `zstd` feature (links libzstd).
|
||||||
|
- Signed checkpoints (`clawhdf5-agent` `signing` module): with
|
||||||
|
`HDF5Memory::set_signing_key` every checkpoint stores an Ed25519-signed
|
||||||
|
manifest (SHA-256 per record in a Merkle tree + settings/sessions/graph
|
||||||
|
hashes; per-record hashes in `/integrity/record_hashes`);
|
||||||
|
`HDF5Memory::verify(path, &pk)` locates edits. The hashes must cover exactly
|
||||||
|
what the file persists in the form the loader returns it (strings lose
|
||||||
|
trailing NULs; an empty WAL mark is not written) or untouched stores stop
|
||||||
|
verifying — `tests/signed_store.rs` round-trips awkward strings. The key is
|
||||||
|
never persisted; a signed store refuses to checkpoint without it
|
||||||
|
(`MemoryError::SigningKeyRequired`, and `MemoryError` is `#[non_exhaustive]`).
|
||||||
|
WAL entries after the checkpoint are not covered.
|
||||||
- `Dataset::verify_provenance()` (clawhdf5 facade, `provenance` feature, on by
|
- `Dataset::verify_provenance()` (clawhdf5 facade, `provenance` feature, on by
|
||||||
default) recomputes a dataset's SHA-256 and compares it against the
|
default) recomputes a dataset's SHA-256 and compares it against the
|
||||||
`_provenance_sha256` attribute written automatically on save when
|
`_provenance_sha256` attribute written automatically on save when
|
||||||
@@ -85,7 +148,7 @@ Cargo workspace with 16 crates under `crates/` (plus `libaec-sys`, an internal F
|
|||||||
Alerts never block a save — drain them with `HDF5Memory::take_anomaly_alerts`.
|
Alerts never block a save — drain them with `HDF5Memory::take_anomaly_alerts`.
|
||||||
`MemorySource` for this bookkeeping is inferred from the caller-supplied
|
`MemorySource` for this bookkeeping is inferred from the caller-supplied
|
||||||
`source_channel` string (a heuristic, not an authenticated trust boundary).
|
`source_channel` string (a heuristic, not an authenticated trust boundary).
|
||||||
- GPU-accelerated batch I/O for large dataset processing
|
- GPU-accelerated vector distance computation (`clawhdf5-gpu`, wgpu); HDF5 I/O itself is CPU-only
|
||||||
- Python and Node.js bindings for cross-language use
|
- Python and Node.js bindings for cross-language use
|
||||||
- NetCDF-4 compatibility for scientific data interop
|
- NetCDF-4 compatibility for scientific data interop
|
||||||
|
|
||||||
@@ -101,6 +164,24 @@ cargo build --release
|
|||||||
cargo test --workspace
|
cargo test --workspace
|
||||||
```
|
```
|
||||||
|
|
||||||
|
### CI
|
||||||
|
`.gitea/workflows/ci.yml` has two jobs, both green as of 2026-09-22:
|
||||||
|
- **`test`** (`ubuntu-latest`, in `rust:latest`) runs `scripts/ci-test.sh` with
|
||||||
|
the h5py/netCDF4 interop suites required (`CLAWHDF5_REQUIRE_INTEROP=1`).
|
||||||
|
Served by the `tank` and `architect` runners.
|
||||||
|
- **`test-arm64`** (`linux_arm64`) lints and tests the aarch64 code — the NEON
|
||||||
|
kernels are `cfg`'d out on x86, so this is the only place they are built.
|
||||||
|
Served by `vision-01` (host mode) and `vision-02` (Docker), so steps must
|
||||||
|
work in both.
|
||||||
|
|
||||||
|
Keep workflows free of JavaScript actions (`actions/checkout`, `actions/cache`,
|
||||||
|
…): `rust:latest` has no `node`, and not every runner reaches GitHub, where
|
||||||
|
they are fetched from. Check out with plain `git` instead. The `test` job
|
||||||
|
installs `cmake` for the opt-in `fast-deflate` (zlib-ng) steps; the default
|
||||||
|
build needs no C toolchain, so `test-arm64` does not.
|
||||||
|
All runners are on `gitea-runner` 3.5.0, from `docker.gitea.com/act_runner`
|
||||||
|
— `gitea/act_runner:latest` on Docker Hub is frozen at 0.6.1.
|
||||||
|
|
||||||
### CLI
|
### CLI
|
||||||
```bash
|
```bash
|
||||||
cargo run -p clawhdf5-cli -- --help
|
cargo run -p clawhdf5-cli -- --help
|
||||||
@@ -115,4 +196,12 @@ python -c "import clawhdf5; print(clawhdf5.__version__)"
|
|||||||
```
|
```
|
||||||
|
|
||||||
## Integration
|
## Integration
|
||||||
ZeroClaw imports this as a Cargo feature (`clawhdf5` feature flag) to persist agent memory with HNSW vector search for context retrieval.
|
- **ClawBrainHub** (`clawverse/clawbrainhub` on git.redclaw.dev) is the one
|
||||||
|
verified consumer: `cbh-core` reads and writes `.brain` files through the
|
||||||
|
facade (`File`, `FileBuilder`, `AttrValue`, `Selection`), `cbh-scanner`
|
||||||
|
uses the facade, and `cbh-cli` uses `clawhdf5_agent::bm25::BM25Index`. It
|
||||||
|
depends on this repo by path (`../clawhdf5`), so it builds against whatever
|
||||||
|
is checked out — changes to those APIs reach it directly. Verified
|
||||||
|
2026-09-25 against main: builds, and its 204 tests pass.
|
||||||
|
- OpenClaw and ZeroClaw were both described as consumers; neither integrates
|
||||||
|
clawhdf5 (see Key Features and `docs/openclaw.md`).
|
||||||
|
|||||||
+4
-1
@@ -21,8 +21,11 @@ members = [
|
|||||||
resolver = "2"
|
resolver = "2"
|
||||||
|
|
||||||
[workspace.package]
|
[workspace.package]
|
||||||
version = "2.4.0"
|
version = "2.7.0"
|
||||||
edition = "2024"
|
edition = "2024"
|
||||||
|
# Oldest toolchain that builds the whole workspace; CI checks it. wgpu (in
|
||||||
|
# clawhdf5-gpu) requires 1.92.
|
||||||
|
rust-version = "1.92"
|
||||||
license = "MIT"
|
license = "MIT"
|
||||||
repository = "https://git.redclaw.dev/quantumclaw/clawhdf5"
|
repository = "https://git.redclaw.dev/quantumclaw/clawhdf5"
|
||||||
|
|
||||||
|
|||||||
@@ -3,24 +3,102 @@
|
|||||||
**The memory layer AI agents deserve. One file. Pure Rust. Zero C dependencies.**
|
**The memory layer AI agents deserve. One file. Pure Rust. Zero C dependencies.**
|
||||||
|
|
||||||
[](LICENSE)
|
[](LICENSE)
|
||||||
[](https://www.rust-lang.org)
|
[](https://www.rust-lang.org)
|
||||||
[](#performance)
|
[](#building)
|
||||||
[](BENCHMARKS.md#longmemeval-results)
|
[](BENCHMARKS.md#longmemeval-results)
|
||||||
[](BENCHMARKS.md#memory-footprint)
|
[](BENCHMARKS.md#memory-footprint-1)
|
||||||
|
|
||||||
ClawHDF5 is a pure-Rust HDF5 implementation combined with a research-grade agent memory engine. It gives AI agents persistent, searchable, cryptographically verifiable memory — all stored in a single portable file.
|
ClawHDF5 is a pure-Rust HDF5 implementation combined with a research-grade agent memory engine. It gives AI agents persistent, searchable, cryptographically verifiable memory (Ed25519-signed checkpoints) — all stored in a single portable file.
|
||||||
|
|
||||||
> **Two things live here:**
|
> **Two things live here:**
|
||||||
> - **A general-purpose, pure-Rust HDF5 library** — zero C dependencies, NetCDF-4 support, SIMD/GPU acceleration. See the **[Crate Map](#crate-map)** and **[BENCHMARKS.md](BENCHMARKS.md)** for the libhdf5 head-to-head numbers.
|
> - **A general-purpose, pure-Rust HDF5 library** — zero C dependencies, NetCDF-4 support, SIMD/GPU acceleration. See the **[Crate Map](#crate-map)** and **[BENCHMARKS.md](BENCHMARKS.md)** for the libhdf5 head-to-head numbers.
|
||||||
> - **An agent memory layer built on top of it** — vector search, knowledge graph, hippocampal-style consolidation, in `clawhdf5-agent`.
|
> - **An agent memory layer built on top of it** — vector search, knowledge graph, hippocampal-style consolidation, in `clawhdf5-agent`.
|
||||||
|
|
||||||
```
|
The crates are not on crates.io yet, so depend on them from git:
|
||||||
cargo add clawhdf5 # core HDF5 read/write, no agent layer
|
|
||||||
cargo add clawhdf5-agent --features agent # + agent memory layer
|
```toml
|
||||||
|
[dependencies]
|
||||||
|
clawhdf5 = { git = "https://git.redclaw.dev/quantumclaw/clawhdf5" } # core HDF5 read/write
|
||||||
|
clawhdf5-agent = { git = "https://git.redclaw.dev/quantumclaw/clawhdf5" } # + agent memory layer
|
||||||
```
|
```
|
||||||
|
|
||||||
|
> **C dependencies, precisely:** the core crates (`clawhdf5`, `clawhdf5-agent`,
|
||||||
|
> `-format`, `-io`, `-filters`, `-ann`, `-accel`, `-netcdf4`, `-cli`) build no C
|
||||||
|
> code by default — no libhdf5, and deflate is the pure-Rust
|
||||||
|
> [zlib-rs](https://github.com/trifectatechfoundation/zlib-rs), which matches
|
||||||
|
> zlib-ng on HDF5 reads and writes and produces byte-identical output
|
||||||
|
> ([BENCHMARKS.md § Deflate backend](BENCHMARKS.md#deflate-backend-zlib-rs-vs-zlib-ng)).
|
||||||
|
> CI fails if a C-building crate enters their default dependency tree. C comes
|
||||||
|
> in only when you ask for it: `fast-deflate` (zlib-ng, needs cmake), `zstd`,
|
||||||
|
> `szip`, the BLAS backends, `clawhdf5-migrate` (bundled SQLite) and the
|
||||||
|
> Node.js bindings.
|
||||||
|
|
||||||
> **New here?** Start with the **[Quickstart Guide](docs/QUICKSTART.md)** · See **[Use Cases](docs/USE_CASES.md)** · Read **[Benchmarks](BENCHMARKS.md)**
|
> **New here?** Start with the **[Quickstart Guide](docs/QUICKSTART.md)** · See **[Use Cases](docs/USE_CASES.md)** · Read **[Benchmarks](BENCHMARKS.md)**
|
||||||
|
|
||||||
|
## What's new (v2.2 → v2.7, and unreleased)
|
||||||
|
|
||||||
|
Five releases in September 2026. Details, including upgrade notes and every
|
||||||
|
breaking change, are in [CHANGELOG.md](CHANGELOG.md).
|
||||||
|
|
||||||
|
**HDF5 correctness (read these if you read files with an earlier release)**
|
||||||
|
- **Extensible Array chunk indexes returned wrong data** past the 36th chunk —
|
||||||
|
any dataset with one unlimited dimension. Silent: plausible numbers from the
|
||||||
|
wrong chunks. Fixed in v2.7.0; re-read affected data.
|
||||||
|
- Fixed and Extensible Array checksums are now verified, so a corrupt chunk
|
||||||
|
index is `ChecksumMismatch` instead of wrong data (v2.7.0).
|
||||||
|
- Compound datatypes written with default libver bounds (plain
|
||||||
|
`h5py.File(path, 'w')`) were mis-parsed; HDF5 2.0 compound v5 and native
|
||||||
|
complex (class 11) types now parse (v2.2.0–v2.3.0).
|
||||||
|
- Committed datatypes, fill values, soft links and `H5T_STD_REF` references now
|
||||||
|
read correctly; external links and external raw data are explicit errors;
|
||||||
|
`attrs()` no longer silently drops attributes (v2.3.0–v2.5.0).
|
||||||
|
- Datasets indexed by a version-2 B-tree now read (v2.5.0).
|
||||||
|
|
||||||
|
**Security and robustness**
|
||||||
|
- A crafted file could abort any reader via B-tree v2 recursion or explode it
|
||||||
|
via shared children; both are now fast errors (v2.7.0).
|
||||||
|
- Virtual-dataset source paths are confined to the file's directory; chunked
|
||||||
|
reads use overflow-checked sizes and fallible allocation, and the facade
|
||||||
|
writes files atomically (v2.3.0).
|
||||||
|
- Agent store: single-writer lock plus `open_read_only`; a crash between
|
||||||
|
checkpoint and WAL truncate no longer duplicates entries; unreadable WALs are
|
||||||
|
quarantined instead of blocking `open()` (v2.3.0).
|
||||||
|
|
||||||
|
**Search quality and speed**
|
||||||
|
- HNSW neighbour selection now uses the paper's diversity heuristic: recall@10
|
||||||
|
at 100K went from 0.31 to 0.98 (v2.4.0).
|
||||||
|
- `hybrid_search` is 79–190× faster than v2.3.0 (p50 0.07 ms at 1K, 4.65 ms at
|
||||||
|
100K). It no longer rebuilds BM25 or rewrites the store per query, and the
|
||||||
|
HNSW graph is persisted (v2.4.0).
|
||||||
|
- Default fusion weights are now the measured 0.4 / 0.6 (v2.5.0). Re-ranking had
|
||||||
|
been discarding the retrieval score, costing the Markdown backend 40.6pp of
|
||||||
|
Hit@1; fixed in v2.6.0.
|
||||||
|
- Selection reads decode only the chunks they touch (a 64×64 window: 105 ms to
|
||||||
|
0.39 ms), and full reads are 1.2–1.9× faster (v2.5.0).
|
||||||
|
|
||||||
|
**Memory**
|
||||||
|
- A loaded store holds ~30% less (embeddings stored once, v2.6.0), and the
|
||||||
|
int8 HNSW index, **on by default for new stores** (unreleased), brings a
|
||||||
|
100K × 384 store to 1.74× the raw vectors. At equal recall it is also faster
|
||||||
|
than `f32`: 1.63× QPS on AVX2, 1.18× on a Raspberry Pi 5 (NEON `SDOT`).
|
||||||
|
|
||||||
|
**Interop and search (unreleased)**
|
||||||
|
- **Files we write now open in h5py and libhdf5.** Every `f32` dataset —
|
||||||
|
including every agent store's embeddings — and every empty dataset was
|
||||||
|
refused by libhdf5. Both were write-side bugs in every release; agent stores
|
||||||
|
fix themselves at their next checkpoint. See
|
||||||
|
[docs/known-issues.md](docs/known-issues.md).
|
||||||
|
- `MemoryConfig::float16` now stores half-precision embeddings (it was
|
||||||
|
ignored), and is on by default for new stores: 48% smaller files, and
|
||||||
|
identical LongMemEval retrieval on real embeddings.
|
||||||
|
- `HDF5Memory::search` with `SearchOptions`: filter by source channel (exact
|
||||||
|
filtered top-k, never slower than unfiltered), and opt-in re-ranking and
|
||||||
|
confidence rejection, which used to be reachable only through `ClawhdfBackend`.
|
||||||
|
|
||||||
|
**Tooling**
|
||||||
|
- CI now runs the h5py/netCDF4 interop suites for real (they had been skipping
|
||||||
|
silently) and runs an aarch64 job for the NEON kernels.
|
||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
## Why ClawhDF5?
|
## Why ClawhDF5?
|
||||||
@@ -33,16 +111,16 @@ Every AI agent needs memory. Today that means scattered Markdown files, SQLite d
|
|||||||
| Keyword search | Separate FTS engine | Integrated BM25 |
|
| Keyword search | Separate FTS engine | Integrated BM25 |
|
||||||
| Knowledge graph | Neo4j or none | In-file graph with spreading activation |
|
| Knowledge graph | Neo4j or none | In-file graph with spreading activation |
|
||||||
| Memory consolidation | Manual pruning | Hippocampal-inspired automatic tiers |
|
| Memory consolidation | Manual pruning | Hippocampal-inspired automatic tiers |
|
||||||
| Temporal queries | Custom code | Native temporal index (716ns) |
|
| Temporal queries | Custom code | Native temporal index (622 ns range query over 10K) |
|
||||||
| Multi-modal | Multiple stores | Unified cross-modal search |
|
| Multi-modal | Multiple stores | Unified cross-modal search (exact scan: 842 µs over 1K records) |
|
||||||
| Security | Hope for the best | Provenance tracking + anomaly detection |
|
| Integrity | Hope for the best | Ed25519-signed checkpoints that pinpoint any edited record, chained-CRC WAL, checksummed chunk indexes, write-anomaly alerts |
|
||||||
| Portability | Config + DB + files | **One `.h5` file. Copy it anywhere.** |
|
| Portability | Config + DB + files | **One `.h5` file. Copy it anywhere.** |
|
||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
## Performance
|
## Performance
|
||||||
|
|
||||||
Vector search and agent-memory operations below are benchmarked on Intel i7-12650H (10C/16T), 384-dim embeddings, Criterion.rs. The HDF5 Core I/O table immediately below is from a separate, independently reproduced run (see its own hardware note).
|
The brute-force/IVF vector search, agent-memory, on-disk footprint and consolidation figures below were measured 2026-09-24 on tank (AMD Ryzen 7 7800X3D, 8C/16T), commit 5c8323c, 384-dim embeddings; the commands are in [BENCHMARKS.md](BENCHMARKS.md). Exceptions are marked where they appear: the HDF5 Core I/O table immediately below is from a separate, independently reproduced run (see its own hardware note), and the HNSW `f32`/`i8` table and the in-memory `i8` column were not re-measured on 2026-09-24.
|
||||||
|
|
||||||
### HDF5 Core I/O (vs libhdf5 1.14.6)
|
### HDF5 Core I/O (vs libhdf5 1.14.6)
|
||||||
|
|
||||||
@@ -58,31 +136,63 @@ Figures below are from an independent reproduction run on a second machine (AMD
|
|||||||
| Sequential read (100K f32) | 23.3 µs | 63.6 µs | **2.7×** |
|
| Sequential read (100K f32) | 23.3 µs | 63.6 µs | **2.7×** |
|
||||||
| Sequential write (100K f32) | 210 µs | 189 µs | **≈ tie** |
|
| Sequential write (100K f32) | 210 µs | 189 µs | **≈ tie** |
|
||||||
|
|
||||||
|
The chunked-write row was re-measured on the same machine on 2026-09-23, after
|
||||||
|
the default deflate backend became pure-Rust zlib-rs: 1.46 ms against
|
||||||
|
libhdf5's 51.4 ms (**35×**), and 1.48 ms with zlib-ng. libhdf5's own time on
|
||||||
|
that machine moved from 65.0 to 51.4 ms between the two dates, which is most
|
||||||
|
of the difference from 45×; compare same-day numbers only.
|
||||||
|
|
||||||
### Vector Search
|
### Vector Search
|
||||||
|
|
||||||
| Scale | Flat | IVF (nprobe=10) | IVF-PQ | vs MemX¹ |
|
**HNSW (the default backend for `hybrid_search`)** — `search_harness`, clustered
|
||||||
|-------|------|-----------------|--------|----------|
|
384-dim data, M = 16, ef_construction = 64, recall measured against an exact scan.
|
||||||
| 1K | **54 µs** | — | — | — |
|
See [BENCHMARKS.md § Search harness](BENCHMARKS.md#search-harness-baseline-v230)
|
||||||
| 10K | 753 µs | **27 µs** | — | — |
|
and [§ Quantising the index copy](BENCHMARKS.md#quantising-the-index-copy-quantized_index):
|
||||||
| 100K | 11.4 ms | 1.32 ms | **1.19 ms** | ~8–76× (see caveat) |
|
|
||||||
|
|
||||||
> Reproduced on the same second machine (Ryzen 7 7800X3D) with a corrected,
|
| N = 100K, ef = 64 | recall@10 | QPS | build |
|
||||||
> apples-to-apples SIMD/scalar/parallel comparison methodology — see
|
|---|---:|---:|---:|
|
||||||
> [BENCHMARKS.md § Independent Validation: tank — LongMemEval & Vector
|
| `f32` index | 0.9945 | 13 399 | 3.2 s |
|
||||||
> Search](BENCHMARKS.md#independent-validation-tank--longmemeval--vector-search-ryzen-7-7800x3d-2026-08-05).
|
| `i8` index + exact re-score (**default for new stores**) | 0.9940 | **21 848** | **1.8 s** |
|
||||||
|
|
||||||
|
Before the v2.4.0 neighbour-selection fix, recall@10 at 100K was 0.31. These
|
||||||
|
two rows are a paired comparison (medians of alternating runs, same binary).
|
||||||
|
A single `f32` run on 2026-09-24 measured recall 0.9945, 19 001 QPS and a
|
||||||
|
2.7 s build; the int8 row was not re-run, so the pair has not been re-checked
|
||||||
|
([§ Quantising the index copy](BENCHMARKS.md#quantising-the-index-copy-quantized_index)).
|
||||||
|
|
||||||
|
**Brute-force and IVF paths** (Criterion, tank, 2026-09-24):
|
||||||
|
|
||||||
|
| Scale | Flat | IVF (nprobe=10) | IVF-PQ | MemX¹ (claimed, end-to-end) |
|
||||||
|
|-------|------|-----------------|--------|----------|
|
||||||
|
| 1K | **47.4 µs** | — | — | — |
|
||||||
|
| 10K | 500.5 µs | **24.8 µs** | — | — |
|
||||||
|
| 100K | 6.58 ms | 592 µs | **869 µs** | <90 ms |
|
||||||
|
|
||||||
|
> These replace figures from the original i7-12650H run (flat 54 µs / 753 µs /
|
||||||
|
> 11.4 ms); a 2026-08-05 run on tank had already matched the new ones — see
|
||||||
|
> [BENCHMARKS.md § Vector Search Latency](BENCHMARKS.md#vector-search-latency).
|
||||||
|
|
||||||
### Agent Memory Operations
|
### Agent Memory Operations
|
||||||
|
|
||||||
| Operation | Latency | Scale |
|
| Operation | Latency | Scale |
|
||||||
|-----------|---------|-------|
|
|-----------|---------|-------|
|
||||||
| Hybrid search (RRF) | **222 µs** | 1K records |
|
| Hybrid search (`HDF5Memory::hybrid_search`, p50) | **0.07 ms** / 0.49 ms / 4.69 ms | 1K / 10K / 100K records |
|
||||||
| BM25 keyword search | **67 µs** | 1K records |
|
| BM25 keyword search | **20.4 µs** | 1K records |
|
||||||
| Knowledge graph BFS | **24 µs** | 1K entities |
|
| Knowledge graph BFS | **23.1 µs** | 1K entities |
|
||||||
| Spreading activation | **17 µs** | 100 entities |
|
| Spreading activation | **10.1 µs** | 100 entities |
|
||||||
| Temporal range query | **716 ns** | 10K timestamps |
|
| Temporal range query | **622 ns** | 10K timestamps |
|
||||||
| Consolidation cycle | **164 µs** | 1K records |
|
| Consolidation cycle | **115.2 µs** | 1K records |
|
||||||
| Memory write (WAL) | **18 µs** | per record (group-commit append; HDF5 batched at flush) |
|
| Cross-modal search (exact scan, 2 embeddings per record) | **842.0 µs** / 8.44 ms | 1K / 10K records |
|
||||||
| Importance gate | **61 ns** | per record |
|
| Memory write (WAL) | **26.1 µs** | per record (group-commit append; HDF5 batched at flush) |
|
||||||
|
| Importance gate | **57.6 ns** | per record (trivial skip) |
|
||||||
|
|
||||||
|
The old 18 µs WAL write was undated, from another machine: v2.3.0 measures
|
||||||
|
24.3 µs on the same hardware as this table, the same as an `f32` store today.
|
||||||
|
`float16` stores (the new default) add ~2 µs for rounding; the int8 index adds
|
||||||
|
nothing. See [BENCHMARKS.md § Write Path](BENCHMARKS.md#write-path).
|
||||||
|
Knowledge-graph traversal was briefly 6.5x slower (155 µs) until this re-run
|
||||||
|
found and fixed an adjacency index rebuilt on every traversal; see
|
||||||
|
[§ Knowledge Graph](BENCHMARKS.md#knowledge-graph).
|
||||||
|
|
||||||
### Chunked Write Throughput (codec comparison)
|
### Chunked Write Throughput (codec comparison)
|
||||||
|
|
||||||
@@ -97,7 +207,7 @@ by default (AoS→SoA byte transpose, +157–204% throughput for float data):
|
|||||||
|
|
||||||
Use `.with_zstd(3)` or `.with_deflate(6)` for write-heavy workloads — both now perform at ~720–750 MiB/s on large matrices. Use `.with_pcodec()` for write-once/read-many workloads where compression ratio matters more than encode speed. Disable auto-shuffle with `.without_shuffle()` for byte arrays that don't benefit from AoS→SoA transposition.
|
Use `.with_zstd(3)` or `.with_deflate(6)` for write-heavy workloads — both now perform at ~720–750 MiB/s on large matrices. Use `.with_pcodec()` for write-once/read-many workloads where compression ratio matters more than encode speed. Disable auto-shuffle with `.without_shuffle()` for byte arrays that don't benefit from AoS→SoA transposition.
|
||||||
|
|
||||||
> ¹ MemX ([arxiv:2603.16171](https://arxiv.org/abs/2603.16171), March 2026): Rust + libSQL, claims <90ms at 100K records. **Not like-for-like:** MemX's figure is *end-to-end* (embeddings + FTS5 + four-factor re-ranking); ours is a *single component* (raw vector search). The ratio overstates the real advantage by an unquantified margin — order-of-magnitude indication only. See [BENCHMARKS.md](BENCHMARKS.md#comparison-to-memx-arxiv260316171).
|
> ¹ MemX ([arxiv:2603.16171](https://arxiv.org/abs/2603.16171), March 2026): Rust + libSQL, claims <90ms at 100K records. **Not like-for-like:** MemX's figure is *end-to-end* (embeddings + FTS5 + four-factor re-ranking); ours is a *single component* (raw vector search), so the two columns are not comparable and no ratio is given. See [BENCHMARKS.md](BENCHMARKS.md#comparison-to-memx-arxiv260316171).
|
||||||
|
|
||||||
### LongMemEval Retrieval Recall
|
### LongMemEval Retrieval Recall
|
||||||
|
|
||||||
@@ -115,13 +225,17 @@ declaration:
|
|||||||
|
|
||||||
Hybrid is the strongest configuration, which is what running two retrieval stages
|
Hybrid is the strongest configuration, which is what running two retrieval stages
|
||||||
is for. The weights matter more than the stages: a sweep of `vector_weight` from
|
is for. The weights matter more than the stages: a sweep of `vector_weight` from
|
||||||
0.0 to 1.0 found the long-standing `0.7/0.3` default is **strictly dominated** by
|
0.0 to 1.0 found the old `0.7/0.3` default is **strictly dominated** by
|
||||||
`0.4/0.6` — better on Hit@1, Hit@5, Hit@10 and MRR at both granularities. Use
|
`0.4/0.6` — better on Hit@1, Hit@5, Hit@10 and MRR at both granularities. Since
|
||||||
`0.4/0.6`, or `0.3/0.7` if rank-1 precision matters most. See
|
v2.5.0 `0.4/0.6` is the default (`hybrid::DEFAULT_FUSION`, used by
|
||||||
[BENCHMARKS.md § Weight sweep](BENCHMARKS.md#longmemeval-results).
|
`unified_search`, `hybrid_search_with` and `ClawhdfBackend`); callers that
|
||||||
|
pass weights to `hybrid_search` explicitly choose their own. Use `0.3/0.7` if
|
||||||
|
rank-1 precision matters most. Reciprocal rank fusion is selectable
|
||||||
|
(`hybrid::Fusion::Rrf`) but measured worse than the weighted sum. See
|
||||||
|
[BENCHMARKS.md § Weight sweep](BENCHMARKS.md#weight-sweep--full-haystack-n500).
|
||||||
|
|
||||||
Vector embeddings require `--features embeddings`; without it the vector stage is
|
The benchmark's vector stage requires `clawhdf5-bench`'s `embeddings` feature
|
||||||
inert and only the BM25 row is produced, which is what every previously published
|
(real MiniLM embeddings); without it the vector stage is inert and only the BM25 row is produced, which is what every previously published
|
||||||
number here measured.
|
number here measured.
|
||||||
|
|
||||||
On the easier `longmemeval_oracle` variant (evidence sessions only) the same
|
On the easier `longmemeval_oracle` variant (evidence sessions only) the same
|
||||||
@@ -146,19 +260,53 @@ retrieval recall reported as QA accuracy typically overstates by 20–30 points.
|
|||||||
|
|
||||||
### Memory Footprint
|
### Memory Footprint
|
||||||
|
|
||||||
| Records | File Size | Bytes/Record | With Compression |
|
**On disk** — 384-dim `float16` embeddings (the default for new stores),
|
||||||
|---------|-----------|--------------|------------------|
|
200-char text, `footprint_bench`
|
||||||
| 1K | ~6.5 MB | ~6.5 KB | ~2.1 MB (3.1x) |
|
([BENCHMARKS.md § Memory Footprint](BENCHMARKS.md#memory-footprint-1)):
|
||||||
| 10K | ~65 MB | ~6.5 KB | ~21 MB (3.1x) |
|
|
||||||
| 100K | ~645 MB | ~6.5 KB | ~208 MB (3.1x) |
|
| Records | File Size | Bytes/Record | Gzip-6 compressed |
|
||||||
|
|---------|-----------|--------------|-------------------|
|
||||||
|
| 1K | 810.4 KB | 829 B | 56.4 KB |
|
||||||
|
| 10K | 7.8 MB | 820 B | 471.3 KB |
|
||||||
|
| 100K | 76.7 MB | 803 B | 4.5 MB |
|
||||||
|
|
||||||
|
The benchmark's synthetic embeddings and text are far more repetitive than
|
||||||
|
real data (only 40 distinct texts), so no column here is an expectation for
|
||||||
|
real data. The compressed column is an upper bound, and the Bytes/Record
|
||||||
|
column is optimistic too: it is not an uncompressed figure, because the store
|
||||||
|
always deflates its text (any string dataset of 4 KiB or more) whatever
|
||||||
|
`MemoryConfig::compression` says. The `float16` embeddings alone are 768 B per
|
||||||
|
record, so 200 characters of real text would take a record above 820 B.
|
||||||
|
This table used to show `f32` stores (1.7 KB per record, 169.8 MB at 100K);
|
||||||
|
those were not re-measured. The float16 study compares the two on the same
|
||||||
|
data: 100K × 384 records take 80.8 MiB as `float16` and 154.0 MiB as `f32`.
|
||||||
|
|
||||||
|
**In memory** — a store reopened from disk, 384-dim `f32`, measured with a
|
||||||
|
counting allocator ([BENCHMARKS.md § Memory footprint](BENCHMARKS.md#memory-footprint)):
|
||||||
|
|
||||||
|
| Records | Raw vectors | Reopened, `f32` index | Reopened, `i8` index (default) |
|
||||||
|
|---------|-------------|-----------------------|--------------------------------|
|
||||||
|
| 1K | 1 MiB | 4 MiB (2.40x) | 2 MiB (1.64x) |
|
||||||
|
| 10K | 15 MiB | 44 MiB (3.03x) | 27 MiB (1.81x) |
|
||||||
|
| 100K | 146 MiB | 399 MiB (2.72x) | **256 MiB (1.74x)** |
|
||||||
|
|
||||||
|
Down from 505 MiB (3.44x) at 100K before v2.6.0, when the cache held every
|
||||||
|
embedding twice. The `f32` column was re-measured on 2026-09-24 and reproduced
|
||||||
|
exactly; the `i8` column was not re-run.
|
||||||
|
|
||||||
### Consolidation Efficiency
|
### Consolidation Efficiency
|
||||||
|
|
||||||
|
1,000 records (10 signal + 990 noise), `working_capacity = 100`
|
||||||
|
([BENCHMARKS.md § Consolidation Efficiency](BENCHMARKS.md#consolidation-efficiency)):
|
||||||
|
|
||||||
| Metric | Before | After | Delta |
|
| Metric | Before | After | Delta |
|
||||||
|--------|--------|-------|-------|
|
|--------|--------|-------|-------|
|
||||||
| Records in store | 1,000 | ~110 | −89% |
|
| Records in store | 1,000 | 100 | −90% |
|
||||||
| Hit@1 recall | ~60% | ~90% | +30% |
|
| Hit@1 recall (signal records) | 100% | 100% | no loss |
|
||||||
| Search latency | ~2.8 ms | ~0.3 ms | **9x faster** |
|
| Search latency (avg) | 2.22 ms | 0.24 ms | **9.3x faster** |
|
||||||
|
|
||||||
|
The consolidation cycle that does this took 0.13 ms; a cycle over 10K records
|
||||||
|
takes 2.81 ms and over 100K 46.7 ms.
|
||||||
|
|
||||||
**Full benchmark details: [BENCHMARKS.md](BENCHMARKS.md)**
|
**Full benchmark details: [BENCHMARKS.md](BENCHMARKS.md)**
|
||||||
|
|
||||||
@@ -166,74 +314,74 @@ retrieval recall reported as QA accuracy typically overstates by 20–30 points.
|
|||||||
|
|
||||||
## Agent Memory Architecture
|
## Agent Memory Architecture
|
||||||
|
|
||||||
ClawhDF5's agent memory engine implements research from 15+ recent papers on agentic memory systems. It's not a toy — it's the real thing.
|
ClawhDF5's agent memory engine draws on 15+ recent papers on agentic memory systems (see [Research Foundation](#research-foundation)).
|
||||||
|
|
||||||
```
|
```
|
||||||
┌─────────────────┐
|
┌─────────────────┐
|
||||||
│ Agent Query │
|
│ Agent Query │
|
||||||
└────────┬────────┘
|
└────────┬────────┘
|
||||||
│
|
│
|
||||||
┌────────────▼────────────┐
|
┌─────────────────▼──────────────────┐
|
||||||
│ Hybrid Retrieval │
|
│ HDF5Memory::search │
|
||||||
│ Vector + BM25 + RRF │
|
│ optional source-channel filter │
|
||||||
└────────────┬────────────┘
|
│ HNSW vector + BM25 keyword │
|
||||||
│
|
│ weighted fusion (0.4 / 0.6) │
|
||||||
┌──────────────────▼──────────────────┐
|
│ × √(Hebbian activation) │
|
||||||
│ Multi-Factor Re-Ranking │
|
└─────────────────┬──────────────────┘
|
||||||
│ temporal · authority · activation │
|
│ opt-in (SearchOptions);
|
||||||
└──────────────────┬──────────────────┘
|
│ ClawhdfBackend turns both on
|
||||||
│
|
┌─────────────────▼──────────────────┐
|
||||||
┌────────────▼────────────┐
|
│ Multi-factor re-ranking │
|
||||||
│ Confidence Rejection │
|
│ relevance · recency · authority · │
|
||||||
│ (suppress bad matches) │
|
│ activation │
|
||||||
└────────────┬────────────┘
|
├────────────────────────────────────┤
|
||||||
│
|
│ Confidence rejection │
|
||||||
┌────────────────────────▼────────────────────────┐
|
│ (suppress bad matches) │
|
||||||
│ Memory Store (HDF5) │
|
└─────────────────┬──────────────────┘
|
||||||
│ │
|
│
|
||||||
│ ┌───────────┐ ┌───────────┐ ┌───────────────┐ │
|
┌────────────────────────────▼────────────────────────────┐
|
||||||
│ │ Working │→│ Episodic │→│ Semantic │ │
|
│ In memory │
|
||||||
│ │ (bounded) │ │ (bounded) │ │ (long-term) │ │
|
│ cache (flat f32 embeddings) · BM25 index · HNSW index │
|
||||||
│ └───────────┘ └───────────┘ └───────────────┘ │
|
│ provenance ledger + anomaly alerts (session-scoped) │
|
||||||
│ │
|
└────────────────────────────┬────────────────────────────┘
|
||||||
│ ┌──────────┐ ┌──────────┐ ┌────────────────┐ │
|
│ WAL append; checkpoint
|
||||||
│ │Knowledge │ │Temporal │ │ Multi-Modal │ │
|
┌────────────────────────────▼────────────────────────────┐
|
||||||
│ │ Graph │ │ Index │ │ Embeddings │ │
|
│ agent_memory.h5 /meta · /memory · /sessions · │
|
||||||
│ └──────────┘ └──────────┘ └────────────────┘ │
|
│ /knowledge_graph │
|
||||||
│ │
|
│ agent_memory.h5.wal chained-CRC write-ahead log │
|
||||||
│ ┌──────────┐ ┌──────────┐ ┌────────────────┐ │
|
│ agent_memory.h5.ann HNSW graph (derived, rebuildable) │
|
||||||
│ │Provenance│ │ Anomaly │ │ Source │ │
|
│ agent_memory.h5.lock single-writer lock │
|
||||||
│ │ Tracking │ │Detection │ │ Isolation │ │
|
└─────────────────────────────────────────────────────────┘
|
||||||
│ └──────────┘ └──────────┘ └────────────────┘ │
|
|
||||||
└─────────────────────────────────────────────────┘
|
|
||||||
│
|
|
||||||
┌────────┴────────┐
|
|
||||||
│ agent_memory.h5 │
|
|
||||||
│ single file │
|
|
||||||
└─────────────────┘
|
|
||||||
```
|
```
|
||||||
|
|
||||||
|
Consolidation tiers (Working → Episodic → Semantic), the knowledge-graph
|
||||||
|
algorithms, temporal and multi-modal indexes are library components you drive
|
||||||
|
directly; the store persists the records, sessions and graph they work over.
|
||||||
|
|
||||||
### Module Overview
|
### Module Overview
|
||||||
|
|
||||||
| Module | What It Does |
|
| Module | What It Does |
|
||||||
|--------|-------------|
|
|--------|-------------|
|
||||||
| **`knowledge`** | Entity/relation graph with BFS traversal, spreading activation, fuzzy entity resolution |
|
| **`knowledge`** | Entity/relation graph with BFS traversal, spreading activation, fuzzy (Levenshtein) entity resolution |
|
||||||
| **`consolidation`** | Three-tier memory (Working → Episodic → Semantic) with importance scoring and time-decay |
|
| **`consolidation`** | Three-tier memory (Working → Episodic → Semantic) with importance scoring, novelty, and time-decay |
|
||||||
| **`hybrid`** | Vector + BM25 fusion with Reciprocal Rank Fusion (RRF, k=60). The vector stage uses the HNSW index by default (`hnsw` feature, on by default); disable with `--no-default-features --features float16` for an exact linear scan |
|
| **`hybrid`** | Vector + BM25 fusion. Default is a min-max-normalised weighted sum, vector 0.4 / keyword 0.6 (`hybrid::DEFAULT_FUSION`, tuned on LongMemEval); RRF is available via `Fusion::Rrf` / `hybrid_search_with`. The vector stage uses the HNSW index by default (`hnsw` feature); disable with `--no-default-features --features float16` for an exact linear scan |
|
||||||
| **`reranker`** | Multi-factor re-ranking: temporal recency, source authority, activation weight |
|
| **`reranker`** | Multi-factor re-ranking: retrieval relevance (leads, weight 1.0), temporal recency, source authority, activation weight. Opt-in via `SearchOptions::with_rerank`; on in `ClawhdfBackend` |
|
||||||
| **`confidence`** | Low-confidence rejection — suppresses spurious recalls when nothing matches |
|
| **`confidence`** | Low-confidence rejection — suppresses spurious recalls when nothing matches. Opt-in via `SearchOptions::with_confidence`; on in `ClawhdfBackend` |
|
||||||
| **`temporal`** | Sorted timestamp index, session DAG, entity timeline, temporal query hints |
|
| **`temporal`** | Sorted timestamp index, session DAG, entity timeline, temporal query hints |
|
||||||
| **`multimodal`** | Cross-modal search across text/image/audio/video embeddings |
|
| **`multimodal`** | Cross-modal search across text/image/audio/video embeddings |
|
||||||
| **`provenance`** | Source attribution, FNV-1a content hashing, integrity verification |
|
| **`signing`** | Ed25519-signed checkpoints: SHA-256 per record in a Merkle tree, plus hashes of settings, sessions and the knowledge graph; `HDF5Memory::verify` names any edited record |
|
||||||
| **`anomaly`** | Write rate limiting, 15 injection pattern detectors, source distribution analysis |
|
| **`provenance`** | Source attribution and an unkeyed FNV-1a content hash per record, held in memory for the session, for detecting accidental corruption (not tamper-proof) |
|
||||||
| **`openclaw`** | OpenClaw integration: MemoryBackend trait, Markdown ↔ HDF5 conversion |
|
| **`anomaly`** | Write rate limiting, 15 injection-pattern detectors, source-distribution analysis. Alerts never block a save; drain them with `take_anomaly_alerts` |
|
||||||
|
| **`openclaw`** | `ClawhdfBackend`: a Markdown-oriented backend (ingest by section, search, read back by path, export). Named for OpenClaw, but **not an OpenClaw plugin** — see [docs/openclaw.md](docs/openclaw.md) |
|
||||||
| **`vector_search`** | Flat cosine, pre-normed, SIMD, BLAS, GPU, parallel search paths |
|
| **`vector_search`** | Flat cosine, pre-normed, SIMD, BLAS, GPU, parallel search paths |
|
||||||
| **`ivf` / `pq`** | IVF-PQ approximate nearest neighbor for billion-scale search |
|
| **`ivf` / `pq`** | Standalone IVF and IVF-PQ indexes (benchmarked to 100K vectors); not used by `HDF5Memory`, whose ANN index is HNSW |
|
||||||
| **`bm25`** | BM25 keyword index with TF-IDF scoring |
|
| **`bm25`** | Incremental Okapi BM25 inverted index, kept for the life of the store; optional stemming |
|
||||||
|
| **`query_expand`** | Synonym / acronym / temporal query expansion |
|
||||||
| **`entity_extract`** | Rule-based entity extraction from text chunks into the knowledge graph |
|
| **`entity_extract`** | Rule-based entity extraction from text chunks into the knowledge graph |
|
||||||
| **`wal`** | Write-ahead log for crash-safe persistence; each entry is CRC32-checked on replay, so a corrupted entry stops replay there instead of loading bad data |
|
| **`wal`** | Write-ahead log (v4) with a chained CRC32 per entry, so a corrupted, reordered, duplicated or spliced entry stops replay; checkpoints record a WAL mark so nothing is applied twice. Appends are not fsynced |
|
||||||
| **`memory_strategy`** | Pluggable strategies: save-every, semantic-shift, user-correction detection |
|
| **`memory_strategy`** | Pluggable strategies: save-every, semantic-shift, user-correction detection |
|
||||||
| **`decision_gate`** | Sub-microsecond trivial/substantive classification |
|
| **`decision_gate`** | Sub-microsecond trivial/substantive classification |
|
||||||
|
| **`ephemeral`** | In-memory TTL/LFU working tier |
|
||||||
| **`async_memory`** | Tokio-based async wrapper over the memory store (`async` feature) |
|
| **`async_memory`** | Tokio-based async wrapper over the memory store (`async` feature) |
|
||||||
|
|
||||||
---
|
---
|
||||||
@@ -265,7 +413,7 @@ assert_eq!(values, vec![22.5, 23.1, 21.8]);
|
|||||||
use clawhdf5_agent::{HDF5Memory, MemoryConfig, MemoryEntry, AgentMemory};
|
use clawhdf5_agent::{HDF5Memory, MemoryConfig, MemoryEntry, AgentMemory};
|
||||||
|
|
||||||
// Create memory store
|
// Create memory store
|
||||||
let config = MemoryConfig::new("agent.h5", "my-agent", 384);
|
let config = MemoryConfig::new("agent.h5".into(), "my-agent", 384);
|
||||||
let mut memory = HDF5Memory::create(config)?;
|
let mut memory = HDF5Memory::create(config)?;
|
||||||
|
|
||||||
// Save a memory
|
// Save a memory
|
||||||
@@ -278,13 +426,67 @@ memory.save(MemoryEntry {
|
|||||||
tags: "preference".into(),
|
tags: "preference".into(),
|
||||||
})?;
|
})?;
|
||||||
|
|
||||||
// Search
|
// Hybrid search: vector + BM25, weighted 0.4 / 0.6 (the measured default)
|
||||||
let results = memory.search(&query_embedding, 5)?;
|
let results = memory.hybrid_search(&query_embedding, "user preferences", 0.4, 0.6, 5);
|
||||||
for result in results {
|
for result in results {
|
||||||
println!("[{:.3}] {}", result.score, result.chunk);
|
println!("[{:.3}] {}", result.score, result.chunk);
|
||||||
}
|
}
|
||||||
```
|
```
|
||||||
|
|
||||||
|
### Search Options
|
||||||
|
|
||||||
|
```rust
|
||||||
|
use clawhdf5_agent::SearchOptions;
|
||||||
|
use clawhdf5_agent::confidence::ConfidenceConfig;
|
||||||
|
use clawhdf5_agent::reranker::ReRankConfig;
|
||||||
|
|
||||||
|
// Only memories from these source channels; still a full page of k results.
|
||||||
|
let work = memory.search(
|
||||||
|
&query_embedding,
|
||||||
|
"deadline",
|
||||||
|
&SearchOptions::new(5).with_sources(["slack", "email"]),
|
||||||
|
);
|
||||||
|
|
||||||
|
// Re-rank by relevance, recency, source authority and activation, then drop
|
||||||
|
// low-confidence results — the pipeline ClawhdfBackend runs.
|
||||||
|
let careful = memory.search(
|
||||||
|
&query_embedding,
|
||||||
|
"user preferences",
|
||||||
|
&SearchOptions::new(5)
|
||||||
|
.with_rerank(ReRankConfig::default())
|
||||||
|
.with_confidence(ConfidenceConfig::default()),
|
||||||
|
);
|
||||||
|
```
|
||||||
|
|
||||||
|
### Signed Checkpoints
|
||||||
|
|
||||||
|
```rust
|
||||||
|
use clawhdf5_agent::signing;
|
||||||
|
|
||||||
|
// Once, somewhere safe: keep the secret key, publish the public key.
|
||||||
|
let key = signing::generate_key();
|
||||||
|
let public = key.verifying_key();
|
||||||
|
|
||||||
|
// Every checkpoint is signed from now on. The key is never written to disk;
|
||||||
|
// a signed store refuses to checkpoint without it.
|
||||||
|
memory.set_signing_key(key);
|
||||||
|
memory.flush_wal()?;
|
||||||
|
|
||||||
|
// Anyone holding the public key can check the file, e.g. after copying it.
|
||||||
|
let report = HDF5Memory::verify(std::path::Path::new("agent.h5"), &public)?;
|
||||||
|
assert!(report.is_valid());
|
||||||
|
// On a tampered file: report.changed_records lists the records that differ.
|
||||||
|
```
|
||||||
|
|
||||||
|
The signature covers every record (text, embedding as stored, channel,
|
||||||
|
timestamp, session, tags, deleted flag, activation), the store's settings,
|
||||||
|
its sessions and its knowledge graph — a change made with any tool is caught.
|
||||||
|
It covers checkpoints, not saves still in the WAL
|
||||||
|
(`report.wal_entries_unsigned` counts those). CLI: `clawhdf5-cli keygen`,
|
||||||
|
`--signing-key <file>` on writing commands, and `verify --public-key`.
|
||||||
|
Signing adds about 20% to a checkpoint and 32 bytes per record to the file
|
||||||
|
([BENCHMARKS.md § Signed checkpoints](BENCHMARKS.md#signed-checkpoints)).
|
||||||
|
|
||||||
### Knowledge Graph
|
### Knowledge Graph
|
||||||
|
|
||||||
```rust
|
```rust
|
||||||
@@ -309,8 +511,8 @@ let neighbors = kg.bfs_neighbors(alice, 2); // 2-hop neighborhood
|
|||||||
let activated = kg.spreading_activation(&[alice], 0.5, 0.01, 5);
|
let activated = kg.spreading_activation(&[alice], 0.5, 0.01, 5);
|
||||||
|
|
||||||
// Entity resolution — fuzzy matching
|
// Entity resolution — fuzzy matching
|
||||||
let resolved = kg.resolve_or_create("alice", "person", -1, 2);
|
let (id, created) = kg.resolve_or_create("alice", "person", -1, 2);
|
||||||
// Returns existing Alice entity (Levenshtein distance ≤ 2)
|
// id == alice, created == false: matched the existing entity (Levenshtein distance ≤ 2)
|
||||||
```
|
```
|
||||||
|
|
||||||
### Memory Consolidation
|
### Memory Consolidation
|
||||||
@@ -321,15 +523,19 @@ use clawhdf5_agent::consolidation::*;
|
|||||||
let config = ConsolidationConfig::default();
|
let config = ConsolidationConfig::default();
|
||||||
let mut engine = ConsolidationEngine::new(config);
|
let mut engine = ConsolidationEngine::new(config);
|
||||||
|
|
||||||
// Add memories — automatically scored for importance
|
let now = 1_700_000_000.0; // seconds since the epoch
|
||||||
engine.add_memory("User prefers dark mode", vec![0.1, 0.2, ...], MemorySource::User);
|
|
||||||
engine.add_memory("ok", vec![0.0, 0.0, ...], MemorySource::System);
|
// Add memories — automatically scored for importance.
|
||||||
|
// Elevated sources (System, …) go through a separate, explicit API.
|
||||||
|
let id = engine.add_memory("User prefers dark mode".into(), vec![0.1, 0.2, ...], UntrustedSource::User, now);
|
||||||
|
engine.add_trusted_memory("ok".into(), vec![0.0, 0.0, ...], TrustedSource::System, now);
|
||||||
|
|
||||||
// Access a memory (reactivates it)
|
// Access a memory (reactivates it)
|
||||||
engine.access_memory(0);
|
engine.access_memory(id, now);
|
||||||
|
|
||||||
// Run consolidation cycle
|
// Run consolidation cycle
|
||||||
let stats = engine.consolidate();
|
engine.consolidate(now);
|
||||||
|
let stats = engine.get_stats();
|
||||||
// Working memories promote to Episodic (if important enough)
|
// Working memories promote to Episodic (if important enough)
|
||||||
// Episodic memories promote to Semantic (if accessed enough)
|
// Episodic memories promote to Semantic (if accessed enough)
|
||||||
// Low-decay memories get evicted when tiers are full
|
// Low-decay memories get evicted when tiers are full
|
||||||
@@ -351,19 +557,25 @@ let ids = index.range_query(1700000000.0, 1700010800.0);
|
|||||||
let recent = index.latest(10);
|
let recent = index.latest(10);
|
||||||
```
|
```
|
||||||
|
|
||||||
### OpenClaw Integration
|
### Markdown Backend
|
||||||
|
|
||||||
|
`ClawhdfBackend` ingests Markdown by section and searches it with the full
|
||||||
|
pipeline. It is a library API — clawhdf5 is **not** an OpenClaw memory plugin
|
||||||
|
([docs/openclaw.md](docs/openclaw.md)). Sections stored this way carry no
|
||||||
|
embedding, so their search is keyword-only unless you save records with
|
||||||
|
vectors through `save_entry`.
|
||||||
|
|
||||||
```rust
|
```rust
|
||||||
use clawhdf5_agent::openclaw::*;
|
use clawhdf5_agent::openclaw::*;
|
||||||
|
|
||||||
// Create backend
|
// Create backend
|
||||||
let mut backend = ClawhdfBackend::create("memory.h5", "agent-1", 384)?;
|
let mut backend = ClawhdfBackend::create(std::path::Path::new("memory.h5"), 384)?;
|
||||||
|
|
||||||
// Ingest existing Markdown memory files
|
// Ingest existing Markdown memory files
|
||||||
let md = std::fs::read_to_string("MEMORY.md")?;
|
let md = std::fs::read_to_string("MEMORY.md")?;
|
||||||
let count = backend.ingest_markdown("MEMORY.md", &md)?;
|
let count = backend.ingest_markdown("MEMORY.md", &md)?;
|
||||||
|
|
||||||
// Search (uses full pipeline: RRF → re-rank → confidence filter)
|
// Search (full pipeline: weighted vector + BM25 fusion → re-rank → confidence filter)
|
||||||
let results = backend.search("user preferences", &query_embedding, 5);
|
let results = backend.search("user preferences", &query_embedding, 5);
|
||||||
|
|
||||||
// Export back to Markdown
|
// Export back to Markdown
|
||||||
@@ -375,22 +587,23 @@ let exported = backend.export_markdown("MEMORY.md")?;
|
|||||||
## Crate Map
|
## Crate Map
|
||||||
|
|
||||||
```
|
```
|
||||||
clawhdf5 workspace (16 crates, ~92K lines of Rust; plus libaec-sys, an
|
clawhdf5 workspace (16 crates, ~86K lines of Rust in src/, ~104K with tests
|
||||||
internal FFI bindings crate for the optional szip feature)
|
and benches; plus libaec-sys, an internal FFI bindings
|
||||||
|
crate for the optional szip feature)
|
||||||
│
|
│
|
||||||
├── Core HDF5
|
├── Core HDF5
|
||||||
│ ├── clawhdf5-format — Binary parser/writer (no_std), shared type definitions
|
│ ├── clawhdf5-format — Binary parser/writer (no_std-capable), shared type definitions
|
||||||
│ ├── clawhdf5-io — I/O abstraction (buffered, mmap, async)
|
│ ├── clawhdf5-io — I/O abstraction (file/memory readers; optional mmap, async, HSDS, MPI)
|
||||||
│ ├── clawhdf5-filters — Fast deflate path (zlib-ng); lz4/zstd/pcodec/szip filters live in clawhdf5-format
|
│ ├── clawhdf5-filters — Fast deflate path (zlib-ng); lz4/zstd/pcodec/szip filters live in clawhdf5-format
|
||||||
│ ├── clawhdf5-derive — Proc macros
|
│ ├── clawhdf5-derive — Proc macros
|
||||||
│ ├── clawhdf5 — High-level API
|
│ ├── clawhdf5 — High-level API
|
||||||
│ ├── clawhdf5-netcdf4 — NetCDF-4 support
|
│ ├── clawhdf5-netcdf4 — NetCDF-4 support
|
||||||
│ ├── clawhdf5-accel — SIMD (NEON, AVX2, AVX-512)
|
│ ├── clawhdf5-accel — SIMD (AVX2, NEON incl. SDOT int8; AVX-512 behind `avx512`)
|
||||||
│ └── clawhdf5-gpu — GPU compute (wgpu, hand-written WGSL compute shaders)
|
│ └── clawhdf5-gpu — GPU compute (wgpu, hand-written WGSL compute shaders)
|
||||||
│
|
│
|
||||||
├── Agent Memory
|
├── Agent Memory
|
||||||
│ ├── clawhdf5-agent — Memory engine (20.9K lines, 32 modules; WAL is CRC32-checked per entry)
|
│ ├── clawhdf5-agent — Memory engine (24.7K lines, 32 modules; chained-CRC WAL)
|
||||||
│ ├── clawhdf5-ann — HNSW approximate nearest neighbor (default backend; optional `parallel` feature)
|
│ ├── clawhdf5-ann — HNSW approximate nearest neighbor (default backend; f32 or int8 storage; `parallel` build)
|
||||||
│ ├── clawhdf5-migrate — SQLite → HDF5 migration
|
│ ├── clawhdf5-migrate — SQLite → HDF5 migration
|
||||||
│ ├── clawhdf5-android — Android JNI bridge
|
│ ├── clawhdf5-android — Android JNI bridge
|
||||||
│ └── clawhdf5-cli — CLI tool
|
│ └── clawhdf5-cli — CLI tool
|
||||||
@@ -411,10 +624,10 @@ ClawhDF5's agent memory design draws from 15+ recent papers:
|
|||||||
|
|
||||||
| Paper | Key Insight | ClawhDF5 Module |
|
| Paper | Key Insight | ClawhDF5 Module |
|
||||||
|-------|-------------|-----------------|
|
|-------|-------------|-----------------|
|
||||||
| **MemX** (2026) | RRF + multi-factor re-ranking | `hybrid`, `reranker` |
|
| **MemX** (2026) | Hybrid fusion + multi-factor re-ranking | `hybrid`, `reranker` |
|
||||||
| **Graph-Native Cognitive Memory** (2026) | Graph-structured belief revision | `knowledge` |
|
| **Graph-Native Cognitive Memory** (2026) | Graph-structured memory (weighted, timestamped relations; entity timelines) | `knowledge`, `temporal` |
|
||||||
| **CraniMem** (2026) | Bounded hippocampal memory | `consolidation` |
|
| **CraniMem** (2026) | Bounded hippocampal memory | `consolidation` |
|
||||||
| **D-MEM** (2026) | Reward prediction error gating | `consolidation` |
|
| **D-MEM** (2026) | Surprise-gated storage (implemented as a novelty score) | `consolidation` |
|
||||||
| **SYNAPSE** (2025) | Spreading activation for recall | `knowledge` |
|
| **SYNAPSE** (2025) | Spreading activation for recall | `knowledge` |
|
||||||
| **RAGdb** (2025) | Zero-dependency edge RAG | Architecture |
|
| **RAGdb** (2025) | Zero-dependency edge RAG | Architecture |
|
||||||
| **MemoryGraft** (2025) | Memory poisoning attacks | `anomaly`, `provenance` |
|
| **MemoryGraft** (2025) | Memory poisoning attacks | `anomaly`, `provenance` |
|
||||||
@@ -429,16 +642,45 @@ ClawhDF5's agent memory design draws from 15+ recent papers:
|
|||||||
|
|
||||||
| Flag | Default | Description |
|
| Flag | Default | Description |
|
||||||
|------|---------|-------------|
|
|------|---------|-------------|
|
||||||
| `agent` | no | Full agent memory layer |
|
| `float16` | **yes** | Half-precision cosine kernel (`cosine_similarity_f16`). Half-precision *storage* is the `MemoryConfig::float16` setting below, and needs no feature |
|
||||||
| `float16` | **yes** | Half-precision embedding storage (2× compression) |
|
|
||||||
| `hnsw` | **yes** | HNSW approximate vector index for `hybrid_search` (via `clawhdf5-ann`); disable for an exact linear scan |
|
| `hnsw` | **yes** | HNSW approximate vector index for `hybrid_search` (via `clawhdf5-ann`); disable for an exact linear scan |
|
||||||
| `parallel` | no | Rayon parallel search |
|
| `parallel` | **yes** | Parallel HNSW bulk build (same graph, ~3× faster on 16 cores) and Rayon brute-force search strategies |
|
||||||
|
| `zstd` | no | Compress embeddings with Zstd instead of deflate when `MemoryConfig::compression` is on (links libzstd) |
|
||||||
| `fast-math` | no | BLAS matrix-vector multiply |
|
| `fast-math` | no | BLAS matrix-vector multiply |
|
||||||
| `accelerate` | no | Apple Accelerate / AMX (macOS) |
|
| `accelerate` | no | Apple Accelerate / AMX (macOS) |
|
||||||
| `openblas` | no | OpenBLAS (Linux) |
|
| `openblas` | no | OpenBLAS (Linux) |
|
||||||
| `gpu` | no | GPU search via wgpu |
|
| `gpu` | no | GPU search via wgpu |
|
||||||
| `async` | no | Tokio async with background flush |
|
| `async` | no | Tokio async with background flush |
|
||||||
|
|
||||||
|
To opt out of the parallel build: `--no-default-features --features float16,hnsw`.
|
||||||
|
For an exact linear cosine scan instead of HNSW: `--no-default-features --features float16`.
|
||||||
|
|
||||||
|
`MemoryConfig::hnsw_m`, `hnsw_ef_construction` and `hnsw_ef_search` tune the
|
||||||
|
vector index (16 / 64 / scale-with-`k` by default) and are stored with the
|
||||||
|
file.
|
||||||
|
|
||||||
|
`MemoryConfig::quantized_index` (**on by default** for new stores) holds the
|
||||||
|
HNSW index's own copy of the embeddings as `i8`, roughly halving a loaded
|
||||||
|
store's memory (2.72x -> 1.74x the raw vectors at 100k x 384). Quantised
|
||||||
|
distances are approximate, so the query path re-scores the candidate pool
|
||||||
|
against the exact embeddings the store already holds, which keeps recall at the
|
||||||
|
`f32` index's level. It is also **faster**: 1.63x the queries per second at
|
||||||
|
equal recall on x86-64 (AVX2) and 1.18x on a Raspberry Pi 5 (NEON `SDOT`), with
|
||||||
|
index builds 1.8x and 2.3x faster respectively. Stores created before the
|
||||||
|
setting existed keep their `f32` index; opt out for new stores with
|
||||||
|
`quantized_index = false` or `clawhdf5-cli create --f32-index`. See
|
||||||
|
[BENCHMARKS.md § Quantising the index copy](BENCHMARKS.md#quantising-the-index-copy-quantized_index).
|
||||||
|
|
||||||
|
`MemoryConfig::float16` (**on by default** for new stores) stores the
|
||||||
|
embeddings on disk as IEEE half precision (numpy `float16`): at 100K × 384 the
|
||||||
|
file drops from 154 to 81 MiB, checkpoints and opens get faster, and on the
|
||||||
|
full LongMemEval haystack with real MiniLM embeddings every retrieval metric
|
||||||
|
matches `f32`. Embeddings are rounded as they are saved, so the store searches
|
||||||
|
the same before and after a reopen; values must lie within ±65504. Existing
|
||||||
|
stores keep their setting. Opt out with `float16 = false` or
|
||||||
|
`clawhdf5-cli create --f32` — e.g. for unnormalised vectors. See
|
||||||
|
[BENCHMARKS.md § float16 embedding storage](BENCHMARKS.md#float16-embedding-storage-memoryconfigfloat16).
|
||||||
|
|
||||||
### `clawhdf5-format`
|
### `clawhdf5-format`
|
||||||
|
|
||||||
| Flag | Default | Description |
|
| Flag | Default | Description |
|
||||||
@@ -447,26 +689,31 @@ ClawhDF5's agent memory design draws from 15+ recent papers:
|
|||||||
| `deflate` | yes | Deflate compression |
|
| `deflate` | yes | Deflate compression |
|
||||||
| `checksum` | yes | Jenkins lookup3 verification |
|
| `checksum` | yes | Jenkins lookup3 verification |
|
||||||
| `provenance` | yes | SHA-256 provenance attributes |
|
| `provenance` | yes | SHA-256 provenance attributes |
|
||||||
| `fast-deflate` | **yes** | zlib-ng backend for faster deflate |
|
| `zlib-rs` | **yes** | Pure-Rust deflate backend ([zlib-rs](https://github.com/trifectatechfoundation/zlib-rs)) |
|
||||||
| `system-zlib-decompress` | **yes** | Use the system zlib for decompression where available |
|
| `fast-deflate` | no | zlib-ng deflate backend instead (C; needs `cmake`). Overrides `zlib-rs` when both are on |
|
||||||
|
| `system-zlib-decompress` | **yes** | Use Apple's system libz for decompression (macOS only; no effect elsewhere) |
|
||||||
| `parallel` | no | Parallel chunk encoding + compression (rayon) |
|
| `parallel` | no | Parallel chunk encoding + compression (rayon) |
|
||||||
| `fast-checksum` | no | crc32fast-accelerated checksums |
|
| `fast-checksum` | no | crc32fast-accelerated checksums |
|
||||||
| `lz4` | no | LZ4 block compression filter (id 32004) |
|
| `lz4` | no | LZ4 block compression filter (id 32004) |
|
||||||
| `zstd` | no | Zstandard compression filter (id 32015) |
|
| `zstd` | no | Zstandard compression filter (id 32015) |
|
||||||
| `pcodec` | no | Pcodec lossless numerical codec (id 32023, via `pco` crate) |
|
| `pcodec` | no | Pcodec lossless numerical codec (via `pco` crate). Private, unregistered filter id 480: **only clawhdf5 can read these datasets** (h5py/libhdf5 cannot). Files from clawhdf5 <= 2.7.0 used id 32023, which is registered to Granular BitRound; they still read. |
|
||||||
| `system-zlib` / `zlib-rs` | no | Alternative zlib backends for deflate |
|
| `system-zlib` | no | System zlib backend for deflate (C) |
|
||||||
| `blake3_hash` | no | BLAKE3 content hashing for provenance |
|
| `blake3_hash` | no | BLAKE3 content hashing for provenance |
|
||||||
|
| `szip` | no | SZIP filter (id 4) via libaec (C, through the internal `libaec-sys` crate) |
|
||||||
|
|
||||||
### `clawhdf5-ann`
|
### `clawhdf5-ann`
|
||||||
|
|
||||||
| Flag | Default | Description |
|
| Flag | Default | Description |
|
||||||
|------|---------|-------------|
|
|------|---------|-------------|
|
||||||
| `parallel` | no | Rayon-parallel neighbor-distance computation during HNSW graph pruning |
|
| `parallel` | no | Batched bulk build runs neighbour planning and back-link pruning on a Rayon pool; the graph is identical with or without it (enabled by `clawhdf5-agent`'s default `parallel`) |
|
||||||
|
|
||||||
### `clawhdf5-io`
|
### `clawhdf5-io`
|
||||||
|
|
||||||
| Flag | Default | Description |
|
| Flag | Default | Description |
|
||||||
|------|---------|-------------|
|
|------|---------|-------------|
|
||||||
|
| `mmap` | no | Memory-mapped reads (`memmap2`) |
|
||||||
|
| `async` | no | Tokio-based async I/O |
|
||||||
|
| `hsds` | no | HSDS (HDF REST service) client |
|
||||||
| `mpi-io` | no | MPI-backed I/O via the `mpi` crate |
|
| `mpi-io` | no | MPI-backed I/O via the `mpi` crate |
|
||||||
|
|
||||||
> **Parallel I/O (MPI) limitation:** `mpi-io`'s read path is a root-rank read
|
> **Parallel I/O (MPI) limitation:** `mpi-io`'s read path is a root-rank read
|
||||||
@@ -480,18 +727,26 @@ ClawhDF5's agent memory design draws from 15+ recent papers:
|
|||||||
## Building
|
## Building
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
# Default
|
# Default (pure Rust: no cmake or C compiler needed)
|
||||||
cargo build --workspace
|
cargo build --workspace
|
||||||
|
|
||||||
# Agent memory with all accelerations (Linux)
|
# Agent memory with all accelerations (Linux)
|
||||||
cargo build -p clawhdf5-agent --features "agent,float16,parallel,fast-math"
|
cargo build -p clawhdf5-agent --features fast-math
|
||||||
|
|
||||||
# Agent memory with Apple Accelerate (macOS)
|
# Agent memory with Apple Accelerate (macOS)
|
||||||
cargo build -p clawhdf5-agent --features "agent,float16,accelerate,parallel,gpu"
|
cargo build -p clawhdf5-agent --features "accelerate,gpu"
|
||||||
|
|
||||||
# Tests
|
# Tests
|
||||||
cargo test --workspace # all 1,650+ tests
|
cargo test --workspace # all 1,850+ tests
|
||||||
cargo test -p clawhdf5-agent # agent memory tests
|
cargo test -p clawhdf5-agent # agent memory tests
|
||||||
|
scripts/ci-test.sh # what CI runs: fmt, clippy matrix, tests,
|
||||||
|
# h5py/netCDF4 interop, no_std
|
||||||
|
|
||||||
|
# The interop suites need a Python with h5py; on a PEP 668 system that has to
|
||||||
|
# be a virtualenv. `ci-test.sh` finds `.venv` on its own, or set
|
||||||
|
# CLAWHDF5_PYTHON. Without one they skip — set CLAWHDF5_REQUIRE_INTEROP=1 to
|
||||||
|
# make that a failure instead.
|
||||||
|
python3 -m venv .venv && .venv/bin/pip install h5py numpy netCDF4 xarray
|
||||||
|
|
||||||
# Benchmarks
|
# Benchmarks
|
||||||
cargo bench -p clawhdf5-agent # agent memory suite
|
cargo bench -p clawhdf5-agent # agent memory suite
|
||||||
@@ -504,25 +759,42 @@ cargo bench -p clawhdf5-bench # h5bench-equivalent I/O suite
|
|||||||
|
|
||||||
```
|
```
|
||||||
agent_memory.h5
|
agent_memory.h5
|
||||||
├── /meta
|
├── /meta (attributes)
|
||||||
│ ├── schema_version: "1.0"
|
│ ├── schema_version: "1.0", edgehdf5_version
|
||||||
│ ├── agent_id, embedder, embedding_dim
|
│ ├── agent_id, embedder, embedding_dim, chunk_size, overlap, created_at
|
||||||
│ └── created_at
|
│ ├── float16, compression, compression_level, compact_threshold,
|
||||||
|
│ │ hebbian_boost, decay_factor, wal_enabled, wal_max_entries
|
||||||
|
│ ├── quantized_index, hnsw_m, hnsw_ef_construction, hnsw_ef_search
|
||||||
|
│ ├── wal_applied_len, wal_applied_crc (WAL mark of the last checkpoint)
|
||||||
|
│ └── ann_generation (ties the .ann sidecar to this checkpoint)
|
||||||
├── /memory
|
├── /memory
|
||||||
│ ├── chunks: string[N]
|
│ ├── chunks: string[N]
|
||||||
│ ├── embeddings: f32[N × D] (or f16 with float16 flag)
|
│ ├── embeddings: f32[N × D], or f16 for a `float16` store
|
||||||
│ ├── tombstones: u8[N]
|
│ │ (chunked; deflate, or Zstd with the `zstd`
|
||||||
│ └── norms: f32[N] (pre-computed L2)
|
│ │ feature, when compression is on)
|
||||||
|
│ ├── source_channel: string[N]
|
||||||
|
│ ├── timestamps: f64[N]
|
||||||
|
│ ├── session_ids: string[N]
|
||||||
|
│ ├── tags: string[N]
|
||||||
|
│ ├── tombstones: u8[N]
|
||||||
|
│ ├── norms: f32[N] (pre-computed L2)
|
||||||
|
│ └── activation_weights: f32[N] (Hebbian)
|
||||||
├── /sessions
|
├── /sessions
|
||||||
│ ├── ids: string[S]
|
│ ├── ids, channels, summaries: string[S]
|
||||||
│ └── summaries: string[S]
|
│ ├── start_idxs, end_idxs: i64[S]
|
||||||
|
│ └── timestamps: f64[S]
|
||||||
└── /knowledge_graph
|
└── /knowledge_graph
|
||||||
├── entity_names: string[E]
|
├── entity_ids, entity_emb_idxs: i64[E]; entity_names, entity_types: string[E]
|
||||||
├── relation_srcs: i64[R]
|
├── relation_srcs, relation_tgts: i64[R]; relation_types: string[R]
|
||||||
├── relation_tgts: i64[R]
|
├── relation_weights: f32[R]; relation_ts: f64[R]
|
||||||
└── relation_types: string[R]
|
└── alias_strings: string[A]; alias_entity_ids: i64[A] (when aliases exist)
|
||||||
```
|
```
|
||||||
|
|
||||||
|
Alongside the store: `<store>.h5.wal` (write-ahead log), `<store>.h5.ann`
|
||||||
|
(HNSW graph; derived, safe to delete) and `<store>.h5.lock` (single-writer
|
||||||
|
lock). A second writer gets `MemoryError::Locked`; use
|
||||||
|
`HDF5Memory::open_read_only` for a lock-free point-in-time view.
|
||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
## Migration
|
## Migration
|
||||||
@@ -541,9 +813,39 @@ Replace in `Cargo.toml` and source:
|
|||||||
|
|
||||||
```bash
|
```bash
|
||||||
cargo install --path crates/clawhdf5-migrate
|
cargo install --path crates/clawhdf5-migrate
|
||||||
clawhdf5-migrate --sqlite old.db --hdf5 memory.h5 --agent-id my-agent --embedding-dim 384
|
clawhdf5-migrate --sqlite old.db --hdf5 memory.h5 --agent-id my-agent --embedder minilm
|
||||||
```
|
```
|
||||||
|
|
||||||
|
The output is an ordinary `clawhdf5-agent` store, written through the agent's
|
||||||
|
own API: open it with `HDF5Memory::open` (or `clawhdf5-cli --path memory.h5 …`)
|
||||||
|
and search it straight away. The source must use the `memory_chunks` / `sessions` / `entities` / `relations` layout (names are
|
||||||
|
configurable with `--*-table`); note that this is not ZeroClaw's schema, and
|
||||||
|
ZeroClaw does not use clawhdf5. What carries over:
|
||||||
|
|
||||||
|
| SQLite | Agent store |
|
||||||
|
|--------|-------------|
|
||||||
|
| `memory_chunks` | memory records (text, embedding, source channel, timestamp, session id, tags); rows with `deleted = 1` become deleted records, or are left out with `--skip-deleted` |
|
||||||
|
| `sessions` | sessions (id, start/end index, channel, summary, timestamp) |
|
||||||
|
| `entities`, `relations` | knowledge graph entities and relations; entities get new ids and relations are re-pointed at them |
|
||||||
|
|
||||||
|
The chunk `id` column has no counterpart in the agent store, so records are
|
||||||
|
written in `id` order and numbered from 0. Embeddings are stored as float16
|
||||||
|
like any new store; `--f32` keeps full precision (and is required for values
|
||||||
|
beyond ±65504). The embedding dimension is detected from the first row unless
|
||||||
|
`--embedding-dim` is given, and every row must have it: a row of another length
|
||||||
|
is an error, never truncated or padded. A source with no memory records (only
|
||||||
|
sessions or the graph) needs `--embedding-dim`, since a store's dimension is
|
||||||
|
fixed when it is created. Every row is checked before the output is created,
|
||||||
|
so a source that cannot be migrated leaves an existing store at `--hdf5` as it
|
||||||
|
was. `--incremental` adds to an existing store only the rows it does not
|
||||||
|
already hold; the source must have the store's dimension, and records already
|
||||||
|
in the store take the source's deleted flag (a row deleted in SQLite since the
|
||||||
|
last run is deleted in the store; one un-deleted there is written again, as
|
||||||
|
the agent has no un-delete). The tool reads the result back with
|
||||||
|
`HDF5Memory::open_read_only`, compares it with the source (every row with
|
||||||
|
`--validate-full`) and checks that a migrated record is found by search;
|
||||||
|
`--dry-run` only counts the rows.
|
||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
## Roadmap
|
## Roadmap
|
||||||
@@ -557,10 +859,10 @@ See [ROADMAP.md](ROADMAP.md) for the full implementation tracker.
|
|||||||
- ✅ Temporal reasoning with sub-µs queries
|
- ✅ Temporal reasoning with sub-µs queries
|
||||||
- ✅ Memory security + anomaly detection
|
- ✅ Memory security + anomaly detection
|
||||||
- ✅ Multi-modal memory (text/image/audio/video)
|
- ✅ Multi-modal memory (text/image/audio/video)
|
||||||
- ✅ OpenClaw integration layer
|
- ✅ Markdown ingest/export backend (`ClawhdfBackend`); an OpenClaw plugin was never built — see [docs/openclaw.md](docs/openclaw.md)
|
||||||
- ✅ Comprehensive Criterion benchmarks
|
- ✅ Comprehensive Criterion benchmarks
|
||||||
|
|
||||||
**Phase 2** — MemoryArena and LongMemEval academic benchmarks are done (see [BENCHMARKS.md](BENCHMARKS.md), reproduced on a second machine); remaining: publish the OpenClaw TypeScript bridge to npm, crates.io/PyPI publishing.
|
**Phase 2** — MemoryArena and LongMemEval academic benchmarks are done (see [BENCHMARKS.md](BENCHMARKS.md), reproduced on a second machine); remaining: crates.io/PyPI publishing. The Node bindings are unpublished and known to be broken ([known issues](docs/known-issues.md)).
|
||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
@@ -577,6 +879,6 @@ MIT
|
|||||||
---
|
---
|
||||||
|
|
||||||
<p align="center">
|
<p align="center">
|
||||||
<em>Built by <a href="https://github.com/redclawsystems">RedClaw Systems</a></em><br>
|
<em>Built by <a href="https://git.redclaw.dev/quantumclaw">RedClaw Systems</a></em><br>
|
||||||
<em>~92,000 lines of Rust. Zero C dependencies. One file to remember everything.</em>
|
<em>~86,000 lines of Rust. Zero C dependencies. One file to remember everything.</em>
|
||||||
</p>
|
</p>
|
||||||
|
|||||||
+14
-8
@@ -105,24 +105,30 @@
|
|||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
## Track 7: OpenClaw Integration
|
## Track 7: OpenClaw Integration — withdrawn (2026-09-25)
|
||||||
**Status:** 🟢 Complete
|
**Status:** ⚪ Withdrawn (the items below were library work; no OpenClaw integration shipped)
|
||||||
**Priority:** Critical (for adoption)
|
**Priority:** Critical (for adoption)
|
||||||
**Crates:** `clawhdf5-agent`, `clawhdf5-napi`
|
**Crates:** `clawhdf5-agent`, `clawhdf5-napi`
|
||||||
|
|
||||||
- [x] **7.1** Memory backend trait — MemoryBackend with search/get/write/ingest/export/stats
|
- [x] **7.1** Memory backend trait — MemoryBackend with search/get/write/ingest/export/stats
|
||||||
- [x] **7.2** Hybrid retrieval pipeline — ClawhdfBackend wires RRF → reranker → confidence rejection
|
- [x] **7.2** Hybrid retrieval pipeline — ClawhdfBackend wires RRF → reranker → confidence rejection
|
||||||
- [x] **7.3** Markdown import/export — MarkdownParser + MarkdownExporter with line tracking + metadata
|
- [x] **7.3** Markdown import/export — MarkdownParser + MarkdownExporter with line tracking + metadata
|
||||||
- [x] **7.4** memory_search tool — backed by full hybrid retrieval pipeline
|
- [x] **7.4** `search()` — backed by the full hybrid retrieval pipeline (a Rust method; no OpenClaw tool was ever registered)
|
||||||
- [x] **7.5** memory_get tool — get() with path + line range support
|
- [x] **7.5** `get()` — read back by path, with a line slice (not an OpenClaw tool either)
|
||||||
- [x] **7.6** Compaction integration — run_compaction() (decay + compact + WAL flush), run_consolidation() (hippocampal engine), tick_session(), flush_wal()
|
- [x] **7.6** Compaction integration — run_compaction() (decay + compact + WAL flush), run_consolidation() (hippocampal engine), tick_session(), flush_wal()
|
||||||
- [x] **7.7** Config surface — `memory.backend = "clawhdf5"` schema documented in docs/openclaw-config.md
|
- [ ] **7.7** ~~Config surface — `memory.backend = "clawhdf5"`~~ — never valid OpenClaw config; docs removed
|
||||||
- [x] **7.8** Documentation + migration guide — docs/migration-guide.md, docs/openclaw-integration.md (architecture, full API reference, code patterns)
|
- [ ] **7.8** ~~Documentation + migration guide~~ — removed: they described an integration that never worked
|
||||||
|
|
||||||
**Node.js bridge:** `clawhdf5-napi` (napi-rs) → `@redclaw/clawhdf5` npm package with full TypeScript types.
|
**Node.js bridge:** `clawhdf5-napi` (napi-rs) and a TypeScript wrapper in `packages/clawhdf5-node` exist but are unpublished, untested in CI and known to be broken (docs/known-issues.md).
|
||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
|
> **Withdrawn.** None of this track produced a working OpenClaw integration: no
|
||||||
|
> plugin was built, the documented `memory.backend = "clawhdf5"` config was never
|
||||||
|
> valid in any OpenClaw release, and the Node package was never published. The
|
||||||
|
> Rust `ClawhdfBackend` remains as a library API. Not pursued for now; see
|
||||||
|
> [docs/openclaw.md](docs/openclaw.md) for what a plugin would need today.
|
||||||
|
|
||||||
## Track 8: Benchmarking & Validation
|
## Track 8: Benchmarking & Validation
|
||||||
**Status:** 🟢 Complete
|
**Status:** 🟢 Complete
|
||||||
**Priority:** High
|
**Priority:** High
|
||||||
@@ -142,7 +148,7 @@
|
|||||||
|
|
||||||
**Phase 1:** ~~Tracks 1, 2, 3 — core memory intelligence~~ 🟢 Complete
|
**Phase 1:** ~~Tracks 1, 2, 3 — core memory intelligence~~ 🟢 Complete
|
||||||
**Phase 2:** ~~Track 4 (temporal) + Track 5 (security)~~ 🟢 Complete
|
**Phase 2:** ~~Track 4 (temporal) + Track 5 (security)~~ 🟢 Complete
|
||||||
**Phase 3:** ~~Track 6 (multi-modal) + Track 7 (OpenClaw integration)~~ 🟢 Complete
|
**Phase 3:** ~~Track 6 (multi-modal)~~ 🟢 Complete; Track 7 (OpenClaw integration) withdrawn
|
||||||
**Phase 4:** ~~Track 8 (benchmarking + validation)~~ 🟢 Complete
|
**Phase 4:** ~~Track 8 (benchmarking + validation)~~ 🟢 Complete
|
||||||
|
|
||||||
All 8 tracks delivered. 1,650+ tests passing, zero clippy warnings.
|
All 8 tracks delivered. 1,650+ tests passing, zero clippy warnings.
|
||||||
|
|||||||
@@ -1,7 +1,8 @@
|
|||||||
[package]
|
[package]
|
||||||
name = "clawhdf5-accel"
|
name = "clawhdf5-accel"
|
||||||
version = "2.4.0"
|
version = "2.7.0"
|
||||||
edition = "2024"
|
edition = "2024"
|
||||||
|
rust-version.workspace = true
|
||||||
description = "SIMD-accelerated operations for rustyhdf5"
|
description = "SIMD-accelerated operations for rustyhdf5"
|
||||||
license = "MIT"
|
license = "MIT"
|
||||||
repository = "https://git.redclaw.dev/quantumclaw/clawhdf5"
|
repository = "https://git.redclaw.dev/quantumclaw/clawhdf5"
|
||||||
|
|||||||
@@ -25,6 +25,55 @@ unsafe fn hsum_256(v: __m256) -> f32 {
|
|||||||
_mm_cvtss_f32(result)
|
_mm_cvtss_f32(result)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// AVX2 dot product of two `i8` slices, widened to `i32`.
|
||||||
|
///
|
||||||
|
/// Each 16-byte half is sign-extended to sixteen `i16` lanes and multiplied
|
||||||
|
/// pairwise with `madd_epi16`, which sums adjacent products straight into
|
||||||
|
/// eight `i32` lanes — the widening that an autovectorised scalar loop does
|
||||||
|
/// in several shuffles is one instruction here. A pair sum is at most
|
||||||
|
/// `2 * 127 * 127`, far inside `i32`.
|
||||||
|
///
|
||||||
|
/// # Safety
|
||||||
|
/// Caller must verify is_x86_feature_detected!("avx2").
|
||||||
|
// SAFETY: Caller must have verified AVX2 via is_x86_feature_detected!.
|
||||||
|
#[target_feature(enable = "avx2")]
|
||||||
|
pub unsafe fn dot_i8(a: &[i8], b: &[i8]) -> i32 {
|
||||||
|
// SAFETY: Caller guarantees AVX2 is available per the # Safety contract;
|
||||||
|
// every load reads 32 bytes at an index checked against `len` first.
|
||||||
|
unsafe {
|
||||||
|
assert_eq!(a.len(), b.len());
|
||||||
|
let len = a.len();
|
||||||
|
let mut i = 0;
|
||||||
|
let mut acc0 = _mm256_setzero_si256();
|
||||||
|
let mut acc1 = _mm256_setzero_si256();
|
||||||
|
|
||||||
|
while i + 32 <= len {
|
||||||
|
let va = _mm256_loadu_si256(a.as_ptr().add(i).cast());
|
||||||
|
let vb = _mm256_loadu_si256(b.as_ptr().add(i).cast());
|
||||||
|
let a_lo = _mm256_cvtepi8_epi16(_mm256_castsi256_si128(va));
|
||||||
|
let b_lo = _mm256_cvtepi8_epi16(_mm256_castsi256_si128(vb));
|
||||||
|
let a_hi = _mm256_cvtepi8_epi16(_mm256_extracti128_si256(va, 1));
|
||||||
|
let b_hi = _mm256_cvtepi8_epi16(_mm256_extracti128_si256(vb, 1));
|
||||||
|
acc0 = _mm256_add_epi32(acc0, _mm256_madd_epi16(a_lo, b_lo));
|
||||||
|
acc1 = _mm256_add_epi32(acc1, _mm256_madd_epi16(a_hi, b_hi));
|
||||||
|
i += 32;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Horizontal sum of the eight i32 lanes.
|
||||||
|
let v = _mm256_add_epi32(acc0, acc1);
|
||||||
|
let s128 = _mm_add_epi32(_mm256_castsi256_si128(v), _mm256_extracti128_si256(v, 1));
|
||||||
|
let s64 = _mm_add_epi32(s128, _mm_unpackhi_epi64(s128, s128));
|
||||||
|
let s32 = _mm_add_epi32(s64, _mm_shuffle_epi32(s64, 0b01));
|
||||||
|
let mut sum = _mm_cvtsi128_si32(s32);
|
||||||
|
|
||||||
|
while i < len {
|
||||||
|
sum += i32::from(a[i]) * i32::from(b[i]);
|
||||||
|
i += 1;
|
||||||
|
}
|
||||||
|
sum
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
/// AVX2 dot product for f32 slices.
|
/// AVX2 dot product for f32 slices.
|
||||||
///
|
///
|
||||||
/// # Safety
|
/// # Safety
|
||||||
|
|||||||
@@ -122,6 +122,36 @@ pub fn dot_product(a: &[f32], b: &[f32]) -> f32 {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Dot product of two `i8` slices, widened to `i32`.
|
||||||
|
///
|
||||||
|
/// The kernel behind int8-quantised vector search. On x86-64 it uses the AVX2
|
||||||
|
/// path whenever AVX2 is present (including on AVX-512 machines, where it is
|
||||||
|
/// what the f32 kernels use too on a default build). On aarch64 it uses the
|
||||||
|
/// ARMv8.2 `SDOT` instruction when the CPU has the dot-product extension, and
|
||||||
|
/// plain NEON otherwise.
|
||||||
|
pub fn dot_i8(a: &[i8], b: &[i8]) -> i32 {
|
||||||
|
match detect_backend() {
|
||||||
|
#[cfg(target_arch = "aarch64")]
|
||||||
|
Backend::Neon => {
|
||||||
|
if std::arch::is_aarch64_feature_detected!("dotprod") {
|
||||||
|
// SAFETY: the dotprod extension was just detected at runtime.
|
||||||
|
unsafe { neon::dot_i8_dotprod(a, b) }
|
||||||
|
} else {
|
||||||
|
// SAFETY: NEON is always available on aarch64.
|
||||||
|
unsafe { neon::dot_i8(a, b) }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(target_arch = "x86_64")]
|
||||||
|
// SAFETY: both variants imply AVX2 was detected at runtime (the
|
||||||
|
// AVX-512 backend is only selected on CPUs that also have AVX2).
|
||||||
|
Backend::Avx2 | Backend::Avx512 if is_x86_feature_detected!("avx2") => unsafe {
|
||||||
|
avx2::dot_i8(a, b)
|
||||||
|
},
|
||||||
|
_ => scalar::dot_i8(a, b),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
/// Compute the L2 norm (magnitude) of a vector.
|
/// Compute the L2 norm (magnitude) of a vector.
|
||||||
pub fn vector_norm(v: &[f32]) -> f32 {
|
pub fn vector_norm(v: &[f32]) -> f32 {
|
||||||
dot_product(v, v).sqrt()
|
dot_product(v, v).sqrt()
|
||||||
@@ -713,3 +743,78 @@ mod tests {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod dot_i8_tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
fn codes(n: usize, seed: u64) -> Vec<i8> {
|
||||||
|
let mut state = seed;
|
||||||
|
(0..n)
|
||||||
|
.map(|_| {
|
||||||
|
state = state
|
||||||
|
.wrapping_mul(6_364_136_223_846_793_005)
|
||||||
|
.wrapping_add(1_442_695_040_888_963_407);
|
||||||
|
// Full range, including the extremes.
|
||||||
|
((state >> 56) as u8) as i8
|
||||||
|
})
|
||||||
|
.collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn dispatched_kernel_matches_scalar_exactly() {
|
||||||
|
// Integer arithmetic: the SIMD path must agree bit for bit, at every
|
||||||
|
// length — including ones that are not multiples of the 32-byte block,
|
||||||
|
// which exercise the tail.
|
||||||
|
for len in [0, 1, 7, 31, 32, 33, 63, 64, 100, 384, 385, 1536] {
|
||||||
|
let a = codes(len, 1 + len as u64);
|
||||||
|
let b = codes(len, 1000 + len as u64);
|
||||||
|
assert_eq!(dot_i8(&a, &b), scalar::dot_i8(&a, &b), "len {len}");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Dispatch only ever takes one path on a given CPU, so on a machine with
|
||||||
|
/// the dot-product extension the plain-NEON kernel would otherwise go
|
||||||
|
/// untested. Check each aarch64 kernel against scalar directly.
|
||||||
|
#[cfg(target_arch = "aarch64")]
|
||||||
|
#[test]
|
||||||
|
fn every_aarch64_kernel_matches_scalar_exactly() {
|
||||||
|
for len in [0, 1, 7, 15, 16, 17, 31, 32, 33, 63, 64, 100, 384, 385, 1536] {
|
||||||
|
let a = codes(len, 7 + len as u64);
|
||||||
|
let b = codes(len, 7000 + len as u64);
|
||||||
|
let want = scalar::dot_i8(&a, &b);
|
||||||
|
// SAFETY: NEON is always available on aarch64.
|
||||||
|
assert_eq!(unsafe { neon::dot_i8(&a, &b) }, want, "neon, len {len}");
|
||||||
|
if std::arch::is_aarch64_feature_detected!("dotprod") {
|
||||||
|
// SAFETY: the dotprod extension was just detected.
|
||||||
|
assert_eq!(
|
||||||
|
unsafe { neon::dot_i8_dotprod(&a, &b) },
|
||||||
|
want,
|
||||||
|
"dotprod, len {len}"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
// The extremes, through both kernels.
|
||||||
|
let lo = vec![-128i8; 4096];
|
||||||
|
let hi = vec![127i8; 4096];
|
||||||
|
// SAFETY: NEON is always available on aarch64.
|
||||||
|
assert_eq!(unsafe { neon::dot_i8(&lo, &lo) }, 4096 * 128 * 128);
|
||||||
|
// SAFETY: NEON is always available on aarch64.
|
||||||
|
assert_eq!(unsafe { neon::dot_i8(&lo, &hi) }, -4096 * 128 * 127);
|
||||||
|
if std::arch::is_aarch64_feature_detected!("dotprod") {
|
||||||
|
// SAFETY: the dotprod extension was just detected.
|
||||||
|
assert_eq!(unsafe { neon::dot_i8_dotprod(&lo, &lo) }, 4096 * 128 * 128);
|
||||||
|
// SAFETY: the dotprod extension was just detected.
|
||||||
|
assert_eq!(unsafe { neon::dot_i8_dotprod(&lo, &hi) }, -4096 * 128 * 127);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn extremes_do_not_overflow() {
|
||||||
|
// -128 * -128 is the largest product; a long run of it must still fit.
|
||||||
|
let a = vec![-128i8; 4096];
|
||||||
|
assert_eq!(dot_i8(&a, &a), 4096 * 128 * 128);
|
||||||
|
let b = vec![127i8; 4096];
|
||||||
|
assert_eq!(dot_i8(&a, &b), -4096 * 128 * 127);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -180,3 +180,130 @@ pub fn checksum_fletcher32(data: &[u8]) -> u32 {
|
|||||||
|
|
||||||
(sum2 << 16) | sum1
|
(sum2 << 16) | sum1
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// NEON dot product of two `i8` slices, widened to `i32`, for any aarch64 CPU.
|
||||||
|
///
|
||||||
|
/// `vmull_s8` multiplies eight lanes into `i16` — even `-128 * -128` is 16 384,
|
||||||
|
/// inside `i16` — and `vpadalq_s16` adds adjacent pairs of those into `i32`
|
||||||
|
/// accumulators, so nothing can overflow before the final horizontal sum.
|
||||||
|
///
|
||||||
|
/// CPUs with the ARMv8.2 dot-product extension should use
|
||||||
|
/// [`dot_i8_dotprod`], which does the multiply and the accumulate in one
|
||||||
|
/// instruction.
|
||||||
|
///
|
||||||
|
/// # Safety
|
||||||
|
/// Caller must ensure aarch64 target (NEON always available).
|
||||||
|
// SAFETY: NEON is always available on aarch64 targets; caller guarantees aarch64.
|
||||||
|
#[target_feature(enable = "neon")]
|
||||||
|
pub unsafe fn dot_i8(a: &[i8], b: &[i8]) -> i32 {
|
||||||
|
assert_eq!(a.len(), b.len());
|
||||||
|
let len = a.len();
|
||||||
|
let mut i = 0;
|
||||||
|
let mut acc0 = vdupq_n_s32(0);
|
||||||
|
let mut acc1 = vdupq_n_s32(0);
|
||||||
|
|
||||||
|
while i + 16 <= len {
|
||||||
|
// SAFETY: NEON is available per the # Safety contract, and both
|
||||||
|
// 16-byte loads start at an index checked against `len` above.
|
||||||
|
unsafe {
|
||||||
|
let va = vld1q_s8(a.as_ptr().add(i));
|
||||||
|
let vb = vld1q_s8(b.as_ptr().add(i));
|
||||||
|
acc0 = vpadalq_s16(acc0, vmull_s8(vget_low_s8(va), vget_low_s8(vb)));
|
||||||
|
acc1 = vpadalq_s16(acc1, vmull_high_s8(va, vb));
|
||||||
|
}
|
||||||
|
i += 16;
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut sum = vaddvq_s32(vaddq_s32(acc0, acc1));
|
||||||
|
while i < len {
|
||||||
|
sum += i32::from(a[i]) * i32::from(b[i]);
|
||||||
|
i += 1;
|
||||||
|
}
|
||||||
|
sum
|
||||||
|
}
|
||||||
|
|
||||||
|
/// One `SDOT`: for each of the four `i32` lanes of `acc`, add the dot
|
||||||
|
/// product of the corresponding four `i8` pairs from `a` and `b`.
|
||||||
|
///
|
||||||
|
/// Written as inline assembly because the `vdotq_s32` intrinsic is still
|
||||||
|
/// behind the unstable `stdarch_neon_dotprod` feature; inline assembly is
|
||||||
|
/// stable on aarch64.
|
||||||
|
///
|
||||||
|
/// # Safety
|
||||||
|
/// Caller must ensure the CPU supports the `dotprod` extension.
|
||||||
|
#[inline]
|
||||||
|
#[target_feature(enable = "neon,dotprod")]
|
||||||
|
unsafe fn sdot(acc: int32x4_t, a: int8x16_t, b: int8x16_t) -> int32x4_t {
|
||||||
|
let mut acc = acc;
|
||||||
|
// SAFETY: `dotprod` is enabled for this function and the caller
|
||||||
|
// guarantees the CPU supports it. The instruction reads only its three
|
||||||
|
// vector registers and touches no memory.
|
||||||
|
unsafe {
|
||||||
|
std::arch::asm!(
|
||||||
|
"sdot {acc:v}.4s, {a:v}.16b, {b:v}.16b",
|
||||||
|
acc = inout(vreg) acc,
|
||||||
|
a = in(vreg) a,
|
||||||
|
b = in(vreg) b,
|
||||||
|
options(pure, nomem, nostack),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
acc
|
||||||
|
}
|
||||||
|
|
||||||
|
/// NEON dot product of two `i8` slices using the ARMv8.2 dot-product
|
||||||
|
/// extension (`SDOT`): sixteen multiply-accumulates per instruction, straight
|
||||||
|
/// into `i32` lanes.
|
||||||
|
///
|
||||||
|
/// Present on the cores this crate actually runs on — Cortex-A76 and later
|
||||||
|
/// (Raspberry Pi 5, current Android phones), Neoverse-N1 (Graviton2, Ampere
|
||||||
|
/// Altra), and every Apple Silicon generation.
|
||||||
|
///
|
||||||
|
/// # Safety
|
||||||
|
/// Caller must verify `is_aarch64_feature_detected!("dotprod")`.
|
||||||
|
// SAFETY: caller has verified the dotprod extension at runtime.
|
||||||
|
#[target_feature(enable = "neon,dotprod")]
|
||||||
|
pub unsafe fn dot_i8_dotprod(a: &[i8], b: &[i8]) -> i32 {
|
||||||
|
assert_eq!(a.len(), b.len());
|
||||||
|
let len = a.len();
|
||||||
|
let mut i = 0;
|
||||||
|
let mut acc0 = vdupq_n_s32(0);
|
||||||
|
let mut acc1 = vdupq_n_s32(0);
|
||||||
|
|
||||||
|
// Two independent accumulators so consecutive SDOTs are not serialised on
|
||||||
|
// one register.
|
||||||
|
while i + 32 <= len {
|
||||||
|
// SAFETY: dotprod is available per the # Safety contract, and every
|
||||||
|
// 16-byte load starts at an index checked against `len` above.
|
||||||
|
unsafe {
|
||||||
|
acc0 = sdot(
|
||||||
|
acc0,
|
||||||
|
vld1q_s8(a.as_ptr().add(i)),
|
||||||
|
vld1q_s8(b.as_ptr().add(i)),
|
||||||
|
);
|
||||||
|
acc1 = sdot(
|
||||||
|
acc1,
|
||||||
|
vld1q_s8(a.as_ptr().add(i + 16)),
|
||||||
|
vld1q_s8(b.as_ptr().add(i + 16)),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
i += 32;
|
||||||
|
}
|
||||||
|
if i + 16 <= len {
|
||||||
|
// SAFETY: as above; the load is bounds-checked by this condition.
|
||||||
|
unsafe {
|
||||||
|
acc0 = sdot(
|
||||||
|
acc0,
|
||||||
|
vld1q_s8(a.as_ptr().add(i)),
|
||||||
|
vld1q_s8(b.as_ptr().add(i)),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
i += 16;
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut sum = vaddvq_s32(vaddq_s32(acc0, acc1));
|
||||||
|
while i < len {
|
||||||
|
sum += i32::from(a[i]) * i32::from(b[i]);
|
||||||
|
i += 1;
|
||||||
|
}
|
||||||
|
sum
|
||||||
|
}
|
||||||
|
|||||||
@@ -140,3 +140,33 @@ fn f16_to_f32_soft(h: u16) -> f32 {
|
|||||||
|
|
||||||
f32::from_bits(f32_bits)
|
f32::from_bits(f32_bits)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Dot product of two `i8` slices, widened to `i32`.
|
||||||
|
///
|
||||||
|
/// `dim` terms of at most `127 * 127` fit an `i32` for any realistic
|
||||||
|
/// dimension (over 130 000 terms before overflow is possible).
|
||||||
|
pub fn dot_i8(a: &[i8], b: &[i8]) -> i32 {
|
||||||
|
assert_eq!(a.len(), b.len());
|
||||||
|
// Four independent accumulators over 32-lane blocks: the widening product
|
||||||
|
// has to sit in a fixed-length chunk for the vectoriser to see it, and the
|
||||||
|
// separate accumulators keep it off one dependency chain.
|
||||||
|
const LANE: usize = 8;
|
||||||
|
let (a_blocks, a_tail) = a.as_chunks::<{ LANE * 4 }>();
|
||||||
|
let (b_blocks, b_tail) = b.as_chunks::<{ LANE * 4 }>();
|
||||||
|
let mut acc = [0i32; 4];
|
||||||
|
for (x, y) in a_blocks.iter().zip(b_blocks) {
|
||||||
|
for (lane, slot) in acc.iter_mut().enumerate() {
|
||||||
|
let mut sum = 0i32;
|
||||||
|
for k in 0..LANE {
|
||||||
|
sum += i32::from(x[lane * LANE + k]) * i32::from(y[lane * LANE + k]);
|
||||||
|
}
|
||||||
|
*slot += sum;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
let tail: i32 = a_tail
|
||||||
|
.iter()
|
||||||
|
.zip(b_tail)
|
||||||
|
.map(|(&x, &y)| i32::from(x) * i32::from(y))
|
||||||
|
.sum();
|
||||||
|
acc[0] + acc[1] + acc[2] + acc[3] + tail
|
||||||
|
}
|
||||||
|
|||||||
@@ -1,7 +1,8 @@
|
|||||||
[package]
|
[package]
|
||||||
name = "clawhdf5-agent"
|
name = "clawhdf5-agent"
|
||||||
version = "2.4.0"
|
version = "2.7.0"
|
||||||
edition = "2024"
|
edition = "2024"
|
||||||
|
rust-version.workspace = true
|
||||||
description = "HDF5-backed persistent memory store for on-device AI agents"
|
description = "HDF5-backed persistent memory store for on-device AI agents"
|
||||||
license = "MIT"
|
license = "MIT"
|
||||||
repository = "https://git.redclaw.dev/quantumclaw/clawhdf5"
|
repository = "https://git.redclaw.dev/quantumclaw/clawhdf5"
|
||||||
@@ -10,14 +11,18 @@ keywords = ["agent", "memory", "hdf5", "vector-search", "embedding"]
|
|||||||
categories = ["database", "science", "algorithms"]
|
categories = ["database", "science", "algorithms"]
|
||||||
|
|
||||||
[dependencies]
|
[dependencies]
|
||||||
clawhdf5-format = { path = "../clawhdf5-format", version = "2.4.0", features = ["parallel", "fast-checksum"] }
|
clawhdf5-format = { path = "../clawhdf5-format", version = "2.7.0", features = ["parallel", "fast-checksum"] }
|
||||||
clawhdf5 = { path = "../clawhdf5", version = "2.4.0" }
|
clawhdf5 = { path = "../clawhdf5", version = "2.7.0" }
|
||||||
clawhdf5-io = { path = "../clawhdf5-io", version = "2.4.0", features = ["mmap"] }
|
clawhdf5-io = { path = "../clawhdf5-io", version = "2.7.0", features = ["mmap"] }
|
||||||
clawhdf5-accel = { path = "../clawhdf5-accel", version = "2.4.0" }
|
clawhdf5-accel = { path = "../clawhdf5-accel", version = "2.7.0" }
|
||||||
clawhdf5-ann = { path = "../clawhdf5-ann", version = "2.4.0", optional = true }
|
clawhdf5-ann = { path = "../clawhdf5-ann", version = "2.7.0", optional = true }
|
||||||
clawhdf5-gpu = { path = "../clawhdf5-gpu", version = "2.4.0", optional = true, default-features = false }
|
clawhdf5-gpu = { path = "../clawhdf5-gpu", version = "2.7.0", optional = true, default-features = false }
|
||||||
serde = { workspace = true }
|
serde = { workspace = true }
|
||||||
byteorder = "1"
|
byteorder = "1"
|
||||||
|
# Signed checkpoints (MemoryConfig-independent; see `signing`). Pure Rust.
|
||||||
|
ed25519-dalek = { version = "2", features = ["rand_core"] }
|
||||||
|
sha2 = "0.10"
|
||||||
|
rand_core = { version = "0.6", features = ["getrandom"] }
|
||||||
half = { workspace = true, optional = true }
|
half = { workspace = true, optional = true }
|
||||||
rayon = { version = "1", optional = true }
|
rayon = { version = "1", optional = true }
|
||||||
matrixmultiply = { version = "0.3", optional = true }
|
matrixmultiply = { version = "0.3", optional = true }
|
||||||
@@ -44,10 +49,16 @@ harness = false
|
|||||||
name = "memory_bench"
|
name = "memory_bench"
|
||||||
harness = false
|
harness = false
|
||||||
|
|
||||||
|
[[bench]]
|
||||||
|
name = "multimodal_bench"
|
||||||
|
harness = false
|
||||||
|
|
||||||
[features]
|
[features]
|
||||||
default = ["float16", "hnsw"]
|
default = ["float16", "hnsw", "parallel"]
|
||||||
float16 = ["half"]
|
float16 = ["half"]
|
||||||
parallel = ["rayon"]
|
# Rayon-parallel brute-force search strategies, and a parallel bulk build of
|
||||||
|
# the HNSW index (same graph, several times faster on a multi-core machine).
|
||||||
|
parallel = ["rayon", "clawhdf5-ann?/parallel"]
|
||||||
# Compress embeddings with Zstd instead of deflate when
|
# Compress embeddings with Zstd instead of deflate when
|
||||||
# `MemoryConfig::compression` is on. Off by default: it links libzstd (C).
|
# `MemoryConfig::compression` is on. Off by default: it links libzstd (C).
|
||||||
zstd = ["clawhdf5/zstd"]
|
zstd = ["clawhdf5/zstd"]
|
||||||
@@ -57,7 +68,6 @@ zstd = ["clawhdf5/zstd"]
|
|||||||
# `--no-default-features` (plus re-enabling other defaults) to force the exact
|
# `--no-default-features` (plus re-enabling other defaults) to force the exact
|
||||||
# linear cosine scan.
|
# linear cosine scan.
|
||||||
hnsw = ["clawhdf5-ann"]
|
hnsw = ["clawhdf5-ann"]
|
||||||
agent = []
|
|
||||||
gpu = ["clawhdf5-gpu/gpu-wgpu"]
|
gpu = ["clawhdf5-gpu/gpu-wgpu"]
|
||||||
fast-math = ["matrixmultiply"]
|
fast-math = ["matrixmultiply"]
|
||||||
accelerate = ["accelerate-src", "cblas-sys"]
|
accelerate = ["accelerate-src", "cblas-sys"]
|
||||||
|
|||||||
@@ -0,0 +1,107 @@
|
|||||||
|
//! Multi-modal memory search benchmarks (`clawhdf5_agent::multimodal`).
|
||||||
|
//!
|
||||||
|
//! Covers `MultiModalStore::search_cross_modal` (every embedding of every
|
||||||
|
//! record, whatever its modality) and, for comparison,
|
||||||
|
//! `MultiModalStore::search_by_modality` restricted to one modality.
|
||||||
|
//!
|
||||||
|
//! Corpus: N records (1K and 10K), each carrying two 384-dim embeddings —
|
||||||
|
//! a text embedding of its caption plus one embedding of its primary modality,
|
||||||
|
//! cycling Image / Audio / Video — so a cross-modal query scores 2N vectors.
|
||||||
|
//! All data comes from a fixed-seed LCG, so every run sees the same corpus.
|
||||||
|
//!
|
||||||
|
//! Run: `cargo bench -p clawhdf5-agent --bench multimodal_bench`
|
||||||
|
|
||||||
|
use std::collections::HashMap;
|
||||||
|
|
||||||
|
use clawhdf5_agent::multimodal::{
|
||||||
|
MediaRef, ModalEmbedding, Modality, MultiModalRecord, MultiModalStore,
|
||||||
|
};
|
||||||
|
use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// Simple deterministic PRNG (LCG), same as the other agent benches
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
struct Rng(u32);
|
||||||
|
|
||||||
|
impl Rng {
|
||||||
|
fn new(seed: u32) -> Self {
|
||||||
|
Self(seed)
|
||||||
|
}
|
||||||
|
fn next_u32(&mut self) -> u32 {
|
||||||
|
self.0 = self.0.wrapping_mul(1103515245).wrapping_add(12345);
|
||||||
|
self.0 >> 16
|
||||||
|
}
|
||||||
|
fn next_f32(&mut self) -> f32 {
|
||||||
|
self.next_u32() as f32 / 65536.0 - 0.5
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn make_vec(rng: &mut Rng, dim: usize) -> Vec<f32> {
|
||||||
|
(0..dim).map(|_| rng.next_f32()).collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// Corpus
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
const DIM: usize = 384;
|
||||||
|
const K: usize = 10;
|
||||||
|
|
||||||
|
const MEDIA: [(Modality, &str, &str); 3] = [
|
||||||
|
(Modality::Image, "image/png", "clip-vit-base"),
|
||||||
|
(Modality::Audio, "audio/wav", "clap-base"),
|
||||||
|
(Modality::Video, "video/mp4", "xclip-base"),
|
||||||
|
];
|
||||||
|
|
||||||
|
fn build_store(n: usize, seed: u32) -> MultiModalStore {
|
||||||
|
let mut rng = Rng::new(seed);
|
||||||
|
let mut store = MultiModalStore::new();
|
||||||
|
for i in 0..n {
|
||||||
|
let (modality, mime, model) = &MEDIA[i % MEDIA.len()];
|
||||||
|
let embeddings = vec![
|
||||||
|
ModalEmbedding::new(Modality::Text, make_vec(&mut rng, DIM), "minilm-l6"),
|
||||||
|
ModalEmbedding::new(modality.clone(), make_vec(&mut rng, DIM), *model),
|
||||||
|
];
|
||||||
|
store.add_record(MultiModalRecord {
|
||||||
|
id: 0,
|
||||||
|
primary_modality: modality.clone(),
|
||||||
|
text_content: Some(format!("{modality} memory {i}")),
|
||||||
|
media_ref: Some(MediaRef::path(format!("/media/{i}"), *mime)),
|
||||||
|
embeddings,
|
||||||
|
observation: None,
|
||||||
|
timestamp: 1_700_000_000.0 + i as f64,
|
||||||
|
metadata: HashMap::new(),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
store
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// Benchmarks
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
fn multimodal_search_benches(c: &mut Criterion) {
|
||||||
|
let query = make_vec(&mut Rng::new(99), DIM);
|
||||||
|
|
||||||
|
let mut group = c.benchmark_group("multimodal_search");
|
||||||
|
group.sample_size(50);
|
||||||
|
|
||||||
|
for (label, n) in [("1k", 1_000usize), ("10k", 10_000)] {
|
||||||
|
let store = build_store(n, 42);
|
||||||
|
assert_eq!(store.count(), n);
|
||||||
|
|
||||||
|
group.bench_with_input(BenchmarkId::new("cross_modal", label), &n, |b, _| {
|
||||||
|
b.iter(|| store.search_cross_modal(&query, K));
|
||||||
|
});
|
||||||
|
|
||||||
|
group.bench_with_input(BenchmarkId::new("by_modality_image", label), &n, |b, _| {
|
||||||
|
b.iter(|| store.search_by_modality(&Modality::Image, &query, K));
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
group.finish();
|
||||||
|
}
|
||||||
|
|
||||||
|
criterion_group!(multimodal_benches, multimodal_search_benches);
|
||||||
|
criterion_main!(multimodal_benches);
|
||||||
@@ -118,6 +118,10 @@ mod tests {
|
|||||||
created_at: "2025-01-01T00:00:00Z".to_string(),
|
created_at: "2025-01-01T00:00:00Z".to_string(),
|
||||||
wal_enabled: false,
|
wal_enabled: false,
|
||||||
wal_max_entries: 500,
|
wal_max_entries: 500,
|
||||||
|
quantized_index: false,
|
||||||
|
hnsw_m: 16,
|
||||||
|
hnsw_ef_construction: 64,
|
||||||
|
hnsw_ef_search: 0,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -37,7 +37,7 @@
|
|||||||
//! let mem = AsyncHDF5Memory::open_with(path, config).await?;
|
//! let mem = AsyncHDF5Memory::open_with(path, config).await?;
|
||||||
//! mem.save(entry).await?; // buffered → background writer
|
//! mem.save(entry).await?; // buffered → background writer
|
||||||
//! mem.save_batch(entries).await?; // also buffered
|
//! mem.save_batch(entries).await?; // also buffered
|
||||||
//! let results = mem.hybrid_search(emb, "query".into(), 0.7, 0.3, 5).await;
|
//! let results = mem.hybrid_search(emb, "query".into(), 0.4, 0.6, 5).await;
|
||||||
//! mem.shutdown().await?; // final flush + stop
|
//! mem.shutdown().await?; // final flush + stop
|
||||||
//! ```
|
//! ```
|
||||||
|
|
||||||
|
|||||||
@@ -59,11 +59,18 @@ pub struct BM25Index {
|
|||||||
k1: f32,
|
k1: f32,
|
||||||
/// BM25 b parameter.
|
/// BM25 b parameter.
|
||||||
b: f32,
|
b: f32,
|
||||||
|
/// Applied to every document and query token, so the two always agree.
|
||||||
|
filter: TokenFilter,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl BM25Index {
|
impl BM25Index {
|
||||||
/// Build a BM25 index from a set of documents, excluding tombstoned entries.
|
/// Build a BM25 index from a set of documents, excluding tombstoned entries.
|
||||||
pub fn build(documents: &[String], tombstones: &[u8]) -> Self {
|
pub fn build(documents: &[String], tombstones: &[u8]) -> Self {
|
||||||
|
Self::build_with(documents, tombstones, TokenFilter::default())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// [`BM25Index::build`] with the token filter chosen explicitly.
|
||||||
|
pub fn build_with(documents: &[String], tombstones: &[u8], filter: TokenFilter) -> Self {
|
||||||
let mut index = Self {
|
let mut index = Self {
|
||||||
inverted: HashMap::new(),
|
inverted: HashMap::new(),
|
||||||
doc_lengths: vec![0; documents.len()],
|
doc_lengths: vec![0; documents.len()],
|
||||||
@@ -72,6 +79,7 @@ impl BM25Index {
|
|||||||
num_docs: 0,
|
num_docs: 0,
|
||||||
k1: DEFAULT_K1,
|
k1: DEFAULT_K1,
|
||||||
b: DEFAULT_B,
|
b: DEFAULT_B,
|
||||||
|
filter,
|
||||||
};
|
};
|
||||||
index.index_documents(documents, tombstones);
|
index.index_documents(documents, tombstones);
|
||||||
index
|
index
|
||||||
@@ -80,8 +88,11 @@ impl BM25Index {
|
|||||||
/// Search the index for a query, returning the top `k` results
|
/// Search the index for a query, returning the top `k` results
|
||||||
/// as `(doc_id, score)` pairs sorted by score descending.
|
/// as `(doc_id, score)` pairs sorted by score descending.
|
||||||
///
|
///
|
||||||
/// Uses Block-Max WAND for early termination when remaining documents
|
/// Scores every matching document exhaustively, then keeps the top `k`.
|
||||||
/// cannot beat the current top-k threshold.
|
/// There is no early termination (WAND, MaxScore): the store's hot path
|
||||||
|
/// is [`scores`](Self::scores), because score fusion normalises over the
|
||||||
|
/// whole matching set and so needs every score, which no pruning scheme
|
||||||
|
/// can skip. This method is for BM25-only callers.
|
||||||
pub fn search(&self, query: &str, k: usize) -> Vec<(usize, f32)> {
|
pub fn search(&self, query: &str, k: usize) -> Vec<(usize, f32)> {
|
||||||
if k == 0 {
|
if k == 0 {
|
||||||
return Vec::new();
|
return Vec::new();
|
||||||
@@ -120,7 +131,7 @@ impl BM25Index {
|
|||||||
// add/remove, and costs one `ln` per query term.
|
// add/remove, and costs one `ln` per query term.
|
||||||
let mut acc = vec![0.0f32; self.doc_lengths.len()];
|
let mut acc = vec![0.0f32; self.doc_lengths.len()];
|
||||||
let mut matched = false;
|
let mut matched = false;
|
||||||
for token in tokenize(query) {
|
for token in tokenize_with(query, self.filter) {
|
||||||
let Some(postings) = self.inverted.get(token.as_str()) else {
|
let Some(postings) = self.inverted.get(token.as_str()) else {
|
||||||
continue;
|
continue;
|
||||||
};
|
};
|
||||||
@@ -146,6 +157,11 @@ impl BM25Index {
|
|||||||
.collect()
|
.collect()
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// The token filter this index was built with.
|
||||||
|
pub fn token_filter(&self) -> TokenFilter {
|
||||||
|
self.filter
|
||||||
|
}
|
||||||
|
|
||||||
/// Number of document slots (live or not) the index covers. Ids are
|
/// Number of document slots (live or not) the index covers. Ids are
|
||||||
/// positions in the document list it mirrors.
|
/// positions in the document list it mirrors.
|
||||||
pub fn len(&self) -> usize {
|
pub fn len(&self) -> usize {
|
||||||
@@ -168,7 +184,7 @@ impl BM25Index {
|
|||||||
}
|
}
|
||||||
debug_assert_eq!(self.doc_lengths[doc_id], 0, "slot {doc_id} is occupied");
|
debug_assert_eq!(self.doc_lengths[doc_id], 0, "slot {doc_id} is occupied");
|
||||||
|
|
||||||
let tokens = tokenize(text);
|
let tokens = tokenize_with(text, self.filter);
|
||||||
let mut term_freqs: HashMap<&str, u32> = HashMap::new();
|
let mut term_freqs: HashMap<&str, u32> = HashMap::new();
|
||||||
for token in &tokens {
|
for token in &tokens {
|
||||||
*term_freqs.entry(token).or_insert(0) += 1;
|
*term_freqs.entry(token).or_insert(0) += 1;
|
||||||
@@ -201,7 +217,7 @@ impl BM25Index {
|
|||||||
/// Remove document `doc_id`, whose indexed text was `text`. The text is
|
/// Remove document `doc_id`, whose indexed text was `text`. The text is
|
||||||
/// needed to find its postings; pass exactly what was added.
|
/// needed to find its postings; pass exactly what was added.
|
||||||
pub fn remove_document(&mut self, doc_id: usize, text: &str) {
|
pub fn remove_document(&mut self, doc_id: usize, text: &str) {
|
||||||
let tokens = tokenize(text);
|
let tokens = tokenize_with(text, self.filter);
|
||||||
let mut seen: std::collections::HashSet<&str> = std::collections::HashSet::new();
|
let mut seen: std::collections::HashSet<&str> = std::collections::HashSet::new();
|
||||||
for token in &tokens {
|
for token in &tokens {
|
||||||
if !seen.insert(token) {
|
if !seen.insert(token) {
|
||||||
@@ -252,7 +268,7 @@ impl BM25Index {
|
|||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
|
|
||||||
let tokens = tokenize(doc);
|
let tokens = tokenize_with(doc, self.filter);
|
||||||
let doc_len = tokens.len() as u32;
|
let doc_len = tokens.len() as u32;
|
||||||
self.doc_lengths[i] = doc_len;
|
self.doc_lengths[i] = doc_len;
|
||||||
total_length += doc_len as u64;
|
total_length += doc_len as u64;
|
||||||
@@ -285,11 +301,86 @@ impl BM25Index {
|
|||||||
|
|
||||||
/// Tokenize a string: lowercase, split on non-alphanumeric characters,
|
/// Tokenize a string: lowercase, split on non-alphanumeric characters,
|
||||||
/// filter empty tokens.
|
/// filter empty tokens.
|
||||||
|
/// What [`tokenize_with`] does to each token after splitting.
|
||||||
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
|
||||||
|
pub enum TokenFilter {
|
||||||
|
/// Lowercase and split only — the original behaviour.
|
||||||
|
#[default]
|
||||||
|
Plain,
|
||||||
|
/// Also strip common English inflections, so "running" and "runs" match
|
||||||
|
/// "run". Conservative on purpose: only plural and past/continuous verb
|
||||||
|
/// endings, and only on tokens long enough that stripping leaves a real
|
||||||
|
/// stem. A stemmer earns its keep by conflating *related* words; an
|
||||||
|
/// aggressive one also conflates unrelated ones ("universe"/"university"),
|
||||||
|
/// which costs precision.
|
||||||
|
Stemmed,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Strip common English inflections from an already-lowercased token.
|
||||||
|
///
|
||||||
|
/// Applied identically to documents and queries, so the pair only has to agree
|
||||||
|
/// with itself — the stem need not be a real word.
|
||||||
|
fn stem(token: &str) -> &str {
|
||||||
|
// Below this, stripping does more harm than good ("bed" -> "b").
|
||||||
|
const MIN_STEM: usize = 4;
|
||||||
|
let strip = |suffix: &str, min_len: usize| -> Option<&str> {
|
||||||
|
let stem = token.strip_suffix(suffix)?;
|
||||||
|
(stem.len() >= min_len).then_some(stem)
|
||||||
|
};
|
||||||
|
|
||||||
|
// Plurals first: "studies" -> "studi", "classes" -> "class", "cats" -> "cat".
|
||||||
|
// "ies" keeps its "i" so the result meets "-ied" ("studied" -> "studi").
|
||||||
|
if let Some(stem) = strip("ies", 2) {
|
||||||
|
return &token[..stem.len() + 1];
|
||||||
|
}
|
||||||
|
for suffix in ["sses", "shes", "ches", "xes", "zes"] {
|
||||||
|
if let Some(stem) = strip(suffix, MIN_STEM - 1) {
|
||||||
|
// Keep the sibilant: "classes" -> "class", not "clas".
|
||||||
|
return &token[..stem.len() + 2];
|
||||||
|
}
|
||||||
|
}
|
||||||
|
// Verb endings before the bare plural, so "raced" doesn't become "raced".
|
||||||
|
if let Some(stem) = strip("ing", MIN_STEM - 1).or_else(|| strip("ed", MIN_STEM - 1)) {
|
||||||
|
return undouble(stem);
|
||||||
|
}
|
||||||
|
if !token.ends_with("ss")
|
||||||
|
&& !token.ends_with("us")
|
||||||
|
&& !token.ends_with("is")
|
||||||
|
&& let Some(stem) = strip("s", MIN_STEM - 1)
|
||||||
|
{
|
||||||
|
return stem;
|
||||||
|
}
|
||||||
|
token
|
||||||
|
}
|
||||||
|
|
||||||
|
/// "runn" -> "run": undo the consonant doubling that "-ing"/"-ed" introduce.
|
||||||
|
fn undouble(stem: &str) -> &str {
|
||||||
|
let mut chars = stem.chars().rev();
|
||||||
|
let (Some(last), Some(prev)) = (chars.next(), chars.next()) else {
|
||||||
|
return stem;
|
||||||
|
};
|
||||||
|
let doubled = last == prev && !"aeiou".contains(last) && last.is_ascii_alphabetic();
|
||||||
|
if doubled && stem.len() > 3 {
|
||||||
|
&stem[..stem.len() - 1]
|
||||||
|
} else {
|
||||||
|
stem
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
fn tokenize(text: &str) -> Vec<String> {
|
fn tokenize(text: &str) -> Vec<String> {
|
||||||
|
tokenize_with(text, TokenFilter::Plain)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Split `text` into scoring tokens under `filter`.
|
||||||
|
pub fn tokenize_with(text: &str, filter: TokenFilter) -> Vec<String> {
|
||||||
text.to_lowercase()
|
text.to_lowercase()
|
||||||
.split(|c: char| !c.is_alphanumeric())
|
.split(|c: char| !c.is_alphanumeric())
|
||||||
.filter(|s| !s.is_empty())
|
.filter(|s| !s.is_empty())
|
||||||
.map(|s| s.to_string())
|
.map(|token| match filter {
|
||||||
|
TokenFilter::Plain => token.to_string(),
|
||||||
|
TokenFilter::Stemmed => stem(token).to_string(),
|
||||||
|
})
|
||||||
.collect()
|
.collect()
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -473,8 +564,9 @@ mod tests {
|
|||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn wand_returns_same_results_as_exhaustive() {
|
fn top_k_search_matches_ranking_every_score() {
|
||||||
// WAND-style search should produce same scores as exhaustive
|
// `search` must agree with ranking the full `scores` set — the
|
||||||
|
// bounded heap is an optimisation over sorting, not an approximation.
|
||||||
let docs: Vec<String> = (0..100)
|
let docs: Vec<String> = (0..100)
|
||||||
.map(|i| {
|
.map(|i| {
|
||||||
if i % 3 == 0 {
|
if i % 3 == 0 {
|
||||||
@@ -599,6 +691,53 @@ mod tests {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn stemming_conflates_inflections_of_the_same_word() {
|
||||||
|
let stem_of = |w: &str| tokenize_with(w, TokenFilter::Stemmed).pop().unwrap();
|
||||||
|
// Pairs that should meet.
|
||||||
|
for (a, b) in [
|
||||||
|
("running", "runs"),
|
||||||
|
("trained", "training"),
|
||||||
|
("miles", "mile"),
|
||||||
|
("studies", "studied"),
|
||||||
|
("mentioned", "mentioning"),
|
||||||
|
("classes", "class"),
|
||||||
|
("planned", "planning"),
|
||||||
|
] {
|
||||||
|
assert_eq!(stem_of(a), stem_of(b), "{a} / {b} should share a stem");
|
||||||
|
}
|
||||||
|
// Pairs that must stay apart. Note which pairs are deliberately absent:
|
||||||
|
// "bed"/"bedding" and "gas"/"gassed" both collapse to one stem, which
|
||||||
|
// is what Porter does too and is right — they are related words.
|
||||||
|
for (a, b) in [
|
||||||
|
("universe", "university"),
|
||||||
|
("business", "busy"),
|
||||||
|
("this", "thing"),
|
||||||
|
] {
|
||||||
|
assert_ne!(stem_of(a), stem_of(b), "{a} / {b} must not be conflated");
|
||||||
|
}
|
||||||
|
// Short words and non-inflections are left alone.
|
||||||
|
for word in ["run", "bus", "is", "his", "data", "gas"] {
|
||||||
|
assert_eq!(stem_of(word), word, "{word} should be untouched");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn stemming_is_off_by_default_and_applied_consistently() {
|
||||||
|
assert_eq!(tokenize("Running miles"), ["running", "miles"]);
|
||||||
|
assert_eq!(
|
||||||
|
tokenize_with("Running miles", TokenFilter::Stemmed),
|
||||||
|
["run", "mile"]
|
||||||
|
);
|
||||||
|
|
||||||
|
// A query inflected differently from the document still matches.
|
||||||
|
let docs = vec!["I ran while training for the marathon".to_string()];
|
||||||
|
let plain = BM25Index::build_with(&docs, &[0], TokenFilter::Plain);
|
||||||
|
let stemmed = BM25Index::build_with(&docs, &[0], TokenFilter::Stemmed);
|
||||||
|
assert!(plain.search("trains", 1).is_empty());
|
||||||
|
assert_eq!(stemmed.search("trains", 1).len(), 1);
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn ties_break_towards_the_lower_doc_id() {
|
fn ties_break_towards_the_lower_doc_id() {
|
||||||
let docs: Vec<String> = (0..6).map(|_| "same text".to_string()).collect();
|
let docs: Vec<String> = (0..6).map(|_| "same text".to_string()).collect();
|
||||||
|
|||||||
@@ -1,17 +1,145 @@
|
|||||||
//! In-memory cache for memory entries, sessions, and knowledge graph.
|
//! In-memory cache for memory entries, sessions, and knowledge graph.
|
||||||
|
|
||||||
use crate::vector_search;
|
use crate::vector_search;
|
||||||
|
use clawhdf5_format::float16::round_to_f16;
|
||||||
|
|
||||||
|
/// Every entry's embedding, in one contiguous `[N x dim]` buffer.
|
||||||
|
///
|
||||||
|
/// Rows are always exactly `dim` long: a shorter one is zero-padded, a longer
|
||||||
|
/// one truncated. The previous `Vec<Vec<f32>>` allowed ragged rows, which
|
||||||
|
/// silently misaligned the flattened copy that the batched kernels read — a
|
||||||
|
/// single wrong-length embedding shifted every row after it. Padding makes
|
||||||
|
/// that unrepresentable. A record stored without an embedding therefore holds
|
||||||
|
/// a zero row, and is told apart by its norm being zero rather than by length.
|
||||||
|
///
|
||||||
|
/// This used to be two fields — a `Vec<Vec<f32>>` and a flattened copy kept in
|
||||||
|
/// lock-step — which stored the whole corpus twice and cost one heap
|
||||||
|
/// allocation per entry on top. At 100k 384-dim entries that duplicate was
|
||||||
|
/// ~150 MiB. Indexing yields a `&[f32]` row, so `embeddings[i]` still reads
|
||||||
|
/// the same way.
|
||||||
|
#[derive(Debug, Clone, Default)]
|
||||||
|
pub struct Embeddings {
|
||||||
|
flat: Vec<f32>,
|
||||||
|
dim: usize,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Embeddings {
|
||||||
|
pub fn new(dim: usize) -> Self {
|
||||||
|
Self {
|
||||||
|
flat: Vec::new(),
|
||||||
|
dim,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Number of embeddings.
|
||||||
|
pub fn len(&self) -> usize {
|
||||||
|
self.flat.len().checked_div(self.dim).unwrap_or(0)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn is_empty(&self) -> bool {
|
||||||
|
self.len() == 0
|
||||||
|
}
|
||||||
|
|
||||||
|
/// The whole buffer, `[N x dim]` row-major — what batched kernels read.
|
||||||
|
pub fn as_flat(&self) -> &[f32] {
|
||||||
|
&self.flat
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn dim(&self) -> usize {
|
||||||
|
self.dim
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Row `i`, or `None` if out of range.
|
||||||
|
pub fn get(&self, i: usize) -> Option<&[f32]> {
|
||||||
|
let start = i.checked_mul(self.dim)?;
|
||||||
|
self.flat.get(start..start.checked_add(self.dim)?)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn iter(&self) -> impl ExactSizeIterator<Item = &[f32]> {
|
||||||
|
self.flat.chunks_exact(self.dim.max(1))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Append one embedding. A row whose length doesn't match `dim` is padded
|
||||||
|
/// or truncated, so the buffer stays rectangular whatever a caller passes.
|
||||||
|
pub fn push(&mut self, embedding: &[f32]) {
|
||||||
|
if self.dim == 0 {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
let take = embedding.len().min(self.dim);
|
||||||
|
self.flat.extend_from_slice(&embedding[..take]);
|
||||||
|
self.flat.resize(self.flat.len() + (self.dim - take), 0.0);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Replace row `i`. Out-of-range indices are ignored.
|
||||||
|
pub fn set(&mut self, i: usize, embedding: &[f32]) {
|
||||||
|
let Some(start) = i.checked_mul(self.dim) else {
|
||||||
|
return;
|
||||||
|
};
|
||||||
|
if start + self.dim > self.flat.len() {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
let take = embedding.len().min(self.dim);
|
||||||
|
self.flat[start..start + take].copy_from_slice(&embedding[..take]);
|
||||||
|
self.flat[start + take..start + self.dim].fill(0.0);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Keep only the rows `keep` returns true for, preserving order.
|
||||||
|
pub fn retain(&mut self, mut keep: impl FnMut(usize) -> bool) {
|
||||||
|
if self.dim == 0 {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
let mut write = 0usize;
|
||||||
|
for read in 0..self.len() {
|
||||||
|
if keep(read) {
|
||||||
|
if write != read {
|
||||||
|
let (dst, src) = (write * self.dim, read * self.dim);
|
||||||
|
self.flat.copy_within(src..src + self.dim, dst);
|
||||||
|
}
|
||||||
|
write += 1;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
self.flat.truncate(write * self.dim);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Replace the contents with `rows`.
|
||||||
|
pub fn reset_from(&mut self, dim: usize, rows: impl IntoIterator<Item = Vec<f32>>) {
|
||||||
|
self.dim = dim;
|
||||||
|
self.flat.clear();
|
||||||
|
for row in rows {
|
||||||
|
self.push(&row);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Adopt an already-flat buffer, trimming any partial trailing row.
|
||||||
|
pub fn set_flat(&mut self, dim: usize, mut flat: Vec<f32>) {
|
||||||
|
self.dim = dim;
|
||||||
|
match flat.len().checked_div(dim) {
|
||||||
|
Some(rows) => flat.truncate(rows * dim),
|
||||||
|
None => flat.clear(),
|
||||||
|
}
|
||||||
|
self.flat = flat;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl PartialEq for Embeddings {
|
||||||
|
fn eq(&self, other: &Self) -> bool {
|
||||||
|
self.dim == other.dim && self.flat == other.flat
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl std::ops::Index<usize> for Embeddings {
|
||||||
|
type Output = [f32];
|
||||||
|
|
||||||
|
fn index(&self, i: usize) -> &[f32] {
|
||||||
|
self.get(i).expect("embedding index out of range")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
/// In-memory cache for the /memory group data.
|
/// In-memory cache for the /memory group data.
|
||||||
#[derive(Debug, Clone)]
|
#[derive(Debug, Clone)]
|
||||||
pub struct MemoryCache {
|
pub struct MemoryCache {
|
||||||
pub chunks: Vec<String>,
|
pub chunks: Vec<String>,
|
||||||
pub embeddings: Vec<Vec<f32>>,
|
pub embeddings: Embeddings,
|
||||||
/// `embeddings` flattened into one contiguous `[N × embedding_dim]`
|
|
||||||
/// buffer, maintained incrementally alongside `embeddings` (push/update/
|
|
||||||
/// compact) so BLAS/Accelerate batch search can read it directly instead
|
|
||||||
/// of re-flattening the whole corpus on every query.
|
|
||||||
pub embeddings_flat: Vec<f32>,
|
|
||||||
pub source_channels: Vec<String>,
|
pub source_channels: Vec<String>,
|
||||||
pub timestamps: Vec<f64>,
|
pub timestamps: Vec<f64>,
|
||||||
pub session_ids: Vec<String>,
|
pub session_ids: Vec<String>,
|
||||||
@@ -22,14 +150,18 @@ pub struct MemoryCache {
|
|||||||
pub norms: Vec<f32>,
|
pub norms: Vec<f32>,
|
||||||
/// Hebbian activation weights (default 1.0 per entry).
|
/// Hebbian activation weights (default 1.0 per entry).
|
||||||
pub activation_weights: Vec<f32>,
|
pub activation_weights: Vec<f32>,
|
||||||
|
/// Round every embedding to IEEE half precision as it enters the cache,
|
||||||
|
/// so the cache holds exactly what a `float16` store writes to disk. Set
|
||||||
|
/// it with [`MemoryCache::set_half_precision`], which also rounds the
|
||||||
|
/// rows already held.
|
||||||
|
pub half_precision: bool,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl MemoryCache {
|
impl MemoryCache {
|
||||||
pub fn new(embedding_dim: usize) -> Self {
|
pub fn new(embedding_dim: usize) -> Self {
|
||||||
Self {
|
Self {
|
||||||
chunks: Vec::new(),
|
chunks: Vec::new(),
|
||||||
embeddings: Vec::new(),
|
embeddings: Embeddings::new(embedding_dim),
|
||||||
embeddings_flat: Vec::new(),
|
|
||||||
source_channels: Vec::new(),
|
source_channels: Vec::new(),
|
||||||
timestamps: Vec::new(),
|
timestamps: Vec::new(),
|
||||||
session_ids: Vec::new(),
|
session_ids: Vec::new(),
|
||||||
@@ -38,18 +170,52 @@ impl MemoryCache {
|
|||||||
embedding_dim,
|
embedding_dim,
|
||||||
norms: Vec::new(),
|
norms: Vec::new(),
|
||||||
activation_weights: Vec::new(),
|
activation_weights: Vec::new(),
|
||||||
|
half_precision: false,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Rebuild `embeddings_flat` from `embeddings` from scratch. Callers that
|
/// Switch half-precision rounding on or off. Turning it on rounds every
|
||||||
/// populate `embeddings` directly (bulk loads) must call this afterward.
|
/// embedding already held (and recomputes norms where one changed) —
|
||||||
pub fn rebuild_flat(&mut self) {
|
/// e.g. a `float16` store whose last checkpoint predates half-precision
|
||||||
self.embeddings_flat.clear();
|
/// storage and so is still `f32` on disk.
|
||||||
self.embeddings_flat
|
pub fn set_half_precision(&mut self, on: bool) {
|
||||||
.reserve(self.embeddings.len() * self.embedding_dim);
|
self.half_precision = on;
|
||||||
for emb in &self.embeddings {
|
if !on {
|
||||||
self.embeddings_flat.extend_from_slice(emb);
|
return;
|
||||||
}
|
}
|
||||||
|
for i in 0..self.embeddings.len() {
|
||||||
|
let row = &self.embeddings[i];
|
||||||
|
if row
|
||||||
|
.iter()
|
||||||
|
.all(|&v| round_to_f16(v).to_bits() == v.to_bits())
|
||||||
|
{
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
let rounded: Vec<f32> = row.iter().map(|&v| round_to_f16(v)).collect();
|
||||||
|
self.norms[i] = vector_search::compute_norm(&rounded);
|
||||||
|
self.embeddings.set(i, &rounded);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// The embedding as the cache will hold it: rounded to half precision
|
||||||
|
/// when [`Self::half_precision`] is on, otherwise unchanged.
|
||||||
|
fn stored_form(&self, mut embedding: Vec<f32>) -> Vec<f32> {
|
||||||
|
if self.half_precision {
|
||||||
|
for v in &mut embedding {
|
||||||
|
*v = round_to_f16(*v);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
embedding
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Kept for callers that used to have to re-flatten after a bulk load.
|
||||||
|
/// The buffer is always flat now, so there is nothing to rebuild.
|
||||||
|
#[deprecated(note = "embeddings are stored flat; this is a no-op")]
|
||||||
|
pub fn rebuild_flat(&mut self) {}
|
||||||
|
|
||||||
|
/// The embeddings as one contiguous `[N x dim]` buffer.
|
||||||
|
pub fn flat_embeddings(&self) -> &[f32] {
|
||||||
|
self.embeddings.as_flat()
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Total number of entries (including tombstoned).
|
/// Total number of entries (including tombstoned).
|
||||||
@@ -77,10 +243,10 @@ impl MemoryCache {
|
|||||||
tags: String,
|
tags: String,
|
||||||
) -> usize {
|
) -> usize {
|
||||||
let idx = self.chunks.len();
|
let idx = self.chunks.len();
|
||||||
|
let embedding = self.stored_form(embedding);
|
||||||
let norm = vector_search::compute_norm(&embedding);
|
let norm = vector_search::compute_norm(&embedding);
|
||||||
self.chunks.push(chunk);
|
self.chunks.push(chunk);
|
||||||
self.embeddings_flat.extend_from_slice(&embedding);
|
self.embeddings.push(&embedding);
|
||||||
self.embeddings.push(embedding);
|
|
||||||
self.source_channels.push(source_channel);
|
self.source_channels.push(source_channel);
|
||||||
self.timestamps.push(timestamp);
|
self.timestamps.push(timestamp);
|
||||||
self.session_ids.push(session_id);
|
self.session_ids.push(session_id);
|
||||||
@@ -116,22 +282,10 @@ impl MemoryCache {
|
|||||||
session_id: String,
|
session_id: String,
|
||||||
) {
|
) {
|
||||||
if idx < self.chunks.len() {
|
if idx < self.chunks.len() {
|
||||||
|
let embedding = self.stored_form(embedding);
|
||||||
let norm = vector_search::compute_norm(&embedding);
|
let norm = vector_search::compute_norm(&embedding);
|
||||||
self.chunks[idx] = chunk;
|
self.chunks[idx] = chunk;
|
||||||
let dim = self.embedding_dim;
|
self.embeddings.set(idx, &embedding);
|
||||||
let flat_start = idx * dim;
|
|
||||||
let matches_dim =
|
|
||||||
embedding.len() == dim && flat_start + dim <= self.embeddings_flat.len();
|
|
||||||
self.embeddings[idx] = embedding;
|
|
||||||
if matches_dim {
|
|
||||||
self.embeddings_flat[flat_start..flat_start + dim]
|
|
||||||
.copy_from_slice(&self.embeddings[idx]);
|
|
||||||
} else {
|
|
||||||
// Embedding length doesn't match embedding_dim (shouldn't
|
|
||||||
// happen in practice) — fall back to a full rebuild rather
|
|
||||||
// than leave embeddings_flat misaligned with embeddings.
|
|
||||||
self.rebuild_flat();
|
|
||||||
}
|
|
||||||
self.source_channels[idx] = source_channel;
|
self.source_channels[idx] = source_channel;
|
||||||
self.timestamps[idx] = timestamp;
|
self.timestamps[idx] = timestamp;
|
||||||
self.session_ids[idx] = session_id;
|
self.session_ids[idx] = session_id;
|
||||||
@@ -183,7 +337,7 @@ impl MemoryCache {
|
|||||||
new_idx += 1;
|
new_idx += 1;
|
||||||
let norm = vector_search::compute_norm(&self.embeddings[i]);
|
let norm = vector_search::compute_norm(&self.embeddings[i]);
|
||||||
new_chunks.push(self.chunks[i].clone());
|
new_chunks.push(self.chunks[i].clone());
|
||||||
new_embeddings.push(self.embeddings[i].clone());
|
new_embeddings.push(self.embeddings[i].to_vec());
|
||||||
new_source_channels.push(self.source_channels[i].clone());
|
new_source_channels.push(self.source_channels[i].clone());
|
||||||
new_timestamps.push(self.timestamps[i]);
|
new_timestamps.push(self.timestamps[i]);
|
||||||
new_session_ids.push(self.session_ids[i].clone());
|
new_session_ids.push(self.session_ids[i].clone());
|
||||||
@@ -196,7 +350,8 @@ impl MemoryCache {
|
|||||||
|
|
||||||
let removed = old_len - new_chunks.len();
|
let removed = old_len - new_chunks.len();
|
||||||
self.chunks = new_chunks;
|
self.chunks = new_chunks;
|
||||||
self.embeddings = new_embeddings;
|
self.embeddings
|
||||||
|
.reset_from(self.embedding_dim, new_embeddings);
|
||||||
self.source_channels = new_source_channels;
|
self.source_channels = new_source_channels;
|
||||||
self.timestamps = new_timestamps;
|
self.timestamps = new_timestamps;
|
||||||
self.session_ids = new_session_ids;
|
self.session_ids = new_session_ids;
|
||||||
@@ -204,16 +359,14 @@ impl MemoryCache {
|
|||||||
self.tombstones = new_tombstones;
|
self.tombstones = new_tombstones;
|
||||||
self.norms = new_norms;
|
self.norms = new_norms;
|
||||||
self.activation_weights = new_activation_weights;
|
self.activation_weights = new_activation_weights;
|
||||||
self.rebuild_flat();
|
|
||||||
|
|
||||||
(removed, index_map)
|
(removed, index_map)
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Flatten all embeddings into a single Vec<f32> for HDF5 storage.
|
/// All embeddings as one owned `[N x dim]` buffer, for HDF5 storage.
|
||||||
/// `embeddings_flat` is already maintained incrementally, so this just
|
/// Prefer [`MemoryCache::flat_embeddings`] where a borrow will do.
|
||||||
/// clones it — kept as a method for callers that want an owned copy.
|
pub fn flat_embeddings_owned(&self) -> Vec<f32> {
|
||||||
pub fn flat_embeddings(&self) -> Vec<f32> {
|
self.embeddings.as_flat().to_vec()
|
||||||
self.embeddings_flat.clone()
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -224,7 +377,7 @@ mod tests {
|
|||||||
/// `embeddings_flat` must always equal a from-scratch flatten of `embeddings`.
|
/// `embeddings_flat` must always equal a from-scratch flatten of `embeddings`.
|
||||||
fn assert_flat_in_sync(cache: &MemoryCache) {
|
fn assert_flat_in_sync(cache: &MemoryCache) {
|
||||||
let expected: Vec<f32> = cache.embeddings.iter().flatten().copied().collect();
|
let expected: Vec<f32> = cache.embeddings.iter().flatten().copied().collect();
|
||||||
assert_eq!(cache.embeddings_flat, expected);
|
assert_eq!(cache.embeddings.as_flat(), expected);
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
@@ -247,7 +400,10 @@ mod tests {
|
|||||||
String::new(),
|
String::new(),
|
||||||
);
|
);
|
||||||
assert_flat_in_sync(&cache);
|
assert_flat_in_sync(&cache);
|
||||||
assert_eq!(cache.embeddings_flat, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
|
assert_eq!(
|
||||||
|
cache.embeddings.as_flat(),
|
||||||
|
vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]
|
||||||
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
@@ -279,7 +435,7 @@ mod tests {
|
|||||||
);
|
);
|
||||||
assert_flat_in_sync(&cache);
|
assert_flat_in_sync(&cache);
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
cache.embeddings_flat,
|
cache.embeddings.as_flat(),
|
||||||
vec![7.0, 8.0, 9.0, 4.0, 5.0, 6.0],
|
vec![7.0, 8.0, 9.0, 4.0, 5.0, 6.0],
|
||||||
"update must overwrite the correct flat slice, not just append"
|
"update must overwrite the correct flat slice, not just append"
|
||||||
);
|
);
|
||||||
@@ -315,14 +471,71 @@ mod tests {
|
|||||||
cache.mark_deleted(1);
|
cache.mark_deleted(1);
|
||||||
cache.compact();
|
cache.compact();
|
||||||
assert_flat_in_sync(&cache);
|
assert_flat_in_sync(&cache);
|
||||||
assert_eq!(cache.embeddings_flat, vec![1.0, 1.0, 3.0, 3.0]);
|
assert_eq!(cache.embeddings.as_flat(), vec![1.0, 1.0, 3.0, 3.0]);
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn rebuild_flat_matches_manual_flatten() {
|
fn rebuild_flat_matches_manual_flatten() {
|
||||||
let mut cache = MemoryCache::new(2);
|
let mut cache = MemoryCache::new(2);
|
||||||
cache.embeddings = vec![vec![1.0, 2.0], vec![3.0, 4.0]];
|
cache
|
||||||
cache.rebuild_flat();
|
.embeddings
|
||||||
assert_eq!(cache.embeddings_flat, vec![1.0, 2.0, 3.0, 4.0]);
|
.reset_from(2, vec![vec![1.0, 2.0], vec![3.0, 4.0]]);
|
||||||
|
assert_eq!(cache.embeddings.as_flat(), vec![1.0, 2.0, 3.0, 4.0]);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn set_half_precision_rounds_existing_rows_and_their_norms() {
|
||||||
|
// A store with float16 set whose checkpoint is still f32 on disk
|
||||||
|
// loads full-precision rows; switching rounding on must bring them to
|
||||||
|
// exactly what the next checkpoint will write.
|
||||||
|
let mut cache = MemoryCache::new(3);
|
||||||
|
cache.push(
|
||||||
|
"a".into(),
|
||||||
|
vec![0.1, 0.2, 0.3],
|
||||||
|
"c".into(),
|
||||||
|
0.0,
|
||||||
|
"s".into(),
|
||||||
|
"".into(),
|
||||||
|
);
|
||||||
|
cache.push(
|
||||||
|
"b".into(),
|
||||||
|
vec![0.5, 0.25, 1.0],
|
||||||
|
"c".into(),
|
||||||
|
0.0,
|
||||||
|
"s".into(),
|
||||||
|
"".into(),
|
||||||
|
);
|
||||||
|
let exact_norm = cache.norms[0];
|
||||||
|
|
||||||
|
cache.set_half_precision(true);
|
||||||
|
let row0: Vec<f32> = [0.1f32, 0.2, 0.3]
|
||||||
|
.iter()
|
||||||
|
.map(|&v| round_to_f16(v))
|
||||||
|
.collect();
|
||||||
|
assert_eq!(&cache.embeddings[0], row0.as_slice());
|
||||||
|
assert_eq!(cache.norms[0], vector_search::compute_norm(&row0));
|
||||||
|
assert_ne!(cache.norms[0], exact_norm);
|
||||||
|
// Already representable: untouched.
|
||||||
|
assert_eq!(&cache.embeddings[1], &[0.5, 0.25, 1.0]);
|
||||||
|
|
||||||
|
// New rows are rounded as they arrive, and updates too.
|
||||||
|
cache.push(
|
||||||
|
"c".into(),
|
||||||
|
vec![0.1, 0.0, 0.0],
|
||||||
|
"c".into(),
|
||||||
|
0.0,
|
||||||
|
"s".into(),
|
||||||
|
"".into(),
|
||||||
|
);
|
||||||
|
assert_eq!(cache.embeddings[2][0], round_to_f16(0.1));
|
||||||
|
cache.update(
|
||||||
|
2,
|
||||||
|
"c".into(),
|
||||||
|
vec![0.3, 0.0, 0.0],
|
||||||
|
"c".into(),
|
||||||
|
0.0,
|
||||||
|
"s".into(),
|
||||||
|
);
|
||||||
|
assert_eq!(cache.embeddings[2][0], round_to_f16(0.3));
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -144,9 +144,44 @@ pub struct ConsolidationStats {
|
|||||||
|
|
||||||
pub struct ImportanceScorer;
|
pub struct ImportanceScorer;
|
||||||
|
|
||||||
|
/// Sum of squares, in 8-wide lanes so it vectorises.
|
||||||
|
fn sum_of_squares(a: &[f32]) -> f32 {
|
||||||
|
let (blocks, tail) = a.as_chunks::<8>();
|
||||||
|
let mut acc = [0.0f32; 8];
|
||||||
|
for b in blocks {
|
||||||
|
for i in 0..8 {
|
||||||
|
acc[i] += b[i] * b[i];
|
||||||
|
}
|
||||||
|
}
|
||||||
|
acc.iter().sum::<f32>() + tail.iter().map(|x| x * x).sum::<f32>()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// `(a · b, |b|²)` in one pass over equal-length slices, in 8-wide lanes.
|
||||||
|
fn dot_and_norm2(a: &[f32], b: &[f32]) -> (f32, f32) {
|
||||||
|
let (a_blocks, a_tail) = a.as_chunks::<8>();
|
||||||
|
let (b_blocks, b_tail) = b.as_chunks::<8>();
|
||||||
|
let mut dot = [0.0f32; 8];
|
||||||
|
let mut nb = [0.0f32; 8];
|
||||||
|
for (x, y) in a_blocks.iter().zip(b_blocks) {
|
||||||
|
for i in 0..8 {
|
||||||
|
dot[i] += x[i] * y[i];
|
||||||
|
nb[i] += y[i] * y[i];
|
||||||
|
}
|
||||||
|
}
|
||||||
|
let mut d = dot.iter().sum::<f32>();
|
||||||
|
let mut n = nb.iter().sum::<f32>();
|
||||||
|
for (x, y) in a_tail.iter().zip(b_tail) {
|
||||||
|
d += x * y;
|
||||||
|
n += y * y;
|
||||||
|
}
|
||||||
|
(d, n)
|
||||||
|
}
|
||||||
|
|
||||||
impl ImportanceScorer {
|
impl ImportanceScorer {
|
||||||
/// Cosine similarity between two embedding slices.
|
/// Cosine similarity between two embedding slices.
|
||||||
/// Returns 0.0 if either norm is zero.
|
/// Returns 0.0 if either norm is zero. The reference that
|
||||||
|
/// [`Self::score_surprise`] is tested against.
|
||||||
|
#[cfg(test)]
|
||||||
fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
|
fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
|
||||||
let len = a.len().min(b.len());
|
let len = a.len().min(b.len());
|
||||||
if len == 0 {
|
if len == 0 {
|
||||||
@@ -167,13 +202,54 @@ impl ImportanceScorer {
|
|||||||
|
|
||||||
/// Novelty score: 1.0 − max cosine similarity against all existing records.
|
/// Novelty score: 1.0 − max cosine similarity against all existing records.
|
||||||
/// Returns 1.0 when there are no existing memories.
|
/// Returns 1.0 when there are no existing memories.
|
||||||
|
///
|
||||||
|
/// Same result as the reference cosine similarity against each record, but the
|
||||||
|
/// new embedding's norm is computed once rather than per record, each
|
||||||
|
/// record costs one fused pass (dot product and its norm together) rather
|
||||||
|
/// than three, and a large working set is scored in parallel. Every insert
|
||||||
|
/// scores against the whole working tier, so this is what an unbounded
|
||||||
|
/// working tier pays for: at 100K records it was the difference between a
|
||||||
|
/// benchmark finishing and not (`BENCHMARKS.md`, "Consolidation Efficiency").
|
||||||
pub fn score_surprise(embedding: &[f32], existing_memories: &[&MemoryRecord]) -> f32 {
|
pub fn score_surprise(embedding: &[f32], existing_memories: &[&MemoryRecord]) -> f32 {
|
||||||
if existing_memories.is_empty() {
|
if existing_memories.is_empty() {
|
||||||
return 1.0;
|
return 1.0;
|
||||||
}
|
}
|
||||||
|
let query_norm2 = sum_of_squares(embedding);
|
||||||
|
let similarity = |r: &&MemoryRecord| -> f32 {
|
||||||
|
let other = &r.embedding;
|
||||||
|
let len = embedding.len().min(other.len());
|
||||||
|
if len == 0 {
|
||||||
|
return 0.0;
|
||||||
|
}
|
||||||
|
let (dot, other_norm2) = dot_and_norm2(&embedding[..len], &other[..len]);
|
||||||
|
// A shorter record compares against the query's matching prefix.
|
||||||
|
let q2 = if len == embedding.len() {
|
||||||
|
query_norm2
|
||||||
|
} else {
|
||||||
|
sum_of_squares(&embedding[..len])
|
||||||
|
};
|
||||||
|
if q2 == 0.0 || other_norm2 == 0.0 {
|
||||||
|
return 0.0;
|
||||||
|
}
|
||||||
|
dot / (q2.sqrt() * other_norm2.sqrt())
|
||||||
|
};
|
||||||
|
#[cfg(feature = "parallel")]
|
||||||
|
let max_sim = if existing_memories.len() >= 4096 {
|
||||||
|
use rayon::prelude::*;
|
||||||
|
existing_memories
|
||||||
|
.par_iter()
|
||||||
|
.map(similarity)
|
||||||
|
.reduce(|| f32::NEG_INFINITY, f32::max)
|
||||||
|
} else {
|
||||||
|
existing_memories
|
||||||
|
.iter()
|
||||||
|
.map(similarity)
|
||||||
|
.fold(f32::NEG_INFINITY, f32::max)
|
||||||
|
};
|
||||||
|
#[cfg(not(feature = "parallel"))]
|
||||||
let max_sim = existing_memories
|
let max_sim = existing_memories
|
||||||
.iter()
|
.iter()
|
||||||
.map(|r| Self::cosine_similarity(embedding, &r.embedding))
|
.map(similarity)
|
||||||
.fold(f32::NEG_INFINITY, f32::max);
|
.fold(f32::NEG_INFINITY, f32::max);
|
||||||
(1.0 - max_sim).clamp(0.0, 1.0)
|
(1.0 - max_sim).clamp(0.0, 1.0)
|
||||||
}
|
}
|
||||||
@@ -471,6 +547,54 @@ impl ConsolidationEngine {
|
|||||||
mod tests {
|
mod tests {
|
||||||
use super::*;
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn score_surprise_matches_the_reference_cosine() {
|
||||||
|
let mut x = 0x2545_F491_4F6C_DD1Du64;
|
||||||
|
let mut next = || {
|
||||||
|
x ^= x << 13;
|
||||||
|
x ^= x >> 7;
|
||||||
|
x ^= x << 17;
|
||||||
|
(x >> 40) as f32 / (1u64 << 24) as f32 - 0.5
|
||||||
|
};
|
||||||
|
let make = |id: u64, v: Vec<f32>| MemoryRecord {
|
||||||
|
id,
|
||||||
|
chunk: String::new(),
|
||||||
|
embedding: v,
|
||||||
|
tier: MemoryTier::Working,
|
||||||
|
importance: 0.0,
|
||||||
|
access_count: 0,
|
||||||
|
last_accessed: 0.0,
|
||||||
|
created_at: 0.0,
|
||||||
|
source: MemorySource::User,
|
||||||
|
};
|
||||||
|
// Ordinary rows, a shorter one, an empty one and a zero vector; and
|
||||||
|
// enough rows to take the parallel path too.
|
||||||
|
for n in [5usize, 5000] {
|
||||||
|
let mut recs: Vec<MemoryRecord> = (0..n as u64)
|
||||||
|
.map(|i| make(i, (0..37).map(|_| next()).collect()))
|
||||||
|
.collect();
|
||||||
|
recs.push(make(9_000, (0..20).map(|_| next()).collect()));
|
||||||
|
recs.push(make(9_001, Vec::new()));
|
||||||
|
recs.push(make(9_002, vec![0.0; 37]));
|
||||||
|
let refs: Vec<&MemoryRecord> = recs.iter().collect();
|
||||||
|
for _ in 0..5 {
|
||||||
|
let q: Vec<f32> = (0..37).map(|_| next()).collect();
|
||||||
|
let expected = (1.0
|
||||||
|
- refs
|
||||||
|
.iter()
|
||||||
|
.map(|r| ImportanceScorer::cosine_similarity(&q, &r.embedding))
|
||||||
|
.fold(f32::NEG_INFINITY, f32::max))
|
||||||
|
.clamp(0.0, 1.0);
|
||||||
|
let got = ImportanceScorer::score_surprise(&q, &refs);
|
||||||
|
assert!((got - expected).abs() < 1e-5, "n={n}: {got} vs {expected}");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
assert_eq!(
|
||||||
|
ImportanceScorer::score_surprise(&[0.0; 4], &[&make(1, vec![1.0; 4])]),
|
||||||
|
1.0
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
// Helper: build a simple normalised embedding of given dimension.
|
// Helper: build a simple normalised embedding of given dimension.
|
||||||
fn unit_vec(dim: usize, hot: usize) -> Vec<f32> {
|
fn unit_vec(dim: usize, hot: usize) -> Vec<f32> {
|
||||||
let mut v = vec![0.0f32; dim];
|
let mut v = vec![0.0f32; dim];
|
||||||
|
|||||||
@@ -28,39 +28,69 @@ use crate::vector_search;
|
|||||||
pub fn hybrid_search(
|
pub fn hybrid_search(
|
||||||
query_embedding: &[f32],
|
query_embedding: &[f32],
|
||||||
query_text: &str,
|
query_text: &str,
|
||||||
vectors: &[Vec<f32>],
|
vectors: &(impl crate::vector_search::VectorSet + Sync + ?Sized),
|
||||||
_chunks: &[String],
|
chunks: &[String],
|
||||||
tombstones: &[u8],
|
tombstones: &[u8],
|
||||||
bm25_index: &BM25Index,
|
bm25_index: &BM25Index,
|
||||||
vector_weight: f32,
|
vector_weight: f32,
|
||||||
keyword_weight: f32,
|
keyword_weight: f32,
|
||||||
k: usize,
|
k: usize,
|
||||||
|
) -> Vec<(usize, f32)> {
|
||||||
|
hybrid_search_fused(
|
||||||
|
query_embedding,
|
||||||
|
query_text,
|
||||||
|
vectors,
|
||||||
|
chunks,
|
||||||
|
tombstones,
|
||||||
|
bm25_index,
|
||||||
|
Fusion::Weighted {
|
||||||
|
vector: vector_weight,
|
||||||
|
keyword: keyword_weight,
|
||||||
|
},
|
||||||
|
k,
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// [`hybrid_search`] with the fusion method chosen explicitly.
|
||||||
|
#[allow(clippy::too_many_arguments)]
|
||||||
|
pub fn hybrid_search_fused(
|
||||||
|
query_embedding: &[f32],
|
||||||
|
query_text: &str,
|
||||||
|
vectors: &(impl crate::vector_search::VectorSet + Sync + ?Sized),
|
||||||
|
_chunks: &[String],
|
||||||
|
tombstones: &[u8],
|
||||||
|
bm25_index: &BM25Index,
|
||||||
|
fusion: Fusion,
|
||||||
|
k: usize,
|
||||||
) -> Vec<(usize, f32)> {
|
) -> Vec<(usize, f32)> {
|
||||||
// Get raw scores from both systems. Request all results so normalization
|
// Get raw scores from both systems. Request all results so normalization
|
||||||
// covers the full distribution.
|
// covers the full distribution.
|
||||||
// Use parallel search when rayon feature is enabled and vector count > 10K.
|
let vec_scores = exact_vector_scores(query_embedding, vectors, tombstones);
|
||||||
let vec_scores = {
|
|
||||||
#[cfg(feature = "parallel")]
|
|
||||||
{
|
|
||||||
if vectors.len() > 10_000 {
|
|
||||||
vector_search::parallel_cosine_batch(
|
|
||||||
query_embedding,
|
|
||||||
vectors,
|
|
||||||
tombstones,
|
|
||||||
vectors.len(),
|
|
||||||
)
|
|
||||||
} else {
|
|
||||||
vector_search::cosine_similarity_batch(query_embedding, vectors, tombstones)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
#[cfg(not(feature = "parallel"))]
|
|
||||||
{
|
|
||||||
vector_search::cosine_similarity_batch(query_embedding, vectors, tombstones)
|
|
||||||
}
|
|
||||||
};
|
|
||||||
let kw_scores = bm25_index.scores(query_text);
|
let kw_scores = bm25_index.scores(query_text);
|
||||||
|
|
||||||
merge_vector_keyword(vec_scores, kw_scores, vector_weight, keyword_weight, k)
|
fuse(vec_scores, kw_scores, fusion, k)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Cosine similarity of `query_embedding` to every vector whose `skip` byte is
|
||||||
|
/// 0 (a tombstone, or any other exclusion mask). Parallel above 10K vectors
|
||||||
|
/// when the `parallel` feature is on.
|
||||||
|
pub fn exact_vector_scores(
|
||||||
|
query_embedding: &[f32],
|
||||||
|
vectors: &(impl crate::vector_search::VectorSet + Sync + ?Sized),
|
||||||
|
skip: &[u8],
|
||||||
|
) -> Vec<(usize, f32)> {
|
||||||
|
#[cfg(feature = "parallel")]
|
||||||
|
{
|
||||||
|
if vectors.count() > 10_000 {
|
||||||
|
return vector_search::parallel_cosine_batch(
|
||||||
|
query_embedding,
|
||||||
|
vectors,
|
||||||
|
skip,
|
||||||
|
vectors.count(),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
vector_search::cosine_similarity_batch(query_embedding, vectors, skip)
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Merge pre-computed vector-similarity and keyword scores into a single ranking.
|
/// Merge pre-computed vector-similarity and keyword scores into a single ranking.
|
||||||
@@ -76,18 +106,92 @@ pub fn merge_vector_keyword(
|
|||||||
keyword_weight: f32,
|
keyword_weight: f32,
|
||||||
k: usize,
|
k: usize,
|
||||||
) -> Vec<(usize, f32)> {
|
) -> Vec<(usize, f32)> {
|
||||||
// Normalize each set to [0, 1].
|
fuse(
|
||||||
let vec_normalized = normalize_scores(&vec_scores);
|
vec_scores,
|
||||||
let kw_normalized = normalize_scores(&kw_scores);
|
kw_scores,
|
||||||
|
Fusion::Weighted {
|
||||||
|
vector: vector_weight,
|
||||||
|
keyword: keyword_weight,
|
||||||
|
},
|
||||||
|
k,
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
// Merge scores with weights.
|
/// How the vector and keyword stages are combined into one ranking.
|
||||||
let mut merged: HashMap<usize, f32> = HashMap::new();
|
#[derive(Debug, Clone, Copy, PartialEq)]
|
||||||
|
pub enum Fusion {
|
||||||
|
/// Min-max normalise each stage over its own candidates, then take a
|
||||||
|
/// weighted sum. Uses the *scores*, so a stage that separates its
|
||||||
|
/// candidates sharply keeps that separation — and a stage whose candidates
|
||||||
|
/// are all near-identical contributes little.
|
||||||
|
Weighted {
|
||||||
|
/// Weight on the vector stage.
|
||||||
|
vector: f32,
|
||||||
|
/// Weight on the keyword stage.
|
||||||
|
keyword: f32,
|
||||||
|
},
|
||||||
|
/// Reciprocal rank fusion: each stage contributes `1 / (k + rank)`,
|
||||||
|
/// ignoring score magnitudes entirely. Robust when the two stages'
|
||||||
|
/// scores aren't comparable, at the cost of discarding confidence.
|
||||||
|
Rrf {
|
||||||
|
/// The rank-damping constant; 60 is the value from the original paper.
|
||||||
|
k: f32,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
for (idx, score) in &vec_normalized {
|
impl Default for Fusion {
|
||||||
*merged.entry(*idx).or_insert(0.0) += vector_weight * score;
|
fn default() -> Self {
|
||||||
|
DEFAULT_FUSION
|
||||||
}
|
}
|
||||||
for (idx, score) in &kw_normalized {
|
}
|
||||||
*merged.entry(*idx).or_insert(0.0) += keyword_weight * score;
|
|
||||||
|
/// The fusion `hybrid_search` uses unless told otherwise.
|
||||||
|
///
|
||||||
|
/// The weights are not a guess: a sweep of every 0.1 step over the full
|
||||||
|
/// LongMemEval haystack (500 questions, real MiniLM embeddings) found the
|
||||||
|
/// long-standing 0.7/0.3 default *strictly dominated* — 0.4/0.6 is better at
|
||||||
|
/// Hit@1, Hit@5, Hit@10 and MRR, at both turn and session granularity. See
|
||||||
|
/// `BENCHMARKS.md`, "Weight sweep".
|
||||||
|
pub const DEFAULT_FUSION: Fusion = Fusion::Weighted {
|
||||||
|
vector: 0.4,
|
||||||
|
keyword: 0.6,
|
||||||
|
};
|
||||||
|
|
||||||
|
/// Combine one ranked candidate list from each stage into a single top-`k`.
|
||||||
|
///
|
||||||
|
/// Neither list need be sorted; both are consumed.
|
||||||
|
pub fn fuse(
|
||||||
|
vec_scores: Vec<(usize, f32)>,
|
||||||
|
kw_scores: Vec<(usize, f32)>,
|
||||||
|
fusion: Fusion,
|
||||||
|
k: usize,
|
||||||
|
) -> Vec<(usize, f32)> {
|
||||||
|
let mut merged: HashMap<usize, f32> = HashMap::new();
|
||||||
|
match fusion {
|
||||||
|
Fusion::Weighted { vector, keyword } => {
|
||||||
|
// Normalize each set to [0, 1].
|
||||||
|
for (idx, score) in &normalize_scores(&vec_scores) {
|
||||||
|
*merged.entry(*idx).or_insert(0.0) += vector * score;
|
||||||
|
}
|
||||||
|
for (idx, score) in &normalize_scores(&kw_scores) {
|
||||||
|
*merged.entry(*idx).or_insert(0.0) += keyword * score;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Fusion::Rrf { k: damping } => {
|
||||||
|
for mut stage in [vec_scores, kw_scores] {
|
||||||
|
// Rank 1 is the best score. Ties break by index so a stage's
|
||||||
|
// contribution doesn't depend on the candidate order it
|
||||||
|
// happened to be produced in.
|
||||||
|
stage.sort_by(|a, b| {
|
||||||
|
b.1.partial_cmp(&a.1)
|
||||||
|
.unwrap_or(std::cmp::Ordering::Equal)
|
||||||
|
.then(a.0.cmp(&b.0))
|
||||||
|
});
|
||||||
|
for (rank, (idx, _)) in stage.iter().enumerate() {
|
||||||
|
*merged.entry(*idx).or_insert(0.0) += 1.0 / (damping + (rank + 1) as f32);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
let mut results: Vec<(usize, f32)> = merged.into_iter().collect();
|
let mut results: Vec<(usize, f32)> = merged.into_iter().collect();
|
||||||
@@ -169,7 +273,7 @@ fn normalize_scores(scores: &[(usize, f32)]) -> Vec<(usize, f32)> {
|
|||||||
pub fn rrf_hybrid_search(
|
pub fn rrf_hybrid_search(
|
||||||
query_embedding: &[f32],
|
query_embedding: &[f32],
|
||||||
query_text: &str,
|
query_text: &str,
|
||||||
vectors: &[Vec<f32>],
|
vectors: &(impl crate::vector_search::VectorSet + Sync + ?Sized),
|
||||||
_chunks: &[String],
|
_chunks: &[String],
|
||||||
tombstones: &[u8],
|
tombstones: &[u8],
|
||||||
bm25_index: &BM25Index,
|
bm25_index: &BM25Index,
|
||||||
@@ -181,12 +285,12 @@ pub fn rrf_hybrid_search(
|
|||||||
let mut vec_scores = {
|
let mut vec_scores = {
|
||||||
#[cfg(feature = "parallel")]
|
#[cfg(feature = "parallel")]
|
||||||
{
|
{
|
||||||
if vectors.len() > 10_000 {
|
if vectors.count() > 10_000 {
|
||||||
vector_search::parallel_cosine_batch(
|
vector_search::parallel_cosine_batch(
|
||||||
query_embedding,
|
query_embedding,
|
||||||
vectors,
|
vectors,
|
||||||
tombstones,
|
tombstones,
|
||||||
vectors.len(),
|
vectors.count(),
|
||||||
)
|
)
|
||||||
} else {
|
} else {
|
||||||
vector_search::cosine_similarity_batch(query_embedding, vectors, tombstones)
|
vector_search::cosine_similarity_batch(query_embedding, vectors, tombstones)
|
||||||
@@ -197,7 +301,7 @@ pub fn rrf_hybrid_search(
|
|||||||
vector_search::cosine_similarity_batch(query_embedding, vectors, tombstones)
|
vector_search::cosine_similarity_batch(query_embedding, vectors, tombstones)
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
let mut kw_scores = bm25_index.search(query_text, vectors.len());
|
let mut kw_scores = bm25_index.search(query_text, vectors.count());
|
||||||
|
|
||||||
// Sort both lists descending so rank 1 = best.
|
// Sort both lists descending so rank 1 = best.
|
||||||
vec_scores.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
|
vec_scores.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
|
||||||
@@ -353,6 +457,49 @@ mod tests {
|
|||||||
assert_eq!(result[0].1, 1.0);
|
assert_eq!(result[0].1, 1.0);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn default_fusion_is_the_tuned_operating_point() {
|
||||||
|
// A sweep over the full LongMemEval haystack found 0.7/0.3 strictly
|
||||||
|
// dominated by 0.4/0.6 (BENCHMARKS.md). This guards the finding
|
||||||
|
// against being quietly undone.
|
||||||
|
assert_eq!(
|
||||||
|
DEFAULT_FUSION,
|
||||||
|
Fusion::Weighted {
|
||||||
|
vector: 0.4,
|
||||||
|
keyword: 0.6
|
||||||
|
}
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn rrf_rewards_agreement_between_the_stages_and_ignores_magnitudes() {
|
||||||
|
// Doc 1 is second-best in both stages; doc 0 is best in one and absent
|
||||||
|
// from the other. RRF prefers the doc both stages liked.
|
||||||
|
let vec_scores = vec![(0, 100.0), (1, 0.9)];
|
||||||
|
let kw_scores = vec![(2, 5.0), (1, 4.9)];
|
||||||
|
let ranked = fuse(vec_scores, kw_scores, Fusion::Rrf { k: 60.0 }, 3);
|
||||||
|
assert_eq!(ranked[0].0, 1, "{ranked:?}");
|
||||||
|
|
||||||
|
// Scaling one stage's scores cannot change an RRF ranking, only the
|
||||||
|
// order within that stage can.
|
||||||
|
let a = fuse(
|
||||||
|
vec![(0, 1.0), (1, 0.5)],
|
||||||
|
vec![(1, 2.0), (0, 1.0)],
|
||||||
|
Fusion::Rrf { k: 60.0 },
|
||||||
|
2,
|
||||||
|
);
|
||||||
|
let b = fuse(
|
||||||
|
vec![(0, 1e6), (1, -3.0)],
|
||||||
|
vec![(1, 0.002), (0, 0.001)],
|
||||||
|
Fusion::Rrf { k: 60.0 },
|
||||||
|
2,
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
a.iter().map(|r| r.0).collect::<Vec<_>>(),
|
||||||
|
b.iter().map(|r| r.0).collect::<Vec<_>>()
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn merge_top_k_matches_a_full_sort() {
|
fn merge_top_k_matches_a_full_sort() {
|
||||||
// Many ties (scores repeat) so the index tie-break is exercised.
|
// Many ties (scores repeat) so the index tie-break is exercised.
|
||||||
|
|||||||
@@ -163,12 +163,13 @@ fn levenshtein(a: &str, b: &str) -> usize {
|
|||||||
/// entities-slice-index map, and an entity-id -> relation-indices map (edges
|
/// entities-slice-index map, and an entity-id -> relation-indices map (edges
|
||||||
/// touching that entity as either source or target).
|
/// touching that entity as either source or target).
|
||||||
///
|
///
|
||||||
/// Built fresh per traversal call rather than cached on `KnowledgeCache`:
|
/// Cached on `KnowledgeCache` and checked against a fingerprint of the graph
|
||||||
/// entities/relations are plain `pub` `Vec`s that get pushed to directly
|
/// on every use ([`graph_fingerprint`]). entities/relations are plain `pub`
|
||||||
/// (e.g. `schema.rs`'s load path bypasses `add_entity`/`add_relation`), so a
|
/// `Vec`s that get changed directly (e.g. `schema.rs`'s load path bypasses
|
||||||
/// persistent index would need extra bookkeeping to avoid drifting stale. A
|
/// `add_entity`/`add_relation`), so the cache cannot rely on being told about
|
||||||
/// one-off O(V+E) build per call is still a large win over the O(V·E) (BFS)
|
/// changes; the fingerprint notices any of them. Rebuilding it on every
|
||||||
/// / O(steps·active·E) (spreading activation) scans it replaces.
|
/// traversal instead made a 2-hop BFS over 1K entities 6.5x slower than the
|
||||||
|
/// scan it replaced (24 -> 155 µs; `BENCHMARKS.md`, "Knowledge Graph").
|
||||||
struct AdjacencyIndex {
|
struct AdjacencyIndex {
|
||||||
entity_index: HashMap<u64, usize>,
|
entity_index: HashMap<u64, usize>,
|
||||||
by_entity: HashMap<u64, Vec<usize>>,
|
by_entity: HashMap<u64, Vec<usize>>,
|
||||||
@@ -204,6 +205,45 @@ impl AdjacencyIndex {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// A hash of everything [`AdjacencyIndex`] depends on — each entity's id and
|
||||||
|
/// position, each relation's endpoints and position. One linear pass, no
|
||||||
|
/// allocation: far cheaper than building the index, which hashes the same
|
||||||
|
/// values into two maps.
|
||||||
|
fn graph_fingerprint(entities: &[Entity], relations: &[Relation]) -> u64 {
|
||||||
|
// splitmix64-style mixing; order matters, so positions are covered.
|
||||||
|
fn mix(h: u64, v: u64) -> u64 {
|
||||||
|
let mut z = (h ^ v).wrapping_add(0x9E37_79B9_7F4A_7C15);
|
||||||
|
z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
|
||||||
|
z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
|
||||||
|
z ^ (z >> 31)
|
||||||
|
}
|
||||||
|
let mut h = mix(entities.len() as u64, relations.len() as u64);
|
||||||
|
for e in entities {
|
||||||
|
h = mix(h, e.id);
|
||||||
|
}
|
||||||
|
for r in relations {
|
||||||
|
h = mix(mix(h, r.src), r.tgt);
|
||||||
|
}
|
||||||
|
h
|
||||||
|
}
|
||||||
|
|
||||||
|
/// The cached [`AdjacencyIndex`] and the fingerprint it was built for.
|
||||||
|
/// Cloning a `KnowledgeCache` starts the clone with an empty cache.
|
||||||
|
#[derive(Default)]
|
||||||
|
struct AdjacencyCache(std::sync::Mutex<Option<(u64, std::sync::Arc<AdjacencyIndex>)>>);
|
||||||
|
|
||||||
|
impl Clone for AdjacencyCache {
|
||||||
|
fn clone(&self) -> Self {
|
||||||
|
Self::default()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl std::fmt::Debug for AdjacencyCache {
|
||||||
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||||
|
f.write_str("AdjacencyCache")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
// ---------------------------------------------------------------------------
|
// ---------------------------------------------------------------------------
|
||||||
// KnowledgeCache
|
// KnowledgeCache
|
||||||
// ---------------------------------------------------------------------------
|
// ---------------------------------------------------------------------------
|
||||||
@@ -216,6 +256,7 @@ pub struct KnowledgeCache {
|
|||||||
pub alias_strings: Vec<String>,
|
pub alias_strings: Vec<String>,
|
||||||
pub alias_entity_ids: Vec<i64>,
|
pub alias_entity_ids: Vec<i64>,
|
||||||
next_entity_id: u64,
|
next_entity_id: u64,
|
||||||
|
adjacency: AdjacencyCache,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl KnowledgeCache {
|
impl KnowledgeCache {
|
||||||
@@ -226,6 +267,7 @@ impl KnowledgeCache {
|
|||||||
alias_strings: Vec::new(),
|
alias_strings: Vec::new(),
|
||||||
alias_entity_ids: Vec::new(),
|
alias_entity_ids: Vec::new(),
|
||||||
next_entity_id: 0,
|
next_entity_id: 0,
|
||||||
|
adjacency: AdjacencyCache::default(),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -236,9 +278,29 @@ impl KnowledgeCache {
|
|||||||
alias_strings: Vec::new(),
|
alias_strings: Vec::new(),
|
||||||
alias_entity_ids: Vec::new(),
|
alias_entity_ids: Vec::new(),
|
||||||
next_entity_id: next_id,
|
next_entity_id: next_id,
|
||||||
|
adjacency: AdjacencyCache::default(),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// The adjacency index for the graph as it is now: the cached one if the
|
||||||
|
/// graph's fingerprint still matches, otherwise rebuilt and cached.
|
||||||
|
fn adjacency_index(&self) -> std::sync::Arc<AdjacencyIndex> {
|
||||||
|
let fp = graph_fingerprint(&self.entities, &self.relations);
|
||||||
|
let mut slot = self
|
||||||
|
.adjacency
|
||||||
|
.0
|
||||||
|
.lock()
|
||||||
|
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
||||||
|
if let Some((cached_fp, idx)) = slot.as_ref()
|
||||||
|
&& *cached_fp == fp
|
||||||
|
{
|
||||||
|
return idx.clone();
|
||||||
|
}
|
||||||
|
let idx = std::sync::Arc::new(AdjacencyIndex::build(&self.entities, &self.relations));
|
||||||
|
*slot = Some((fp, idx.clone()));
|
||||||
|
idx
|
||||||
|
}
|
||||||
|
|
||||||
// -----------------------------------------------------------------------
|
// -----------------------------------------------------------------------
|
||||||
// Entity management
|
// Entity management
|
||||||
// -----------------------------------------------------------------------
|
// -----------------------------------------------------------------------
|
||||||
@@ -397,7 +459,7 @@ impl KnowledgeCache {
|
|||||||
/// together with their discovered depth. The seed entity itself is NOT
|
/// together with their discovered depth. The seed entity itself is NOT
|
||||||
/// included. Traversal follows both outgoing and incoming relation edges.
|
/// included. Traversal follows both outgoing and incoming relation edges.
|
||||||
pub fn bfs_neighbors(&self, entity_id: u64, max_depth: usize) -> Vec<(Entity, usize)> {
|
pub fn bfs_neighbors(&self, entity_id: u64, max_depth: usize) -> Vec<(Entity, usize)> {
|
||||||
let idx = AdjacencyIndex::build(&self.entities, &self.relations);
|
let idx = self.adjacency_index();
|
||||||
let mut visited: HashSet<u64> = HashSet::new();
|
let mut visited: HashSet<u64> = HashSet::new();
|
||||||
let mut queue: VecDeque<(u64, usize)> = VecDeque::new();
|
let mut queue: VecDeque<(u64, usize)> = VecDeque::new();
|
||||||
let mut results: Vec<(Entity, usize)> = Vec::new();
|
let mut results: Vec<(Entity, usize)> = Vec::new();
|
||||||
@@ -502,7 +564,7 @@ impl KnowledgeCache {
|
|||||||
min_activation: f32,
|
min_activation: f32,
|
||||||
max_steps: usize,
|
max_steps: usize,
|
||||||
) -> Vec<(u64, f32)> {
|
) -> Vec<(u64, f32)> {
|
||||||
let idx = AdjacencyIndex::build(&self.entities, &self.relations);
|
let idx = self.adjacency_index();
|
||||||
let mut activation: HashMap<u64, f32> = HashMap::new();
|
let mut activation: HashMap<u64, f32> = HashMap::new();
|
||||||
|
|
||||||
// Initialise seeds with activation 1.0.
|
// Initialise seeds with activation 1.0.
|
||||||
@@ -631,6 +693,51 @@ impl Default for KnowledgeCache {
|
|||||||
mod tests {
|
mod tests {
|
||||||
use super::*;
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn cached_adjacency_sees_direct_changes_to_the_graph() {
|
||||||
|
// The index is cached across traversals, but entities/relations are
|
||||||
|
// pub Vecs anyone can edit; every kind of edit must be seen.
|
||||||
|
let mut kg = KnowledgeCache::new();
|
||||||
|
let a = kg.add_entity("a", "t", -1);
|
||||||
|
let b = kg.add_entity("b", "t", -1);
|
||||||
|
let c = kg.add_entity("c", "t", -1);
|
||||||
|
kg.add_relation(a, b, "r", 1.0);
|
||||||
|
let ids = |kg: &KnowledgeCache| -> Vec<u64> {
|
||||||
|
let mut v: Vec<u64> = kg.bfs_neighbors(a, 3).iter().map(|(e, _)| e.id).collect();
|
||||||
|
v.sort();
|
||||||
|
v
|
||||||
|
};
|
||||||
|
assert_eq!(ids(&kg), vec![b]);
|
||||||
|
assert_eq!(ids(&kg), vec![b], "cached index reused");
|
||||||
|
|
||||||
|
// Pushed directly, bypassing add_relation.
|
||||||
|
kg.relations.push(Relation {
|
||||||
|
src: b,
|
||||||
|
tgt: c,
|
||||||
|
..Relation::default()
|
||||||
|
});
|
||||||
|
assert_eq!(ids(&kg), vec![b, c]);
|
||||||
|
|
||||||
|
// Rewired in place: same lengths, different edge.
|
||||||
|
kg.relations[1].tgt = a;
|
||||||
|
assert_eq!(ids(&kg), vec![b]);
|
||||||
|
|
||||||
|
// Removed and replaced: same lengths again.
|
||||||
|
kg.relations.pop();
|
||||||
|
kg.relations.push(Relation {
|
||||||
|
src: a,
|
||||||
|
tgt: c,
|
||||||
|
..Relation::default()
|
||||||
|
});
|
||||||
|
assert_eq!(ids(&kg), vec![b, c]);
|
||||||
|
let act: Vec<u64> = kg
|
||||||
|
.spreading_activation(&[a], 0.5, 0.0, 2)
|
||||||
|
.iter()
|
||||||
|
.map(|(id, _)| *id)
|
||||||
|
.collect();
|
||||||
|
assert!(act.contains(&c));
|
||||||
|
}
|
||||||
|
|
||||||
// -----------------------------------------------------------------------
|
// -----------------------------------------------------------------------
|
||||||
// Original tests — must remain passing
|
// Original tests — must remain passing
|
||||||
// -----------------------------------------------------------------------
|
// -----------------------------------------------------------------------
|
||||||
|
|||||||
@@ -1,4 +1,4 @@
|
|||||||
//! ZeroClaw agent memory HDF5 backend.
|
//! Agent memory stored in a single HDF5 file.
|
||||||
//!
|
//!
|
||||||
//! Provides persistent memory storage for AI agents using HDF5 files.
|
//! Provides persistent memory storage for AI agents using HDF5 files.
|
||||||
//! All data is cached in-memory for fast access and flushed to disk
|
//! All data is cached in-memory for fast access and flushed to disk
|
||||||
@@ -36,6 +36,7 @@ pub mod reranker;
|
|||||||
pub mod schema;
|
pub mod schema;
|
||||||
pub mod search;
|
pub mod search;
|
||||||
pub mod session;
|
pub mod session;
|
||||||
|
pub mod signing;
|
||||||
pub mod storage;
|
pub mod storage;
|
||||||
mod store_lock;
|
mod store_lock;
|
||||||
pub mod temporal;
|
pub mod temporal;
|
||||||
@@ -62,26 +63,23 @@ use std::path::{Path, PathBuf};
|
|||||||
|
|
||||||
use cache::MemoryCache;
|
use cache::MemoryCache;
|
||||||
#[cfg(feature = "hnsw")]
|
#[cfg(feature = "hnsw")]
|
||||||
use clawhdf5_ann::{DistanceMetric, HnswIndex};
|
use clawhdf5_ann::{DistanceMetric, HnswIndex, Storage};
|
||||||
|
use clawhdf5_format::float16::round_to_f16;
|
||||||
use ephemeral::{EphemeralConfig, EphemeralStore};
|
use ephemeral::{EphemeralConfig, EphemeralStore};
|
||||||
|
|
||||||
/// HNSW construction parameters used for the agent's vector index. Cosine is the
|
|
||||||
/// agent's similarity metric, so the index is built with cosine distance.
|
|
||||||
#[cfg(feature = "hnsw")]
|
|
||||||
const HNSW_M: usize = 16;
|
|
||||||
#[cfg(feature = "hnsw")]
|
|
||||||
const HNSW_EF_CONSTRUCTION: usize = 64;
|
|
||||||
// EphemeralEntry and EphemeralStats are part of the crate public API via
|
// EphemeralEntry and EphemeralStats are part of the crate public API via
|
||||||
// the `ephemeral` module; they are not needed directly in lib.rs internals.
|
// the `ephemeral` module; they are not needed directly in lib.rs internals.
|
||||||
#[allow(unused_imports)]
|
#[allow(unused_imports)]
|
||||||
pub use ephemeral::{EphemeralEntry, EphemeralStats};
|
pub use ephemeral::{EphemeralEntry, EphemeralStats};
|
||||||
use knowledge::KnowledgeCache;
|
use knowledge::KnowledgeCache;
|
||||||
use memory_strategy::{Exchange, MemoryStrategy, StrategyOutput};
|
use memory_strategy::{Exchange, MemoryStrategy, StrategyOutput};
|
||||||
use session::SessionCache;
|
pub use search::SearchOptions;
|
||||||
|
pub use session::{SessionCache, SessionEntry};
|
||||||
|
|
||||||
// --- Error type ---
|
// --- Error type ---
|
||||||
|
|
||||||
#[derive(Debug)]
|
#[derive(Debug)]
|
||||||
|
#[non_exhaustive]
|
||||||
pub enum MemoryError {
|
pub enum MemoryError {
|
||||||
Io(std::io::Error),
|
Io(std::io::Error),
|
||||||
Hdf5(String),
|
Hdf5(String),
|
||||||
@@ -89,6 +87,14 @@ pub enum MemoryError {
|
|||||||
NotFound(String),
|
NotFound(String),
|
||||||
/// Another `HDF5Memory` (in this or another process) has the store open.
|
/// Another `HDF5Memory` (in this or another process) has the store open.
|
||||||
Locked(String),
|
Locked(String),
|
||||||
|
/// A record the store cannot hold as given, e.g. an embedding value
|
||||||
|
/// outside the half-precision range of a `float16` store.
|
||||||
|
InvalidEntry(String),
|
||||||
|
/// The store's checkpoints are signed and no signing key is set, so a
|
||||||
|
/// checkpoint would leave it unsigned. Set the key with
|
||||||
|
/// [`HDF5Memory::set_signing_key`], or drop the signature on purpose with
|
||||||
|
/// [`HDF5Memory::remove_signature`].
|
||||||
|
SigningKeyRequired(String),
|
||||||
}
|
}
|
||||||
|
|
||||||
impl std::fmt::Display for MemoryError {
|
impl std::fmt::Display for MemoryError {
|
||||||
@@ -99,6 +105,8 @@ impl std::fmt::Display for MemoryError {
|
|||||||
MemoryError::Schema(e) => write!(f, "schema error: {e}"),
|
MemoryError::Schema(e) => write!(f, "schema error: {e}"),
|
||||||
MemoryError::NotFound(e) => write!(f, "not found: {e}"),
|
MemoryError::NotFound(e) => write!(f, "not found: {e}"),
|
||||||
MemoryError::Locked(e) => write!(f, "store is locked: {e}"),
|
MemoryError::Locked(e) => write!(f, "store is locked: {e}"),
|
||||||
|
MemoryError::InvalidEntry(e) => write!(f, "invalid entry: {e}"),
|
||||||
|
MemoryError::SigningKeyRequired(e) => write!(f, "signing key required: {e}"),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -130,6 +138,19 @@ pub struct MemoryConfig {
|
|||||||
pub embedding_dim: usize,
|
pub embedding_dim: usize,
|
||||||
pub chunk_size: usize,
|
pub chunk_size: usize,
|
||||||
pub overlap: usize,
|
pub overlap: usize,
|
||||||
|
/// Store embeddings as IEEE half precision (numpy `float16`): half the
|
||||||
|
/// bytes of the embeddings dataset on disk. Every embedding is rounded to
|
||||||
|
/// the nearest half as it enters the store, in memory as well as on disk,
|
||||||
|
/// so search results are the same before and after a reopen. Values must
|
||||||
|
/// lie within ±65504; a save outside that is `MemoryError::InvalidEntry`.
|
||||||
|
/// Fixed when the store is created (persisted in `/meta`).
|
||||||
|
///
|
||||||
|
/// **On by default for new stores**: on the full LongMemEval haystack with
|
||||||
|
/// real MiniLM embeddings every retrieval metric matched `f32`, and at
|
||||||
|
/// 100K records the file is 48% smaller (`BENCHMARKS.md`). Existing
|
||||||
|
/// stores keep the setting they were created with. Set it to `false` for
|
||||||
|
/// full-precision embeddings, e.g. for unnormalised vectors that may
|
||||||
|
/// exceed the half-precision range.
|
||||||
pub float16: bool,
|
pub float16: bool,
|
||||||
pub compression: bool,
|
pub compression: bool,
|
||||||
pub compression_level: u32,
|
pub compression_level: u32,
|
||||||
@@ -139,6 +160,40 @@ pub struct MemoryConfig {
|
|||||||
pub created_at: String,
|
pub created_at: String,
|
||||||
pub wal_enabled: bool,
|
pub wal_enabled: bool,
|
||||||
pub wal_max_entries: usize,
|
pub wal_max_entries: usize,
|
||||||
|
/// Store the vector index's own copy of the embeddings as int8 rather than
|
||||||
|
/// f32, a quarter of the memory. **On by default** for new stores.
|
||||||
|
///
|
||||||
|
/// The index's copy is the single largest part of a loaded store's
|
||||||
|
/// footprint. Quantised distances are approximate, so the candidate pool
|
||||||
|
/// is re-scored against the cache's exact embeddings before fusion, which
|
||||||
|
/// holds recall at the f32 index's level. It is also faster, not slower:
|
||||||
|
/// at equal recall, 1.63x the queries per second on x86-64 (AVX2) and
|
||||||
|
/// 1.18x on a Raspberry Pi 5 (NEON `SDOT`), with builds 1.8x and 2.3x
|
||||||
|
/// faster. See `BENCHMARKS.md`.
|
||||||
|
///
|
||||||
|
/// Persisted with the store. Stores written before this setting existed
|
||||||
|
/// have no stored value and open as `false`, so reopening an old store
|
||||||
|
/// never changes how its index is held.
|
||||||
|
///
|
||||||
|
/// Has no effect without the `hnsw` feature.
|
||||||
|
pub quantized_index: bool,
|
||||||
|
/// HNSW graph degree. Higher means a denser graph: better recall, more
|
||||||
|
/// memory and slower builds. Clamped to at least 2 when the index is
|
||||||
|
/// built, since a graph with fewer connections is not one.
|
||||||
|
///
|
||||||
|
/// Has no effect without the `hnsw` feature.
|
||||||
|
pub hnsw_m: usize,
|
||||||
|
/// Candidate list size while building the HNSW graph. Higher means a
|
||||||
|
/// better graph and a slower build; it does not affect query cost.
|
||||||
|
///
|
||||||
|
/// Has no effect without the `hnsw` feature.
|
||||||
|
pub hnsw_ef_construction: usize,
|
||||||
|
/// Candidate list size for a query, trading throughput for recall. `0`
|
||||||
|
/// keeps the default, which scales with the requested `k`
|
||||||
|
/// (`max(k * 8, 64)`) so that fusion still sees a useful pool.
|
||||||
|
///
|
||||||
|
/// Has no effect without the `hnsw` feature.
|
||||||
|
pub hnsw_ef_search: usize,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl MemoryConfig {
|
impl MemoryConfig {
|
||||||
@@ -151,7 +206,7 @@ impl MemoryConfig {
|
|||||||
embedding_dim,
|
embedding_dim,
|
||||||
chunk_size: 512,
|
chunk_size: 512,
|
||||||
overlap: 50,
|
overlap: 50,
|
||||||
float16: false,
|
float16: true,
|
||||||
compression: false,
|
compression: false,
|
||||||
compression_level: 0,
|
compression_level: 0,
|
||||||
compact_threshold: 0.3,
|
compact_threshold: 0.3,
|
||||||
@@ -160,6 +215,10 @@ impl MemoryConfig {
|
|||||||
created_at,
|
created_at,
|
||||||
wal_enabled: true,
|
wal_enabled: true,
|
||||||
wal_max_entries: 500,
|
wal_max_entries: 500,
|
||||||
|
quantized_index: true,
|
||||||
|
hnsw_m: 16,
|
||||||
|
hnsw_ef_construction: 64,
|
||||||
|
hnsw_ef_search: 0,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -258,11 +317,20 @@ pub struct HDF5Memory {
|
|||||||
/// every record, on every single query. Built lazily on first use; see
|
/// every record, on every single query. Built lazily on first use; see
|
||||||
/// [`HDF5Memory::ensure_bm25_fresh`] for how it stays in sync.
|
/// [`HDF5Memory::ensure_bm25_fresh`] for how it stays in sync.
|
||||||
bm25: Option<bm25::BM25Index>,
|
bm25: Option<bm25::BM25Index>,
|
||||||
|
/// Token filter the keyword index is built with. Changing it drops the
|
||||||
|
/// index; it is not persisted, because the index is not either.
|
||||||
|
bm25_filter: bm25::TokenFilter,
|
||||||
/// Activation weights changed since the last checkpoint (searches boost
|
/// Activation weights changed since the last checkpoint (searches boost
|
||||||
/// the records they return). Cleared by `flush`.
|
/// the records they return). Cleared by `flush`.
|
||||||
activations_dirty: bool,
|
activations_dirty: bool,
|
||||||
/// Opened with [`HDF5Memory::open_read_only`]: nothing may reach the disk.
|
/// Opened with [`HDF5Memory::open_read_only`]: nothing may reach the disk.
|
||||||
read_only: bool,
|
read_only: bool,
|
||||||
|
/// Key that signs every checkpoint; never persisted. See
|
||||||
|
/// [`HDF5Memory::set_signing_key`].
|
||||||
|
signing_key: Option<signing::SigningKey>,
|
||||||
|
/// Checkpoints of this store are signed: the file on disk is, or a key
|
||||||
|
/// has been set. A checkpoint without a key is then refused.
|
||||||
|
signed: bool,
|
||||||
/// A WAL that `open()` could not read and moved aside; see
|
/// A WAL that `open()` could not read and moved aside; see
|
||||||
/// [`HDF5Memory::quarantined_wal`].
|
/// [`HDF5Memory::quarantined_wal`].
|
||||||
quarantined_wal: Option<PathBuf>,
|
quarantined_wal: Option<PathBuf>,
|
||||||
@@ -282,7 +350,8 @@ impl HDF5Memory {
|
|||||||
/// Create a new HDF5 memory file with the given configuration.
|
/// Create a new HDF5 memory file with the given configuration.
|
||||||
pub fn create(config: MemoryConfig) -> Result<Self> {
|
pub fn create(config: MemoryConfig) -> Result<Self> {
|
||||||
let lock = store_lock::StoreLock::acquire(&config.path)?;
|
let lock = store_lock::StoreLock::acquire(&config.path)?;
|
||||||
let cache = MemoryCache::new(config.embedding_dim);
|
let mut cache = MemoryCache::new(config.embedding_dim);
|
||||||
|
cache.set_half_precision(config.float16);
|
||||||
let sessions = SessionCache::new();
|
let sessions = SessionCache::new();
|
||||||
let knowledge = KnowledgeCache::new();
|
let knowledge = KnowledgeCache::new();
|
||||||
|
|
||||||
@@ -314,8 +383,11 @@ impl HDF5Memory {
|
|||||||
anomaly: anomaly::WriteAnomalyDetector::new(anomaly::AnomalyConfig::default()),
|
anomaly: anomaly::WriteAnomalyDetector::new(anomaly::AnomalyConfig::default()),
|
||||||
anomaly_alerts: Vec::new(),
|
anomaly_alerts: Vec::new(),
|
||||||
bm25: None,
|
bm25: None,
|
||||||
|
bm25_filter: bm25::TokenFilter::default(),
|
||||||
activations_dirty: false,
|
activations_dirty: false,
|
||||||
read_only: false,
|
read_only: false,
|
||||||
|
signing_key: None,
|
||||||
|
signed: false,
|
||||||
quarantined_wal: None,
|
quarantined_wal: None,
|
||||||
_lock: Some(lock),
|
_lock: Some(lock),
|
||||||
})
|
})
|
||||||
@@ -440,7 +512,17 @@ impl HDF5Memory {
|
|||||||
|
|
||||||
#[cfg(feature = "hnsw")]
|
#[cfg(feature = "hnsw")]
|
||||||
let loaded_index = if replay_only_appended {
|
let loaded_index = if replay_only_appended {
|
||||||
Self::load_vector_index(path, checkpoint.ann_generation, &cache, n_checkpoint)
|
Self::load_vector_index(
|
||||||
|
path,
|
||||||
|
checkpoint.ann_generation,
|
||||||
|
&cache,
|
||||||
|
n_checkpoint,
|
||||||
|
if config.quantized_index {
|
||||||
|
Storage::Int8
|
||||||
|
} else {
|
||||||
|
Storage::Float32
|
||||||
|
},
|
||||||
|
)
|
||||||
} else {
|
} else {
|
||||||
None
|
None
|
||||||
};
|
};
|
||||||
@@ -481,8 +563,11 @@ impl HDF5Memory {
|
|||||||
anomaly: anomaly::WriteAnomalyDetector::new(anomaly::AnomalyConfig::default()),
|
anomaly: anomaly::WriteAnomalyDetector::new(anomaly::AnomalyConfig::default()),
|
||||||
anomaly_alerts: Vec::new(),
|
anomaly_alerts: Vec::new(),
|
||||||
bm25: None,
|
bm25: None,
|
||||||
|
bm25_filter: bm25::TokenFilter::default(),
|
||||||
activations_dirty: false,
|
activations_dirty: false,
|
||||||
read_only,
|
read_only,
|
||||||
|
signing_key: None,
|
||||||
|
signed: checkpoint.signed,
|
||||||
quarantined_wal,
|
quarantined_wal,
|
||||||
_lock: lock,
|
_lock: lock,
|
||||||
})
|
})
|
||||||
@@ -553,6 +638,7 @@ impl HDF5Memory {
|
|||||||
generation: Option<u64>,
|
generation: Option<u64>,
|
||||||
cache: &MemoryCache,
|
cache: &MemoryCache,
|
||||||
n_checkpoint: usize,
|
n_checkpoint: usize,
|
||||||
|
storage: Storage,
|
||||||
) -> Option<HnswIndex> {
|
) -> Option<HnswIndex> {
|
||||||
let generation = generation?;
|
let generation = generation?;
|
||||||
let bytes = std::fs::read(Self::vector_index_path(store)).ok()?;
|
let bytes = std::fs::read(Self::vector_index_path(store)).ok()?;
|
||||||
@@ -560,15 +646,17 @@ impl HDF5Memory {
|
|||||||
if u64::from_le_bytes(stamp.try_into().ok()?) != generation {
|
if u64::from_le_bytes(stamp.try_into().ok()?) != generation {
|
||||||
return None;
|
return None;
|
||||||
}
|
}
|
||||||
let vectors = cache.embeddings.get(..n_checkpoint)?.to_vec();
|
let vectors: Vec<Vec<f32>> = (0..n_checkpoint)
|
||||||
let mut index = HnswIndex::from_graph_bytes(graph, vectors).ok()?;
|
.map(|i| cache.embeddings.get(i).map(<[f32]>::to_vec))
|
||||||
|
.collect::<Option<_>>()?;
|
||||||
|
let mut index = HnswIndex::from_graph_bytes_with(graph, vectors, storage).ok()?;
|
||||||
if index.dimension() != cache.embedding_dim {
|
if index.dimension() != cache.embedding_dim {
|
||||||
return None;
|
return None;
|
||||||
}
|
}
|
||||||
// Records appended since (replayed from the WAL) join incrementally.
|
// Records appended since (replayed from the WAL) join incrementally.
|
||||||
for id in n_checkpoint..cache.embeddings.len() {
|
for id in n_checkpoint..cache.embeddings.len() {
|
||||||
if cache.embeddings[id].len() != index.dimension()
|
if cache.embeddings[id].len() != index.dimension()
|
||||||
|| index.insert(cache.embeddings[id].clone()) != id
|
|| index.insert(cache.embeddings[id].to_vec()) != id
|
||||||
{
|
{
|
||||||
return None;
|
return None;
|
||||||
}
|
}
|
||||||
@@ -591,7 +679,7 @@ impl HDF5Memory {
|
|||||||
pub(crate) fn ensure_bm25_fresh(&mut self) -> &bm25::BM25Index {
|
pub(crate) fn ensure_bm25_fresh(&mut self) -> &bm25::BM25Index {
|
||||||
let n = self.cache.chunks.len();
|
let n = self.cache.chunks.len();
|
||||||
let bm25 = match self.bm25.take() {
|
let bm25 = match self.bm25.take() {
|
||||||
Some(index) if index.len() <= n => {
|
Some(index) if index.len() <= n && index.token_filter() == self.bm25_filter => {
|
||||||
let mut index = index;
|
let mut index = index;
|
||||||
for id in index.len()..n {
|
for id in index.len()..n {
|
||||||
if self.cache.tombstones[id] == 0 {
|
if self.cache.tombstones[id] == 0 {
|
||||||
@@ -601,11 +689,26 @@ impl HDF5Memory {
|
|||||||
index.pad_to(n);
|
index.pad_to(n);
|
||||||
index
|
index
|
||||||
}
|
}
|
||||||
_ => bm25::BM25Index::build(&self.cache.chunks, &self.cache.tombstones),
|
_ => bm25::BM25Index::build_with(
|
||||||
|
&self.cache.chunks,
|
||||||
|
&self.cache.tombstones,
|
||||||
|
self.bm25_filter,
|
||||||
|
),
|
||||||
};
|
};
|
||||||
self.bm25.insert(bm25)
|
self.bm25.insert(bm25)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Choose how keyword-search tokens are normalised, rebuilding the index
|
||||||
|
/// on next use. [`bm25::TokenFilter::Stemmed`] matches inflections of the
|
||||||
|
/// same word at some cost in precision; measure before adopting it (see
|
||||||
|
/// `BENCHMARKS.md`).
|
||||||
|
pub fn set_token_filter(&mut self, filter: bm25::TokenFilter) {
|
||||||
|
if filter != self.bm25_filter {
|
||||||
|
self.bm25_filter = filter;
|
||||||
|
self.bm25 = None;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
/// Record `id` was tombstoned; its text is still in the cache.
|
/// Record `id` was tombstoned; its text is still in the cache.
|
||||||
fn bm25_on_delete(&mut self, id: usize) {
|
fn bm25_on_delete(&mut self, id: usize) {
|
||||||
if let Some(index) = self.bm25.as_mut()
|
if let Some(index) = self.bm25.as_mut()
|
||||||
@@ -625,6 +728,39 @@ impl HDF5Memory {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Sign every checkpoint from now on with `key` (Ed25519). The key is
|
||||||
|
/// never written anywhere; set it again after every `open`. Once a store
|
||||||
|
/// is signed, a checkpoint without the key is refused
|
||||||
|
/// ([`MemoryError::SigningKeyRequired`]) rather than silently leaving it
|
||||||
|
/// unsigned. Setting a different key re-signs the store under that key
|
||||||
|
/// from the next checkpoint; a verifier trusting the old key will then
|
||||||
|
/// reject it, which is the point. Call [`AgentMemory::flush_wal`] to sign
|
||||||
|
/// right away.
|
||||||
|
pub fn set_signing_key(&mut self, key: signing::SigningKey) {
|
||||||
|
self.signing_key = Some(key);
|
||||||
|
self.signed = true;
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Stop signing: the next checkpoint writes the store unsigned. The
|
||||||
|
/// deliberate way out of [`MemoryError::SigningKeyRequired`].
|
||||||
|
pub fn remove_signature(&mut self) {
|
||||||
|
self.signing_key = None;
|
||||||
|
self.signed = false;
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Checkpoints of this store are signed (on disk, or from the next
|
||||||
|
/// checkpoint because a key has been set).
|
||||||
|
pub fn is_signed(&self) -> bool {
|
||||||
|
self.signed
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Check the checkpoint at `path` against the public key the caller
|
||||||
|
/// trusts; see [`signing::verify_store`]. Reads the file only: it works
|
||||||
|
/// on a store another process has open.
|
||||||
|
pub fn verify(path: &Path, trusted: &signing::VerifyingKey) -> Result<signing::VerifyReport> {
|
||||||
|
signing::verify_store(path, trusted)
|
||||||
|
}
|
||||||
|
|
||||||
/// Flush current state to disk and truncate the WAL.
|
/// Flush current state to disk and truncate the WAL.
|
||||||
///
|
///
|
||||||
/// Every code path that persists the full cache to the .h5 file must
|
/// Every code path that persists the full cache to the .h5 file must
|
||||||
@@ -640,10 +776,28 @@ impl HDF5Memory {
|
|||||||
// Record which WAL prefix this checkpoint contains, so a crash before
|
// Record which WAL prefix this checkpoint contains, so a crash before
|
||||||
// the truncate below can't replay those entries a second time.
|
// the truncate below can't replay those entries a second time.
|
||||||
let wal_applied = self.wal.as_ref().map(|w| w.mark());
|
let wal_applied = self.wal.as_ref().map(|w| w.mark());
|
||||||
|
let signature = match &self.signing_key {
|
||||||
|
Some(key) => Some(signing::sign(
|
||||||
|
key,
|
||||||
|
&self.config,
|
||||||
|
&self.cache,
|
||||||
|
&self.sessions,
|
||||||
|
&self.knowledge,
|
||||||
|
wal_applied,
|
||||||
|
)),
|
||||||
|
None if self.signed => {
|
||||||
|
return Err(MemoryError::SigningKeyRequired(format!(
|
||||||
|
"{} is signed; set its signing key before a checkpoint \
|
||||||
|
(saves so far are held in the WAL or in memory)",
|
||||||
|
self.config.path.display()
|
||||||
|
)));
|
||||||
|
}
|
||||||
|
None => None,
|
||||||
|
};
|
||||||
// Written before the .h5 so a crash in between leaves a sidecar whose
|
// Written before the .h5 so a crash in between leaves a sidecar whose
|
||||||
// generation matches no checkpoint (ignored), never the reverse.
|
// generation matches no checkpoint (ignored), never the reverse.
|
||||||
let ann_generation = self.persist_vector_index();
|
let ann_generation = self.persist_vector_index();
|
||||||
storage::write_to_disk_with_meta(
|
storage::write_to_disk_signed(
|
||||||
&self.config.path,
|
&self.config.path,
|
||||||
&self.config,
|
&self.config,
|
||||||
&self.cache,
|
&self.cache,
|
||||||
@@ -652,7 +806,9 @@ impl HDF5Memory {
|
|||||||
&schema::CheckpointMeta {
|
&schema::CheckpointMeta {
|
||||||
wal_applied,
|
wal_applied,
|
||||||
ann_generation,
|
ann_generation,
|
||||||
|
signed: signature.is_some(),
|
||||||
},
|
},
|
||||||
|
signature.as_ref(),
|
||||||
)?;
|
)?;
|
||||||
if let Some(ref mut w) = self.wal {
|
if let Some(ref mut w) = self.wal {
|
||||||
w.truncate()?;
|
w.truncate()?;
|
||||||
@@ -781,6 +937,42 @@ impl HDF5Memory {
|
|||||||
// the index length drifts from the cache length (covering any mutation path
|
// the index length drifts from the cache length (covering any mutation path
|
||||||
// that doesn't call a hook, e.g. consolidation pushes).
|
// that doesn't call a hook, e.g. consolidation pushes).
|
||||||
|
|
||||||
|
/// Graph degree for the index, never below the 2 the builder requires:
|
||||||
|
/// a config value of 0 or 1 would otherwise panic inside `clawhdf5-ann`.
|
||||||
|
#[cfg(feature = "hnsw")]
|
||||||
|
fn hnsw_m(&self) -> usize {
|
||||||
|
self.config.hnsw_m.max(2)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Build-time candidate list size, never below the graph degree — a
|
||||||
|
/// smaller one cannot fill a node's connections.
|
||||||
|
#[cfg(feature = "hnsw")]
|
||||||
|
fn hnsw_ef_construction(&self) -> usize {
|
||||||
|
self.config.hnsw_ef_construction.max(self.hnsw_m())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Query-time candidate list size for a `k`-result search. `0` means the
|
||||||
|
/// default, which scales with `k`.
|
||||||
|
#[cfg(feature = "hnsw")]
|
||||||
|
pub(crate) fn hnsw_ef_search(&self, k: usize) -> usize {
|
||||||
|
let default = (k * 8).max(64);
|
||||||
|
if self.config.hnsw_ef_search == 0 {
|
||||||
|
default
|
||||||
|
} else {
|
||||||
|
self.config.hnsw_ef_search.max(k)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// How the index should store its copy of the vectors, per the config.
|
||||||
|
#[cfg(feature = "hnsw")]
|
||||||
|
fn index_storage(&self) -> Storage {
|
||||||
|
if self.config.quantized_index {
|
||||||
|
Storage::Int8
|
||||||
|
} else {
|
||||||
|
Storage::Float32
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
/// Build an HNSW index over the entire cache, re-applying tombstones as
|
/// Build an HNSW index over the entire cache, re-applying tombstones as
|
||||||
/// soft-deletions so node ids stay aligned with cache indices.
|
/// soft-deletions so node ids stay aligned with cache indices.
|
||||||
///
|
///
|
||||||
@@ -796,11 +988,15 @@ impl HDF5Memory {
|
|||||||
if self.cache.embeddings.iter().any(|e| e.len() != dim) {
|
if self.cache.embeddings.iter().any(|e| e.len() != dim) {
|
||||||
return None;
|
return None;
|
||||||
}
|
}
|
||||||
let mut index = HnswIndex::build_with_metric(
|
// The index owns its vectors, so it needs rows rather than the cache's
|
||||||
&self.cache.embeddings,
|
// flat buffer. This copy is the index's own; the cache keeps one.
|
||||||
HNSW_M,
|
let rows: Vec<Vec<f32>> = self.cache.embeddings.iter().map(<[f32]>::to_vec).collect();
|
||||||
HNSW_EF_CONSTRUCTION,
|
let mut index = HnswIndex::build_with(
|
||||||
|
&rows,
|
||||||
|
self.hnsw_m(),
|
||||||
|
self.hnsw_ef_construction(),
|
||||||
DistanceMetric::Cosine,
|
DistanceMetric::Cosine,
|
||||||
|
self.index_storage(),
|
||||||
);
|
);
|
||||||
for (i, &t) in self.cache.tombstones.iter().enumerate() {
|
for (i, &t) in self.cache.tombstones.iter().enumerate() {
|
||||||
if t != 0 {
|
if t != 0 {
|
||||||
@@ -826,7 +1022,7 @@ impl HDF5Memory {
|
|||||||
let dim = index.dimension();
|
let dim = index.dimension();
|
||||||
let appended = (self.hnsw_synced_len..n).all(|id| {
|
let appended = (self.hnsw_synced_len..n).all(|id| {
|
||||||
self.cache.embeddings[id].len() == dim
|
self.cache.embeddings[id].len() == dim
|
||||||
&& index.insert(self.cache.embeddings[id].clone()) == id
|
&& index.insert(self.cache.embeddings[id].to_vec()) == id
|
||||||
});
|
});
|
||||||
if appended {
|
if appended {
|
||||||
for id in self.hnsw_synced_len..n {
|
for id in self.hnsw_synced_len..n {
|
||||||
@@ -857,7 +1053,7 @@ impl HDF5Memory {
|
|||||||
let emb_len = self.cache.embeddings[idx].len();
|
let emb_len = self.cache.embeddings[idx].len();
|
||||||
match self.hnsw.as_mut() {
|
match self.hnsw.as_mut() {
|
||||||
Some(index) if emb_len == index.dimension() => {
|
Some(index) if emb_len == index.dimension() => {
|
||||||
let id = index.insert(self.cache.embeddings[idx].clone());
|
let id = index.insert(self.cache.embeddings[idx].to_vec());
|
||||||
if id == idx {
|
if id == idx {
|
||||||
self.hnsw_synced_len = self.cache.embeddings.len();
|
self.hnsw_synced_len = self.cache.embeddings.len();
|
||||||
} else {
|
} else {
|
||||||
@@ -903,6 +1099,18 @@ impl HDF5Memory {
|
|||||||
&self.config
|
&self.config
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// The sessions recorded in this store.
|
||||||
|
pub fn sessions(&self) -> &SessionCache {
|
||||||
|
&self.sessions
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Mutable access to the sessions, e.g. to add many at once. Changes
|
||||||
|
/// reach the disk at the next checkpoint (any flushing call, such as
|
||||||
|
/// [`HDF5Memory::flush_wal`] or `save_batch`), not immediately.
|
||||||
|
pub fn sessions_mut(&mut self) -> &mut SessionCache {
|
||||||
|
&mut self.sessions
|
||||||
|
}
|
||||||
|
|
||||||
/// Get a reference to the knowledge cache.
|
/// Get a reference to the knowledge cache.
|
||||||
pub fn knowledge(&self) -> &KnowledgeCache {
|
pub fn knowledge(&self) -> &KnowledgeCache {
|
||||||
&self.knowledge
|
&self.knowledge
|
||||||
@@ -965,7 +1173,29 @@ impl HDF5Memory {
|
|||||||
/// Upsert: if an active entry with the same tags (key) exists, update it in-place.
|
/// Upsert: if an active entry with the same tags (key) exists, update it in-place.
|
||||||
/// Otherwise append a new entry. Use this for key-based memory stores where
|
/// Otherwise append a new entry. Use this for key-based memory stores where
|
||||||
/// the same key should not create duplicates.
|
/// the same key should not create duplicates.
|
||||||
|
/// A `float16` store holds embeddings as IEEE half precision, which has no
|
||||||
|
/// finite value beyond ±65504. Refuse such an embedding rather than
|
||||||
|
/// silently store infinity. (Values that are already infinite or NaN are
|
||||||
|
/// stored as they are, as in an `f32` store.)
|
||||||
|
fn check_embedding(&self, embedding: &[f32]) -> Result<()> {
|
||||||
|
if !self.config.float16 {
|
||||||
|
return Ok(());
|
||||||
|
}
|
||||||
|
let overflow = embedding
|
||||||
|
.iter()
|
||||||
|
.enumerate()
|
||||||
|
.find(|&(_, &v)| v.is_finite() && round_to_f16(v).is_infinite());
|
||||||
|
match overflow {
|
||||||
|
None => Ok(()),
|
||||||
|
Some((i, v)) => Err(MemoryError::InvalidEntry(format!(
|
||||||
|
"embedding[{i}] = {v} is outside the half-precision range (±65504) \
|
||||||
|
of this float16 store"
|
||||||
|
))),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
pub fn save_or_update(&mut self, entry: MemoryEntry) -> Result<usize> {
|
pub fn save_or_update(&mut self, entry: MemoryEntry) -> Result<usize> {
|
||||||
|
self.check_embedding(&entry.embedding)?;
|
||||||
if let Some(existing_idx) = self.cache.find_by_tags(&entry.tags) {
|
if let Some(existing_idx) = self.cache.find_by_tags(&entry.tags) {
|
||||||
if let Some(ref mut w) = self.wal {
|
if let Some(ref mut w) = self.wal {
|
||||||
let wal_entry = wal::WalEntry {
|
let wal_entry = wal::WalEntry {
|
||||||
@@ -1021,6 +1251,7 @@ impl HDF5Memory {
|
|||||||
|
|
||||||
impl AgentMemory for HDF5Memory {
|
impl AgentMemory for HDF5Memory {
|
||||||
fn save(&mut self, entry: MemoryEntry) -> Result<usize> {
|
fn save(&mut self, entry: MemoryEntry) -> Result<usize> {
|
||||||
|
self.check_embedding(&entry.embedding)?;
|
||||||
if let Some(ref mut w) = self.wal {
|
if let Some(ref mut w) = self.wal {
|
||||||
let wal_entry = wal::WalEntry {
|
let wal_entry = wal::WalEntry {
|
||||||
entry_type: wal::WalEntryType::Save,
|
entry_type: wal::WalEntryType::Save,
|
||||||
@@ -1062,6 +1293,10 @@ impl AgentMemory for HDF5Memory {
|
|||||||
}
|
}
|
||||||
|
|
||||||
fn save_batch(&mut self, entries: Vec<MemoryEntry>) -> Result<Vec<usize>> {
|
fn save_batch(&mut self, entries: Vec<MemoryEntry>) -> Result<Vec<usize>> {
|
||||||
|
// All or nothing: check every entry before storing any.
|
||||||
|
for entry in &entries {
|
||||||
|
self.check_embedding(&entry.embedding)?;
|
||||||
|
}
|
||||||
let mut indices = Vec::with_capacity(entries.len());
|
let mut indices = Vec::with_capacity(entries.len());
|
||||||
for entry in entries {
|
for entry in entries {
|
||||||
let idx = self.cache.push(
|
let idx = self.cache.push(
|
||||||
@@ -1222,6 +1457,9 @@ impl HDF5Memory {
|
|||||||
})?;
|
})?;
|
||||||
let view = memory_strategy::CacheStoreView::new(&self.cache, &self.knowledge);
|
let view = memory_strategy::CacheStoreView::new(&self.cache, &self.knowledge);
|
||||||
let output = strat.evaluate(&exchange, &view);
|
let output = strat.evaluate(&exchange, &view);
|
||||||
|
for e in &output.entries {
|
||||||
|
self.check_embedding(&e.embedding)?;
|
||||||
|
}
|
||||||
for e in &output.entries {
|
for e in &output.entries {
|
||||||
self.cache.push(
|
self.cache.push(
|
||||||
e.chunk.clone(),
|
e.chunk.clone(),
|
||||||
@@ -1246,6 +1484,35 @@ impl HDF5Memory {
|
|||||||
}
|
}
|
||||||
|
|
||||||
impl HDF5Memory {
|
impl HDF5Memory {
|
||||||
|
/// Delete many records with a single checkpoint, where
|
||||||
|
/// [`AgentMemory::delete`] checkpoints once per record.
|
||||||
|
///
|
||||||
|
/// All or nothing: if any id is out of range or already deleted (or
|
||||||
|
/// repeated), nothing is deleted and `MemoryError::NotFound` is returned.
|
||||||
|
/// Unlike `delete`, this never auto-compacts, so the records stay in the
|
||||||
|
/// store as tombstones (their indices unchanged) until [`AgentMemory::compact`]
|
||||||
|
/// is called — importers use it to carry over records that were already
|
||||||
|
/// deleted in the source.
|
||||||
|
pub fn delete_batch(&mut self, ids: &[usize]) -> Result<()> {
|
||||||
|
let mut seen = std::collections::HashSet::with_capacity(ids.len());
|
||||||
|
for &id in ids {
|
||||||
|
if self.cache.tombstones.get(id).copied() != Some(0) || !seen.insert(id) {
|
||||||
|
return Err(MemoryError::NotFound(format!(
|
||||||
|
"entry {id} not found or already deleted"
|
||||||
|
)));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if ids.is_empty() {
|
||||||
|
return Ok(());
|
||||||
|
}
|
||||||
|
for &id in ids {
|
||||||
|
self.cache.mark_deleted(id);
|
||||||
|
self.hnsw_on_delete(id);
|
||||||
|
self.bm25_on_delete(id);
|
||||||
|
}
|
||||||
|
self.flush()
|
||||||
|
}
|
||||||
|
|
||||||
pub fn tick_session(&mut self) -> Result<()> {
|
pub fn tick_session(&mut self) -> Result<()> {
|
||||||
let d = self.config.decay_factor;
|
let d = self.config.decay_factor;
|
||||||
for w in self.cache.activation_weights.iter_mut() {
|
for w in self.cache.activation_weights.iter_mut() {
|
||||||
@@ -1306,6 +1573,16 @@ impl HDF5Memory {
|
|||||||
let mut promoted = 0;
|
let mut promoted = 0;
|
||||||
|
|
||||||
for key in candidates {
|
for key in candidates {
|
||||||
|
// Check before taking, so a rejected entry stays in the ephemeral
|
||||||
|
// tier rather than being lost.
|
||||||
|
if let Some(emb) = self
|
||||||
|
.ephemeral
|
||||||
|
.as_ref()
|
||||||
|
.and_then(|s| s.get_entry(&key))
|
||||||
|
.and_then(|e| e.embedding.as_deref())
|
||||||
|
{
|
||||||
|
self.check_embedding(emb)?;
|
||||||
|
}
|
||||||
let entry = match self
|
let entry = match self
|
||||||
.ephemeral
|
.ephemeral
|
||||||
.as_mut()
|
.as_mut()
|
||||||
@@ -1350,7 +1627,8 @@ impl HDF5Memory {
|
|||||||
k: usize,
|
k: usize,
|
||||||
) -> Vec<SearchResult> {
|
) -> Vec<SearchResult> {
|
||||||
// Persistent tier.
|
// Persistent tier.
|
||||||
let persistent = self.hybrid_search(query_embedding, query_text, 0.7, 0.3, k);
|
let persistent =
|
||||||
|
self.hybrid_search_with(query_embedding, query_text, hybrid::DEFAULT_FUSION, k);
|
||||||
const EPHEMERAL_BOOST: f32 = 1.2;
|
const EPHEMERAL_BOOST: f32 = 1.2;
|
||||||
let mut results = persistent;
|
let mut results = persistent;
|
||||||
|
|
||||||
@@ -1433,6 +1711,79 @@ mod tests {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn delete_batch_tombstones_without_compacting() {
|
||||||
|
let dir = TempDir::new().unwrap();
|
||||||
|
let path = dir.path().join("test.h5");
|
||||||
|
let mut mem = HDF5Memory::create(make_config(&dir)).unwrap();
|
||||||
|
mem.save_batch(
|
||||||
|
(0..4)
|
||||||
|
.map(|i| make_entry(&format!("record {i}"), &[i as f32, 1.0, 0.0, 0.0]))
|
||||||
|
.collect(),
|
||||||
|
)
|
||||||
|
.unwrap();
|
||||||
|
// 3 of 4 is far past compact_threshold (0.3): delete() would compact.
|
||||||
|
mem.delete_batch(&[0, 1, 3]).unwrap();
|
||||||
|
assert_eq!(mem.count(), 4);
|
||||||
|
assert_eq!(mem.count_active(), 1);
|
||||||
|
drop(mem);
|
||||||
|
|
||||||
|
let mut mem = HDF5Memory::open(&path).unwrap();
|
||||||
|
assert_eq!(mem.cache.tombstones, vec![1, 1, 0, 1]);
|
||||||
|
let hits = mem.hybrid_search(&[0.0, 1.0, 0.0, 0.0], "record", 0.5, 0.5, 10);
|
||||||
|
assert!(
|
||||||
|
hits.iter().all(|r| r.index == 2),
|
||||||
|
"tombstoned record returned"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn delete_batch_is_all_or_nothing() {
|
||||||
|
let dir = TempDir::new().unwrap();
|
||||||
|
let mut mem = HDF5Memory::create(make_config(&dir)).unwrap();
|
||||||
|
mem.save_batch(vec![
|
||||||
|
make_entry("a", &[1.0, 0.0, 0.0, 0.0]),
|
||||||
|
make_entry("b", &[0.0, 1.0, 0.0, 0.0]),
|
||||||
|
])
|
||||||
|
.unwrap();
|
||||||
|
for bad in [&[0, 5][..], &[1, 1][..]] {
|
||||||
|
assert!(matches!(
|
||||||
|
mem.delete_batch(bad),
|
||||||
|
Err(MemoryError::NotFound(_))
|
||||||
|
));
|
||||||
|
assert_eq!(mem.count_active(), 2, "{bad:?} deleted something");
|
||||||
|
}
|
||||||
|
mem.delete_batch(&[]).unwrap();
|
||||||
|
assert_eq!(mem.count_active(), 2);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn sessions_mut_add_at_keeps_timestamp_across_reopen() {
|
||||||
|
let dir = TempDir::new().unwrap();
|
||||||
|
let path = dir.path().join("test.h5");
|
||||||
|
let mut mem = HDF5Memory::create(make_config(&dir)).unwrap();
|
||||||
|
mem.sessions_mut()
|
||||||
|
.add_at("s-old", 2, 7, "discord", "old summary", 1.7e15);
|
||||||
|
mem.flush_wal().unwrap();
|
||||||
|
drop(mem);
|
||||||
|
|
||||||
|
let mem = HDF5Memory::open_read_only(&path).unwrap();
|
||||||
|
let s = mem.sessions();
|
||||||
|
assert_eq!(s.len(), 1);
|
||||||
|
let e = &s.entries[0];
|
||||||
|
assert_eq!(
|
||||||
|
(
|
||||||
|
e.id.as_str(),
|
||||||
|
e.start_idx,
|
||||||
|
e.end_idx,
|
||||||
|
e.channel.as_str(),
|
||||||
|
e.ts
|
||||||
|
),
|
||||||
|
("s-old", 2, 7, "discord", 1.7e15)
|
||||||
|
);
|
||||||
|
assert_eq!(s.summaries[0], "old summary");
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn create_new_file() {
|
fn create_new_file() {
|
||||||
let dir = TempDir::new().unwrap();
|
let dir = TempDir::new().unwrap();
|
||||||
@@ -1954,6 +2305,34 @@ mod tests {
|
|||||||
assert_eq!(top_ids(&mut reopened, &q), expected_after);
|
assert_eq!(top_ids(&mut reopened, &q), expected_after);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn set_token_filter_rebuilds_the_keyword_index() {
|
||||||
|
let dir = TempDir::new().unwrap();
|
||||||
|
let mut mem = HDF5Memory::create(make_config(&dir)).unwrap();
|
||||||
|
mem.save(make_entry(
|
||||||
|
"I was training for a marathon",
|
||||||
|
&[1.0, 0.0, 0.0, 0.0],
|
||||||
|
))
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
// Count only genuine keyword matches: `hybrid_search` also returns
|
||||||
|
// zero-score filler when fewer than k records are relevant.
|
||||||
|
let hits = |mem: &mut HDF5Memory| {
|
||||||
|
mem.hybrid_search(&[0.0, 0.0, 0.0, 0.0], "trains", 0.0, 1.0, 5)
|
||||||
|
.iter()
|
||||||
|
.filter(|r| r.score > 0.0)
|
||||||
|
.count()
|
||||||
|
};
|
||||||
|
assert_eq!(hits(&mut mem), 0);
|
||||||
|
|
||||||
|
mem.set_token_filter(bm25::TokenFilter::Stemmed);
|
||||||
|
assert_eq!(hits(&mut mem), 1, "index should have been rebuilt stemmed");
|
||||||
|
|
||||||
|
// And back, rebuilding again.
|
||||||
|
mem.set_token_filter(bm25::TokenFilter::Plain);
|
||||||
|
assert_eq!(hits(&mut mem), 0);
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn keyword_index_stays_in_sync_through_every_mutation() {
|
fn keyword_index_stays_in_sync_through_every_mutation() {
|
||||||
let dir = TempDir::new().unwrap();
|
let dir = TempDir::new().unwrap();
|
||||||
|
|||||||
@@ -1,10 +1,12 @@
|
|||||||
//! OpenClaw Integration Layer.
|
//! A Markdown-oriented memory backend over [`crate::HDF5Memory`].
|
||||||
//!
|
//!
|
||||||
//! Bridge between OpenClaw agent gateway (Markdown + sqlite-vec) and the
|
//! Named for OpenClaw, whose workspace memory is Markdown, but **not an
|
||||||
//! clawhdf5 HDF5-backed memory backend. Provides:
|
//! OpenClaw plugin**: nothing here registers with OpenClaw, and the
|
||||||
|
//! integration it was written for never worked (see `docs/openclaw.md`).
|
||||||
|
//! Provides:
|
||||||
//!
|
//!
|
||||||
//! - [`MemoryBackend`] — the trait OpenClaw implements against.
|
//! - [`MemoryBackend`] — search / read back / write / ingest / export.
|
||||||
//! - [`ClawhdfBackend`] — concrete HDF5-backed implementation.
|
//! - [`ClawhdfBackend`] — the HDF5-backed implementation.
|
||||||
//! - [`MarkdownParser`] — splits Markdown into [`MarkdownSection`] records.
|
//! - [`MarkdownParser`] — splits Markdown into [`MarkdownSection`] records.
|
||||||
//! - [`MarkdownExporter`] — renders sections back to Markdown text.
|
//! - [`MarkdownExporter`] — renders sections back to Markdown text.
|
||||||
|
|
||||||
@@ -13,9 +15,8 @@ use std::path::{Path, PathBuf};
|
|||||||
use std::time::{SystemTime, UNIX_EPOCH};
|
use std::time::{SystemTime, UNIX_EPOCH};
|
||||||
|
|
||||||
use crate::{
|
use crate::{
|
||||||
AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry,
|
AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry, SearchOptions,
|
||||||
confidence::{ConfidenceConfig, ScoredResult, reject_low_confidence},
|
confidence::ConfidenceConfig, reranker::ReRankConfig,
|
||||||
reranker::{ReRankConfig, RerankInput, rerank},
|
|
||||||
};
|
};
|
||||||
|
|
||||||
// ─────────────────────────────────────────────────────────────────────────────
|
// ─────────────────────────────────────────────────────────────────────────────
|
||||||
@@ -62,7 +63,8 @@ pub struct BackendStats {
|
|||||||
// MemoryBackend trait
|
// MemoryBackend trait
|
||||||
// ─────────────────────────────────────────────────────────────────────────────
|
// ─────────────────────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
/// Interface that OpenClaw uses to interact with a memory backend.
|
/// A Markdown-oriented memory backend: search, read back by path, write,
|
||||||
|
/// ingest and export.
|
||||||
///
|
///
|
||||||
/// Implementors provide persistent storage, full-text + vector search,
|
/// Implementors provide persistent storage, full-text + vector search,
|
||||||
/// Markdown ingestion / export, and statistics.
|
/// Markdown ingestion / export, and statistics.
|
||||||
@@ -319,7 +321,7 @@ impl MarkdownExporter {
|
|||||||
///
|
///
|
||||||
/// # Path mapping
|
/// # Path mapping
|
||||||
///
|
///
|
||||||
/// OpenClaw addresses memories by file path (e.g. `"memory/user.md"`).
|
/// Memories are addressed by file path (e.g. `"memory/user.md"`).
|
||||||
/// Internally every [`MemoryEntry`] stores the originating path as its
|
/// Internally every [`MemoryEntry`] stores the originating path as its
|
||||||
/// `source_channel`. Section sub-paths are stored as
|
/// `source_channel`. Section sub-paths are stored as
|
||||||
/// `"<path>::<heading>"`.
|
/// `"<path>::<heading>"`.
|
||||||
@@ -422,7 +424,7 @@ impl ClawhdfBackend {
|
|||||||
|
|
||||||
// ── Compaction & Consolidation hooks (7.6) ────────────────────────────
|
// ── Compaction & Consolidation hooks (7.6) ────────────────────────────
|
||||||
|
|
||||||
/// Run a compaction cycle — called by OpenClaw during session compaction.
|
/// Run a compaction cycle (decay, compaction, WAL flush).
|
||||||
///
|
///
|
||||||
/// Sequence:
|
/// Sequence:
|
||||||
/// 1. `tick_session()` — apply Hebbian decay to all activation weights.
|
/// 1. `tick_session()` — apply Hebbian decay to all activation weights.
|
||||||
@@ -466,7 +468,7 @@ impl ClawhdfBackend {
|
|||||||
let record = MemoryRecord {
|
let record = MemoryRecord {
|
||||||
id: i as u64,
|
id: i as u64,
|
||||||
chunk: cache.chunks[i].clone(),
|
chunk: cache.chunks[i].clone(),
|
||||||
embedding: cache.embeddings[i].clone(),
|
embedding: cache.embeddings[i].to_vec(),
|
||||||
tier: MemoryTier::Working,
|
tier: MemoryTier::Working,
|
||||||
importance: cache.activation_weights[i],
|
importance: cache.activation_weights[i],
|
||||||
access_count: 0,
|
access_count: 0,
|
||||||
@@ -524,67 +526,27 @@ impl ClawhdfBackend {
|
|||||||
|
|
||||||
impl MemoryBackend for ClawhdfBackend {
|
impl MemoryBackend for ClawhdfBackend {
|
||||||
/// Search using hybrid vector + BM25 retrieval, then re-rank and
|
/// Search using hybrid vector + BM25 retrieval, then re-rank and
|
||||||
/// confidence-filter.
|
/// confidence-filter — [`HDF5Memory::search`] with both stages on.
|
||||||
fn search(
|
fn search(
|
||||||
&mut self,
|
&mut self,
|
||||||
query_text: &str,
|
query_text: &str,
|
||||||
query_embedding: &[f32],
|
query_embedding: &[f32],
|
||||||
k: usize,
|
k: usize,
|
||||||
) -> Vec<MemorySearchResult> {
|
) -> Vec<MemorySearchResult> {
|
||||||
// 1. Hybrid retrieval (RRF-blended vector + BM25).
|
let options = SearchOptions::new(k)
|
||||||
let candidates = k.saturating_mul(3).max(10);
|
.with_rerank(self.rerank_config)
|
||||||
let raw = self
|
.with_confidence(self.confidence_config.clone())
|
||||||
.memory
|
.at_time(Self::now_secs());
|
||||||
.hybrid_search(query_embedding, query_text, 0.7, 0.3, candidates);
|
self.memory
|
||||||
|
.search(query_embedding, query_text, &options)
|
||||||
if raw.is_empty() {
|
|
||||||
return Vec::new();
|
|
||||||
}
|
|
||||||
|
|
||||||
let now = Self::now_secs();
|
|
||||||
|
|
||||||
// 2. Re-rank using temporal recency, source authority, Hebbian weight.
|
|
||||||
let rerank_inputs: Vec<RerankInput> = raw
|
|
||||||
.iter()
|
|
||||||
.map(|r| RerankInput {
|
|
||||||
index: r.index,
|
|
||||||
timestamp: r.timestamp,
|
|
||||||
source_channel: r.source_channel.clone(),
|
|
||||||
raw_activation: r.activation,
|
|
||||||
})
|
|
||||||
.collect();
|
|
||||||
|
|
||||||
let reranked = rerank(&rerank_inputs, &self.rerank_config, now);
|
|
||||||
|
|
||||||
// 3. Confidence rejection.
|
|
||||||
let scored: Vec<ScoredResult> = reranked
|
|
||||||
.iter()
|
|
||||||
.map(|r| ScoredResult {
|
|
||||||
index: r.index,
|
|
||||||
score: r.combined_score,
|
|
||||||
})
|
|
||||||
.collect();
|
|
||||||
|
|
||||||
let confident = reject_low_confidence(&scored, &self.confidence_config);
|
|
||||||
|
|
||||||
// 4. Map back to MemorySearchResult; preserve raw text via index lookup.
|
|
||||||
let raw_by_idx: HashMap<usize, &crate::SearchResult> =
|
|
||||||
raw.iter().map(|r| (r.index, r)).collect();
|
|
||||||
|
|
||||||
confident
|
|
||||||
.into_iter()
|
.into_iter()
|
||||||
.take(k)
|
.map(|r| MemorySearchResult {
|
||||||
.filter_map(|sr| {
|
text: r.chunk,
|
||||||
let r = raw_by_idx.get(&sr.index)?;
|
score: r.score,
|
||||||
let path = r.source_channel.clone();
|
path: r.source_channel.clone(),
|
||||||
Some(MemorySearchResult {
|
line_range: None,
|
||||||
text: r.chunk.clone(),
|
timestamp: Some(r.timestamp),
|
||||||
score: sr.score,
|
source: r.source_channel,
|
||||||
path: path.clone(),
|
|
||||||
line_range: None,
|
|
||||||
timestamp: Some(r.timestamp),
|
|
||||||
source: path,
|
|
||||||
})
|
|
||||||
})
|
})
|
||||||
.collect()
|
.collect()
|
||||||
}
|
}
|
||||||
@@ -713,11 +675,13 @@ impl MemoryBackend for ClawhdfBackend {
|
|||||||
|
|
||||||
let total_records = cache.count_active();
|
let total_records = cache.count_active();
|
||||||
|
|
||||||
|
// A record saved without an embedding occupies a zero row, so "has an
|
||||||
|
// embedding" is "has a non-zero norm" rather than "row is non-empty".
|
||||||
let total_embeddings = cache
|
let total_embeddings = cache
|
||||||
.embeddings
|
.norms
|
||||||
.iter()
|
.iter()
|
||||||
.enumerate()
|
.enumerate()
|
||||||
.filter(|(i, emb)| cache.tombstones[*i] == 0 && !emb.is_empty())
|
.filter(|(i, norm)| cache.tombstones[*i] == 0 && **norm > 0.0)
|
||||||
.count();
|
.count();
|
||||||
|
|
||||||
let file_size_bytes = std::fs::metadata(&self.hdf5_path)
|
let file_size_bytes = std::fs::metadata(&self.hdf5_path)
|
||||||
|
|||||||
@@ -2,7 +2,7 @@
|
|||||||
//!
|
//!
|
||||||
//! Records the origin, authorship, and a content hash of every memory chunk
|
//! Records the origin, authorship, and a content hash of every memory chunk
|
||||||
//! so the system can detect *accidental* corruption and trace data lineage.
|
//! so the system can detect *accidental* corruption and trace data lineage.
|
||||||
//! The hash is unkeyed (see [`fnv1a_64`]) — this is not a tamper-evidence or
|
//! The hash is unkeyed (FNV-1a) — this is not a tamper-evidence or
|
||||||
//! authenticity guarantee.
|
//! authenticity guarantee.
|
||||||
|
|
||||||
use std::collections::HashMap;
|
use std::collections::HashMap;
|
||||||
|
|||||||
@@ -6,6 +6,11 @@
|
|||||||
//! - Temporal expansion (time-related rewrites)
|
//! - Temporal expansion (time-related rewrites)
|
||||||
//! - Morphological variants (stemming-like transforms)
|
//! - Morphological variants (stemming-like transforms)
|
||||||
//! - Knowledge graph expansion (entity aliases and neighbors)
|
//! - Knowledge graph expansion (entity aliases and neighbors)
|
||||||
|
//!
|
||||||
|
//! The morphological rules are crude suffix swaps, so some variants are not
|
||||||
|
//! words ("during" -> "dured"). That is tolerable for a BM25 stage, which
|
||||||
|
//! simply finds no postings for a nonsense term, but it means expansion is not
|
||||||
|
//! free: measure before enabling it on a retrieval path.
|
||||||
|
|
||||||
use crate::knowledge::KnowledgeCache;
|
use crate::knowledge::KnowledgeCache;
|
||||||
|
|
||||||
@@ -340,20 +345,87 @@ fn contains_phrase(text: &str, phrase: &str) -> bool {
|
|||||||
|
|
||||||
/// Replace a phrase in `text` case-insensitively, preserving surrounding case.
|
/// Replace a phrase in `text` case-insensitively, preserving surrounding case.
|
||||||
fn replace_word_case_insensitive(text: &str, from: &str, to: &str) -> String {
|
fn replace_word_case_insensitive(text: &str, from: &str, to: &str) -> String {
|
||||||
case_insensitive_replace(text, from, to)
|
replace_first(text, from, to, MatchKind::WholeWord)
|
||||||
}
|
}
|
||||||
|
|
||||||
fn case_insensitive_replace(text: &str, from: &str, to: &str) -> String {
|
fn case_insensitive_replace(text: &str, from: &str, to: &str) -> String {
|
||||||
let lower = text.to_lowercase();
|
replace_first(text, from, to, MatchKind::Substring)
|
||||||
let lower_from = from.to_lowercase();
|
}
|
||||||
if let Some(pos) = lower.find(&lower_from) {
|
|
||||||
let end = pos + from.len();
|
/// Whether a match may fall inside a larger word.
|
||||||
format!("{}{}{}", &text[..pos], to, &text[end..])
|
#[derive(Clone, Copy, PartialEq)]
|
||||||
} else {
|
enum MatchKind {
|
||||||
text.to_string()
|
/// Match anywhere, including inside another word.
|
||||||
|
Substring,
|
||||||
|
/// Match only when both ends sit on a word boundary.
|
||||||
|
WholeWord,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Replace the first case-insensitive match of `from` in `text` with `to`.
|
||||||
|
///
|
||||||
|
/// Matching walks the *original* string rather than a lowercased copy. The
|
||||||
|
/// previous implementation searched `text.to_lowercase()` and then sliced
|
||||||
|
/// `text` with the offsets it found, which only holds while lowercasing
|
||||||
|
/// preserves byte length. It does not: Turkish `İ` (2 bytes) lowercases to
|
||||||
|
/// `i` + U+0307 (3 bytes), so every later offset was wrong — silently
|
||||||
|
/// corrupting the output, or panicking when an offset landed inside a
|
||||||
|
/// character or past the end. `"İ AI"` was enough to panic.
|
||||||
|
fn replace_first(text: &str, from: &str, to: &str, kind: MatchKind) -> String {
|
||||||
|
match find_case_insensitive(text, from, kind) {
|
||||||
|
Some((start, end)) => {
|
||||||
|
let mut out = String::with_capacity(text.len() - (end - start) + to.len());
|
||||||
|
out.push_str(&text[..start]);
|
||||||
|
out.push_str(to);
|
||||||
|
out.push_str(&text[end..]);
|
||||||
|
out
|
||||||
|
}
|
||||||
|
None => text.to_string(),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Byte range of the first case-insensitive match of `needle` in `haystack`.
|
||||||
|
fn find_case_insensitive(haystack: &str, needle: &str, kind: MatchKind) -> Option<(usize, usize)> {
|
||||||
|
if needle.is_empty() {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
let lowered: Vec<char> = needle.chars().flat_map(char::to_lowercase).collect();
|
||||||
|
let is_word = |c: char| c.is_alphanumeric() || c == '_';
|
||||||
|
|
||||||
|
for (start, _) in haystack.char_indices() {
|
||||||
|
if kind == MatchKind::WholeWord
|
||||||
|
&& haystack[..start].chars().next_back().is_some_and(is_word)
|
||||||
|
{
|
||||||
|
continue; // mid-word: "ai" inside "training"
|
||||||
|
}
|
||||||
|
let mut matched = 0usize;
|
||||||
|
let mut end = start;
|
||||||
|
for (offset, ch) in haystack[start..].char_indices() {
|
||||||
|
if matched == lowered.len() {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
let mut consumed_all = true;
|
||||||
|
for lc in ch.to_lowercase() {
|
||||||
|
if lowered.get(matched) != Some(&lc) {
|
||||||
|
consumed_all = false;
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
matched += 1;
|
||||||
|
}
|
||||||
|
if !consumed_all {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
end = start + offset + ch.len_utf8();
|
||||||
|
}
|
||||||
|
if matched == lowered.len()
|
||||||
|
&& !(kind == MatchKind::WholeWord
|
||||||
|
&& haystack[end..].chars().next().is_some_and(is_word))
|
||||||
|
{
|
||||||
|
return Some((start, end));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
None
|
||||||
|
}
|
||||||
|
|
||||||
/// Simple whitespace/punctuation tokenizer.
|
/// Simple whitespace/punctuation tokenizer.
|
||||||
fn tokenize(text: &str) -> Vec<String> {
|
fn tokenize(text: &str) -> Vec<String> {
|
||||||
text.split(|c: char| !c.is_alphanumeric())
|
text.split(|c: char| !c.is_alphanumeric())
|
||||||
@@ -637,4 +709,86 @@ mod tests {
|
|||||||
expanded.iter().map(|x| &x.text).collect::<Vec<_>>()
|
expanded.iter().map(|x| &x.text).collect::<Vec<_>>()
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
#[test]
|
||||||
|
fn acronyms_only_match_whole_words() {
|
||||||
|
let ex = QueryExpander::new(QueryExpansionConfig::default());
|
||||||
|
// "training" contains "ai", "programming" contains "pr". These used to
|
||||||
|
// be rewritten to "trArtificial Intelligencening" and
|
||||||
|
// "Pull Requestogramming".
|
||||||
|
for query in [
|
||||||
|
"How many miles during my marathon training?",
|
||||||
|
"Which programming language did I pick?",
|
||||||
|
"I updated the maintainer list",
|
||||||
|
] {
|
||||||
|
for expansion in ex.expand(query) {
|
||||||
|
assert!(
|
||||||
|
expansion.expansion_type != "acronym",
|
||||||
|
"{query:?} produced {expansion:?}"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
// A real acronym still expands, in both directions.
|
||||||
|
let texts: Vec<String> = ex
|
||||||
|
.expand("What about the API and the database?")
|
||||||
|
.into_iter()
|
||||||
|
.filter(|e| e.expansion_type == "acronym")
|
||||||
|
.map(|e| e.text)
|
||||||
|
.collect();
|
||||||
|
assert!(
|
||||||
|
texts
|
||||||
|
.iter()
|
||||||
|
.any(|t| t.contains("Application Programming Interface")),
|
||||||
|
"{texts:?}"
|
||||||
|
);
|
||||||
|
assert!(texts.iter().any(|t| t.contains("DB")), "{texts:?}");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn non_ascii_queries_do_not_panic_or_corrupt() {
|
||||||
|
let ex = QueryExpander::new(QueryExpansionConfig::default());
|
||||||
|
// Turkish 'İ' is 2 bytes but lowercases to 3, so offsets taken from a
|
||||||
|
// lowercased copy no longer line up with the original. `"İ AI"` used
|
||||||
|
// to panic; `"İstanbul AI trip"` used to silently eat a character.
|
||||||
|
for query in ["İ AI", "İé AI", "İİ ML", "İstanbul AI trip", "ǰ ML notes"] {
|
||||||
|
for expansion in ex.expand(query) {
|
||||||
|
assert!(
|
||||||
|
expansion.text.contains('İ') || expansion.text.contains('ǰ'),
|
||||||
|
"{query:?} lost its leading character: {expansion:?}"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
let expanded = ex.expand("İstanbul AI trip");
|
||||||
|
assert!(
|
||||||
|
expanded
|
||||||
|
.iter()
|
||||||
|
.any(|e| e.text == "İstanbul Artificial Intelligence trip"),
|
||||||
|
"{expanded:?}"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn whole_word_matching_handles_string_edges_and_case() {
|
||||||
|
assert_eq!(
|
||||||
|
replace_word_case_insensitive("ai tools", "AI", "Artificial Intelligence"),
|
||||||
|
"Artificial Intelligence tools"
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
replace_word_case_insensitive("tools for ai", "AI", "Artificial Intelligence"),
|
||||||
|
"tools for Artificial Intelligence"
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
replace_word_case_insensitive("the aim", "AI", "Artificial Intelligence"),
|
||||||
|
"the aim",
|
||||||
|
"must not match inside a word"
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
replace_word_case_insensitive("no match here", "xyz", "abc"),
|
||||||
|
"no match here"
|
||||||
|
);
|
||||||
|
// Only the first occurrence is replaced, as before.
|
||||||
|
assert_eq!(
|
||||||
|
replace_word_case_insensitive("ai and ai", "ai", "ML"),
|
||||||
|
"ML and ai"
|
||||||
|
);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -4,8 +4,10 @@
|
|||||||
//! into a single composite score for each retrieved result.
|
//! into a single composite score for each retrieved result.
|
||||||
|
|
||||||
/// Configuration for the multi-factor re-ranker.
|
/// Configuration for the multi-factor re-ranker.
|
||||||
#[derive(Debug, Clone)]
|
#[derive(Debug, Clone, Copy)]
|
||||||
pub struct ReRankConfig {
|
pub struct ReRankConfig {
|
||||||
|
/// Weight applied to the retrieval score the candidate arrived with.
|
||||||
|
pub relevance_weight: f32,
|
||||||
/// Weight applied to the temporal decay score (0.0–1.0).
|
/// Weight applied to the temporal decay score (0.0–1.0).
|
||||||
pub temporal_weight: f32,
|
pub temporal_weight: f32,
|
||||||
/// Weight applied to the source authority score (0.0–1.0).
|
/// Weight applied to the source authority score (0.0–1.0).
|
||||||
@@ -20,6 +22,9 @@ pub struct ReRankConfig {
|
|||||||
impl Default for ReRankConfig {
|
impl Default for ReRankConfig {
|
||||||
fn default() -> Self {
|
fn default() -> Self {
|
||||||
Self {
|
Self {
|
||||||
|
// Relevance leads: the metadata signals break ties and nudge, they
|
||||||
|
// do not decide. See `BENCHMARKS.md`, "Recency discrimination".
|
||||||
|
relevance_weight: 1.0,
|
||||||
temporal_weight: 0.3,
|
temporal_weight: 0.3,
|
||||||
authority_weight: 0.2,
|
authority_weight: 0.2,
|
||||||
activation_weight: 0.5,
|
activation_weight: 0.5,
|
||||||
@@ -41,6 +46,8 @@ pub struct ReRankResult {
|
|||||||
pub authority_score: f32,
|
pub authority_score: f32,
|
||||||
/// Normalised Hebbian activation score in [0, 1].
|
/// Normalised Hebbian activation score in [0, 1].
|
||||||
pub activation_score: f32,
|
pub activation_score: f32,
|
||||||
|
/// The retrieval score carried through from the input.
|
||||||
|
pub relevance_score: f32,
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Compute an exponential decay temporal score.
|
/// Compute an exponential decay temporal score.
|
||||||
@@ -105,6 +112,15 @@ pub struct RerankInput {
|
|||||||
pub source_channel: String,
|
pub source_channel: String,
|
||||||
/// Raw Hebbian activation weight for this entry.
|
/// Raw Hebbian activation weight for this entry.
|
||||||
pub raw_activation: f32,
|
pub raw_activation: f32,
|
||||||
|
/// The retrieval score that put this entry in the candidate list.
|
||||||
|
///
|
||||||
|
/// Re-ranking is meant to *adjust* the retriever's ordering with signals
|
||||||
|
/// it does not have, not to replace it. Without this the combined score
|
||||||
|
/// was made of recency, authority and activation alone, so a candidate
|
||||||
|
/// pool came back ordered by age with its relevance ordering discarded.
|
||||||
|
/// Callers with no meaningful score can pass the same value for every
|
||||||
|
/// entry, which reduces to the old behaviour.
|
||||||
|
pub relevance: f32,
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Re-rank a list of retrieval results using multi-factor scoring.
|
/// Re-rank a list of retrieval results using multi-factor scoring.
|
||||||
@@ -138,7 +154,8 @@ pub fn rerank(
|
|||||||
let auth = source_authority_score(&inp.source_channel);
|
let auth = source_authority_score(&inp.source_channel);
|
||||||
let act = activation_score(inp.raw_activation);
|
let act = activation_score(inp.raw_activation);
|
||||||
|
|
||||||
let combined = config.temporal_weight * ts
|
let combined = config.relevance_weight * inp.relevance
|
||||||
|
+ config.temporal_weight * ts
|
||||||
+ config.authority_weight * auth
|
+ config.authority_weight * auth
|
||||||
+ config.activation_weight * act;
|
+ config.activation_weight * act;
|
||||||
|
|
||||||
@@ -148,6 +165,7 @@ pub fn rerank(
|
|||||||
temporal_score: ts,
|
temporal_score: ts,
|
||||||
authority_score: auth,
|
authority_score: auth,
|
||||||
activation_score: act,
|
activation_score: act,
|
||||||
|
relevance_score: inp.relevance,
|
||||||
}
|
}
|
||||||
})
|
})
|
||||||
.collect();
|
.collect();
|
||||||
@@ -253,22 +271,51 @@ mod tests {
|
|||||||
timestamp: 0.0, // very old
|
timestamp: 0.0, // very old
|
||||||
source_channel: "other".to_string(),
|
source_channel: "other".to_string(),
|
||||||
raw_activation: 0.1,
|
raw_activation: 0.1,
|
||||||
|
relevance: 0.0,
|
||||||
},
|
},
|
||||||
RerankInput {
|
RerankInput {
|
||||||
index: 1,
|
index: 1,
|
||||||
timestamp: 86_400.0, // one day ago
|
timestamp: 86_400.0, // one day ago
|
||||||
source_channel: "conversation".to_string(),
|
source_channel: "conversation".to_string(),
|
||||||
raw_activation: 0.5,
|
raw_activation: 0.5,
|
||||||
|
relevance: 0.0,
|
||||||
},
|
},
|
||||||
RerankInput {
|
RerankInput {
|
||||||
index: 2,
|
index: 2,
|
||||||
timestamp: 172_800.0, // "now"
|
timestamp: 172_800.0, // "now"
|
||||||
source_channel: "user_correction".to_string(),
|
source_channel: "user_correction".to_string(),
|
||||||
raw_activation: 1.0,
|
raw_activation: 1.0,
|
||||||
|
relevance: 0.0,
|
||||||
},
|
},
|
||||||
]
|
]
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn relevance_leads_but_recency_breaks_near_ties() {
|
||||||
|
let entry = |index, timestamp, relevance| RerankInput {
|
||||||
|
index,
|
||||||
|
timestamp,
|
||||||
|
source_channel: "conversation".to_string(),
|
||||||
|
raw_activation: 1.0,
|
||||||
|
relevance,
|
||||||
|
};
|
||||||
|
let now = 10.0 * 86_400.0;
|
||||||
|
let config = ReRankConfig::default();
|
||||||
|
|
||||||
|
// A clearly better match wins despite being much older. Before
|
||||||
|
// `relevance` existed the combined score ignored it entirely, so this
|
||||||
|
// returned the newer, irrelevant entry.
|
||||||
|
let ranked = rerank(&[entry(0, 0.0, 1.0), entry(1, now, 0.1)], &config, now);
|
||||||
|
assert_eq!(ranked[0].index, 0, "{ranked:?}");
|
||||||
|
|
||||||
|
// Between near-equal matches, the newer one wins.
|
||||||
|
let ranked = rerank(&[entry(0, 0.0, 0.80), entry(1, now, 0.79)], &config, now);
|
||||||
|
assert_eq!(ranked[0].index, 1, "{ranked:?}");
|
||||||
|
|
||||||
|
// The breakdown carries the relevance through.
|
||||||
|
assert_eq!(ranked[0].relevance_score, 0.79);
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn rerank_returns_all_entries() {
|
fn rerank_returns_all_entries() {
|
||||||
let inputs = make_inputs();
|
let inputs = make_inputs();
|
||||||
@@ -302,6 +349,7 @@ mod tests {
|
|||||||
#[test]
|
#[test]
|
||||||
fn rerank_score_breakdown_matches_manual_calculation() {
|
fn rerank_score_breakdown_matches_manual_calculation() {
|
||||||
let config = ReRankConfig {
|
let config = ReRankConfig {
|
||||||
|
relevance_weight: 0.0,
|
||||||
temporal_weight: 1.0,
|
temporal_weight: 1.0,
|
||||||
authority_weight: 0.0,
|
authority_weight: 0.0,
|
||||||
activation_weight: 0.0,
|
activation_weight: 0.0,
|
||||||
@@ -312,6 +360,7 @@ mod tests {
|
|||||||
timestamp: 0.0,
|
timestamp: 0.0,
|
||||||
source_channel: "other".to_string(),
|
source_channel: "other".to_string(),
|
||||||
raw_activation: 0.5,
|
raw_activation: 0.5,
|
||||||
|
relevance: 0.0,
|
||||||
}];
|
}];
|
||||||
let now = 3600.0_f64; // exactly one half-life later
|
let now = 3600.0_f64; // exactly one half-life later
|
||||||
let results = rerank(&inputs, &config, now);
|
let results = rerank(&inputs, &config, now);
|
||||||
|
|||||||
@@ -15,6 +15,9 @@ use crate::session::SessionCache;
|
|||||||
use crate::wal::WalMark;
|
use crate::wal::WalMark;
|
||||||
|
|
||||||
pub const SCHEMA_VERSION: &str = "1.0";
|
pub const SCHEMA_VERSION: &str = "1.0";
|
||||||
|
/// Writer-version tag stored in `/meta` as `edgehdf5_version`. Kept for file
|
||||||
|
/// compatibility; despite the name it has nothing to do with ZeroClaw, which
|
||||||
|
/// does not use clawhdf5.
|
||||||
pub const ZEROCLAW_VERSION: &str = "0.8.0";
|
pub const ZEROCLAW_VERSION: &str = "0.8.0";
|
||||||
|
|
||||||
/// `/meta` attributes holding the [`WalMark`] of the WAL prefix already folded
|
/// `/meta` attributes holding the [`WalMark`] of the WAL prefix already folded
|
||||||
@@ -23,6 +26,7 @@ pub const ZEROCLAW_VERSION: &str = "0.8.0";
|
|||||||
const WAL_APPLIED_LEN_ATTR: &str = "wal_applied_len";
|
const WAL_APPLIED_LEN_ATTR: &str = "wal_applied_len";
|
||||||
const WAL_APPLIED_CRC_ATTR: &str = "wal_applied_crc";
|
const WAL_APPLIED_CRC_ATTR: &str = "wal_applied_crc";
|
||||||
const ANN_GENERATION_ATTR: &str = "ann_generation";
|
const ANN_GENERATION_ATTR: &str = "ann_generation";
|
||||||
|
const SIG_VERSION_ATTR: &str = "sig_version";
|
||||||
|
|
||||||
/// Build a complete HDF5 file from the in-memory state.
|
/// Build a complete HDF5 file from the in-memory state.
|
||||||
pub fn build_hdf5_file(
|
pub fn build_hdf5_file(
|
||||||
@@ -46,7 +50,7 @@ pub fn build_hdf5_file_with_mark(
|
|||||||
) -> Result<Vec<u8>, MemoryError> {
|
) -> Result<Vec<u8>, MemoryError> {
|
||||||
let meta = CheckpointMeta {
|
let meta = CheckpointMeta {
|
||||||
wal_applied,
|
wal_applied,
|
||||||
ann_generation: None,
|
..CheckpointMeta::default()
|
||||||
};
|
};
|
||||||
build_hdf5_file_with_meta(config, cache, sessions, knowledge, &meta)
|
build_hdf5_file_with_meta(config, cache, sessions, knowledge, &meta)
|
||||||
}
|
}
|
||||||
@@ -61,6 +65,10 @@ pub struct CheckpointMeta {
|
|||||||
/// one left over from another checkpoint can never be attached to records
|
/// one left over from another checkpoint can never be attached to records
|
||||||
/// it wasn't built from.
|
/// it wasn't built from.
|
||||||
pub ann_generation: Option<u64>,
|
pub ann_generation: Option<u64>,
|
||||||
|
/// The checkpoint carries an Ed25519 signature (see [`crate::signing`]).
|
||||||
|
/// Read-only: whether a checkpoint is *written* signed is decided by the
|
||||||
|
/// signature passed to [`build_hdf5_file_signed`].
|
||||||
|
pub signed: bool,
|
||||||
}
|
}
|
||||||
|
|
||||||
/// [`build_hdf5_file`] with checkpoint bookkeeping.
|
/// [`build_hdf5_file`] with checkpoint bookkeeping.
|
||||||
@@ -70,6 +78,19 @@ pub fn build_hdf5_file_with_meta(
|
|||||||
sessions: &SessionCache,
|
sessions: &SessionCache,
|
||||||
knowledge: &KnowledgeCache,
|
knowledge: &KnowledgeCache,
|
||||||
checkpoint: &CheckpointMeta,
|
checkpoint: &CheckpointMeta,
|
||||||
|
) -> Result<Vec<u8>, MemoryError> {
|
||||||
|
build_hdf5_file_signed(config, cache, sessions, knowledge, checkpoint, None)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// [`build_hdf5_file_with_meta`], plus a signed manifest of the contents
|
||||||
|
/// (see [`crate::signing`]).
|
||||||
|
pub fn build_hdf5_file_signed(
|
||||||
|
config: &MemoryConfig,
|
||||||
|
cache: &MemoryCache,
|
||||||
|
sessions: &SessionCache,
|
||||||
|
knowledge: &KnowledgeCache,
|
||||||
|
checkpoint: &CheckpointMeta,
|
||||||
|
signature: Option<&crate::signing::StoredSignature>,
|
||||||
) -> Result<Vec<u8>, MemoryError> {
|
) -> Result<Vec<u8>, MemoryError> {
|
||||||
let wal_applied = checkpoint.wal_applied;
|
let wal_applied = checkpoint.wal_applied;
|
||||||
let mut builder = clawhdf5::FileBuilder::new();
|
let mut builder = clawhdf5::FileBuilder::new();
|
||||||
@@ -104,6 +125,19 @@ pub fn build_hdf5_file_with_meta(
|
|||||||
"wal_max_entries",
|
"wal_max_entries",
|
||||||
AttrValue::I64(config.wal_max_entries as i64),
|
AttrValue::I64(config.wal_max_entries as i64),
|
||||||
);
|
);
|
||||||
|
meta.set_attr(
|
||||||
|
"quantized_index",
|
||||||
|
AttrValue::I64(config.quantized_index.into()),
|
||||||
|
);
|
||||||
|
meta.set_attr("hnsw_m", AttrValue::I64(config.hnsw_m as i64));
|
||||||
|
meta.set_attr(
|
||||||
|
"hnsw_ef_construction",
|
||||||
|
AttrValue::I64(config.hnsw_ef_construction as i64),
|
||||||
|
);
|
||||||
|
meta.set_attr(
|
||||||
|
"hnsw_ef_search",
|
||||||
|
AttrValue::I64(config.hnsw_ef_search as i64),
|
||||||
|
);
|
||||||
meta.set_attr(
|
meta.set_attr(
|
||||||
"edgehdf5_version",
|
"edgehdf5_version",
|
||||||
AttrValue::String(ZEROCLAW_VERSION.into()),
|
AttrValue::String(ZEROCLAW_VERSION.into()),
|
||||||
@@ -117,11 +151,42 @@ pub fn build_hdf5_file_with_meta(
|
|||||||
// round trip through every reader.
|
// round trip through every reader.
|
||||||
meta.set_attr(ANN_GENERATION_ATTR, AttrValue::I64(generation as i64));
|
meta.set_attr(ANN_GENERATION_ATTR, AttrValue::I64(generation as i64));
|
||||||
}
|
}
|
||||||
|
if let Some(sig) = signature {
|
||||||
|
use crate::signing::to_hex;
|
||||||
|
let m = &sig.manifest;
|
||||||
|
meta.set_attr(
|
||||||
|
SIG_VERSION_ATTR,
|
||||||
|
AttrValue::I64(crate::signing::MANIFEST_VERSION),
|
||||||
|
);
|
||||||
|
meta.set_attr("sig_algorithm", AttrValue::String("ed25519".into()));
|
||||||
|
meta.set_attr("sig_public_key", AttrValue::String(to_hex(&sig.public_key)));
|
||||||
|
meta.set_attr("sig_signature", AttrValue::String(to_hex(&sig.signature)));
|
||||||
|
meta.set_attr("sig_record_count", AttrValue::I64(m.record_count as i64));
|
||||||
|
meta.set_attr(
|
||||||
|
"sig_records_root",
|
||||||
|
AttrValue::String(to_hex(&m.records_root)),
|
||||||
|
);
|
||||||
|
meta.set_attr("sig_settings", AttrValue::String(to_hex(&m.settings)));
|
||||||
|
meta.set_attr("sig_sessions", AttrValue::String(to_hex(&m.sessions)));
|
||||||
|
meta.set_attr("sig_graph", AttrValue::String(to_hex(&m.graph)));
|
||||||
|
}
|
||||||
// Need at least one dataset in the group for it to be a proper group
|
// Need at least one dataset in the group for it to be a proper group
|
||||||
meta.create_dataset("_marker").with_u8_data(&[1]).compact();
|
meta.create_dataset("_marker").with_u8_data(&[1]).compact();
|
||||||
let finished_meta = meta.finish();
|
let finished_meta = meta.finish();
|
||||||
builder.add_group(finished_meta);
|
builder.add_group(finished_meta);
|
||||||
|
|
||||||
|
// /integrity: the signed per-record hashes, so verification can say
|
||||||
|
// which records changed.
|
||||||
|
if let Some(sig) = signature {
|
||||||
|
let mut group = builder.create_group("integrity");
|
||||||
|
let flat: Vec<u8> = sig.record_hashes.iter().flatten().copied().collect();
|
||||||
|
group
|
||||||
|
.create_dataset("record_hashes")
|
||||||
|
.with_u8_data(&flat)
|
||||||
|
.with_shape(&[sig.record_hashes.len() as u64, 32]);
|
||||||
|
builder.add_group(group.finish());
|
||||||
|
}
|
||||||
|
|
||||||
// /memory group
|
// /memory group
|
||||||
build_memory_group(&mut builder, config, cache)?;
|
build_memory_group(&mut builder, config, cache)?;
|
||||||
|
|
||||||
@@ -146,20 +211,27 @@ fn build_memory_group(
|
|||||||
// chunks: fixed-length string array
|
// chunks: fixed-length string array
|
||||||
write_string_dataset(&mut group, "chunks", &cache.chunks);
|
write_string_dataset(&mut group, "chunks", &cache.chunks);
|
||||||
|
|
||||||
// embeddings: f32 [N x D]
|
// embeddings: [N x D], f32 — or IEEE half precision for a `float16`
|
||||||
|
// store. The cache already holds half-rounded values then, so this
|
||||||
|
// conversion is exact and a reopened store sees the same numbers.
|
||||||
let n = cache.embeddings.len() as u64;
|
let n = cache.embeddings.len() as u64;
|
||||||
let d = cache.embedding_dim as u64;
|
let d = cache.embedding_dim as u64;
|
||||||
let flat = cache.flat_embeddings();
|
let flat = cache.flat_embeddings();
|
||||||
{
|
{
|
||||||
let ds = group
|
let ds = group.create_dataset("embeddings");
|
||||||
.create_dataset("embeddings")
|
let elem_bytes: u64 = if config.float16 {
|
||||||
.with_f32_data(&flat)
|
ds.with_f16_data(flat);
|
||||||
.with_shape(&[n, d]);
|
2
|
||||||
|
} else {
|
||||||
|
ds.with_f32_data(flat);
|
||||||
|
4
|
||||||
|
};
|
||||||
|
ds.with_shape(&[n, d]);
|
||||||
|
|
||||||
// Chunk size tuning: target ~256KB per chunk for optimal I/O
|
// Chunk size tuning: target ~256KB per chunk for optimal I/O
|
||||||
if n > 0 && d > 0 {
|
if n > 0 && d > 0 {
|
||||||
let target_chunk_bytes: u64 = 256 * 1024;
|
let target_chunk_bytes: u64 = 256 * 1024;
|
||||||
let rows_per_chunk = (target_chunk_bytes / (d * 4)).max(1).min(n);
|
let rows_per_chunk = (target_chunk_bytes / (d * elem_bytes)).max(1).min(n);
|
||||||
ds.with_chunks(&[rows_per_chunk, d]);
|
ds.with_chunks(&[rows_per_chunk, d]);
|
||||||
|
|
||||||
// Compression. Shuffle is applied automatically (auto-shuffle
|
// Compression. Shuffle is applied automatically (auto-shuffle
|
||||||
@@ -420,6 +492,64 @@ pub fn read_wal_mark(file: &clawhdf5::File) -> Option<WalMark> {
|
|||||||
Some(WalMark { len, crc })
|
Some(WalMark { len, crc })
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Read a checkpoint's signature, if it has one. A signature whose
|
||||||
|
/// attributes are present but malformed is an error, not "unsigned".
|
||||||
|
pub fn read_signature(
|
||||||
|
file: &clawhdf5::File,
|
||||||
|
) -> Result<Option<crate::signing::StoredSignature>, MemoryError> {
|
||||||
|
use crate::signing::{Manifest, StoredSignature, from_hex};
|
||||||
|
let attrs = file
|
||||||
|
.group("meta")
|
||||||
|
.and_then(|g| g.attrs())
|
||||||
|
.map_err(|e| MemoryError::Schema(format!("cannot read /meta attrs: {e}")))?;
|
||||||
|
let version = match attrs.get(SIG_VERSION_ATTR) {
|
||||||
|
None => return Ok(None),
|
||||||
|
Some(AttrValue::I64(v)) => *v,
|
||||||
|
Some(_) => return Err(MemoryError::Schema("malformed sig_version".into())),
|
||||||
|
};
|
||||||
|
if version != crate::signing::MANIFEST_VERSION {
|
||||||
|
return Err(MemoryError::Schema(format!(
|
||||||
|
"unsupported signature version {version}"
|
||||||
|
)));
|
||||||
|
}
|
||||||
|
fn hex<const N: usize>(
|
||||||
|
attrs: &std::collections::HashMap<String, AttrValue>,
|
||||||
|
name: &str,
|
||||||
|
) -> Result<[u8; N], MemoryError> {
|
||||||
|
match attrs.get(name) {
|
||||||
|
Some(AttrValue::String(s)) => from_hex::<N>(s),
|
||||||
|
_ => None,
|
||||||
|
}
|
||||||
|
.ok_or_else(|| MemoryError::Schema(format!("malformed or missing {name}")))
|
||||||
|
}
|
||||||
|
let record_count = match attrs.get("sig_record_count") {
|
||||||
|
Some(AttrValue::I64(v)) if *v >= 0 => *v as u64,
|
||||||
|
_ => return Err(MemoryError::Schema("malformed sig_record_count".into())),
|
||||||
|
};
|
||||||
|
let group = file
|
||||||
|
.group("integrity")
|
||||||
|
.map_err(|e| MemoryError::Schema(format!("signed checkpoint without /integrity: {e}")))?;
|
||||||
|
let flat = read_u8_dataset(&group, "record_hashes")?;
|
||||||
|
if flat.len() % 32 != 0 {
|
||||||
|
return Err(MemoryError::Schema(
|
||||||
|
"/integrity/record_hashes is not a whole number of hashes".into(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
let record_hashes = flat.as_chunks::<32>().0.to_vec();
|
||||||
|
Ok(Some(StoredSignature {
|
||||||
|
manifest: Manifest {
|
||||||
|
record_count,
|
||||||
|
records_root: hex::<32>(&attrs, "sig_records_root")?,
|
||||||
|
settings: hex::<32>(&attrs, "sig_settings")?,
|
||||||
|
sessions: hex::<32>(&attrs, "sig_sessions")?,
|
||||||
|
graph: hex::<32>(&attrs, "sig_graph")?,
|
||||||
|
},
|
||||||
|
record_hashes,
|
||||||
|
public_key: hex::<32>(&attrs, "sig_public_key")?,
|
||||||
|
signature: hex::<64>(&attrs, "sig_signature")?,
|
||||||
|
}))
|
||||||
|
}
|
||||||
|
|
||||||
/// Read the checkpoint bookkeeping from `/meta`.
|
/// Read the checkpoint bookkeeping from `/meta`.
|
||||||
pub fn read_checkpoint_meta(file: &clawhdf5::File) -> CheckpointMeta {
|
pub fn read_checkpoint_meta(file: &clawhdf5::File) -> CheckpointMeta {
|
||||||
let ann_generation = file
|
let ann_generation = file
|
||||||
@@ -430,9 +560,14 @@ pub fn read_checkpoint_meta(file: &clawhdf5::File) -> CheckpointMeta {
|
|||||||
Some(AttrValue::I64(v)) => Some(*v as u64),
|
Some(AttrValue::I64(v)) => Some(*v as u64),
|
||||||
_ => None,
|
_ => None,
|
||||||
});
|
});
|
||||||
|
let signed = file
|
||||||
|
.group("meta")
|
||||||
|
.and_then(|g| g.attrs())
|
||||||
|
.is_ok_and(|attrs| attrs.contains_key(SIG_VERSION_ATTR));
|
||||||
CheckpointMeta {
|
CheckpointMeta {
|
||||||
wal_applied: read_wal_mark(file),
|
wal_applied: read_wal_mark(file),
|
||||||
ann_generation,
|
ann_generation,
|
||||||
|
signed,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -484,10 +619,32 @@ pub fn validate_and_load(
|
|||||||
wal_max_entries: optional_i64_attr(&attrs, "wal_max_entries")
|
wal_max_entries: optional_i64_attr(&attrs, "wal_max_entries")
|
||||||
.and_then(|v| usize::try_from(v).ok())
|
.and_then(|v| usize::try_from(v).ok())
|
||||||
.unwrap_or(500),
|
.unwrap_or(500),
|
||||||
|
// `false`, not the new-store default: a store written before this
|
||||||
|
// setting existed was built with an f32 index, and reopening it must
|
||||||
|
// not silently change that.
|
||||||
|
quantized_index: optional_bool_attr(&attrs, "quantized_index", false),
|
||||||
|
hnsw_m: optional_i64_attr(&attrs, "hnsw_m")
|
||||||
|
.and_then(|v| usize::try_from(v).ok())
|
||||||
|
.unwrap_or(16),
|
||||||
|
hnsw_ef_construction: optional_i64_attr(&attrs, "hnsw_ef_construction")
|
||||||
|
.and_then(|v| usize::try_from(v).ok())
|
||||||
|
.unwrap_or(64),
|
||||||
|
hnsw_ef_search: optional_i64_attr(&attrs, "hnsw_ef_search")
|
||||||
|
.and_then(|v| usize::try_from(v).ok())
|
||||||
|
.unwrap_or(0),
|
||||||
};
|
};
|
||||||
|
|
||||||
// Load /memory group
|
// Load /memory group
|
||||||
let memory_cache = load_memory_group(file, embedding_dim)?;
|
let mut memory_cache = load_memory_group(file, embedding_dim)?;
|
||||||
|
// A float16 store's cache holds half-rounded embeddings. Embeddings read
|
||||||
|
// from an f16 dataset already are; a float16 store whose last checkpoint
|
||||||
|
// predates half-precision storage is still f32 on disk and is rounded
|
||||||
|
// here.
|
||||||
|
if config.float16 && embeddings_are_f16(file) {
|
||||||
|
memory_cache.half_precision = true;
|
||||||
|
} else {
|
||||||
|
memory_cache.set_half_precision(config.float16);
|
||||||
|
}
|
||||||
|
|
||||||
// Load /sessions group
|
// Load /sessions group
|
||||||
let session_cache = load_sessions_group(file)?;
|
let session_cache = load_sessions_group(file)?;
|
||||||
@@ -563,12 +720,7 @@ fn load_memory_group(
|
|||||||
.collect(),
|
.collect(),
|
||||||
};
|
};
|
||||||
|
|
||||||
// Unflatten embeddings
|
// No unflattening: the cache stores the buffer as it is on disk.
|
||||||
let embeddings: Vec<Vec<f32>> = flat_embeddings
|
|
||||||
.chunks(embedding_dim)
|
|
||||||
.map(|c| c.to_vec())
|
|
||||||
.collect();
|
|
||||||
|
|
||||||
// Read activation_weights if present, default to vec![1.0; N] for backward compat
|
// Read activation_weights if present, default to vec![1.0; N] for backward compat
|
||||||
let activation_weights = match read_f32_dataset(&group, "activation_weights") {
|
let activation_weights = match read_f32_dataset(&group, "activation_weights") {
|
||||||
Ok(w) if w.len() == n => w,
|
Ok(w) if w.len() == n => w,
|
||||||
@@ -576,7 +728,7 @@ fn load_memory_group(
|
|||||||
};
|
};
|
||||||
|
|
||||||
cache.chunks = chunks;
|
cache.chunks = chunks;
|
||||||
cache.embeddings = embeddings;
|
cache.embeddings.set_flat(embedding_dim, flat_embeddings);
|
||||||
cache.source_channels = source_channels;
|
cache.source_channels = source_channels;
|
||||||
cache.timestamps = timestamps;
|
cache.timestamps = timestamps;
|
||||||
cache.session_ids = session_ids;
|
cache.session_ids = session_ids;
|
||||||
@@ -584,7 +736,6 @@ fn load_memory_group(
|
|||||||
cache.tombstones = tombstones;
|
cache.tombstones = tombstones;
|
||||||
cache.norms = norms;
|
cache.norms = norms;
|
||||||
cache.activation_weights = activation_weights;
|
cache.activation_weights = activation_weights;
|
||||||
cache.rebuild_flat();
|
|
||||||
|
|
||||||
Ok(cache)
|
Ok(cache)
|
||||||
}
|
}
|
||||||
@@ -736,6 +887,13 @@ fn read_string_dataset_from_group(
|
|||||||
.map_err(|e| MemoryError::Hdf5(format!("cannot read strings from {name}: {e}")))
|
.map_err(|e| MemoryError::Hdf5(format!("cannot read strings from {name}: {e}")))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Whether `/memory/embeddings` is stored as IEEE half precision.
|
||||||
|
fn embeddings_are_f16(file: &clawhdf5::File) -> bool {
|
||||||
|
file.dataset("memory/embeddings")
|
||||||
|
.and_then(|ds| ds.dtype())
|
||||||
|
.is_ok_and(|dt| matches!(dt, clawhdf5::DType::Other(ref s) if s == "float16"))
|
||||||
|
}
|
||||||
|
|
||||||
fn read_f32_dataset(group: &clawhdf5::Group<'_>, name: &str) -> Result<Vec<f32>, MemoryError> {
|
fn read_f32_dataset(group: &clawhdf5::Group<'_>, name: &str) -> Result<Vec<f32>, MemoryError> {
|
||||||
let ds = group
|
let ds = group
|
||||||
.dataset(name)
|
.dataset(name)
|
||||||
|
|||||||
@@ -2,61 +2,196 @@
|
|||||||
|
|
||||||
use std::path::Path;
|
use std::path::Path;
|
||||||
|
|
||||||
|
use std::collections::HashSet;
|
||||||
|
|
||||||
use crate::bm25;
|
use crate::bm25;
|
||||||
|
use crate::confidence::{ConfidenceConfig, ScoredResult, reject_low_confidence};
|
||||||
use crate::hybrid;
|
use crate::hybrid;
|
||||||
|
use crate::reranker::{ReRankConfig, RerankInput, rerank};
|
||||||
use crate::{HDF5Memory, MAX_ACTIVATION_WEIGHT, MemoryError, Result, SearchResult};
|
use crate::{HDF5Memory, MAX_ACTIVATION_WEIGHT, MemoryError, Result, SearchResult};
|
||||||
|
|
||||||
|
/// Options for [`HDF5Memory::search`].
|
||||||
|
///
|
||||||
|
/// [`SearchOptions::new`] is plain hybrid search with the tuned default
|
||||||
|
/// fusion — the same as `hybrid_search_with(.., hybrid::DEFAULT_FUSION, k)`.
|
||||||
|
/// Every stage beyond that is opt-in.
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
pub struct SearchOptions {
|
||||||
|
/// Number of results to return.
|
||||||
|
pub k: usize,
|
||||||
|
/// How the vector and keyword stages are combined.
|
||||||
|
pub fusion: hybrid::Fusion,
|
||||||
|
/// Only consider records whose `source_channel` is one of these. The
|
||||||
|
/// filter applies *before* ranking, so a filtered search still returns up
|
||||||
|
/// to `k` results and scores are normalised over the records it can
|
||||||
|
/// return. `None` searches everything; an empty list matches nothing.
|
||||||
|
pub source_channels: Option<Vec<String>>,
|
||||||
|
/// Re-rank a candidate pool by retrieval relevance, recency, source
|
||||||
|
/// authority and activation — the pipeline the OpenClaw backend runs.
|
||||||
|
pub rerank: Option<ReRankConfig>,
|
||||||
|
/// Candidates retrieved for re-ranking; 0 means `max(3k, 10)`.
|
||||||
|
pub rerank_pool: usize,
|
||||||
|
/// Drop low-confidence results (after re-ranking, when that is on).
|
||||||
|
pub confidence: Option<ConfidenceConfig>,
|
||||||
|
/// The time recency is measured from, in seconds since the epoch.
|
||||||
|
/// `None` uses the system clock.
|
||||||
|
pub now: Option<f64>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl SearchOptions {
|
||||||
|
pub fn new(k: usize) -> Self {
|
||||||
|
Self {
|
||||||
|
k,
|
||||||
|
fusion: hybrid::DEFAULT_FUSION,
|
||||||
|
source_channels: None,
|
||||||
|
rerank: None,
|
||||||
|
rerank_pool: 0,
|
||||||
|
confidence: None,
|
||||||
|
now: None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn with_fusion(mut self, fusion: hybrid::Fusion) -> Self {
|
||||||
|
self.fusion = fusion;
|
||||||
|
self
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Search only records from these source channels.
|
||||||
|
pub fn with_sources<S: Into<String>>(mut self, channels: impl IntoIterator<Item = S>) -> Self {
|
||||||
|
self.source_channels = Some(channels.into_iter().map(Into::into).collect());
|
||||||
|
self
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn with_rerank(mut self, config: ReRankConfig) -> Self {
|
||||||
|
self.rerank = Some(config);
|
||||||
|
self
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn with_confidence(mut self, config: ConfidenceConfig) -> Self {
|
||||||
|
self.confidence = Some(config);
|
||||||
|
self
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Measure recency from `now` (seconds since the epoch) instead of the
|
||||||
|
/// system clock — for reproducible results and tests.
|
||||||
|
pub fn at_time(mut self, now: f64) -> Self {
|
||||||
|
self.now = Some(now);
|
||||||
|
self
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Default for SearchOptions {
|
||||||
|
fn default() -> Self {
|
||||||
|
Self::new(10)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
impl HDF5Memory {
|
impl HDF5Memory {
|
||||||
/// Vector + keyword scoring stage of [`HDF5Memory::hybrid_search`].
|
/// Vector + keyword scoring stage of [`HDF5Memory::search`].
|
||||||
///
|
///
|
||||||
/// Without the `hnsw` feature this is a full linear cosine scan (the exact
|
/// Without the `hnsw` feature this is a full linear cosine scan (the exact
|
||||||
/// previous behaviour, also used as the correctness oracle in tests). With
|
/// previous behaviour, also used as the correctness oracle in tests). With
|
||||||
/// `hnsw` enabled and an index available, the vector candidates come from an
|
/// `hnsw` enabled and an index available, the vector candidates come from an
|
||||||
/// approximate-nearest-neighbour search over an over-fetched pool, then merge
|
/// approximate-nearest-neighbour search over an over-fetched pool, then merge
|
||||||
/// with BM25 via the shared [`hybrid::merge_vector_keyword`].
|
/// with BM25 via the shared [`hybrid::merge_vector_keyword`].
|
||||||
|
///
|
||||||
|
/// `exclude`, when given, marks records that must not be returned (1 =
|
||||||
|
/// excluded; it covers tombstones too). The index is over-fetched in
|
||||||
|
/// proportion to how much the mask removes. Surfacing `pool` candidates
|
||||||
|
/// costs the index roughly `pool × M` distance evaluations, while an exact
|
||||||
|
/// scan of the allowed records costs one each — so whenever that scan is
|
||||||
|
/// the cheaper of the two it is used instead, and it is also the fallback
|
||||||
|
/// if the pool comes back with too few allowed hits (the allowed records
|
||||||
|
/// sit away from the query). A filtered search never comes back short.
|
||||||
#[cfg(feature = "hnsw")]
|
#[cfg(feature = "hnsw")]
|
||||||
fn vector_keyword_search(
|
fn vector_keyword_search(
|
||||||
&mut self,
|
&mut self,
|
||||||
query_embedding: &[f32],
|
query_embedding: &[f32],
|
||||||
query_text: &str,
|
query_text: &str,
|
||||||
bm25: &bm25::BM25Index,
|
bm25: &bm25::BM25Index,
|
||||||
vector_weight: f32,
|
fusion: hybrid::Fusion,
|
||||||
keyword_weight: f32,
|
|
||||||
k: usize,
|
k: usize,
|
||||||
|
exclude: Option<&[u8]>,
|
||||||
) -> Vec<(usize, f32)> {
|
) -> Vec<(usize, f32)> {
|
||||||
self.ensure_hnsw_fresh();
|
self.ensure_hnsw_fresh();
|
||||||
|
let n = self.cache.len();
|
||||||
|
// Over-fetch so the merge sees a useful vector pool. `ef` is
|
||||||
|
// configurable, but the pool the fusion stage sees is not tied to it:
|
||||||
|
// a caller lowering `ef` for speed should not silently narrow what
|
||||||
|
// fusion has to work with.
|
||||||
|
let mut pool = (k * 8).max(64);
|
||||||
|
let mut allowed = n;
|
||||||
|
if let Some(ex) = exclude {
|
||||||
|
allowed = ex.iter().filter(|&&e| e == 0).count();
|
||||||
|
if allowed == 0 {
|
||||||
|
return Vec::new();
|
||||||
|
}
|
||||||
|
// Expect `pool` allowed hits if the filter is independent of the
|
||||||
|
// query's neighbourhood.
|
||||||
|
pool = pool.saturating_mul(n).div_ceil(allowed);
|
||||||
|
if allowed <= pool.saturating_mul(self.hnsw_m()) {
|
||||||
|
return self.exact_masked_search(query_embedding, query_text, bm25, fusion, k, ex);
|
||||||
|
}
|
||||||
|
}
|
||||||
match self.hnsw.as_ref() {
|
match self.hnsw.as_ref() {
|
||||||
Some(index) if !index.is_empty() && index.dimension() == query_embedding.len() => {
|
Some(index) if !index.is_empty() && index.dimension() == query_embedding.len() => {
|
||||||
// Over-fetch so the merge sees a useful vector pool; cosine
|
let ef = self.hnsw_ef_search(k).max(pool);
|
||||||
// distance from the index converts back to similarity (1 - d).
|
let candidates = index.search(query_embedding, pool, ef);
|
||||||
let pool = (k * 8).max(64);
|
// A quantised index returns approximate distances, and no
|
||||||
let vec_scores: Vec<(usize, f32)> = index
|
// amount of `ef` fixes that — the loss is in the distances,
|
||||||
.search(query_embedding, pool, pool)
|
// not the graph. Re-score the pool against the cache's exact
|
||||||
|
// embeddings, which cost nothing extra to keep: recall then
|
||||||
|
// matches an f32 index. See `BENCHMARKS.md`.
|
||||||
|
let exact = index.storage() == clawhdf5_ann::Storage::Int8;
|
||||||
|
let vec_scores: Vec<(usize, f32)> = candidates
|
||||||
.into_iter()
|
.into_iter()
|
||||||
.map(|(id, dist)| (id, 1.0 - dist))
|
.filter(|(id, _)| exclude.is_none_or(|ex| ex[*id] == 0))
|
||||||
|
.map(|(id, dist)| {
|
||||||
|
let score = if exact {
|
||||||
|
crate::vector_search::cosine_similarity(
|
||||||
|
query_embedding,
|
||||||
|
&self.cache.embeddings[id],
|
||||||
|
)
|
||||||
|
} else {
|
||||||
|
1.0 - dist
|
||||||
|
};
|
||||||
|
(id, score)
|
||||||
|
})
|
||||||
.collect();
|
.collect();
|
||||||
// Fusion normalises over every keyword match, so it needs all
|
// Fusion normalises over every keyword match, so it needs all
|
||||||
// the scores — but not ranked.
|
// the scores — but not ranked.
|
||||||
let kw_scores = bm25.scores(query_text);
|
let mut kw_scores = bm25.scores(query_text);
|
||||||
hybrid::merge_vector_keyword(
|
if let Some(ex) = exclude {
|
||||||
vec_scores,
|
if vec_scores.len() < k.min(allowed) {
|
||||||
kw_scores,
|
// The allowed records are not where the index looked.
|
||||||
vector_weight,
|
return self.exact_masked_search(
|
||||||
keyword_weight,
|
query_embedding,
|
||||||
k,
|
query_text,
|
||||||
)
|
bm25,
|
||||||
|
fusion,
|
||||||
|
k,
|
||||||
|
ex,
|
||||||
|
);
|
||||||
|
}
|
||||||
|
kw_scores.retain(|(id, _)| ex[*id] == 0);
|
||||||
|
}
|
||||||
|
hybrid::fuse(vec_scores, kw_scores, fusion, k)
|
||||||
}
|
}
|
||||||
_ => hybrid::hybrid_search(
|
_ => match exclude {
|
||||||
query_embedding,
|
Some(ex) => {
|
||||||
query_text,
|
self.exact_masked_search(query_embedding, query_text, bm25, fusion, k, ex)
|
||||||
&self.cache.embeddings,
|
}
|
||||||
&self.cache.chunks,
|
None => hybrid::hybrid_search_fused(
|
||||||
&self.cache.tombstones,
|
query_embedding,
|
||||||
bm25,
|
query_text,
|
||||||
vector_weight,
|
&self.cache.embeddings,
|
||||||
keyword_weight,
|
&self.cache.chunks,
|
||||||
k,
|
&self.cache.tombstones,
|
||||||
),
|
bm25,
|
||||||
|
fusion,
|
||||||
|
k,
|
||||||
|
),
|
||||||
|
},
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -66,21 +201,52 @@ impl HDF5Memory {
|
|||||||
query_embedding: &[f32],
|
query_embedding: &[f32],
|
||||||
query_text: &str,
|
query_text: &str,
|
||||||
bm25: &bm25::BM25Index,
|
bm25: &bm25::BM25Index,
|
||||||
vector_weight: f32,
|
fusion: hybrid::Fusion,
|
||||||
keyword_weight: f32,
|
|
||||||
k: usize,
|
k: usize,
|
||||||
|
exclude: Option<&[u8]>,
|
||||||
) -> Vec<(usize, f32)> {
|
) -> Vec<(usize, f32)> {
|
||||||
hybrid::hybrid_search(
|
match exclude {
|
||||||
query_embedding,
|
Some(ex) => self.exact_masked_search(query_embedding, query_text, bm25, fusion, k, ex),
|
||||||
query_text,
|
None => hybrid::hybrid_search_fused(
|
||||||
&self.cache.embeddings,
|
query_embedding,
|
||||||
&self.cache.chunks,
|
query_text,
|
||||||
&self.cache.tombstones,
|
&self.cache.embeddings,
|
||||||
bm25,
|
&self.cache.chunks,
|
||||||
vector_weight,
|
&self.cache.tombstones,
|
||||||
keyword_weight,
|
bm25,
|
||||||
k,
|
fusion,
|
||||||
)
|
k,
|
||||||
|
),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Exact hybrid search over the records `exclude` leaves (0 = allowed).
|
||||||
|
fn exact_masked_search(
|
||||||
|
&self,
|
||||||
|
query_embedding: &[f32],
|
||||||
|
query_text: &str,
|
||||||
|
bm25: &bm25::BM25Index,
|
||||||
|
fusion: hybrid::Fusion,
|
||||||
|
k: usize,
|
||||||
|
exclude: &[u8],
|
||||||
|
) -> Vec<(usize, f32)> {
|
||||||
|
let vec_scores =
|
||||||
|
hybrid::exact_vector_scores(query_embedding, &self.cache.embeddings, exclude);
|
||||||
|
let mut kw_scores = bm25.scores(query_text);
|
||||||
|
kw_scores.retain(|(id, _)| exclude.get(*id) == Some(&0));
|
||||||
|
hybrid::fuse(vec_scores, kw_scores, fusion, k)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// The exclusion mask for a source-channel filter: 1 for a tombstoned
|
||||||
|
/// record or one from a channel not in `channels`.
|
||||||
|
fn source_mask(&self, channels: &[String]) -> Vec<u8> {
|
||||||
|
let allowed: HashSet<&str> = channels.iter().map(String::as_str).collect();
|
||||||
|
self.cache
|
||||||
|
.source_channels
|
||||||
|
.iter()
|
||||||
|
.zip(&self.cache.tombstones)
|
||||||
|
.map(|(ch, &t)| u8::from(t != 0 || !allowed.contains(ch.as_str())))
|
||||||
|
.collect()
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Perform hybrid search combining cosine vector similarity and BM25 keyword search.
|
/// Perform hybrid search combining cosine vector similarity and BM25 keyword search.
|
||||||
@@ -92,6 +258,61 @@ impl HDF5Memory {
|
|||||||
keyword_weight: f32,
|
keyword_weight: f32,
|
||||||
k: usize,
|
k: usize,
|
||||||
) -> Vec<SearchResult> {
|
) -> Vec<SearchResult> {
|
||||||
|
self.hybrid_search_with(
|
||||||
|
query_embedding,
|
||||||
|
query_text,
|
||||||
|
hybrid::Fusion::Weighted {
|
||||||
|
vector: vector_weight,
|
||||||
|
keyword: keyword_weight,
|
||||||
|
},
|
||||||
|
k,
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// [`HDF5Memory::hybrid_search`] with the fusion method chosen explicitly.
|
||||||
|
///
|
||||||
|
/// [`hybrid::DEFAULT_FUSION`] is what the weighted form defaults to;
|
||||||
|
/// [`hybrid::Fusion::Rrf`] combines the two stages by rank instead of by
|
||||||
|
/// score.
|
||||||
|
pub fn hybrid_search_with(
|
||||||
|
&mut self,
|
||||||
|
query_embedding: &[f32],
|
||||||
|
query_text: &str,
|
||||||
|
fusion: hybrid::Fusion,
|
||||||
|
k: usize,
|
||||||
|
) -> Vec<SearchResult> {
|
||||||
|
self.search(
|
||||||
|
query_embedding,
|
||||||
|
query_text,
|
||||||
|
&SearchOptions::new(k).with_fusion(fusion),
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Hybrid search with optional source filtering, re-ranking and
|
||||||
|
/// confidence rejection — see [`SearchOptions`].
|
||||||
|
///
|
||||||
|
/// Stages, in order: vector + keyword retrieval over the records the
|
||||||
|
/// source filter allows; fusion; scaling by Hebbian activation; re-ranking
|
||||||
|
/// (if on) of a `rerank_pool` of candidates; confidence rejection (if on);
|
||||||
|
/// the top `k`. The records returned with a positive score get their
|
||||||
|
/// Hebbian boost.
|
||||||
|
pub fn search(
|
||||||
|
&mut self,
|
||||||
|
query_embedding: &[f32],
|
||||||
|
query_text: &str,
|
||||||
|
options: &SearchOptions,
|
||||||
|
) -> Vec<SearchResult> {
|
||||||
|
let k = options.k;
|
||||||
|
let fetch = match options.rerank {
|
||||||
|
Some(_) if options.rerank_pool > 0 => options.rerank_pool.max(k),
|
||||||
|
Some(_) => k.saturating_mul(3).max(10),
|
||||||
|
None => k,
|
||||||
|
};
|
||||||
|
let exclude = options
|
||||||
|
.source_channels
|
||||||
|
.as_deref()
|
||||||
|
.map(|channels| self.source_mask(channels));
|
||||||
|
|
||||||
// The keyword index lives for the life of the store and is updated
|
// The keyword index lives for the life of the store and is updated
|
||||||
// incrementally. Take it out for the duration of the call so the
|
// incrementally. Take it out for the duration of the call so the
|
||||||
// vector stage can borrow `self` mutably, then put it back.
|
// vector stage can borrow `self` mutably, then put it back.
|
||||||
@@ -101,10 +322,12 @@ impl HDF5Memory {
|
|||||||
query_embedding,
|
query_embedding,
|
||||||
query_text,
|
query_text,
|
||||||
&bm25,
|
&bm25,
|
||||||
vector_weight,
|
options.fusion,
|
||||||
keyword_weight,
|
fetch,
|
||||||
k,
|
exclude.as_deref(),
|
||||||
);
|
);
|
||||||
|
self.bm25 = Some(bm25);
|
||||||
|
|
||||||
let mut results: Vec<SearchResult> = scored
|
let mut results: Vec<SearchResult> = scored
|
||||||
.into_iter()
|
.into_iter()
|
||||||
.map(|(idx, score)| {
|
.map(|(idx, score)| {
|
||||||
@@ -128,6 +351,25 @@ impl HDF5Memory {
|
|||||||
.then(a.index.cmp(&b.index))
|
.then(a.index.cmp(&b.index))
|
||||||
});
|
});
|
||||||
|
|
||||||
|
if let Some(config) = &options.rerank {
|
||||||
|
results = Self::rerank_results(results, config, options.now);
|
||||||
|
}
|
||||||
|
if let Some(config) = &options.confidence {
|
||||||
|
let scored: Vec<ScoredResult> = results
|
||||||
|
.iter()
|
||||||
|
.map(|r| ScoredResult {
|
||||||
|
index: r.index,
|
||||||
|
score: r.score,
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
let keep: HashSet<usize> = reject_low_confidence(&scored, config)
|
||||||
|
.into_iter()
|
||||||
|
.map(|r| r.index)
|
||||||
|
.collect();
|
||||||
|
results.retain(|r| keep.contains(&r.index));
|
||||||
|
}
|
||||||
|
results.truncate(k);
|
||||||
|
|
||||||
// Only reinforce records that actually matched. When fewer than `k`
|
// Only reinforce records that actually matched. When fewer than `k`
|
||||||
// records are relevant, the rest of the list is zero-score filler;
|
// records are relevant, the rest of the list is zero-score filler;
|
||||||
// boosting it would teach the store that arbitrary records are
|
// boosting it would teach the store that arbitrary records are
|
||||||
@@ -138,11 +380,45 @@ impl HDF5Memory {
|
|||||||
.map(|r| r.index)
|
.map(|r| r.index)
|
||||||
.collect();
|
.collect();
|
||||||
self.apply_hebbian_boost(&hit_indices);
|
self.apply_hebbian_boost(&hit_indices);
|
||||||
self.bm25 = Some(bm25);
|
|
||||||
|
|
||||||
results
|
results
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Reorder by the re-ranker's combined score, which also becomes each
|
||||||
|
/// result's `score`.
|
||||||
|
fn rerank_results(
|
||||||
|
results: Vec<SearchResult>,
|
||||||
|
config: &ReRankConfig,
|
||||||
|
now: Option<f64>,
|
||||||
|
) -> Vec<SearchResult> {
|
||||||
|
let now = now.unwrap_or_else(|| {
|
||||||
|
std::time::SystemTime::now()
|
||||||
|
.duration_since(std::time::UNIX_EPOCH)
|
||||||
|
.map(|d| d.as_secs_f64())
|
||||||
|
.unwrap_or(0.0)
|
||||||
|
});
|
||||||
|
let inputs: Vec<RerankInput> = results
|
||||||
|
.iter()
|
||||||
|
.map(|r| RerankInput {
|
||||||
|
index: r.index,
|
||||||
|
timestamp: r.timestamp,
|
||||||
|
source_channel: r.source_channel.clone(),
|
||||||
|
raw_activation: r.activation,
|
||||||
|
relevance: r.score,
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
let mut by_index: std::collections::HashMap<usize, SearchResult> =
|
||||||
|
results.into_iter().map(|r| (r.index, r)).collect();
|
||||||
|
rerank(&inputs, config, now)
|
||||||
|
.into_iter()
|
||||||
|
.filter_map(|rr| {
|
||||||
|
let mut r = by_index.remove(&rr.index)?;
|
||||||
|
r.score = rr.combined_score;
|
||||||
|
Some(r)
|
||||||
|
})
|
||||||
|
.collect()
|
||||||
|
}
|
||||||
|
|
||||||
/// Reinforce the records a query returned. The new weights are persisted by
|
/// Reinforce the records a query returned. The new weights are persisted by
|
||||||
/// the next checkpoint (any write that flushes, `flush_wal`, or drop) — not
|
/// the next checkpoint (any write that flushes, `flush_wal`, or drop) — not
|
||||||
/// by rewriting the whole store inside the query, which is what made
|
/// by rewriting the whole store inside the query, which is what made
|
||||||
|
|||||||
@@ -33,7 +33,7 @@ impl SessionCache {
|
|||||||
self.entries.is_empty()
|
self.entries.is_empty()
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Add a new session with its summary.
|
/// Add a new session with its summary, timestamped now.
|
||||||
pub fn add(
|
pub fn add(
|
||||||
&mut self,
|
&mut self,
|
||||||
id: &str,
|
id: &str,
|
||||||
@@ -47,6 +47,21 @@ impl SessionCache {
|
|||||||
.unwrap_or_default()
|
.unwrap_or_default()
|
||||||
.as_secs_f64()
|
.as_secs_f64()
|
||||||
* 1_000_000.0; // microseconds
|
* 1_000_000.0; // microseconds
|
||||||
|
self.add_at(id, start_idx, end_idx, channel, summary, ts);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Add a session with an explicit timestamp (Unix **microseconds**, the
|
||||||
|
/// unit [`SessionEntry::ts`] uses) — for importers carrying sessions over
|
||||||
|
/// from another store, whose original time should be kept.
|
||||||
|
pub fn add_at(
|
||||||
|
&mut self,
|
||||||
|
id: &str,
|
||||||
|
start_idx: usize,
|
||||||
|
end_idx: usize,
|
||||||
|
channel: &str,
|
||||||
|
summary: &str,
|
||||||
|
ts: f64,
|
||||||
|
) {
|
||||||
self.entries.push(SessionEntry {
|
self.entries.push(SessionEntry {
|
||||||
id: id.to_string(),
|
id: id.to_string(),
|
||||||
start_idx: start_idx as u64,
|
start_idx: start_idx as u64,
|
||||||
|
|||||||
@@ -0,0 +1,419 @@
|
|||||||
|
//! Ed25519-signed checkpoints.
|
||||||
|
//!
|
||||||
|
//! When a signing key is set ([`crate::HDF5Memory::set_signing_key`]), every
|
||||||
|
//! checkpoint writes a signed manifest of the store: a SHA-256 per memory
|
||||||
|
//! record rolled into a Merkle root, plus hashes of the store's settings, its
|
||||||
|
//! sessions and its knowledge graph. [`verify_store`] recomputes all of it from
|
||||||
|
//! the file and checks the signature against a public key the caller trusts,
|
||||||
|
//! so any change to the checkpointed file — a record's text or embedding, a
|
||||||
|
//! setting, a session, a graph edge, made through this crate or any other HDF5
|
||||||
|
//! tool — is detected, and the per-record hashes say which records changed.
|
||||||
|
//!
|
||||||
|
//! What it does not cover: saves still only in the WAL (made since the last
|
||||||
|
//! checkpoint). [`VerifyReport::wal_entries_unsigned`] counts them.
|
||||||
|
//!
|
||||||
|
//! The hashes cover exactly what the file persists, in the form the loader
|
||||||
|
//! returns it, so a store verifies after any number of reopen/checkpoint
|
||||||
|
//! cycles. Derived data (L2 norms, the vector index) is not covered; it is
|
||||||
|
//! recomputed from covered data.
|
||||||
|
|
||||||
|
use ed25519_dalek::{Signature, Signer, Verifier};
|
||||||
|
pub use ed25519_dalek::{SigningKey, VerifyingKey};
|
||||||
|
use sha2::{Digest, Sha256};
|
||||||
|
|
||||||
|
use crate::MemoryConfig;
|
||||||
|
use crate::cache::MemoryCache;
|
||||||
|
use crate::knowledge::KnowledgeCache;
|
||||||
|
use crate::session::SessionCache;
|
||||||
|
use crate::wal::WalMark;
|
||||||
|
|
||||||
|
/// Version of the manifest encoding; part of what is signed.
|
||||||
|
pub const MANIFEST_VERSION: i64 = 1;
|
||||||
|
|
||||||
|
type Hash = [u8; 32];
|
||||||
|
|
||||||
|
/// The hashes a signature covers.
|
||||||
|
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||||
|
pub struct Manifest {
|
||||||
|
pub record_count: u64,
|
||||||
|
/// Merkle root over the per-record hashes.
|
||||||
|
pub records_root: Hash,
|
||||||
|
/// Settings persisted in `/meta`, plus the checkpoint's WAL mark.
|
||||||
|
pub settings: Hash,
|
||||||
|
pub sessions: Hash,
|
||||||
|
pub graph: Hash,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Manifest {
|
||||||
|
/// The exact bytes that are signed.
|
||||||
|
pub fn signed_bytes(&self) -> Vec<u8> {
|
||||||
|
let mut m = Vec::with_capacity(160);
|
||||||
|
m.extend_from_slice(b"clawhdf5-agent signed checkpoint\0");
|
||||||
|
m.extend_from_slice(&MANIFEST_VERSION.to_le_bytes());
|
||||||
|
m.extend_from_slice(&self.record_count.to_le_bytes());
|
||||||
|
m.extend_from_slice(&self.records_root);
|
||||||
|
m.extend_from_slice(&self.settings);
|
||||||
|
m.extend_from_slice(&self.sessions);
|
||||||
|
m.extend_from_slice(&self.graph);
|
||||||
|
m
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// A signature as stored in a checkpoint.
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
pub struct StoredSignature {
|
||||||
|
pub manifest: Manifest,
|
||||||
|
pub record_hashes: Vec<Hash>,
|
||||||
|
pub public_key: [u8; 32],
|
||||||
|
pub signature: [u8; 64],
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Build the manifest (and per-record hashes) for the state about to be
|
||||||
|
/// checkpointed, and sign it.
|
||||||
|
pub fn sign(
|
||||||
|
key: &SigningKey,
|
||||||
|
config: &MemoryConfig,
|
||||||
|
cache: &MemoryCache,
|
||||||
|
sessions: &SessionCache,
|
||||||
|
knowledge: &KnowledgeCache,
|
||||||
|
wal_applied: Option<WalMark>,
|
||||||
|
) -> StoredSignature {
|
||||||
|
let (manifest, record_hashes) = manifest(config, cache, sessions, knowledge, wal_applied);
|
||||||
|
let signature = key.sign(&manifest.signed_bytes()).to_bytes();
|
||||||
|
StoredSignature {
|
||||||
|
manifest,
|
||||||
|
record_hashes,
|
||||||
|
public_key: key.verifying_key().to_bytes(),
|
||||||
|
signature,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Compute the manifest of a store's state.
|
||||||
|
pub fn manifest(
|
||||||
|
config: &MemoryConfig,
|
||||||
|
cache: &MemoryCache,
|
||||||
|
sessions: &SessionCache,
|
||||||
|
knowledge: &KnowledgeCache,
|
||||||
|
wal_applied: Option<WalMark>,
|
||||||
|
) -> (Manifest, Vec<Hash>) {
|
||||||
|
let record_hashes: Vec<Hash> = (0..cache.len()).map(|i| record_hash(cache, i)).collect();
|
||||||
|
let manifest = Manifest {
|
||||||
|
record_count: cache.len() as u64,
|
||||||
|
records_root: merkle_root(&record_hashes),
|
||||||
|
settings: settings_hash(config, wal_applied),
|
||||||
|
sessions: sessions_hash(sessions),
|
||||||
|
graph: graph_hash(knowledge),
|
||||||
|
};
|
||||||
|
(manifest, record_hashes)
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// Canonical encoding
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
/// A SHA-256 over length-prefixed fields, so no two different field lists
|
||||||
|
/// hash the same bytes.
|
||||||
|
struct Fields(Sha256);
|
||||||
|
|
||||||
|
impl Fields {
|
||||||
|
fn new(domain: &str) -> Self {
|
||||||
|
let mut h = Sha256::new();
|
||||||
|
h.update((domain.len() as u64).to_le_bytes());
|
||||||
|
h.update(domain.as_bytes());
|
||||||
|
Self(h)
|
||||||
|
}
|
||||||
|
fn bytes(&mut self, b: &[u8]) -> &mut Self {
|
||||||
|
self.0.update((b.len() as u64).to_le_bytes());
|
||||||
|
self.0.update(b);
|
||||||
|
self
|
||||||
|
}
|
||||||
|
/// Strings as the loader returns them: stored null-padded, so a trailing
|
||||||
|
/// NUL cannot survive a round trip and must not be part of the hash.
|
||||||
|
fn str(&mut self, s: &str) -> &mut Self {
|
||||||
|
self.bytes(s.trim_end_matches('\0').as_bytes())
|
||||||
|
}
|
||||||
|
fn u64(&mut self, v: u64) -> &mut Self {
|
||||||
|
self.0.update(v.to_le_bytes());
|
||||||
|
self
|
||||||
|
}
|
||||||
|
fn f64(&mut self, v: f64) -> &mut Self {
|
||||||
|
self.0.update(v.to_bits().to_le_bytes());
|
||||||
|
self
|
||||||
|
}
|
||||||
|
fn f32(&mut self, v: f32) -> &mut Self {
|
||||||
|
self.0.update(v.to_bits().to_le_bytes());
|
||||||
|
self
|
||||||
|
}
|
||||||
|
fn finish(self) -> Hash {
|
||||||
|
self.0.finalize().into()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Everything persisted about record `i`, including its position. The
|
||||||
|
/// embedding is hashed as the cache holds it — for a `float16` store that is
|
||||||
|
/// the half-rounded value the file holds.
|
||||||
|
fn record_hash(cache: &MemoryCache, i: usize) -> Hash {
|
||||||
|
let mut f = Fields::new("clawhdf5-agent/record");
|
||||||
|
f.u64(i as u64).str(&cache.chunks[i]);
|
||||||
|
let emb: Vec<u8> = cache.embeddings[i]
|
||||||
|
.iter()
|
||||||
|
.flat_map(|v| v.to_bits().to_le_bytes())
|
||||||
|
.collect();
|
||||||
|
f.bytes(&emb)
|
||||||
|
.str(&cache.source_channels[i])
|
||||||
|
.f64(cache.timestamps[i])
|
||||||
|
.str(&cache.session_ids[i])
|
||||||
|
.str(&cache.tags[i])
|
||||||
|
.u64(u64::from(cache.tombstones[i]))
|
||||||
|
.f32(cache.activation_weights[i]);
|
||||||
|
f.finish()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Binary Merkle tree: leaves are the record hashes; a parent hashes its two
|
||||||
|
/// children with a node prefix; an odd node is carried up unchanged.
|
||||||
|
fn merkle_root(leaves: &[Hash]) -> Hash {
|
||||||
|
if leaves.is_empty() {
|
||||||
|
return Fields::new("clawhdf5-agent/merkle-empty").finish();
|
||||||
|
}
|
||||||
|
let mut level: Vec<Hash> = leaves.to_vec();
|
||||||
|
while level.len() > 1 {
|
||||||
|
level = level
|
||||||
|
.chunks(2)
|
||||||
|
.map(|pair| match pair {
|
||||||
|
[l, r] => {
|
||||||
|
let mut h = Sha256::new();
|
||||||
|
h.update([1u8]);
|
||||||
|
h.update(l);
|
||||||
|
h.update(r);
|
||||||
|
h.finalize().into()
|
||||||
|
}
|
||||||
|
[only] => *only,
|
||||||
|
_ => unreachable!(),
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
}
|
||||||
|
level[0]
|
||||||
|
}
|
||||||
|
|
||||||
|
fn settings_hash(c: &MemoryConfig, wal_applied: Option<WalMark>) -> Hash {
|
||||||
|
let mut f = Fields::new("clawhdf5-agent/settings");
|
||||||
|
f.str(crate::schema::SCHEMA_VERSION)
|
||||||
|
.str(&c.created_at)
|
||||||
|
.str(&c.agent_id)
|
||||||
|
.str(&c.embedder)
|
||||||
|
.u64(c.embedding_dim as u64)
|
||||||
|
.u64(c.chunk_size as u64)
|
||||||
|
.u64(c.overlap as u64)
|
||||||
|
.u64(u64::from(c.float16))
|
||||||
|
.u64(u64::from(c.compression))
|
||||||
|
.u64(u64::from(c.compression_level))
|
||||||
|
.f32(c.compact_threshold)
|
||||||
|
.f32(c.hebbian_boost)
|
||||||
|
.f32(c.decay_factor)
|
||||||
|
.u64(u64::from(c.wal_enabled))
|
||||||
|
.u64(c.wal_max_entries as u64)
|
||||||
|
.u64(u64::from(c.quantized_index))
|
||||||
|
.u64(c.hnsw_m as u64)
|
||||||
|
.u64(c.hnsw_ef_construction as u64)
|
||||||
|
.u64(c.hnsw_ef_search as u64);
|
||||||
|
// An empty mark is not written to the file, so it must hash as none.
|
||||||
|
match wal_applied.filter(|m| m.len > 0) {
|
||||||
|
Some(m) => f.u64(1).u64(m.len).u64(u64::from(m.crc)),
|
||||||
|
None => f.u64(0),
|
||||||
|
};
|
||||||
|
f.finish()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn sessions_hash(s: &SessionCache) -> Hash {
|
||||||
|
let mut f = Fields::new("clawhdf5-agent/sessions");
|
||||||
|
f.u64(s.entries.len() as u64);
|
||||||
|
for (i, e) in s.entries.iter().enumerate() {
|
||||||
|
f.str(&e.id)
|
||||||
|
.u64(e.start_idx)
|
||||||
|
.u64(e.end_idx)
|
||||||
|
.str(&e.channel)
|
||||||
|
.f64(e.ts)
|
||||||
|
.str(s.summaries.get(i).map(String::as_str).unwrap_or(""));
|
||||||
|
}
|
||||||
|
f.finish()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn graph_hash(k: &KnowledgeCache) -> Hash {
|
||||||
|
let mut f = Fields::new("clawhdf5-agent/graph");
|
||||||
|
f.u64(k.entities.len() as u64);
|
||||||
|
for e in &k.entities {
|
||||||
|
f.u64(e.id)
|
||||||
|
.str(&e.name)
|
||||||
|
.str(&e.entity_type)
|
||||||
|
.u64(e.embedding_idx as u64);
|
||||||
|
}
|
||||||
|
f.u64(k.relations.len() as u64);
|
||||||
|
for r in &k.relations {
|
||||||
|
f.u64(r.src)
|
||||||
|
.u64(r.tgt)
|
||||||
|
.str(&r.relation)
|
||||||
|
.f32(r.weight)
|
||||||
|
.f64(r.ts);
|
||||||
|
}
|
||||||
|
f.u64(k.alias_strings.len() as u64);
|
||||||
|
for (s, id) in k.alias_strings.iter().zip(&k.alias_entity_ids) {
|
||||||
|
f.str(s).u64(*id as u64);
|
||||||
|
}
|
||||||
|
f.finish()
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// Verification
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
/// The outcome of [`verify_store`].
|
||||||
|
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||||
|
pub struct VerifyReport {
|
||||||
|
/// The checkpoint carries a signature.
|
||||||
|
pub signed: bool,
|
||||||
|
/// The signature was made by the key the caller trusts.
|
||||||
|
pub key_matches: bool,
|
||||||
|
/// The signature over the stored manifest is valid.
|
||||||
|
pub signature_valid: bool,
|
||||||
|
/// The file's current contents match the signed manifest.
|
||||||
|
pub records_match: bool,
|
||||||
|
pub settings_match: bool,
|
||||||
|
pub sessions_match: bool,
|
||||||
|
pub graph_match: bool,
|
||||||
|
/// Records whose contents differ from what was signed (by position),
|
||||||
|
/// when the stored per-record hashes are themselves authentic.
|
||||||
|
pub changed_records: Vec<usize>,
|
||||||
|
/// Records in the file versus in the signed manifest.
|
||||||
|
pub record_count: u64,
|
||||||
|
pub signed_record_count: u64,
|
||||||
|
/// The public key the checkpoint claims to be signed by.
|
||||||
|
pub public_key: Option<[u8; 32]>,
|
||||||
|
/// Saves in the WAL after the checkpoint: not covered by the signature.
|
||||||
|
pub wal_entries_unsigned: usize,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl VerifyReport {
|
||||||
|
/// Signed by the trusted key, signature valid, and every part of the
|
||||||
|
/// file unchanged since it was signed.
|
||||||
|
pub fn is_valid(&self) -> bool {
|
||||||
|
self.signed
|
||||||
|
&& self.key_matches
|
||||||
|
&& self.signature_valid
|
||||||
|
&& self.records_match
|
||||||
|
&& self.settings_match
|
||||||
|
&& self.sessions_match
|
||||||
|
&& self.graph_match
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Check a store file against the public key the caller trusts.
|
||||||
|
///
|
||||||
|
/// Reads the checkpoint (not the WAL), recomputes every hash from its
|
||||||
|
/// contents and checks the signature. Never writes.
|
||||||
|
pub fn verify_store(
|
||||||
|
path: &std::path::Path,
|
||||||
|
trusted: &VerifyingKey,
|
||||||
|
) -> Result<VerifyReport, crate::MemoryError> {
|
||||||
|
let file = clawhdf5::File::open(path)
|
||||||
|
.map_err(|e| crate::MemoryError::Hdf5(format!("cannot open {}: {e}", path.display())))?;
|
||||||
|
let (config, cache, sessions, knowledge) = crate::schema::validate_and_load(&file)?;
|
||||||
|
let checkpoint = crate::schema::read_checkpoint_meta(&file);
|
||||||
|
let stored = crate::schema::read_signature(&file)?;
|
||||||
|
let wal_entries_unsigned = count_wal_entries_after(path, checkpoint.wal_applied);
|
||||||
|
|
||||||
|
let (current, current_hashes) = manifest(
|
||||||
|
&config,
|
||||||
|
&cache,
|
||||||
|
&sessions,
|
||||||
|
&knowledge,
|
||||||
|
checkpoint.wal_applied,
|
||||||
|
);
|
||||||
|
|
||||||
|
let Some(stored) = stored else {
|
||||||
|
return Ok(VerifyReport {
|
||||||
|
signed: false,
|
||||||
|
key_matches: false,
|
||||||
|
signature_valid: false,
|
||||||
|
records_match: false,
|
||||||
|
settings_match: false,
|
||||||
|
sessions_match: false,
|
||||||
|
graph_match: false,
|
||||||
|
changed_records: Vec::new(),
|
||||||
|
record_count: current.record_count,
|
||||||
|
signed_record_count: 0,
|
||||||
|
public_key: None,
|
||||||
|
wal_entries_unsigned,
|
||||||
|
});
|
||||||
|
};
|
||||||
|
|
||||||
|
let key_matches = stored.public_key == trusted.to_bytes();
|
||||||
|
let signature_valid = trusted
|
||||||
|
.verify(
|
||||||
|
&stored.manifest.signed_bytes(),
|
||||||
|
&Signature::from_bytes(&stored.signature),
|
||||||
|
)
|
||||||
|
.is_ok();
|
||||||
|
// The stored per-record hashes can localise a change only if they are
|
||||||
|
// the ones that were signed.
|
||||||
|
let hashes_authentic = signature_valid
|
||||||
|
&& stored.record_hashes.len() as u64 == stored.manifest.record_count
|
||||||
|
&& merkle_root(&stored.record_hashes) == stored.manifest.records_root;
|
||||||
|
let changed_records = if hashes_authentic {
|
||||||
|
let n = current_hashes.len().max(stored.record_hashes.len());
|
||||||
|
(0..n)
|
||||||
|
.filter(|&i| current_hashes.get(i) != stored.record_hashes.get(i))
|
||||||
|
.collect()
|
||||||
|
} else {
|
||||||
|
Vec::new()
|
||||||
|
};
|
||||||
|
|
||||||
|
Ok(VerifyReport {
|
||||||
|
signed: true,
|
||||||
|
key_matches,
|
||||||
|
signature_valid,
|
||||||
|
records_match: signature_valid
|
||||||
|
&& current.record_count == stored.manifest.record_count
|
||||||
|
&& current.records_root == stored.manifest.records_root,
|
||||||
|
settings_match: signature_valid && current.settings == stored.manifest.settings,
|
||||||
|
sessions_match: signature_valid && current.sessions == stored.manifest.sessions,
|
||||||
|
graph_match: signature_valid && current.graph == stored.manifest.graph,
|
||||||
|
changed_records,
|
||||||
|
record_count: current.record_count,
|
||||||
|
signed_record_count: stored.manifest.record_count,
|
||||||
|
public_key: Some(stored.public_key),
|
||||||
|
wal_entries_unsigned,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
fn count_wal_entries_after(store: &std::path::Path, mark: Option<WalMark>) -> usize {
|
||||||
|
let wal = store.with_extension("h5.wal");
|
||||||
|
if !wal.exists() {
|
||||||
|
return 0;
|
||||||
|
}
|
||||||
|
crate::wal::WalFile::read_entries_for_migration(&wal, mark)
|
||||||
|
.map(|e| e.len())
|
||||||
|
.unwrap_or(0)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// A new random signing key from the operating system's RNG.
|
||||||
|
pub fn generate_key() -> SigningKey {
|
||||||
|
SigningKey::generate(&mut rand_core::OsRng)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Hex encoding for keys and signatures in attributes and the CLI.
|
||||||
|
pub fn to_hex(bytes: &[u8]) -> String {
|
||||||
|
bytes.iter().map(|b| format!("{b:02x}")).collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Parse hex into exactly `N` bytes.
|
||||||
|
pub fn from_hex<const N: usize>(s: &str) -> Option<[u8; N]> {
|
||||||
|
let s = s.trim();
|
||||||
|
if s.len() != 2 * N {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
let mut out = [0u8; N];
|
||||||
|
for (i, byte) in out.iter_mut().enumerate() {
|
||||||
|
*byte = u8::from_str_radix(&s[2 * i..2 * i + 2], 16).ok()?;
|
||||||
|
}
|
||||||
|
Some(out)
|
||||||
|
}
|
||||||
@@ -36,7 +36,7 @@ pub fn write_to_disk_with_mark(
|
|||||||
) -> Result<(), MemoryError> {
|
) -> Result<(), MemoryError> {
|
||||||
let meta = schema::CheckpointMeta {
|
let meta = schema::CheckpointMeta {
|
||||||
wal_applied,
|
wal_applied,
|
||||||
ann_generation: None,
|
..schema::CheckpointMeta::default()
|
||||||
};
|
};
|
||||||
write_to_disk_with_meta(path, config, cache, sessions, knowledge, &meta)
|
write_to_disk_with_meta(path, config, cache, sessions, knowledge, &meta)
|
||||||
}
|
}
|
||||||
@@ -50,7 +50,21 @@ pub fn write_to_disk_with_meta(
|
|||||||
knowledge: &KnowledgeCache,
|
knowledge: &KnowledgeCache,
|
||||||
checkpoint: &schema::CheckpointMeta,
|
checkpoint: &schema::CheckpointMeta,
|
||||||
) -> Result<(), MemoryError> {
|
) -> Result<(), MemoryError> {
|
||||||
let bytes = schema::build_hdf5_file_with_meta(config, cache, sessions, knowledge, checkpoint)?;
|
write_to_disk_signed(path, config, cache, sessions, knowledge, checkpoint, None)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// [`write_to_disk_with_meta`] with a signed manifest of the contents.
|
||||||
|
pub fn write_to_disk_signed(
|
||||||
|
path: &Path,
|
||||||
|
config: &MemoryConfig,
|
||||||
|
cache: &MemoryCache,
|
||||||
|
sessions: &SessionCache,
|
||||||
|
knowledge: &KnowledgeCache,
|
||||||
|
checkpoint: &schema::CheckpointMeta,
|
||||||
|
signature: Option<&crate::signing::StoredSignature>,
|
||||||
|
) -> Result<(), MemoryError> {
|
||||||
|
let bytes =
|
||||||
|
schema::build_hdf5_file_signed(config, cache, sessions, knowledge, checkpoint, signature)?;
|
||||||
|
|
||||||
if bytes.is_empty() {
|
if bytes.is_empty() {
|
||||||
return Err(MemoryError::Hdf5("build_hdf5_file produced 0 bytes".into()));
|
return Err(MemoryError::Hdf5("build_hdf5_file produced 0 bytes".into()));
|
||||||
@@ -113,13 +127,11 @@ pub type StoreState = (MemoryConfig, MemoryCache, SessionCache, KnowledgeCache);
|
|||||||
/// [`read_from_disk`], plus the checkpoint's [`WalMark`] (if any) so the
|
/// [`read_from_disk`], plus the checkpoint's [`WalMark`] (if any) so the
|
||||||
/// caller can skip WAL entries this file already contains.
|
/// caller can skip WAL entries this file already contains.
|
||||||
pub fn read_from_disk_with_mark(path: &Path) -> Result<(StoreState, Option<WalMark>), MemoryError> {
|
pub fn read_from_disk_with_mark(path: &Path) -> Result<(StoreState, Option<WalMark>), MemoryError> {
|
||||||
let mmap = clawhdf5_io::MmapReader::open(path).map_err(MemoryError::Io)?;
|
// `File::open` memory-maps the file itself (the facade's `mmap` feature is
|
||||||
|
// on by default). Mapping it here and handing over `as_bytes().to_vec()`
|
||||||
// Advise the OS we'll need the whole file for parsing
|
// did the same work and then copied the whole store — a second full copy
|
||||||
mmap.advise_willneed(0, mmap.len());
|
// of the file, live for the whole parse, on top of the mapping.
|
||||||
|
let file = clawhdf5::File::open(path)
|
||||||
// Parse the HDF5 file from the mmap'd bytes
|
|
||||||
let file = clawhdf5::File::from_bytes(mmap.as_bytes().to_vec())
|
|
||||||
.map_err(|e| MemoryError::Hdf5(format!("cannot open {}: {e}", path.display())))?;
|
.map_err(|e| MemoryError::Hdf5(format!("cannot open {}: {e}", path.display())))?;
|
||||||
|
|
||||||
let (mut config, cache, sessions, knowledge) = schema::validate_and_load(&file)?;
|
let (mut config, cache, sessions, knowledge) = schema::validate_and_load(&file)?;
|
||||||
@@ -133,9 +145,7 @@ pub fn read_from_disk_with_mark(path: &Path) -> Result<(StoreState, Option<WalMa
|
|||||||
pub fn read_from_disk_with_meta(
|
pub fn read_from_disk_with_meta(
|
||||||
path: &Path,
|
path: &Path,
|
||||||
) -> Result<(StoreState, schema::CheckpointMeta), MemoryError> {
|
) -> Result<(StoreState, schema::CheckpointMeta), MemoryError> {
|
||||||
let mmap = clawhdf5_io::MmapReader::open(path).map_err(MemoryError::Io)?;
|
let file = clawhdf5::File::open(path)
|
||||||
mmap.advise_willneed(0, mmap.len());
|
|
||||||
let file = clawhdf5::File::from_bytes(mmap.as_bytes().to_vec())
|
|
||||||
.map_err(|e| MemoryError::Hdf5(format!("cannot open {}: {e}", path.display())))?;
|
.map_err(|e| MemoryError::Hdf5(format!("cannot open {}: {e}", path.display())))?;
|
||||||
let (mut config, cache, sessions, knowledge) = schema::validate_and_load(&file)?;
|
let (mut config, cache, sessions, knowledge) = schema::validate_and_load(&file)?;
|
||||||
config.path = path.to_path_buf();
|
config.path = path.to_path_buf();
|
||||||
|
|||||||
@@ -4,6 +4,44 @@
|
|||||||
//! `clawhdf5_accel`, with optional float16 support via the `half` crate.
|
//! `clawhdf5_accel`, with optional float16 support via the `half` crate.
|
||||||
//! Supports pre-computed norms for eliminating redundant norm computations.
|
//! Supports pre-computed norms for eliminating redundant norm computations.
|
||||||
|
|
||||||
|
/// A corpus of equal-length embeddings addressable by index.
|
||||||
|
///
|
||||||
|
/// Lets the batch kernels read either the cache's flat `[N x dim]` buffer or a
|
||||||
|
/// plain `Vec<Vec<f32>>` without either side owning a second copy.
|
||||||
|
pub trait VectorSet {
|
||||||
|
/// Number of embeddings.
|
||||||
|
fn count(&self) -> usize;
|
||||||
|
/// Embedding `i`; callers only index below [`VectorSet::count`].
|
||||||
|
fn row(&self, i: usize) -> &[f32];
|
||||||
|
}
|
||||||
|
|
||||||
|
impl VectorSet for [Vec<f32>] {
|
||||||
|
fn count(&self) -> usize {
|
||||||
|
self.len()
|
||||||
|
}
|
||||||
|
fn row(&self, i: usize) -> &[f32] {
|
||||||
|
&self[i]
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl VectorSet for Vec<Vec<f32>> {
|
||||||
|
fn count(&self) -> usize {
|
||||||
|
self.len()
|
||||||
|
}
|
||||||
|
fn row(&self, i: usize) -> &[f32] {
|
||||||
|
&self[i]
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl VectorSet for crate::cache::Embeddings {
|
||||||
|
fn count(&self) -> usize {
|
||||||
|
self.len()
|
||||||
|
}
|
||||||
|
fn row(&self, i: usize) -> &[f32] {
|
||||||
|
&self[i]
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
/// Compute cosine similarity between two f32 slices.
|
/// Compute cosine similarity between two f32 slices.
|
||||||
///
|
///
|
||||||
/// Returns 0.0 if either vector has zero magnitude.
|
/// Returns 0.0 if either vector has zero magnitude.
|
||||||
@@ -22,7 +60,7 @@ pub fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
|
|||||||
/// Returns `(index, score)` pairs sorted by score descending.
|
/// Returns `(index, score)` pairs sorted by score descending.
|
||||||
pub fn cosine_similarity_batch(
|
pub fn cosine_similarity_batch(
|
||||||
query: &[f32],
|
query: &[f32],
|
||||||
vectors: &[Vec<f32>],
|
vectors: &(impl VectorSet + ?Sized),
|
||||||
tombstones: &[u8],
|
tombstones: &[u8],
|
||||||
) -> Vec<(usize, f32)> {
|
) -> Vec<(usize, f32)> {
|
||||||
let query_norm = clawhdf5_accel::vector_norm(query);
|
let query_norm = clawhdf5_accel::vector_norm(query);
|
||||||
@@ -30,7 +68,7 @@ pub fn cosine_similarity_batch(
|
|||||||
return Vec::new();
|
return Vec::new();
|
||||||
}
|
}
|
||||||
|
|
||||||
let n = vectors.len();
|
let n = vectors.count();
|
||||||
let mut results: Vec<(usize, f32)> = Vec::with_capacity(n);
|
let mut results: Vec<(usize, f32)> = Vec::with_capacity(n);
|
||||||
|
|
||||||
// Process 4 vectors at a time where possible
|
// Process 4 vectors at a time where possible
|
||||||
@@ -42,8 +80,9 @@ pub fn cosine_similarity_batch(
|
|||||||
if i < tombstones.len() && tombstones[i] != 0 {
|
if i < tombstones.len() && tombstones[i] != 0 {
|
||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
let vec_norm = clawhdf5_accel::vector_norm(&vectors[i]);
|
let vec_norm = clawhdf5_accel::vector_norm(vectors.row(i));
|
||||||
let score = crate::cosine_similarity_prenorm(query, query_norm, &vectors[i], vec_norm);
|
let score =
|
||||||
|
crate::cosine_similarity_prenorm(query, query_norm, vectors.row(i), vec_norm);
|
||||||
results.push((i, score));
|
results.push((i, score));
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -53,8 +92,8 @@ pub fn cosine_similarity_batch(
|
|||||||
if i < tombstones.len() && tombstones[i] != 0 {
|
if i < tombstones.len() && tombstones[i] != 0 {
|
||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
let vec_norm = clawhdf5_accel::vector_norm(&vectors[i]);
|
let vec_norm = clawhdf5_accel::vector_norm(vectors.row(i));
|
||||||
let score = crate::cosine_similarity_prenorm(query, query_norm, &vectors[i], vec_norm);
|
let score = crate::cosine_similarity_prenorm(query, query_norm, vectors.row(i), vec_norm);
|
||||||
results.push((i, score));
|
results.push((i, score));
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -68,7 +107,7 @@ pub fn cosine_similarity_batch(
|
|||||||
/// collections. Uses `score = dot(query, vec) / (query_norm * stored_norm)`.
|
/// collections. Uses `score = dot(query, vec) / (query_norm * stored_norm)`.
|
||||||
pub fn cosine_similarity_batch_prenorm(
|
pub fn cosine_similarity_batch_prenorm(
|
||||||
query: &[f32],
|
query: &[f32],
|
||||||
vectors: &[Vec<f32>],
|
vectors: &(impl VectorSet + ?Sized),
|
||||||
norms: &[f32],
|
norms: &[f32],
|
||||||
tombstones: &[u8],
|
tombstones: &[u8],
|
||||||
) -> Vec<(usize, f32)> {
|
) -> Vec<(usize, f32)> {
|
||||||
@@ -77,7 +116,7 @@ pub fn cosine_similarity_batch_prenorm(
|
|||||||
return Vec::new();
|
return Vec::new();
|
||||||
}
|
}
|
||||||
|
|
||||||
let n = vectors.len();
|
let n = vectors.count();
|
||||||
let mut results: Vec<(usize, f32)> = Vec::with_capacity(n);
|
let mut results: Vec<(usize, f32)> = Vec::with_capacity(n);
|
||||||
|
|
||||||
for i in 0..n {
|
for i in 0..n {
|
||||||
@@ -85,7 +124,7 @@ pub fn cosine_similarity_batch_prenorm(
|
|||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
let vec_norm = norms[i];
|
let vec_norm = norms[i];
|
||||||
let score = crate::cosine_similarity_prenorm(query, query_norm, &vectors[i], vec_norm);
|
let score = crate::cosine_similarity_prenorm(query, query_norm, vectors.row(i), vec_norm);
|
||||||
results.push((i, score));
|
results.push((i, score));
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -162,7 +201,7 @@ pub fn cosine_similarity_f16(
|
|||||||
#[cfg(feature = "parallel")]
|
#[cfg(feature = "parallel")]
|
||||||
pub fn parallel_cosine_batch(
|
pub fn parallel_cosine_batch(
|
||||||
query: &[f32],
|
query: &[f32],
|
||||||
vectors: &[Vec<f32>],
|
vectors: &(impl VectorSet + Sync + ?Sized),
|
||||||
tombstones: &[u8],
|
tombstones: &[u8],
|
||||||
k: usize,
|
k: usize,
|
||||||
) -> Vec<(usize, f32)> {
|
) -> Vec<(usize, f32)> {
|
||||||
@@ -174,24 +213,27 @@ pub fn parallel_cosine_batch(
|
|||||||
}
|
}
|
||||||
|
|
||||||
let num_cores = rayon::current_num_threads().max(1);
|
let num_cores = rayon::current_num_threads().max(1);
|
||||||
let chunk_size = vectors.len().div_ceil(num_cores);
|
let chunk_size = vectors.count().div_ceil(num_cores);
|
||||||
if chunk_size == 0 {
|
if chunk_size == 0 {
|
||||||
return Vec::new();
|
return Vec::new();
|
||||||
}
|
}
|
||||||
|
|
||||||
let mut all_results: Vec<(usize, f32)> = vectors
|
// Chunk over index ranges: the corpus may be one flat buffer rather than
|
||||||
.par_chunks(chunk_size)
|
// a slice of rows, so there is nothing to `par_chunks` over.
|
||||||
.enumerate()
|
let n = vectors.count();
|
||||||
.flat_map(|(chunk_idx, chunk)| {
|
let mut all_results: Vec<(usize, f32)> = (0..n.div_ceil(chunk_size))
|
||||||
|
.into_par_iter()
|
||||||
|
.flat_map(|chunk_idx| {
|
||||||
let base = chunk_idx * chunk_size;
|
let base = chunk_idx * chunk_size;
|
||||||
let mut local: Vec<(usize, f32)> = Vec::with_capacity(chunk.len());
|
let end = (base + chunk_size).min(n);
|
||||||
for (j, vec) in chunk.iter().enumerate() {
|
let mut local: Vec<(usize, f32)> = Vec::with_capacity(end - base);
|
||||||
let i = base + j;
|
for i in base..end {
|
||||||
if i < tombstones.len() && tombstones[i] != 0 {
|
if i < tombstones.len() && tombstones[i] != 0 {
|
||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
let vec_norm = clawhdf5_accel::vector_norm(vec);
|
let vec_norm = clawhdf5_accel::vector_norm(vectors.row(i));
|
||||||
let score = crate::cosine_similarity_prenorm(query, query_norm, vec, vec_norm);
|
let score =
|
||||||
|
crate::cosine_similarity_prenorm(query, query_norm, vectors.row(i), vec_norm);
|
||||||
local.push((i, score));
|
local.push((i, score));
|
||||||
}
|
}
|
||||||
local.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
|
local.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
|
||||||
@@ -209,7 +251,7 @@ pub fn parallel_cosine_batch(
|
|||||||
#[cfg(feature = "parallel")]
|
#[cfg(feature = "parallel")]
|
||||||
pub fn parallel_cosine_batch_prenorm(
|
pub fn parallel_cosine_batch_prenorm(
|
||||||
query: &[f32],
|
query: &[f32],
|
||||||
vectors: &[Vec<f32>],
|
vectors: &(impl VectorSet + Sync + ?Sized),
|
||||||
norms: &[f32],
|
norms: &[f32],
|
||||||
tombstones: &[u8],
|
tombstones: &[u8],
|
||||||
k: usize,
|
k: usize,
|
||||||
@@ -222,23 +264,26 @@ pub fn parallel_cosine_batch_prenorm(
|
|||||||
}
|
}
|
||||||
|
|
||||||
let num_cores = rayon::current_num_threads().max(1);
|
let num_cores = rayon::current_num_threads().max(1);
|
||||||
let chunk_size = vectors.len().div_ceil(num_cores);
|
let chunk_size = vectors.count().div_ceil(num_cores);
|
||||||
if chunk_size == 0 {
|
if chunk_size == 0 {
|
||||||
return Vec::new();
|
return Vec::new();
|
||||||
}
|
}
|
||||||
|
|
||||||
let mut all_results: Vec<(usize, f32)> = vectors
|
// Chunk over index ranges: the corpus may be one flat buffer rather than
|
||||||
.par_chunks(chunk_size)
|
// a slice of rows, so there is nothing to `par_chunks` over.
|
||||||
.enumerate()
|
let n = vectors.count();
|
||||||
.flat_map(|(chunk_idx, chunk)| {
|
let mut all_results: Vec<(usize, f32)> = (0..n.div_ceil(chunk_size))
|
||||||
|
.into_par_iter()
|
||||||
|
.flat_map(|chunk_idx| {
|
||||||
let base = chunk_idx * chunk_size;
|
let base = chunk_idx * chunk_size;
|
||||||
let mut local: Vec<(usize, f32)> = Vec::with_capacity(chunk.len());
|
let end = (base + chunk_size).min(n);
|
||||||
for (j, vec) in chunk.iter().enumerate() {
|
let mut local: Vec<(usize, f32)> = Vec::with_capacity(end - base);
|
||||||
let i = base + j;
|
for i in base..end {
|
||||||
if i < tombstones.len() && tombstones[i] != 0 {
|
if i < tombstones.len() && tombstones[i] != 0 {
|
||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
let score = crate::cosine_similarity_prenorm(query, query_norm, vec, norms[i]);
|
let score =
|
||||||
|
crate::cosine_similarity_prenorm(query, query_norm, vectors.row(i), norms[i]);
|
||||||
local.push((i, score));
|
local.push((i, score));
|
||||||
}
|
}
|
||||||
local.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
|
local.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
|
||||||
|
|||||||
Binary file not shown.
@@ -0,0 +1,260 @@
|
|||||||
|
//! `MemoryConfig::float16`: embeddings stored as IEEE half precision.
|
||||||
|
//!
|
||||||
|
//! The setting used to be recorded in `/meta` and otherwise ignored — the
|
||||||
|
//! embeddings dataset was always `f32`. These tests pin what it now does: the
|
||||||
|
//! dataset is `float16`, the in-memory cache holds exactly the values the file
|
||||||
|
//! holds (so search results survive a reopen bit for bit), and a value half
|
||||||
|
//! precision cannot represent is refused rather than stored as infinity.
|
||||||
|
|
||||||
|
use std::path::{Path, PathBuf};
|
||||||
|
|
||||||
|
use clawhdf5_agent::{AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry, MemoryError};
|
||||||
|
use clawhdf5_format::float16::round_to_f16;
|
||||||
|
use tempfile::TempDir;
|
||||||
|
|
||||||
|
const DIM: usize = 64;
|
||||||
|
|
||||||
|
/// Deterministic, embedding-like unit vectors.
|
||||||
|
fn embedding(seed: u64) -> Vec<f32> {
|
||||||
|
let mut x = seed.wrapping_mul(0x9E37_79B9_7F4A_7C15) | 1;
|
||||||
|
let v: Vec<f32> = (0..DIM)
|
||||||
|
.map(|_| {
|
||||||
|
x ^= x << 13;
|
||||||
|
x ^= x >> 7;
|
||||||
|
x ^= x << 17;
|
||||||
|
(x >> 40) as f32 / (1u64 << 24) as f32 - 0.5
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
let norm = v.iter().map(|a| a * a).sum::<f32>().sqrt();
|
||||||
|
v.iter().map(|a| a / norm).collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn entry(i: u64) -> MemoryEntry {
|
||||||
|
MemoryEntry {
|
||||||
|
chunk: format!("memory number {i} about topic {}", i % 7),
|
||||||
|
embedding: embedding(i),
|
||||||
|
source_channel: "test".into(),
|
||||||
|
timestamp: i as f64,
|
||||||
|
session_id: "s".into(),
|
||||||
|
tags: format!("t{i}"),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn config(dir: &TempDir, name: &str, float16: bool) -> MemoryConfig {
|
||||||
|
let mut c = MemoryConfig::new(dir.path().join(name), "agent", DIM);
|
||||||
|
c.float16 = float16;
|
||||||
|
c
|
||||||
|
}
|
||||||
|
|
||||||
|
fn embeddings_dtype_and_values(path: &Path) -> (String, Vec<f32>) {
|
||||||
|
let file = clawhdf5::File::open(path).unwrap();
|
||||||
|
let ds = file.dataset("memory/embeddings").unwrap();
|
||||||
|
(format!("{:?}", ds.dtype().unwrap()), ds.read_f32().unwrap())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn search_bits(m: &mut HDF5Memory, q: u64) -> Vec<(usize, u32)> {
|
||||||
|
m.hybrid_search(&embedding(q), "memory topic 3", 0.4, 0.6, 10)
|
||||||
|
.iter()
|
||||||
|
.map(|r| (r.index, r.score.to_bits()))
|
||||||
|
.collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn float16_store_writes_half_precision_and_reopens_identically() {
|
||||||
|
let dir = TempDir::new().unwrap();
|
||||||
|
// Two identical stores. Search is not read-only (it boosts the Hebbian
|
||||||
|
// activation of what it returns, and checkpoints persist that), so each
|
||||||
|
// is queried exactly once: one live, one after a checkpoint and reopen.
|
||||||
|
let live_cfg = config(&dir, "live.h5", true);
|
||||||
|
let cfg = config(&dir, "f16.h5", true);
|
||||||
|
let path: PathBuf = cfg.path.clone();
|
||||||
|
|
||||||
|
let mut live = HDF5Memory::create(live_cfg).unwrap();
|
||||||
|
live.save_batch((0..200).map(entry).collect()).unwrap();
|
||||||
|
let mut m = HDF5Memory::create(cfg).unwrap();
|
||||||
|
m.save_batch((0..200).map(entry).collect()).unwrap();
|
||||||
|
drop(m);
|
||||||
|
|
||||||
|
// On disk: a genuine float16 dataset holding the rounded inputs.
|
||||||
|
let (dtype, values) = embeddings_dtype_and_values(&path);
|
||||||
|
assert_eq!(dtype, "Other(\"float16\")");
|
||||||
|
let expected: Vec<u32> = (0..200)
|
||||||
|
.flat_map(|i| embedding(i).into_iter().map(|v| round_to_f16(v).to_bits()))
|
||||||
|
.collect();
|
||||||
|
let got: Vec<u32> = values.iter().map(|v| v.to_bits()).collect();
|
||||||
|
assert_eq!(got, expected);
|
||||||
|
|
||||||
|
// Reopened, the store answers exactly as the live one does: the cache
|
||||||
|
// held the half-rounded values before the checkpoint.
|
||||||
|
let mut reopened = HDF5Memory::open(&path).unwrap();
|
||||||
|
for q in 0..5 {
|
||||||
|
assert_eq!(
|
||||||
|
search_bits(&mut live, 1000 + q),
|
||||||
|
search_bits(&mut reopened, 1000 + q),
|
||||||
|
"query {q}"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn float16_halves_the_embeddings_on_disk() {
|
||||||
|
let dir = TempDir::new().unwrap();
|
||||||
|
let mut sizes = Vec::new();
|
||||||
|
for float16 in [false, true] {
|
||||||
|
let cfg = config(&dir, &format!("s{float16}.h5"), float16);
|
||||||
|
let path = cfg.path.clone();
|
||||||
|
let mut m = HDF5Memory::create(cfg).unwrap();
|
||||||
|
m.save_batch((0..2000).map(entry).collect()).unwrap();
|
||||||
|
drop(m);
|
||||||
|
sizes.push(std::fs::metadata(&path).unwrap().len());
|
||||||
|
}
|
||||||
|
let embedding_bytes_f32 = (2000 * DIM * 4) as u64;
|
||||||
|
let saved = sizes[0] - sizes[1];
|
||||||
|
// Half of the f32 embeddings, give or take metadata and alignment.
|
||||||
|
assert!(
|
||||||
|
saved.abs_diff(embedding_bytes_f32 / 2) < 16 * 1024,
|
||||||
|
"f32 {} B, f16 {} B, saved {saved} B, expected ~{} B",
|
||||||
|
sizes[0],
|
||||||
|
sizes[1],
|
||||||
|
embedding_bytes_f32 / 2
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn f32_store_is_unchanged() {
|
||||||
|
let dir = TempDir::new().unwrap();
|
||||||
|
let cfg = config(&dir, "f32.h5", false);
|
||||||
|
let path = cfg.path.clone();
|
||||||
|
let mut m = HDF5Memory::create(cfg).unwrap();
|
||||||
|
m.save_batch((0..50).map(entry).collect()).unwrap();
|
||||||
|
drop(m);
|
||||||
|
let (dtype, values) = embeddings_dtype_and_values(&path);
|
||||||
|
assert_eq!(dtype, "F32");
|
||||||
|
let expected: Vec<f32> = (0..50).flat_map(embedding).collect();
|
||||||
|
assert_eq!(values, expected);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn out_of_range_values_are_refused_not_stored_as_infinity() {
|
||||||
|
let dir = TempDir::new().unwrap();
|
||||||
|
let mut cfg = config(&dir, "range.h5", true);
|
||||||
|
cfg.wal_enabled = true;
|
||||||
|
let path = cfg.path.clone();
|
||||||
|
let mut m = HDF5Memory::create(cfg).unwrap();
|
||||||
|
m.save(entry(1)).unwrap();
|
||||||
|
|
||||||
|
let mut bad = entry(2);
|
||||||
|
bad.embedding[5] = 70_000.0;
|
||||||
|
match m.save(bad.clone()) {
|
||||||
|
Err(MemoryError::InvalidEntry(msg)) => assert!(msg.contains("embedding[5]"), "{msg}"),
|
||||||
|
other => panic!("expected InvalidEntry, got {other:?}"),
|
||||||
|
}
|
||||||
|
assert!(matches!(
|
||||||
|
m.save_or_update(bad.clone()),
|
||||||
|
Err(MemoryError::InvalidEntry(_))
|
||||||
|
));
|
||||||
|
// A batch is all or nothing.
|
||||||
|
assert!(matches!(
|
||||||
|
m.save_batch(vec![entry(3), bad.clone(), entry(4)]),
|
||||||
|
Err(MemoryError::InvalidEntry(_))
|
||||||
|
));
|
||||||
|
assert_eq!(m.count(), 1);
|
||||||
|
|
||||||
|
// The largest finite half, and values that round down to it, are fine.
|
||||||
|
let mut edge = entry(5);
|
||||||
|
edge.embedding[0] = 65504.0;
|
||||||
|
edge.embedding[1] = -65519.0;
|
||||||
|
m.save(edge).unwrap();
|
||||||
|
assert_eq!(m.count(), 2);
|
||||||
|
drop(m);
|
||||||
|
|
||||||
|
// Nothing rejected reached the WAL or the file.
|
||||||
|
let m = HDF5Memory::open(&path).unwrap();
|
||||||
|
assert_eq!(m.count(), 2);
|
||||||
|
|
||||||
|
// An f32 store takes the same value as it always did.
|
||||||
|
let mut m32 = HDF5Memory::create(config(&dir, "range32.h5", false)).unwrap();
|
||||||
|
m32.save(bad).unwrap();
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn wal_replay_rounds_like_a_live_save() {
|
||||||
|
let dir = TempDir::new().unwrap();
|
||||||
|
let mut cfg = config(&dir, "wal.h5", true);
|
||||||
|
cfg.wal_enabled = true;
|
||||||
|
cfg.wal_max_entries = 10_000; // keep everything in the WAL
|
||||||
|
let path = cfg.path.clone();
|
||||||
|
let mut m = HDF5Memory::create(cfg).unwrap();
|
||||||
|
for i in 0..30 {
|
||||||
|
m.save(entry(i)).unwrap();
|
||||||
|
}
|
||||||
|
let live = search_bits(&mut m, 77);
|
||||||
|
|
||||||
|
// Crash image: the .h5 is still the empty checkpoint; everything is in
|
||||||
|
// the WAL, which holds the caller's f32 values.
|
||||||
|
let crash = TempDir::new().unwrap();
|
||||||
|
let image = crash.path().join("image.h5");
|
||||||
|
std::fs::copy(&path, &image).unwrap();
|
||||||
|
std::fs::copy(
|
||||||
|
path.with_extension("h5.wal"),
|
||||||
|
image.with_extension("h5.wal"),
|
||||||
|
)
|
||||||
|
.unwrap();
|
||||||
|
drop(m);
|
||||||
|
|
||||||
|
let mut recovered = HDF5Memory::open(&image).unwrap();
|
||||||
|
assert_eq!(recovered.count(), 30);
|
||||||
|
assert_eq!(search_bits(&mut recovered, 77), live);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn new_stores_default_to_float16() {
|
||||||
|
let dir = TempDir::new().unwrap();
|
||||||
|
let path = dir.path().join("default.h5");
|
||||||
|
let mut m = HDF5Memory::create(MemoryConfig::new(path.clone(), "agent", DIM)).unwrap();
|
||||||
|
assert!(m.config().float16);
|
||||||
|
m.save_batch((0..10).map(entry).collect()).unwrap();
|
||||||
|
drop(m);
|
||||||
|
assert_eq!(embeddings_dtype_and_values(&path).0, "Other(\"float16\")");
|
||||||
|
assert!(HDF5Memory::open(&path).unwrap().config().float16);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn an_existing_f32_store_stays_f32() {
|
||||||
|
// Written by the v2.5.0 CLI, with `float16 = 0` in /meta (every agent
|
||||||
|
// store has recorded it). Flipping the default for new stores must not
|
||||||
|
// reach back and round an existing store's embeddings.
|
||||||
|
let dir = TempDir::new().unwrap();
|
||||||
|
let path = dir.path().join("legacy.h5");
|
||||||
|
std::fs::copy(
|
||||||
|
concat!(
|
||||||
|
env!("CARGO_MANIFEST_DIR"),
|
||||||
|
"/tests/fixtures/store_v2_5_0.h5"
|
||||||
|
),
|
||||||
|
&path,
|
||||||
|
)
|
||||||
|
.unwrap();
|
||||||
|
let before = embeddings_dtype_and_values(&path);
|
||||||
|
assert_eq!(before.0, "F32");
|
||||||
|
|
||||||
|
let mut m = HDF5Memory::open(&path).unwrap();
|
||||||
|
assert!(!m.config().float16, "an old store must reopen as f32");
|
||||||
|
let dim = m.config().embedding_dim;
|
||||||
|
let odd: Vec<f32> = (0..dim).map(|i| 0.1 + i as f32 * 1e-4).collect();
|
||||||
|
m.save_batch(vec![MemoryEntry {
|
||||||
|
chunk: "added after the upgrade".into(),
|
||||||
|
embedding: odd.clone(),
|
||||||
|
source_channel: "test".into(),
|
||||||
|
timestamp: 1.0,
|
||||||
|
session_id: "s".into(),
|
||||||
|
tags: String::new(),
|
||||||
|
}])
|
||||||
|
.unwrap();
|
||||||
|
drop(m);
|
||||||
|
|
||||||
|
// Checkpointed: still f32, the old rows untouched and the new one exact.
|
||||||
|
let (dtype, values) = embeddings_dtype_and_values(&path);
|
||||||
|
assert_eq!(dtype, "F32");
|
||||||
|
assert_eq!(&values[..before.1.len()], before.1.as_slice());
|
||||||
|
assert_eq!(&values[before.1.len()..], odd.as_slice());
|
||||||
|
}
|
||||||
@@ -0,0 +1,155 @@
|
|||||||
|
//! An agent store is a standard HDF5 file: h5py can open it and read every
|
||||||
|
//! dataset.
|
||||||
|
//!
|
||||||
|
//! It could not: the float datatype's sign-bit position was hard-coded for
|
||||||
|
//! f64, so every f32 dataset (embeddings, norms, activation weights) made
|
||||||
|
//! libhdf5 refuse the file with "sign bit position out of bounds".
|
||||||
|
|
||||||
|
use std::process::Command;
|
||||||
|
|
||||||
|
use clawhdf5_agent::{AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry};
|
||||||
|
|
||||||
|
fn python() -> String {
|
||||||
|
std::env::var("CLAWHDF5_PYTHON").unwrap_or_else(|_| "python3".to_string())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn h5py_available() -> bool {
|
||||||
|
Command::new(python())
|
||||||
|
.args(["-c", "import h5py"])
|
||||||
|
.output()
|
||||||
|
.map(|o| o.status.success())
|
||||||
|
.unwrap_or(false)
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn h5py_reads_every_dataset_of_an_agent_store() {
|
||||||
|
if !h5py_available() {
|
||||||
|
assert!(
|
||||||
|
std::env::var("CLAWHDF5_REQUIRE_INTEROP").as_deref() != Ok("1"),
|
||||||
|
"CLAWHDF5_REQUIRE_INTEROP=1 but python3 with h5py is not available"
|
||||||
|
);
|
||||||
|
eprintln!("SKIP: python3 with h5py not available");
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
let dir = tempfile::tempdir().unwrap();
|
||||||
|
for float16 in [false, true] {
|
||||||
|
let path = dir.path().join(format!("store_{float16}.h5"));
|
||||||
|
let mut cfg = MemoryConfig::new(path.clone(), "agent", 8);
|
||||||
|
cfg.float16 = float16;
|
||||||
|
let mut m = HDF5Memory::create(cfg).unwrap();
|
||||||
|
// save_batch checkpoints, so the records are in the .h5, not the WAL.
|
||||||
|
m.save_batch(
|
||||||
|
(0..20)
|
||||||
|
.map(|i| MemoryEntry {
|
||||||
|
chunk: format!("memory {i}"),
|
||||||
|
embedding: (0..8).map(|j| ((i * 8 + j) as f32).sin()).collect(),
|
||||||
|
source_channel: "test".into(),
|
||||||
|
timestamp: i as f64,
|
||||||
|
session_id: "s".into(),
|
||||||
|
tags: String::new(),
|
||||||
|
})
|
||||||
|
.collect(),
|
||||||
|
)
|
||||||
|
.unwrap();
|
||||||
|
drop(m);
|
||||||
|
|
||||||
|
// Exact expected values, as bits: numpy's sin need not match Rust's
|
||||||
|
// to the last place.
|
||||||
|
let bits = (0..160)
|
||||||
|
.map(|k| (k as f32).sin().to_bits().to_string())
|
||||||
|
.collect::<Vec<_>>()
|
||||||
|
.join(",");
|
||||||
|
let script = format!(
|
||||||
|
r#"
|
||||||
|
import h5py, numpy as np
|
||||||
|
want = np.float16 if {py_bool} else np.float32
|
||||||
|
with h5py.File("{path}", "r") as f:
|
||||||
|
names = []
|
||||||
|
f.visititems(lambda n, o: names.append(n) if isinstance(o, h5py.Dataset) else None)
|
||||||
|
for n in names:
|
||||||
|
f[n][()] # every dataset must decode
|
||||||
|
e = f["memory/embeddings"]
|
||||||
|
assert e.dtype == want, e.dtype
|
||||||
|
assert e.shape == (20, 8), e.shape
|
||||||
|
ref = np.array([{bits}], dtype=np.uint32).view(np.float32).astype(want).reshape(20, 8)
|
||||||
|
assert (e[()] == ref).all()
|
||||||
|
assert f["memory/norms"].dtype == np.float32
|
||||||
|
print(len(names))
|
||||||
|
"#,
|
||||||
|
py_bool = if float16 { "True" } else { "False" },
|
||||||
|
path = path.display()
|
||||||
|
);
|
||||||
|
let out = Command::new(python())
|
||||||
|
.args(["-c", &script])
|
||||||
|
.output()
|
||||||
|
.unwrap();
|
||||||
|
assert!(
|
||||||
|
out.status.success(),
|
||||||
|
"float16={float16}: {}",
|
||||||
|
String::from_utf8_lossy(&out.stderr)
|
||||||
|
);
|
||||||
|
let n: usize = String::from_utf8_lossy(&out.stdout).trim().parse().unwrap();
|
||||||
|
assert!(n >= 10, "only {n} datasets");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn an_edit_made_with_h5py_breaks_the_signature_and_names_the_record() {
|
||||||
|
if !h5py_available() {
|
||||||
|
assert!(
|
||||||
|
std::env::var("CLAWHDF5_REQUIRE_INTEROP").as_deref() != Ok("1"),
|
||||||
|
"CLAWHDF5_REQUIRE_INTEROP=1 but python3 with h5py is not available"
|
||||||
|
);
|
||||||
|
eprintln!("SKIP: python3 with h5py not available");
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
use clawhdf5_agent::signing::SigningKey;
|
||||||
|
let dir = tempfile::tempdir().unwrap();
|
||||||
|
let path = dir.path().join("signed.h5");
|
||||||
|
let key = SigningKey::from_bytes(&[42; 32]);
|
||||||
|
let mut m = HDF5Memory::create(MemoryConfig::new(path.clone(), "agent", 8)).unwrap();
|
||||||
|
m.set_signing_key(key.clone());
|
||||||
|
m.save_batch(
|
||||||
|
(0..10)
|
||||||
|
.map(|i| MemoryEntry {
|
||||||
|
chunk: format!("memory {i}"),
|
||||||
|
embedding: (0..8).map(|j| ((i * 8 + j) as f32).cos()).collect(),
|
||||||
|
source_channel: "test".into(),
|
||||||
|
timestamp: i as f64,
|
||||||
|
session_id: "s".into(),
|
||||||
|
tags: String::new(),
|
||||||
|
})
|
||||||
|
.collect(),
|
||||||
|
)
|
||||||
|
.unwrap();
|
||||||
|
drop(m);
|
||||||
|
assert!(
|
||||||
|
HDF5Memory::verify(&path, &key.verifying_key())
|
||||||
|
.unwrap()
|
||||||
|
.is_valid()
|
||||||
|
);
|
||||||
|
|
||||||
|
// Someone edits one timestamp in place with h5py.
|
||||||
|
let script = format!(
|
||||||
|
r#"
|
||||||
|
import h5py
|
||||||
|
with h5py.File("{}", "r+") as f:
|
||||||
|
ts = f["memory/timestamps"]
|
||||||
|
ts[3] = 12345.0
|
||||||
|
"#,
|
||||||
|
path.display()
|
||||||
|
);
|
||||||
|
let out = Command::new(python())
|
||||||
|
.args(["-c", &script])
|
||||||
|
.output()
|
||||||
|
.unwrap();
|
||||||
|
assert!(
|
||||||
|
out.status.success(),
|
||||||
|
"{}",
|
||||||
|
String::from_utf8_lossy(&out.stderr)
|
||||||
|
);
|
||||||
|
|
||||||
|
let r = HDF5Memory::verify(&path, &key.verifying_key()).unwrap();
|
||||||
|
assert!(r.signature_valid && !r.is_valid(), "{r:?}");
|
||||||
|
assert_eq!(r.changed_records, vec![3]);
|
||||||
|
}
|
||||||
@@ -165,3 +165,182 @@ fn save_batch_then_search_is_consistent() {
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn quantized_index_matches_the_f32_index_after_re_scoring() {
|
||||||
|
// A quantised index holds approximate vectors, but the store still has the
|
||||||
|
// exact ones, so the query path re-scores the candidate pool before
|
||||||
|
// fusion. The results a caller sees should therefore be the same.
|
||||||
|
let dim = 64;
|
||||||
|
let n = 400;
|
||||||
|
let mut seed = 0x5EED_1234_5678_9ABC;
|
||||||
|
let vectors: Vec<Vec<f32>> = (0..n).map(|_| make_vector(&mut seed, dim)).collect();
|
||||||
|
let queries: Vec<Vec<f32>> = (0..20).map(|_| make_vector(&mut seed, dim)).collect();
|
||||||
|
|
||||||
|
let build = |dir: &TempDir, quantized: bool| {
|
||||||
|
let mut config = MemoryConfig::new(dir.path().join("mem.h5"), "agent", dim);
|
||||||
|
config.quantized_index = quantized;
|
||||||
|
let mut mem = HDF5Memory::create(config).unwrap();
|
||||||
|
for (i, v) in vectors.iter().enumerate() {
|
||||||
|
mem.save(entry(&format!("chunk {i}"), v.clone(), &format!("k{i}")))
|
||||||
|
.unwrap();
|
||||||
|
}
|
||||||
|
mem
|
||||||
|
};
|
||||||
|
|
||||||
|
let exact_dir = TempDir::new().unwrap();
|
||||||
|
let quant_dir = TempDir::new().unwrap();
|
||||||
|
let mut exact = build(&exact_dir, false);
|
||||||
|
let mut quantized = build(&quant_dir, true);
|
||||||
|
|
||||||
|
let k = 10;
|
||||||
|
let mut agree = 0;
|
||||||
|
for q in &queries {
|
||||||
|
let want: Vec<usize> = exact
|
||||||
|
.hybrid_search(q, "", 1.0, 0.0, k)
|
||||||
|
.iter()
|
||||||
|
.map(|r| r.index)
|
||||||
|
.collect();
|
||||||
|
agree += quantized
|
||||||
|
.hybrid_search(q, "", 1.0, 0.0, k)
|
||||||
|
.iter()
|
||||||
|
.filter(|r| want.contains(&r.index))
|
||||||
|
.count();
|
||||||
|
}
|
||||||
|
let overlap = agree as f64 / (k * queries.len()) as f64;
|
||||||
|
assert!(
|
||||||
|
overlap >= 0.95,
|
||||||
|
"quantised store should match the f32 one: {overlap}"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn quantized_index_setting_survives_a_reopen() {
|
||||||
|
let dir = TempDir::new().unwrap();
|
||||||
|
let path = dir.path().join("mem.h5");
|
||||||
|
let mut config = MemoryConfig::new(path.clone(), "agent", 8);
|
||||||
|
config.quantized_index = true;
|
||||||
|
let mut mem = HDF5Memory::create(config).unwrap();
|
||||||
|
let mut seed = 7;
|
||||||
|
for i in 0..30 {
|
||||||
|
mem.save(entry(&format!("c{i}"), make_vector(&mut seed, 8), "t"))
|
||||||
|
.unwrap();
|
||||||
|
}
|
||||||
|
mem.flush_wal().unwrap();
|
||||||
|
drop(mem);
|
||||||
|
|
||||||
|
// Reopening must not silently quadruple the index's memory, so the flag
|
||||||
|
// is part of the stored config rather than a per-session choice.
|
||||||
|
let reopened = HDF5Memory::open(&path).unwrap();
|
||||||
|
assert!(reopened.config().quantized_index);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn hnsw_parameters_are_configurable_and_persisted() {
|
||||||
|
// The graph degree and both candidate-list sizes used to be constants, so
|
||||||
|
// a deployment could not trade recall against memory or speed at all.
|
||||||
|
let dir = TempDir::new().unwrap();
|
||||||
|
let path = dir.path().join("mem.h5");
|
||||||
|
let mut config = MemoryConfig::new(path.clone(), "agent", 16);
|
||||||
|
config.hnsw_m = 8;
|
||||||
|
config.hnsw_ef_construction = 32;
|
||||||
|
config.hnsw_ef_search = 128;
|
||||||
|
let mut mem = HDF5Memory::create(config).unwrap();
|
||||||
|
|
||||||
|
let mut seed = 99;
|
||||||
|
let vectors: Vec<Vec<f32>> = (0..300).map(|_| make_vector(&mut seed, 16)).collect();
|
||||||
|
for (i, v) in vectors.iter().enumerate() {
|
||||||
|
mem.save(entry(&format!("c{i}"), v.clone(), "t")).unwrap();
|
||||||
|
}
|
||||||
|
// Still correct with a smaller graph: an exact match must rank first.
|
||||||
|
let top = mem.hybrid_search(&vectors[42], "", 1.0, 0.0, 1);
|
||||||
|
assert_eq!(top[0].index, 42);
|
||||||
|
|
||||||
|
mem.flush_wal().unwrap();
|
||||||
|
drop(mem);
|
||||||
|
let reopened = HDF5Memory::open(&path).unwrap();
|
||||||
|
assert_eq!(reopened.config().hnsw_m, 8);
|
||||||
|
assert_eq!(reopened.config().hnsw_ef_construction, 32);
|
||||||
|
assert_eq!(reopened.config().hnsw_ef_search, 128);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn degenerate_hnsw_parameters_do_not_panic() {
|
||||||
|
// `clawhdf5-ann` asserts m >= 2, so a zero from a config file — or from a
|
||||||
|
// caller who assumed 0 meant "default" — would abort the process inside
|
||||||
|
// the index builder. The store clamps instead.
|
||||||
|
let dir = TempDir::new().unwrap();
|
||||||
|
let mut config = MemoryConfig::new(dir.path().join("mem.h5"), "agent", 8);
|
||||||
|
config.hnsw_m = 0;
|
||||||
|
config.hnsw_ef_construction = 0;
|
||||||
|
config.hnsw_ef_search = 1;
|
||||||
|
let mut mem = HDF5Memory::create(config).unwrap();
|
||||||
|
|
||||||
|
let mut seed = 5;
|
||||||
|
let vectors: Vec<Vec<f32>> = (0..50).map(|_| make_vector(&mut seed, 8)).collect();
|
||||||
|
for (i, v) in vectors.iter().enumerate() {
|
||||||
|
mem.save(entry(&format!("c{i}"), v.clone(), "t")).unwrap();
|
||||||
|
}
|
||||||
|
let results = mem.hybrid_search(&vectors[7], "", 1.0, 0.0, 5);
|
||||||
|
assert_eq!(results[0].index, 7, "exact match should still rank first");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn new_stores_default_to_the_quantized_index() {
|
||||||
|
// int8 is the default because it is smaller and, with an exact re-score,
|
||||||
|
// faster at equal recall on every platform measured (see BENCHMARKS.md).
|
||||||
|
let dir = TempDir::new().unwrap();
|
||||||
|
let config = MemoryConfig::new(dir.path().join("mem.h5"), "agent", 8);
|
||||||
|
assert!(config.quantized_index);
|
||||||
|
|
||||||
|
let path = config.path.clone();
|
||||||
|
let mut mem = HDF5Memory::create(config).unwrap();
|
||||||
|
let mut seed = 3;
|
||||||
|
let vectors: Vec<Vec<f32>> = (0..40).map(|_| make_vector(&mut seed, 8)).collect();
|
||||||
|
for (i, v) in vectors.iter().enumerate() {
|
||||||
|
mem.save(entry(&format!("c{i}"), v.clone(), "t")).unwrap();
|
||||||
|
}
|
||||||
|
assert_eq!(
|
||||||
|
mem.hybrid_search(&vectors[11], "", 1.0, 0.0, 1)[0].index,
|
||||||
|
11
|
||||||
|
);
|
||||||
|
mem.flush_wal().unwrap();
|
||||||
|
drop(mem);
|
||||||
|
assert!(HDF5Memory::open(&path).unwrap().config().quantized_index);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn a_store_written_before_the_setting_existed_stays_f32() {
|
||||||
|
// `store_v2_5_0.h5` was written by the v2.5.0 CLI, before
|
||||||
|
// `quantized_index` or the HNSW parameters were persisted, so it carries
|
||||||
|
// none of them. Flipping the default for new stores must not reach back
|
||||||
|
// and change how an existing store's index is held.
|
||||||
|
let dir = TempDir::new().unwrap();
|
||||||
|
let path = dir.path().join("legacy.h5");
|
||||||
|
std::fs::copy(
|
||||||
|
concat!(
|
||||||
|
env!("CARGO_MANIFEST_DIR"),
|
||||||
|
"/tests/fixtures/store_v2_5_0.h5"
|
||||||
|
),
|
||||||
|
&path,
|
||||||
|
)
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
let bytes = std::fs::read(&path).unwrap();
|
||||||
|
assert!(
|
||||||
|
!bytes.windows(15).any(|w| w == b"quantized_index"),
|
||||||
|
"the fixture must predate the setting, or it tests nothing"
|
||||||
|
);
|
||||||
|
|
||||||
|
let mut mem = HDF5Memory::open(&path).unwrap();
|
||||||
|
assert!(
|
||||||
|
!mem.config().quantized_index,
|
||||||
|
"an old store must reopen with an f32 index"
|
||||||
|
);
|
||||||
|
assert_eq!(mem.config().hnsw_m, 16);
|
||||||
|
assert_eq!(mem.config().hnsw_ef_construction, 64);
|
||||||
|
assert_eq!(mem.count(), 6);
|
||||||
|
// And it still searches: entry 3's own embedding finds it first.
|
||||||
|
let hit = mem.hybrid_search(&[3.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], "", 1.0, 0.0, 1);
|
||||||
|
assert_eq!(hit[0].index, 3);
|
||||||
|
}
|
||||||
|
|||||||
@@ -0,0 +1,344 @@
|
|||||||
|
//! `HDF5Memory::search` with `SearchOptions`: source filtering, re-ranking and
|
||||||
|
//! confidence rejection in the store's own search path.
|
||||||
|
|
||||||
|
use std::collections::HashSet;
|
||||||
|
|
||||||
|
use clawhdf5_agent::confidence::ConfidenceConfig;
|
||||||
|
use clawhdf5_agent::reranker::ReRankConfig;
|
||||||
|
use clawhdf5_agent::{AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry, SearchOptions, hybrid};
|
||||||
|
use tempfile::TempDir;
|
||||||
|
|
||||||
|
const DIM: usize = 32;
|
||||||
|
const N: usize = 3000;
|
||||||
|
const CLUSTERS: usize = 20;
|
||||||
|
|
||||||
|
struct Rng(u64);
|
||||||
|
impl Rng {
|
||||||
|
fn next(&mut self) -> u64 {
|
||||||
|
self.0 = self.0.wrapping_add(0x9E37_79B9_7F4A_7C15);
|
||||||
|
let mut z = self.0;
|
||||||
|
z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
|
||||||
|
z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
|
||||||
|
z ^ (z >> 31)
|
||||||
|
}
|
||||||
|
fn unit(&mut self) -> f32 {
|
||||||
|
(self.next() >> 40) as f32 / (1u64 << 24) as f32 - 0.5
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn normalize(v: &mut [f32]) {
|
||||||
|
let n = v.iter().map(|x| x * x).sum::<f32>().sqrt();
|
||||||
|
v.iter_mut().for_each(|x| *x /= n);
|
||||||
|
}
|
||||||
|
|
||||||
|
struct Data {
|
||||||
|
vectors: Vec<Vec<f32>>,
|
||||||
|
cluster: Vec<usize>,
|
||||||
|
centres: Vec<Vec<f32>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
fn data() -> Data {
|
||||||
|
let mut rng = Rng(42);
|
||||||
|
let centres: Vec<Vec<f32>> = (0..CLUSTERS)
|
||||||
|
.map(|_| {
|
||||||
|
let mut c: Vec<f32> = (0..DIM).map(|_| rng.unit()).collect();
|
||||||
|
normalize(&mut c);
|
||||||
|
c
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
let mut vectors = Vec::new();
|
||||||
|
let mut cluster = Vec::new();
|
||||||
|
for i in 0..N {
|
||||||
|
let c = i % CLUSTERS;
|
||||||
|
let mut v: Vec<f32> = centres[c].iter().map(|x| x + rng.unit() * 0.3).collect();
|
||||||
|
normalize(&mut v);
|
||||||
|
vectors.push(v);
|
||||||
|
cluster.push(c);
|
||||||
|
}
|
||||||
|
Data {
|
||||||
|
vectors,
|
||||||
|
cluster,
|
||||||
|
centres,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Channel of record `i` for a filter keeping `percent`% of the store at
|
||||||
|
/// random (independent of the vectors).
|
||||||
|
fn random_channel(i: usize, rng_seed: u64, percent: u64) -> String {
|
||||||
|
let mut r = Rng(rng_seed ^ (i as u64 * 7919));
|
||||||
|
if r.next() % 100 < percent {
|
||||||
|
"keep".into()
|
||||||
|
} else {
|
||||||
|
"other".into()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn build(data: &Data, channel: impl Fn(usize) -> String) -> (TempDir, HDF5Memory) {
|
||||||
|
let dir = TempDir::new().unwrap();
|
||||||
|
let mut cfg = MemoryConfig::new(dir.path().join("s.h5"), "agent", DIM);
|
||||||
|
cfg.hebbian_boost = 0.0; // every query sees the same store
|
||||||
|
let mut m = HDF5Memory::create(cfg).unwrap();
|
||||||
|
let entries = data
|
||||||
|
.vectors
|
||||||
|
.iter()
|
||||||
|
.enumerate()
|
||||||
|
.map(|(i, v)| MemoryEntry {
|
||||||
|
chunk: format!("record {i} cluster {}", data.cluster[i]),
|
||||||
|
embedding: v.clone(),
|
||||||
|
source_channel: channel(i),
|
||||||
|
timestamp: i as f64,
|
||||||
|
session_id: "s".into(),
|
||||||
|
tags: format!("t{i}"),
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
m.save_batch(entries).unwrap();
|
||||||
|
(dir, m)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Exact top-k by cosine among the records `allowed` keeps.
|
||||||
|
fn exact_top(data: &Data, q: &[f32], k: usize, allowed: impl Fn(usize) -> bool) -> Vec<usize> {
|
||||||
|
let mut s: Vec<(usize, f32)> = (0..N)
|
||||||
|
.filter(|&i| allowed(i))
|
||||||
|
.map(|i| (i, data.vectors[i].iter().zip(q).map(|(a, b)| a * b).sum()))
|
||||||
|
.collect();
|
||||||
|
s.sort_by(|a, b| b.1.total_cmp(&a.1).then(a.0.cmp(&b.0)));
|
||||||
|
s.into_iter().take(k).map(|(i, _)| i).collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn query(data: &Data, i: usize) -> Vec<f32> {
|
||||||
|
let mut rng = Rng(1000 + i as u64);
|
||||||
|
let mut q: Vec<f32> = data.centres[i % CLUSTERS]
|
||||||
|
.iter()
|
||||||
|
.map(|x| x + rng.unit() * 0.3)
|
||||||
|
.collect();
|
||||||
|
normalize(&mut q);
|
||||||
|
q
|
||||||
|
}
|
||||||
|
|
||||||
|
fn vector_only(k: usize) -> SearchOptions {
|
||||||
|
SearchOptions::new(k).with_fusion(hybrid::Fusion::Weighted {
|
||||||
|
vector: 1.0,
|
||||||
|
keyword: 0.0,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn source_filter_returns_only_allowed_records_and_a_full_page() {
|
||||||
|
let d = data();
|
||||||
|
// At N = 3000 and k = 10 the index serves a filter only when that is
|
||||||
|
// cheaper than scanning the allowed records: pool = 80 * N / allowed
|
||||||
|
// candidates at ~M = 16 distances each, against `allowed` distances. So
|
||||||
|
// 90% goes through the index, 50% and 1% to the exact scan.
|
||||||
|
for percent in [90, 50, 1] {
|
||||||
|
let (_dir, mut m) = build(&d, |i| random_channel(i, 5, percent));
|
||||||
|
let allowed = |i: usize| random_channel(i, 5, percent) == "keep";
|
||||||
|
let mut hits = 0;
|
||||||
|
for qi in 0..40 {
|
||||||
|
let q = query(&d, qi);
|
||||||
|
let got = m.search(&q, "", &vector_only(10).with_sources(["keep"]));
|
||||||
|
assert_eq!(got.len(), 10, "{percent}%: short page");
|
||||||
|
assert!(got.iter().all(|r| r.source_channel == "keep"));
|
||||||
|
let want: HashSet<usize> = exact_top(&d, &q, 10, allowed).into_iter().collect();
|
||||||
|
hits += got.iter().filter(|r| want.contains(&r.index)).count();
|
||||||
|
}
|
||||||
|
let recall = hits as f64 / 400.0;
|
||||||
|
let floor = if percent == 90 { 0.95 } else { 1.0 };
|
||||||
|
assert!(recall >= floor, "{percent}%: recall@10 {recall}");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn filter_away_from_the_query_falls_back_to_an_exact_scan() {
|
||||||
|
// Channel = cluster, and the filter keeps two clusters (10% of the
|
||||||
|
// store) that are not the query's: the index's neighbourhood of the
|
||||||
|
// query holds none of them. The search must still return the exact
|
||||||
|
// top 10 among the allowed records, not a short or empty page.
|
||||||
|
let d = data();
|
||||||
|
let (_dir, mut m) = build(&d, |i| format!("c{}", d.cluster[i]));
|
||||||
|
for qi in 0..20 {
|
||||||
|
let q = query(&d, qi);
|
||||||
|
let a = format!("c{}", (qi + 7) % CLUSTERS);
|
||||||
|
let b = format!("c{}", (qi + 13) % CLUSTERS);
|
||||||
|
let got: Vec<usize> = m
|
||||||
|
.search(
|
||||||
|
&q,
|
||||||
|
"",
|
||||||
|
&vector_only(10).with_sources([a.clone(), b.clone()]),
|
||||||
|
)
|
||||||
|
.iter()
|
||||||
|
.map(|r| r.index)
|
||||||
|
.collect();
|
||||||
|
let want = exact_top(&d, &q, 10, |i| {
|
||||||
|
let c = format!("c{}", d.cluster[i]);
|
||||||
|
c == a || c == b
|
||||||
|
});
|
||||||
|
assert_eq!(got, want, "query {qi}");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn filter_edge_cases() {
|
||||||
|
let d = data();
|
||||||
|
let (_dir, mut m) = build(&d, |i| random_channel(i, 9, 50));
|
||||||
|
let q = query(&d, 0);
|
||||||
|
assert!(
|
||||||
|
m.search(
|
||||||
|
&q,
|
||||||
|
"cluster",
|
||||||
|
&SearchOptions::new(10).with_sources(Vec::<String>::new())
|
||||||
|
)
|
||||||
|
.is_empty()
|
||||||
|
);
|
||||||
|
assert!(
|
||||||
|
m.search(
|
||||||
|
&q,
|
||||||
|
"cluster",
|
||||||
|
&SearchOptions::new(10).with_sources(["nope"])
|
||||||
|
)
|
||||||
|
.is_empty()
|
||||||
|
);
|
||||||
|
// Keyword matches from other channels are filtered too.
|
||||||
|
let got = m.search(
|
||||||
|
&q,
|
||||||
|
"record cluster",
|
||||||
|
&SearchOptions::new(50).with_sources(["keep"]),
|
||||||
|
);
|
||||||
|
assert_eq!(got.len(), 50);
|
||||||
|
assert!(got.iter().all(|r| r.source_channel == "keep"));
|
||||||
|
// Deleted records never come back, filtered or not.
|
||||||
|
let first = got[0].index;
|
||||||
|
m.delete(first).unwrap();
|
||||||
|
let again = m.search(
|
||||||
|
&q,
|
||||||
|
"record cluster",
|
||||||
|
&SearchOptions::new(50).with_sources(["keep"]),
|
||||||
|
);
|
||||||
|
assert!(again.iter().all(|r| r.index != first));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn plain_options_equal_hybrid_search_with() {
|
||||||
|
// Two identical stores, so neither query sees the other's boosts.
|
||||||
|
let d = data();
|
||||||
|
let (_a, mut a) = build(&d, |i| random_channel(i, 3, 50));
|
||||||
|
let (_b, mut b) = build(&d, |i| random_channel(i, 3, 50));
|
||||||
|
for qi in 0..10 {
|
||||||
|
let q = query(&d, qi);
|
||||||
|
let x: Vec<(usize, u32)> = a
|
||||||
|
.search(&q, "record cluster 3", &SearchOptions::new(10))
|
||||||
|
.iter()
|
||||||
|
.map(|r| (r.index, r.score.to_bits()))
|
||||||
|
.collect();
|
||||||
|
let y: Vec<(usize, u32)> = b
|
||||||
|
.hybrid_search_with(&q, "record cluster 3", hybrid::DEFAULT_FUSION, 10)
|
||||||
|
.iter()
|
||||||
|
.map(|r| (r.index, r.score.to_bits()))
|
||||||
|
.collect();
|
||||||
|
assert_eq!(x, y);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn small_store(entries: &[(&str, &str, f64)]) -> (TempDir, HDF5Memory) {
|
||||||
|
let dir = TempDir::new().unwrap();
|
||||||
|
let mut m = HDF5Memory::create(MemoryConfig::new(dir.path().join("r.h5"), "a", 4)).unwrap();
|
||||||
|
m.save_batch(
|
||||||
|
entries
|
||||||
|
.iter()
|
||||||
|
.map(|(chunk, channel, ts)| MemoryEntry {
|
||||||
|
chunk: chunk.to_string(),
|
||||||
|
embedding: vec![1.0, 0.0, 0.0, 0.0],
|
||||||
|
source_channel: channel.to_string(),
|
||||||
|
timestamp: *ts,
|
||||||
|
session_id: "s".into(),
|
||||||
|
tags: String::new(),
|
||||||
|
})
|
||||||
|
.collect(),
|
||||||
|
)
|
||||||
|
.unwrap();
|
||||||
|
(dir, m)
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn rerank_breaks_relevance_ties_by_recency() {
|
||||||
|
// Identical text and vectors, so retrieval ties; re-ranking must put the
|
||||||
|
// newer record first and report the combined score.
|
||||||
|
let now = 1_000_000.0;
|
||||||
|
let (_d, mut m) = small_store(&[
|
||||||
|
("user prefers dark mode", "chat", now - 30.0 * 86_400.0),
|
||||||
|
("user prefers dark mode", "chat", now - 60.0),
|
||||||
|
]);
|
||||||
|
let q = [1.0, 0.0, 0.0, 0.0];
|
||||||
|
let plain = m.search(&q, "dark mode", &SearchOptions::new(2));
|
||||||
|
assert_eq!(plain[0].index, 0, "ties break by index without re-ranking");
|
||||||
|
let reranked = m.search(
|
||||||
|
&q,
|
||||||
|
"dark mode",
|
||||||
|
&SearchOptions::new(2)
|
||||||
|
.with_rerank(ReRankConfig::default())
|
||||||
|
.at_time(now),
|
||||||
|
);
|
||||||
|
assert_eq!(reranked[0].index, 1);
|
||||||
|
assert!(reranked[0].score > reranked[1].score);
|
||||||
|
assert_ne!(reranked[0].score.to_bits(), plain[0].score.to_bits());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn confidence_rejects_when_nothing_is_good_enough() {
|
||||||
|
let (_d, mut m) = small_store(&[("alpha", "chat", 0.0), ("beta", "chat", 0.0)]);
|
||||||
|
let q = [1.0, 0.0, 0.0, 0.0];
|
||||||
|
let strict = ConfidenceConfig {
|
||||||
|
min_score: 10.0,
|
||||||
|
..ConfidenceConfig::default()
|
||||||
|
};
|
||||||
|
assert!(
|
||||||
|
m.search(&q, "alpha", &SearchOptions::new(2).with_confidence(strict))
|
||||||
|
.is_empty()
|
||||||
|
);
|
||||||
|
let lenient = ConfidenceConfig {
|
||||||
|
min_score: 0.0,
|
||||||
|
min_gap: f32::INFINITY,
|
||||||
|
max_results: 1,
|
||||||
|
};
|
||||||
|
assert_eq!(
|
||||||
|
m.search(&q, "alpha", &SearchOptions::new(2).with_confidence(lenient))
|
||||||
|
.len(),
|
||||||
|
1
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn only_returned_results_are_reinforced() {
|
||||||
|
// With re-ranking, a pool of max(3k, 10) candidates is retrieved; only
|
||||||
|
// the k returned should gain activation.
|
||||||
|
let d = data();
|
||||||
|
let dir = TempDir::new().unwrap();
|
||||||
|
let path = dir.path().join("h.h5");
|
||||||
|
let mut m = HDF5Memory::create(MemoryConfig::new(path, "a", DIM)).unwrap();
|
||||||
|
m.save_batch(
|
||||||
|
(0..200)
|
||||||
|
.map(|i| MemoryEntry {
|
||||||
|
chunk: format!("record {i}"),
|
||||||
|
embedding: d.vectors[i].clone(),
|
||||||
|
source_channel: "chat".into(),
|
||||||
|
timestamp: i as f64,
|
||||||
|
session_id: "s".into(),
|
||||||
|
tags: String::new(),
|
||||||
|
})
|
||||||
|
.collect(),
|
||||||
|
)
|
||||||
|
.unwrap();
|
||||||
|
let q = query(&d, 0);
|
||||||
|
let got = m.search(
|
||||||
|
&q,
|
||||||
|
"record",
|
||||||
|
&SearchOptions::new(3).with_rerank(ReRankConfig::default()),
|
||||||
|
);
|
||||||
|
assert_eq!(got.len(), 3);
|
||||||
|
let returned: HashSet<usize> = got.iter().map(|r| r.index).collect();
|
||||||
|
// A second plain search reports each record's current activation.
|
||||||
|
let all = m.search(&q, "record", &SearchOptions::new(200));
|
||||||
|
for r in &all {
|
||||||
|
let boosted = r.activation > 1.0;
|
||||||
|
assert_eq!(boosted, returned.contains(&r.index), "record {}", r.index);
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,330 @@
|
|||||||
|
//! Ed25519-signed checkpoints: `HDF5Memory::set_signing_key` and
|
||||||
|
//! `HDF5Memory::verify`.
|
||||||
|
|
||||||
|
use std::path::Path;
|
||||||
|
|
||||||
|
use clawhdf5_agent::signing::{SigningKey, VerifyReport, VerifyingKey};
|
||||||
|
use clawhdf5_agent::storage;
|
||||||
|
use clawhdf5_agent::{AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry, MemoryError, schema};
|
||||||
|
use tempfile::TempDir;
|
||||||
|
|
||||||
|
const DIM: usize = 16;
|
||||||
|
|
||||||
|
fn key(seed: u8) -> SigningKey {
|
||||||
|
SigningKey::from_bytes(&[seed; 32])
|
||||||
|
}
|
||||||
|
|
||||||
|
fn entry(i: usize, chunk: &str) -> MemoryEntry {
|
||||||
|
MemoryEntry {
|
||||||
|
chunk: chunk.to_string(),
|
||||||
|
embedding: (0..DIM)
|
||||||
|
.map(|j| ((i * DIM + j) as f32 * 0.37).sin())
|
||||||
|
.collect(),
|
||||||
|
source_channel: "chat".into(),
|
||||||
|
timestamp: 1_700_000_000.0 + i as f64,
|
||||||
|
session_id: format!("s{}", i % 3),
|
||||||
|
tags: format!("t{i}"),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Awkward strings on purpose: they must hash the same after a round trip.
|
||||||
|
const TEXTS: [&str; 6] = [
|
||||||
|
"plain text",
|
||||||
|
"ünïcödé — 日本語 🙂",
|
||||||
|
"",
|
||||||
|
"trailing spaces ",
|
||||||
|
"tab\tand\nnewline",
|
||||||
|
"x",
|
||||||
|
];
|
||||||
|
|
||||||
|
fn signed_store(dir: &TempDir, float16: bool, k: &SigningKey) -> std::path::PathBuf {
|
||||||
|
let mut cfg = MemoryConfig::new(dir.path().join("s.h5"), "agent", DIM);
|
||||||
|
cfg.float16 = float16;
|
||||||
|
let path = cfg.path.clone();
|
||||||
|
let mut m = HDF5Memory::create(cfg).unwrap();
|
||||||
|
m.set_signing_key(k.clone());
|
||||||
|
let entries = (0..30).map(|i| entry(i, TEXTS[i % TEXTS.len()])).collect();
|
||||||
|
m.save_batch(entries).unwrap();
|
||||||
|
// Some graph and a deleted record, so every part of the manifest is used.
|
||||||
|
let a = m.knowledge_mut().add_entity("Alice", "person", 0);
|
||||||
|
let b = m.knowledge_mut().add_entity("Acme", "org", -1);
|
||||||
|
m.knowledge_mut().add_relation(a, b, "works_at", 0.75);
|
||||||
|
m.sessions_mut()
|
||||||
|
.add_at("s0", 0, 9, "chat", "first session", 1_700_000_000.0);
|
||||||
|
m.delete(4).unwrap();
|
||||||
|
m.flush_wal().unwrap();
|
||||||
|
path
|
||||||
|
}
|
||||||
|
|
||||||
|
fn verify(path: &Path, k: &SigningKey) -> VerifyReport {
|
||||||
|
HDF5Memory::verify(path, &k.verifying_key()).unwrap()
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn a_signed_store_verifies_through_reopen_and_checkpoint_cycles() {
|
||||||
|
for float16 in [true, false] {
|
||||||
|
let dir = TempDir::new().unwrap();
|
||||||
|
let k = key(7);
|
||||||
|
let path = signed_store(&dir, float16, &k);
|
||||||
|
let r = verify(&path, &k);
|
||||||
|
assert!(r.is_valid(), "float16={float16}: {r:?}");
|
||||||
|
assert_eq!(r.public_key, Some(k.verifying_key().to_bytes()));
|
||||||
|
assert_eq!(r.record_count, 30);
|
||||||
|
assert!(r.changed_records.is_empty());
|
||||||
|
|
||||||
|
// Reopen, change nothing, checkpoint again (with the key): still valid.
|
||||||
|
for _ in 0..3 {
|
||||||
|
let mut m = HDF5Memory::open(&path).unwrap();
|
||||||
|
assert!(m.is_signed());
|
||||||
|
m.set_signing_key(k.clone());
|
||||||
|
m.flush_wal().unwrap();
|
||||||
|
drop(m);
|
||||||
|
assert!(verify(&path, &k).is_valid());
|
||||||
|
}
|
||||||
|
// And after real changes, re-signed.
|
||||||
|
let mut m = HDF5Memory::open(&path).unwrap();
|
||||||
|
m.set_signing_key(k.clone());
|
||||||
|
m.save(entry(99, "added later")).unwrap();
|
||||||
|
m.hybrid_search(&entry(1, "").embedding, "text", 0.4, 0.6, 5);
|
||||||
|
m.flush_wal().unwrap();
|
||||||
|
drop(m);
|
||||||
|
let r = verify(&path, &k);
|
||||||
|
assert!(r.is_valid(), "{r:?}");
|
||||||
|
assert_eq!(r.record_count, 31);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn a_signed_store_refuses_to_checkpoint_without_its_key() {
|
||||||
|
let dir = TempDir::new().unwrap();
|
||||||
|
let k = key(1);
|
||||||
|
let path = signed_store(&dir, true, &k);
|
||||||
|
|
||||||
|
let mut m = HDF5Memory::open(&path).unwrap();
|
||||||
|
m.save(entry(50, "pending")).unwrap();
|
||||||
|
match m.flush_wal() {
|
||||||
|
Err(MemoryError::SigningKeyRequired(msg)) => assert!(msg.contains("signed"), "{msg}"),
|
||||||
|
other => panic!("expected SigningKeyRequired, got {other:?}"),
|
||||||
|
}
|
||||||
|
// The file is untouched and still valid; the save is still in the WAL.
|
||||||
|
let r = verify(&path, &k);
|
||||||
|
assert!(r.is_valid());
|
||||||
|
assert_eq!(r.wal_entries_unsigned, 1);
|
||||||
|
|
||||||
|
// Supplying the key lets the checkpoint through, signed.
|
||||||
|
m.set_signing_key(k.clone());
|
||||||
|
m.flush_wal().unwrap();
|
||||||
|
drop(m);
|
||||||
|
let r = verify(&path, &k);
|
||||||
|
assert!(r.is_valid());
|
||||||
|
assert_eq!((r.record_count, r.wal_entries_unsigned), (31, 0));
|
||||||
|
|
||||||
|
// Removing the signature on purpose writes it unsigned.
|
||||||
|
let mut m = HDF5Memory::open(&path).unwrap();
|
||||||
|
m.remove_signature();
|
||||||
|
m.flush_wal().unwrap();
|
||||||
|
drop(m);
|
||||||
|
let r = verify(&path, &k);
|
||||||
|
assert!(!r.signed && !r.is_valid());
|
||||||
|
assert!(!HDF5Memory::open(&path).unwrap().is_signed());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn the_wrong_key_does_not_verify_and_a_new_key_re_signs() {
|
||||||
|
let dir = TempDir::new().unwrap();
|
||||||
|
let (a, b) = (key(1), key(2));
|
||||||
|
let path = signed_store(&dir, true, &a);
|
||||||
|
let r = verify(&path, &b);
|
||||||
|
assert!(r.signed && !r.key_matches && !r.signature_valid && !r.is_valid());
|
||||||
|
|
||||||
|
let mut m = HDF5Memory::open(&path).unwrap();
|
||||||
|
m.set_signing_key(b.clone());
|
||||||
|
m.flush_wal().unwrap();
|
||||||
|
drop(m);
|
||||||
|
assert!(verify(&path, &b).is_valid());
|
||||||
|
assert!(!verify(&path, &a).is_valid());
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Rewrite the store with changed contents but the *old* signature — what
|
||||||
|
/// someone with write access to the file, but not the key, can do.
|
||||||
|
fn tamper(path: &Path, change: impl FnOnce(&mut Tampered)) {
|
||||||
|
let file = clawhdf5::File::open(path).unwrap();
|
||||||
|
let (config, cache, sessions, knowledge) = schema::validate_and_load(&file).unwrap();
|
||||||
|
let checkpoint = schema::read_checkpoint_meta(&file);
|
||||||
|
let signature = schema::read_signature(&file).unwrap().unwrap();
|
||||||
|
drop(file);
|
||||||
|
let mut t = Tampered {
|
||||||
|
config,
|
||||||
|
cache,
|
||||||
|
sessions,
|
||||||
|
knowledge,
|
||||||
|
};
|
||||||
|
change(&mut t);
|
||||||
|
storage::write_to_disk_signed(
|
||||||
|
path,
|
||||||
|
&t.config,
|
||||||
|
&t.cache,
|
||||||
|
&t.sessions,
|
||||||
|
&t.knowledge,
|
||||||
|
&checkpoint,
|
||||||
|
Some(&signature),
|
||||||
|
)
|
||||||
|
.unwrap();
|
||||||
|
}
|
||||||
|
|
||||||
|
struct Tampered {
|
||||||
|
config: MemoryConfig,
|
||||||
|
cache: clawhdf5_agent::cache::MemoryCache,
|
||||||
|
sessions: clawhdf5_agent::SessionCache,
|
||||||
|
knowledge: clawhdf5_agent::knowledge::KnowledgeCache,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn every_kind_of_edit_is_detected_and_located() {
|
||||||
|
let k = key(3);
|
||||||
|
type Edit = Box<dyn FnOnce(&mut Tampered)>;
|
||||||
|
type Case = (&'static str, Edit, fn(&VerifyReport) -> bool);
|
||||||
|
let cases: Vec<Case> = vec![
|
||||||
|
(
|
||||||
|
"record text",
|
||||||
|
Box::new(|t: &mut Tampered| t.cache.chunks[7] = "rewritten".into()),
|
||||||
|
|r| !r.records_match && r.changed_records == vec![7],
|
||||||
|
),
|
||||||
|
(
|
||||||
|
"one embedding value",
|
||||||
|
Box::new(|t: &mut Tampered| {
|
||||||
|
let mut e = t.cache.embeddings[12].to_vec();
|
||||||
|
e[3] = 0.5;
|
||||||
|
t.cache.embeddings.set(12, &e);
|
||||||
|
}),
|
||||||
|
|r| r.changed_records == vec![12],
|
||||||
|
),
|
||||||
|
(
|
||||||
|
"undelete",
|
||||||
|
Box::new(|t: &mut Tampered| t.cache.tombstones[4] = 0),
|
||||||
|
|r| r.changed_records == vec![4],
|
||||||
|
),
|
||||||
|
(
|
||||||
|
"timestamp",
|
||||||
|
Box::new(|t: &mut Tampered| t.cache.timestamps[20] += 1.0),
|
||||||
|
|r| r.changed_records == vec![20],
|
||||||
|
),
|
||||||
|
(
|
||||||
|
"record appended",
|
||||||
|
Box::new(|t: &mut Tampered| {
|
||||||
|
t.cache.push(
|
||||||
|
"new".into(),
|
||||||
|
vec![0.1; DIM],
|
||||||
|
"x".into(),
|
||||||
|
1.0,
|
||||||
|
"s".into(),
|
||||||
|
"".into(),
|
||||||
|
);
|
||||||
|
}),
|
||||||
|
|r| !r.records_match && r.changed_records == vec![30] && r.record_count == 31,
|
||||||
|
),
|
||||||
|
(
|
||||||
|
"setting",
|
||||||
|
Box::new(|t: &mut Tampered| t.config.agent_id = "someone-else".into()),
|
||||||
|
|r| !r.settings_match && r.records_match,
|
||||||
|
),
|
||||||
|
(
|
||||||
|
"session summary",
|
||||||
|
Box::new(|t: &mut Tampered| t.sessions.summaries[0] = "edited".into()),
|
||||||
|
|r| !r.sessions_match && r.records_match,
|
||||||
|
),
|
||||||
|
(
|
||||||
|
"graph edge",
|
||||||
|
Box::new(|t: &mut Tampered| t.knowledge.relations[0].weight = 1.0),
|
||||||
|
|r| !r.graph_match && r.records_match,
|
||||||
|
),
|
||||||
|
];
|
||||||
|
for (name, edit, check) in cases {
|
||||||
|
let dir = TempDir::new().unwrap();
|
||||||
|
let path = signed_store(&dir, true, &k);
|
||||||
|
tamper(&path, edit);
|
||||||
|
let r = verify(&path, &k);
|
||||||
|
assert!(
|
||||||
|
r.signed && r.key_matches && r.signature_valid,
|
||||||
|
"{name}: {r:?}"
|
||||||
|
);
|
||||||
|
assert!(!r.is_valid(), "{name}: edit not detected: {r:?}");
|
||||||
|
assert!(check(&r), "{name}: {r:?}");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn a_forged_manifest_fails_the_signature() {
|
||||||
|
// Recomputing the hashes for tampered contents does not help without the
|
||||||
|
// key: the signature no longer matches the manifest.
|
||||||
|
let dir = TempDir::new().unwrap();
|
||||||
|
let k = key(5);
|
||||||
|
let path = signed_store(&dir, true, &k);
|
||||||
|
let file = clawhdf5::File::open(&path).unwrap();
|
||||||
|
let (config, mut cache, sessions, knowledge) = schema::validate_and_load(&file).unwrap();
|
||||||
|
let checkpoint = schema::read_checkpoint_meta(&file);
|
||||||
|
let mut sig = schema::read_signature(&file).unwrap().unwrap();
|
||||||
|
drop(file);
|
||||||
|
cache.chunks[0] = "forged".into();
|
||||||
|
// Re-sign with an attacker key, then splice the victim's public key back.
|
||||||
|
let forged = clawhdf5_agent::signing::sign(
|
||||||
|
&key(66),
|
||||||
|
&config,
|
||||||
|
&cache,
|
||||||
|
&sessions,
|
||||||
|
&knowledge,
|
||||||
|
checkpoint.wal_applied,
|
||||||
|
);
|
||||||
|
sig.manifest = forged.manifest;
|
||||||
|
sig.record_hashes = forged.record_hashes;
|
||||||
|
storage::write_to_disk_signed(
|
||||||
|
&path,
|
||||||
|
&config,
|
||||||
|
&cache,
|
||||||
|
&sessions,
|
||||||
|
&knowledge,
|
||||||
|
&checkpoint,
|
||||||
|
Some(&sig),
|
||||||
|
)
|
||||||
|
.unwrap();
|
||||||
|
let r = verify(&path, &k);
|
||||||
|
assert!(
|
||||||
|
r.key_matches && !r.signature_valid && !r.is_valid(),
|
||||||
|
"{r:?}"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn an_unsigned_store_reports_unsigned() {
|
||||||
|
let dir = TempDir::new().unwrap();
|
||||||
|
let mut m = HDF5Memory::create(MemoryConfig::new(dir.path().join("u.h5"), "a", DIM)).unwrap();
|
||||||
|
m.save_batch(vec![entry(0, "hello")]).unwrap();
|
||||||
|
drop(m);
|
||||||
|
let r = HDF5Memory::verify(&dir.path().join("u.h5"), &VerifyingKey::from(&key(1))).unwrap();
|
||||||
|
assert!(!r.signed && !r.is_valid());
|
||||||
|
assert_eq!(r.record_count, 1);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn nul_bytes_in_text_still_verify() {
|
||||||
|
// Strings are stored null-padded; the hash must follow what a reopened
|
||||||
|
// store actually holds, or an untouched store would fail to verify.
|
||||||
|
let dir = TempDir::new().unwrap();
|
||||||
|
let k = key(9);
|
||||||
|
let mut m = HDF5Memory::create(MemoryConfig::new(dir.path().join("n.h5"), "a", DIM)).unwrap();
|
||||||
|
m.set_signing_key(k.clone());
|
||||||
|
m.save_batch(vec![
|
||||||
|
entry(0, "inner\0nul"),
|
||||||
|
entry(1, "trailing nul\0"),
|
||||||
|
entry(2, "\0leading"),
|
||||||
|
])
|
||||||
|
.unwrap();
|
||||||
|
drop(m);
|
||||||
|
let r = verify(&dir.path().join("n.h5"), &k);
|
||||||
|
assert!(r.is_valid(), "{r:?}");
|
||||||
|
let m = HDF5Memory::open(&dir.path().join("n.h5")).unwrap();
|
||||||
|
eprintln!(
|
||||||
|
"reloaded: {:?}",
|
||||||
|
(0..3).map(|i| m.get_chunk(i)).collect::<Vec<_>>()
|
||||||
|
);
|
||||||
|
}
|
||||||
@@ -1,7 +1,8 @@
|
|||||||
[package]
|
[package]
|
||||||
name = "clawhdf5-android"
|
name = "clawhdf5-android"
|
||||||
version = "2.4.0"
|
version = "2.7.0"
|
||||||
edition = "2024"
|
edition = "2024"
|
||||||
|
rust-version.workspace = true
|
||||||
description = "Android JNI bridge for edgehdf5-memory HDF5 backend"
|
description = "Android JNI bridge for edgehdf5-memory HDF5 backend"
|
||||||
license = "MIT"
|
license = "MIT"
|
||||||
|
|
||||||
|
|||||||
@@ -1,7 +1,8 @@
|
|||||||
[package]
|
[package]
|
||||||
name = "clawhdf5-ann"
|
name = "clawhdf5-ann"
|
||||||
version = "2.4.0"
|
version = "2.7.0"
|
||||||
edition = "2024"
|
edition = "2024"
|
||||||
|
rust-version.workspace = true
|
||||||
description = "HNSW approximate nearest neighbor index stored as HDF5"
|
description = "HNSW approximate nearest neighbor index stored as HDF5"
|
||||||
license = "MIT"
|
license = "MIT"
|
||||||
repository = "https://git.redclaw.dev/quantumclaw/clawhdf5"
|
repository = "https://git.redclaw.dev/quantumclaw/clawhdf5"
|
||||||
@@ -10,9 +11,9 @@ keywords = ["hdf5", "ann", "hnsw", "nearest-neighbor"]
|
|||||||
categories = ["algorithms", "science"]
|
categories = ["algorithms", "science"]
|
||||||
|
|
||||||
[dependencies]
|
[dependencies]
|
||||||
clawhdf5-format = { path = "../clawhdf5-format", version = "2.4.0" }
|
clawhdf5-format = { path = "../clawhdf5-format", version = "2.7.0" }
|
||||||
clawhdf5-io = { path = "../clawhdf5-io", version = "2.4.0" }
|
clawhdf5-io = { path = "../clawhdf5-io", version = "2.7.0" }
|
||||||
clawhdf5-accel = { path = "../clawhdf5-accel", version = "2.4.0" }
|
clawhdf5-accel = { path = "../clawhdf5-accel", version = "2.7.0" }
|
||||||
rayon = { version = "1", optional = true }
|
rayon = { version = "1", optional = true }
|
||||||
|
|
||||||
[features]
|
[features]
|
||||||
|
|||||||
+664
-127
File diff suppressed because it is too large
Load Diff
@@ -5,4 +5,4 @@
|
|||||||
|
|
||||||
mod hnsw;
|
mod hnsw;
|
||||||
|
|
||||||
pub use hnsw::{DistanceMetric, HnswIndex};
|
pub use hnsw::{DistanceMetric, HnswIndex, Storage};
|
||||||
|
|||||||
@@ -1,7 +1,8 @@
|
|||||||
[package]
|
[package]
|
||||||
name = "clawhdf5-bench"
|
name = "clawhdf5-bench"
|
||||||
version = "2.4.0"
|
version = "2.7.0"
|
||||||
edition = "2024"
|
edition = "2024"
|
||||||
|
rust-version.workspace = true
|
||||||
description = "Benchmark harnesses for clawhdf5-agent (Track 8)"
|
description = "Benchmark harnesses for clawhdf5-agent (Track 8)"
|
||||||
license = "MIT"
|
license = "MIT"
|
||||||
|
|
||||||
@@ -13,6 +14,10 @@ path = "src/bin/longmemeval_bench.rs"
|
|||||||
name = "memory_arena"
|
name = "memory_arena"
|
||||||
path = "src/bin/memory_arena.rs"
|
path = "src/bin/memory_arena.rs"
|
||||||
|
|
||||||
|
[[bin]]
|
||||||
|
name = "read_harness"
|
||||||
|
path = "src/bin/read_harness.rs"
|
||||||
|
|
||||||
[[bin]]
|
[[bin]]
|
||||||
name = "search_harness"
|
name = "search_harness"
|
||||||
path = "src/bin/search_harness.rs"
|
path = "src/bin/search_harness.rs"
|
||||||
@@ -53,6 +58,8 @@ harness = false
|
|||||||
[dependencies]
|
[dependencies]
|
||||||
clawhdf5-agent = { path = "../clawhdf5-agent" }
|
clawhdf5-agent = { path = "../clawhdf5-agent" }
|
||||||
clawhdf5-ann = { path = "../clawhdf5-ann" }
|
clawhdf5-ann = { path = "../clawhdf5-ann" }
|
||||||
|
clawhdf5 = { path = "../clawhdf5" }
|
||||||
|
clawhdf5-format = { path = "../clawhdf5-format" }
|
||||||
clawhdf5-io = { path = "../clawhdf5-io" }
|
clawhdf5-io = { path = "../clawhdf5-io" }
|
||||||
mpi = { version = "0.8", optional = true }
|
mpi = { version = "0.8", optional = true }
|
||||||
serde = { workspace = true }
|
serde = { workspace = true }
|
||||||
|
|||||||
@@ -396,7 +396,7 @@ fn run_memory_reduction_benchmark() {
|
|||||||
println!();
|
println!();
|
||||||
println!(
|
println!(
|
||||||
"{:>8} {:>10} {:>10} {:>10} {:>12}",
|
"{:>8} {:>10} {:>10} {:>10} {:>12}",
|
||||||
"Initial", "Remaining", "Eviction%", "Signal OK?", "BM25 Speedup"
|
"Initial", "Remaining", "Eviction%", "Signal OK?", "Records ÷"
|
||||||
);
|
);
|
||||||
println!("{}", "-".repeat(58));
|
println!("{}", "-".repeat(58));
|
||||||
|
|
||||||
@@ -440,7 +440,8 @@ fn run_memory_reduction_benchmark() {
|
|||||||
// Check all signal records survived
|
// Check all signal records survived
|
||||||
let signal_survived = signal_ids.iter().all(|&id| engine.get_by_id(id).is_some());
|
let signal_survived = signal_ids.iter().all(|&id| engine.get_by_id(id).is_some());
|
||||||
|
|
||||||
// Rough speedup: BM25 scales roughly linearly with record count
|
// How many times fewer records there are. Not a measured speedup —
|
||||||
|
// Part 1 measures search latency before and after.
|
||||||
let speedup = before_count as f64 / after_count.max(1) as f64;
|
let speedup = before_count as f64 / after_count.max(1) as f64;
|
||||||
|
|
||||||
println!(
|
println!(
|
||||||
@@ -480,7 +481,7 @@ fn main() {
|
|||||||
println!(" 3. Reducing search latency proportional to record reduction");
|
println!(" 3. Reducing search latency proportional to record reduction");
|
||||||
println!();
|
println!();
|
||||||
println!(
|
println!(
|
||||||
"Cycle time scales sub-linearly: 100 records ~microseconds, 100K records ~tens of ms."
|
"Cycle time grows a little faster than linearly: 100 records ~microseconds, 100K records ~tens of ms."
|
||||||
);
|
);
|
||||||
println!("Signal records with Correction source + high access_count survive eviction.");
|
println!("Signal records with Correction source + high access_count survive eviction.");
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -11,12 +11,14 @@
|
|||||||
//!
|
//!
|
||||||
//! Configuration matrix:
|
//! Configuration matrix:
|
||||||
//! - Text lengths: short (50 chars), medium (200 chars), long (1000 chars)
|
//! - Text lengths: short (50 chars), medium (200 chars), long (1000 chars)
|
||||||
//! - Embedding: 384-dim f32 (1536 bytes raw per record)
|
//! - Embedding: 384-dim, stored as float16 (the default for new stores) or
|
||||||
|
//! f32 with `--f32`; "raw" bytes are counted as f32 input either way
|
||||||
//! - WAL: enabled and disabled
|
//! - WAL: enabled and disabled
|
||||||
//!
|
//!
|
||||||
//! # Usage
|
//! # Usage
|
||||||
//! ```
|
//! ```
|
||||||
//! cargo run --release --bin footprint_bench
|
//! cargo run --release --bin footprint_bench # float16 stores
|
||||||
|
//! cargo run --release --bin footprint_bench -- --f32 # f32 stores
|
||||||
//! ```
|
//! ```
|
||||||
|
|
||||||
use std::time::Instant;
|
use std::time::Instant;
|
||||||
@@ -24,6 +26,9 @@ use std::time::Instant;
|
|||||||
use clawhdf5_agent::{AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry};
|
use clawhdf5_agent::{AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry};
|
||||||
use tempfile::TempDir;
|
use tempfile::TempDir;
|
||||||
|
|
||||||
|
/// `--f32`: build f32 stores instead of the library's float16 default.
|
||||||
|
static F32: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
|
||||||
|
|
||||||
const EMBEDDING_DIM: usize = 384;
|
const EMBEDDING_DIM: usize = 384;
|
||||||
|
|
||||||
// Raw bytes per record: 384 f32 embeddings + median text + overhead
|
// Raw bytes per record: 384 f32 embeddings + median text + overhead
|
||||||
@@ -152,6 +157,9 @@ fn measure_footprint(
|
|||||||
config.compression = compression;
|
config.compression = compression;
|
||||||
config.compression_level = if compression { 6 } else { 0 };
|
config.compression_level = if compression { 6 } else { 0 };
|
||||||
config.compact_threshold = 0.0;
|
config.compact_threshold = 0.0;
|
||||||
|
if F32.load(std::sync::atomic::Ordering::Relaxed) {
|
||||||
|
config.float16 = false;
|
||||||
|
}
|
||||||
|
|
||||||
let mut memory = HDF5Memory::create(config).expect("HDF5Memory::create failed");
|
let mut memory = HDF5Memory::create(config).expect("HDF5Memory::create failed");
|
||||||
|
|
||||||
@@ -241,11 +249,19 @@ fn fmt_n(n: usize) -> String {
|
|||||||
// ---------------------------------------------------------------------------
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
fn main() {
|
fn main() {
|
||||||
|
if std::env::args().skip(1).any(|a| a == "--f32") {
|
||||||
|
F32.store(true, std::sync::atomic::Ordering::Relaxed);
|
||||||
|
}
|
||||||
|
let stored = if F32.load(std::sync::atomic::Ordering::Relaxed) {
|
||||||
|
"f32 (1,536 bytes per record)"
|
||||||
|
} else {
|
||||||
|
"float16 (768 bytes per record; the default for new stores)"
|
||||||
|
};
|
||||||
println!("=================================================================");
|
println!("=================================================================");
|
||||||
println!(" ClawhDF5 Memory Footprint Benchmark");
|
println!(" ClawhDF5 Memory Footprint Benchmark");
|
||||||
println!("=================================================================");
|
println!("=================================================================");
|
||||||
println!();
|
println!();
|
||||||
println!("Embedding: 384-dim f32 = 1,536 bytes raw per record");
|
println!("Embedding: 384-dim, stored as {stored}; raw input counted as f32");
|
||||||
println!("Text lengths: short=50 chars, medium=200 chars, long=1000 chars");
|
println!("Text lengths: short=50 chars, medium=200 chars, long=1000 chars");
|
||||||
println!();
|
println!();
|
||||||
|
|
||||||
|
|||||||
@@ -55,44 +55,155 @@ use std::time::{Duration, Instant};
|
|||||||
#[path = "longmemeval_bench/embedder.rs"]
|
#[path = "longmemeval_bench/embedder.rs"]
|
||||||
mod embedder;
|
mod embedder;
|
||||||
|
|
||||||
use clawhdf5_agent::{AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry};
|
use clawhdf5_agent::bm25::TokenFilter;
|
||||||
|
use clawhdf5_agent::hybrid::Fusion;
|
||||||
|
use clawhdf5_agent::reranker::{ReRankConfig, RerankInput, rerank};
|
||||||
|
use clawhdf5_agent::{AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry, SearchResult};
|
||||||
use serde::Deserialize;
|
use serde::Deserialize;
|
||||||
use tempfile::TempDir;
|
use tempfile::TempDir;
|
||||||
|
|
||||||
const EMBEDDING_DIM: usize = 384;
|
const EMBEDDING_DIM: usize = 384;
|
||||||
|
|
||||||
|
/// `--float16`: build every per-question store with `MemoryConfig::float16`,
|
||||||
|
/// so embeddings are rounded to half precision as they are saved — exactly
|
||||||
|
/// what such a store searches over.
|
||||||
|
static FLOAT16: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
|
||||||
|
|
||||||
|
/// A mode's fusion, as one short string for the reports.
|
||||||
|
fn describe(mode: Mode) -> String {
|
||||||
|
let fusion = match mode.fusion {
|
||||||
|
Fusion::Weighted { vector, keyword } => format!("vector_{vector:.1}_keyword_{keyword:.1}"),
|
||||||
|
Fusion::Rrf { k } => format!("rrf_k{k:.0}"),
|
||||||
|
};
|
||||||
|
let tokens = match mode.tokens {
|
||||||
|
TokenFilter::Plain => fusion,
|
||||||
|
TokenFilter::Stemmed => format!("{fusion}_stemmed"),
|
||||||
|
};
|
||||||
|
match mode.rerank {
|
||||||
|
None => tokens,
|
||||||
|
Some(cfg) if cfg.relevance_weight == 0.0 => format!("{tokens}_rerank_metadata"),
|
||||||
|
Some(cfg) => format!(
|
||||||
|
"{tokens}_rerank_blended_hl{:.0}d",
|
||||||
|
cfg.temporal_half_life_secs / 86_400.0
|
||||||
|
),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
/// A retrieval configuration: how much of the score comes from each stage.
|
/// A retrieval configuration: how much of the score comes from each stage.
|
||||||
#[derive(Clone, Copy)]
|
#[derive(Clone, Copy)]
|
||||||
struct Mode {
|
struct Mode {
|
||||||
label: &'static str,
|
label: &'static str,
|
||||||
vector_weight: f32,
|
/// How the two retrieval stages are combined into one ranking.
|
||||||
keyword_weight: f32,
|
fusion: Fusion,
|
||||||
|
/// How keyword tokens are normalised before indexing and querying.
|
||||||
|
tokens: TokenFilter,
|
||||||
|
/// Re-rank the retrieved candidates with recency and friends, relative to
|
||||||
|
/// the question's own date.
|
||||||
|
rerank: Option<ReRankConfig>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Mode {
|
||||||
|
const fn weighted(label: &'static str, vector: f32, keyword: f32) -> Self {
|
||||||
|
Self {
|
||||||
|
label,
|
||||||
|
fusion: Fusion::Weighted { vector, keyword },
|
||||||
|
tokens: TokenFilter::Plain,
|
||||||
|
rerank: None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg_attr(not(feature = "embeddings"), allow(dead_code))]
|
||||||
|
fn reranked(mut self, label: &'static str, rerank: ReRankConfig) -> Self {
|
||||||
|
self.label = label;
|
||||||
|
self.rerank = Some(rerank);
|
||||||
|
self
|
||||||
|
}
|
||||||
|
|
||||||
|
const fn stemmed(mut self, label: &'static str) -> Self {
|
||||||
|
self.label = label;
|
||||||
|
self.tokens = TokenFilter::Stemmed;
|
||||||
|
self
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// The only mode available without real embeddings. Passing zero vectors with
|
/// The only mode available without real embeddings. Passing zero vectors with
|
||||||
/// `vector_weight = 0.0` is what made the vector stage inert.
|
/// `vector_weight = 0.0` is what made the vector stage inert.
|
||||||
const BM25_ONLY: Mode = Mode {
|
const BM25_ONLY: Mode = Mode::weighted("BM25 only (vector stage inert)", 0.0, 1.0);
|
||||||
label: "BM25 only (vector stage inert)",
|
|
||||||
vector_weight: 0.0,
|
|
||||||
keyword_weight: 1.0,
|
|
||||||
};
|
|
||||||
#[cfg(feature = "embeddings")]
|
#[cfg(feature = "embeddings")]
|
||||||
const VECTOR_ONLY: Mode = Mode {
|
const VECTOR_ONLY: Mode = Mode::weighted("Vector only (MiniLM + HNSW)", 1.0, 0.0);
|
||||||
label: "Vector only (MiniLM + HNSW)",
|
|
||||||
vector_weight: 1.0,
|
|
||||||
keyword_weight: 0.0,
|
|
||||||
};
|
|
||||||
/// Tuned by `--sweep` over the full haystack. The former 0.7/0.3 was a
|
/// Tuned by `--sweep` over the full haystack. The former 0.7/0.3 was a
|
||||||
/// documented default that had never been searched, and the sweep found it
|
/// documented default that had never been searched, and the sweep found it
|
||||||
/// strictly dominated: 0.4/0.6 is better on Hit@1, Hit@5, Hit@10 and MRR at
|
/// strictly dominated: 0.4/0.6 is better on Hit@1, Hit@5, Hit@10 and MRR at
|
||||||
/// both granularities.
|
/// both granularities.
|
||||||
#[cfg(feature = "embeddings")]
|
#[cfg(feature = "embeddings")]
|
||||||
const HYBRID: Mode = Mode {
|
const HYBRID: Mode = Mode::weighted("Hybrid (0.4 vector / 0.6 BM25, tuned)", 0.4, 0.6);
|
||||||
label: "Hybrid (0.4 vector / 0.6 BM25, tuned)",
|
|
||||||
vector_weight: 0.4,
|
/// Reciprocal rank fusion, the documented alternative to the weighted sum.
|
||||||
keyword_weight: 0.6,
|
/// It ignores score magnitudes, so there is nothing to tune — which is the
|
||||||
|
/// claim being tested.
|
||||||
|
#[cfg(feature = "embeddings")]
|
||||||
|
const RRF: Mode = Mode {
|
||||||
|
label: "Hybrid (reciprocal rank fusion, k=60)",
|
||||||
|
fusion: Fusion::Rrf { k: 60.0 },
|
||||||
|
tokens: TokenFilter::Plain,
|
||||||
|
rerank: None,
|
||||||
};
|
};
|
||||||
|
|
||||||
|
/// The same two configurations with stemmed keyword tokens, so the tokenizer's
|
||||||
|
/// effect is isolated from everything else.
|
||||||
|
const BM25_STEMMED: Mode = BM25_ONLY.stemmed("BM25 only, stemmed tokens");
|
||||||
|
|
||||||
|
/// Re-ranking as it behaved before `relevance` was an input: the combined
|
||||||
|
/// score was recency + authority + activation only, so the retriever's own
|
||||||
|
/// ordering was discarded.
|
||||||
|
#[cfg(feature = "embeddings")]
|
||||||
|
fn hybrid_rerank_metadata_only() -> Mode {
|
||||||
|
HYBRID.reranked(
|
||||||
|
"Hybrid + rerank (metadata only, pre-fix)",
|
||||||
|
ReRankConfig {
|
||||||
|
relevance_weight: 0.0,
|
||||||
|
..ReRankConfig::default()
|
||||||
|
},
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Re-ranking as it behaves now: relevance leads, recency nudges.
|
||||||
|
#[cfg(feature = "embeddings")]
|
||||||
|
fn hybrid_rerank_blended() -> Mode {
|
||||||
|
HYBRID.reranked(
|
||||||
|
"Hybrid + rerank (relevance + recency)",
|
||||||
|
ReRankConfig::default(),
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// The same blend at several half-lives. Decay is `2^(-age / half_life)`, so a
|
||||||
|
/// half-life far shorter than the gaps between memories sends every score to
|
||||||
|
/// zero and the signal vanishes; far longer and everything scores ~1 and it
|
||||||
|
/// vanishes the other way. The right value tracks how far apart the memories
|
||||||
|
/// actually are.
|
||||||
|
#[cfg(feature = "embeddings")]
|
||||||
|
fn hybrid_rerank_half_lives() -> Vec<Mode> {
|
||||||
|
[
|
||||||
|
("1 day", 86_400.0),
|
||||||
|
("7 days", 7.0 * 86_400.0),
|
||||||
|
("30 days", 30.0 * 86_400.0),
|
||||||
|
("90 days", 90.0 * 86_400.0),
|
||||||
|
]
|
||||||
|
.into_iter()
|
||||||
|
.map(|(label, half_life)| {
|
||||||
|
HYBRID.reranked(
|
||||||
|
Box::leak(format!("Hybrid + rerank, half-life {label}").into_boxed_str()),
|
||||||
|
ReRankConfig {
|
||||||
|
temporal_half_life_secs: half_life,
|
||||||
|
..ReRankConfig::default()
|
||||||
|
},
|
||||||
|
)
|
||||||
|
})
|
||||||
|
.collect()
|
||||||
|
}
|
||||||
|
#[cfg(feature = "embeddings")]
|
||||||
|
const HYBRID_STEMMED: Mode = HYBRID.stemmed("Hybrid 0.4/0.6, stemmed tokens");
|
||||||
|
|
||||||
/// Every 0.1 step of vector weight, keyword weight taking the remainder.
|
/// Every 0.1 step of vector weight, keyword weight taking the remainder.
|
||||||
///
|
///
|
||||||
/// Labels are leaked to `&'static str` because `Mode::label` is a `&'static
|
/// Labels are leaked to `&'static str` because `Mode::label` is a `&'static
|
||||||
@@ -104,11 +215,11 @@ fn sweep_modes() -> Vec<Mode> {
|
|||||||
(0..=10)
|
(0..=10)
|
||||||
.map(|i| {
|
.map(|i| {
|
||||||
let v = i as f32 / 10.0;
|
let v = i as f32 / 10.0;
|
||||||
Mode {
|
Mode::weighted(
|
||||||
label: Box::leak(format!("sweep v={v:.1} / k={:.1}", 1.0 - v).into_boxed_str()),
|
Box::leak(format!("sweep v={v:.1} / k={:.1}", 1.0 - v).into_boxed_str()),
|
||||||
vector_weight: v,
|
v,
|
||||||
keyword_weight: 1.0 - v,
|
1.0 - v,
|
||||||
}
|
)
|
||||||
})
|
})
|
||||||
.collect()
|
.collect()
|
||||||
}
|
}
|
||||||
@@ -181,6 +292,37 @@ struct Question {
|
|||||||
haystack_session_ids: Vec<String>,
|
haystack_session_ids: Vec<String>,
|
||||||
haystack_sessions: Vec<Vec<Turn>>,
|
haystack_sessions: Vec<Vec<Turn>>,
|
||||||
answer_session_ids: Vec<String>,
|
answer_session_ids: Vec<String>,
|
||||||
|
/// One timestamp per haystack session, e.g. "2023/05/25 (Thu) 20:21".
|
||||||
|
#[serde(default)]
|
||||||
|
haystack_dates: Vec<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Seconds since the epoch for a LongMemEval session date, which looks like
|
||||||
|
/// `2023/05/25 (Thu) 20:21`. Sessions are stored in chronological order, so a
|
||||||
|
/// date that cannot be parsed falls back to its position — order is preserved
|
||||||
|
/// even if the interval is not.
|
||||||
|
fn session_time(date: &str, position: usize) -> f64 {
|
||||||
|
let stamp = |y: i64, mo: i64, d: i64, h: i64, mi: i64| -> f64 {
|
||||||
|
// Days since 1970-01-01 via the civil-from-days algorithm.
|
||||||
|
let (y, mo) = if mo <= 2 { (y - 1, mo + 12) } else { (y, mo) };
|
||||||
|
let era = y.div_euclid(400);
|
||||||
|
let yoe = y - era * 400;
|
||||||
|
let doy = (153 * (mo - 3) + 2) / 5 + d - 1;
|
||||||
|
let doe = yoe * 365 + yoe / 4 - yoe / 100 + doy;
|
||||||
|
let days = era * 146_097 + doe - 719_468;
|
||||||
|
(days * 86_400 + h * 3_600 + mi * 60) as f64
|
||||||
|
};
|
||||||
|
let parse = || -> Option<f64> {
|
||||||
|
let (ymd, rest) = date.split_once(' ')?;
|
||||||
|
let mut ymd = ymd.split('/');
|
||||||
|
let y = ymd.next()?.parse().ok()?;
|
||||||
|
let mo = ymd.next()?.parse().ok()?;
|
||||||
|
let d = ymd.next()?.parse().ok()?;
|
||||||
|
let hm = rest.rsplit(' ').next()?;
|
||||||
|
let (h, mi) = hm.split_once(':')?;
|
||||||
|
Some(stamp(y, mo, d, h.parse().ok()?, mi.parse().ok()?))
|
||||||
|
};
|
||||||
|
parse().unwrap_or(1_000_000.0 + position as f64 * 86_400.0)
|
||||||
}
|
}
|
||||||
|
|
||||||
// ---------------------------------------------------------------------------
|
// ---------------------------------------------------------------------------
|
||||||
@@ -199,11 +341,21 @@ struct Metrics {
|
|||||||
rr_turn: f64,
|
rr_turn: f64,
|
||||||
abstention_correct: u32,
|
abstention_correct: u32,
|
||||||
abstention_total: u32,
|
abstention_total: u32,
|
||||||
|
/// Questions where the newest gold session outranked the older ones, out
|
||||||
|
/// of those with more than one gold session and at least one retrieved.
|
||||||
|
newest_gold_first: u32,
|
||||||
|
newest_gold_total: u32,
|
||||||
latency_ns: Vec<u64>,
|
latency_ns: Vec<u64>,
|
||||||
count: u32,
|
count: u32,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl Metrics {
|
impl Metrics {
|
||||||
|
/// `None` when no question in this bucket had multiple gold sessions.
|
||||||
|
fn newest_gold_first_pct(&self) -> Option<f64> {
|
||||||
|
(self.newest_gold_total > 0)
|
||||||
|
.then(|| self.newest_gold_first as f64 / self.newest_gold_total as f64 * 100.0)
|
||||||
|
}
|
||||||
|
|
||||||
fn hit1_session_pct(&self) -> f64 {
|
fn hit1_session_pct(&self) -> f64 {
|
||||||
self.hit1_session as f64 / self.count.max(1) as f64 * 100.0
|
self.hit1_session as f64 / self.count.max(1) as f64 * 100.0
|
||||||
}
|
}
|
||||||
@@ -261,6 +413,16 @@ struct EvalResult {
|
|||||||
hit5_turn: bool,
|
hit5_turn: bool,
|
||||||
hit10_turn: bool,
|
hit10_turn: bool,
|
||||||
rr_turn: Option<f64>,
|
rr_turn: Option<f64>,
|
||||||
|
/// For a question whose evidence spans several dated sessions (a
|
||||||
|
/// `knowledge-update`, where an earlier fact is superseded by a later
|
||||||
|
/// one): did the *newest* gold session outrank every older gold session
|
||||||
|
/// that was returned? `None` when the question has one gold session, or
|
||||||
|
/// when none were retrieved, so there is nothing to discriminate.
|
||||||
|
///
|
||||||
|
/// Plain recall cannot see this. LongMemEval labels *both* the stale and
|
||||||
|
/// the updated session as gold, so returning either counts as a hit — yet
|
||||||
|
/// only one of them answers the question correctly.
|
||||||
|
newest_gold_first: Option<bool>,
|
||||||
latency: Duration,
|
latency: Duration,
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -274,21 +436,29 @@ fn evaluate_question(
|
|||||||
let mut config = MemoryConfig::new(dir.path().join("lme.h5"), "lme-bench", EMBEDDING_DIM);
|
let mut config = MemoryConfig::new(dir.path().join("lme.h5"), "lme-bench", EMBEDDING_DIM);
|
||||||
config.wal_enabled = false;
|
config.wal_enabled = false;
|
||||||
config.compact_threshold = 0.0;
|
config.compact_threshold = 0.0;
|
||||||
|
config.float16 = FLOAT16.load(std::sync::atomic::Ordering::Relaxed);
|
||||||
|
|
||||||
let mut memory = HDF5Memory::create(config).expect("failed to create HDF5Memory");
|
let mut memory = HDF5Memory::create(config).expect("failed to create HDF5Memory");
|
||||||
|
memory.set_token_filter(mode.tokens);
|
||||||
|
|
||||||
// Build MemoryEntry list from all haystack sessions
|
// Build MemoryEntry list from all haystack sessions
|
||||||
let mut entries: Vec<MemoryEntry> = Vec::new();
|
let mut entries: Vec<MemoryEntry> = Vec::new();
|
||||||
let mut turn_has_answer: Vec<bool> = Vec::new();
|
let mut turn_has_answer: Vec<bool> = Vec::new();
|
||||||
let mut ts = 1_000_000.0f64;
|
|
||||||
|
|
||||||
for (sess_idx, session) in q.haystack_sessions.iter().enumerate() {
|
for (sess_idx, session) in q.haystack_sessions.iter().enumerate() {
|
||||||
let sess_id = q
|
let sess_id = q
|
||||||
.haystack_session_ids
|
.haystack_session_ids
|
||||||
.get(sess_idx)
|
.get(sess_idx)
|
||||||
.map(String::as_str)
|
.map(String::as_str)
|
||||||
.unwrap_or("unknown");
|
.unwrap_or("unknown");
|
||||||
for turn in session {
|
// Real session dates, not a synthetic counter: anything that decays
|
||||||
|
// with age needs true intervals, not just the right order.
|
||||||
|
let session_start = q
|
||||||
|
.haystack_dates
|
||||||
|
.get(sess_idx)
|
||||||
|
.map_or(sess_idx as f64 * 86_400.0, |d| session_time(d, sess_idx));
|
||||||
|
for (turn_idx, turn) in session.iter().enumerate() {
|
||||||
|
// Spread a session's turns over the minutes following its start.
|
||||||
|
let ts = session_start + turn_idx as f64 * 60.0;
|
||||||
entries.push(MemoryEntry {
|
entries.push(MemoryEntry {
|
||||||
chunk: turn.content.clone(),
|
chunk: turn.content.clone(),
|
||||||
embedding: embedding_for(embeddings, &turn.content),
|
embedding: embedding_for(embeddings, &turn.content),
|
||||||
@@ -302,7 +472,6 @@ fn evaluate_question(
|
|||||||
},
|
},
|
||||||
});
|
});
|
||||||
turn_has_answer.push(turn.has_answer);
|
turn_has_answer.push(turn.has_answer);
|
||||||
ts += 1.0;
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -319,17 +488,87 @@ fn evaluate_question(
|
|||||||
// Set of session IDs that contain the answer
|
// Set of session IDs that contain the answer
|
||||||
let answer_sess_set: HashSet<&str> = q.answer_session_ids.iter().map(String::as_str).collect();
|
let answer_sess_set: HashSet<&str> = q.answer_session_ids.iter().map(String::as_str).collect();
|
||||||
|
|
||||||
|
// When each gold session was recorded, so "newest" is by date rather than
|
||||||
|
// by position (the two agree in this dataset, but the metric should not
|
||||||
|
// depend on that).
|
||||||
|
let gold_times: HashMap<&str, f64> = q
|
||||||
|
.haystack_session_ids
|
||||||
|
.iter()
|
||||||
|
.enumerate()
|
||||||
|
.filter(|(_, sid)| answer_sess_set.contains(sid.as_str()))
|
||||||
|
.map(|(i, sid)| {
|
||||||
|
let t = q
|
||||||
|
.haystack_dates
|
||||||
|
.get(i)
|
||||||
|
.map_or(i as f64 * 86_400.0, |d| session_time(d, i));
|
||||||
|
(sid.as_str(), t)
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
|
||||||
let query_emb = embedding_for(embeddings, &q.question);
|
let query_emb = embedding_for(embeddings, &q.question);
|
||||||
let t0 = Instant::now();
|
let t0 = Instant::now();
|
||||||
let results = memory.hybrid_search(
|
// Re-ranking only reorders; it needs a candidate pool larger than `top_k`
|
||||||
&query_emb,
|
// to have anything to promote.
|
||||||
&q.question,
|
let pool = if mode.rerank.is_some() {
|
||||||
mode.vector_weight,
|
top_k * 4
|
||||||
mode.keyword_weight,
|
} else {
|
||||||
top_k,
|
top_k
|
||||||
);
|
};
|
||||||
|
let mut results = memory.hybrid_search_with(&query_emb, &q.question, mode.fusion, pool);
|
||||||
|
if let Some(config) = mode.rerank {
|
||||||
|
// "Now" is the moment the question was asked, so decay measures how
|
||||||
|
// stale each memory was at that point.
|
||||||
|
let now = session_time(&q.question_date, q.haystack_sessions.len());
|
||||||
|
let inputs: Vec<RerankInput> = results
|
||||||
|
.iter()
|
||||||
|
.map(|r| RerankInput {
|
||||||
|
index: r.index,
|
||||||
|
timestamp: r.timestamp,
|
||||||
|
source_channel: r.source_channel.clone(),
|
||||||
|
raw_activation: r.activation,
|
||||||
|
relevance: r.score,
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
let order: Vec<usize> = rerank(&inputs, &config, now)
|
||||||
|
.into_iter()
|
||||||
|
.map(|r| r.index)
|
||||||
|
.collect();
|
||||||
|
let by_index: HashMap<usize, SearchResult> =
|
||||||
|
results.into_iter().map(|r| (r.index, r)).collect();
|
||||||
|
results = order
|
||||||
|
.into_iter()
|
||||||
|
.filter_map(|i| by_index.get(&i).cloned())
|
||||||
|
.collect();
|
||||||
|
}
|
||||||
|
results.truncate(top_k);
|
||||||
let latency = t0.elapsed();
|
let latency = t0.elapsed();
|
||||||
|
|
||||||
|
// Rank of the best-placed result from each gold session.
|
||||||
|
let mut first_rank: HashMap<&str, usize> = HashMap::new();
|
||||||
|
for (rank, result) in results.iter().enumerate() {
|
||||||
|
let sid = memory.cache.session_ids[result.index].as_str();
|
||||||
|
if let Some((gold_sid, _)) = gold_times.get_key_value(sid) {
|
||||||
|
first_rank.entry(gold_sid).or_insert(rank);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
let newest_gold_first = if gold_times.len() < 2 || first_rank.is_empty() {
|
||||||
|
None
|
||||||
|
} else {
|
||||||
|
// The newest gold session must be retrieved, and no older gold session
|
||||||
|
// may outrank it.
|
||||||
|
let newest = gold_times
|
||||||
|
.iter()
|
||||||
|
.max_by(|a, b| a.1.total_cmp(b.1))
|
||||||
|
.map(|(sid, _)| *sid)
|
||||||
|
.expect("at least two gold sessions");
|
||||||
|
Some(match first_rank.get(newest) {
|
||||||
|
Some(&newest_rank) => first_rank
|
||||||
|
.iter()
|
||||||
|
.all(|(sid, &rank)| *sid == newest || rank > newest_rank),
|
||||||
|
None => false,
|
||||||
|
})
|
||||||
|
};
|
||||||
|
|
||||||
// Session-level recall
|
// Session-level recall
|
||||||
let mut hit1_session = false;
|
let mut hit1_session = false;
|
||||||
let mut hit5_session = false;
|
let mut hit5_session = false;
|
||||||
@@ -384,6 +623,7 @@ fn evaluate_question(
|
|||||||
hit5_turn,
|
hit5_turn,
|
||||||
hit10_turn,
|
hit10_turn,
|
||||||
rr_turn,
|
rr_turn,
|
||||||
|
newest_gold_first,
|
||||||
latency,
|
latency,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -472,10 +712,7 @@ fn print_report(
|
|||||||
println!(" LongMemEval Benchmark — {}", mode.label);
|
println!(" LongMemEval Benchmark — {}", mode.label);
|
||||||
println!("=================================================================");
|
println!("=================================================================");
|
||||||
println!();
|
println!();
|
||||||
println!(
|
println!("Mode: {}", describe(mode));
|
||||||
"Mode: vector_weight={:.1} / keyword_weight={:.1}",
|
|
||||||
mode.vector_weight, mode.keyword_weight
|
|
||||||
);
|
|
||||||
println!();
|
println!();
|
||||||
println!("Scoring target: RETRIEVAL RECALL (did the gold memory land in top-k).");
|
println!("Scoring target: RETRIEVAL RECALL (did the gold memory land in top-k).");
|
||||||
println!(" No answer is generated or scored. This is NOT the official");
|
println!(" No answer is generated or scored. This is NOT the official");
|
||||||
@@ -538,6 +775,24 @@ fn print_report(
|
|||||||
);
|
);
|
||||||
println!();
|
println!();
|
||||||
|
|
||||||
|
if let Some(pct) = overall.newest_gold_first_pct() {
|
||||||
|
println!(
|
||||||
|
"## Recency Discrimination (n={})",
|
||||||
|
overall.newest_gold_total
|
||||||
|
);
|
||||||
|
println!(
|
||||||
|
" Newest gold session ranked first: {}/{} ({pct:.1}%)",
|
||||||
|
overall.newest_gold_first, overall.newest_gold_total
|
||||||
|
);
|
||||||
|
println!(
|
||||||
|
" Questions whose evidence spans several dated sessions — a fact and\n \
|
||||||
|
its later correction. Both sessions are labelled gold, so recall\n \
|
||||||
|
scores either as a hit; this asks whether the *current* one came\n \
|
||||||
|
first. A retriever with no sense of time scores near chance."
|
||||||
|
);
|
||||||
|
println!();
|
||||||
|
}
|
||||||
|
|
||||||
if overall.abstention_total > 0 {
|
if overall.abstention_total > 0 {
|
||||||
println!("## Abstention Accuracy");
|
println!("## Abstention Accuracy");
|
||||||
println!(
|
println!(
|
||||||
@@ -602,10 +857,7 @@ fn print_report(
|
|||||||
println!("```json");
|
println!("```json");
|
||||||
println!("{{");
|
println!("{{");
|
||||||
println!(" \"benchmark\": \"longmemeval\",");
|
println!(" \"benchmark\": \"longmemeval\",");
|
||||||
println!(
|
println!(" \"mode\": \"{}\",", describe(mode));
|
||||||
" \"mode\": \"vector_{:.1}_keyword_{:.1}\",",
|
|
||||||
mode.vector_weight, mode.keyword_weight
|
|
||||||
);
|
|
||||||
println!(" \"dataset_variant\": \"{}\",", profile.variant());
|
println!(" \"dataset_variant\": \"{}\",", profile.variant());
|
||||||
println!(" \"scoring_target\": \"retrieval_recall\",");
|
println!(" \"scoring_target\": \"retrieval_recall\",");
|
||||||
println!(" \"k\": 10,");
|
println!(" \"k\": 10,");
|
||||||
@@ -654,6 +906,14 @@ fn print_report(
|
|||||||
} else {
|
} else {
|
||||||
println!(" \"abstention_accuracy\": null,");
|
println!(" \"abstention_accuracy\": null,");
|
||||||
}
|
}
|
||||||
|
match overall.newest_gold_first_pct() {
|
||||||
|
Some(pct) => println!(
|
||||||
|
" \"newest_gold_first\": {:.4}, \"newest_gold_n\": {},",
|
||||||
|
pct / 100.0,
|
||||||
|
overall.newest_gold_total
|
||||||
|
),
|
||||||
|
None => println!(" \"newest_gold_first\": null,"),
|
||||||
|
}
|
||||||
println!(" \"latency_us\": {{");
|
println!(" \"latency_us\": {{");
|
||||||
println!(
|
println!(
|
||||||
" \"avg\": {:.1}, \"p50\": {:.1}, \"p95\": {:.1}, \"p99\": {:.1}",
|
" \"avg\": {:.1}, \"p50\": {:.1}, \"p95\": {:.1}, \"p99\": {:.1}",
|
||||||
@@ -676,6 +936,8 @@ fn main() {
|
|||||||
let mut limit: Option<usize> = None;
|
let mut limit: Option<usize> = None;
|
||||||
let mut weights_dir: Option<String> = None;
|
let mut weights_dir: Option<String> = None;
|
||||||
let mut sweep = false;
|
let mut sweep = false;
|
||||||
|
#[cfg_attr(not(feature = "embeddings"), allow(unused_mut, unused_variables))]
|
||||||
|
let mut rerank_sweep = false;
|
||||||
let mut args = std::env::args().skip(1);
|
let mut args = std::env::args().skip(1);
|
||||||
while let Some(arg) = args.next() {
|
while let Some(arg) = args.next() {
|
||||||
match arg.as_str() {
|
match arg.as_str() {
|
||||||
@@ -684,6 +946,20 @@ fn main() {
|
|||||||
limit = Some(v.parse().expect("--limit must be a positive integer"));
|
limit = Some(v.parse().expect("--limit must be a positive integer"));
|
||||||
}
|
}
|
||||||
"--sweep" => sweep = true,
|
"--sweep" => sweep = true,
|
||||||
|
"--float16" => {
|
||||||
|
FLOAT16.store(true, std::sync::atomic::Ordering::Relaxed);
|
||||||
|
eprintln!("Stores use MemoryConfig::float16 (half-precision embeddings)");
|
||||||
|
}
|
||||||
|
"--rerank-sweep" => {
|
||||||
|
// Re-ranking needs the vector stage to have candidates worth
|
||||||
|
// reordering, so this is an embeddings-only comparison.
|
||||||
|
#[cfg(feature = "embeddings")]
|
||||||
|
{
|
||||||
|
rerank_sweep = true;
|
||||||
|
}
|
||||||
|
#[cfg(not(feature = "embeddings"))]
|
||||||
|
eprintln!("warning: --rerank-sweep needs --features embeddings; ignoring");
|
||||||
|
}
|
||||||
"--embeddings" => {
|
"--embeddings" => {
|
||||||
weights_dir = Some(args.next().expect("--embeddings needs a directory"));
|
weights_dir = Some(args.next().expect("--embeddings needs a directory"));
|
||||||
}
|
}
|
||||||
@@ -702,6 +978,12 @@ fn main() {
|
|||||||
BM25-only, vector-only, and hybrid separately. Requires\n\
|
BM25-only, vector-only, and hybrid separately. Requires\n\
|
||||||
--features embeddings; without it the vector stage is\n\
|
--features embeddings; without it the vector stage is\n\
|
||||||
inert and only the BM25 row is produced.\n\
|
inert and only the BM25 row is produced.\n\
|
||||||
|
--rerank-sweep\n\
|
||||||
|
compare re-ranking off, metadata-only (the old\n\
|
||||||
|
behaviour) and blended at several half-lives.\n\
|
||||||
|
--float16\n\
|
||||||
|
build each store with MemoryConfig::float16, to\n\
|
||||||
|
compare retrieval on half-precision embeddings.\n\
|
||||||
--sweep instead of the three named modes, sweep vector_weight\n\
|
--sweep instead of the three named modes, sweep vector_weight\n\
|
||||||
from 0.0 to 1.0 in 0.1 steps. The 0.7/0.3 default was\n\
|
from 0.0 to 1.0 in 0.1 steps. The 0.7/0.3 default was\n\
|
||||||
never searched; this is what searches it."
|
never searched; this is what searches it."
|
||||||
@@ -767,19 +1049,34 @@ fn main() {
|
|||||||
{
|
{
|
||||||
if sweep {
|
if sweep {
|
||||||
sweep_modes()
|
sweep_modes()
|
||||||
|
} else if rerank_sweep {
|
||||||
|
let mut modes = vec![HYBRID, hybrid_rerank_metadata_only()];
|
||||||
|
modes.extend(hybrid_rerank_half_lives());
|
||||||
|
modes
|
||||||
} else {
|
} else {
|
||||||
vec![BM25_ONLY, VECTOR_ONLY, HYBRID]
|
vec![
|
||||||
|
BM25_ONLY,
|
||||||
|
VECTOR_ONLY,
|
||||||
|
HYBRID,
|
||||||
|
RRF,
|
||||||
|
BM25_STEMMED,
|
||||||
|
HYBRID_STEMMED,
|
||||||
|
hybrid_rerank_metadata_only(),
|
||||||
|
hybrid_rerank_blended(),
|
||||||
|
]
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
#[cfg(not(feature = "embeddings"))]
|
#[cfg(not(feature = "embeddings"))]
|
||||||
{
|
{
|
||||||
vec![BM25_ONLY]
|
vec![BM25_ONLY, BM25_STEMMED]
|
||||||
}
|
}
|
||||||
} else {
|
} else {
|
||||||
if sweep {
|
if sweep {
|
||||||
eprintln!("warning: --sweep needs --embeddings; running BM25 only");
|
eprintln!("warning: --sweep needs --embeddings; running BM25 only");
|
||||||
}
|
}
|
||||||
vec![BM25_ONLY]
|
// Stemming is a property of the keyword stage, so it can be compared
|
||||||
|
// without a model.
|
||||||
|
vec![BM25_ONLY, BM25_STEMMED]
|
||||||
};
|
};
|
||||||
|
|
||||||
for (mode_idx, mode) in modes.iter().enumerate() {
|
for (mode_idx, mode) in modes.iter().enumerate() {
|
||||||
@@ -882,6 +1179,14 @@ fn run_mode(
|
|||||||
entry.rr_turn += rr;
|
entry.rr_turn += rr;
|
||||||
overall.rr_turn += rr;
|
overall.rr_turn += rr;
|
||||||
}
|
}
|
||||||
|
if let Some(newest_first) = result.newest_gold_first {
|
||||||
|
entry.newest_gold_total += 1;
|
||||||
|
overall.newest_gold_total += 1;
|
||||||
|
if newest_first {
|
||||||
|
entry.newest_gold_first += 1;
|
||||||
|
overall.newest_gold_first += 1;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
let ns = result.latency.as_nanos() as u64;
|
let ns = result.latency.as_nanos() as u64;
|
||||||
entry.latency_ns.push(ns);
|
entry.latency_ns.push(ns);
|
||||||
@@ -893,3 +1198,30 @@ fn run_mode(
|
|||||||
eprintln!();
|
eprintln!();
|
||||||
print_report(&overall, &by_type, profile, mode);
|
print_report(&overall, &by_type, profile, mode);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::session_time;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn session_dates_parse_to_the_right_instant() {
|
||||||
|
// Reference values from Python's datetime, UTC.
|
||||||
|
for (date, expected) in [
|
||||||
|
("2023/05/25 (Thu) 20:21", 1_685_046_060.0),
|
||||||
|
("1970/01/01 (Thu) 00:00", 0.0),
|
||||||
|
("2000/02/29 (Tue) 12:00", 951_825_600.0),
|
||||||
|
("2023/12/31 (Sun) 23:59", 1_704_067_140.0),
|
||||||
|
("2024/03/01 (Fri) 00:00", 1_709_251_200.0),
|
||||||
|
] {
|
||||||
|
assert_eq!(session_time(date, 0), expected, "{date}");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn unparseable_dates_fall_back_to_position_order() {
|
||||||
|
let a = session_time("not a date", 0);
|
||||||
|
let b = session_time("", 1);
|
||||||
|
let c = session_time("2023/13/99 (???) 99:99", 2);
|
||||||
|
assert!(a < b && b < c, "fallback must preserve session order");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -0,0 +1,176 @@
|
|||||||
|
//! HDF5 read-path measurement harness: full reads vs. hyperslab selections on
|
||||||
|
//! a chunked 2-D dataset, compressed and uncompressed, plus a contiguous one.
|
||||||
|
//!
|
||||||
|
//! The question it answers for every read-path change: does the cost of a
|
||||||
|
//! selection scale with the *selection*, or with the whole dataset?
|
||||||
|
//!
|
||||||
|
//! ```text
|
||||||
|
//! cargo run --release -p clawhdf5-bench --bin read_harness
|
||||||
|
//! cargo run --release -p clawhdf5-bench --bin read_harness -- --large # 512 MB
|
||||||
|
//! ```
|
||||||
|
|
||||||
|
use std::time::{Duration, Instant};
|
||||||
|
|
||||||
|
use clawhdf5::{File, FileBuilder};
|
||||||
|
use clawhdf5_format::selection::Selection;
|
||||||
|
|
||||||
|
const CHUNK: u64 = 256;
|
||||||
|
|
||||||
|
struct Layout {
|
||||||
|
name: &'static str,
|
||||||
|
chunked: bool,
|
||||||
|
deflate: bool,
|
||||||
|
}
|
||||||
|
|
||||||
|
const LAYOUTS: [Layout; 3] = [
|
||||||
|
Layout {
|
||||||
|
name: "chunked + deflate",
|
||||||
|
chunked: true,
|
||||||
|
deflate: true,
|
||||||
|
},
|
||||||
|
Layout {
|
||||||
|
name: "chunked",
|
||||||
|
chunked: true,
|
||||||
|
deflate: false,
|
||||||
|
},
|
||||||
|
Layout {
|
||||||
|
name: "contiguous",
|
||||||
|
chunked: false,
|
||||||
|
deflate: false,
|
||||||
|
},
|
||||||
|
];
|
||||||
|
|
||||||
|
/// Smooth-ish, compressible data whose value encodes its position, so a read
|
||||||
|
/// can be verified exactly.
|
||||||
|
fn value(row: u64, col: u64) -> f64 {
|
||||||
|
(row * 100_003 + col) as f64 * 0.5
|
||||||
|
}
|
||||||
|
|
||||||
|
fn write_file(path: &std::path::Path, rows: u64, cols: u64) {
|
||||||
|
let data: Vec<f64> = (0..rows)
|
||||||
|
.flat_map(|r| (0..cols).map(move |c| value(r, c)))
|
||||||
|
.collect();
|
||||||
|
let mut builder = FileBuilder::new();
|
||||||
|
for (i, layout) in LAYOUTS.iter().enumerate() {
|
||||||
|
let ds = builder.create_dataset(&format!("d{i}"));
|
||||||
|
ds.with_f64_data(&data).with_shape(&[rows, cols]);
|
||||||
|
if layout.chunked {
|
||||||
|
ds.with_chunks(&[CHUNK, CHUNK]);
|
||||||
|
}
|
||||||
|
if layout.deflate {
|
||||||
|
ds.with_deflate(4);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
builder.write(path).unwrap();
|
||||||
|
}
|
||||||
|
|
||||||
|
fn median(mut samples: Vec<Duration>) -> Duration {
|
||||||
|
samples.sort();
|
||||||
|
samples[samples.len() / 2]
|
||||||
|
}
|
||||||
|
|
||||||
|
fn time<T>(reps: usize, mut f: impl FnMut() -> T) -> Duration {
|
||||||
|
median(
|
||||||
|
(0..reps)
|
||||||
|
.map(|_| {
|
||||||
|
let t = Instant::now();
|
||||||
|
std::hint::black_box(f());
|
||||||
|
t.elapsed()
|
||||||
|
})
|
||||||
|
.collect(),
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn slab(start: [u64; 2], count: [u64; 2]) -> Selection {
|
||||||
|
Selection::Hyperslab {
|
||||||
|
start: start.to_vec(),
|
||||||
|
stride: vec![1, 1],
|
||||||
|
count: count.to_vec(),
|
||||||
|
block: vec![1, 1],
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn main() {
|
||||||
|
let large = std::env::args().any(|a| a == "--large");
|
||||||
|
let (rows, cols) = if large { (8192, 8192) } else { (4096, 2048) };
|
||||||
|
let total_mb = (rows * cols * 8) as f64 / (1 << 20) as f64;
|
||||||
|
if cfg!(debug_assertions) {
|
||||||
|
eprintln!("warning: debug build — numbers are meaningless. Use --release.");
|
||||||
|
}
|
||||||
|
|
||||||
|
let dir = tempfile::TempDir::new().unwrap();
|
||||||
|
let path = dir.path().join("read_harness.h5");
|
||||||
|
write_file(&path, rows, cols);
|
||||||
|
let file_mb = std::fs::metadata(&path).unwrap().len() as f64 / (1 << 20) as f64;
|
||||||
|
|
||||||
|
println!("## Read harness");
|
||||||
|
println!(
|
||||||
|
"\n{rows} x {cols} f64 ({total_mb:.0} MB per dataset), chunks {CHUNK} x {CHUNK}, file {file_mb:.0} MB\n"
|
||||||
|
);
|
||||||
|
|
||||||
|
// (label, selection, elements selected)
|
||||||
|
let selections: Vec<(&str, Selection, u64)> = vec![
|
||||||
|
(
|
||||||
|
"64 x 64 window (1 chunk)",
|
||||||
|
slab([300, 300], [64, 64]),
|
||||||
|
64 * 64,
|
||||||
|
),
|
||||||
|
(
|
||||||
|
"512 x 512 window (4-9 chunks)",
|
||||||
|
slab([1000, 700], [512, 512]),
|
||||||
|
512 * 512,
|
||||||
|
),
|
||||||
|
("one row", slab([rows / 2, 0], [1, cols]), cols),
|
||||||
|
("one column", slab([0, cols / 2], [rows, 1]), rows),
|
||||||
|
];
|
||||||
|
|
||||||
|
println!("| layout | read | selected | time ms | MB/s of selection | vs full read |");
|
||||||
|
println!("|---|---|---:|---:|---:|---:|");
|
||||||
|
for (i, layout) in LAYOUTS.iter().enumerate() {
|
||||||
|
// Fresh handle per layout so one dataset's cached chunks don't help
|
||||||
|
// (or evict) another's.
|
||||||
|
let file = File::open(&path).unwrap();
|
||||||
|
let ds = file.dataset(&format!("d{i}")).unwrap();
|
||||||
|
|
||||||
|
let full_cold = time(1, || ds.read_f64().unwrap());
|
||||||
|
let full = time(3, || ds.read_f64().unwrap());
|
||||||
|
println!(
|
||||||
|
"| {} | full (first) | {total_mb:.0} MB | {:.1} | {:.0} | |",
|
||||||
|
layout.name,
|
||||||
|
full_cold.as_secs_f64() * 1e3,
|
||||||
|
total_mb / full_cold.as_secs_f64()
|
||||||
|
);
|
||||||
|
println!(
|
||||||
|
"| {} | full (repeat) | {total_mb:.0} MB | {:.1} | {:.0} | 1.00x |",
|
||||||
|
layout.name,
|
||||||
|
full.as_secs_f64() * 1e3,
|
||||||
|
total_mb / full.as_secs_f64()
|
||||||
|
);
|
||||||
|
|
||||||
|
for (label, selection, elements) in &selections {
|
||||||
|
// A fresh handle again: measure the selection on its own, not
|
||||||
|
// served from chunks the full read just cached.
|
||||||
|
let file = File::open(&path).unwrap();
|
||||||
|
let ds = file.dataset(&format!("d{i}")).unwrap();
|
||||||
|
let got = ds.read_f64_selection(selection).unwrap();
|
||||||
|
assert_eq!(got.len() as u64, *elements, "{label}");
|
||||||
|
if let Selection::Hyperslab { start, .. } = selection {
|
||||||
|
assert_eq!(got[0], value(start[0], start[1]), "{label}: wrong data");
|
||||||
|
}
|
||||||
|
let took = time(5, || {
|
||||||
|
let file = File::open(&path).unwrap();
|
||||||
|
let ds = file.dataset(&format!("d{i}")).unwrap();
|
||||||
|
ds.read_f64_selection(selection).unwrap()
|
||||||
|
});
|
||||||
|
let mb = (*elements * 8) as f64 / (1 << 20) as f64;
|
||||||
|
println!(
|
||||||
|
"| {} | {label} | {:.2} MB | {:.2} | {:.0} | {:.3}x |",
|
||||||
|
layout.name,
|
||||||
|
mb,
|
||||||
|
took.as_secs_f64() * 1e3,
|
||||||
|
mb / took.as_secs_f64(),
|
||||||
|
took.as_secs_f64() / full_cold.as_secs_f64()
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -19,12 +19,15 @@
|
|||||||
//! cargo run --release -p clawhdf5-bench --bin search_harness -- --full # + 100K
|
//! cargo run --release -p clawhdf5-bench --bin search_harness -- --full # + 100K
|
||||||
//! cargo run --release -p clawhdf5-bench --bin search_harness -- --json out.json
|
//! cargo run --release -p clawhdf5-bench --bin search_harness -- --json out.json
|
||||||
//! cargo run --release -p clawhdf5-bench --bin search_harness -- --ann-only --uniform
|
//! cargo run --release -p clawhdf5-bench --bin search_harness -- --ann-only --uniform
|
||||||
|
//! cargo run --release -p clawhdf5-bench --bin search_harness -- --float16-study --full
|
||||||
|
//! cargo run --release -p clawhdf5-bench --bin search_harness -- --options-study --full
|
||||||
|
//! cargo run --release -p clawhdf5-bench --bin search_harness -- --signing-study --full
|
||||||
//! ```
|
//! ```
|
||||||
|
|
||||||
use std::time::{Duration, Instant};
|
use std::time::{Duration, Instant};
|
||||||
|
|
||||||
use clawhdf5_agent::{AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry};
|
use clawhdf5_agent::{AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry};
|
||||||
use clawhdf5_ann::{DistanceMetric, HnswIndex};
|
use clawhdf5_ann::{DistanceMetric, HnswIndex, Storage};
|
||||||
|
|
||||||
const DIM: usize = 384;
|
const DIM: usize = 384;
|
||||||
const K: usize = 10;
|
const K: usize = 10;
|
||||||
@@ -84,6 +87,25 @@ struct Dataset {
|
|||||||
/// that appears only on clustered data points at graph connectivity.
|
/// that appears only on clustered data points at graph connectivity.
|
||||||
static UNIFORM: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
|
static UNIFORM: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
|
||||||
|
|
||||||
|
/// `--int8`: build the HNSW index over int8-quantised vectors (a quarter of
|
||||||
|
/// the memory) instead of f32, to price the recall it costs.
|
||||||
|
static INT8: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
|
||||||
|
|
||||||
|
/// `--f16-first`: in `--float16-study`, run the float16 store first.
|
||||||
|
static F16_FIRST: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
|
||||||
|
|
||||||
|
/// `--rerank`: re-score the candidate pool against the exact vectors before
|
||||||
|
/// taking the top K.
|
||||||
|
static RERANK: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
|
||||||
|
|
||||||
|
fn storage() -> Storage {
|
||||||
|
if INT8.load(std::sync::atomic::Ordering::Relaxed) {
|
||||||
|
Storage::Int8
|
||||||
|
} else {
|
||||||
|
Storage::Float32
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
fn make_dataset(n: usize, seed: u64) -> Dataset {
|
fn make_dataset(n: usize, seed: u64) -> Dataset {
|
||||||
let mut rng = Rng(seed);
|
let mut rng = Rng(seed);
|
||||||
if UNIFORM.load(std::sync::atomic::Ordering::Relaxed) {
|
if UNIFORM.load(std::sync::atomic::Ordering::Relaxed) {
|
||||||
@@ -169,6 +191,11 @@ fn text_for(cluster: usize, i: usize, rng: &mut Rng) -> String {
|
|||||||
// Measurement helpers
|
// Measurement helpers
|
||||||
// ---------------------------------------------------------------------------
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
/// Exact cosine distance between unit-length vectors.
|
||||||
|
fn exact_dist(a: &[f32], b: &[f32]) -> f32 {
|
||||||
|
1.0 - a.iter().zip(b).map(|(x, y)| x * y).sum::<f32>()
|
||||||
|
}
|
||||||
|
|
||||||
fn exact_top_k(vectors: &[Vec<f32>], query: &[f32], k: usize) -> Vec<usize> {
|
fn exact_top_k(vectors: &[Vec<f32>], query: &[f32], k: usize) -> Vec<usize> {
|
||||||
// Vectors are unit length, so cosine order == dot-product order.
|
// Vectors are unit length, so cosine order == dot-product order.
|
||||||
let mut scored: Vec<(usize, f32)> = vectors
|
let mut scored: Vec<(usize, f32)> = vectors
|
||||||
@@ -198,6 +225,83 @@ fn summarize(mut samples: Vec<Duration>) -> Latency {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Counts live heap bytes, so a structure's cost can be measured by
|
||||||
|
/// difference.
|
||||||
|
///
|
||||||
|
/// RSS cannot do this from inside one process: freeing a large structure
|
||||||
|
/// returns its pages to the allocator's pool rather than to the OS, so
|
||||||
|
/// allocating the next one shows no change. Measured that way, a store that
|
||||||
|
/// holds the corpus twice and one that holds it once look identical.
|
||||||
|
struct CountingAllocator;
|
||||||
|
|
||||||
|
static LIVE_BYTES: std::sync::atomic::AtomicI64 = std::sync::atomic::AtomicI64::new(0);
|
||||||
|
|
||||||
|
/// High-water mark of [`LIVE_BYTES`] since it was last reset.
|
||||||
|
///
|
||||||
|
/// Live bytes at a checkpoint cannot see a buffer that was allocated and
|
||||||
|
/// freed in between, and that is exactly the shape of a transient copy —
|
||||||
|
/// which still has to fit in memory while it exists.
|
||||||
|
static PEAK_BYTES: std::sync::atomic::AtomicI64 = std::sync::atomic::AtomicI64::new(0);
|
||||||
|
|
||||||
|
fn note_peak(live: i64) {
|
||||||
|
PEAK_BYTES.fetch_max(live, std::sync::atomic::Ordering::Relaxed);
|
||||||
|
}
|
||||||
|
|
||||||
|
// SAFETY: every method forwards to the system allocator with the same layout
|
||||||
|
// it was given, and only adds bookkeeping around it.
|
||||||
|
unsafe impl std::alloc::GlobalAlloc for CountingAllocator {
|
||||||
|
unsafe fn alloc(&self, layout: std::alloc::Layout) -> *mut u8 {
|
||||||
|
let ptr = unsafe { std::alloc::System.alloc(layout) };
|
||||||
|
if !ptr.is_null() {
|
||||||
|
let live = LIVE_BYTES
|
||||||
|
.fetch_add(layout.size() as i64, std::sync::atomic::Ordering::Relaxed)
|
||||||
|
+ layout.size() as i64;
|
||||||
|
note_peak(live);
|
||||||
|
}
|
||||||
|
ptr
|
||||||
|
}
|
||||||
|
|
||||||
|
unsafe fn dealloc(&self, ptr: *mut u8, layout: std::alloc::Layout) {
|
||||||
|
LIVE_BYTES.fetch_sub(layout.size() as i64, std::sync::atomic::Ordering::Relaxed);
|
||||||
|
unsafe { std::alloc::System.dealloc(ptr, layout) }
|
||||||
|
}
|
||||||
|
|
||||||
|
unsafe fn realloc(&self, ptr: *mut u8, layout: std::alloc::Layout, new_size: usize) -> *mut u8 {
|
||||||
|
let new_ptr = unsafe { std::alloc::System.realloc(ptr, layout, new_size) };
|
||||||
|
if !new_ptr.is_null() {
|
||||||
|
let delta = new_size as i64 - layout.size() as i64;
|
||||||
|
let live = LIVE_BYTES.fetch_add(delta, std::sync::atomic::Ordering::Relaxed) + delta;
|
||||||
|
note_peak(live);
|
||||||
|
}
|
||||||
|
new_ptr
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[global_allocator]
|
||||||
|
static ALLOCATOR: CountingAllocator = CountingAllocator;
|
||||||
|
|
||||||
|
/// Live heap bytes right now.
|
||||||
|
fn heap_bytes() -> u64 {
|
||||||
|
LIVE_BYTES.load(std::sync::atomic::Ordering::Relaxed).max(0) as u64
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Start watching for a new high-water mark from the current live total.
|
||||||
|
fn reset_peak() {
|
||||||
|
PEAK_BYTES.store(
|
||||||
|
LIVE_BYTES.load(std::sync::atomic::Ordering::Relaxed),
|
||||||
|
std::sync::atomic::Ordering::Relaxed,
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// The highest live total seen since [`reset_peak`].
|
||||||
|
fn peak_bytes() -> u64 {
|
||||||
|
PEAK_BYTES.load(std::sync::atomic::Ordering::Relaxed).max(0) as u64
|
||||||
|
}
|
||||||
|
|
||||||
|
fn mib(bytes: u64) -> f64 {
|
||||||
|
bytes as f64 / (1 << 20) as f64
|
||||||
|
}
|
||||||
|
|
||||||
fn micros(d: Duration) -> f64 {
|
fn micros(d: Duration) -> f64 {
|
||||||
d.as_secs_f64() * 1e6
|
d.as_secs_f64() * 1e6
|
||||||
}
|
}
|
||||||
@@ -219,11 +323,12 @@ fn bench_ann(n: usize, json: &mut Vec<serde_json::Value>) {
|
|||||||
.collect();
|
.collect();
|
||||||
|
|
||||||
let started = Instant::now();
|
let started = Instant::now();
|
||||||
let index = HnswIndex::build_with_metric(
|
let index = HnswIndex::build_with(
|
||||||
&data.vectors,
|
&data.vectors,
|
||||||
HNSW_M,
|
HNSW_M,
|
||||||
HNSW_EF_CONSTRUCTION,
|
HNSW_EF_CONSTRUCTION,
|
||||||
DistanceMetric::Cosine,
|
DistanceMetric::Cosine,
|
||||||
|
storage(),
|
||||||
);
|
);
|
||||||
let build = started.elapsed();
|
let build = started.elapsed();
|
||||||
|
|
||||||
@@ -240,7 +345,8 @@ fn bench_ann(n: usize, json: &mut Vec<serde_json::Value>) {
|
|||||||
);
|
);
|
||||||
|
|
||||||
println!(
|
println!(
|
||||||
"\n### HNSW, N = {n}, dim = {DIM}, M = {HNSW_M}, ef_construction = {HNSW_EF_CONSTRUCTION}\n"
|
"\n### HNSW, N = {n}, dim = {DIM}, M = {HNSW_M}, ef_construction = {HNSW_EF_CONSTRUCTION}, storage = {:?}\n",
|
||||||
|
index.storage()
|
||||||
);
|
);
|
||||||
println!(
|
println!(
|
||||||
"build: {:.1} ms ({:.0} vectors/s) · exact scan: {:.0} QPS, p50 {:.0} µs\n",
|
"build: {:.1} ms ({:.0} vectors/s) · exact scan: {:.0} QPS, p50 {:.0} µs\n",
|
||||||
@@ -251,12 +357,26 @@ fn bench_ann(n: usize, json: &mut Vec<serde_json::Value>) {
|
|||||||
);
|
);
|
||||||
println!("| ef | recall@{K} | QPS | p50 µs | p99 µs |");
|
println!("| ef | recall@{K} | QPS | p50 µs | p99 µs |");
|
||||||
println!("|---:|---:|---:|---:|---:|");
|
println!("|---:|---:|---:|---:|---:|");
|
||||||
|
// With a quantised index the distances it returns are approximate, so
|
||||||
|
// the candidates are re-scored against the exact vectors the caller
|
||||||
|
// already holds (in the agent, the embedding cache) before taking the
|
||||||
|
// top K. `--rerank` prices that: it costs one exact distance per
|
||||||
|
// candidate and is what decides whether int8 is usable.
|
||||||
|
let rerank = RERANK.load(std::sync::atomic::Ordering::Relaxed);
|
||||||
|
let pool = if rerank { K * 4 } else { K };
|
||||||
for ef in EF_VALUES {
|
for ef in EF_VALUES {
|
||||||
let mut hits = 0usize;
|
let mut hits = 0usize;
|
||||||
let mut samples = Vec::with_capacity(data.queries.len());
|
let mut samples = Vec::with_capacity(data.queries.len());
|
||||||
for (q, want) in data.queries.iter().zip(&truth) {
|
for (q, want) in data.queries.iter().zip(&truth) {
|
||||||
let t = Instant::now();
|
let t = Instant::now();
|
||||||
let got = index.search(q, K, ef);
|
let mut got = index.search(q, pool, ef.max(pool));
|
||||||
|
if rerank {
|
||||||
|
for cand in &mut got {
|
||||||
|
cand.1 = exact_dist(&data.vectors[cand.0], q);
|
||||||
|
}
|
||||||
|
got.select_nth_unstable_by(K - 1, |a, b| a.1.total_cmp(&b.1));
|
||||||
|
got.truncate(K);
|
||||||
|
}
|
||||||
samples.push(t.elapsed());
|
samples.push(t.elapsed());
|
||||||
hits += got.iter().filter(|(id, _)| want.contains(id)).count();
|
hits += got.iter().filter(|(id, _)| want.contains(id)).count();
|
||||||
}
|
}
|
||||||
@@ -369,6 +489,394 @@ fn bench_end_to_end(n: usize, json: &mut Vec<serde_json::Value>) {
|
|||||||
}));
|
}));
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// Signing study: what does an Ed25519-signed checkpoint cost?
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
/// `--signing-study`: checkpoint time unsigned vs signed, `verify` time, and
|
||||||
|
/// the file-size cost of the stored per-record hashes. Default store
|
||||||
|
/// settings (float16, int8 index). Medians of five checkpoints / three
|
||||||
|
/// verifies.
|
||||||
|
fn signing_study(n: usize) {
|
||||||
|
use clawhdf5_agent::signing::SigningKey;
|
||||||
|
let data = make_dataset(n, 0x516 ^ n as u64);
|
||||||
|
let mut rng = Rng(9);
|
||||||
|
let entries: Vec<MemoryEntry> = data
|
||||||
|
.vectors
|
||||||
|
.iter()
|
||||||
|
.enumerate()
|
||||||
|
.map(|(i, v)| MemoryEntry {
|
||||||
|
chunk: text_for(data.cluster_of[i], i, &mut rng),
|
||||||
|
embedding: v.clone(),
|
||||||
|
source_channel: "bench".into(),
|
||||||
|
timestamp: i as f64,
|
||||||
|
session_id: format!("s{}", i % 50),
|
||||||
|
tags: format!("t{i}"),
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
let dir = tempfile::tempdir().unwrap();
|
||||||
|
let path = dir.path().join("sign.h5");
|
||||||
|
let mut mem = HDF5Memory::create(MemoryConfig::new(path.clone(), "bench", DIM)).unwrap();
|
||||||
|
mem.save_batch(entries).unwrap();
|
||||||
|
std::hint::black_box(mem.hybrid_search(&data.queries[0], "", 1.0, 0.0, K));
|
||||||
|
|
||||||
|
let median = |mut v: Vec<Duration>| {
|
||||||
|
v.sort();
|
||||||
|
v[v.len() / 2]
|
||||||
|
};
|
||||||
|
let checkpoint = |mem: &mut HDF5Memory| {
|
||||||
|
median(
|
||||||
|
(0..5)
|
||||||
|
.map(|_| {
|
||||||
|
let t = Instant::now();
|
||||||
|
mem.flush_wal().unwrap();
|
||||||
|
t.elapsed()
|
||||||
|
})
|
||||||
|
.collect(),
|
||||||
|
)
|
||||||
|
};
|
||||||
|
let unsigned = checkpoint(&mut mem);
|
||||||
|
let unsigned_bytes = std::fs::metadata(&path).unwrap().len();
|
||||||
|
let key = SigningKey::from_bytes(&[7; 32]);
|
||||||
|
mem.set_signing_key(key.clone());
|
||||||
|
let signed = checkpoint(&mut mem);
|
||||||
|
let signed_bytes = std::fs::metadata(&path).unwrap().len();
|
||||||
|
drop(mem);
|
||||||
|
let vk = key.verifying_key();
|
||||||
|
let verify = median(
|
||||||
|
(0..3)
|
||||||
|
.map(|_| {
|
||||||
|
let t = Instant::now();
|
||||||
|
let r = HDF5Memory::verify(&path, &vk).unwrap();
|
||||||
|
let d = t.elapsed();
|
||||||
|
assert!(r.is_valid());
|
||||||
|
d
|
||||||
|
})
|
||||||
|
.collect(),
|
||||||
|
);
|
||||||
|
println!(
|
||||||
|
"| {n} | {:.1} | {:.1} | {:+.1} | {:.1} | {:+.2} |",
|
||||||
|
millis(unsigned),
|
||||||
|
millis(signed),
|
||||||
|
millis(signed) - millis(unsigned),
|
||||||
|
millis(verify),
|
||||||
|
(signed_bytes as f64 - unsigned_bytes as f64) / (1024.0 * 1024.0),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// Search options study: source filters, re-ranking, confidence rejection
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
/// `--options-study`: what `HDF5Memory::search`'s options cost and whether a
|
||||||
|
/// filtered search finds the right records. Filters keep 50%, 10% or 1% of
|
||||||
|
/// the store at random, or two whole clusters away from the query (the case
|
||||||
|
/// the index cannot serve, which falls back to an exact scan). Recall is
|
||||||
|
/// vector-only against an exact scan of the allowed records; latency is full
|
||||||
|
/// hybrid search. Hebbian boosting is off.
|
||||||
|
fn options_study(n: usize) {
|
||||||
|
use clawhdf5_agent::SearchOptions;
|
||||||
|
use clawhdf5_agent::confidence::ConfidenceConfig;
|
||||||
|
use clawhdf5_agent::hybrid::Fusion;
|
||||||
|
use clawhdf5_agent::reranker::ReRankConfig;
|
||||||
|
|
||||||
|
let data = make_dataset(n, 0x0B7 ^ n as u64);
|
||||||
|
let n_clusters = data.cluster_of.iter().max().map_or(1, |m| m + 1);
|
||||||
|
let mut rng = Rng(5);
|
||||||
|
let bucket_of: Vec<usize> = (0..n).map(|_| rng.below(100)).collect();
|
||||||
|
let bucket = &bucket_of;
|
||||||
|
let query_texts: Vec<String> = data
|
||||||
|
.query_cluster
|
||||||
|
.iter()
|
||||||
|
.enumerate()
|
||||||
|
.map(|(i, c)| text_for(*c, i, &mut rng))
|
||||||
|
.collect();
|
||||||
|
let exact_top = |q: &[f32], allowed: &dyn Fn(usize) -> bool| -> Vec<usize> {
|
||||||
|
let mut s: Vec<(usize, f32)> = (0..n)
|
||||||
|
.filter(|&i| allowed(i))
|
||||||
|
.map(|i| (i, data.vectors[i].iter().zip(q).map(|(a, b)| a * b).sum()))
|
||||||
|
.collect();
|
||||||
|
s.sort_by(|a, b| b.1.total_cmp(&a.1).then(a.0.cmp(&b.0)));
|
||||||
|
s.into_iter().take(K).map(|(i, _)| i).collect()
|
||||||
|
};
|
||||||
|
|
||||||
|
// Two stores: channel = random bucket, and channel = cluster.
|
||||||
|
let dir = tempfile::tempdir().unwrap();
|
||||||
|
let mut stores = Vec::new();
|
||||||
|
for by_cluster in [false, true] {
|
||||||
|
let mut rng = Rng(3);
|
||||||
|
let entries: Vec<MemoryEntry> = data
|
||||||
|
.vectors
|
||||||
|
.iter()
|
||||||
|
.enumerate()
|
||||||
|
.map(|(i, v)| MemoryEntry {
|
||||||
|
chunk: text_for(data.cluster_of[i], i, &mut rng),
|
||||||
|
embedding: v.clone(),
|
||||||
|
source_channel: if by_cluster {
|
||||||
|
format!("c{}", data.cluster_of[i])
|
||||||
|
} else {
|
||||||
|
format!("b{}", bucket[i])
|
||||||
|
},
|
||||||
|
timestamp: i as f64,
|
||||||
|
session_id: format!("s{}", i % 50),
|
||||||
|
tags: format!("t{i}"),
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
let mut config = MemoryConfig::new(
|
||||||
|
dir.path().join(format!("opt_{by_cluster}.h5")),
|
||||||
|
"bench",
|
||||||
|
DIM,
|
||||||
|
);
|
||||||
|
config.hebbian_boost = 0.0;
|
||||||
|
let mut mem = HDF5Memory::create(config).unwrap();
|
||||||
|
mem.save_batch(entries).unwrap();
|
||||||
|
std::hint::black_box(mem.search(&data.queries[0], "", &SearchOptions::new(K)));
|
||||||
|
stores.push(mem);
|
||||||
|
}
|
||||||
|
|
||||||
|
let vector_only = SearchOptions::new(K).with_fusion(Fusion::Weighted {
|
||||||
|
vector: 1.0,
|
||||||
|
keyword: 0.0,
|
||||||
|
});
|
||||||
|
// (label, store, channels for query i, allowed(i, record))
|
||||||
|
type Case<'a> = (
|
||||||
|
String,
|
||||||
|
usize,
|
||||||
|
Box<dyn Fn(usize) -> Option<Vec<String>> + 'a>,
|
||||||
|
Box<dyn Fn(usize, usize) -> bool + 'a>,
|
||||||
|
);
|
||||||
|
let mut cases: Vec<Case> = vec![(
|
||||||
|
"no filter".into(),
|
||||||
|
0,
|
||||||
|
Box::new(|_| None),
|
||||||
|
Box::new(|_, _| true),
|
||||||
|
)];
|
||||||
|
for pct in [50usize, 10, 1] {
|
||||||
|
cases.push((
|
||||||
|
format!("random {pct}%"),
|
||||||
|
0,
|
||||||
|
Box::new(move |_| Some((0..pct).map(|b| format!("b{b}")).collect())),
|
||||||
|
Box::new(move |_, i| bucket[i] < pct),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
let d = &data;
|
||||||
|
let away = move |qi: usize| {
|
||||||
|
let qc = d.query_cluster[qi];
|
||||||
|
[
|
||||||
|
(qc + n_clusters / 3) % n_clusters,
|
||||||
|
(qc + 2 * n_clusters / 3) % n_clusters,
|
||||||
|
]
|
||||||
|
};
|
||||||
|
cases.push((
|
||||||
|
"2 clusters away from the query".into(),
|
||||||
|
1,
|
||||||
|
Box::new(move |qi| Some(away(qi).iter().map(|c| format!("c{c}")).collect())),
|
||||||
|
Box::new(move |qi, i| away(qi).contains(&d.cluster_of[i])),
|
||||||
|
));
|
||||||
|
|
||||||
|
for (label, store, channels, allowed) in &cases {
|
||||||
|
let mem = &mut stores[*store];
|
||||||
|
let mut hits = 0;
|
||||||
|
let mut kept = 0;
|
||||||
|
for (qi, q) in data.queries.iter().enumerate() {
|
||||||
|
let mut opts = vector_only.clone();
|
||||||
|
opts.source_channels = channels(qi);
|
||||||
|
let got = mem.search(q, "", &opts);
|
||||||
|
let want = exact_top(q, &|i| allowed(qi, i));
|
||||||
|
kept += want.len();
|
||||||
|
hits += got.iter().filter(|r| want.contains(&r.index)).count();
|
||||||
|
}
|
||||||
|
let latency = summarize(
|
||||||
|
(0..N_QUERIES)
|
||||||
|
.map(|qi| {
|
||||||
|
let mut opts = SearchOptions::new(K);
|
||||||
|
opts.source_channels = channels(qi);
|
||||||
|
let t = Instant::now();
|
||||||
|
std::hint::black_box(mem.search(&data.queries[qi], &query_texts[qi], &opts));
|
||||||
|
t.elapsed()
|
||||||
|
})
|
||||||
|
.collect(),
|
||||||
|
);
|
||||||
|
println!(
|
||||||
|
"| {n} | {label} | {:.4} | {:.3} | {:.3} |",
|
||||||
|
hits as f64 / kept.max(1) as f64,
|
||||||
|
millis(latency.p50),
|
||||||
|
millis(latency.p99),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
let mem = &mut stores[0];
|
||||||
|
for (label, opts) in [
|
||||||
|
(
|
||||||
|
"re-rank",
|
||||||
|
SearchOptions::new(K).with_rerank(ReRankConfig::default()),
|
||||||
|
),
|
||||||
|
(
|
||||||
|
"re-rank + confidence",
|
||||||
|
SearchOptions::new(K)
|
||||||
|
.with_rerank(ReRankConfig::default())
|
||||||
|
.with_confidence(ConfidenceConfig::default()),
|
||||||
|
),
|
||||||
|
] {
|
||||||
|
let latency = summarize(
|
||||||
|
(0..N_QUERIES)
|
||||||
|
.map(|qi| {
|
||||||
|
let t = Instant::now();
|
||||||
|
std::hint::black_box(mem.search(&data.queries[qi], &query_texts[qi], &opts));
|
||||||
|
t.elapsed()
|
||||||
|
})
|
||||||
|
.collect(),
|
||||||
|
);
|
||||||
|
println!(
|
||||||
|
"| {n} | {label} | — | {:.3} | {:.3} |",
|
||||||
|
millis(latency.p50),
|
||||||
|
millis(latency.p99)
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// float16 study: what does half-precision embedding storage cost?
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
/// `--float16-study`: the same data in an `f32` store and a `float16` store.
|
||||||
|
/// Reports file size, checkpoint and open time, vector-search recall@10
|
||||||
|
/// against an exact scan of the *original* f32 vectors, how often the two
|
||||||
|
/// stores return the same top 10, and `hybrid_search` latency. Hebbian
|
||||||
|
/// boosting is off, so every query sees the same store.
|
||||||
|
fn float16_study(n: usize) {
|
||||||
|
let data = make_dataset(n, 0xF16 ^ n as u64);
|
||||||
|
let mut rng = Rng(11);
|
||||||
|
let query_texts: Vec<String> = data
|
||||||
|
.query_cluster
|
||||||
|
.iter()
|
||||||
|
.enumerate()
|
||||||
|
.map(|(i, c)| text_for(*c, i, &mut rng))
|
||||||
|
.collect();
|
||||||
|
|
||||||
|
// Exact top K by cosine (the vectors are unit length) on the f32 inputs.
|
||||||
|
let exact: Vec<Vec<usize>> = data
|
||||||
|
.queries
|
||||||
|
.iter()
|
||||||
|
.map(|q| {
|
||||||
|
let mut scored: Vec<(usize, f32)> = data
|
||||||
|
.vectors
|
||||||
|
.iter()
|
||||||
|
.enumerate()
|
||||||
|
.map(|(i, v)| (i, v.iter().zip(q).map(|(a, b)| a * b).sum()))
|
||||||
|
.collect();
|
||||||
|
scored.sort_by(|a, b| b.1.total_cmp(&a.1).then(a.0.cmp(&b.0)));
|
||||||
|
scored.into_iter().take(K).map(|(i, _)| i).collect()
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
|
||||||
|
let dir = tempfile::tempdir().unwrap();
|
||||||
|
let mut per_variant: Vec<(bool, Vec<Vec<usize>>)> = Vec::new();
|
||||||
|
// `--f16-first` swaps the order, to check the numbers do not depend on
|
||||||
|
// which store runs first (page cache, allocator, CPU frequency).
|
||||||
|
let order = if F16_FIRST.load(std::sync::atomic::Ordering::Relaxed) {
|
||||||
|
[true, false]
|
||||||
|
} else {
|
||||||
|
[false, true]
|
||||||
|
};
|
||||||
|
for float16 in order {
|
||||||
|
let path = dir.path().join(format!("f16study_{float16}.h5"));
|
||||||
|
let mut rng = Rng(3);
|
||||||
|
let entries: Vec<MemoryEntry> = data
|
||||||
|
.vectors
|
||||||
|
.iter()
|
||||||
|
.enumerate()
|
||||||
|
.map(|(i, v)| MemoryEntry {
|
||||||
|
chunk: text_for(data.cluster_of[i], i, &mut rng),
|
||||||
|
embedding: v.clone(),
|
||||||
|
source_channel: "bench".into(),
|
||||||
|
timestamp: i as f64,
|
||||||
|
session_id: format!("s{}", i % 50),
|
||||||
|
tags: format!("t{i}"),
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
let mut config = MemoryConfig::new(path.clone(), "bench", DIM);
|
||||||
|
config.float16 = float16;
|
||||||
|
config.hebbian_boost = 0.0;
|
||||||
|
let mut mem = HDF5Memory::create(config).unwrap();
|
||||||
|
mem.save_batch(entries).unwrap();
|
||||||
|
// Build the indexes, then time a checkpoint that writes everything.
|
||||||
|
std::hint::black_box(mem.hybrid_search(&data.queries[0], "", 1.0, 0.0, K));
|
||||||
|
let t = Instant::now();
|
||||||
|
mem.flush_wal().unwrap();
|
||||||
|
let checkpoint = t.elapsed();
|
||||||
|
drop(mem);
|
||||||
|
let file_bytes = std::fs::metadata(&path).unwrap().len();
|
||||||
|
|
||||||
|
// Median of three opens.
|
||||||
|
let mut opens: Vec<Duration> = (0..3)
|
||||||
|
.map(|_| {
|
||||||
|
let t = Instant::now();
|
||||||
|
let m = HDF5Memory::open(&path).unwrap();
|
||||||
|
let d = t.elapsed();
|
||||||
|
drop(m);
|
||||||
|
d
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
opens.sort();
|
||||||
|
let mut mem = HDF5Memory::open(&path).unwrap();
|
||||||
|
|
||||||
|
// Vector-only search: empty text, all weight on the vector stage.
|
||||||
|
let results: Vec<Vec<usize>> = data
|
||||||
|
.queries
|
||||||
|
.iter()
|
||||||
|
.map(|q| {
|
||||||
|
mem.hybrid_search(q, "", 1.0, 0.0, K)
|
||||||
|
.iter()
|
||||||
|
.map(|r| r.index)
|
||||||
|
.collect()
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
let hits: usize = results
|
||||||
|
.iter()
|
||||||
|
.zip(&exact)
|
||||||
|
.map(|(got, want)| got.iter().filter(|i| want.contains(i)).count())
|
||||||
|
.sum();
|
||||||
|
let recall = hits as f64 / (K * data.queries.len()) as f64;
|
||||||
|
|
||||||
|
let latency = summarize(
|
||||||
|
(0..N_QUERIES)
|
||||||
|
.map(|i| {
|
||||||
|
let t = Instant::now();
|
||||||
|
std::hint::black_box(mem.hybrid_search(
|
||||||
|
&data.queries[i],
|
||||||
|
&query_texts[i],
|
||||||
|
0.4,
|
||||||
|
0.6,
|
||||||
|
K,
|
||||||
|
));
|
||||||
|
t.elapsed()
|
||||||
|
})
|
||||||
|
.collect(),
|
||||||
|
);
|
||||||
|
let overlap = match per_variant.first() {
|
||||||
|
Some((_, other)) => {
|
||||||
|
let same: usize = results
|
||||||
|
.iter()
|
||||||
|
.zip(other)
|
||||||
|
.map(|(a, b)| a.iter().filter(|i| b.contains(i)).count())
|
||||||
|
.sum();
|
||||||
|
format!("{:.4}", same as f64 / (K * data.queries.len()) as f64)
|
||||||
|
}
|
||||||
|
None => "—".into(),
|
||||||
|
};
|
||||||
|
println!(
|
||||||
|
"| {n} | {} | {:.1} | {:.0} | {:.1} | {recall:.4} | {overlap} | {:.3} |",
|
||||||
|
if float16 { "float16" } else { "f32" },
|
||||||
|
mib(file_bytes),
|
||||||
|
millis(checkpoint),
|
||||||
|
millis(opens[1]),
|
||||||
|
millis(latency.p50),
|
||||||
|
);
|
||||||
|
per_variant.push((float16, results));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
// ---------------------------------------------------------------------------
|
// ---------------------------------------------------------------------------
|
||||||
// Fusion study: does capping the keyword candidate pool change the ranking?
|
// Fusion study: does capping the keyword candidate pool change the ranking?
|
||||||
// ---------------------------------------------------------------------------
|
// ---------------------------------------------------------------------------
|
||||||
@@ -394,11 +902,12 @@ fn fusion_study(n: usize) {
|
|||||||
.map(|(i, c)| text_for(*c, i, &mut rng))
|
.map(|(i, c)| text_for(*c, i, &mut rng))
|
||||||
.collect();
|
.collect();
|
||||||
let bm25 = BM25Index::build(&texts, &vec![0u8; n]);
|
let bm25 = BM25Index::build(&texts, &vec![0u8; n]);
|
||||||
let index = HnswIndex::build_with_metric(
|
let index = HnswIndex::build_with(
|
||||||
&data.vectors,
|
&data.vectors,
|
||||||
HNSW_M,
|
HNSW_M,
|
||||||
HNSW_EF_CONSTRUCTION,
|
HNSW_EF_CONSTRUCTION,
|
||||||
DistanceMetric::Cosine,
|
DistanceMetric::Cosine,
|
||||||
|
storage(),
|
||||||
);
|
);
|
||||||
|
|
||||||
let vec_pool = (K * 8).max(64);
|
let vec_pool = (K * 8).max(64);
|
||||||
@@ -450,6 +959,69 @@ fn fusion_study(n: usize) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// What an in-memory store costs, stage by stage. The vectors are the floor:
|
||||||
|
/// everything above it is bookkeeping that could in principle be shared.
|
||||||
|
fn bench_footprint(n: usize) {
|
||||||
|
let data = make_dataset(n, 0xF007 ^ n as u64);
|
||||||
|
let mut rng = Rng(11);
|
||||||
|
let dir = tempfile::TempDir::new().unwrap();
|
||||||
|
let path = dir.path().join("footprint.h5");
|
||||||
|
|
||||||
|
let base = heap_bytes();
|
||||||
|
let entries: Vec<MemoryEntry> = data
|
||||||
|
.vectors
|
||||||
|
.iter()
|
||||||
|
.enumerate()
|
||||||
|
.map(|(i, v)| MemoryEntry {
|
||||||
|
chunk: text_for(data.cluster_of[i], i, &mut rng),
|
||||||
|
embedding: v.clone(),
|
||||||
|
source_channel: "bench".into(),
|
||||||
|
timestamp: i as f64,
|
||||||
|
session_id: format!("s{}", i % 50),
|
||||||
|
tags: format!("t{i}"),
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
let after_entries = heap_bytes();
|
||||||
|
|
||||||
|
let mut config = MemoryConfig::new(path, "bench", DIM);
|
||||||
|
config.quantized_index = INT8.load(std::sync::atomic::Ordering::Relaxed);
|
||||||
|
let mut mem = HDF5Memory::create(config).unwrap();
|
||||||
|
mem.save_batch(entries).unwrap();
|
||||||
|
let after_store = heap_bytes();
|
||||||
|
|
||||||
|
// First query builds the vector and keyword indexes.
|
||||||
|
std::hint::black_box(mem.hybrid_search(&data.queries[0], "record", 0.7, 0.3, K));
|
||||||
|
let after_indexes = heap_bytes();
|
||||||
|
|
||||||
|
// Reopening is the figure that matters for a long-lived process, and the
|
||||||
|
// only one RSS reports honestly: memory freed when the ingest buffers went
|
||||||
|
// away stays in the allocator's pool, so the stage deltas above understate
|
||||||
|
// what was given back.
|
||||||
|
let path = mem.config().path.clone();
|
||||||
|
drop(mem);
|
||||||
|
let before_open = heap_bytes();
|
||||||
|
reset_peak();
|
||||||
|
let reopened = HDF5Memory::open(&path).unwrap();
|
||||||
|
let after_open = heap_bytes();
|
||||||
|
let loaded = after_open.saturating_sub(before_open);
|
||||||
|
// Peak over the open, not just what it leaves behind: a buffer allocated
|
||||||
|
// and freed during the parse never shows up in the live total.
|
||||||
|
let peak = peak_bytes().saturating_sub(before_open);
|
||||||
|
drop(reopened);
|
||||||
|
|
||||||
|
let raw = (n * DIM * 4) as u64;
|
||||||
|
println!(
|
||||||
|
"| {n} | {:.0} | {:.0} | {:.0} | {:.0} | {:.0} | {:.0} | {:.2}x |",
|
||||||
|
mib(raw),
|
||||||
|
mib(after_entries.saturating_sub(base)),
|
||||||
|
mib(after_store.saturating_sub(after_entries)),
|
||||||
|
mib(after_indexes.saturating_sub(after_store)),
|
||||||
|
mib(loaded),
|
||||||
|
mib(peak),
|
||||||
|
loaded as f64 / raw as f64,
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
fn main() {
|
fn main() {
|
||||||
let args: Vec<String> = std::env::args().skip(1).collect();
|
let args: Vec<String> = std::env::args().skip(1).collect();
|
||||||
let full = args.iter().any(|a| a == "--full");
|
let full = args.iter().any(|a| a == "--full");
|
||||||
@@ -464,6 +1036,60 @@ fn main() {
|
|||||||
}
|
}
|
||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
|
if args.iter().any(|a| a == "--signing-study") {
|
||||||
|
println!("## Signed checkpoints ({DIM}-dim, float16, int8 index)\n");
|
||||||
|
println!(
|
||||||
|
"| N | checkpoint ms, unsigned | checkpoint ms, signed | signing adds ms | verify ms | file MiB added |"
|
||||||
|
);
|
||||||
|
println!("|---:|---:|---:|---:|---:|---:|");
|
||||||
|
for &n in if full {
|
||||||
|
&[1_000, 10_000, 100_000][..]
|
||||||
|
} else {
|
||||||
|
&[1_000, 10_000][..]
|
||||||
|
} {
|
||||||
|
signing_study(n);
|
||||||
|
}
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
if args.iter().any(|a| a == "--options-study") {
|
||||||
|
println!("## Search options ({DIM}-dim, k = {K}, Hebbian boost off)\n");
|
||||||
|
println!("| N | options | filtered recall@10 | p50 ms | p99 ms |");
|
||||||
|
println!("|---:|---|---:|---:|---:|");
|
||||||
|
for &n in if full {
|
||||||
|
&[10_000, 100_000][..]
|
||||||
|
} else {
|
||||||
|
&[10_000][..]
|
||||||
|
} {
|
||||||
|
options_study(n);
|
||||||
|
}
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
if args.iter().any(|a| a == "--f16-first") {
|
||||||
|
F16_FIRST.store(true, std::sync::atomic::Ordering::Relaxed);
|
||||||
|
}
|
||||||
|
if args.iter().any(|a| a == "--float16-study") {
|
||||||
|
println!("## float16 embedding storage ({DIM}-dim, int8 index, Hebbian boost off)\n");
|
||||||
|
println!(
|
||||||
|
"| N | embeddings | file MiB | checkpoint ms | open ms | recall@10 | top-10 overlap with the other | hybrid p50 ms |"
|
||||||
|
);
|
||||||
|
println!("|---:|---|---:|---:|---:|---:|---:|---:|");
|
||||||
|
for &n in if full {
|
||||||
|
&[1_000, 10_000, 100_000][..]
|
||||||
|
} else {
|
||||||
|
&[1_000, 10_000][..]
|
||||||
|
} {
|
||||||
|
float16_study(n);
|
||||||
|
}
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
if args.iter().any(|a| a == "--int8") {
|
||||||
|
INT8.store(true, std::sync::atomic::Ordering::Relaxed);
|
||||||
|
println!("(int8-quantised index vectors)");
|
||||||
|
}
|
||||||
|
if args.iter().any(|a| a == "--rerank") {
|
||||||
|
RERANK.store(true, std::sync::atomic::Ordering::Relaxed);
|
||||||
|
println!("(candidates re-scored against exact vectors)");
|
||||||
|
}
|
||||||
if args.iter().any(|a| a == "--uniform") {
|
if args.iter().any(|a| a == "--uniform") {
|
||||||
UNIFORM.store(true, std::sync::atomic::Ordering::Relaxed);
|
UNIFORM.store(true, std::sync::atomic::Ordering::Relaxed);
|
||||||
println!("(uniform random data)");
|
println!("(uniform random data)");
|
||||||
@@ -485,8 +1111,24 @@ fn main() {
|
|||||||
|
|
||||||
let mut json = Vec::new();
|
let mut json = Vec::new();
|
||||||
println!("## Search harness");
|
println!("## Search harness");
|
||||||
for &n in sizes {
|
|
||||||
bench_ann(n, &mut json);
|
if args.iter().any(|a| a == "--footprint") {
|
||||||
|
println!("\n### Resident memory, {DIM}-dim f32\n");
|
||||||
|
println!(
|
||||||
|
"| N | vectors (raw) | entries MiB | store MiB | indexes MiB | reopened MiB | peak during open MiB | reopened / raw |"
|
||||||
|
);
|
||||||
|
println!("|---:|---:|---:|---:|---:|---:|---:|---:|");
|
||||||
|
for &n in sizes {
|
||||||
|
bench_footprint(n);
|
||||||
|
}
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
// `--e2e-only` skips the index benchmarks, so the end-to-end section runs
|
||||||
|
// in a process that has not already spun up a thread pool.
|
||||||
|
if !args.iter().any(|a| a == "--e2e-only") {
|
||||||
|
for &n in sizes {
|
||||||
|
bench_ann(n, &mut json);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
if ann_only {
|
if ann_only {
|
||||||
|
|||||||
@@ -1,7 +1,8 @@
|
|||||||
[package]
|
[package]
|
||||||
name = "clawhdf5-cli"
|
name = "clawhdf5-cli"
|
||||||
version = "2.4.0"
|
version = "2.7.0"
|
||||||
edition = "2024"
|
edition = "2024"
|
||||||
|
rust-version.workspace = true
|
||||||
license = "MIT"
|
license = "MIT"
|
||||||
description = "CLI for clawhdf5 agent memory — create, save, search, recall, stats"
|
description = "CLI for clawhdf5 agent memory — create, save, search, recall, stats"
|
||||||
repository = "https://git.redclaw.dev/quantumclaw/clawhdf5"
|
repository = "https://git.redclaw.dev/quantumclaw/clawhdf5"
|
||||||
@@ -14,7 +15,7 @@ name = "clawhdf5"
|
|||||||
path = "src/main.rs"
|
path = "src/main.rs"
|
||||||
|
|
||||||
[dependencies]
|
[dependencies]
|
||||||
clawhdf5-agent = { path = "../clawhdf5-agent", version = "2.4.0" }
|
clawhdf5-agent = { path = "../clawhdf5-agent", version = "2.7.0" }
|
||||||
clap = { version = "4", features = ["derive", "env"] }
|
clap = { version = "4", features = ["derive", "env"] }
|
||||||
serde_json = "1"
|
serde_json = "1"
|
||||||
serde = { workspace = true }
|
serde = { workspace = true }
|
||||||
|
|||||||
+165
-17
@@ -1,15 +1,22 @@
|
|||||||
use std::path::PathBuf;
|
use std::path::{Path, PathBuf};
|
||||||
|
|
||||||
use clap::{Parser, Subcommand};
|
use clap::{Parser, Subcommand};
|
||||||
|
use clawhdf5_agent::signing::{self, SigningKey, VerifyingKey};
|
||||||
use clawhdf5_agent::{AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry};
|
use clawhdf5_agent::{AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry};
|
||||||
|
|
||||||
/// ClawhDF5 — HDF5-backed cognitive memory for AI agents
|
/// ClawhDF5 — HDF5-backed cognitive memory for AI agents
|
||||||
#[derive(Parser)]
|
#[derive(Parser)]
|
||||||
#[command(name = "clawhdf5", version, about)]
|
#[command(name = "clawhdf5", version, about)]
|
||||||
struct Cli {
|
struct Cli {
|
||||||
/// Path to the .h5 memory file
|
/// Path to the .h5 memory file (not needed for `keygen`)
|
||||||
#[arg(short, long, env = "CLAWHDF5_PATH")]
|
#[arg(short, long, env = "CLAWHDF5_PATH")]
|
||||||
path: PathBuf,
|
path: Option<PathBuf>,
|
||||||
|
|
||||||
|
/// File holding an Ed25519 signing key (64 hex characters, from
|
||||||
|
/// `keygen`). Every checkpoint this command makes is then signed; a
|
||||||
|
/// signed store refuses to checkpoint without it.
|
||||||
|
#[arg(long, env = "CLAWHDF5_SIGNING_KEY", global = true)]
|
||||||
|
signing_key: Option<PathBuf>,
|
||||||
|
|
||||||
#[command(subcommand)]
|
#[command(subcommand)]
|
||||||
command: Commands,
|
command: Commands,
|
||||||
@@ -28,6 +35,22 @@ enum Commands {
|
|||||||
/// Enable write-ahead log
|
/// Enable write-ahead log
|
||||||
#[arg(long)]
|
#[arg(long)]
|
||||||
wal: bool,
|
wal: bool,
|
||||||
|
/// Hold the vector index's copy of the embeddings as f32 instead of
|
||||||
|
/// the default int8 (which uses a quarter of the memory and is faster
|
||||||
|
/// at equal recall)
|
||||||
|
#[arg(long)]
|
||||||
|
f32_index: bool,
|
||||||
|
/// Accepted for compatibility; int8 is now the default
|
||||||
|
#[arg(long, hide = true, conflicts_with = "f32_index")]
|
||||||
|
quantized_index: bool,
|
||||||
|
/// Store embeddings as full-precision f32 instead of the default
|
||||||
|
/// half precision (float16: half the bytes, about three significant
|
||||||
|
/// digits, values within ±65504)
|
||||||
|
#[arg(long)]
|
||||||
|
f32: bool,
|
||||||
|
/// Accepted for compatibility; float16 is now the default
|
||||||
|
#[arg(long, hide = true, conflicts_with = "f32")]
|
||||||
|
float16: bool,
|
||||||
},
|
},
|
||||||
/// Save a memory entry (reads JSON from stdin or --json)
|
/// Save a memory entry (reads JSON from stdin or --json)
|
||||||
Save {
|
Save {
|
||||||
@@ -75,6 +98,38 @@ enum Commands {
|
|||||||
/// Destination path
|
/// Destination path
|
||||||
dest: PathBuf,
|
dest: PathBuf,
|
||||||
},
|
},
|
||||||
|
/// Generate an Ed25519 signing key for signed checkpoints
|
||||||
|
Keygen {
|
||||||
|
/// Where to write the secret key (created new, owner-only on Unix)
|
||||||
|
#[arg(long)]
|
||||||
|
out: PathBuf,
|
||||||
|
},
|
||||||
|
/// Verify a signed store against a public key; exit status 2 if not valid
|
||||||
|
Verify {
|
||||||
|
/// The trusted public key: 64 hex characters, or a file holding them
|
||||||
|
#[arg(long)]
|
||||||
|
public_key: String,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
fn read_signing_key(path: &Path) -> Result<SigningKey, Box<dyn std::error::Error>> {
|
||||||
|
let text = std::fs::read_to_string(path)
|
||||||
|
.map_err(|e| format!("cannot read signing key {}: {e}", path.display()))?;
|
||||||
|
let bytes = signing::from_hex::<32>(&text)
|
||||||
|
.ok_or_else(|| format!("{} is not a 64-hex-character key", path.display()))?;
|
||||||
|
Ok(SigningKey::from_bytes(&bytes))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Open for writing, with the signing key applied if one was given.
|
||||||
|
fn open_writable(
|
||||||
|
path: &Path,
|
||||||
|
key: &Option<SigningKey>,
|
||||||
|
) -> Result<HDF5Memory, Box<dyn std::error::Error>> {
|
||||||
|
let mut mem = HDF5Memory::open(path)?;
|
||||||
|
if let Some(k) = key {
|
||||||
|
mem.set_signing_key(k.clone());
|
||||||
|
}
|
||||||
|
Ok(mem)
|
||||||
}
|
}
|
||||||
|
|
||||||
fn main() {
|
fn main() {
|
||||||
@@ -87,17 +142,76 @@ fn main() {
|
|||||||
}
|
}
|
||||||
|
|
||||||
fn run(cli: Cli) -> Result<(), Box<dyn std::error::Error>> {
|
fn run(cli: Cli) -> Result<(), Box<dyn std::error::Error>> {
|
||||||
|
if let Commands::Keygen { out } = &cli.command {
|
||||||
|
let key = signing::generate_key();
|
||||||
|
let mut opts = std::fs::OpenOptions::new();
|
||||||
|
opts.write(true).create_new(true);
|
||||||
|
#[cfg(unix)]
|
||||||
|
{
|
||||||
|
use std::os::unix::fs::OpenOptionsExt;
|
||||||
|
opts.mode(0o600);
|
||||||
|
}
|
||||||
|
use std::io::Write;
|
||||||
|
let mut f = opts
|
||||||
|
.open(out)
|
||||||
|
.map_err(|e| format!("cannot create {}: {e}", out.display()))?;
|
||||||
|
writeln!(f, "{}", signing::to_hex(&key.to_bytes()))?;
|
||||||
|
let j = serde_json::json!({
|
||||||
|
"status": "generated",
|
||||||
|
"secret_key_file": out.display().to_string(),
|
||||||
|
"public_key": signing::to_hex(&key.verifying_key().to_bytes()),
|
||||||
|
});
|
||||||
|
println!("{}", serde_json::to_string_pretty(&j)?);
|
||||||
|
return Ok(());
|
||||||
|
}
|
||||||
|
let path = cli
|
||||||
|
.path
|
||||||
|
.clone()
|
||||||
|
.ok_or("--path (or CLAWHDF5_PATH) is required")?;
|
||||||
|
let key = cli
|
||||||
|
.signing_key
|
||||||
|
.as_deref()
|
||||||
|
.map(read_signing_key)
|
||||||
|
.transpose()?;
|
||||||
match cli.command {
|
match cli.command {
|
||||||
Commands::Create { agent_id, dim, wal } => {
|
Commands::Create {
|
||||||
let mut config = MemoryConfig::new(cli.path.clone(), &agent_id, dim);
|
agent_id,
|
||||||
|
dim,
|
||||||
|
wal,
|
||||||
|
f32_index,
|
||||||
|
quantized_index: _,
|
||||||
|
f32,
|
||||||
|
float16: _,
|
||||||
|
} => {
|
||||||
|
let mut config = MemoryConfig::new(path.clone(), &agent_id, dim);
|
||||||
config.wal_enabled = wal;
|
config.wal_enabled = wal;
|
||||||
let mem = HDF5Memory::create(config)?;
|
// As with --f32-index: only ever switch the library default off.
|
||||||
|
if f32 {
|
||||||
|
config.float16 = false;
|
||||||
|
}
|
||||||
|
let config_float16 = config.float16;
|
||||||
|
// Only ever switch *off* the library default: assigning the flag
|
||||||
|
// outright would force every CLI-created store back to f32 unless
|
||||||
|
// the caller knew to ask for int8.
|
||||||
|
if f32_index {
|
||||||
|
config.quantized_index = false;
|
||||||
|
}
|
||||||
|
let config_quantized = config.quantized_index;
|
||||||
|
let mut mem = HDF5Memory::create(config)?;
|
||||||
|
// Sign straight away, so the store is never on disk unsigned.
|
||||||
|
if let Some(k) = &key {
|
||||||
|
mem.set_signing_key(k.clone());
|
||||||
|
mem.flush_wal()?;
|
||||||
|
}
|
||||||
let j = serde_json::json!({
|
let j = serde_json::json!({
|
||||||
"status": "created",
|
"status": "created",
|
||||||
"path": cli.path.display().to_string(),
|
"path": path.display().to_string(),
|
||||||
"agent_id": agent_id,
|
"agent_id": agent_id,
|
||||||
"embedding_dim": dim,
|
"embedding_dim": dim,
|
||||||
"wal_enabled": wal,
|
"wal_enabled": wal,
|
||||||
|
"quantized_index": config_quantized,
|
||||||
|
"float16": config_float16,
|
||||||
|
"signed": mem.is_signed(),
|
||||||
"count": mem.count(),
|
"count": mem.count(),
|
||||||
});
|
});
|
||||||
println!("{}", serde_json::to_string_pretty(&j)?);
|
println!("{}", serde_json::to_string_pretty(&j)?);
|
||||||
@@ -114,7 +228,7 @@ fn run(cli: Cli) -> Result<(), Box<dyn std::error::Error>> {
|
|||||||
}
|
}
|
||||||
};
|
};
|
||||||
let entry: MemoryEntry = serde_json::from_str(&input)?;
|
let entry: MemoryEntry = serde_json::from_str(&input)?;
|
||||||
let mut mem = HDF5Memory::open(&cli.path)?;
|
let mut mem = open_writable(&path, &key)?;
|
||||||
let idx = mem.save(entry)?;
|
let idx = mem.save(entry)?;
|
||||||
let j = serde_json::json!({ "status": "saved", "index": idx, "count": mem.count() });
|
let j = serde_json::json!({ "status": "saved", "index": idx, "count": mem.count() });
|
||||||
println!("{}", serde_json::to_string(&j)?);
|
println!("{}", serde_json::to_string(&j)?);
|
||||||
@@ -128,7 +242,7 @@ fn run(cli: Cli) -> Result<(), Box<dyn std::error::Error>> {
|
|||||||
keyword_weight,
|
keyword_weight,
|
||||||
} => {
|
} => {
|
||||||
let emb: Vec<f32> = serde_json::from_str(&embedding)?;
|
let emb: Vec<f32> = serde_json::from_str(&embedding)?;
|
||||||
let mut mem = HDF5Memory::open(&cli.path)?;
|
let mut mem = open_writable(&path, &key)?;
|
||||||
let results = mem.hybrid_search(&emb, &query, vector_weight, keyword_weight, top_k);
|
let results = mem.hybrid_search(&emb, &query, vector_weight, keyword_weight, top_k);
|
||||||
let j: Vec<serde_json::Value> = results
|
let j: Vec<serde_json::Value> = results
|
||||||
.iter()
|
.iter()
|
||||||
@@ -146,7 +260,7 @@ fn run(cli: Cli) -> Result<(), Box<dyn std::error::Error>> {
|
|||||||
}
|
}
|
||||||
|
|
||||||
Commands::Recall { index } => {
|
Commands::Recall { index } => {
|
||||||
let mem = HDF5Memory::open_read_only(&cli.path)?;
|
let mem = HDF5Memory::open_read_only(&path)?;
|
||||||
match mem.get_chunk(index) {
|
match mem.get_chunk(index) {
|
||||||
Some(content) => {
|
Some(content) => {
|
||||||
let j = serde_json::json!({ "index": index, "chunk": content });
|
let j = serde_json::json!({ "index": index, "chunk": content });
|
||||||
@@ -160,22 +274,23 @@ fn run(cli: Cli) -> Result<(), Box<dyn std::error::Error>> {
|
|||||||
}
|
}
|
||||||
|
|
||||||
Commands::Stats => {
|
Commands::Stats => {
|
||||||
let mem = HDF5Memory::open_read_only(&cli.path)?;
|
let mem = HDF5Memory::open_read_only(&path)?;
|
||||||
let cfg = mem.config();
|
let cfg = mem.config();
|
||||||
let j = serde_json::json!({
|
let j = serde_json::json!({
|
||||||
"path": cli.path.display().to_string(),
|
"path": path.display().to_string(),
|
||||||
"agent_id": cfg.agent_id,
|
"agent_id": cfg.agent_id,
|
||||||
"embedding_dim": cfg.embedding_dim,
|
"embedding_dim": cfg.embedding_dim,
|
||||||
"count": mem.count(),
|
"count": mem.count(),
|
||||||
"active": mem.count_active(),
|
"active": mem.count_active(),
|
||||||
"wal_enabled": cfg.wal_enabled,
|
"wal_enabled": cfg.wal_enabled,
|
||||||
"wal_pending": mem.wal_pending_count(),
|
"wal_pending": mem.wal_pending_count(),
|
||||||
|
"signed": mem.is_signed(),
|
||||||
});
|
});
|
||||||
println!("{}", serde_json::to_string_pretty(&j)?);
|
println!("{}", serde_json::to_string_pretty(&j)?);
|
||||||
}
|
}
|
||||||
|
|
||||||
Commands::FlushWal => {
|
Commands::FlushWal => {
|
||||||
let mut mem = HDF5Memory::open(&cli.path)?;
|
let mut mem = open_writable(&path, &key)?;
|
||||||
let before = mem.wal_pending_count();
|
let before = mem.wal_pending_count();
|
||||||
mem.flush_wal()?;
|
mem.flush_wal()?;
|
||||||
let j = serde_json::json!({
|
let j = serde_json::json!({
|
||||||
@@ -187,7 +302,7 @@ fn run(cli: Cli) -> Result<(), Box<dyn std::error::Error>> {
|
|||||||
}
|
}
|
||||||
|
|
||||||
Commands::AgentsMd { output } => {
|
Commands::AgentsMd { output } => {
|
||||||
let mem = HDF5Memory::open_read_only(&cli.path)?;
|
let mem = HDF5Memory::open_read_only(&path)?;
|
||||||
let md = mem.generate_agents_md();
|
let md = mem.generate_agents_md();
|
||||||
match output {
|
match output {
|
||||||
Some(p) => {
|
Some(p) => {
|
||||||
@@ -199,7 +314,7 @@ fn run(cli: Cli) -> Result<(), Box<dyn std::error::Error>> {
|
|||||||
}
|
}
|
||||||
|
|
||||||
Commands::Export => {
|
Commands::Export => {
|
||||||
let mem = HDF5Memory::open_read_only(&cli.path)?;
|
let mem = HDF5Memory::open_read_only(&path)?;
|
||||||
for i in 0..mem.count() {
|
for i in 0..mem.count() {
|
||||||
if let Some(chunk) = mem.get_chunk(i) {
|
if let Some(chunk) = mem.get_chunk(i) {
|
||||||
let j = serde_json::json!({ "index": i, "chunk": chunk });
|
let j = serde_json::json!({ "index": i, "chunk": chunk });
|
||||||
@@ -208,11 +323,44 @@ fn run(cli: Cli) -> Result<(), Box<dyn std::error::Error>> {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
Commands::Keygen { .. } => unreachable!("handled before opening a store"),
|
||||||
|
|
||||||
|
Commands::Verify { public_key } => {
|
||||||
|
let text = if Path::new(&public_key).is_file() {
|
||||||
|
std::fs::read_to_string(&public_key)?
|
||||||
|
} else {
|
||||||
|
public_key
|
||||||
|
};
|
||||||
|
let bytes = signing::from_hex::<32>(&text)
|
||||||
|
.ok_or("--public-key must be 64 hex characters or a file holding them")?;
|
||||||
|
let trusted = VerifyingKey::from_bytes(&bytes)?;
|
||||||
|
let r = HDF5Memory::verify(&path, &trusted)?;
|
||||||
|
let j = serde_json::json!({
|
||||||
|
"valid": r.is_valid(),
|
||||||
|
"signed": r.signed,
|
||||||
|
"key_matches": r.key_matches,
|
||||||
|
"signature_valid": r.signature_valid,
|
||||||
|
"records_match": r.records_match,
|
||||||
|
"settings_match": r.settings_match,
|
||||||
|
"sessions_match": r.sessions_match,
|
||||||
|
"graph_match": r.graph_match,
|
||||||
|
"changed_records": r.changed_records,
|
||||||
|
"record_count": r.record_count,
|
||||||
|
"signed_record_count": r.signed_record_count,
|
||||||
|
"signed_by": r.public_key.map(|k| signing::to_hex(&k)),
|
||||||
|
"wal_entries_unsigned": r.wal_entries_unsigned,
|
||||||
|
});
|
||||||
|
println!("{}", serde_json::to_string_pretty(&j)?);
|
||||||
|
if !r.is_valid() {
|
||||||
|
std::process::exit(2);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
Commands::Snapshot { dest } => {
|
Commands::Snapshot { dest } => {
|
||||||
let _result = clawhdf5_agent::storage::snapshot_file(&cli.path, &dest)?;
|
let _result = clawhdf5_agent::storage::snapshot_file(&path, &dest)?;
|
||||||
let j = serde_json::json!({
|
let j = serde_json::json!({
|
||||||
"status": "snapshot_created",
|
"status": "snapshot_created",
|
||||||
"source": cli.path.display().to_string(),
|
"source": path.display().to_string(),
|
||||||
"dest": dest.display().to_string(),
|
"dest": dest.display().to_string(),
|
||||||
});
|
});
|
||||||
println!("{}", serde_json::to_string(&j)?);
|
println!("{}", serde_json::to_string(&j)?);
|
||||||
|
|||||||
@@ -1,7 +1,8 @@
|
|||||||
[package]
|
[package]
|
||||||
name = "clawhdf5-derive"
|
name = "clawhdf5-derive"
|
||||||
version = "2.4.0"
|
version = "2.7.0"
|
||||||
edition = "2024"
|
edition = "2024"
|
||||||
|
rust-version.workspace = true
|
||||||
description = "Derive macros for rustyhdf5 HDF5 traits"
|
description = "Derive macros for rustyhdf5 HDF5 traits"
|
||||||
license = "MIT"
|
license = "MIT"
|
||||||
repository = "https://git.redclaw.dev/quantumclaw/clawhdf5"
|
repository = "https://git.redclaw.dev/quantumclaw/clawhdf5"
|
||||||
|
|||||||
@@ -1,7 +1,8 @@
|
|||||||
[package]
|
[package]
|
||||||
name = "clawhdf5-filters"
|
name = "clawhdf5-filters"
|
||||||
version = "2.4.0"
|
version = "2.7.0"
|
||||||
edition = "2024"
|
edition = "2024"
|
||||||
|
rust-version.workspace = true
|
||||||
description = "Filter and compression pipeline for clawhdf5"
|
description = "Filter and compression pipeline for clawhdf5"
|
||||||
license = "MIT"
|
license = "MIT"
|
||||||
repository = "https://git.redclaw.dev/quantumclaw/clawhdf5"
|
repository = "https://git.redclaw.dev/quantumclaw/clawhdf5"
|
||||||
@@ -25,8 +26,12 @@ name = "compression_bench"
|
|||||||
harness = false
|
harness = false
|
||||||
|
|
||||||
[features]
|
[features]
|
||||||
default = ["fast-deflate"]
|
# Pure-Rust zlib-rs by default; `fast-deflate` (zlib-ng, C) overrides it.
|
||||||
|
default = ["zlib-rs"]
|
||||||
fast-deflate = ["flate2/zlib-ng"]
|
fast-deflate = ["flate2/zlib-ng"]
|
||||||
system-zlib = ["flate2/zlib-default"]
|
system-zlib = ["flate2/zlib-default"]
|
||||||
zlib-rs = ["flate2/zlib-rs"]
|
# `runtime_detection` gives zlib-rs `std`, which it needs to detect and use
|
||||||
|
# SIMD at runtime. flate2 enables it by default, but we build flate2 with
|
||||||
|
# default-features = false, and without it zlib-rs inflates 3.5x slower.
|
||||||
|
zlib-rs = ["flate2/zlib-rs", "flate2/runtime_detection"]
|
||||||
apple-compression = []
|
apple-compression = []
|
||||||
|
|||||||
@@ -8,16 +8,18 @@ Filter and compression pipeline for clawhdf5.
|
|||||||
## Features
|
## Features
|
||||||
|
|
||||||
- DEFLATE compression/decompression
|
- DEFLATE compression/decompression
|
||||||
- Fast deflate via zlib-ng (`fast-deflate` feature)
|
- Pure-Rust deflate via zlib-rs (default, `zlib-rs` feature)
|
||||||
|
- zlib-ng instead, if you want it (`fast-deflate` feature; C, needs cmake)
|
||||||
- Apple Compression framework support (`apple-compression` feature)
|
- Apple Compression framework support (`apple-compression` feature)
|
||||||
|
|
||||||
## Usage
|
## Usage
|
||||||
|
|
||||||
```rust
|
```rust
|
||||||
use clawhdf5_filters::{deflate_decode, deflate_encode};
|
use clawhdf5_filters::{deflate_compress, deflate_decompress};
|
||||||
|
|
||||||
let compressed = deflate_encode(&data, 6).unwrap();
|
let compressed = deflate_compress(&data, 6).unwrap();
|
||||||
let decompressed = deflate_decode(&compressed).unwrap();
|
// The second argument bounds the output: the expected decompressed size.
|
||||||
|
let decompressed = deflate_decompress(&compressed, data.len()).unwrap();
|
||||||
```
|
```
|
||||||
|
|
||||||
## License
|
## License
|
||||||
|
|||||||
@@ -1,12 +1,13 @@
|
|||||||
//! Fast deflate backends: Apple Compression Framework and zlib-ng.
|
//! Deflate backends: Apple Compression Framework, zlib-ng and zlib-rs.
|
||||||
//!
|
//!
|
||||||
//! Backend selection priority (decompression & compression):
|
//! Backend selection priority (decompression & compression):
|
||||||
//! 1. Apple Compression Framework (macOS only, `apple-compression` feature)
|
//! 1. Apple Compression Framework (macOS only, `apple-compression` feature)
|
||||||
//! 2. flate2 with zlib-ng backend (`fast-deflate` feature) or miniz_oxide (default)
|
//! 2. flate2 with zlib-ng (`fast-deflate`), else zlib-rs (`zlib-rs`, the
|
||||||
|
//! default), else miniz_oxide
|
||||||
//!
|
//!
|
||||||
//! The Apple Compression Framework uses hardware-accelerated zlib on Apple Silicon
|
//! The Apple Compression Framework uses hardware-accelerated zlib on Apple Silicon
|
||||||
//! and is typically the fastest option on macOS. zlib-ng is the fastest portable
|
//! and is typically the fastest option on macOS. zlib-rs is a pure-Rust port of
|
||||||
//! option and what C HDF5 uses internally.
|
//! zlib-ng; see `BENCHMARKS.md` for how the two compare.
|
||||||
|
|
||||||
// ---------------------------------------------------------------------------
|
// ---------------------------------------------------------------------------
|
||||||
// Apple Compression Framework FFI (macOS only)
|
// Apple Compression Framework FFI (macOS only)
|
||||||
@@ -243,65 +244,117 @@ mod apple {
|
|||||||
}
|
}
|
||||||
|
|
||||||
// ---------------------------------------------------------------------------
|
// ---------------------------------------------------------------------------
|
||||||
// Streaming decompression via flate2 (uses zlib-ng when fast-deflate enabled)
|
// One-shot (de)compression via flate2 (whichever backend flate2 was built with)
|
||||||
|
//
|
||||||
|
// The whole input goes to the codec in one call, into an output buffer sized
|
||||||
|
// up front. `flate2::read::ZlibDecoder` / `write::ZlibEncoder` stream through a
|
||||||
|
// 32 KiB buffer instead, which cost zlib-rs up to 3.7x against zlib-ng on a
|
||||||
|
// 1 MB chunk. clawhdf5-format's deflate filter does the same; see
|
||||||
|
// `BENCHMARKS.md`, "Deflate backend".
|
||||||
// ---------------------------------------------------------------------------
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
/// Streaming decompress with pre-allocated output buffer.
|
/// Decompress into a buffer pre-sized to `output_size`, the expected
|
||||||
///
|
/// decompressed length (known for HDF5 chunks). Output longer than that is an
|
||||||
/// When the output size is known (typical for HDF5 chunks), this avoids
|
/// error, as is a stream that ends early.
|
||||||
/// dynamic reallocation by writing directly into a pre-sized buffer.
|
|
||||||
pub(crate) fn flate2_decompress_preallocated(
|
pub(crate) fn flate2_decompress_preallocated(
|
||||||
data: &[u8],
|
data: &[u8],
|
||||||
output_size: usize,
|
output_size: usize,
|
||||||
) -> Result<Vec<u8>, String> {
|
) -> Result<Vec<u8>, String> {
|
||||||
use std::io::Read;
|
inflate_bounded(data, output_size, output_size)
|
||||||
let mut decoder = flate2::read::ZlibDecoder::new(data);
|
|
||||||
let mut output = vec![0u8; output_size];
|
|
||||||
let mut total_read = 0;
|
|
||||||
|
|
||||||
loop {
|
|
||||||
match decoder.read(&mut output[total_read..]) {
|
|
||||||
Ok(0) => break,
|
|
||||||
Ok(n) => total_read += n,
|
|
||||||
Err(e) => return Err(e.to_string()),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
output.truncate(total_read);
|
|
||||||
Ok(output)
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Absolute ceiling on decompressed output when the caller has no size hint,
|
/// Absolute ceiling on decompressed output when the caller has no size hint,
|
||||||
/// preventing unbounded allocation from a hostile/corrupted zlib stream.
|
/// preventing unbounded allocation from a hostile/corrupted zlib stream.
|
||||||
const MAX_DECOMPRESS_SIZE: usize = 256 * 1024 * 1024;
|
const MAX_DECOMPRESS_SIZE: usize = 256 * 1024 * 1024;
|
||||||
|
|
||||||
/// Streaming decompress with dynamic sizing (when output size is unknown).
|
/// Decompress with no size hint, bounded by [`MAX_DECOMPRESS_SIZE`] so a
|
||||||
///
|
/// hostile zlib stream cannot force arbitrarily large allocation (a "zlib
|
||||||
/// Bounded by [`MAX_DECOMPRESS_SIZE`] since there is no chunk-size hint to
|
/// bomb").
|
||||||
/// 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;
|
let hint = data.len().saturating_mul(4).min(1 << 20);
|
||||||
let decoder = flate2::read::ZlibDecoder::new(data);
|
inflate_bounded(data, hint, MAX_DECOMPRESS_SIZE).map_err(|e| {
|
||||||
let mut result = Vec::new();
|
if e.ends_with("exceeds size limit") {
|
||||||
decoder
|
format!(
|
||||||
.take(MAX_DECOMPRESS_SIZE as u64 + 1)
|
"decompressed output exceeds {} MiB limit",
|
||||||
.read_to_end(&mut result)
|
MAX_DECOMPRESS_SIZE / 1024 / 1024
|
||||||
.map_err(|e| e.to_string())?;
|
)
|
||||||
if result.len() > MAX_DECOMPRESS_SIZE {
|
} else {
|
||||||
return Err(format!(
|
e
|
||||||
"decompressed output exceeds {} MiB limit",
|
}
|
||||||
MAX_DECOMPRESS_SIZE / 1024 / 1024
|
})
|
||||||
));
|
|
||||||
}
|
|
||||||
Ok(result)
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Compress data using flate2 (zlib-ng when fast-deflate enabled, else miniz_oxide).
|
/// Inflate a zlib stream, starting from `size_hint` bytes of output and
|
||||||
|
/// failing past `limit`.
|
||||||
|
fn inflate_bounded(data: &[u8], size_hint: usize, limit: usize) -> Result<Vec<u8>, String> {
|
||||||
|
use flate2::{Decompress, FlushDecompress, Status};
|
||||||
|
|
||||||
|
// One byte of headroom past the limit distinguishes an over-size stream
|
||||||
|
// from one that legitimately ends exactly at the limit.
|
||||||
|
let max_capacity = limit.saturating_add(1);
|
||||||
|
let mut out = Vec::new();
|
||||||
|
out.try_reserve_exact(size_hint.clamp(1, max_capacity))
|
||||||
|
.map_err(|e| format!("deflate: cannot allocate output: {e}"))?;
|
||||||
|
|
||||||
|
let mut inflater = Decompress::new(true);
|
||||||
|
loop {
|
||||||
|
let (in_before, out_before) = (inflater.total_in(), inflater.total_out());
|
||||||
|
let status = inflater
|
||||||
|
.decompress_vec(
|
||||||
|
&data[in_before as usize..],
|
||||||
|
&mut out,
|
||||||
|
FlushDecompress::Finish,
|
||||||
|
)
|
||||||
|
.map_err(|e| format!("deflate: {e}"))?;
|
||||||
|
if out.len() > limit {
|
||||||
|
return Err("deflate: output exceeds size limit".into());
|
||||||
|
}
|
||||||
|
match status {
|
||||||
|
Status::StreamEnd => return Ok(out),
|
||||||
|
Status::Ok | Status::BufError if out.len() == out.capacity() => {
|
||||||
|
let grow = out.capacity().min(max_capacity - out.capacity()).max(1);
|
||||||
|
out.try_reserve_exact(grow)
|
||||||
|
.map_err(|e| format!("deflate: cannot allocate output: {e}"))?;
|
||||||
|
}
|
||||||
|
Status::Ok | Status::BufError => {
|
||||||
|
if inflater.total_in() as usize >= data.len()
|
||||||
|
|| (inflater.total_in(), inflater.total_out()) == (in_before, out_before)
|
||||||
|
{
|
||||||
|
return Err("deflate: truncated stream".into());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Compress data using flate2 (zlib-ng, zlib-rs or miniz_oxide; see module docs).
|
||||||
pub(crate) fn flate2_compress(data: &[u8], level: u32) -> Result<Vec<u8>, String> {
|
pub(crate) fn flate2_compress(data: &[u8], level: u32) -> Result<Vec<u8>, String> {
|
||||||
use std::io::Write;
|
use flate2::{Compress, Compression, FlushCompress, Status};
|
||||||
let mut encoder = flate2::write::ZlibEncoder::new(Vec::new(), flate2::Compression::new(level));
|
|
||||||
encoder.write_all(data).map_err(|e| e.to_string())?;
|
// zlib's compressBound, plus the zlib header and trailer.
|
||||||
encoder.finish().map_err(|e| e.to_string())
|
let bound = data.len() + (data.len() >> 12) + (data.len() >> 14) + (data.len() >> 25) + 13 + 6;
|
||||||
|
let mut out = Vec::new();
|
||||||
|
out.try_reserve_exact(bound)
|
||||||
|
.map_err(|e| format!("deflate: cannot allocate output: {e}"))?;
|
||||||
|
|
||||||
|
let mut deflater = Compress::new(Compression::new(level), true);
|
||||||
|
loop {
|
||||||
|
let (in_before, out_before) = (deflater.total_in(), deflater.total_out());
|
||||||
|
let status = deflater
|
||||||
|
.compress_vec(&data[in_before as usize..], &mut out, FlushCompress::Finish)
|
||||||
|
.map_err(|e| format!("deflate: {e}"))?;
|
||||||
|
match status {
|
||||||
|
Status::StreamEnd => return Ok(out),
|
||||||
|
Status::Ok | Status::BufError if out.len() == out.capacity() => out
|
||||||
|
.try_reserve(out.capacity().max(4096))
|
||||||
|
.map_err(|e| format!("deflate: cannot allocate output: {e}"))?,
|
||||||
|
Status::Ok | Status::BufError => {
|
||||||
|
if (deflater.total_in(), deflater.total_out()) == (in_before, out_before) {
|
||||||
|
return Err("deflate: encoder made no progress".into());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// ---------------------------------------------------------------------------
|
// ---------------------------------------------------------------------------
|
||||||
@@ -312,7 +365,7 @@ pub(crate) fn flate2_compress(data: &[u8], level: u32) -> Result<Vec<u8>, String
|
|||||||
///
|
///
|
||||||
/// Selection order:
|
/// Selection order:
|
||||||
/// 1. Apple Compression Framework (macOS + `apple-compression` feature)
|
/// 1. Apple Compression Framework (macOS + `apple-compression` feature)
|
||||||
/// 2. flate2 (zlib-ng with `fast-deflate`, otherwise miniz_oxide)
|
/// 2. flate2 (zlib-ng with `fast-deflate`, else zlib-rs, else miniz_oxide)
|
||||||
///
|
///
|
||||||
/// When `output_hint` > 0, pre-allocates the output buffer for zero-copy
|
/// When `output_hint` > 0, pre-allocates the output buffer for zero-copy
|
||||||
/// decompression (avoids reallocation).
|
/// decompression (avoids reallocation).
|
||||||
@@ -344,7 +397,7 @@ pub fn decompress(data: &[u8], output_hint: usize) -> Result<Vec<u8>, String> {
|
|||||||
///
|
///
|
||||||
/// Selection order:
|
/// Selection order:
|
||||||
/// 1. Apple Compression Framework (macOS + `apple-compression` feature)
|
/// 1. Apple Compression Framework (macOS + `apple-compression` feature)
|
||||||
/// 2. flate2 (zlib-ng with `fast-deflate`, otherwise miniz_oxide)
|
/// 2. flate2 (zlib-ng with `fast-deflate`, else zlib-rs, else miniz_oxide)
|
||||||
pub fn compress(data: &[u8], level: u32) -> Result<Vec<u8>, String> {
|
pub fn compress(data: &[u8], level: u32) -> Result<Vec<u8>, String> {
|
||||||
#[cfg(all(target_os = "macos", feature = "apple-compression"))]
|
#[cfg(all(target_os = "macos", feature = "apple-compression"))]
|
||||||
{
|
{
|
||||||
@@ -377,9 +430,19 @@ pub fn active_backend() -> &'static str {
|
|||||||
{
|
{
|
||||||
"zlib-ng"
|
"zlib-ng"
|
||||||
}
|
}
|
||||||
|
// flate2 prefers a C zlib over zlib-rs when both are enabled.
|
||||||
|
#[cfg(all(
|
||||||
|
not(all(target_os = "macos", feature = "apple-compression")),
|
||||||
|
not(feature = "fast-deflate"),
|
||||||
|
feature = "zlib-rs"
|
||||||
|
))]
|
||||||
|
{
|
||||||
|
"zlib-rs"
|
||||||
|
}
|
||||||
#[cfg(not(any(
|
#[cfg(not(any(
|
||||||
all(target_os = "macos", feature = "apple-compression"),
|
all(target_os = "macos", feature = "apple-compression"),
|
||||||
feature = "fast-deflate"
|
feature = "fast-deflate",
|
||||||
|
feature = "zlib-rs"
|
||||||
)))]
|
)))]
|
||||||
{
|
{
|
||||||
"miniz_oxide"
|
"miniz_oxide"
|
||||||
@@ -436,7 +499,7 @@ mod tests {
|
|||||||
fn backend_name_is_set() {
|
fn backend_name_is_set() {
|
||||||
let name = active_backend();
|
let name = active_backend();
|
||||||
assert!(
|
assert!(
|
||||||
["miniz_oxide", "zlib-ng", "apple-compression"].contains(&name),
|
["miniz_oxide", "zlib-rs", "zlib-ng", "apple-compression"].contains(&name),
|
||||||
"unexpected backend: {name}"
|
"unexpected backend: {name}"
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -2,12 +2,14 @@
|
|||||||
//!
|
//!
|
||||||
//! Provides deflate (zlib) decompression/compression with multiple backend options:
|
//! Provides deflate (zlib) decompression/compression with multiple backend options:
|
||||||
//!
|
//!
|
||||||
//! - **Default**: `miniz_oxide` (pure Rust, no C dependencies)
|
//! - **Default (`zlib-rs` feature)**: `zlib-rs` via flate2 (pure Rust, no C
|
||||||
//! - **`fast-deflate` feature**: `zlib-ng` via flate2 (~2-3x faster, matches C HDF5)
|
//! dependencies)
|
||||||
|
//! - **`fast-deflate` feature**: `zlib-ng` via flate2 (C, built with cmake)
|
||||||
//! - **`apple-compression` feature**: Apple Compression Framework on macOS
|
//! - **`apple-compression` feature**: Apple Compression Framework on macOS
|
||||||
//! (hardware-accelerated on Apple Silicon)
|
//! (hardware-accelerated on Apple Silicon)
|
||||||
|
//! - With none of the above: `miniz_oxide` (pure Rust, slower)
|
||||||
//!
|
//!
|
||||||
//! Backend priority: apple-compression > zlib-ng > miniz_oxide.
|
//! Backend priority: apple-compression > zlib-ng > zlib-rs > miniz_oxide.
|
||||||
|
|
||||||
pub mod fast_deflate;
|
pub mod fast_deflate;
|
||||||
|
|
||||||
@@ -115,7 +117,7 @@ mod tests {
|
|||||||
fn backend_reports_name() {
|
fn backend_reports_name() {
|
||||||
let name = deflate_backend();
|
let name = deflate_backend();
|
||||||
assert!(
|
assert!(
|
||||||
["miniz_oxide", "zlib-ng", "apple-compression"].contains(&name),
|
["miniz_oxide", "zlib-rs", "zlib-ng", "apple-compression"].contains(&name),
|
||||||
"unexpected backend: {name}"
|
"unexpected backend: {name}"
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -1,7 +1,8 @@
|
|||||||
[package]
|
[package]
|
||||||
name = "clawhdf5-format"
|
name = "clawhdf5-format"
|
||||||
version = "2.4.0"
|
version = "2.7.0"
|
||||||
edition = "2024"
|
edition = "2024"
|
||||||
|
rust-version.workspace = true
|
||||||
description = "Pure-Rust HDF5 binary format parsing and writing — no C dependencies"
|
description = "Pure-Rust HDF5 binary format parsing and writing — no C dependencies"
|
||||||
license = "MIT"
|
license = "MIT"
|
||||||
repository = "https://git.redclaw.dev/quantumclaw/clawhdf5"
|
repository = "https://git.redclaw.dev/quantumclaw/clawhdf5"
|
||||||
@@ -23,16 +24,20 @@ libaec-sys = { path = "../libaec-sys", version = "0.1", optional = true }
|
|||||||
pco = { version = "1.0", optional = true }
|
pco = { version = "1.0", optional = true }
|
||||||
|
|
||||||
[dev-dependencies]
|
[dev-dependencies]
|
||||||
|
half = { workspace = true }
|
||||||
serde_json = "1"
|
serde_json = "1"
|
||||||
criterion = { workspace = true }
|
criterion = { workspace = true }
|
||||||
clawhdf5-derive = { path = "../clawhdf5-derive", version = "2.4.0" }
|
clawhdf5-derive = { path = "../clawhdf5-derive", version = "2.7.0" }
|
||||||
|
|
||||||
[[bench]]
|
[[bench]]
|
||||||
name = "bench"
|
name = "bench"
|
||||||
harness = false
|
harness = false
|
||||||
|
|
||||||
[features]
|
[features]
|
||||||
default = ["std", "checksum", "deflate", "provenance", "fast-deflate", "system-zlib-decompress"]
|
# Deflate backend: `zlib-rs` (pure Rust) by default. `fast-deflate` selects
|
||||||
|
# zlib-ng instead (C, built with cmake); flate2 prefers a C zlib whenever one
|
||||||
|
# is enabled, so turning it on anywhere in the build overrides the default.
|
||||||
|
default = ["std", "checksum", "deflate", "provenance", "zlib-rs", "system-zlib-decompress"]
|
||||||
std = []
|
std = []
|
||||||
checksum = []
|
checksum = []
|
||||||
deflate = ["flate2"]
|
deflate = ["flate2"]
|
||||||
@@ -42,7 +47,10 @@ fast-checksum = ["crc32fast"]
|
|||||||
fast-deflate = ["flate2/zlib-ng"]
|
fast-deflate = ["flate2/zlib-ng"]
|
||||||
system-zlib = ["flate2/zlib-default"]
|
system-zlib = ["flate2/zlib-default"]
|
||||||
system-zlib-decompress = []
|
system-zlib-decompress = []
|
||||||
zlib-rs = ["flate2/zlib-rs"]
|
# `runtime_detection` gives zlib-rs `std`, which it needs to detect and use
|
||||||
|
# SIMD at runtime. flate2 enables it by default, but we build flate2 with
|
||||||
|
# default-features = false, and without it zlib-rs inflates 3.5x slower.
|
||||||
|
zlib-rs = ["flate2/zlib-rs", "flate2/runtime_detection"]
|
||||||
lz4 = ["lz4_flex"]
|
lz4 = ["lz4_flex"]
|
||||||
zstd = ["dep:zstd"]
|
zstd = ["dep:zstd"]
|
||||||
blake3_hash = ["blake3"]
|
blake3_hash = ["blake3"]
|
||||||
|
|||||||
@@ -1 +1,4 @@
|
|||||||
target/
|
target/
|
||||||
|
corpus/
|
||||||
|
artifacts/
|
||||||
|
coverage/
|
||||||
|
|||||||
Binary file not shown.
@@ -1,15 +1,36 @@
|
|||||||
#![no_main]
|
#![no_main]
|
||||||
|
use clawhdf5_format::btree_v2::{BTreeV2Header, collect_btree_v2_records};
|
||||||
use libfuzzer_sys::fuzz_target;
|
use libfuzzer_sys::fuzz_target;
|
||||||
|
|
||||||
fuzz_target!(|data: &[u8]| {
|
fuzz_target!(|data: &[u8]| {
|
||||||
for &offset_size in &[4u8, 8] {
|
for &offset_size in &[4u8, 8] {
|
||||||
for &length_size in &[4u8, 8] {
|
for &length_size in &[4u8, 8] {
|
||||||
let _ = clawhdf5_format::btree_v2::BTreeV2Header::parse(
|
if let Ok(header) = BTreeV2Header::parse(data, 0, offset_size, length_size) {
|
||||||
data,
|
let _ = collect_btree_v2_records(data, &header, offset_size, length_size);
|
||||||
0,
|
}
|
||||||
offset_size,
|
|
||||||
length_size,
|
|
||||||
);
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Parsing a header requires a valid checksum, which random input almost
|
||||||
|
// never has, so the traversal behind it went unfuzzed — and that is where
|
||||||
|
// a node listing itself as its own child overflowed the stack. Take the
|
||||||
|
// header fields straight from the input instead and walk the rest.
|
||||||
|
let Some((fields, file)) = data.split_first_chunk::<20>() else {
|
||||||
|
return;
|
||||||
|
};
|
||||||
|
let header = BTreeV2Header {
|
||||||
|
tree_type: fields[0],
|
||||||
|
node_size: u32::from_le_bytes([fields[1], fields[2], fields[3], fields[4]]),
|
||||||
|
record_size: u16::from_le_bytes([fields[5], fields[6]]),
|
||||||
|
depth: u16::from_le_bytes([fields[7], fields[8]]),
|
||||||
|
root_node_address: u64::from(u32::from_le_bytes([
|
||||||
|
fields[9], fields[10], fields[11], fields[12],
|
||||||
|
])),
|
||||||
|
num_records_in_root: u16::from_le_bytes([fields[13], fields[14]]),
|
||||||
|
total_records: u64::from(u32::from_le_bytes([
|
||||||
|
fields[15], fields[16], fields[17], fields[18],
|
||||||
|
])),
|
||||||
|
};
|
||||||
|
let offset_size = if fields[19] & 1 == 0 { 4 } else { 8 };
|
||||||
|
let _ = collect_btree_v2_records(file, &header, offset_size, 8);
|
||||||
});
|
});
|
||||||
|
|||||||
@@ -172,6 +172,17 @@ fn max_records_leaf(node_size: u32, record_size: u16) -> u64 {
|
|||||||
((node_size - overhead) / record_size as u32) as u64
|
((node_size - overhead) / record_size as u32) as u64
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Deepest B-tree v2 accepted. See [`collect_btree_v2_records`].
|
||||||
|
const MAX_DEPTH: u16 = 64;
|
||||||
|
|
||||||
|
/// Take `n` records from the traversal's budget, or refuse the tree.
|
||||||
|
fn spend(budget: &mut usize, n: usize) -> Result<(), FormatError> {
|
||||||
|
*budget = budget
|
||||||
|
.checked_sub(n)
|
||||||
|
.ok_or(FormatError::NestingDepthExceeded)?;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
/// Collect all records from a B-tree v2 by traversing from the root.
|
/// Collect all records from a B-tree v2 by traversing from the root.
|
||||||
pub fn collect_btree_v2_records(
|
pub fn collect_btree_v2_records(
|
||||||
file_data: &[u8],
|
file_data: &[u8],
|
||||||
@@ -182,6 +193,22 @@ pub fn collect_btree_v2_records(
|
|||||||
if header.total_records == 0 || header.num_records_in_root == 0 {
|
if header.total_records == 0 || header.num_records_in_root == 0 {
|
||||||
return Ok(Vec::new());
|
return Ok(Vec::new());
|
||||||
}
|
}
|
||||||
|
// Recursion is one frame per level, and the depth is read from the file:
|
||||||
|
// a crafted header claiming 65 535 levels over a node that is its own
|
||||||
|
// child overflowed the stack. 64 matches the fractal heap's guard, and no
|
||||||
|
// real tree comes close — even at the minimum fan-out of two it would
|
||||||
|
// hold more than 2^64 records.
|
||||||
|
if header.depth > MAX_DEPTH {
|
||||||
|
return Err(FormatError::NestingDepthExceeded);
|
||||||
|
}
|
||||||
|
// A valid tree stores each record once, in its own bytes, so it cannot
|
||||||
|
// hold more records than the file has room for. Children are addresses,
|
||||||
|
// though, and nothing makes them distinct: levels whose children all
|
||||||
|
// point at one shared node below reach it fan-out^depth times, which is
|
||||||
|
// millions of records from a few kilobytes. Counting against what the
|
||||||
|
// file could physically contain bounds that without trusting the
|
||||||
|
// header's own `total_records`.
|
||||||
|
let mut budget = file_data.len() / usize::from(header.record_size.max(1));
|
||||||
|
|
||||||
let max_leaf_nrec = max_records_leaf(header.node_size, header.record_size);
|
let max_leaf_nrec = max_records_leaf(header.node_size, header.record_size);
|
||||||
|
|
||||||
@@ -206,6 +233,7 @@ pub fn collect_btree_v2_records(
|
|||||||
offset_size,
|
offset_size,
|
||||||
length_size,
|
length_size,
|
||||||
max_leaf_nrec,
|
max_leaf_nrec,
|
||||||
|
&mut budget,
|
||||||
&mut records,
|
&mut records,
|
||||||
)?;
|
)?;
|
||||||
Ok(records)
|
Ok(records)
|
||||||
@@ -273,6 +301,7 @@ fn collect_internal_records(
|
|||||||
offset_size: u8,
|
offset_size: u8,
|
||||||
length_size: u8,
|
length_size: u8,
|
||||||
max_leaf_nrec: u64,
|
max_leaf_nrec: u64,
|
||||||
|
budget: &mut usize,
|
||||||
out: &mut Vec<BTreeV2Record>,
|
out: &mut Vec<BTreeV2Record>,
|
||||||
) -> Result<(), FormatError> {
|
) -> Result<(), FormatError> {
|
||||||
// signature(4) + version(1) + type(1) = 6
|
// signature(4) + version(1) + type(1) = 6
|
||||||
@@ -350,6 +379,8 @@ fn collect_internal_records(
|
|||||||
// We collect child[0] records, then record[0], then child[1], etc.
|
// We collect child[0] records, then record[0], then child[1], etc.
|
||||||
for (i, &(child_addr, child_nrec)) in children.iter().enumerate() {
|
for (i, &(child_addr, child_nrec)) in children.iter().enumerate() {
|
||||||
if child_depth == 0 {
|
if child_depth == 0 {
|
||||||
|
// Before parsing, so a refused tree is not also a large allocation.
|
||||||
|
spend(budget, usize::from(child_nrec))?;
|
||||||
let leaf_recs =
|
let leaf_recs =
|
||||||
parse_leaf_records(file_data, child_addr as usize, child_nrec, record_size)?;
|
parse_leaf_records(file_data, child_addr as usize, child_nrec, record_size)?;
|
||||||
out.extend(leaf_recs);
|
out.extend(leaf_recs);
|
||||||
@@ -364,6 +395,7 @@ fn collect_internal_records(
|
|||||||
offset_size,
|
offset_size,
|
||||||
length_size,
|
length_size,
|
||||||
max_leaf_nrec,
|
max_leaf_nrec,
|
||||||
|
budget,
|
||||||
out,
|
out,
|
||||||
)?;
|
)?;
|
||||||
}
|
}
|
||||||
@@ -393,6 +425,7 @@ fn collect_internal_records(
|
|||||||
available: file_data.len(),
|
available: file_data.len(),
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
|
spend(budget, 1)?;
|
||||||
out.push(BTreeV2Record {
|
out.push(BTreeV2Record {
|
||||||
data: file_data[rec_start..rec_end].to_vec(),
|
data: file_data[rec_start..rec_end].to_vec(),
|
||||||
});
|
});
|
||||||
@@ -466,6 +499,124 @@ mod tests {
|
|||||||
buf
|
buf
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// An internal node laid out exactly as `collect_internal_records` will
|
||||||
|
/// read it at `depth`: `records` zeroed records, then `children` pointers,
|
||||||
|
/// all to `child_addr` claiming `child_nrec` records.
|
||||||
|
fn internal_node(
|
||||||
|
depth: u16,
|
||||||
|
node_size: u32,
|
||||||
|
record_size: u16,
|
||||||
|
records: usize,
|
||||||
|
children: usize,
|
||||||
|
child_addr: u64,
|
||||||
|
child_nrec: u64,
|
||||||
|
) -> Vec<u8> {
|
||||||
|
let max_leaf = max_records_leaf(node_size, record_size);
|
||||||
|
let nrec_width = bytes_for_max_records(if depth == 1 { max_leaf } else { max_leaf * 2 });
|
||||||
|
let total_width = if depth > 1 {
|
||||||
|
bytes_for_max_records(header_max_total_records(max_leaf, depth - 1))
|
||||||
|
} else {
|
||||||
|
0
|
||||||
|
};
|
||||||
|
let mut buf = b"BTIN".to_vec();
|
||||||
|
buf.extend_from_slice(&[0, 5]);
|
||||||
|
buf.resize(buf.len() + records * record_size as usize, 0);
|
||||||
|
for _ in 0..children {
|
||||||
|
buf.extend_from_slice(&child_addr.to_le_bytes());
|
||||||
|
buf.extend_from_slice(&child_nrec.to_le_bytes()[..nrec_width]);
|
||||||
|
buf.resize(buf.len() + total_width, 0);
|
||||||
|
}
|
||||||
|
buf
|
||||||
|
}
|
||||||
|
|
||||||
|
fn header(depth: u16, root: u64, root_nrec: u16, total: u64) -> BTreeV2Header {
|
||||||
|
BTreeV2Header {
|
||||||
|
tree_type: 5,
|
||||||
|
node_size: 512,
|
||||||
|
record_size: 8,
|
||||||
|
depth,
|
||||||
|
root_node_address: root,
|
||||||
|
num_records_in_root: root_nrec,
|
||||||
|
total_records: total,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn a_node_that_is_its_own_child_is_rejected_not_recursed() {
|
||||||
|
// One internal node whose two children are itself, under a header
|
||||||
|
// claiming the deepest tree a u16 allows. The layout stops depending
|
||||||
|
// on depth once the subtree-total width saturates, so every level
|
||||||
|
// parses cleanly and recursion runs ~65 000 frames deep: before the
|
||||||
|
// cap this overflowed the stack and aborted the process, from a file
|
||||||
|
// of under 100 bytes.
|
||||||
|
let mut data = internal_node(u16::MAX, 512, 8, 1, 2, 0, 1);
|
||||||
|
data.resize(4096, 0);
|
||||||
|
let result = collect_btree_v2_records(&data, &header(u16::MAX, 0, 1, 1), 8, 8);
|
||||||
|
assert!(result.is_err(), "{result:?}");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn a_shared_subtree_cannot_multiply_the_work() {
|
||||||
|
// A chain of distinct levels, each node's children all pointing at the
|
||||||
|
// single node below, ending in a real leaf. Every node parses and
|
||||||
|
// nothing is cyclic, yet the leaf is reached fan-out^depth times: 62
|
||||||
|
// children over 4 levels is ~15 million leaf visits from a few
|
||||||
|
// kilobytes. A valid tree cannot hold more records than the file has
|
||||||
|
// room for, so that bounds the traversal instead.
|
||||||
|
let (node_size, record_size) = (512u32, 8u16);
|
||||||
|
let fanout = 62usize;
|
||||||
|
let depth = 4u16;
|
||||||
|
let leaf = build_leaf_node(5, &[&[0u8; 8][..]]);
|
||||||
|
|
||||||
|
// Lay out root first, then each lower level, then the leaf.
|
||||||
|
let mut nodes: Vec<Vec<u8>> = Vec::new();
|
||||||
|
let mut addrs = Vec::new();
|
||||||
|
let mut at = 0u64;
|
||||||
|
let mut sizes = Vec::new();
|
||||||
|
for d in (1..=depth).rev() {
|
||||||
|
let n = internal_node(d, node_size, record_size, fanout - 1, fanout, 0, 0);
|
||||||
|
sizes.push(n.len());
|
||||||
|
}
|
||||||
|
for size in &sizes {
|
||||||
|
addrs.push(at);
|
||||||
|
at += *size as u64;
|
||||||
|
}
|
||||||
|
let leaf_addr = at;
|
||||||
|
for (i, d) in (1..=depth).rev().enumerate() {
|
||||||
|
let (child, child_nrec) = if d == 1 {
|
||||||
|
(leaf_addr, 1)
|
||||||
|
} else {
|
||||||
|
(addrs[i + 1], fanout as u64 - 1)
|
||||||
|
};
|
||||||
|
nodes.push(internal_node(
|
||||||
|
d,
|
||||||
|
node_size,
|
||||||
|
record_size,
|
||||||
|
fanout - 1,
|
||||||
|
fanout,
|
||||||
|
child,
|
||||||
|
child_nrec,
|
||||||
|
));
|
||||||
|
}
|
||||||
|
let mut data: Vec<u8> = nodes.concat();
|
||||||
|
data.extend_from_slice(&leaf);
|
||||||
|
data.resize(data.len() + 64, 0);
|
||||||
|
|
||||||
|
let started = std::time::Instant::now();
|
||||||
|
let result =
|
||||||
|
collect_btree_v2_records(&data, &header(depth, 0, fanout as u16 - 1, u64::MAX), 8, 8);
|
||||||
|
assert!(
|
||||||
|
result.is_err(),
|
||||||
|
"expected a refusal, got {} records",
|
||||||
|
result.map_or(0, |r| r.len())
|
||||||
|
);
|
||||||
|
assert!(
|
||||||
|
started.elapsed() < std::time::Duration::from_secs(2),
|
||||||
|
"took {:?}",
|
||||||
|
started.elapsed()
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn parse_header() {
|
fn parse_header() {
|
||||||
let data = build_btree_v2_header(5, 512, 11, 0, 0x1000, 3, 3, 8, 8);
|
let data = build_btree_v2_header(5, 512, 11, 0, 0x1000, 3, 3, 8, 8);
|
||||||
|
|||||||
@@ -223,13 +223,32 @@ pub const DEFAULT_CACHE_BYTES: usize = 16 * 1024 * 1024; // 16 MiB
|
|||||||
/// coordinate map and reduces collision chains compared to power-of-two sizes.
|
/// coordinate map and reduces collision chains compared to power-of-two sizes.
|
||||||
pub const DEFAULT_MAX_SLOTS: usize = 521;
|
pub const DEFAULT_MAX_SLOTS: usize = 521;
|
||||||
|
|
||||||
|
/// Most datasets whose chunk index a [`ChunkCache`] keeps at once.
|
||||||
|
pub const MAX_INDEXED_DATASETS: usize = 64;
|
||||||
|
|
||||||
|
/// Most chunk-index entries, summed over all datasets, a [`ChunkCache`] keeps.
|
||||||
|
/// Least-recently-used datasets' indexes are dropped past this (the dataset
|
||||||
|
/// being read is always kept), so a file with many or huge chunked datasets
|
||||||
|
/// cannot grow the cache without bound.
|
||||||
|
pub const MAX_INDEXED_CHUNKS: usize = 1 << 20;
|
||||||
|
|
||||||
|
/// The dataset key the address-less (legacy) methods use when
|
||||||
|
/// [`ChunkCache::ensure_dataset`] has not been called.
|
||||||
|
#[cfg(feature = "std")]
|
||||||
|
const UNBOUND_DATASET: u64 = u64::MAX;
|
||||||
|
|
||||||
// ---------------------------------------------------------------------------
|
// ---------------------------------------------------------------------------
|
||||||
// LRU entry
|
// LRU entry
|
||||||
// ---------------------------------------------------------------------------
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
/// Decompressed chunks are keyed by dataset *and* coordinate: every chunked
|
||||||
|
/// dataset has a chunk at (0, 0, ...), so the coordinate alone is ambiguous.
|
||||||
|
#[cfg(feature = "std")]
|
||||||
|
type SlotKey = (u64, ChunkCoord);
|
||||||
|
|
||||||
#[cfg(feature = "std")]
|
#[cfg(feature = "std")]
|
||||||
struct CachedChunk {
|
struct CachedChunk {
|
||||||
coord: ChunkCoord,
|
key: SlotKey,
|
||||||
/// Shared so a cache hit is a refcount bump, not a copy of the whole
|
/// Shared so a cache hit is a refcount bump, not a copy of the whole
|
||||||
/// (potentially large) decompressed chunk.
|
/// (potentially large) decompressed chunk.
|
||||||
data: Arc<CacheAlignedBuffer>,
|
data: Arc<CacheAlignedBuffer>,
|
||||||
@@ -237,21 +256,48 @@ struct CachedChunk {
|
|||||||
last_access: u64,
|
last_access: u64,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Per-dataset index state.
|
||||||
|
#[cfg(feature = "std")]
|
||||||
|
#[derive(Default)]
|
||||||
|
struct DatasetEntry {
|
||||||
|
/// Chunk coordinate -> ChunkInfo (offset + size in file).
|
||||||
|
index: Option<Arc<HashMap<ChunkCoord, ChunkInfo>>>,
|
||||||
|
/// Pre-built chunk index for O(1) coordinate lookups.
|
||||||
|
chunk_index: Option<Arc<ChunkIndex>>,
|
||||||
|
/// Pre-computed chunk layout for fast assembly.
|
||||||
|
chunk_layout: Option<Arc<ChunkLayout>>,
|
||||||
|
/// Tick of the last use, for dropping the least recently used dataset.
|
||||||
|
last_used: u64,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(feature = "std")]
|
||||||
|
impl DatasetEntry {
|
||||||
|
fn weight(&self) -> usize {
|
||||||
|
self.index.as_ref().map_or(0, |m| m.len())
|
||||||
|
+ self.chunk_index.as_ref().map_or(0, |c| c.num_chunks())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
// ---------------------------------------------------------------------------
|
// ---------------------------------------------------------------------------
|
||||||
// ChunkCache
|
// ChunkCache
|
||||||
// ---------------------------------------------------------------------------
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
/// A per-dataset chunk cache with hash-based index and LRU eviction.
|
/// A per-file chunk cache: chunk indexes per dataset, plus an LRU of
|
||||||
|
/// decompressed chunks, all keyed by dataset.
|
||||||
///
|
///
|
||||||
/// # Usage
|
/// A dataset is identified by the address of its chunk index (B-tree, fixed
|
||||||
|
/// or extensible array, ...), which is unique within a file. Every method
|
||||||
|
/// that takes an `addr` works on that dataset only, so threads reading
|
||||||
|
/// different datasets through one shared cache never see each other's
|
||||||
|
/// chunks. The address-less methods (`has_index`, `populate_index`,
|
||||||
|
/// `get_decompressed`, ...) act on the dataset last bound with
|
||||||
|
/// [`Self::ensure_dataset`]; that binding is shared state, so concurrent
|
||||||
|
/// readers must use the `*_in` / `*_for` methods instead (the chunked
|
||||||
|
/// readers in [`crate::chunked_read`] do).
|
||||||
///
|
///
|
||||||
/// ```ignore
|
/// Memory is bounded: decompressed data by `max_bytes`/`max_slots` across
|
||||||
/// let cache = ChunkCache::new();
|
/// all datasets, indexes by [`MAX_INDEXED_DATASETS`] and
|
||||||
/// // Pass &cache to read_chunked_data — it will populate the index lazily.
|
/// [`MAX_INDEXED_CHUNKS`].
|
||||||
/// ```
|
|
||||||
///
|
|
||||||
/// The cache is wrapped in `Mutex` internally so it can be mutated through
|
|
||||||
/// shared references (thread-safe).
|
|
||||||
///
|
///
|
||||||
/// Only available with the `std` feature because it requires `std::sync::Mutex`.
|
/// Only available with the `std` feature because it requires `std::sync::Mutex`.
|
||||||
#[cfg(feature = "std")]
|
#[cfg(feature = "std")]
|
||||||
@@ -261,26 +307,20 @@ pub struct ChunkCache {
|
|||||||
|
|
||||||
#[cfg(feature = "std")]
|
#[cfg(feature = "std")]
|
||||||
struct CacheInner {
|
struct CacheInner {
|
||||||
/// Hash index: chunk coordinate -> ChunkInfo (offset + size in file).
|
/// Per-dataset chunk indexes, keyed by chunk-index address.
|
||||||
/// Populated once per dataset on first access.
|
datasets: HashMap<u64, DatasetEntry>,
|
||||||
index: Option<HashMap<ChunkCoord, ChunkInfo>>,
|
|
||||||
|
|
||||||
/// Address of the dataset (its chunk-index base address) that the cached
|
/// Dataset the address-less methods act on (see `ensure_dataset`).
|
||||||
/// index, chunk index, layout, and decompressed slots currently belong to.
|
current: Option<u64>,
|
||||||
/// The cache is shared per file across datasets, so every cached-read entry
|
|
||||||
/// checks this and resets the per-dataset state when the dataset changes —
|
|
||||||
/// otherwise one dataset's chunk index (with its own rank) would be reused
|
|
||||||
/// for another, corrupting reads.
|
|
||||||
index_addr: Option<u64>,
|
|
||||||
|
|
||||||
/// 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
|
/// Key -> index into `slots`, for O(1) lookup instead of a linear
|
||||||
/// scan. Kept in sync with `slots` on every insert/evict/clear — in
|
/// scan. Kept in sync with `slots` on every insert/evict/clear — in
|
||||||
/// particular, `slots.swap_remove(i)` moves the last element into slot
|
/// particular, `slots.swap_remove(i)` moves the last element into slot
|
||||||
/// `i`, so the moved element's index entry must be updated too.
|
/// `i`, so the moved element's index entry must be updated too.
|
||||||
slot_index: HashMap<ChunkCoord, usize>,
|
slot_index: HashMap<SlotKey, usize>,
|
||||||
|
|
||||||
/// Current total bytes of cached decompressed data.
|
/// Current total bytes of cached decompressed data.
|
||||||
current_bytes: usize,
|
current_bytes: usize,
|
||||||
@@ -294,17 +334,145 @@ struct CacheInner {
|
|||||||
/// Monotonic counter for LRU ordering.
|
/// Monotonic counter for LRU ordering.
|
||||||
tick: u64,
|
tick: u64,
|
||||||
|
|
||||||
/// Last accessed chunk coordinate (for sequential detection).
|
/// Last accessed chunk (for sequential detection).
|
||||||
last_coord: Option<ChunkCoord>,
|
last_coord: Option<SlotKey>,
|
||||||
|
|
||||||
/// Access pattern statistics.
|
/// Access pattern statistics.
|
||||||
stats: AccessStats,
|
stats: AccessStats,
|
||||||
|
}
|
||||||
|
|
||||||
/// Pre-built chunk index for O(1) coordinate lookups.
|
#[cfg(feature = "std")]
|
||||||
chunk_index: Option<ChunkIndex>,
|
impl CacheInner {
|
||||||
|
fn current(&self) -> u64 {
|
||||||
|
self.current.unwrap_or(UNBOUND_DATASET)
|
||||||
|
}
|
||||||
|
|
||||||
/// Pre-computed chunk layout for fast assembly.
|
fn touch(&mut self, addr: u64) -> &mut DatasetEntry {
|
||||||
chunk_layout: Option<ChunkLayout>,
|
self.tick += 1;
|
||||||
|
let tick = self.tick;
|
||||||
|
let entry = self.datasets.entry(addr).or_default();
|
||||||
|
entry.last_used = tick;
|
||||||
|
entry
|
||||||
|
}
|
||||||
|
|
||||||
|
fn entry(&self, addr: u64) -> Option<&DatasetEntry> {
|
||||||
|
self.datasets.get(&addr)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Drop least-recently-used datasets' indexes (never `keep`'s) until the
|
||||||
|
/// dataset and chunk-entry budgets hold.
|
||||||
|
fn trim_datasets(&mut self, keep: u64) {
|
||||||
|
loop {
|
||||||
|
let total: usize = self.datasets.values().map(DatasetEntry::weight).sum();
|
||||||
|
if self.datasets.len() <= MAX_INDEXED_DATASETS && total <= MAX_INDEXED_CHUNKS {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
let victim = self
|
||||||
|
.datasets
|
||||||
|
.iter()
|
||||||
|
.filter(|(a, _)| **a != keep)
|
||||||
|
.min_by_key(|(_, e)| e.last_used)
|
||||||
|
.map(|(a, _)| *a);
|
||||||
|
match victim {
|
||||||
|
Some(a) => {
|
||||||
|
self.datasets.remove(&a);
|
||||||
|
}
|
||||||
|
None => return,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn get_decompressed(&mut self, addr: u64, coord: &[u64]) -> Option<Arc<CacheAlignedBuffer>> {
|
||||||
|
self.tick += 1;
|
||||||
|
let tick = self.tick;
|
||||||
|
|
||||||
|
// Track sequential vs random access
|
||||||
|
let is_sequential = self.last_coord.as_ref().is_some_and(|(prev_addr, prev)| {
|
||||||
|
// Sequential if exactly one dimension changed
|
||||||
|
let changes: usize = prev
|
||||||
|
.iter()
|
||||||
|
.zip(coord.iter())
|
||||||
|
.filter(|(a, b)| a != b)
|
||||||
|
.count();
|
||||||
|
*prev_addr == addr && changes <= 1
|
||||||
|
});
|
||||||
|
if is_sequential {
|
||||||
|
self.stats.sequential_count += 1;
|
||||||
|
} else if self.last_coord.is_some() {
|
||||||
|
self.stats.random_count += 1;
|
||||||
|
}
|
||||||
|
let key: SlotKey = (addr, coord.to_vec());
|
||||||
|
let found = if let Some(&idx) = self.slot_index.get(&key) {
|
||||||
|
self.slots[idx].last_access = tick;
|
||||||
|
Some(Arc::clone(&self.slots[idx].data))
|
||||||
|
} else {
|
||||||
|
None
|
||||||
|
};
|
||||||
|
self.last_coord = Some(key);
|
||||||
|
if let Some(ref data) = found {
|
||||||
|
self.stats.hits += 1;
|
||||||
|
self.stats.bytes_read += data.len() as u64;
|
||||||
|
} else {
|
||||||
|
self.stats.misses += 1;
|
||||||
|
}
|
||||||
|
found
|
||||||
|
}
|
||||||
|
|
||||||
|
fn put_decompressed(
|
||||||
|
&mut self,
|
||||||
|
key: SlotKey,
|
||||||
|
data: Arc<CacheAlignedBuffer>,
|
||||||
|
) -> Arc<CacheAlignedBuffer> {
|
||||||
|
let data_len = data.len();
|
||||||
|
|
||||||
|
// Don't cache if single chunk exceeds budget — still return the data
|
||||||
|
// to the caller, just don't retain it.
|
||||||
|
if data_len > self.max_bytes {
|
||||||
|
return data;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Check if already present
|
||||||
|
self.tick += 1;
|
||||||
|
let tick = self.tick;
|
||||||
|
if let Some(&idx) = self.slot_index.get(&key) {
|
||||||
|
self.slots[idx].last_access = tick;
|
||||||
|
return Arc::clone(&self.slots[idx].data); // already cached
|
||||||
|
}
|
||||||
|
|
||||||
|
// Evict until we have room
|
||||||
|
while self.slots.len() >= self.max_slots
|
||||||
|
|| (self.current_bytes + data_len > self.max_bytes && !self.slots.is_empty())
|
||||||
|
{
|
||||||
|
// Find LRU slot
|
||||||
|
let lru_idx = self
|
||||||
|
.slots
|
||||||
|
.iter()
|
||||||
|
.enumerate()
|
||||||
|
.min_by_key(|(_, s)| s.last_access)
|
||||||
|
.map(|(i, _)| i)
|
||||||
|
.unwrap();
|
||||||
|
let removed = self.slots.swap_remove(lru_idx);
|
||||||
|
self.slot_index.remove(&removed.key);
|
||||||
|
// 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 < self.slots.len() {
|
||||||
|
let moved_key = self.slots[lru_idx].key.clone();
|
||||||
|
self.slot_index.insert(moved_key, lru_idx);
|
||||||
|
}
|
||||||
|
self.current_bytes -= removed.data.len();
|
||||||
|
self.stats.evictions += 1;
|
||||||
|
}
|
||||||
|
|
||||||
|
self.current_bytes += data_len;
|
||||||
|
let new_idx = self.slots.len();
|
||||||
|
self.slot_index.insert(key.clone(), new_idx);
|
||||||
|
self.slots.push(CachedChunk {
|
||||||
|
key,
|
||||||
|
data: Arc::clone(&data),
|
||||||
|
last_access: tick,
|
||||||
|
});
|
||||||
|
data
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Access pattern statistics tracked by the chunk cache.
|
/// Access pattern statistics tracked by the chunk cache.
|
||||||
@@ -356,8 +524,8 @@ impl ChunkCache {
|
|||||||
pub fn with_capacity(max_bytes: usize, max_slots: usize) -> Self {
|
pub fn with_capacity(max_bytes: usize, max_slots: usize) -> Self {
|
||||||
Self {
|
Self {
|
||||||
inner: std::sync::Mutex::new(CacheInner {
|
inner: std::sync::Mutex::new(CacheInner {
|
||||||
index: None,
|
datasets: HashMap::new(),
|
||||||
index_addr: None,
|
current: 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)),
|
slot_index: HashMap::with_capacity(max_slots.min(64)),
|
||||||
current_bytes: 0,
|
current_bytes: 0,
|
||||||
@@ -366,335 +534,331 @@ impl ChunkCache {
|
|||||||
tick: 0,
|
tick: 0,
|
||||||
last_coord: None,
|
last_coord: None,
|
||||||
stats: AccessStats::default(),
|
stats: AccessStats::default(),
|
||||||
chunk_index: None,
|
|
||||||
chunk_layout: None,
|
|
||||||
}),
|
}),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// ----- Index operations -----
|
fn lock(&self) -> std::sync::MutexGuard<'_, CacheInner> {
|
||||||
|
self.inner.lock().unwrap_or_else(|e| e.into_inner())
|
||||||
|
}
|
||||||
|
|
||||||
/// Bind the cache to the dataset at chunk-index address `addr`.
|
/// The most decompressed bytes this cache will hold.
|
||||||
|
pub fn max_bytes(&self) -> usize {
|
||||||
|
self.lock().max_bytes
|
||||||
|
}
|
||||||
|
|
||||||
|
// ----- Dataset-keyed operations (safe to use concurrently) -----
|
||||||
|
|
||||||
|
/// The chunk list of the dataset whose chunk index is at `addr`.
|
||||||
///
|
///
|
||||||
/// The cache is shared per file across all of its datasets. If the cache
|
/// On the first call for a dataset, `build` scans its chunk index; the
|
||||||
/// currently holds state for a different dataset, all per-dataset state
|
/// result is kept (offsets truncated to `rank` for the lookup key), so
|
||||||
/// (chunk index, chunk-index map, layout, and decompressed slots) is
|
/// later calls skip the scan. `build` runs without the cache lock held;
|
||||||
/// dropped so the next access rebuilds it for this dataset. Reading the
|
/// if two threads race to build the same dataset's index, the first
|
||||||
/// same dataset again is a no-op, preserving the cache's benefit for
|
/// stored one wins and both return equivalent lists.
|
||||||
/// repeated/sequential access. Returns `true` if a reset occurred.
|
pub fn chunks_for<E>(
|
||||||
pub fn ensure_dataset(&self, addr: u64) -> bool {
|
&self,
|
||||||
let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
|
addr: u64,
|
||||||
if inner.index_addr == Some(addr) {
|
rank: usize,
|
||||||
return false;
|
build: impl FnOnce() -> Result<Vec<ChunkInfo>, E>,
|
||||||
|
) -> Result<Vec<ChunkInfo>, E> {
|
||||||
|
Ok(self
|
||||||
|
.index_for(addr, rank, build)?
|
||||||
|
.values()
|
||||||
|
.cloned()
|
||||||
|
.collect())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn index_for<E>(
|
||||||
|
&self,
|
||||||
|
addr: u64,
|
||||||
|
rank: usize,
|
||||||
|
build: impl FnOnce() -> Result<Vec<ChunkInfo>, E>,
|
||||||
|
) -> Result<Arc<HashMap<ChunkCoord, ChunkInfo>>, E> {
|
||||||
|
if let Some(index) = self.lock().touch(addr).index.clone() {
|
||||||
|
return Ok(index);
|
||||||
}
|
}
|
||||||
inner.index = None;
|
let chunks = build()?;
|
||||||
inner.chunk_index = None;
|
let map: HashMap<ChunkCoord, ChunkInfo> = chunks
|
||||||
inner.chunk_layout = None;
|
.into_iter()
|
||||||
inner.slots.clear();
|
.map(|ci| (ci.offsets.iter().take(rank).copied().collect(), ci))
|
||||||
inner.slot_index.clear();
|
.collect();
|
||||||
inner.current_bytes = 0;
|
let mut inner = self.lock();
|
||||||
inner.last_coord = None;
|
let entry = inner.touch(addr);
|
||||||
inner.index_addr = Some(addr);
|
let index = Arc::clone(entry.index.get_or_insert_with(|| Arc::new(map)));
|
||||||
true
|
inner.trim_datasets(addr);
|
||||||
|
Ok(index)
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Returns `true` if the chunk index has been built.
|
/// The pre-computed assembly layout of the dataset at `addr`, building
|
||||||
|
/// its chunk index (via `build`, as in [`Self::chunks_for`]) and layout on
|
||||||
|
/// first use.
|
||||||
|
pub fn chunk_layout_for<E>(
|
||||||
|
&self,
|
||||||
|
addr: u64,
|
||||||
|
rank: usize,
|
||||||
|
build: impl FnOnce() -> Result<Vec<ChunkInfo>, E>,
|
||||||
|
ds_dims: &[usize],
|
||||||
|
chunk_dims: &[usize],
|
||||||
|
elem_size: usize,
|
||||||
|
) -> Result<Arc<ChunkLayout>, E> {
|
||||||
|
let (layout, chunk_index) = {
|
||||||
|
let mut inner = self.lock();
|
||||||
|
let entry = inner.touch(addr);
|
||||||
|
(entry.chunk_layout.clone(), entry.chunk_index.clone())
|
||||||
|
};
|
||||||
|
if let Some(layout) = layout {
|
||||||
|
return Ok(layout);
|
||||||
|
}
|
||||||
|
let chunk_index = match chunk_index {
|
||||||
|
Some(ci) => ci,
|
||||||
|
None => {
|
||||||
|
let index = self.index_for(addr, rank, build)?;
|
||||||
|
let chunks: Vec<ChunkInfo> = index.values().cloned().collect();
|
||||||
|
Arc::new(ChunkIndex::build(&chunks, rank))
|
||||||
|
}
|
||||||
|
};
|
||||||
|
let layout = ChunkLayout::build(&chunk_index, ds_dims, chunk_dims, elem_size);
|
||||||
|
let mut inner = self.lock();
|
||||||
|
let entry = inner.touch(addr);
|
||||||
|
entry.chunk_index.get_or_insert(chunk_index);
|
||||||
|
let layout = Arc::clone(entry.chunk_layout.get_or_insert_with(|| Arc::new(layout)));
|
||||||
|
inner.trim_datasets(addr);
|
||||||
|
Ok(layout)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Cached decompressed chunk at `coord` of the dataset at `addr`.
|
||||||
|
///
|
||||||
|
/// O(1) lookup; the clone is an `Arc` refcount bump, not a copy of the
|
||||||
|
/// underlying decompressed data.
|
||||||
|
pub fn get_decompressed_in(&self, addr: u64, coord: &[u64]) -> Option<Arc<CacheAlignedBuffer>> {
|
||||||
|
self.lock().get_decompressed(addr, coord)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Cache decompressed chunk data for `coord` of the dataset at `addr`.
|
||||||
|
/// Returns the `Arc`-shared buffer now cached (or already cached).
|
||||||
|
pub fn put_decompressed_in(
|
||||||
|
&self,
|
||||||
|
addr: u64,
|
||||||
|
coord: ChunkCoord,
|
||||||
|
data: Vec<u8>,
|
||||||
|
) -> Arc<CacheAlignedBuffer> {
|
||||||
|
self.put_decompressed_aligned_in(addr, coord, CacheAlignedBuffer::from_vec(data))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// [`Self::put_decompressed_in`] for an already-aligned buffer.
|
||||||
|
pub fn put_decompressed_aligned_in(
|
||||||
|
&self,
|
||||||
|
addr: u64,
|
||||||
|
coord: ChunkCoord,
|
||||||
|
data: CacheAlignedBuffer,
|
||||||
|
) -> Arc<CacheAlignedBuffer> {
|
||||||
|
let data = Arc::new(data);
|
||||||
|
self.lock().put_decompressed((addr, coord), data)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Record that the given chunk coordinates of the dataset at `addr` are
|
||||||
|
/// predicted to be accessed soon (bookkeeping only).
|
||||||
|
///
|
||||||
|
/// This does **not** prefetch or pre-decompress anything — it only
|
||||||
|
/// checks whether each coordinate is already in the chunk index and
|
||||||
|
/// updates access-pattern stats accordingly.
|
||||||
|
pub fn prefetch_hint_in(&self, addr: u64, next_coords: &[ChunkCoord]) {
|
||||||
|
let mut inner = self.lock();
|
||||||
|
let Some(index) = inner.entry(addr).and_then(|e| e.index.clone()) else {
|
||||||
|
return;
|
||||||
|
};
|
||||||
|
let known = next_coords
|
||||||
|
.iter()
|
||||||
|
.filter(|c| index.contains_key(*c))
|
||||||
|
.count();
|
||||||
|
inner.stats.sequential_count += known as u64;
|
||||||
|
}
|
||||||
|
|
||||||
|
// ----- Address-less operations on the bound dataset -----
|
||||||
|
|
||||||
|
/// Bind the address-less methods to the dataset at chunk-index address
|
||||||
|
/// `addr`. Returns `true` if this changed the bound dataset.
|
||||||
|
///
|
||||||
|
/// Each dataset's state is kept separately, so switching loses nothing
|
||||||
|
/// and never exposes one dataset's index or chunks to another. The
|
||||||
|
/// binding itself is shared, though: concurrent readers should use the
|
||||||
|
/// `addr`-taking methods rather than bind and then call these.
|
||||||
|
pub fn ensure_dataset(&self, addr: u64) -> bool {
|
||||||
|
let mut inner = self.lock();
|
||||||
|
let changed = inner.current != Some(addr);
|
||||||
|
inner.current = Some(addr);
|
||||||
|
changed
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Returns `true` if the bound dataset's chunk index has been built.
|
||||||
pub fn has_index(&self) -> bool {
|
pub fn has_index(&self) -> bool {
|
||||||
self.inner
|
let inner = self.lock();
|
||||||
.lock()
|
inner
|
||||||
.unwrap_or_else(|e| e.into_inner())
|
.entry(inner.current())
|
||||||
.index
|
.is_some_and(|e| e.index.is_some())
|
||||||
.is_some()
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Build the chunk index from a pre-collected list of `ChunkInfo`.
|
/// Build the bound dataset's chunk index from a pre-collected list of
|
||||||
|
/// `ChunkInfo`.
|
||||||
///
|
///
|
||||||
/// The `rank` parameter is used to truncate offsets to spatial dims only
|
/// The `rank` parameter is used to truncate offsets to spatial dims only
|
||||||
/// (B-tree v1 stores rank+1 offsets).
|
/// (B-tree v1 stores rank+1 offsets).
|
||||||
pub fn populate_index(&self, chunks: &[ChunkInfo], rank: usize) {
|
pub fn populate_index(&self, chunks: &[ChunkInfo], rank: usize) {
|
||||||
let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
|
let addr = self.lock().current();
|
||||||
if inner.index.is_some() {
|
let _ = self.index_for::<core::convert::Infallible>(addr, rank, || Ok(chunks.to_vec()));
|
||||||
return; // already populated
|
|
||||||
}
|
|
||||||
let mut map = HashMap::with_capacity(chunks.len());
|
|
||||||
|
|
||||||
for ci in chunks {
|
|
||||||
let coord: ChunkCoord = ci.offsets.iter().take(rank).copied().collect();
|
|
||||||
map.insert(coord, ci.clone());
|
|
||||||
}
|
|
||||||
inner.index = Some(map);
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Look up a chunk by its spatial coordinate in the index.
|
/// Look up a chunk by its spatial coordinate in the bound dataset's index.
|
||||||
pub fn lookup_index(&self, coord: &[u64]) -> Option<ChunkInfo> {
|
pub fn lookup_index(&self, coord: &[u64]) -> Option<ChunkInfo> {
|
||||||
let inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
|
let inner = self.lock();
|
||||||
inner.index.as_ref()?.get(coord).cloned()
|
inner
|
||||||
|
.entry(inner.current())?
|
||||||
|
.index
|
||||||
|
.as_ref()?
|
||||||
|
.get(coord)
|
||||||
|
.cloned()
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Return all indexed chunks as a `Vec<ChunkInfo>` (order unspecified).
|
/// Return all of the bound dataset's indexed chunks (order unspecified).
|
||||||
pub fn all_indexed_chunks(&self) -> Option<Vec<ChunkInfo>> {
|
pub fn all_indexed_chunks(&self) -> Option<Vec<ChunkInfo>> {
|
||||||
let inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
|
let inner = self.lock();
|
||||||
inner.index.as_ref().map(|m| m.values().cloned().collect())
|
let index = inner.entry(inner.current())?.index.as_ref()?;
|
||||||
|
Some(index.values().cloned().collect())
|
||||||
}
|
}
|
||||||
|
|
||||||
// ----- Chunk index (pre-built coordinate → ChunkInfo map) -----
|
/// Returns `true` if the bound dataset's `ChunkIndex` has been built.
|
||||||
|
|
||||||
/// Returns `true` if the chunk B-tree index has been built.
|
|
||||||
pub fn has_chunk_index(&self) -> bool {
|
pub fn has_chunk_index(&self) -> bool {
|
||||||
self.inner
|
let inner = self.lock();
|
||||||
.lock()
|
inner
|
||||||
.unwrap_or_else(|e| e.into_inner())
|
.entry(inner.current())
|
||||||
.chunk_index
|
.is_some_and(|e| e.chunk_index.is_some())
|
||||||
.is_some()
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Build and store the chunk B-tree index from a pre-collected list of `ChunkInfo`.
|
/// Build and store the bound dataset's `ChunkIndex`.
|
||||||
pub fn populate_chunk_index(&self, chunks: &[ChunkInfo], rank: usize) {
|
pub fn populate_chunk_index(&self, chunks: &[ChunkInfo], rank: usize) {
|
||||||
let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
|
let built = Arc::new(ChunkIndex::build(chunks, rank));
|
||||||
if inner.chunk_index.is_some() {
|
let mut inner = self.lock();
|
||||||
return;
|
let addr = inner.current();
|
||||||
}
|
inner.touch(addr).chunk_index.get_or_insert(built);
|
||||||
inner.chunk_index = Some(ChunkIndex::build(chunks, rank));
|
inner.trim_datasets(addr);
|
||||||
}
|
}
|
||||||
|
|
||||||
// ----- Chunk layout (pre-computed assembly plan) -----
|
/// Returns `true` if the bound dataset's chunk layout has been computed.
|
||||||
|
|
||||||
/// Returns `true` if the chunk layout has been computed.
|
|
||||||
pub fn has_chunk_layout(&self) -> bool {
|
pub fn has_chunk_layout(&self) -> bool {
|
||||||
self.inner
|
let inner = self.lock();
|
||||||
.lock()
|
inner
|
||||||
.unwrap_or_else(|e| e.into_inner())
|
.entry(inner.current())
|
||||||
.chunk_layout
|
.is_some_and(|e| e.chunk_layout.is_some())
|
||||||
.is_some()
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Build and store the pre-computed chunk layout for fast assembly.
|
/// Build and store the bound dataset's chunk layout (needs its
|
||||||
|
/// `ChunkIndex`; does nothing without one).
|
||||||
pub fn populate_chunk_layout(&self, ds_dims: &[usize], chunk_dims: &[usize], elem_size: usize) {
|
pub fn populate_chunk_layout(&self, ds_dims: &[usize], chunk_dims: &[usize], elem_size: usize) {
|
||||||
let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
|
let mut inner = self.lock();
|
||||||
if inner.chunk_layout.is_some() {
|
let addr = inner.current();
|
||||||
|
let entry = inner.touch(addr);
|
||||||
|
if entry.chunk_layout.is_some() {
|
||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
if let Some(ref idx) = inner.chunk_index {
|
if let Some(idx) = entry.chunk_index.clone() {
|
||||||
inner.chunk_layout = Some(ChunkLayout::build(idx, ds_dims, chunk_dims, elem_size));
|
entry.chunk_layout = Some(Arc::new(ChunkLayout::build(
|
||||||
|
&idx, ds_dims, chunk_dims, elem_size,
|
||||||
|
)));
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Execute a function with a reference to the chunk layout.
|
/// Execute a function with a reference to the bound dataset's chunk
|
||||||
///
|
/// layout. Returns `None` if the layout hasn't been computed yet.
|
||||||
/// Returns `None` if the layout hasn't been computed yet.
|
|
||||||
pub fn with_chunk_layout<F, R>(&self, f: F) -> Option<R>
|
pub fn with_chunk_layout<F, R>(&self, f: F) -> Option<R>
|
||||||
where
|
where
|
||||||
F: FnOnce(&ChunkLayout) -> R,
|
F: FnOnce(&ChunkLayout) -> R,
|
||||||
{
|
{
|
||||||
let inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
|
let layout = {
|
||||||
inner.chunk_layout.as_ref().map(f)
|
let inner = self.lock();
|
||||||
|
inner.entry(inner.current())?.chunk_layout.clone()?
|
||||||
|
};
|
||||||
|
Some(f(&layout))
|
||||||
}
|
}
|
||||||
|
|
||||||
// ----- Decompressed data cache (LRU) -----
|
/// Try to get cached decompressed data for a chunk of the bound dataset.
|
||||||
|
|
||||||
/// Try to get cached decompressed data for a chunk coordinate.
|
|
||||||
///
|
///
|
||||||
/// O(1) lookup. Returns an owned copy for API compatibility with callers
|
/// Returns an owned copy; prefer [`Self::get_decompressed_aligned`] when
|
||||||
/// that need a `Vec<u8>`; prefer [`Self::get_decompressed_aligned`] when
|
/// an `Arc`-shared buffer works for the caller.
|
||||||
/// 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)
|
self.get_decompressed_aligned(coord)
|
||||||
.map(|arc| arc.as_slice().to_vec())
|
.map(|arc| arc.as_slice().to_vec())
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Try to get a reference-counted clone of the aligned buffer for a chunk.
|
/// Reference-counted cached buffer for a chunk of the bound dataset.
|
||||||
///
|
|
||||||
/// 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>> {
|
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.lock();
|
||||||
inner.tick += 1;
|
let addr = inner.current();
|
||||||
let tick = inner.tick;
|
inner.get_decompressed(addr, coord)
|
||||||
|
|
||||||
// Track sequential vs random access
|
|
||||||
let is_sequential = inner.last_coord.as_ref().is_some_and(|prev| {
|
|
||||||
// Sequential if exactly one dimension changed
|
|
||||||
let changes: usize = prev
|
|
||||||
.iter()
|
|
||||||
.zip(coord.iter())
|
|
||||||
.filter(|(a, b)| a != b)
|
|
||||||
.count();
|
|
||||||
changes <= 1
|
|
||||||
});
|
|
||||||
if is_sequential {
|
|
||||||
inner.stats.sequential_count += 1;
|
|
||||||
} else if inner.last_coord.is_some() {
|
|
||||||
inner.stats.random_count += 1;
|
|
||||||
}
|
|
||||||
inner.last_coord = Some(coord.to_vec());
|
|
||||||
|
|
||||||
let found = if let Some(&idx) = inner.slot_index.get(coord) {
|
|
||||||
inner.slots[idx].last_access = tick;
|
|
||||||
Some(Arc::clone(&inner.slots[idx].data))
|
|
||||||
} else {
|
|
||||||
None
|
|
||||||
};
|
|
||||||
if let Some(ref data) = found {
|
|
||||||
inner.stats.hits += 1;
|
|
||||||
inner.stats.bytes_read += data.len() as u64;
|
|
||||||
} else {
|
|
||||||
inner.stats.misses += 1;
|
|
||||||
}
|
|
||||||
found
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Insert decompressed chunk data into the LRU cache.
|
/// Insert decompressed chunk data for the bound dataset into the LRU
|
||||||
///
|
/// cache, returning the `Arc`-shared buffer now cached.
|
||||||
/// The data is stored in a [`CacheAlignedBuffer`] so subsequent reads
|
|
||||||
/// return cache-line-aligned memory. Returns the `Arc`-shared buffer that
|
|
||||||
/// is now cached (or already was), so the caller can reuse it directly
|
|
||||||
/// instead of holding a separate copy of the same data.
|
|
||||||
pub fn put_decompressed(&self, coord: ChunkCoord, data: Vec<u8>) -> Arc<CacheAlignedBuffer> {
|
pub fn put_decompressed(&self, coord: ChunkCoord, data: Vec<u8>) -> Arc<CacheAlignedBuffer> {
|
||||||
let aligned = CacheAlignedBuffer::from_vec(data);
|
self.put_decompressed_aligned(coord, CacheAlignedBuffer::from_vec(data))
|
||||||
self.put_decompressed_aligned(coord, aligned)
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Insert an already-aligned buffer into the LRU cache.
|
/// Insert an already-aligned buffer for the bound dataset.
|
||||||
///
|
|
||||||
/// 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(
|
pub fn put_decompressed_aligned(
|
||||||
&self,
|
&self,
|
||||||
coord: ChunkCoord,
|
coord: ChunkCoord,
|
||||||
data: CacheAlignedBuffer,
|
data: CacheAlignedBuffer,
|
||||||
) -> Arc<CacheAlignedBuffer> {
|
) -> Arc<CacheAlignedBuffer> {
|
||||||
let data = Arc::new(data);
|
let data = Arc::new(data);
|
||||||
let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
|
let mut inner = self.lock();
|
||||||
let data_len = data.len();
|
let addr = inner.current();
|
||||||
|
inner.put_decompressed((addr, coord), data)
|
||||||
// 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 {
|
|
||||||
return data;
|
|
||||||
}
|
|
||||||
|
|
||||||
// Check if already present
|
|
||||||
inner.tick += 1;
|
|
||||||
let tick = inner.tick;
|
|
||||||
if let Some(&idx) = inner.slot_index.get(&coord) {
|
|
||||||
inner.slots[idx].last_access = tick;
|
|
||||||
return Arc::clone(&inner.slots[idx].data); // already cached
|
|
||||||
}
|
|
||||||
|
|
||||||
// Evict until we have room
|
|
||||||
while inner.slots.len() >= inner.max_slots
|
|
||||||
|| (inner.current_bytes + data_len > inner.max_bytes && !inner.slots.is_empty())
|
|
||||||
{
|
|
||||||
// Find LRU slot
|
|
||||||
let lru_idx = inner
|
|
||||||
.slots
|
|
||||||
.iter()
|
|
||||||
.enumerate()
|
|
||||||
.min_by_key(|(_, s)| s.last_access)
|
|
||||||
.map(|(i, _)| i)
|
|
||||||
.unwrap();
|
|
||||||
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.stats.evictions += 1;
|
|
||||||
}
|
|
||||||
|
|
||||||
inner.current_bytes += data_len;
|
|
||||||
let new_idx = inner.slots.len();
|
|
||||||
inner.slot_index.insert(coord.clone(), new_idx);
|
|
||||||
inner.slots.push(CachedChunk {
|
|
||||||
coord,
|
|
||||||
data: Arc::clone(&data),
|
|
||||||
last_access: tick,
|
|
||||||
});
|
|
||||||
data
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Clear the entire cache (index + decompressed data).
|
/// [`Self::prefetch_hint_in`] for the bound dataset.
|
||||||
|
pub fn prefetch_hint(&self, next_coords: &[ChunkCoord]) {
|
||||||
|
let addr = self.lock().current();
|
||||||
|
self.prefetch_hint_in(addr, next_coords);
|
||||||
|
}
|
||||||
|
|
||||||
|
// ----- Whole-cache operations -----
|
||||||
|
|
||||||
|
/// Clear the entire cache (indexes + decompressed data + stats).
|
||||||
pub fn clear(&self) {
|
pub fn clear(&self) {
|
||||||
let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
|
let mut inner = self.lock();
|
||||||
inner.index = None;
|
inner.datasets.clear();
|
||||||
inner.index_addr = None;
|
inner.current = None;
|
||||||
inner.slots.clear();
|
inner.slots.clear();
|
||||||
inner.slot_index.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;
|
||||||
inner.stats = AccessStats::default();
|
inner.stats = AccessStats::default();
|
||||||
inner.chunk_index = None;
|
|
||||||
inner.chunk_layout = None;
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Record that the given chunk coordinates are predicted to be accessed
|
|
||||||
/// soon (bookkeeping only).
|
|
||||||
///
|
|
||||||
/// This does **not** prefetch or pre-decompress anything — it only
|
|
||||||
/// checks whether each coordinate is already in the chunk index and
|
|
||||||
/// updates access-pattern stats accordingly. Real prefetching (e.g.
|
|
||||||
/// background pre-decompression) is not implemented.
|
|
||||||
pub fn prefetch_hint(&self, next_coords: &[ChunkCoord]) {
|
|
||||||
let inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
|
|
||||||
if inner.index.is_none() {
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
drop(inner);
|
|
||||||
// For each predicted coordinate, verify it exists in the index.
|
|
||||||
// The index is already populated, so this is a no-op for known chunks.
|
|
||||||
// The purpose is to signal intent — callers can pre-decompress if needed.
|
|
||||||
// We touch the stats to record that prefetch hints were issued.
|
|
||||||
let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
|
|
||||||
for coord in next_coords {
|
|
||||||
let exists = inner
|
|
||||||
.index
|
|
||||||
.as_ref()
|
|
||||||
.map(|idx| idx.contains_key(coord))
|
|
||||||
.unwrap_or(false);
|
|
||||||
if exists {
|
|
||||||
inner.stats.sequential_count += 1;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Return the current access pattern statistics.
|
/// Return the current access pattern statistics.
|
||||||
pub fn access_stats(&self) -> AccessStats {
|
pub fn access_stats(&self) -> AccessStats {
|
||||||
self.inner
|
self.lock().stats.clone()
|
||||||
.lock()
|
|
||||||
.unwrap_or_else(|e| e.into_inner())
|
|
||||||
.stats
|
|
||||||
.clone()
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Update the sweep direction label in the access stats.
|
/// Update the sweep direction label in the access stats.
|
||||||
pub fn set_sweep_direction(&self, direction: &'static str) {
|
pub fn set_sweep_direction(&self, direction: &'static str) {
|
||||||
self.inner
|
self.lock().stats.sweep_direction = Some(direction);
|
||||||
.lock()
|
|
||||||
.unwrap_or_else(|e| e.into_inner())
|
|
||||||
.stats
|
|
||||||
.sweep_direction = Some(direction);
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Number of decompressed chunks currently cached.
|
/// Number of decompressed chunks currently cached (all datasets).
|
||||||
pub fn cached_chunk_count(&self) -> usize {
|
pub fn cached_chunk_count(&self) -> usize {
|
||||||
self.inner
|
self.lock().slots.len()
|
||||||
.lock()
|
|
||||||
.unwrap_or_else(|e| e.into_inner())
|
|
||||||
.slots
|
|
||||||
.len()
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Total bytes of decompressed data currently cached.
|
/// Total bytes of decompressed data currently cached (all datasets).
|
||||||
pub fn cached_bytes(&self) -> usize {
|
pub fn cached_bytes(&self) -> usize {
|
||||||
self.inner
|
self.lock().current_bytes
|
||||||
.lock()
|
}
|
||||||
.unwrap_or_else(|e| e.into_inner())
|
|
||||||
.current_bytes
|
/// Number of datasets whose chunk index is currently kept.
|
||||||
|
pub fn indexed_dataset_count(&self) -> usize {
|
||||||
|
self.lock().datasets.len()
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -803,6 +967,92 @@ mod tests {
|
|||||||
assert_eq!(cache.cached_bytes(), 0);
|
assert_eq!(cache.cached_bytes(), 0);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn datasets_sharing_coordinates_stay_separate() {
|
||||||
|
let cache = ChunkCache::new();
|
||||||
|
let a = vec![make_chunk(vec![0, 0], 0x100, 8)];
|
||||||
|
let b = vec![make_chunk(vec![0, 0], 0x900, 8)];
|
||||||
|
let got_a = cache.chunks_for::<()>(1, 1, || Ok(a.clone())).unwrap();
|
||||||
|
let got_b = cache.chunks_for::<()>(2, 1, || Ok(b.clone())).unwrap();
|
||||||
|
assert_eq!(got_a[0].address, 0x100);
|
||||||
|
assert_eq!(got_b[0].address, 0x900);
|
||||||
|
// Built once per dataset: a second lookup doesn't call the builder.
|
||||||
|
let again = cache
|
||||||
|
.chunks_for::<()>(1, 1, || panic!("index rebuilt"))
|
||||||
|
.unwrap();
|
||||||
|
assert_eq!(again[0].address, 0x100);
|
||||||
|
|
||||||
|
cache.put_decompressed_in(1, vec![0], vec![1; 4]);
|
||||||
|
cache.put_decompressed_in(2, vec![0], vec![2; 4]);
|
||||||
|
assert_eq!(
|
||||||
|
cache.get_decompressed_in(1, &[0]).unwrap().as_slice(),
|
||||||
|
&[1; 4]
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
cache.get_decompressed_in(2, &[0]).unwrap().as_slice(),
|
||||||
|
&[2; 4]
|
||||||
|
);
|
||||||
|
assert!(cache.get_decompressed_in(3, &[0]).is_none());
|
||||||
|
assert_eq!(cache.cached_chunk_count(), 2);
|
||||||
|
|
||||||
|
// The bound-dataset methods see only the bound dataset.
|
||||||
|
cache.ensure_dataset(2);
|
||||||
|
assert_eq!(cache.lookup_index(&[0]).unwrap().address, 0x900);
|
||||||
|
assert_eq!(cache.get_decompressed(&[0]).unwrap(), vec![2; 4]);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn dataset_indexes_are_bounded() {
|
||||||
|
let cache = ChunkCache::new();
|
||||||
|
for addr in 0..(MAX_INDEXED_DATASETS as u64 + 10) {
|
||||||
|
cache
|
||||||
|
.chunks_for::<()>(addr, 1, || Ok(vec![make_chunk(vec![0], addr, 8)]))
|
||||||
|
.unwrap();
|
||||||
|
}
|
||||||
|
assert_eq!(cache.indexed_dataset_count(), MAX_INDEXED_DATASETS);
|
||||||
|
|
||||||
|
// One huge index evicts the others but is itself kept.
|
||||||
|
let huge: Vec<ChunkInfo> = (0..MAX_INDEXED_CHUNKS as u64)
|
||||||
|
.map(|i| make_chunk(vec![i], i, 8))
|
||||||
|
.collect();
|
||||||
|
let got = cache.chunks_for::<()>(9999, 1, || Ok(huge)).unwrap();
|
||||||
|
assert_eq!(got.len(), MAX_INDEXED_CHUNKS);
|
||||||
|
assert_eq!(cache.indexed_dataset_count(), 1);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn concurrent_readers_of_different_datasets_see_their_own_chunks() {
|
||||||
|
let cache = std::sync::Arc::new(ChunkCache::with_capacity(1 << 20, 64));
|
||||||
|
let handles: Vec<_> = (0..8u64)
|
||||||
|
.map(|t| {
|
||||||
|
let cache = std::sync::Arc::clone(&cache);
|
||||||
|
std::thread::spawn(move || {
|
||||||
|
for round in 0..500u64 {
|
||||||
|
let addr = (t + round) % 16;
|
||||||
|
let coord = vec![round % 4];
|
||||||
|
let chunks = cache
|
||||||
|
.chunks_for::<()>(addr, 1, || {
|
||||||
|
Ok((0..4).map(|c| make_chunk(vec![c], addr, 8)).collect())
|
||||||
|
})
|
||||||
|
.unwrap();
|
||||||
|
assert!(chunks.iter().all(|c| c.address == addr));
|
||||||
|
let want = vec![addr as u8; 8];
|
||||||
|
let got = match cache.get_decompressed_in(addr, &coord) {
|
||||||
|
Some(hit) => hit.to_vec(),
|
||||||
|
None => cache
|
||||||
|
.put_decompressed_in(addr, coord, want.clone())
|
||||||
|
.to_vec(),
|
||||||
|
};
|
||||||
|
assert_eq!(got, want);
|
||||||
|
}
|
||||||
|
})
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
for h in handles {
|
||||||
|
h.join().unwrap();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn duplicate_insert_is_noop() {
|
fn duplicate_insert_is_noop() {
|
||||||
let cache = ChunkCache::new();
|
let cache = ChunkCache::new();
|
||||||
|
|||||||
@@ -0,0 +1,200 @@
|
|||||||
|
//! Chunk-index linearisation shared by the Fixed Array and Extensible Array
|
||||||
|
//! chunk indexes (reader and writer).
|
||||||
|
//!
|
||||||
|
//! Both indexes store one element per chunk at a *linear* index, and the
|
||||||
|
//! library derives that index from the chunk's scaled coordinates
|
||||||
|
//! (`offset / chunk_dim`) using the dataset's **maximum** dimensions, not its
|
||||||
|
//! current ones (`H5D__farray_idx_get_addr` / `H5D__earray_idx_get_addr`,
|
||||||
|
//! via `layout->max_down_chunks`). A dataset whose current shape is smaller
|
||||||
|
//! than its maxshape therefore has gaps in the index, and laying it out by the
|
||||||
|
//! current shape puts every chunk after the first row in the wrong place.
|
||||||
|
//!
|
||||||
|
//! The Extensible Array adds one more step: its one unlimited dimension has no
|
||||||
|
//! finite chunk count, so the library *swizzles* the coordinates to make that
|
||||||
|
//! dimension the slowest-varying one (`H5VM_swizzle_coords`, which moves
|
||||||
|
//! `coords[unlim_dim]` to the front and shifts the dimensions before it right
|
||||||
|
//! by one) before linearising with `swizzled_max_down_chunks`. When the
|
||||||
|
//! unlimited dimension is already dimension 0 no swizzle happens.
|
||||||
|
|
||||||
|
#[cfg(not(feature = "std"))]
|
||||||
|
extern crate alloc;
|
||||||
|
|
||||||
|
#[cfg(not(feature = "std"))]
|
||||||
|
use alloc::{vec, vec::Vec};
|
||||||
|
|
||||||
|
use crate::error::FormatError;
|
||||||
|
|
||||||
|
/// How a chunk index maps linear element indexes to chunk coordinates.
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
pub(crate) struct ChunkGrid {
|
||||||
|
/// Spatial chunk dimensions, in dataset order.
|
||||||
|
chunk_dims: Vec<u64>,
|
||||||
|
/// Chunks per dimension covering the *current* extent, in dataset order.
|
||||||
|
cur_chunks: Vec<u64>,
|
||||||
|
/// Dataset dimension stored at each linearisation position (slowest
|
||||||
|
/// first). The identity except for a swizzled Extensible Array.
|
||||||
|
order: Vec<usize>,
|
||||||
|
/// Linear stride of each linearisation position.
|
||||||
|
down: Vec<u64>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ChunkGrid {
|
||||||
|
/// Grid for a Fixed Array index: row-major over the chunk counts of the
|
||||||
|
/// maximum dimensions (`max_dims`, falling back to the current dimensions
|
||||||
|
/// when the dataspace records none).
|
||||||
|
pub(crate) fn fixed_array(
|
||||||
|
cur_dims: &[u64],
|
||||||
|
max_dims: Option<&[u64]>,
|
||||||
|
chunk_dims: &[u64],
|
||||||
|
) -> Result<Self, FormatError> {
|
||||||
|
Self::build(cur_dims, max_dims, chunk_dims, None)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Grid for an Extensible Array index: like the Fixed Array, but the
|
||||||
|
/// unlimited dimension (the one whose maximum is `H5S_UNLIMITED`) is moved
|
||||||
|
/// to the slowest-varying position first.
|
||||||
|
pub(crate) fn extensible_array(
|
||||||
|
cur_dims: &[u64],
|
||||||
|
max_dims: Option<&[u64]>,
|
||||||
|
chunk_dims: &[u64],
|
||||||
|
) -> Result<Self, FormatError> {
|
||||||
|
let unlim = max_dims.and_then(|m| m.iter().position(|&d| d == u64::MAX));
|
||||||
|
Self::build(cur_dims, max_dims, chunk_dims, unlim)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn build(
|
||||||
|
cur_dims: &[u64],
|
||||||
|
max_dims: Option<&[u64]>,
|
||||||
|
chunk_dims: &[u64],
|
||||||
|
unlim: Option<usize>,
|
||||||
|
) -> Result<Self, FormatError> {
|
||||||
|
let rank = chunk_dims.len();
|
||||||
|
if cur_dims.len() != rank || max_dims.is_some_and(|m| m.len() != rank) {
|
||||||
|
return Err(FormatError::ChunkedReadError(
|
||||||
|
"chunk index rank does not match the dataspace".into(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
if chunk_dims.contains(&0) {
|
||||||
|
return Err(FormatError::ChunkedReadError(
|
||||||
|
"chunk dimension is zero".into(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
let cur_chunks: Vec<u64> = cur_dims
|
||||||
|
.iter()
|
||||||
|
.zip(chunk_dims)
|
||||||
|
.map(|(&d, &c)| d.div_ceil(c))
|
||||||
|
.collect();
|
||||||
|
// Chunk counts of the maximum extent. An unlimited dimension has no
|
||||||
|
// finite count; it only ever sits in the slowest position, where its
|
||||||
|
// count never enters a stride. A (corrupt) maximum smaller than the
|
||||||
|
// current extent is widened so no allocated chunk becomes unreachable.
|
||||||
|
let max_chunks: Vec<u64> = (0..rank)
|
||||||
|
.map(|d| {
|
||||||
|
let max = max_dims.map_or(cur_dims[d], |m| m[d]);
|
||||||
|
if max == u64::MAX {
|
||||||
|
u64::MAX
|
||||||
|
} else {
|
||||||
|
max.div_ceil(chunk_dims[d]).max(cur_chunks[d])
|
||||||
|
}
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
|
||||||
|
let mut order: Vec<usize> = (0..rank).collect();
|
||||||
|
if let Some(u) = unlim {
|
||||||
|
order.remove(u);
|
||||||
|
order.insert(0, u);
|
||||||
|
}
|
||||||
|
let mut down = vec![1u64; rank];
|
||||||
|
for p in (0..rank.saturating_sub(1)).rev() {
|
||||||
|
let next = max_chunks[order[p + 1]];
|
||||||
|
if next == u64::MAX {
|
||||||
|
// Only reachable with more than one unlimited dimension, which
|
||||||
|
// neither index type can describe.
|
||||||
|
return Err(FormatError::ChunkedReadError(
|
||||||
|
"array chunk index with more than one unlimited dimension".into(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
down[p] = down[p + 1].checked_mul(next).ok_or_else(|| {
|
||||||
|
FormatError::Overflow("chunk index linear stride overflows u64".into())
|
||||||
|
})?;
|
||||||
|
}
|
||||||
|
Ok(Self {
|
||||||
|
chunk_dims: chunk_dims.to_vec(),
|
||||||
|
cur_chunks,
|
||||||
|
order,
|
||||||
|
down,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Dataset-space offsets of the chunk stored at linear `index`, or `None`
|
||||||
|
/// when that chunk lies outside the current extent (the index still has a
|
||||||
|
/// slot for it; the library ignores such chunks on read).
|
||||||
|
pub(crate) fn offsets(&self, index: u64) -> Option<Vec<u64>> {
|
||||||
|
let rank = self.chunk_dims.len();
|
||||||
|
let mut offsets = vec![0u64; rank];
|
||||||
|
let mut rem = index;
|
||||||
|
for p in 0..rank {
|
||||||
|
let d = self.order[p];
|
||||||
|
let scaled = rem / self.down[p];
|
||||||
|
rem %= self.down[p];
|
||||||
|
if scaled >= self.cur_chunks[d] {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
offsets[d] = scaled * self.chunk_dims[d];
|
||||||
|
}
|
||||||
|
Some(offsets)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Linear index of the chunk with scaled coordinates `scaled`
|
||||||
|
/// (`offset / chunk_dim` per dimension, in dataset order).
|
||||||
|
pub(crate) fn linear_index(&self, scaled: &[u64]) -> u64 {
|
||||||
|
self.order
|
||||||
|
.iter()
|
||||||
|
.zip(&self.down)
|
||||||
|
.map(|(&d, &stride)| scaled[d] * stride)
|
||||||
|
.sum()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn fixed_array_uses_max_dims() {
|
||||||
|
// shape (4, 6), chunks (2, 3), maxshape (20, 10): 10 x 4 chunk grid.
|
||||||
|
let g = ChunkGrid::fixed_array(&[4, 6], Some(&[20, 10]), &[2, 3]).unwrap();
|
||||||
|
assert_eq!(g.offsets(0), Some(vec![0, 0]));
|
||||||
|
assert_eq!(g.offsets(1), Some(vec![0, 3]));
|
||||||
|
assert_eq!(g.offsets(2), None); // column chunk 2 is beyond the extent
|
||||||
|
assert_eq!(g.offsets(4), Some(vec![2, 0]));
|
||||||
|
assert_eq!(g.offsets(5), Some(vec![2, 3]));
|
||||||
|
assert_eq!(g.offsets(8), None); // row chunk 2 is beyond the extent
|
||||||
|
assert_eq!(g.linear_index(&[1, 1]), 5);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn extensible_array_swizzles_unlimited_dim() {
|
||||||
|
// maxshape (10, None): dim 1 is unlimited and becomes slowest.
|
||||||
|
let g = ChunkGrid::extensible_array(&[4, 6], Some(&[10, u64::MAX]), &[2, 3]).unwrap();
|
||||||
|
// max chunks of dim 0 = 5, so index = c1 * 5 + c0.
|
||||||
|
assert_eq!(g.linear_index(&[1, 0]), 1);
|
||||||
|
assert_eq!(g.linear_index(&[0, 1]), 5);
|
||||||
|
assert_eq!(g.offsets(5), Some(vec![0, 3]));
|
||||||
|
assert_eq!(g.offsets(6), Some(vec![2, 3]));
|
||||||
|
assert_eq!(g.offsets(2), None);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn extensible_array_unlimited_first_is_row_major() {
|
||||||
|
let g = ChunkGrid::extensible_array(&[4, 6], Some(&[u64::MAX, 30]), &[2, 3]).unwrap();
|
||||||
|
// max chunks of dim 1 = 10.
|
||||||
|
assert_eq!(g.linear_index(&[1, 1]), 11);
|
||||||
|
assert_eq!(g.offsets(11), Some(vec![2, 3]));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn rejects_two_unlimited_dims_after_the_first() {
|
||||||
|
assert!(ChunkGrid::fixed_array(&[4, 6], Some(&[u64::MAX, u64::MAX]), &[2, 3]).is_err());
|
||||||
|
}
|
||||||
|
}
|
||||||
File diff suppressed because it is too large
Load Diff
@@ -4,18 +4,18 @@
|
|||||||
extern crate alloc;
|
extern crate alloc;
|
||||||
|
|
||||||
#[cfg(not(feature = "std"))]
|
#[cfg(not(feature = "std"))]
|
||||||
use alloc::{vec, vec::Vec};
|
use alloc::{format, vec, vec::Vec};
|
||||||
|
|
||||||
use crate::checksum::jenkins_lookup3;
|
use crate::checksum::jenkins_lookup3;
|
||||||
use crate::chunk_cache::{CACHE_LINE_SIZE, align_to_cache_line};
|
use crate::chunk_cache::{CACHE_LINE_SIZE, align_to_cache_line};
|
||||||
|
use crate::chunk_grid::ChunkGrid;
|
||||||
use crate::ea_writer;
|
use crate::ea_writer;
|
||||||
use crate::error::FormatError;
|
use crate::error::FormatError;
|
||||||
use crate::filter_pipeline::{
|
use crate::filter_pipeline::{
|
||||||
FILTER_DEFLATE, FILTER_FLETCHER32, FILTER_LZ4, FILTER_PCODEC, FILTER_SHUFFLE, FILTER_ZSTD,
|
FILTER_DEFLATE, FILTER_FLETCHER32, FILTER_LZ4, FILTER_PCODEC, FILTER_PCODEC_NAME,
|
||||||
FilterDescription, FilterPipeline,
|
FILTER_SHUFFLE, FILTER_ZSTD, FilterDescription, FilterPipeline,
|
||||||
};
|
};
|
||||||
use crate::filters::compress_chunk;
|
use crate::filters::compress_chunk;
|
||||||
|
|
||||||
/// Round a file offset up to the next cache-line boundary.
|
/// Round a file offset up to the next cache-line boundary.
|
||||||
///
|
///
|
||||||
/// This ensures chunk data starts at an address that is a multiple of the
|
/// This ensures chunk data starts at an address that is a multiple of the
|
||||||
@@ -45,10 +45,43 @@ pub struct ChunkOptions {
|
|||||||
pub lz4: bool,
|
pub lz4: bool,
|
||||||
/// Zstandard compression level (1-22), None = no zstd. Filter ID 32015.
|
/// Zstandard compression level (1-22), None = no zstd. Filter ID 32015.
|
||||||
pub zstd_level: Option<u32>,
|
pub zstd_level: Option<u32>,
|
||||||
/// Pcodec lossless numerical compression. Filter ID 32023.
|
/// Pcodec lossless numerical compression. Private, unregistered filter
|
||||||
|
/// ID [`FILTER_PCODEC`] (480): only clawhdf5 can read it.
|
||||||
pub pcodec: bool,
|
pub pcodec: bool,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Largest chunk the automatic choice produces, in bytes.
|
||||||
|
const AUTO_CHUNK_TARGET_BYTES: u64 = 1 << 20;
|
||||||
|
|
||||||
|
/// Extent assumed for a dimension that is currently empty (an unlimited
|
||||||
|
/// dimension not yet written to) — the same stand-in h5py uses.
|
||||||
|
const AUTO_CHUNK_EMPTY_DIM: u64 = 1024;
|
||||||
|
|
||||||
|
/// Choose chunk dimensions for a dataset nobody specified them for.
|
||||||
|
///
|
||||||
|
/// Asking for compression (or any filter) without chunk dimensions used to
|
||||||
|
/// make the whole dataset one chunk. That defeats the point of chunking: any
|
||||||
|
/// read — even a single row — must decompress everything, and a large dataset
|
||||||
|
/// cannot be decompressed in parallel. Datasets up to the target size stay a
|
||||||
|
/// single chunk, exactly as before; larger ones are split by halving the
|
||||||
|
/// dimensions in turn (so chunks keep roughly the dataset's proportions, the
|
||||||
|
/// approach h5py takes) until a chunk fits the target.
|
||||||
|
pub fn auto_chunk_dims(shape: &[u64], elem_size: usize) -> Vec<u64> {
|
||||||
|
let mut dims: Vec<u64> = shape
|
||||||
|
.iter()
|
||||||
|
.map(|&d| if d == 0 { AUTO_CHUNK_EMPTY_DIM } else { d })
|
||||||
|
.collect();
|
||||||
|
let elem = elem_size.max(1) as u64;
|
||||||
|
let bytes = |dims: &[u64]| dims.iter().fold(elem, |acc, &d| acc.saturating_mul(d));
|
||||||
|
let mut axis = 0;
|
||||||
|
while bytes(&dims) > AUTO_CHUNK_TARGET_BYTES && dims.iter().any(|&d| d > 1) {
|
||||||
|
let i = axis % dims.len();
|
||||||
|
dims[i] = dims[i].div_ceil(2);
|
||||||
|
axis += 1;
|
||||||
|
}
|
||||||
|
dims
|
||||||
|
}
|
||||||
|
|
||||||
impl ChunkOptions {
|
impl ChunkOptions {
|
||||||
/// Whether any chunking option is enabled.
|
/// Whether any chunking option is enabled.
|
||||||
pub fn is_chunked(&self) -> bool {
|
pub fn is_chunked(&self) -> bool {
|
||||||
@@ -84,7 +117,7 @@ impl ChunkOptions {
|
|||||||
if self.pcodec {
|
if self.pcodec {
|
||||||
filters.push(FilterDescription {
|
filters.push(FilterDescription {
|
||||||
filter_id: FILTER_PCODEC,
|
filter_id: FILTER_PCODEC,
|
||||||
name: Some("pcodec".into()),
|
name: Some(FILTER_PCODEC_NAME.into()),
|
||||||
flags: 0,
|
flags: 0,
|
||||||
client_data: vec![element_size],
|
client_data: vec![element_size],
|
||||||
});
|
});
|
||||||
@@ -135,11 +168,17 @@ impl ChunkOptions {
|
|||||||
|
|
||||||
/// Determine chunk dimensions, using user-specified or auto-computing.
|
/// Determine chunk dimensions, using user-specified or auto-computing.
|
||||||
pub fn resolve_chunk_dims(&self, shape: &[u64]) -> Vec<u64> {
|
pub fn resolve_chunk_dims(&self, shape: &[u64]) -> Vec<u64> {
|
||||||
if let Some(ref dims) = self.chunk_dims {
|
// Without the element size, assume 8 bytes (the widest common scalar);
|
||||||
dims.clone()
|
// the writer uses `resolve_chunk_dims_for`.
|
||||||
} else {
|
self.resolve_chunk_dims_for(shape, 8)
|
||||||
// Auto chunk: use the full dataset shape (single chunk)
|
}
|
||||||
shape.to_vec()
|
|
||||||
|
/// Chunk dimensions for a dataset of `shape` whose elements are `elem_size`
|
||||||
|
/// bytes: the caller's if given, otherwise chosen automatically.
|
||||||
|
pub fn resolve_chunk_dims_for(&self, shape: &[u64], elem_size: usize) -> Vec<u64> {
|
||||||
|
match self.chunk_dims {
|
||||||
|
Some(ref dims) => dims.clone(),
|
||||||
|
None => auto_chunk_dims(shape, elem_size),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -406,6 +445,27 @@ fn serialize_v4_fixed_array(
|
|||||||
element_size: u32,
|
element_size: u32,
|
||||||
max_bits: u8,
|
max_bits: u8,
|
||||||
) -> Vec<u8> {
|
) -> Vec<u8> {
|
||||||
|
let mut buf = layout_v4_chunked_prefix(chunk_dims, element_size);
|
||||||
|
|
||||||
|
// chunk index type = 3 (Fixed Array)
|
||||||
|
buf.push(3);
|
||||||
|
|
||||||
|
// max_dblk_page_nelmts_bits — must match FAHD max_nelmts_bits
|
||||||
|
buf.push(max_bits);
|
||||||
|
|
||||||
|
// Fixed Array header address
|
||||||
|
match offset_size {
|
||||||
|
4 => buf.extend_from_slice(&(fixed_array_address as u32).to_le_bytes()),
|
||||||
|
8 => buf.extend_from_slice(&fixed_array_address.to_le_bytes()),
|
||||||
|
_ => {}
|
||||||
|
}
|
||||||
|
|
||||||
|
buf
|
||||||
|
}
|
||||||
|
|
||||||
|
/// The part of a v4 chunked layout message before the chunk index type:
|
||||||
|
/// version, class, flags and the chunk dimensions (plus the element size).
|
||||||
|
fn layout_v4_chunked_prefix(chunk_dims: &[u32], element_size: u32) -> Vec<u8> {
|
||||||
let mut buf = Vec::new();
|
let mut buf = Vec::new();
|
||||||
buf.push(4); // version
|
buf.push(4); // version
|
||||||
buf.push(2); // class = chunked
|
buf.push(2); // class = chunked
|
||||||
@@ -445,125 +505,143 @@ fn serialize_v4_fixed_array(
|
|||||||
4 => buf.extend_from_slice(&element_size.to_le_bytes()),
|
4 => buf.extend_from_slice(&element_size.to_le_bytes()),
|
||||||
_ => {}
|
_ => {}
|
||||||
}
|
}
|
||||||
|
|
||||||
// chunk index type = 3 (Fixed Array)
|
|
||||||
buf.push(3);
|
|
||||||
|
|
||||||
// max_dblk_page_nelmts_bits — must match FAHD max_nelmts_bits
|
|
||||||
buf.push(max_bits);
|
|
||||||
|
|
||||||
// Fixed Array header address
|
|
||||||
match offset_size {
|
|
||||||
4 => buf.extend_from_slice(&(fixed_array_address as u32).to_le_bytes()),
|
|
||||||
8 => buf.extend_from_slice(&fixed_array_address.to_le_bytes()),
|
|
||||||
_ => {}
|
|
||||||
}
|
|
||||||
|
|
||||||
buf
|
buf
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// log2 of the elements per Fixed Array data block page (the library's
|
||||||
|
/// default, `H5D_FARRAY_MAX_DBLK_PAGE_NELMTS_BITS`).
|
||||||
|
const FA_PAGE_BITS: u8 = 10;
|
||||||
|
|
||||||
|
pub(crate) fn push_addr(buf: &mut Vec<u8>, addr: u64, offset_size: u8) {
|
||||||
|
match offset_size {
|
||||||
|
4 => buf.extend_from_slice(&(addr as u32).to_le_bytes()),
|
||||||
|
_ => buf.extend_from_slice(&addr.to_le_bytes()),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Width of the chunk-size field of a filtered chunk index element. Must
|
||||||
|
/// match the library's `H5D_FARRAY_FILT_COMPUTE_CHUNK_SIZE_LEN` (the EA and
|
||||||
|
/// B-tree v2 indexes use the same formula):
|
||||||
|
/// `1 + ((log2(unfiltered chunk bytes) + 8) / 8)`, capped at 8.
|
||||||
|
pub(crate) fn filtered_chunk_size_len(slots: &[Option<WrittenChunk>]) -> usize {
|
||||||
|
let max_raw = slots
|
||||||
|
.iter()
|
||||||
|
.flatten()
|
||||||
|
.map(|c| c.raw_size)
|
||||||
|
.max()
|
||||||
|
.unwrap_or(1);
|
||||||
|
let log2_val = if max_raw <= 1 {
|
||||||
|
0
|
||||||
|
} else {
|
||||||
|
63 - max_raw.leading_zeros()
|
||||||
|
};
|
||||||
|
(1 + ((log2_val + 8) / 8) as usize).min(8)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Append one chunk index element: the chunk's address, plus its stored size
|
||||||
|
/// and filter mask when the dataset is filtered. `None` is an unallocated
|
||||||
|
/// chunk (undefined address, zero size and mask).
|
||||||
|
pub(crate) fn push_index_element(
|
||||||
|
buf: &mut Vec<u8>,
|
||||||
|
slot: Option<&WrittenChunk>,
|
||||||
|
offset_size: u8,
|
||||||
|
chunk_size_bytes: Option<usize>,
|
||||||
|
) {
|
||||||
|
match slot {
|
||||||
|
Some(c) => {
|
||||||
|
push_addr(buf, c.address, offset_size);
|
||||||
|
if let Some(n) = chunk_size_bytes {
|
||||||
|
buf.extend_from_slice(&c.compressed_size.to_le_bytes()[..n]);
|
||||||
|
buf.extend_from_slice(&c.filter_mask.to_le_bytes());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
None => {
|
||||||
|
buf.extend(core::iter::repeat_n(0xFF, offset_size as usize));
|
||||||
|
if let Some(n) = chunk_size_bytes {
|
||||||
|
buf.extend(core::iter::repeat_n(0x00, n + 4));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
/// Build a complete Fixed Array at a known absolute address.
|
/// Build a complete Fixed Array at a known absolute address.
|
||||||
|
///
|
||||||
|
/// `slots` holds one entry per element of the array, i.e. per chunk of the
|
||||||
|
/// dataset's *maximum* extent in the order [`crate::chunk_grid`] defines;
|
||||||
|
/// `None` marks a chunk that is not allocated. An array with more elements
|
||||||
|
/// than fit in one page (`2^FA_PAGE_BITS`) gets a paged data block: a
|
||||||
|
/// page-init bitmap after the prefix, then one checksummed page per
|
||||||
|
/// `2^FA_PAGE_BITS` elements, the last one short (`H5FA__dblock_create`).
|
||||||
pub fn build_fixed_array_at(
|
pub fn build_fixed_array_at(
|
||||||
chunks: &[WrittenChunk],
|
slots: &[Option<WrittenChunk>],
|
||||||
offset_size: u8,
|
offset_size: u8,
|
||||||
length_size: u8,
|
length_size: u8,
|
||||||
has_filters: bool,
|
has_filters: bool,
|
||||||
fa_base_address: u64,
|
fa_base_address: u64,
|
||||||
) -> Vec<u8> {
|
) -> Vec<u8> {
|
||||||
let os = offset_size as usize;
|
let os = offset_size as usize;
|
||||||
let num_elements = chunks.len();
|
let num_elements = slots.len();
|
||||||
|
|
||||||
// For filtered chunks, compute chunk_size encoding width.
|
|
||||||
// Must match the HDF5 C library's H5D_FARRAY_FILT_COMPUTE_CHUNK_SIZE_LEN macro:
|
|
||||||
// chunk_size_len = 1 + ((H5VM_log2_gen(chunk.size) + 8) / 8)
|
|
||||||
// where chunk.size is the unfiltered chunk size in bytes (product of all chunk dims).
|
|
||||||
let chunk_size_bytes: usize = if has_filters {
|
|
||||||
let max_raw = chunks.iter().map(|c| c.raw_size).max().unwrap_or(1);
|
|
||||||
let log2_val = if max_raw <= 1 {
|
|
||||||
0
|
|
||||||
} else {
|
|
||||||
63 - max_raw.leading_zeros()
|
|
||||||
};
|
|
||||||
let len = 1 + ((log2_val + 8) / 8) as usize;
|
|
||||||
len.min(8)
|
|
||||||
} else {
|
|
||||||
0
|
|
||||||
};
|
|
||||||
|
|
||||||
let elem_size = if has_filters {
|
|
||||||
os + chunk_size_bytes + 4
|
|
||||||
} else {
|
|
||||||
os
|
|
||||||
};
|
|
||||||
|
|
||||||
|
let chunk_size_bytes = has_filters.then(|| filtered_chunk_size_len(slots));
|
||||||
|
let elem_size = os + chunk_size_bytes.map_or(0, |n| n + 4);
|
||||||
let client_id: u8 = if has_filters { 1 } else { 0 };
|
let client_id: u8 = if has_filters { 1 } else { 0 };
|
||||||
|
|
||||||
// FAHD total size
|
// FAHD total size
|
||||||
let nelmts_field_size = length_size as usize;
|
let fahd_total_size = 4 + 1 + 1 + 1 + 1 + length_size as usize + os + 4;
|
||||||
let fahd_total_size = 4 + 1 + 1 + 1 + 1 + nelmts_field_size + os + 4;
|
|
||||||
let fadb_address = fa_base_address + fahd_total_size as u64;
|
let fadb_address = fa_base_address + fahd_total_size as u64;
|
||||||
|
|
||||||
// Build FAHD
|
|
||||||
let mut fahd = Vec::with_capacity(fahd_total_size);
|
let mut fahd = Vec::with_capacity(fahd_total_size);
|
||||||
fahd.extend_from_slice(b"FAHD");
|
fahd.extend_from_slice(b"FAHD");
|
||||||
fahd.push(0); // version
|
fahd.push(0); // version
|
||||||
fahd.push(client_id);
|
fahd.push(client_id);
|
||||||
fahd.push(elem_size as u8);
|
fahd.push(elem_size as u8);
|
||||||
|
fahd.push(FA_PAGE_BITS);
|
||||||
// max_nelmts_bits: use 10 as default (page_size = 1024), matching h5py convention
|
|
||||||
let max_bits: u8 = 10;
|
|
||||||
fahd.push(max_bits);
|
|
||||||
|
|
||||||
match length_size {
|
match length_size {
|
||||||
4 => fahd.extend_from_slice(&(num_elements as u32).to_le_bytes()),
|
4 => fahd.extend_from_slice(&(num_elements as u32).to_le_bytes()),
|
||||||
8 => fahd.extend_from_slice(&(num_elements as u64).to_le_bytes()),
|
|
||||||
_ => fahd.extend_from_slice(&(num_elements as u64).to_le_bytes()),
|
_ => fahd.extend_from_slice(&(num_elements as u64).to_le_bytes()),
|
||||||
}
|
}
|
||||||
|
push_addr(&mut fahd, fadb_address, offset_size);
|
||||||
match offset_size {
|
|
||||||
4 => fahd.extend_from_slice(&(fadb_address as u32).to_le_bytes()),
|
|
||||||
8 => fahd.extend_from_slice(&fadb_address.to_le_bytes()),
|
|
||||||
_ => fahd.extend_from_slice(&fadb_address.to_le_bytes()),
|
|
||||||
}
|
|
||||||
|
|
||||||
// Checksum
|
|
||||||
let checksum = jenkins_lookup3(&fahd);
|
let checksum = jenkins_lookup3(&fahd);
|
||||||
fahd.extend_from_slice(&checksum.to_le_bytes());
|
fahd.extend_from_slice(&checksum.to_le_bytes());
|
||||||
|
|
||||||
assert_eq!(fahd.len(), fahd_total_size);
|
assert_eq!(fahd.len(), fahd_total_size);
|
||||||
|
|
||||||
// Build FADB
|
// FADB prefix
|
||||||
let mut fadb = Vec::new();
|
let mut fadb = Vec::new();
|
||||||
fadb.extend_from_slice(b"FADB");
|
fadb.extend_from_slice(b"FADB");
|
||||||
fadb.push(0); // version
|
fadb.push(0); // version
|
||||||
fadb.push(client_id);
|
fadb.push(client_id);
|
||||||
|
push_addr(&mut fadb, fa_base_address, offset_size);
|
||||||
|
|
||||||
// header address
|
let page_nelmts = 1usize << FA_PAGE_BITS;
|
||||||
match offset_size {
|
if num_elements <= page_nelmts {
|
||||||
4 => fadb.extend_from_slice(&(fa_base_address as u32).to_le_bytes()),
|
// Unpaged: the elements follow the prefix, one checksum over both.
|
||||||
8 => fadb.extend_from_slice(&fa_base_address.to_le_bytes()),
|
for slot in slots {
|
||||||
_ => fadb.extend_from_slice(&fa_base_address.to_le_bytes()),
|
push_index_element(&mut fadb, slot.as_ref(), offset_size, chunk_size_bytes);
|
||||||
}
|
|
||||||
|
|
||||||
// Element data
|
|
||||||
for chunk in chunks {
|
|
||||||
match offset_size {
|
|
||||||
4 => fadb.extend_from_slice(&(chunk.address as u32).to_le_bytes()),
|
|
||||||
8 => fadb.extend_from_slice(&chunk.address.to_le_bytes()),
|
|
||||||
_ => fadb.extend_from_slice(&chunk.address.to_le_bytes()),
|
|
||||||
}
|
}
|
||||||
if has_filters {
|
let fadb_checksum = jenkins_lookup3(&fadb);
|
||||||
// Write compressed size using chunk_size_bytes (variable width)
|
fadb.extend_from_slice(&fadb_checksum.to_le_bytes());
|
||||||
let cs_bytes = chunk.compressed_size.to_le_bytes();
|
} else {
|
||||||
fadb.extend_from_slice(&cs_bytes[..chunk_size_bytes]);
|
// Paged: every page is written, so every page-init bit is set
|
||||||
fadb.extend_from_slice(&chunk.filter_mask.to_le_bytes());
|
// (MSB-first, as `H5VM_bit_set` packs them). The prefix and bitmap
|
||||||
|
// share a checksum; each page carries its own.
|
||||||
|
let npages = num_elements.div_ceil(page_nelmts);
|
||||||
|
let mut bitmap = vec![0u8; npages.div_ceil(8)];
|
||||||
|
for p in 0..npages {
|
||||||
|
bitmap[p / 8] |= 0x80 >> (p % 8);
|
||||||
|
}
|
||||||
|
fadb.extend_from_slice(&bitmap);
|
||||||
|
let prefix_checksum = jenkins_lookup3(&fadb);
|
||||||
|
fadb.extend_from_slice(&prefix_checksum.to_le_bytes());
|
||||||
|
for page in slots.chunks(page_nelmts) {
|
||||||
|
let start = fadb.len();
|
||||||
|
for slot in page {
|
||||||
|
push_index_element(&mut fadb, slot.as_ref(), offset_size, chunk_size_bytes);
|
||||||
|
}
|
||||||
|
let page_checksum = jenkins_lookup3(&fadb[start..]);
|
||||||
|
fadb.extend_from_slice(&page_checksum.to_le_bytes());
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// FADB checksum
|
|
||||||
let fadb_checksum = jenkins_lookup3(&fadb);
|
|
||||||
fadb.extend_from_slice(&fadb_checksum.to_le_bytes());
|
|
||||||
|
|
||||||
let mut combined = fahd;
|
let mut combined = fahd;
|
||||||
combined.extend_from_slice(&fadb);
|
combined.extend_from_slice(&fadb);
|
||||||
combined
|
combined
|
||||||
@@ -630,7 +708,8 @@ pub fn build_chunked_data_from_precompressed(
|
|||||||
pre: &PrecompressedChunks,
|
pre: &PrecompressedChunks,
|
||||||
base_address: u64,
|
base_address: u64,
|
||||||
maxshape: Option<&[u64]>,
|
maxshape: Option<&[u64]>,
|
||||||
) -> ChunkedDataResult {
|
) -> Result<ChunkedDataResult, FormatError> {
|
||||||
|
let index = ChunkIndexPlan::new(&pre.shape, maxshape, &pre.chunk_dims)?;
|
||||||
let offset_size: u8 = 8;
|
let offset_size: u8 = 8;
|
||||||
let length_size: u8 = 8;
|
let length_size: u8 = 8;
|
||||||
let num_chunks = pre.chunks.len();
|
let num_chunks = pre.chunks.len();
|
||||||
@@ -656,71 +735,329 @@ pub fn build_chunked_data_from_precompressed(
|
|||||||
}
|
}
|
||||||
|
|
||||||
let chunk_dims_u32: Vec<u32> = pre.chunk_dims.iter().map(|&d| d as u32).collect();
|
let chunk_dims_u32: Vec<u32> = pre.chunk_dims.iter().map(|&d| d as u32).collect();
|
||||||
let use_extensible = maxshape.is_some_and(|ms| ms.contains(&u64::MAX));
|
|
||||||
|
|
||||||
let aligned_idx = align_to_cache_line(data_buf.len());
|
let aligned_idx = align_to_cache_line(data_buf.len());
|
||||||
if aligned_idx > data_buf.len() {
|
if aligned_idx > data_buf.len() {
|
||||||
data_buf.resize(aligned_idx, 0u8);
|
data_buf.resize(aligned_idx, 0u8);
|
||||||
}
|
}
|
||||||
|
|
||||||
let layout_message = if use_extensible {
|
let layout_message = match &index {
|
||||||
let ea_address = base_address + data_buf.len() as u64;
|
ChunkIndexPlan::ExtensibleArray(grid) => {
|
||||||
let ea_bytes = ea_writer::build_extensible_array_at(
|
let ea_address = base_address + data_buf.len() as u64;
|
||||||
&written_chunks,
|
let slots = index_slots(grid, &pre.shape, &pre.chunk_dims, &written_chunks, None)?;
|
||||||
offset_size,
|
let ea_bytes = ea_writer::build_extensible_array_at(
|
||||||
length_size,
|
&slots,
|
||||||
pre.has_filters,
|
offset_size,
|
||||||
ea_address,
|
length_size,
|
||||||
);
|
pre.has_filters,
|
||||||
data_buf.extend_from_slice(&ea_bytes);
|
ea_address,
|
||||||
ea_writer::serialize_v4_extensible_array(
|
);
|
||||||
&chunk_dims_u32,
|
data_buf.extend_from_slice(&ea_bytes);
|
||||||
ea_address,
|
ea_writer::serialize_v4_extensible_array(
|
||||||
offset_size,
|
&chunk_dims_u32,
|
||||||
element_size as u32,
|
ea_address,
|
||||||
)
|
offset_size,
|
||||||
} else if num_chunks == 1 {
|
element_size as u32,
|
||||||
let chunk_addr = written_chunks[0].address;
|
)
|
||||||
let filtered_size = if pre.has_filters {
|
}
|
||||||
Some(written_chunks[0].compressed_size)
|
ChunkIndexPlan::SingleChunk => {
|
||||||
} else {
|
let chunk_addr = written_chunks[0].address;
|
||||||
None
|
let filtered_size = if pre.has_filters {
|
||||||
};
|
Some(written_chunks[0].compressed_size)
|
||||||
let filter_mask = if pre.has_filters { Some(0u32) } else { None };
|
} else {
|
||||||
serialize_v4_single_chunk(
|
None
|
||||||
&chunk_dims_u32,
|
};
|
||||||
chunk_addr,
|
let filter_mask = if pre.has_filters { Some(0u32) } else { None };
|
||||||
filtered_size,
|
serialize_v4_single_chunk(
|
||||||
filter_mask,
|
&chunk_dims_u32,
|
||||||
offset_size,
|
chunk_addr,
|
||||||
element_size as u32,
|
filtered_size,
|
||||||
)
|
filter_mask,
|
||||||
} else {
|
offset_size,
|
||||||
let fa_address = base_address + data_buf.len() as u64;
|
element_size as u32,
|
||||||
let fa_bytes = build_fixed_array_at(
|
)
|
||||||
&written_chunks,
|
}
|
||||||
offset_size,
|
ChunkIndexPlan::FixedArray(grid, nslots) => {
|
||||||
length_size,
|
let fa_address = base_address + data_buf.len() as u64;
|
||||||
pre.has_filters,
|
let slots = index_slots(
|
||||||
fa_address,
|
grid,
|
||||||
);
|
&pre.shape,
|
||||||
data_buf.extend_from_slice(&fa_bytes);
|
&pre.chunk_dims,
|
||||||
serialize_v4_fixed_array(
|
&written_chunks,
|
||||||
&chunk_dims_u32,
|
Some(*nslots),
|
||||||
fa_address,
|
)?;
|
||||||
offset_size,
|
let fa_bytes = build_fixed_array_at(
|
||||||
element_size as u32,
|
&slots,
|
||||||
10, // max_nelmts_bits — matches h5py convention
|
offset_size,
|
||||||
)
|
length_size,
|
||||||
|
pre.has_filters,
|
||||||
|
fa_address,
|
||||||
|
);
|
||||||
|
data_buf.extend_from_slice(&fa_bytes);
|
||||||
|
serialize_v4_fixed_array(
|
||||||
|
&chunk_dims_u32,
|
||||||
|
fa_address,
|
||||||
|
offset_size,
|
||||||
|
element_size as u32,
|
||||||
|
FA_PAGE_BITS,
|
||||||
|
)
|
||||||
|
}
|
||||||
|
ChunkIndexPlan::BTreeV2 => {
|
||||||
|
let bt_address = base_address + data_buf.len() as u64;
|
||||||
|
let records: Vec<(Vec<u64>, &WrittenChunk)> = written_chunks
|
||||||
|
.iter()
|
||||||
|
.enumerate()
|
||||||
|
.map(|(i, c)| (scaled_coords(&pre.shape, &pre.chunk_dims, i), c))
|
||||||
|
.collect();
|
||||||
|
let (bt_bytes, node_size) = build_btree_v2_chunk_index_at(
|
||||||
|
pre.shape.len(),
|
||||||
|
&records,
|
||||||
|
offset_size,
|
||||||
|
length_size,
|
||||||
|
pre.has_filters,
|
||||||
|
bt_address,
|
||||||
|
)?;
|
||||||
|
data_buf.extend_from_slice(&bt_bytes);
|
||||||
|
serialize_v4_btree_v2(
|
||||||
|
&chunk_dims_u32,
|
||||||
|
bt_address,
|
||||||
|
offset_size,
|
||||||
|
element_size as u32,
|
||||||
|
node_size,
|
||||||
|
)
|
||||||
|
}
|
||||||
};
|
};
|
||||||
|
|
||||||
ChunkedDataResult {
|
Ok(ChunkedDataResult {
|
||||||
data_bytes: data_buf,
|
data_bytes: data_buf,
|
||||||
layout_message,
|
layout_message,
|
||||||
pipeline_message: pre.pipeline_message.clone(),
|
pipeline_message: pre.pipeline_message.clone(),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Most slots a Fixed Array index may have before we refuse to build it: its
|
||||||
|
/// data block holds one element per chunk of the *maximum* extent, so a huge
|
||||||
|
/// finite maxshape with small chunks would otherwise exhaust memory.
|
||||||
|
const MAX_FIXED_ARRAY_SLOTS: u64 = 1 << 26;
|
||||||
|
|
||||||
|
/// Which chunk index a dataset gets, following the library's choice in
|
||||||
|
/// `H5D__layout_set_latest_indexing`: version-2 B-tree for more than one
|
||||||
|
/// unlimited dimension, Extensible Array for exactly one, Fixed Array for a
|
||||||
|
/// finite maxshape, Single Chunk when the whole maximum extent is one chunk.
|
||||||
|
enum ChunkIndexPlan {
|
||||||
|
SingleChunk,
|
||||||
|
/// The grid and the number of array elements (chunks of the max extent).
|
||||||
|
FixedArray(ChunkGrid, usize),
|
||||||
|
ExtensibleArray(ChunkGrid),
|
||||||
|
BTreeV2,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ChunkIndexPlan {
|
||||||
|
fn new(
|
||||||
|
shape: &[u64],
|
||||||
|
maxshape: Option<&[u64]>,
|
||||||
|
chunk_dims: &[u64],
|
||||||
|
) -> Result<Self, FormatError> {
|
||||||
|
let bad = |what: &str| FormatError::ChunkedReadError(format!("maxshape: {what}"));
|
||||||
|
if let Some(ms) = maxshape {
|
||||||
|
if ms.len() != shape.len() {
|
||||||
|
return Err(bad("rank differs from the shape"));
|
||||||
|
}
|
||||||
|
if ms.iter().zip(shape).any(|(&m, &s)| m < s) {
|
||||||
|
return Err(bad("smaller than the shape"));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
let max = maxshape.unwrap_or(shape);
|
||||||
|
let nunlim = max.iter().filter(|&&d| d == u64::MAX).count();
|
||||||
|
match nunlim {
|
||||||
|
0 => {
|
||||||
|
let nslots = max
|
||||||
|
.iter()
|
||||||
|
.zip(chunk_dims)
|
||||||
|
.try_fold(1u64, |acc, (&m, &c)| acc.checked_mul(m.div_ceil(c.max(1))))
|
||||||
|
.filter(|&n| n <= MAX_FIXED_ARRAY_SLOTS)
|
||||||
|
.ok_or_else(|| {
|
||||||
|
bad("too many chunks for a Fixed Array index; \
|
||||||
|
use larger chunks or an unlimited dimension")
|
||||||
|
})?;
|
||||||
|
// A Single Chunk index needs that one chunk to exist; an
|
||||||
|
// empty dataset gets an all-unallocated Fixed Array instead.
|
||||||
|
let empty = shape.contains(&0);
|
||||||
|
if nslots == 1 && !empty {
|
||||||
|
Ok(Self::SingleChunk)
|
||||||
|
} else {
|
||||||
|
let grid = ChunkGrid::fixed_array(shape, Some(max), chunk_dims)?;
|
||||||
|
Ok(Self::FixedArray(grid, nslots as usize))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
1 => Ok(Self::ExtensibleArray(ChunkGrid::extensible_array(
|
||||||
|
shape,
|
||||||
|
Some(max),
|
||||||
|
chunk_dims,
|
||||||
|
)?)),
|
||||||
|
_ => Ok(Self::BTreeV2),
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Place each written chunk at its linear index in `grid`. `chunks` are in
|
||||||
|
/// row-major order over the chunks of the current extent (`split_into_chunks`).
|
||||||
|
/// `len` fixes the slot count (Fixed Array); otherwise it is one past the
|
||||||
|
/// highest index used.
|
||||||
|
fn index_slots(
|
||||||
|
grid: &ChunkGrid,
|
||||||
|
shape: &[u64],
|
||||||
|
chunk_dims: &[u64],
|
||||||
|
chunks: &[WrittenChunk],
|
||||||
|
len: Option<usize>,
|
||||||
|
) -> Result<Vec<Option<WrittenChunk>>, FormatError> {
|
||||||
|
let mut placed: Vec<(usize, &WrittenChunk)> = Vec::with_capacity(chunks.len());
|
||||||
|
for (i, chunk) in chunks.iter().enumerate() {
|
||||||
|
let scaled = scaled_coords(shape, chunk_dims, i);
|
||||||
|
let idx = usize::try_from(grid.linear_index(&scaled))
|
||||||
|
.map_err(|_| FormatError::Overflow("chunk index slot".into()))?;
|
||||||
|
placed.push((idx, chunk));
|
||||||
|
}
|
||||||
|
let n = len.unwrap_or_else(|| placed.iter().map(|&(i, _)| i + 1).max().unwrap_or(0));
|
||||||
|
let mut slots = vec![None; n];
|
||||||
|
for (idx, chunk) in placed {
|
||||||
|
*slots
|
||||||
|
.get_mut(idx)
|
||||||
|
.ok_or_else(|| FormatError::Overflow("chunk index slot".into()))? = Some(chunk.clone());
|
||||||
|
}
|
||||||
|
Ok(slots)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Scaled coordinates (`offset / chunk_dim`) of the `i`-th chunk in the
|
||||||
|
/// row-major order `split_into_chunks` produces over the current extent.
|
||||||
|
fn scaled_coords(shape: &[u64], chunk_dims: &[u64], i: usize) -> Vec<u64> {
|
||||||
|
let rank = shape.len();
|
||||||
|
let mut scaled = vec![0u64; rank];
|
||||||
|
let mut rem = i as u64;
|
||||||
|
for d in (0..rank).rev() {
|
||||||
|
let n = shape[d].div_ceil(chunk_dims[d]);
|
||||||
|
scaled[d] = rem % n;
|
||||||
|
rem /= n;
|
||||||
|
}
|
||||||
|
scaled
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Node size the library gives a chunk index B-tree (`H5D_BT2_NODE_SIZE`),
|
||||||
|
/// with its split and merge percentages.
|
||||||
|
const BT2_NODE_SIZE: u32 = 2048;
|
||||||
|
const BT2_SPLIT_PERCENT: u8 = 100;
|
||||||
|
const BT2_MERGE_PERCENT: u8 = 40;
|
||||||
|
/// B-tree v2 record types for chunk indexes (`H5B2_CDSET_ID`,
|
||||||
|
/// `H5B2_CDSET_FILT_ID`).
|
||||||
|
const BT2_CHUNK_UNFILTERED: u8 = 10;
|
||||||
|
const BT2_CHUNK_FILTERED: u8 = 11;
|
||||||
|
|
||||||
|
/// Build a version-2 B-tree chunk index (the library's index for datasets
|
||||||
|
/// with more than one unlimited dimension) at a known absolute address.
|
||||||
|
///
|
||||||
|
/// `records` are `(scaled coordinates, chunk)` in lexicographic order of the
|
||||||
|
/// coordinates, which is the order the library's comparator
|
||||||
|
/// (`H5VM_vector_cmp_u`) keeps them in. The tree is a single leaf: the
|
||||||
|
/// library's 2048-byte node when the records fit, otherwise a leaf node
|
||||||
|
/// sized to hold them all (the root's record count is 16-bit, so at most
|
||||||
|
/// 65535 chunks). Returns the bytes and the node size the layout message
|
||||||
|
/// must record.
|
||||||
|
fn build_btree_v2_chunk_index_at(
|
||||||
|
rank: usize,
|
||||||
|
records: &[(Vec<u64>, &WrittenChunk)],
|
||||||
|
offset_size: u8,
|
||||||
|
length_size: u8,
|
||||||
|
has_filters: bool,
|
||||||
|
base_address: u64,
|
||||||
|
) -> Result<(Vec<u8>, u32), FormatError> {
|
||||||
|
let os = offset_size as usize;
|
||||||
|
let nrec = u16::try_from(records.len()).map_err(|_| {
|
||||||
|
FormatError::ChunkedReadError(
|
||||||
|
"more than 65535 chunks with more than one unlimited dimension: \
|
||||||
|
use larger chunks"
|
||||||
|
.into(),
|
||||||
|
)
|
||||||
|
})?;
|
||||||
|
let chunk_size_bytes = has_filters.then(|| {
|
||||||
|
let slots: Vec<Option<WrittenChunk>> =
|
||||||
|
records.iter().map(|(_, c)| Some((*c).clone())).collect();
|
||||||
|
filtered_chunk_size_len(&slots)
|
||||||
|
});
|
||||||
|
let record_size = os + chunk_size_bytes.map_or(0, |n| n + 4) + 8 * rank;
|
||||||
|
// Leaf: signature, version, type, records, checksum.
|
||||||
|
let leaf_len = 4 + 1 + 1 + records.len() * record_size + 4;
|
||||||
|
let node_size = u32::try_from(leaf_len)
|
||||||
|
.map_err(|_| FormatError::Overflow("B-tree v2 leaf size".into()))?
|
||||||
|
.max(BT2_NODE_SIZE);
|
||||||
|
let tree_type = if has_filters {
|
||||||
|
BT2_CHUNK_FILTERED
|
||||||
|
} else {
|
||||||
|
BT2_CHUNK_UNFILTERED
|
||||||
|
};
|
||||||
|
|
||||||
|
let hdr_len = 4 + 1 + 1 + 4 + 2 + 2 + 1 + 1 + os + 2 + length_size as usize + 4;
|
||||||
|
let leaf_address = base_address + hdr_len as u64;
|
||||||
|
|
||||||
|
let mut out = Vec::with_capacity(hdr_len + node_size as usize);
|
||||||
|
out.extend_from_slice(b"BTHD");
|
||||||
|
out.push(0); // version
|
||||||
|
out.push(tree_type);
|
||||||
|
out.extend_from_slice(&node_size.to_le_bytes());
|
||||||
|
out.extend_from_slice(&(record_size as u16).to_le_bytes());
|
||||||
|
out.extend_from_slice(&0u16.to_le_bytes()); // depth
|
||||||
|
out.push(BT2_SPLIT_PERCENT);
|
||||||
|
out.push(BT2_MERGE_PERCENT);
|
||||||
|
if records.is_empty() {
|
||||||
|
out.extend(core::iter::repeat_n(0xFF, os));
|
||||||
|
} else {
|
||||||
|
push_addr(&mut out, leaf_address, offset_size);
|
||||||
|
}
|
||||||
|
out.extend_from_slice(&nrec.to_le_bytes());
|
||||||
|
match length_size {
|
||||||
|
4 => out.extend_from_slice(&(records.len() as u32).to_le_bytes()),
|
||||||
|
_ => out.extend_from_slice(&(records.len() as u64).to_le_bytes()),
|
||||||
|
}
|
||||||
|
let sum = jenkins_lookup3(&out);
|
||||||
|
out.extend_from_slice(&sum.to_le_bytes());
|
||||||
|
debug_assert_eq!(out.len(), hdr_len);
|
||||||
|
if records.is_empty() {
|
||||||
|
return Ok((out, node_size));
|
||||||
|
}
|
||||||
|
|
||||||
|
let leaf_start = out.len();
|
||||||
|
out.extend_from_slice(b"BTLF");
|
||||||
|
out.push(0); // version
|
||||||
|
out.push(tree_type);
|
||||||
|
for (scaled, chunk) in records {
|
||||||
|
push_index_element(&mut out, Some(chunk), offset_size, chunk_size_bytes);
|
||||||
|
for &c in scaled {
|
||||||
|
out.extend_from_slice(&c.to_le_bytes());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
let sum = jenkins_lookup3(&out[leaf_start..]);
|
||||||
|
out.extend_from_slice(&sum.to_le_bytes());
|
||||||
|
// The library reads whole nodes; pad the leaf out to the node size.
|
||||||
|
out.resize(leaf_start + node_size as usize, 0);
|
||||||
|
Ok((out, node_size))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Serialize a v4 layout message for a version-2 B-tree chunk index.
|
||||||
|
fn serialize_v4_btree_v2(
|
||||||
|
chunk_dims: &[u32],
|
||||||
|
btree_address: u64,
|
||||||
|
offset_size: u8,
|
||||||
|
element_size: u32,
|
||||||
|
node_size: u32,
|
||||||
|
) -> Vec<u8> {
|
||||||
|
let mut buf = layout_v4_chunked_prefix(chunk_dims, element_size);
|
||||||
|
buf.push(5); // chunk index type = 5 (version-2 B-tree)
|
||||||
|
buf.extend_from_slice(&node_size.to_le_bytes());
|
||||||
|
buf.push(BT2_SPLIT_PERCENT);
|
||||||
|
buf.push(BT2_MERGE_PERCENT);
|
||||||
|
push_addr(&mut buf, btree_address, offset_size);
|
||||||
|
buf
|
||||||
|
}
|
||||||
|
|
||||||
/// Build chunked data with absolute addresses.
|
/// Build chunked data with absolute addresses.
|
||||||
/// If `maxshape` has unlimited dims, uses Extensible Array index.
|
/// If `maxshape` has unlimited dims, uses Extensible Array index.
|
||||||
pub fn build_chunked_data_at(
|
pub fn build_chunked_data_at(
|
||||||
@@ -753,11 +1090,7 @@ 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(
|
build_chunked_data_from_precompressed(&pre, base_address, maxshape)
|
||||||
&pre,
|
|
||||||
base_address,
|
|
||||||
maxshape,
|
|
||||||
))
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Write selected elements into an existing in-memory dataset buffer.
|
/// Write selected elements into an existing in-memory dataset buffer.
|
||||||
@@ -890,6 +1223,7 @@ pub fn write_selection_to_buffer(
|
|||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
|
|
||||||
use super::*;
|
use super::*;
|
||||||
use crate::chunked_read::read_chunked_data;
|
use crate::chunked_read::read_chunked_data;
|
||||||
use crate::data_layout::DataLayout;
|
use crate::data_layout::DataLayout;
|
||||||
@@ -1143,6 +1477,45 @@ mod tests {
|
|||||||
assert_eq!(dims, vec![100, 50]);
|
assert_eq!(dims, vec![100, 50]);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn auto_chunking_splits_only_large_datasets() {
|
||||||
|
let bytes = |dims: &[u64], elem: u64| dims.iter().product::<u64>() * elem;
|
||||||
|
// Up to the target: one chunk, as before.
|
||||||
|
assert_eq!(auto_chunk_dims(&[100, 50], 8), [100, 50]);
|
||||||
|
assert_eq!(auto_chunk_dims(&[131_072], 8), [131_072]); // exactly 1 MiB
|
||||||
|
// Larger: split, keeping proportions, never above the target.
|
||||||
|
let big = auto_chunk_dims(&[4096, 2048], 8);
|
||||||
|
assert!(bytes(&big, 8) <= AUTO_CHUNK_TARGET_BYTES, "{big:?}");
|
||||||
|
assert!(bytes(&big, 8) > AUTO_CHUNK_TARGET_BYTES / 4, "{big:?}");
|
||||||
|
assert_eq!(big[0] / big[1], 2, "proportions kept: {big:?}");
|
||||||
|
// Every dimension stays within the dataset and at least 1.
|
||||||
|
for shape in [
|
||||||
|
vec![10_000_000u64],
|
||||||
|
vec![3, 5_000_000],
|
||||||
|
vec![1, 1, 9_000_000],
|
||||||
|
vec![7; 9],
|
||||||
|
] {
|
||||||
|
let dims = auto_chunk_dims(&shape, 4);
|
||||||
|
assert!(
|
||||||
|
dims.iter().zip(&shape).all(|(c, s)| *c >= 1 && c <= s),
|
||||||
|
"{shape:?} -> {dims:?}"
|
||||||
|
);
|
||||||
|
assert!(
|
||||||
|
bytes(&dims, 4) <= AUTO_CHUNK_TARGET_BYTES,
|
||||||
|
"{shape:?} -> {dims:?}"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
// An empty (unlimited, unwritten) dimension still gets a usable chunk.
|
||||||
|
let growable = auto_chunk_dims(&[0, 128], 8);
|
||||||
|
assert!(growable[0] >= 1 && bytes(&growable, 8) <= AUTO_CHUNK_TARGET_BYTES);
|
||||||
|
// Explicit dimensions always win.
|
||||||
|
let explicit = ChunkOptions {
|
||||||
|
chunk_dims: Some(vec![10, 10]),
|
||||||
|
..Default::default()
|
||||||
|
};
|
||||||
|
assert_eq!(explicit.resolve_chunk_dims_for(&[4096, 2048], 8), [10, 10]);
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn chunk_options_pipeline_deflate() {
|
fn chunk_options_pipeline_deflate() {
|
||||||
// Auto-shuffle is applied before compression by default (matches h5py).
|
// Auto-shuffle is applied before compression by default (matches h5py).
|
||||||
@@ -1237,6 +1610,7 @@ mod tests {
|
|||||||
chunk_index_type,
|
chunk_index_type,
|
||||||
single_chunk_filtered_size,
|
single_chunk_filtered_size,
|
||||||
single_chunk_filter_mask,
|
single_chunk_filter_mask,
|
||||||
|
..
|
||||||
} => {
|
} => {
|
||||||
assert_eq!(version, 4);
|
assert_eq!(version, 4);
|
||||||
assert_eq!(chunk_index_type, Some(1));
|
assert_eq!(chunk_index_type, Some(1));
|
||||||
@@ -1305,7 +1679,8 @@ mod tests {
|
|||||||
filter_mask: 0,
|
filter_mask: 0,
|
||||||
},
|
},
|
||||||
];
|
];
|
||||||
let fa = build_fixed_array_at(&chunks, 8, 8, false, 0x2000);
|
let slots: Vec<_> = chunks.into_iter().map(Some).collect();
|
||||||
|
let fa = build_fixed_array_at(&slots, 8, 8, false, 0x2000);
|
||||||
// Should start with FAHD
|
// Should start with FAHD
|
||||||
assert_eq!(&fa[0..4], b"FAHD");
|
assert_eq!(&fa[0..4], b"FAHD");
|
||||||
// FAHD size = 4+1+1+1+1+8+8+4 = 28
|
// FAHD size = 4+1+1+1+1+8+8+4 = 28
|
||||||
@@ -1352,7 +1727,8 @@ mod tests {
|
|||||||
filter_mask: 0,
|
filter_mask: 0,
|
||||||
},
|
},
|
||||||
];
|
];
|
||||||
let ea = ea_writer::build_extensible_array_at(&chunks, 8, 8, false, 0x2000);
|
let slots: Vec<_> = chunks.into_iter().map(Some).collect();
|
||||||
|
let ea = ea_writer::build_extensible_array_at(&slots, 8, 8, false, 0x2000);
|
||||||
assert_eq!(&ea[0..4], b"EAHD");
|
assert_eq!(&ea[0..4], b"EAHD");
|
||||||
// Find EAIB after EAHD: 12 fixed + 6*8 stats + 8 addr + 4 checksum = 72
|
// Find EAIB after EAHD: 12 fixed + 6*8 stats + 8 addr + 4 checksum = 72
|
||||||
let aehd_size = 4 + 1 + 1 + 1 + 1 + 1 + 1 + 1 + 1 + 6 * 8 + 8 + 4;
|
let aehd_size = 4 + 1 + 1 + 1 + 1 + 1 + 1 + 1 + 1 + 6 * 8 + 8 + 4;
|
||||||
@@ -1435,9 +1811,20 @@ mod tests {
|
|||||||
|
|
||||||
// ---- h5py round-trip tests for chunked writes ----
|
// ---- h5py round-trip tests for chunked writes ----
|
||||||
|
|
||||||
|
/// The Python interpreter to drive interop checks with.
|
||||||
|
///
|
||||||
|
/// `CLAWHDF5_PYTHON` lets these run against a virtualenv holding h5py,
|
||||||
|
/// which on a PEP 668 "externally managed" system is the only place it
|
||||||
|
/// can be installed. Without it the suite silently skips, and a silent
|
||||||
|
/// skip here is how a datatype bug once reached a release.
|
||||||
|
#[cfg(feature = "std")]
|
||||||
|
fn python() -> String {
|
||||||
|
std::env::var("CLAWHDF5_PYTHON").unwrap_or_else(|_| "python3".to_string())
|
||||||
|
}
|
||||||
|
|
||||||
#[cfg(feature = "std")]
|
#[cfg(feature = "std")]
|
||||||
fn h5py_available() -> bool {
|
fn h5py_available() -> bool {
|
||||||
std::process::Command::new("python3")
|
std::process::Command::new(python())
|
||||||
.args(["-c", "import h5py"])
|
.args(["-c", "import h5py"])
|
||||||
.output()
|
.output()
|
||||||
.map(|o| o.status.success())
|
.map(|o| o.status.success())
|
||||||
@@ -1449,10 +1836,10 @@ mod tests {
|
|||||||
if !h5py_available() {
|
if !h5py_available() {
|
||||||
panic!("h5py not installed — skipping interop test");
|
panic!("h5py not installed — skipping interop test");
|
||||||
}
|
}
|
||||||
let o = std::process::Command::new("python3")
|
let o = std::process::Command::new(python())
|
||||||
.args(["-c", script])
|
.args(["-c", script])
|
||||||
.output()
|
.output()
|
||||||
.expect("python3");
|
.expect("python interpreter");
|
||||||
if !o.status.success() {
|
if !o.status.success() {
|
||||||
panic!("h5py: {}", String::from_utf8_lossy(&o.stderr));
|
panic!("h5py: {}", String::from_utf8_lossy(&o.stderr));
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -53,6 +53,11 @@ pub enum DataLayout {
|
|||||||
single_chunk_filtered_size: Option<u64>,
|
single_chunk_filtered_size: Option<u64>,
|
||||||
/// Filter mask for v4 single chunk with filters.
|
/// Filter mask for v4 single chunk with filters.
|
||||||
single_chunk_filter_mask: Option<u32>,
|
single_chunk_filter_mask: Option<u32>,
|
||||||
|
/// Layout v4 flag bit 0 (`H5D_CHUNK_DONT_FILTER_PARTIAL_CHUNKS`):
|
||||||
|
/// partial edge chunks — those extending past the dataset's current
|
||||||
|
/// extent in some dimension — are stored without the filter pipeline,
|
||||||
|
/// even though their filter mask is 0. Always `false` for v3.
|
||||||
|
dont_filter_partial_edge_chunks: bool,
|
||||||
},
|
},
|
||||||
/// Virtual dataset layout (v4 only).
|
/// Virtual dataset layout (v4 only).
|
||||||
Virtual {
|
Virtual {
|
||||||
@@ -322,6 +327,7 @@ impl DataLayout {
|
|||||||
chunk_index_type: None,
|
chunk_index_type: None,
|
||||||
single_chunk_filtered_size: None,
|
single_chunk_filtered_size: None,
|
||||||
single_chunk_filter_mask: None,
|
single_chunk_filter_mask: None,
|
||||||
|
dont_filter_partial_edge_chunks: false,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
_ => Err(FormatError::InvalidLayoutClass(layout_class)),
|
_ => Err(FormatError::InvalidLayoutClass(layout_class)),
|
||||||
@@ -505,6 +511,7 @@ impl DataLayout {
|
|||||||
chunk_index_type: Some(chunk_index_type),
|
chunk_index_type: Some(chunk_index_type),
|
||||||
single_chunk_filtered_size,
|
single_chunk_filtered_size,
|
||||||
single_chunk_filter_mask,
|
single_chunk_filter_mask,
|
||||||
|
dont_filter_partial_edge_chunks: flags & 0x01 != 0,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
3 => {
|
3 => {
|
||||||
@@ -602,6 +609,7 @@ mod tests {
|
|||||||
chunk_index_type: None,
|
chunk_index_type: None,
|
||||||
single_chunk_filtered_size: None,
|
single_chunk_filtered_size: None,
|
||||||
single_chunk_filter_mask: None,
|
single_chunk_filter_mask: None,
|
||||||
|
dont_filter_partial_edge_chunks: false,
|
||||||
}
|
}
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
@@ -679,10 +687,35 @@ mod tests {
|
|||||||
chunk_index_type: Some(1),
|
chunk_index_type: Some(1),
|
||||||
single_chunk_filtered_size: None,
|
single_chunk_filtered_size: None,
|
||||||
single_chunk_filter_mask: None,
|
single_chunk_filter_mask: None,
|
||||||
|
dont_filter_partial_edge_chunks: false,
|
||||||
}
|
}
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn v4_chunked_dont_filter_partial_edge_chunks_flag() {
|
||||||
|
let mut buf = vec![4u8, 2]; // version=4, class=2
|
||||||
|
buf.push(0x01); // flags bit 0 = don't filter partial edge chunks
|
||||||
|
buf.push(2); // dimensionality=2
|
||||||
|
buf.push(4); // dim_size_encoded_length=4
|
||||||
|
buf.extend_from_slice(&5u32.to_le_bytes());
|
||||||
|
buf.extend_from_slice(&4u32.to_le_bytes());
|
||||||
|
buf.push(3); // Fixed Array
|
||||||
|
buf.push(10); // max_dblk_page_nelmts_bits
|
||||||
|
buf.extend_from_slice(&0x3000u64.to_le_bytes());
|
||||||
|
match DataLayout::parse(&buf, 8, 8).unwrap() {
|
||||||
|
DataLayout::Chunked {
|
||||||
|
dont_filter_partial_edge_chunks,
|
||||||
|
btree_address,
|
||||||
|
..
|
||||||
|
} => {
|
||||||
|
assert!(dont_filter_partial_edge_chunks);
|
||||||
|
assert_eq!(btree_address, Some(0x3000));
|
||||||
|
}
|
||||||
|
other => panic!("expected Chunked, got {other:?}"),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn v4_chunked_single_chunk_with_filters() {
|
fn v4_chunked_single_chunk_with_filters() {
|
||||||
let mut buf = vec![4u8, 2]; // version=4, class=2
|
let mut buf = vec![4u8, 2]; // version=4, class=2
|
||||||
@@ -705,6 +738,7 @@ mod tests {
|
|||||||
chunk_index_type: Some(1),
|
chunk_index_type: Some(1),
|
||||||
single_chunk_filtered_size: Some(1024),
|
single_chunk_filtered_size: Some(1024),
|
||||||
single_chunk_filter_mask: Some(0),
|
single_chunk_filter_mask: Some(0),
|
||||||
|
dont_filter_partial_edge_chunks: false,
|
||||||
}
|
}
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -307,6 +307,24 @@ pub fn read_raw_data_selection(
|
|||||||
) -> Result<Vec<u8>, FormatError> {
|
) -> Result<Vec<u8>, FormatError> {
|
||||||
use crate::selection::Selection;
|
use crate::selection::Selection;
|
||||||
|
|
||||||
|
crate::partial_read::validate(selection, &dataspace.dimensions)?;
|
||||||
|
|
||||||
|
// Read only what the selection's bounding box touches when that is
|
||||||
|
// possible; everything below is the decode-everything-then-pick path,
|
||||||
|
// kept for the cases `partial_read` declines.
|
||||||
|
if let Some(selected) = crate::partial_read::read_selection(
|
||||||
|
file_data,
|
||||||
|
layout,
|
||||||
|
dataspace,
|
||||||
|
datatype.type_size() as usize,
|
||||||
|
pipeline,
|
||||||
|
offset_size,
|
||||||
|
length_size,
|
||||||
|
selection,
|
||||||
|
)? {
|
||||||
|
return Ok(selected);
|
||||||
|
}
|
||||||
|
|
||||||
match selection {
|
match selection {
|
||||||
Selection::All => {
|
Selection::All => {
|
||||||
return read_raw_data_full(
|
return read_raw_data_full(
|
||||||
@@ -755,14 +773,7 @@ pub fn read_as_f64_zerocopy<'a>(raw: &'a [u8], datatype: &Datatype) -> Option<&'
|
|||||||
// Only native LE f64 is eligible
|
// Only native LE f64 is eligible
|
||||||
#[cfg(target_endian = "little")]
|
#[cfg(target_endian = "little")]
|
||||||
{
|
{
|
||||||
if !matches!(
|
if !is_native_le_float(datatype, FloatFormat::Double) {
|
||||||
datatype,
|
|
||||||
Datatype::FloatingPoint {
|
|
||||||
size: 8,
|
|
||||||
byte_order: DatatypeByteOrder::LittleEndian,
|
|
||||||
..
|
|
||||||
}
|
|
||||||
) {
|
|
||||||
return None;
|
return None;
|
||||||
}
|
}
|
||||||
if !raw.len().is_multiple_of(8) {
|
if !raw.len().is_multiple_of(8) {
|
||||||
@@ -791,14 +802,7 @@ pub fn read_as_f64_zerocopy<'a>(raw: &'a [u8], datatype: &Datatype) -> Option<&'
|
|||||||
pub fn read_as_f32_zerocopy<'a>(raw: &'a [u8], datatype: &Datatype) -> Option<&'a [f32]> {
|
pub fn read_as_f32_zerocopy<'a>(raw: &'a [u8], datatype: &Datatype) -> Option<&'a [f32]> {
|
||||||
#[cfg(target_endian = "little")]
|
#[cfg(target_endian = "little")]
|
||||||
{
|
{
|
||||||
if !matches!(
|
if !is_native_le_float(datatype, FloatFormat::Single) {
|
||||||
datatype,
|
|
||||||
Datatype::FloatingPoint {
|
|
||||||
size: 4,
|
|
||||||
byte_order: DatatypeByteOrder::LittleEndian,
|
|
||||||
..
|
|
||||||
}
|
|
||||||
) {
|
|
||||||
return None;
|
return None;
|
||||||
}
|
}
|
||||||
if !raw.len().is_multiple_of(4) {
|
if !raw.len().is_multiple_of(4) {
|
||||||
@@ -858,11 +862,35 @@ fn get_size(dt: &Datatype) -> usize {
|
|||||||
dt.type_size() as usize
|
dt.type_size() as usize
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Reinterpret little-endian bytes as `count` native values of `T` on a
|
||||||
|
/// little-endian target, in one copy.
|
||||||
|
///
|
||||||
|
/// The buffer is allocated uninitialised and filled by the copy. It used to be
|
||||||
|
/// `vec![0; count]` first, which for a large dataset meant writing every page
|
||||||
|
/// twice (zero it, then overwrite it) — about as expensive as the copy itself.
|
||||||
|
#[cfg(target_endian = "little")]
|
||||||
|
fn native_le_to_vec<T: Copy>(raw: &[u8], count: usize) -> Vec<T> {
|
||||||
|
let bytes = count * core::mem::size_of::<T>();
|
||||||
|
debug_assert!(bytes <= raw.len());
|
||||||
|
let mut result: Vec<T> = Vec::with_capacity(count);
|
||||||
|
// SAFETY: `result` has capacity for `count` values of `T`, i.e. `bytes`
|
||||||
|
// bytes; `raw` holds at least `bytes` bytes (callers derive `count` from
|
||||||
|
// `raw.len() / size_of::<T>()`); the regions cannot overlap because
|
||||||
|
// `result` was just allocated. Every `T` used here (f32/f64/i32/i64) is
|
||||||
|
// valid for any bit pattern, so after the copy all `count` values are
|
||||||
|
// initialised and `set_len` is sound.
|
||||||
|
unsafe {
|
||||||
|
core::ptr::copy_nonoverlapping(raw.as_ptr(), result.as_mut_ptr().cast::<u8>(), bytes);
|
||||||
|
result.set_len(count);
|
||||||
|
}
|
||||||
|
result
|
||||||
|
}
|
||||||
|
|
||||||
/// Convert raw bytes to `f64` values.
|
/// Convert raw bytes to `f64` values.
|
||||||
pub fn read_as_f64(raw: &[u8], datatype: &Datatype) -> Result<Vec<f64>, FormatError> {
|
pub fn read_as_f64(raw: &[u8], datatype: &Datatype) -> Result<Vec<f64>, FormatError> {
|
||||||
// Array datatypes (e.g. an array-typed compound member) are read as a flat
|
// Array datatypes read as a flat sequence of their base elements, and
|
||||||
// sequence of their base elements.
|
// enumerations (h5py's bool among them) as their integer values.
|
||||||
if let Datatype::Array { base_type, .. } = datatype {
|
if let Datatype::Array { base_type, .. } | Datatype::Enumeration { base_type, .. } = datatype {
|
||||||
return read_as_f64(raw, base_type);
|
return read_as_f64(raw, base_type);
|
||||||
}
|
}
|
||||||
ensure_numeric(datatype, "FloatingPoint or FixedPoint")?;
|
ensure_numeric(datatype, "FloatingPoint or FixedPoint")?;
|
||||||
@@ -877,27 +905,19 @@ pub fn read_as_f64(raw: &[u8], datatype: &Datatype) -> Result<Vec<f64>, FormatEr
|
|||||||
|
|
||||||
// Fast path: native-endian f64 — single bulk memcpy
|
// Fast path: native-endian f64 — single bulk memcpy
|
||||||
#[cfg(target_endian = "little")]
|
#[cfg(target_endian = "little")]
|
||||||
if matches!(
|
if is_native_le_float(datatype, FloatFormat::Double) {
|
||||||
datatype,
|
return Ok(native_le_to_vec::<f64>(raw, count));
|
||||||
Datatype::FloatingPoint {
|
|
||||||
size: 8,
|
|
||||||
byte_order: DatatypeByteOrder::LittleEndian,
|
|
||||||
..
|
|
||||||
}
|
|
||||||
) {
|
|
||||||
let mut result = vec![0.0f64; count];
|
|
||||||
// SAFETY: On LE platforms, f64 in-memory representation matches LE bytes.
|
|
||||||
// We copy raw bytes directly into the f64 buffer.
|
|
||||||
// SAFETY: The byte slice is properly aligned for this type and the length is divisible by size_of::<T>().
|
|
||||||
unsafe {
|
|
||||||
core::ptr::copy_nonoverlapping(raw.as_ptr(), result.as_mut_ptr() as *mut u8, raw.len());
|
|
||||||
}
|
|
||||||
return Ok(result);
|
|
||||||
}
|
}
|
||||||
|
|
||||||
let order = get_byte_order(datatype);
|
let order = get_byte_order(datatype);
|
||||||
let mut result = Vec::with_capacity(count);
|
let mut result = Vec::with_capacity(count);
|
||||||
|
if let Datatype::FloatingPoint { .. } = datatype {
|
||||||
|
let format = FloatFormat::of(datatype)?;
|
||||||
|
for chunk in raw.chunks_exact(elem_size) {
|
||||||
|
result.push(format.decode(chunk, &order));
|
||||||
|
}
|
||||||
|
return Ok(result);
|
||||||
|
}
|
||||||
for i in 0..count {
|
for i in 0..count {
|
||||||
let chunk = &raw[i * elem_size..(i + 1) * elem_size];
|
let chunk = &raw[i * elem_size..(i + 1) * elem_size];
|
||||||
let val = convert_to_f64(chunk, datatype, &order)?;
|
let val = convert_to_f64(chunk, datatype, &order)?;
|
||||||
@@ -912,18 +932,7 @@ fn convert_to_f64(
|
|||||||
order: &DatatypeByteOrder,
|
order: &DatatypeByteOrder,
|
||||||
) -> Result<f64, FormatError> {
|
) -> Result<f64, FormatError> {
|
||||||
match dt {
|
match dt {
|
||||||
Datatype::FloatingPoint { size, .. } => match size {
|
Datatype::FloatingPoint { .. } => Ok(FloatFormat::of(dt)?.decode(bytes, order)),
|
||||||
4 => {
|
|
||||||
let v = read_f32_bytes(bytes, order);
|
|
||||||
Ok(v as f64)
|
|
||||||
}
|
|
||||||
8 => Ok(read_f64_bytes(bytes, order)),
|
|
||||||
2 => Ok(read_f16_bytes(bytes, order) as f64),
|
|
||||||
_ => Err(FormatError::DataSizeMismatch {
|
|
||||||
expected: 8,
|
|
||||||
actual: *size as usize,
|
|
||||||
}),
|
|
||||||
},
|
|
||||||
Datatype::FixedPoint {
|
Datatype::FixedPoint {
|
||||||
size,
|
size,
|
||||||
signed,
|
signed,
|
||||||
@@ -947,9 +956,83 @@ fn convert_to_f64(
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// One numeric element as stored, before conversion to the caller's type.
|
||||||
|
#[derive(Debug, Clone, Copy, PartialEq)]
|
||||||
|
enum Scalar {
|
||||||
|
Signed(i64),
|
||||||
|
Unsigned(u64),
|
||||||
|
Float(f64),
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Scalar {
|
||||||
|
// Every conversion follows libhdf5's default (hard) conversions: a value
|
||||||
|
// outside the target type's range saturates to its minimum or maximum —
|
||||||
|
// including a negative value read as unsigned, which reads as 0 — rather
|
||||||
|
// than being truncated to its low bits. Floats truncate toward zero; NaN
|
||||||
|
// converts to 0 (libhdf5 leaves that case to the C cast, whose result is
|
||||||
|
// platform-dependent).
|
||||||
|
|
||||||
|
fn to_i64(self) -> i64 {
|
||||||
|
match self {
|
||||||
|
Scalar::Signed(v) => v,
|
||||||
|
Scalar::Unsigned(v) => i64::try_from(v).unwrap_or(i64::MAX),
|
||||||
|
Scalar::Float(v) => v as i64,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn to_u64(self) -> u64 {
|
||||||
|
match self {
|
||||||
|
Scalar::Signed(v) => u64::try_from(v).unwrap_or(0),
|
||||||
|
Scalar::Unsigned(v) => v,
|
||||||
|
Scalar::Float(v) => v as u64,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn to_i32(self) -> i32 {
|
||||||
|
match self {
|
||||||
|
Scalar::Signed(v) => v.clamp(i32::MIN.into(), i32::MAX.into()) as i32,
|
||||||
|
Scalar::Unsigned(v) => i32::try_from(v).unwrap_or(i32::MAX),
|
||||||
|
Scalar::Float(v) => v as i32,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Decode one element of a numeric datatype.
|
||||||
|
fn decode_scalar(
|
||||||
|
bytes: &[u8],
|
||||||
|
dt: &Datatype,
|
||||||
|
order: &DatatypeByteOrder,
|
||||||
|
) -> Result<Scalar, FormatError> {
|
||||||
|
match dt {
|
||||||
|
Datatype::FixedPoint {
|
||||||
|
size,
|
||||||
|
signed,
|
||||||
|
bit_offset,
|
||||||
|
bit_precision,
|
||||||
|
..
|
||||||
|
} => {
|
||||||
|
let full = read_unsigned_int(bytes, *size as usize, order);
|
||||||
|
let (off, prec) = effective_bits(*size as usize, *bit_offset, *bit_precision);
|
||||||
|
Ok(if *signed {
|
||||||
|
Scalar::Signed(extract_signed(full, off, prec))
|
||||||
|
} else {
|
||||||
|
Scalar::Unsigned(extract_unsigned(full, off, prec))
|
||||||
|
})
|
||||||
|
}
|
||||||
|
_ => convert_to_f64(bytes, dt, order).map(Scalar::Float),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
/// Convert raw bytes to `i64` values.
|
/// Convert raw bytes to `i64` values.
|
||||||
|
///
|
||||||
|
/// Values are converted the way libhdf5 converts them: integers outside the
|
||||||
|
/// target range saturate at its minimum or maximum (a negative value read as
|
||||||
|
/// unsigned is 0), and floating-point data is truncated toward zero and
|
||||||
|
/// saturated, with NaN read as 0.
|
||||||
pub fn read_as_i64(raw: &[u8], datatype: &Datatype) -> Result<Vec<i64>, FormatError> {
|
pub fn read_as_i64(raw: &[u8], datatype: &Datatype) -> Result<Vec<i64>, FormatError> {
|
||||||
if let Datatype::Array { base_type, .. } = datatype {
|
// Array datatypes read as a flat sequence of their base elements, and
|
||||||
|
// enumerations (h5py's bool among them) as their integer values.
|
||||||
|
if let Datatype::Array { base_type, .. } | Datatype::Enumeration { base_type, .. } = datatype {
|
||||||
return read_as_i64(raw, base_type);
|
return read_as_i64(raw, base_type);
|
||||||
}
|
}
|
||||||
ensure_numeric(datatype, "FixedPoint (signed)")?;
|
ensure_numeric(datatype, "FixedPoint (signed)")?;
|
||||||
@@ -975,28 +1058,28 @@ pub fn read_as_i64(raw: &[u8], datatype: &Datatype) -> Result<Vec<i64>, FormatEr
|
|||||||
}
|
}
|
||||||
)
|
)
|
||||||
{
|
{
|
||||||
let mut result = vec![0i64; count];
|
return Ok(native_le_to_vec::<i64>(raw, count));
|
||||||
// SAFETY: The byte slice is properly aligned for this type and the length is divisible by size_of::<T>().
|
|
||||||
unsafe {
|
|
||||||
core::ptr::copy_nonoverlapping(raw.as_ptr(), result.as_mut_ptr() as *mut u8, raw.len());
|
|
||||||
}
|
|
||||||
return Ok(result);
|
|
||||||
}
|
}
|
||||||
|
|
||||||
let order = get_byte_order(datatype);
|
let order = get_byte_order(datatype);
|
||||||
let (off, prec) = fixed_bits(datatype);
|
|
||||||
let mut result = Vec::with_capacity(count);
|
let mut result = Vec::with_capacity(count);
|
||||||
for i in 0..count {
|
for i in 0..count {
|
||||||
let chunk = &raw[i * elem_size..(i + 1) * elem_size];
|
let chunk = &raw[i * elem_size..(i + 1) * elem_size];
|
||||||
let full = read_unsigned_int(chunk, elem_size, &order);
|
result.push(decode_scalar(chunk, datatype, &order)?.to_i64());
|
||||||
result.push(extract_signed(full, off, prec));
|
|
||||||
}
|
}
|
||||||
Ok(result)
|
Ok(result)
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Convert raw bytes to `u64` values.
|
/// Convert raw bytes to `u64` values.
|
||||||
|
///
|
||||||
|
/// Values are converted the way libhdf5 converts them: integers outside the
|
||||||
|
/// target range saturate at its minimum or maximum (a negative value read as
|
||||||
|
/// unsigned is 0), and floating-point data is truncated toward zero and
|
||||||
|
/// saturated, with NaN read as 0.
|
||||||
pub fn read_as_u64(raw: &[u8], datatype: &Datatype) -> Result<Vec<u64>, FormatError> {
|
pub fn read_as_u64(raw: &[u8], datatype: &Datatype) -> Result<Vec<u64>, FormatError> {
|
||||||
if let Datatype::Array { base_type, .. } = datatype {
|
// Array datatypes read as a flat sequence of their base elements, and
|
||||||
|
// enumerations (h5py's bool among them) as their integer values.
|
||||||
|
if let Datatype::Array { base_type, .. } | Datatype::Enumeration { base_type, .. } = datatype {
|
||||||
return read_as_u64(raw, base_type);
|
return read_as_u64(raw, base_type);
|
||||||
}
|
}
|
||||||
ensure_numeric(datatype, "FixedPoint (unsigned)")?;
|
ensure_numeric(datatype, "FixedPoint (unsigned)")?;
|
||||||
@@ -1009,19 +1092,19 @@ pub fn read_as_u64(raw: &[u8], datatype: &Datatype) -> Result<Vec<u64>, FormatEr
|
|||||||
}
|
}
|
||||||
let count = raw.len() / elem_size;
|
let count = raw.len() / elem_size;
|
||||||
let order = get_byte_order(datatype);
|
let order = get_byte_order(datatype);
|
||||||
let (off, prec) = fixed_bits(datatype);
|
|
||||||
let mut result = Vec::with_capacity(count);
|
let mut result = Vec::with_capacity(count);
|
||||||
for i in 0..count {
|
for i in 0..count {
|
||||||
let chunk = &raw[i * elem_size..(i + 1) * elem_size];
|
let chunk = &raw[i * elem_size..(i + 1) * elem_size];
|
||||||
let full = read_unsigned_int(chunk, elem_size, &order);
|
result.push(decode_scalar(chunk, datatype, &order)?.to_u64());
|
||||||
result.push(extract_unsigned(full, off, prec));
|
|
||||||
}
|
}
|
||||||
Ok(result)
|
Ok(result)
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Convert raw bytes to `f32` values.
|
/// Convert raw bytes to `f32` values.
|
||||||
pub fn read_as_f32(raw: &[u8], datatype: &Datatype) -> Result<Vec<f32>, FormatError> {
|
pub fn read_as_f32(raw: &[u8], datatype: &Datatype) -> Result<Vec<f32>, FormatError> {
|
||||||
if let Datatype::Array { base_type, .. } = datatype {
|
// Array datatypes read as a flat sequence of their base elements, and
|
||||||
|
// enumerations (h5py's bool among them) as their integer values.
|
||||||
|
if let Datatype::Array { base_type, .. } | Datatype::Enumeration { base_type, .. } = datatype {
|
||||||
return read_as_f32(raw, base_type);
|
return read_as_f32(raw, base_type);
|
||||||
}
|
}
|
||||||
ensure_numeric(datatype, "FloatingPoint")?;
|
ensure_numeric(datatype, "FloatingPoint")?;
|
||||||
@@ -1036,36 +1119,36 @@ pub fn read_as_f32(raw: &[u8], datatype: &Datatype) -> Result<Vec<f32>, FormatEr
|
|||||||
|
|
||||||
// Fast path: native-endian f32 — single bulk memcpy
|
// Fast path: native-endian f32 — single bulk memcpy
|
||||||
#[cfg(target_endian = "little")]
|
#[cfg(target_endian = "little")]
|
||||||
if matches!(
|
if is_native_le_float(datatype, FloatFormat::Single) {
|
||||||
datatype,
|
return Ok(native_le_to_vec::<f32>(raw, count));
|
||||||
Datatype::FloatingPoint {
|
}
|
||||||
size: 4,
|
// Little-endian IEEE half precision (numpy float16): widen directly.
|
||||||
byte_order: DatatypeByteOrder::LittleEndian,
|
if is_native_le_float(datatype, FloatFormat::Half) {
|
||||||
..
|
let (halves, _) = raw[..count * 2].as_chunks::<2>();
|
||||||
}
|
return Ok(halves
|
||||||
) {
|
.iter()
|
||||||
let mut result = vec![0.0f32; count];
|
.map(|&b| f16_bits_to_f32(u16::from_le_bytes(b)))
|
||||||
// SAFETY: The byte slice is properly aligned for this type and the length is divisible by size_of::<T>().
|
.collect());
|
||||||
unsafe {
|
|
||||||
core::ptr::copy_nonoverlapping(raw.as_ptr(), result.as_mut_ptr() as *mut u8, raw.len());
|
|
||||||
}
|
|
||||||
return Ok(result);
|
|
||||||
}
|
}
|
||||||
|
|
||||||
let order = get_byte_order(datatype);
|
let order = get_byte_order(datatype);
|
||||||
let mut result = Vec::with_capacity(count);
|
let mut result = Vec::with_capacity(count);
|
||||||
|
if let Datatype::FloatingPoint { .. } = datatype {
|
||||||
|
let format = FloatFormat::of(datatype)?;
|
||||||
|
for chunk in raw.chunks_exact(elem_size) {
|
||||||
|
result.push(match format {
|
||||||
|
FloatFormat::Single => read_f32_bytes(chunk, &order),
|
||||||
|
FloatFormat::Half => read_f16_bytes(chunk, &order),
|
||||||
|
// Double rounds; every other supported layout (bfloat16, FP8)
|
||||||
|
// is exact in f32.
|
||||||
|
_ => format.decode(chunk, &order) as f32,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
return Ok(result);
|
||||||
|
}
|
||||||
for i in 0..count {
|
for i in 0..count {
|
||||||
let chunk = &raw[i * elem_size..(i + 1) * elem_size];
|
let chunk = &raw[i * elem_size..(i + 1) * elem_size];
|
||||||
match datatype {
|
match datatype {
|
||||||
Datatype::FloatingPoint { size: 4, .. } => {
|
|
||||||
result.push(read_f32_bytes(chunk, &order));
|
|
||||||
}
|
|
||||||
Datatype::FloatingPoint { size: 8, .. } => {
|
|
||||||
result.push(read_f64_bytes(chunk, &order) as f32);
|
|
||||||
}
|
|
||||||
Datatype::FloatingPoint { size: 2, .. } => {
|
|
||||||
result.push(read_f16_bytes(chunk, &order));
|
|
||||||
}
|
|
||||||
Datatype::FixedPoint {
|
Datatype::FixedPoint {
|
||||||
signed: true,
|
signed: true,
|
||||||
size,
|
size,
|
||||||
@@ -1100,8 +1183,15 @@ pub fn read_as_f32(raw: &[u8], datatype: &Datatype) -> Result<Vec<f32>, FormatEr
|
|||||||
}
|
}
|
||||||
|
|
||||||
/// Convert raw bytes to `i32` values.
|
/// Convert raw bytes to `i32` values.
|
||||||
|
///
|
||||||
|
/// Values are converted the way libhdf5 converts them: integers outside the
|
||||||
|
/// target range saturate at its minimum or maximum (a negative value read as
|
||||||
|
/// unsigned is 0), and floating-point data is truncated toward zero and
|
||||||
|
/// saturated, with NaN read as 0.
|
||||||
pub fn read_as_i32(raw: &[u8], datatype: &Datatype) -> Result<Vec<i32>, FormatError> {
|
pub fn read_as_i32(raw: &[u8], datatype: &Datatype) -> Result<Vec<i32>, FormatError> {
|
||||||
if let Datatype::Array { base_type, .. } = datatype {
|
// Array datatypes read as a flat sequence of their base elements, and
|
||||||
|
// enumerations (h5py's bool among them) as their integer values.
|
||||||
|
if let Datatype::Array { base_type, .. } | Datatype::Enumeration { base_type, .. } = datatype {
|
||||||
return read_as_i32(raw, base_type);
|
return read_as_i32(raw, base_type);
|
||||||
}
|
}
|
||||||
ensure_numeric(datatype, "FixedPoint")?;
|
ensure_numeric(datatype, "FixedPoint")?;
|
||||||
@@ -1122,25 +1212,19 @@ pub fn read_as_i32(raw: &[u8], datatype: &Datatype) -> Result<Vec<i32>, FormatEr
|
|||||||
datatype,
|
datatype,
|
||||||
Datatype::FixedPoint {
|
Datatype::FixedPoint {
|
||||||
byte_order: DatatypeByteOrder::LittleEndian,
|
byte_order: DatatypeByteOrder::LittleEndian,
|
||||||
|
signed: true,
|
||||||
..
|
..
|
||||||
}
|
}
|
||||||
)
|
)
|
||||||
{
|
{
|
||||||
let mut result = vec![0i32; count];
|
return Ok(native_le_to_vec::<i32>(raw, count));
|
||||||
// SAFETY: The byte slice is properly aligned for this type and the length is divisible by size_of::<T>().
|
|
||||||
unsafe {
|
|
||||||
core::ptr::copy_nonoverlapping(raw.as_ptr(), result.as_mut_ptr() as *mut u8, raw.len());
|
|
||||||
}
|
|
||||||
return Ok(result);
|
|
||||||
}
|
}
|
||||||
|
|
||||||
let order = get_byte_order(datatype);
|
let order = get_byte_order(datatype);
|
||||||
let (off, prec) = fixed_bits(datatype);
|
|
||||||
let mut result = Vec::with_capacity(count);
|
let mut result = Vec::with_capacity(count);
|
||||||
for i in 0..count {
|
for i in 0..count {
|
||||||
let chunk = &raw[i * elem_size..(i + 1) * elem_size];
|
let chunk = &raw[i * elem_size..(i + 1) * elem_size];
|
||||||
let full = read_unsigned_int(chunk, elem_size, &order);
|
result.push(decode_scalar(chunk, datatype, &order)?.to_i32());
|
||||||
result.push(extract_signed(full, off, prec) as i32);
|
|
||||||
}
|
}
|
||||||
Ok(result)
|
Ok(result)
|
||||||
}
|
}
|
||||||
@@ -1407,6 +1491,26 @@ pub fn read_object_references(
|
|||||||
}
|
}
|
||||||
Ok(result)
|
Ok(result)
|
||||||
}
|
}
|
||||||
|
Datatype::Reference {
|
||||||
|
ref_type: crate::datatype::ReferenceType::Object2,
|
||||||
|
size,
|
||||||
|
} => {
|
||||||
|
let elem_size = *size as usize;
|
||||||
|
if elem_size == 0 {
|
||||||
|
return Ok(Vec::new());
|
||||||
|
}
|
||||||
|
if !raw.len().is_multiple_of(elem_size) {
|
||||||
|
return Err(FormatError::DataSizeMismatch {
|
||||||
|
expected: 0,
|
||||||
|
actual: raw.len(),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
raw.chunks_exact(elem_size)
|
||||||
|
.map(|element| {
|
||||||
|
decode_std_object_ref(element).map(|address| ObjectReference { address })
|
||||||
|
})
|
||||||
|
.collect()
|
||||||
|
}
|
||||||
_ => Err(FormatError::TypeMismatch {
|
_ => Err(FormatError::TypeMismatch {
|
||||||
expected: "Reference(Object)",
|
expected: "Reference(Object)",
|
||||||
actual: datatype_name(datatype),
|
actual: datatype_name(datatype),
|
||||||
@@ -1414,6 +1518,46 @@ pub fn read_object_references(
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Decode one `H5T_STD_REF` object reference as stored in a dataset:
|
||||||
|
/// `type(1) flags(1) token_size(1) token(token_size)`, zero-padded to the
|
||||||
|
/// element size. For a reference within the same file the token is the target
|
||||||
|
/// object's header address. An all-zero element is a null reference and
|
||||||
|
/// decodes to the undefined address (`u64::MAX`).
|
||||||
|
fn decode_std_object_ref(element: &[u8]) -> Result<u64, FormatError> {
|
||||||
|
const STD_REF_OBJECT: u8 = 2;
|
||||||
|
const FLAG_EXTERNAL: u8 = 0x01;
|
||||||
|
if element.iter().all(|&b| b == 0) {
|
||||||
|
return Ok(u64::MAX);
|
||||||
|
}
|
||||||
|
let [ref_type, flags, token_size, token @ ..] = element else {
|
||||||
|
return Err(FormatError::UnexpectedEof {
|
||||||
|
expected: 3,
|
||||||
|
available: element.len(),
|
||||||
|
});
|
||||||
|
};
|
||||||
|
if *ref_type != STD_REF_OBJECT {
|
||||||
|
return Err(FormatError::InvalidReferenceType(*ref_type));
|
||||||
|
}
|
||||||
|
if flags & FLAG_EXTERNAL != 0 {
|
||||||
|
// Carries a file name as well; nothing here follows those.
|
||||||
|
return Err(FormatError::TypeMismatch {
|
||||||
|
expected: "object reference within this file",
|
||||||
|
actual: "external object reference",
|
||||||
|
});
|
||||||
|
}
|
||||||
|
let n = *token_size as usize;
|
||||||
|
if n == 0 || n > 8 || n > token.len() {
|
||||||
|
return Err(FormatError::UnexpectedEof {
|
||||||
|
expected: 3 + n,
|
||||||
|
available: element.len(),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
Ok(token[..n]
|
||||||
|
.iter()
|
||||||
|
.rev()
|
||||||
|
.fold(0u64, |addr, &byte| (addr << 8) | u64::from(byte)))
|
||||||
|
}
|
||||||
|
|
||||||
/// Read region references from raw bytes.
|
/// Read region references from raw bytes.
|
||||||
///
|
///
|
||||||
/// Region references encode a dataset selection (hyperslab, point list, etc.)
|
/// Region references encode a dataset selection (hyperslab, point list, etc.)
|
||||||
@@ -1521,6 +1665,174 @@ fn reorder_bytes(bytes: &[u8], order: &DatatypeByteOrder) -> [u8; 8] {
|
|||||||
buf
|
buf
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// How the bits of a floating-point datatype are laid out, read from the
|
||||||
|
/// datatype message's fields rather than assumed from its size (a 2-byte
|
||||||
|
/// float may be IEEE half or bfloat16).
|
||||||
|
#[derive(Debug, Clone, Copy, PartialEq)]
|
||||||
|
enum FloatFormat {
|
||||||
|
/// IEEE-754 binary16.
|
||||||
|
Half,
|
||||||
|
/// IEEE-754 binary32.
|
||||||
|
Single,
|
||||||
|
/// IEEE-754 binary64.
|
||||||
|
Double,
|
||||||
|
/// Any other IEEE-style layout (implied leading mantissa bit, all-ones
|
||||||
|
/// exponent for infinity/NaN) whose values are all exact in `f64`:
|
||||||
|
/// bfloat16, the FP8 formats, and similar.
|
||||||
|
Other(FloatLayout),
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Copy, PartialEq)]
|
||||||
|
struct FloatLayout {
|
||||||
|
exponent_location: u32,
|
||||||
|
exponent_size: u32,
|
||||||
|
mantissa_location: u32,
|
||||||
|
mantissa_size: u32,
|
||||||
|
exponent_bias: u32,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl FloatFormat {
|
||||||
|
fn of(dt: &Datatype) -> Result<FloatFormat, FormatError> {
|
||||||
|
let Datatype::FloatingPoint {
|
||||||
|
size,
|
||||||
|
exponent_location,
|
||||||
|
exponent_size,
|
||||||
|
mantissa_location,
|
||||||
|
mantissa_size,
|
||||||
|
exponent_bias,
|
||||||
|
..
|
||||||
|
} = dt
|
||||||
|
else {
|
||||||
|
return Err(FormatError::TypeMismatch {
|
||||||
|
expected: "FloatingPoint",
|
||||||
|
actual: datatype_name(dt),
|
||||||
|
});
|
||||||
|
};
|
||||||
|
let layout = FloatLayout {
|
||||||
|
exponent_location: u32::from(*exponent_location),
|
||||||
|
exponent_size: u32::from(*exponent_size),
|
||||||
|
mantissa_location: u32::from(*mantissa_location),
|
||||||
|
mantissa_size: u32::from(*mantissa_size),
|
||||||
|
exponent_bias: *exponent_bias,
|
||||||
|
};
|
||||||
|
let fields = (
|
||||||
|
layout.exponent_location,
|
||||||
|
layout.exponent_size,
|
||||||
|
layout.mantissa_location,
|
||||||
|
layout.mantissa_size,
|
||||||
|
layout.exponent_bias,
|
||||||
|
);
|
||||||
|
let bits = size.saturating_mul(8);
|
||||||
|
// The sign bit is not kept in `Datatype`; every standard layout has it
|
||||||
|
// directly above the exponent, with the mantissa below.
|
||||||
|
let well_formed = layout.exponent_size > 0
|
||||||
|
&& layout.mantissa_size > 0
|
||||||
|
&& layout.mantissa_location + layout.mantissa_size <= layout.exponent_location
|
||||||
|
&& layout.exponent_location + layout.exponent_size < bits;
|
||||||
|
match (size, fields) {
|
||||||
|
(2, (10, 5, 0, 10, 15)) => Ok(FloatFormat::Half),
|
||||||
|
(4, (23, 8, 0, 23, 127)) => Ok(FloatFormat::Single),
|
||||||
|
(8, (52, 11, 0, 52, 1023)) => Ok(FloatFormat::Double),
|
||||||
|
_ if well_formed
|
||||||
|
&& *size <= 8
|
||||||
|
&& layout.exponent_size <= 11
|
||||||
|
&& layout.mantissa_size <= 52 =>
|
||||||
|
{
|
||||||
|
Ok(FloatFormat::Other(layout))
|
||||||
|
}
|
||||||
|
// Fields that cannot describe any float (e.g. left zeroed by a
|
||||||
|
// hand-built datatype): fall back to the IEEE type of that size.
|
||||||
|
(2, _) if !well_formed => Ok(FloatFormat::Half),
|
||||||
|
(4, _) if !well_formed => Ok(FloatFormat::Single),
|
||||||
|
(8, _) if !well_formed => Ok(FloatFormat::Double),
|
||||||
|
// x87 80-bit extended, binary128, ...: not representable in f64.
|
||||||
|
_ => Err(FormatError::TypeMismatch {
|
||||||
|
expected: "floating point of at most 64 bits (IEEE-style layout)",
|
||||||
|
actual: "FloatingPoint",
|
||||||
|
}),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn decode(self, bytes: &[u8], order: &DatatypeByteOrder) -> f64 {
|
||||||
|
match self {
|
||||||
|
FloatFormat::Half => f64::from(read_f16_bytes(bytes, order)),
|
||||||
|
FloatFormat::Single => f64::from(read_f32_bytes(bytes, order)),
|
||||||
|
FloatFormat::Double => read_f64_bytes(bytes, order),
|
||||||
|
FloatFormat::Other(layout) => {
|
||||||
|
layout.decode(read_unsigned_int(bytes, bytes.len(), order))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl FloatLayout {
|
||||||
|
/// Decode the value held in the low `size * 8` bits of `bits`.
|
||||||
|
fn decode(self, bits: u64) -> f64 {
|
||||||
|
let field = |location: u32, size: u32| (bits >> location) & ((1u64 << size) - 1);
|
||||||
|
let exponent = field(self.exponent_location, self.exponent_size);
|
||||||
|
let mantissa = field(self.mantissa_location, self.mantissa_size);
|
||||||
|
let negative = field(self.exponent_location + self.exponent_size, 1) == 1;
|
||||||
|
let max_exponent = (1u64 << self.exponent_size) - 1;
|
||||||
|
let magnitude = if exponent == max_exponent {
|
||||||
|
if mantissa == 0 {
|
||||||
|
f64::INFINITY
|
||||||
|
} else {
|
||||||
|
f64::NAN
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
let bias = i64::from(self.exponent_bias);
|
||||||
|
let msize = i64::from(self.mantissa_size);
|
||||||
|
// value = significand * 2^power, with an implied leading 1 unless
|
||||||
|
// the number is subnormal (exponent field 0).
|
||||||
|
let (significand, power) = if exponent == 0 {
|
||||||
|
(mantissa, 1 - bias - msize)
|
||||||
|
} else {
|
||||||
|
(
|
||||||
|
mantissa | (1u64 << self.mantissa_size),
|
||||||
|
exponent as i64 - bias - msize,
|
||||||
|
)
|
||||||
|
};
|
||||||
|
scale_by_pow2(significand as f64, power)
|
||||||
|
};
|
||||||
|
if negative { -magnitude } else { magnitude }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// `x * 2^power` without `std` (no `powi`/`libm`). `x` is a non-negative
|
||||||
|
/// integer below 2^53, so it is exact.
|
||||||
|
fn scale_by_pow2(x: f64, power: i64) -> f64 {
|
||||||
|
if x == 0.0 || power < -1200 {
|
||||||
|
return 0.0;
|
||||||
|
}
|
||||||
|
if power > 1100 {
|
||||||
|
return f64::INFINITY;
|
||||||
|
}
|
||||||
|
let pow2 = |p: i64| f64::from_bits(((p + 1023) as u64) << 52);
|
||||||
|
let mut x = x;
|
||||||
|
let mut power = power;
|
||||||
|
while power > 1023 {
|
||||||
|
x *= pow2(1023);
|
||||||
|
power -= 1023;
|
||||||
|
}
|
||||||
|
while power < -1022 {
|
||||||
|
x *= pow2(-1022);
|
||||||
|
power += 1022;
|
||||||
|
}
|
||||||
|
x * pow2(power)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Whether `datatype` is the little-endian IEEE float `format`, whose bytes
|
||||||
|
/// can be copied straight into native values on a little-endian target.
|
||||||
|
fn is_native_le_float(datatype: &Datatype, format: FloatFormat) -> bool {
|
||||||
|
matches!(
|
||||||
|
datatype,
|
||||||
|
Datatype::FloatingPoint {
|
||||||
|
byte_order: DatatypeByteOrder::LittleEndian,
|
||||||
|
..
|
||||||
|
}
|
||||||
|
) && FloatFormat::of(datatype).is_ok_and(|f| f == format)
|
||||||
|
}
|
||||||
|
|
||||||
fn read_f64_bytes(bytes: &[u8], order: &DatatypeByteOrder) -> f64 {
|
fn read_f64_bytes(bytes: &[u8], order: &DatatypeByteOrder) -> f64 {
|
||||||
let buf = reorder_bytes(bytes, order);
|
let buf = reorder_bytes(bytes, order);
|
||||||
f64::from_le_bytes(buf)
|
f64::from_le_bytes(buf)
|
||||||
@@ -1542,36 +1854,7 @@ fn read_f16_bytes(bytes: &[u8], order: &DatatypeByteOrder) -> f32 {
|
|||||||
f16_bits_to_f32(u16::from_le_bytes(buf))
|
f16_bits_to_f32(u16::from_le_bytes(buf))
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Convert the bit pattern of an IEEE-754 half (binary16) to an `f32`.
|
use crate::float16::f16_bits_to_f32;
|
||||||
fn f16_bits_to_f32(h: u16) -> f32 {
|
|
||||||
let h = h as u32;
|
|
||||||
let sign = (h & 0x8000) << 16;
|
|
||||||
let exp = (h >> 10) & 0x1f;
|
|
||||||
let mant = h & 0x3ff;
|
|
||||||
let bits = if exp == 0 {
|
|
||||||
if mant == 0 {
|
|
||||||
sign // signed zero
|
|
||||||
} else {
|
|
||||||
// Subnormal: normalize into an f32 normal.
|
|
||||||
let mut e: i32 = -1;
|
|
||||||
let mut m = mant;
|
|
||||||
loop {
|
|
||||||
e += 1;
|
|
||||||
m <<= 1;
|
|
||||||
if m & 0x400 != 0 {
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
let m = m & 0x3ff;
|
|
||||||
sign | (((127 - 15 - e) as u32) << 23) | (m << 13)
|
|
||||||
}
|
|
||||||
} else if exp == 0x1f {
|
|
||||||
sign | 0x7f80_0000 | (mant << 13) // inf / NaN
|
|
||||||
} else {
|
|
||||||
sign | ((exp + (127 - 15)) << 23) | (mant << 13)
|
|
||||||
};
|
|
||||||
f32::from_bits(bits)
|
|
||||||
}
|
|
||||||
|
|
||||||
fn read_f32_bytes(bytes: &[u8], order: &DatatypeByteOrder) -> f32 {
|
fn read_f32_bytes(bytes: &[u8], order: &DatatypeByteOrder) -> f32 {
|
||||||
let mut buf = [0u8; 4];
|
let mut buf = [0u8; 4];
|
||||||
@@ -1600,20 +1883,6 @@ fn effective_bits(size: usize, bit_offset: u16, bit_precision: u16) -> (u32, u32
|
|||||||
(bit_offset as u32, prec)
|
(bit_offset as u32, prec)
|
||||||
}
|
}
|
||||||
|
|
||||||
/// `(bit_offset, bit_precision)` for a fixed-point datatype, full width for
|
|
||||||
/// other types.
|
|
||||||
fn fixed_bits(datatype: &Datatype) -> (u32, u32) {
|
|
||||||
match datatype {
|
|
||||||
Datatype::FixedPoint {
|
|
||||||
size,
|
|
||||||
bit_offset,
|
|
||||||
bit_precision,
|
|
||||||
..
|
|
||||||
} => effective_bits(*size as usize, *bit_offset, *bit_precision),
|
|
||||||
_ => (0, 0),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Whether a datatype occupies its full storage width (bit offset 0, precision
|
/// Whether a datatype occupies its full storage width (bit offset 0, precision
|
||||||
/// == size·8), in which case the bulk-copy fast read paths apply. Non
|
/// == size·8), in which case the bulk-copy fast read paths apply. Non
|
||||||
/// fixed-point types are treated as full width.
|
/// fixed-point types are treated as full width.
|
||||||
@@ -1826,6 +2095,67 @@ mod tests {
|
|||||||
assert_eq!(read_as_u64(&raw, &dt).unwrap(), vec![4095, 1, 2048]);
|
assert_eq!(read_as_u64(&raw, &dt).unwrap(), vec![4095, 1, 2048]);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn float_to_int_truncates_and_saturates() {
|
||||||
|
// Values libhdf5 hands to an undefined C cast: NaN reads as 0 and
|
||||||
|
// exactly 2^63 saturates instead of wrapping to i64::MIN.
|
||||||
|
let dt = make_f64_le_type();
|
||||||
|
let vals = [f64::NAN, 2f64.powi(63), -2.5, 2.0f64.powi(64)];
|
||||||
|
let raw: Vec<u8> = vals.iter().flat_map(|v| v.to_le_bytes()).collect();
|
||||||
|
assert_eq!(
|
||||||
|
read_as_i64(&raw, &dt).unwrap(),
|
||||||
|
vec![0, i64::MAX, -2, i64::MAX]
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
read_as_u64(&raw, &dt).unwrap(),
|
||||||
|
vec![0, 1 << 63, 0, u64::MAX]
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
read_as_i32(&raw, &dt).unwrap(),
|
||||||
|
vec![0, i32::MAX, -2, i32::MAX]
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn bfloat16_and_fp8_decode_by_fields() {
|
||||||
|
// bfloat16 is a 2-byte float that is not IEEE half.
|
||||||
|
let bf16 = Datatype::FloatingPoint {
|
||||||
|
size: 2,
|
||||||
|
byte_order: DatatypeByteOrder::LittleEndian,
|
||||||
|
bit_offset: 0,
|
||||||
|
bit_precision: 16,
|
||||||
|
exponent_location: 7,
|
||||||
|
exponent_size: 8,
|
||||||
|
mantissa_location: 0,
|
||||||
|
mantissa_size: 7,
|
||||||
|
exponent_bias: 127,
|
||||||
|
};
|
||||||
|
let raw: Vec<u8> = [0x3FC0u16, 0xC010, 0x7F80, 0x0001]
|
||||||
|
.iter()
|
||||||
|
.flat_map(|v| v.to_le_bytes())
|
||||||
|
.collect();
|
||||||
|
let got = read_as_f64(&raw, &bf16).unwrap();
|
||||||
|
assert_eq!(&got[..3], &[1.5, -2.25, f64::INFINITY]);
|
||||||
|
assert_eq!(got[3], 2f64.powi(-133)); // smallest subnormal
|
||||||
|
assert_eq!(read_as_f32(&raw, &bf16).unwrap()[..2], [1.5, -2.25]);
|
||||||
|
|
||||||
|
// FP8 E4M3: 1, -1, 2, 0, NaN (IEEE-style, as libhdf5 treats it).
|
||||||
|
let e4m3 = Datatype::FloatingPoint {
|
||||||
|
size: 1,
|
||||||
|
byte_order: DatatypeByteOrder::LittleEndian,
|
||||||
|
bit_offset: 0,
|
||||||
|
bit_precision: 8,
|
||||||
|
exponent_location: 3,
|
||||||
|
exponent_size: 4,
|
||||||
|
mantissa_location: 0,
|
||||||
|
mantissa_size: 3,
|
||||||
|
exponent_bias: 7,
|
||||||
|
};
|
||||||
|
let got = read_as_f64(&[0x38, 0xB8, 0x40, 0x00, 0x7E], &e4m3).unwrap();
|
||||||
|
assert_eq!(&got[..4], &[1.0, -1.0, 2.0, 0.0]);
|
||||||
|
assert!(got[4].is_nan());
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn full_width_signed_unchanged() {
|
fn full_width_signed_unchanged() {
|
||||||
// Regression: full-width 32-bit signed must be unaffected.
|
// Regression: full-width 32-bit signed must be unaffected.
|
||||||
|
|||||||
@@ -4,7 +4,7 @@
|
|||||||
//! for compound, enumeration, variable-length, and array types.
|
//! for compound, enumeration, variable-length, and array types.
|
||||||
|
|
||||||
#[cfg(not(feature = "std"))]
|
#[cfg(not(feature = "std"))]
|
||||||
use alloc::{boxed::Box, string::String, vec, vec::Vec};
|
use alloc::{boxed::Box, format, string::String, vec, vec::Vec};
|
||||||
|
|
||||||
use byteorder::{ByteOrder, LittleEndian};
|
use byteorder::{ByteOrder, LittleEndian};
|
||||||
|
|
||||||
@@ -36,8 +36,18 @@ pub enum CharacterSet {
|
|||||||
/// Reference type.
|
/// Reference type.
|
||||||
#[derive(Debug, Clone, PartialEq)]
|
#[derive(Debug, Clone, PartialEq)]
|
||||||
pub enum ReferenceType {
|
pub enum ReferenceType {
|
||||||
|
/// Legacy object reference: the target's object header address.
|
||||||
Object,
|
Object,
|
||||||
|
/// Legacy dataset region reference.
|
||||||
DatasetRegion,
|
DatasetRegion,
|
||||||
|
/// `H5T_STD_REF` object reference (HDF5 1.12+, datatype message version
|
||||||
|
/// 4): a small header followed by an object token. Decoded by
|
||||||
|
/// `data_read::read_object_references`.
|
||||||
|
Object2,
|
||||||
|
/// `H5T_STD_REF` dataset region reference.
|
||||||
|
DatasetRegion2,
|
||||||
|
/// `H5T_STD_REF` attribute reference.
|
||||||
|
Attribute,
|
||||||
}
|
}
|
||||||
|
|
||||||
/// A member of a compound datatype.
|
/// A member of a compound datatype.
|
||||||
@@ -127,6 +137,17 @@ pub enum Datatype {
|
|||||||
},
|
},
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Longest opaque tag that can be stored: its NUL-padded length must fit
|
||||||
|
/// the 8-bit length in the datatype's class bits.
|
||||||
|
pub const MAX_OPAQUE_TAG_LEN: usize = 248;
|
||||||
|
|
||||||
|
/// An opaque tag up to (not including) its first NUL.
|
||||||
|
fn opaque_tag_text(tag: &[u8]) -> &[u8] {
|
||||||
|
tag.iter()
|
||||||
|
.position(|&b| b == 0)
|
||||||
|
.map_or(tag, |end| &tag[..end])
|
||||||
|
}
|
||||||
|
|
||||||
fn ensure_len(data: &[u8], offset: usize, needed: usize) -> Result<(), FormatError> {
|
fn ensure_len(data: &[u8], offset: usize, needed: usize) -> Result<(), FormatError> {
|
||||||
match offset.checked_add(needed) {
|
match offset.checked_add(needed) {
|
||||||
Some(end) if end <= data.len() => Ok(()),
|
Some(end) if end <= data.len() => Ok(()),
|
||||||
@@ -351,7 +372,10 @@ impl Datatype {
|
|||||||
// Opaque
|
// Opaque
|
||||||
let tag_len = bf0 as usize;
|
let tag_len = bf0 as usize;
|
||||||
ensure_len(data, pos, tag_len)?;
|
ensure_len(data, pos, tag_len)?;
|
||||||
let tag = data[pos..pos + tag_len].to_vec();
|
// The stored tag is NUL-padded to a multiple of 8 bytes; the
|
||||||
|
// tag itself ends at the first NUL (libhdf5 reads it with
|
||||||
|
// `strndup`).
|
||||||
|
let tag = opaque_tag_text(&data[pos..pos + tag_len]).to_vec();
|
||||||
// Tags are padded to multiple of 8 bytes
|
// Tags are padded to multiple of 8 bytes
|
||||||
let padded = (tag_len + 7) & !7;
|
let padded = (tag_len + 7) & !7;
|
||||||
let pos = 8 + padded; // from start of properties
|
let pos = 8 + padded; // from start of properties
|
||||||
@@ -424,9 +448,15 @@ impl Datatype {
|
|||||||
7 => {
|
7 => {
|
||||||
// Reference
|
// Reference
|
||||||
let ref_type_val = bf0 & 0x0F;
|
let ref_type_val = bf0 & 0x0F;
|
||||||
let ref_type = match ref_type_val {
|
// Datatype message version 4 (HDF5 1.12) revised this class:
|
||||||
0 => ReferenceType::Object,
|
// types 2-4 are the new `H5T_STD_REF` references, and the high
|
||||||
1 => ReferenceType::DatasetRegion,
|
// nibble of the first flag byte carries their encoding version.
|
||||||
|
let ref_type = match (ref_type_val, version) {
|
||||||
|
(0, _) => ReferenceType::Object,
|
||||||
|
(1, _) => ReferenceType::DatasetRegion,
|
||||||
|
(2, 4..) => ReferenceType::Object2,
|
||||||
|
(3, 4..) => ReferenceType::DatasetRegion2,
|
||||||
|
(4, 4..) => ReferenceType::Attribute,
|
||||||
_ => return Err(FormatError::InvalidReferenceType(ref_type_val)),
|
_ => return Err(FormatError::InvalidReferenceType(ref_type_val)),
|
||||||
};
|
};
|
||||||
Ok((Datatype::Reference { size, ref_type }, pos))
|
Ok((Datatype::Reference { size, ref_type }, pos))
|
||||||
@@ -624,7 +654,8 @@ impl Datatype {
|
|||||||
mantissa_size,
|
mantissa_size,
|
||||||
exponent_bias,
|
exponent_bias,
|
||||||
} => {
|
} => {
|
||||||
let mut bf0 = 0x20u8; // bit 5: sign location bit (standard IEEE 754)
|
// Bits 4-5: mantissa normalization = 2 (implied leading 1, IEEE 754).
|
||||||
|
let mut bf0 = 0x20u8;
|
||||||
match byte_order {
|
match byte_order {
|
||||||
DatatypeByteOrder::BigEndian => {
|
DatatypeByteOrder::BigEndian => {
|
||||||
bf0 |= 0x01;
|
bf0 |= 0x01;
|
||||||
@@ -634,9 +665,14 @@ impl Datatype {
|
|||||||
}
|
}
|
||||||
_ => {}
|
_ => {}
|
||||||
}
|
}
|
||||||
// bf[1] bits 0-1: mantissa normalization = 2 (MSB not stored, IEEE 754)
|
// Bits 8-15: the sign bit's position, the top bit of the value.
|
||||||
let bf1 = 0x3fu8; // matching what h5py generates
|
// This was hard-coded to 63, which is right only for f64: the
|
||||||
let mut buf = Self::build_header(1, 1, [bf0, bf1, 0], *size);
|
// HDF5 library rejects any other float with "sign bit position
|
||||||
|
// out of bounds", so every f32 dataset and attribute we wrote
|
||||||
|
// was unreadable by h5py and libhdf5.
|
||||||
|
let sign_location =
|
||||||
|
(u32::from(*bit_offset) + u32::from(*bit_precision)).saturating_sub(1) as u8;
|
||||||
|
let mut buf = Self::build_header(1, 1, [bf0, sign_location, 0], *size);
|
||||||
buf.extend_from_slice(&bit_offset.to_le_bytes());
|
buf.extend_from_slice(&bit_offset.to_le_bytes());
|
||||||
buf.extend_from_slice(&bit_precision.to_le_bytes());
|
buf.extend_from_slice(&bit_precision.to_le_bytes());
|
||||||
buf.push(*exponent_location);
|
buf.push(*exponent_location);
|
||||||
@@ -745,7 +781,77 @@ impl Datatype {
|
|||||||
buf.extend_from_slice(&base_type.serialize());
|
buf.extend_from_slice(&base_type.serialize());
|
||||||
buf
|
buf
|
||||||
}
|
}
|
||||||
_ => Vec::new(),
|
Datatype::Time {
|
||||||
|
size,
|
||||||
|
bit_precision,
|
||||||
|
} => {
|
||||||
|
// Byte order is not modelled for time types; write little-endian.
|
||||||
|
let mut buf = Self::build_header(2, 1, [0, 0, 0], *size);
|
||||||
|
buf.extend_from_slice(&bit_precision.to_le_bytes());
|
||||||
|
buf
|
||||||
|
}
|
||||||
|
Datatype::BitField {
|
||||||
|
size,
|
||||||
|
byte_order,
|
||||||
|
bit_offset,
|
||||||
|
bit_precision,
|
||||||
|
} => {
|
||||||
|
let bf0 = u8::from(matches!(byte_order, DatatypeByteOrder::BigEndian));
|
||||||
|
let mut buf = Self::build_header(4, 1, [bf0, 0, 0], *size);
|
||||||
|
buf.extend_from_slice(&bit_offset.to_le_bytes());
|
||||||
|
buf.extend_from_slice(&bit_precision.to_le_bytes());
|
||||||
|
buf
|
||||||
|
}
|
||||||
|
Datatype::Opaque { size, tag } => {
|
||||||
|
// The tag is stored NUL-padded to a multiple of 8 bytes and the
|
||||||
|
// padded length goes in the class bits, as libhdf5 writes it.
|
||||||
|
// A tag longer than MAX_OPAQUE_TAG_LEN cannot be encoded;
|
||||||
|
// `check_encodable` rejects it before a file is written.
|
||||||
|
let tag = opaque_tag_text(tag);
|
||||||
|
let tag = &tag[..tag.len().min(MAX_OPAQUE_TAG_LEN)];
|
||||||
|
let padded = tag.len().div_ceil(8) * 8;
|
||||||
|
let mut buf = Self::build_header(5, 1, [padded as u8, 0, 0], *size);
|
||||||
|
buf.extend_from_slice(tag);
|
||||||
|
buf.resize(8 + padded, 0);
|
||||||
|
buf
|
||||||
|
}
|
||||||
|
Datatype::Reference { size, ref_type } => {
|
||||||
|
// Legacy references are datatype version 1; the H5T_STD_REF
|
||||||
|
// kinds only exist from version 4, which also carries their
|
||||||
|
// encoding version (1) in the high nibble.
|
||||||
|
let (version, bf0) = match ref_type {
|
||||||
|
ReferenceType::Object => (1, 0),
|
||||||
|
ReferenceType::DatasetRegion => (1, 1),
|
||||||
|
ReferenceType::Object2 => (4, 0x12),
|
||||||
|
ReferenceType::DatasetRegion2 => (4, 0x13),
|
||||||
|
ReferenceType::Attribute => (4, 0x14),
|
||||||
|
};
|
||||||
|
Self::build_header(7, version, [bf0, 0, 0], *size)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Check that this datatype can be written: every part of it has an
|
||||||
|
/// on-disk encoding. [`Self::serialize`] cannot report errors, so the
|
||||||
|
/// writer calls this first.
|
||||||
|
pub fn check_encodable(&self) -> Result<(), FormatError> {
|
||||||
|
match self {
|
||||||
|
Datatype::Opaque { tag, .. } if opaque_tag_text(tag).len() > MAX_OPAQUE_TAG_LEN => {
|
||||||
|
Err(FormatError::SerializationError(format!(
|
||||||
|
"opaque tag is {} bytes; at most {MAX_OPAQUE_TAG_LEN} can be stored",
|
||||||
|
opaque_tag_text(tag).len()
|
||||||
|
)))
|
||||||
|
}
|
||||||
|
Datatype::String { size: 0, .. } => Err(FormatError::SerializationError(
|
||||||
|
"fixed-length string datatype of size 0 (libhdf5 requires at least 1 byte)".into(),
|
||||||
|
)),
|
||||||
|
Datatype::Compound { members, .. } => members
|
||||||
|
.iter()
|
||||||
|
.try_for_each(|m| m.datatype.check_encodable()),
|
||||||
|
Datatype::Enumeration { base_type, .. }
|
||||||
|
| Datatype::VariableLength { base_type, .. }
|
||||||
|
| Datatype::Array { base_type, .. } => base_type.check_encodable(),
|
||||||
|
_ => Ok(()),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -802,6 +908,24 @@ fn build_dt_header(class: u8, version: u8, bf: [u8; 3], size: u32) -> Vec<u8> {
|
|||||||
mod tests {
|
mod tests {
|
||||||
use super::*;
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn float_sign_location_is_the_top_bit_of_the_value() {
|
||||||
|
// The HDF5 library rejects a float whose sign position is not inside
|
||||||
|
// its precision; this was hard-coded to 63, so every f32 we wrote was
|
||||||
|
// unreadable by h5py. Byte 2 of the message is the sign position.
|
||||||
|
use crate::type_builders::{make_f16_type, make_f32_type, make_f64_type};
|
||||||
|
for (dt, sign) in [
|
||||||
|
(make_f16_type(), 15),
|
||||||
|
(make_f32_type(), 31),
|
||||||
|
(make_f64_type(), 63),
|
||||||
|
] {
|
||||||
|
let bytes = dt.serialize();
|
||||||
|
assert_eq!(bytes[2], sign, "{dt:?}");
|
||||||
|
let (parsed, _) = Datatype::parse(&bytes).unwrap();
|
||||||
|
assert_eq!(parsed, dt);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
// Helper to build a fixed-point datatype message
|
// Helper to build a fixed-point datatype message
|
||||||
fn build_fixed_point(
|
fn build_fixed_point(
|
||||||
size: u32,
|
size: u32,
|
||||||
@@ -1563,6 +1687,144 @@ mod tests {
|
|||||||
assert_eq!(err, FormatError::InvalidCharacterSet(2));
|
assert_eq!(err, FormatError::InvalidCharacterSet(2));
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_reference_v4_std_ref_from_hdf5_2_0() {
|
||||||
|
// Datatype message of an H5T_STD_REF dataset written by HDF5 2.0:
|
||||||
|
// class 7, version 4, type 2 (object), encoding version 1, 18 bytes.
|
||||||
|
let bytes = [0x47, 0x12, 0x00, 0x00, 0x12, 0x00, 0x00, 0x00];
|
||||||
|
let (dt, consumed) = Datatype::parse(&bytes).unwrap();
|
||||||
|
assert_eq!(consumed, 8);
|
||||||
|
assert_eq!(
|
||||||
|
dt,
|
||||||
|
Datatype::Reference {
|
||||||
|
size: 18,
|
||||||
|
ref_type: ReferenceType::Object2
|
||||||
|
}
|
||||||
|
);
|
||||||
|
// The new types are only valid from datatype version 4.
|
||||||
|
let old_version = [0x37, 0x12, 0x00, 0x00, 0x12, 0x00, 0x00, 0x00];
|
||||||
|
assert_eq!(
|
||||||
|
Datatype::parse(&old_version).unwrap_err(),
|
||||||
|
FormatError::InvalidReferenceType(2)
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
fn hex(s: &str) -> Vec<u8> {
|
||||||
|
(0..s.len())
|
||||||
|
.step_by(2)
|
||||||
|
.map(|i| u8::from_str_radix(&s[i..i + 2], 16).unwrap())
|
||||||
|
.collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// `serialize` used to return an empty message for these four classes,
|
||||||
|
/// which libhdf5 rejects ("ran off end of input buffer while decoding").
|
||||||
|
/// Expected bytes are libhdf5's own encoding (HDF5 2.0 `H5Tencode`, or the
|
||||||
|
/// datatype message of an HDF5 2.0 file for `H5T_STD_REF`).
|
||||||
|
#[test]
|
||||||
|
fn serialize_matches_libhdf5_for_time_bitfield_opaque_reference() {
|
||||||
|
let cases = [
|
||||||
|
(
|
||||||
|
Datatype::Reference {
|
||||||
|
size: 8,
|
||||||
|
ref_type: ReferenceType::Object,
|
||||||
|
},
|
||||||
|
"1700000008000000",
|
||||||
|
),
|
||||||
|
(
|
||||||
|
Datatype::Reference {
|
||||||
|
size: 12,
|
||||||
|
ref_type: ReferenceType::DatasetRegion,
|
||||||
|
},
|
||||||
|
"170100000c000000",
|
||||||
|
),
|
||||||
|
(
|
||||||
|
Datatype::Reference {
|
||||||
|
size: 18,
|
||||||
|
ref_type: ReferenceType::Object2,
|
||||||
|
},
|
||||||
|
"4712000012000000",
|
||||||
|
),
|
||||||
|
(
|
||||||
|
Datatype::BitField {
|
||||||
|
size: 1,
|
||||||
|
byte_order: DatatypeByteOrder::LittleEndian,
|
||||||
|
bit_offset: 0,
|
||||||
|
bit_precision: 8,
|
||||||
|
},
|
||||||
|
"140000000100000000000800",
|
||||||
|
),
|
||||||
|
(
|
||||||
|
Datatype::BitField {
|
||||||
|
size: 2,
|
||||||
|
byte_order: DatatypeByteOrder::BigEndian,
|
||||||
|
bit_offset: 0,
|
||||||
|
bit_precision: 16,
|
||||||
|
},
|
||||||
|
"140100000200000000001000",
|
||||||
|
),
|
||||||
|
(
|
||||||
|
Datatype::Opaque {
|
||||||
|
size: 4,
|
||||||
|
tag: b"mytag".to_vec(),
|
||||||
|
},
|
||||||
|
"15080000040000006d79746167000000",
|
||||||
|
),
|
||||||
|
(
|
||||||
|
Datatype::Opaque {
|
||||||
|
size: 4,
|
||||||
|
tag: b"12345678".to_vec(),
|
||||||
|
},
|
||||||
|
"15080000040000003132333435363738",
|
||||||
|
),
|
||||||
|
(
|
||||||
|
Datatype::Opaque {
|
||||||
|
size: 4,
|
||||||
|
tag: vec![],
|
||||||
|
},
|
||||||
|
"1500000004000000",
|
||||||
|
),
|
||||||
|
(
|
||||||
|
Datatype::Time {
|
||||||
|
size: 4,
|
||||||
|
bit_precision: 32,
|
||||||
|
},
|
||||||
|
"12000000040000002000",
|
||||||
|
),
|
||||||
|
];
|
||||||
|
for (dt, expected) in cases {
|
||||||
|
let bytes = dt.serialize();
|
||||||
|
assert_eq!(bytes, hex(expected), "{dt:?}");
|
||||||
|
let (parsed, consumed) = Datatype::parse(&bytes).unwrap();
|
||||||
|
assert_eq!(parsed, dt);
|
||||||
|
assert_eq!(consumed, bytes.len());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn opaque_tag_padding_is_not_part_of_the_tag() {
|
||||||
|
// libhdf5 pads "mytag" to 8 bytes; parsing must not return the NULs,
|
||||||
|
// or copying the type would grow the tag.
|
||||||
|
let (dt, _) = Datatype::parse(&hex("15080000040000006d79746167000000")).unwrap();
|
||||||
|
assert_eq!(
|
||||||
|
dt,
|
||||||
|
Datatype::Opaque {
|
||||||
|
size: 4,
|
||||||
|
tag: b"mytag".to_vec()
|
||||||
|
}
|
||||||
|
);
|
||||||
|
let long = Datatype::Opaque {
|
||||||
|
size: 1,
|
||||||
|
tag: vec![b'x'; MAX_OPAQUE_TAG_LEN + 1],
|
||||||
|
};
|
||||||
|
assert!(long.check_encodable().is_err());
|
||||||
|
let ok = Datatype::Opaque {
|
||||||
|
size: 1,
|
||||||
|
tag: vec![b'x'; MAX_OPAQUE_TAG_LEN],
|
||||||
|
};
|
||||||
|
assert!(ok.check_encodable().is_ok());
|
||||||
|
assert_eq!(Datatype::parse(&ok.serialize()).unwrap().0, ok);
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn test_error_invalid_reference_type() {
|
fn test_error_invalid_reference_type() {
|
||||||
let buf = build_dt_header(7, 1, [5, 0, 0], 8);
|
let buf = build_dt_header(7, 1, [5, 0, 0], 8);
|
||||||
|
|||||||
@@ -7,7 +7,7 @@ extern crate alloc;
|
|||||||
use alloc::{vec, vec::Vec};
|
use alloc::{vec, vec::Vec};
|
||||||
|
|
||||||
use crate::checksum::jenkins_lookup3;
|
use crate::checksum::jenkins_lookup3;
|
||||||
use crate::chunked_write::WrittenChunk;
|
use crate::chunked_write::{WrittenChunk, filtered_chunk_size_len, push_addr, push_index_element};
|
||||||
|
|
||||||
/// Serialize a v4 Extensible Array layout message.
|
/// Serialize a v4 Extensible Array layout message.
|
||||||
pub(crate) fn serialize_v4_extensible_array(
|
pub(crate) fn serialize_v4_extensible_array(
|
||||||
@@ -58,11 +58,11 @@ pub(crate) fn serialize_v4_extensible_array(
|
|||||||
buf.push(4);
|
buf.push(4);
|
||||||
|
|
||||||
// EA creation parameters (must match AEHD and HDF5 C library defaults)
|
// EA creation parameters (must match AEHD and HDF5 C library defaults)
|
||||||
buf.push(32); // max_nelmts_bits
|
buf.push(MAX_NELMTS_BITS);
|
||||||
buf.push(4); // idx_blk_elmts
|
buf.push(IDX_BLK_ELMTS);
|
||||||
buf.push(4); // super_blk_min_data_ptrs
|
buf.push(SUP_BLK_MIN_DATA_PTRS);
|
||||||
buf.push(16); // data_blk_min_elmts
|
buf.push(DATA_BLK_MIN_ELMTS);
|
||||||
buf.push(10); // max_dblk_page_nelmts_bits
|
buf.push(MAX_DBLK_PAGE_NELMTS_BITS);
|
||||||
|
|
||||||
// EA header address
|
// EA header address
|
||||||
match offset_size {
|
match offset_size {
|
||||||
@@ -74,304 +74,281 @@ pub(crate) fn serialize_v4_extensible_array(
|
|||||||
buf
|
buf
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// EA creation parameters — the HDF5 library's defaults for chunk indexes
|
||||||
|
// (`H5D_EARRAY_*`); the layout message above and the header must agree.
|
||||||
|
const MAX_NELMTS_BITS: u8 = 32;
|
||||||
|
const IDX_BLK_ELMTS: u8 = 4;
|
||||||
|
const SUP_BLK_MIN_DATA_PTRS: u8 = 4;
|
||||||
|
const DATA_BLK_MIN_ELMTS: u8 = 16;
|
||||||
|
const MAX_DBLK_PAGE_NELMTS_BITS: u8 = 10;
|
||||||
|
|
||||||
|
/// One data block of the array: its first element (relative to the end of
|
||||||
|
/// the index block's own elements), element count, and address when it is
|
||||||
|
/// allocated.
|
||||||
|
struct DataBlock {
|
||||||
|
start: usize,
|
||||||
|
nelmts: usize,
|
||||||
|
addr: Option<u64>,
|
||||||
|
}
|
||||||
|
|
||||||
/// Build a complete Extensible Array at a known absolute address.
|
/// Build a complete Extensible Array at a known absolute address.
|
||||||
///
|
///
|
||||||
/// For simplicity, we put all elements inline in the index block when the
|
/// `slots[i]` is the element at linear index `i` (see `chunk_grid`); `None`
|
||||||
/// number of chunks is small (up to idx_blk_elmts), otherwise use inline +
|
/// marks an unallocated chunk. The first `IDX_BLK_ELMTS` elements live in
|
||||||
/// direct data blocks.
|
/// the index block, the rest in data blocks grouped by super block level
|
||||||
|
/// exactly as `H5EA__hdr_init` sizes them: level `u` has `2^(u/2)` data
|
||||||
|
/// blocks of `DATA_BLK_MIN_ELMTS * 2^ceil(u/2)` elements. The data blocks of
|
||||||
|
/// the first levels are addressed straight from the index block; later
|
||||||
|
/// levels go through a super block (EASB). Data blocks larger than a page
|
||||||
|
/// (`2^MAX_DBLK_PAGE_NELMTS_BITS` elements) are paged, with their page-init
|
||||||
|
/// bits kept in the owning super block. Only blocks holding a defined element
|
||||||
|
/// are allocated; the rest keep the undefined address, as in a file the
|
||||||
|
/// library wrote.
|
||||||
pub fn build_extensible_array_at(
|
pub fn build_extensible_array_at(
|
||||||
chunks: &[WrittenChunk],
|
slots: &[Option<WrittenChunk>],
|
||||||
offset_size: u8,
|
offset_size: u8,
|
||||||
length_size: u8,
|
length_size: u8,
|
||||||
has_filters: bool,
|
has_filters: bool,
|
||||||
ea_base_address: u64,
|
ea_base_address: u64,
|
||||||
) -> Vec<u8> {
|
) -> Vec<u8> {
|
||||||
let os = offset_size as usize;
|
let os = offset_size as usize;
|
||||||
let num_elements = chunks.len();
|
let chunk_size_bytes = has_filters.then(|| filtered_chunk_size_len(slots));
|
||||||
|
let elem_size = os + chunk_size_bytes.map_or(0, |n| n + 4);
|
||||||
// Compute element encoding size (same logic as Fixed Array)
|
|
||||||
let chunk_size_bytes: usize = if has_filters {
|
|
||||||
let max_raw = chunks.iter().map(|c| c.raw_size).max().unwrap_or(1);
|
|
||||||
let log2_val = if max_raw <= 1 {
|
|
||||||
0
|
|
||||||
} else {
|
|
||||||
63 - max_raw.leading_zeros()
|
|
||||||
};
|
|
||||||
let len = 1 + ((log2_val + 8) / 8) as usize;
|
|
||||||
len.min(8)
|
|
||||||
} else {
|
|
||||||
0
|
|
||||||
};
|
|
||||||
|
|
||||||
let elem_size = if has_filters {
|
|
||||||
os + chunk_size_bytes + 4
|
|
||||||
} else {
|
|
||||||
os
|
|
||||||
};
|
|
||||||
|
|
||||||
let client_id: u8 = if has_filters { 1 } else { 0 };
|
let client_id: u8 = if has_filters { 1 } else { 0 };
|
||||||
|
let arr_off_size = (MAX_NELMTS_BITS as usize).div_ceil(8);
|
||||||
|
let page_nelmts = 1usize << MAX_DBLK_PAGE_NELMTS_BITS;
|
||||||
|
let idx_blk = IDX_BLK_ELMTS as usize;
|
||||||
|
|
||||||
// EA creation parameters — must match HDF5 C library defaults exactly
|
// Elements past the last defined one are never realised
|
||||||
let max_nelmts_bits: u8 = 32;
|
// (`max_idx_set` is one past the highest index ever set).
|
||||||
let idx_blk_elmts: u8 = 4;
|
let max_idx_set = slots.iter().rposition(Option::is_some).map_or(0, |i| i + 1);
|
||||||
let min_dblk_nelmts: u8 = 16;
|
let slots = &slots[..max_idx_set];
|
||||||
let super_blk_min_nelmts: u8 = 4;
|
let defined_in = |start: usize, n: usize| -> bool {
|
||||||
let max_dblk_nelmts_bits: u8 = 10;
|
let lo = idx_blk.saturating_add(start).min(slots.len());
|
||||||
|
let hi = idx_blk
|
||||||
|
.saturating_add(start)
|
||||||
|
.saturating_add(n)
|
||||||
|
.min(slots.len());
|
||||||
|
slots[lo..hi].iter().any(Option::is_some)
|
||||||
|
};
|
||||||
|
|
||||||
// EAHD size: fixed(12) + 6 stats(6*length_size) + addr(offset_size) + checksum(4)
|
// Super block levels: (ndblks, dblk_nelmts, first element).
|
||||||
|
let log2_dmin = (DATA_BLK_MIN_ELMTS as u32).trailing_zeros() as usize;
|
||||||
|
let nsblks = 1 + MAX_NELMTS_BITS as usize - log2_dmin;
|
||||||
|
let ndblk_addrs = 2 * (SUP_BLK_MIN_DATA_PTRS as usize - 1);
|
||||||
|
let mut levels: Vec<(usize, usize, usize)> = Vec::with_capacity(nsblks);
|
||||||
|
let mut start = 0usize;
|
||||||
|
for u in 0..nsblks {
|
||||||
|
let ndblks = 1usize << (u / 2);
|
||||||
|
let nelmts = (DATA_BLK_MIN_ELMTS as usize) << u.div_ceil(2);
|
||||||
|
levels.push((ndblks, nelmts, start));
|
||||||
|
// Saturate: on 32-bit targets the last levels only need to compare
|
||||||
|
// as "beyond the end".
|
||||||
|
start = start.saturating_add(ndblks.saturating_mul(nelmts));
|
||||||
|
}
|
||||||
|
// Levels whose data blocks the index block addresses directly.
|
||||||
|
let mut direct_levels = 0;
|
||||||
|
let mut n = 0;
|
||||||
|
while n < ndblk_addrs {
|
||||||
|
n += levels[direct_levels].0;
|
||||||
|
direct_levels += 1;
|
||||||
|
}
|
||||||
|
let nsblk_addrs = nsblks - direct_levels;
|
||||||
|
|
||||||
|
let dblk_size = |nelmts: usize| -> usize {
|
||||||
|
let prefix = 4 + 1 + 1 + os + arr_off_size + 4;
|
||||||
|
if nelmts > page_nelmts {
|
||||||
|
prefix + (nelmts / page_nelmts) * (page_nelmts * elem_size + 4)
|
||||||
|
} else {
|
||||||
|
prefix + nelmts * elem_size
|
||||||
|
}
|
||||||
|
};
|
||||||
|
let sblk_bitmap_len = |ndblks: usize, nelmts: usize| -> usize {
|
||||||
|
if nelmts > page_nelmts {
|
||||||
|
ndblks * (nelmts / page_nelmts).div_ceil(8)
|
||||||
|
} else {
|
||||||
|
0
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
// Plan addresses: header, index block, the direct data blocks, then each
|
||||||
|
// allocated super block followed by its allocated data blocks.
|
||||||
let aehd_size = 4 + 1 + 1 + 1 + 1 + 1 + 1 + 1 + 1 + 6 * length_size as usize + os + 4;
|
let aehd_size = 4 + 1 + 1 + 1 + 1 + 1 + 1 + 1 + 1 + 6 * length_size as usize + os + 4;
|
||||||
let aeib_address = ea_base_address + aehd_size as u64;
|
let aeib_address = ea_base_address + aehd_size as u64;
|
||||||
|
let aeib_size = 4 + 1 + 1 + os + idx_blk * elem_size + ndblk_addrs * os + nsblk_addrs * os + 4;
|
||||||
|
let mut cursor = aeib_address + aeib_size as u64;
|
||||||
|
|
||||||
// Determine how many elements go inline vs data blocks
|
let mut ndata_blks = 0u64;
|
||||||
let n_inline = (idx_blk_elmts as usize).min(num_elements);
|
let mut data_blk_size = 0u64;
|
||||||
let remaining_after_inline = num_elements.saturating_sub(n_inline);
|
let mut nsuper_blks = 0u64;
|
||||||
|
let mut super_blk_size = 0u64;
|
||||||
|
let mut realized = idx_blk as u64;
|
||||||
|
|
||||||
// Compute super block layout per HDF5 spec
|
let mut plan_dblk = |cursor: &mut u64, start: usize, nelmts: usize| -> DataBlock {
|
||||||
let sblk_min = super_blk_min_nelmts as usize;
|
let addr = defined_in(start, nelmts).then(|| {
|
||||||
let log2_dblk_min = if min_dblk_nelmts <= 1 {
|
let a = *cursor;
|
||||||
0
|
let size = dblk_size(nelmts) as u64;
|
||||||
} else {
|
*cursor += size;
|
||||||
(min_dblk_nelmts as u32).trailing_zeros() as usize
|
ndata_blks += 1;
|
||||||
|
data_blk_size += size;
|
||||||
|
realized += nelmts as u64;
|
||||||
|
a
|
||||||
|
});
|
||||||
|
DataBlock {
|
||||||
|
start,
|
||||||
|
nelmts,
|
||||||
|
addr,
|
||||||
|
}
|
||||||
};
|
};
|
||||||
let nsblks = (max_nelmts_bits as usize).saturating_sub(log2_dblk_min) + 1;
|
|
||||||
|
|
||||||
// Direct data block addresses (from super blocks 0..sblk_min-1)
|
let mut direct: Vec<DataBlock> = Vec::with_capacity(ndblk_addrs);
|
||||||
let mut dblk_sizes: Vec<usize> = Vec::new();
|
for &(ndblks, nelmts, first) in &levels[..direct_levels] {
|
||||||
for sblk_idx in 0..sblk_min.min(nsblks) {
|
for k in 0..ndblks {
|
||||||
let ndblks = 1usize << (sblk_idx / 2);
|
direct.push(plan_dblk(&mut cursor, first + k * nelmts, nelmts));
|
||||||
let dblk_nelmts = (min_dblk_nelmts as usize) * (1 << sblk_idx.div_ceil(2));
|
|
||||||
for _ in 0..ndblks {
|
|
||||||
dblk_sizes.push(dblk_nelmts);
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
let n_direct_dblks = dblk_sizes.len();
|
// (super block address, level, its data blocks)
|
||||||
|
let mut supers: Vec<(Option<u64>, usize, Vec<DataBlock>)> = Vec::with_capacity(nsblk_addrs);
|
||||||
// Super block addresses (for super blocks sblk_min..nsblks-1)
|
for (u, &(ndblks, nelmts, first)) in levels.iter().enumerate().skip(direct_levels) {
|
||||||
let n_sblk_addrs = nsblks.saturating_sub(sblk_min);
|
if !defined_in(first, ndblks.saturating_mul(nelmts)) {
|
||||||
|
supers.push((None, u, Vec::new()));
|
||||||
// EAIB size
|
continue;
|
||||||
let aeib_size = 4
|
|
||||||
+ 1
|
|
||||||
+ 1
|
|
||||||
+ os
|
|
||||||
+ idx_blk_elmts as usize * elem_size
|
|
||||||
+ n_direct_dblks * os
|
|
||||||
+ n_sblk_addrs * os
|
|
||||||
+ 4;
|
|
||||||
|
|
||||||
// Build AEHD
|
|
||||||
let mut aehd = Vec::with_capacity(aehd_size);
|
|
||||||
aehd.extend_from_slice(b"EAHD");
|
|
||||||
aehd.push(0); // version
|
|
||||||
aehd.push(client_id);
|
|
||||||
aehd.push(elem_size as u8);
|
|
||||||
aehd.push(max_nelmts_bits);
|
|
||||||
aehd.push(idx_blk_elmts);
|
|
||||||
aehd.push(min_dblk_nelmts);
|
|
||||||
aehd.push(super_blk_min_nelmts);
|
|
||||||
aehd.push(max_dblk_nelmts_bits);
|
|
||||||
|
|
||||||
// Count data blocks that will have chunks
|
|
||||||
let n_active_dblks: u64 = if remaining_after_inline > 0 {
|
|
||||||
let mut count = 0u64;
|
|
||||||
let mut ci = n_inline;
|
|
||||||
for &sz in &dblk_sizes {
|
|
||||||
if ci < num_elements {
|
|
||||||
count += 1;
|
|
||||||
ci += sz;
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
count
|
let sb_size =
|
||||||
} else {
|
4 + 1 + 1 + os + arr_off_size + sblk_bitmap_len(ndblks, nelmts) + ndblks * os + 4;
|
||||||
0
|
let sb_addr = cursor;
|
||||||
};
|
cursor += sb_size as u64;
|
||||||
let blk_off_size = (max_nelmts_bits as usize).div_ceil(8);
|
nsuper_blks += 1;
|
||||||
let aedb_header_overhead = 4 + 1 + 1 + os + blk_off_size + 4;
|
super_blk_size += sb_size as u64;
|
||||||
let data_blk_total_size: u64 = if remaining_after_inline > 0 {
|
let dblks = (0..ndblks)
|
||||||
let mut total = 0u64;
|
.map(|k| plan_dblk(&mut cursor, first + k * nelmts, nelmts))
|
||||||
let mut ci = n_inline;
|
.collect();
|
||||||
for &sz in &dblk_sizes {
|
supers.push((Some(sb_addr), u, dblks));
|
||||||
if ci < num_elements {
|
}
|
||||||
total += (aedb_header_overhead + sz * elem_size) as u64;
|
|
||||||
ci += sz;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
total
|
|
||||||
} else {
|
|
||||||
0
|
|
||||||
};
|
|
||||||
let max_idx_set: u64 = if remaining_after_inline > 0 {
|
|
||||||
let mut max_set = idx_blk_elmts as u64;
|
|
||||||
let mut ci = n_inline;
|
|
||||||
for &sz in &dblk_sizes {
|
|
||||||
if ci < num_elements {
|
|
||||||
max_set += sz as u64;
|
|
||||||
ci += sz;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
max_set
|
|
||||||
} else {
|
|
||||||
idx_blk_elmts as u64
|
|
||||||
};
|
|
||||||
|
|
||||||
|
let slot = |i: usize| slots.get(i).and_then(Option::as_ref);
|
||||||
let write_length = |buf: &mut Vec<u8>, val: u64| match length_size {
|
let write_length = |buf: &mut Vec<u8>, val: u64| match length_size {
|
||||||
4 => buf.extend_from_slice(&(val as u32).to_le_bytes()),
|
4 => buf.extend_from_slice(&(val as u32).to_le_bytes()),
|
||||||
_ => buf.extend_from_slice(&val.to_le_bytes()),
|
_ => buf.extend_from_slice(&val.to_le_bytes()),
|
||||||
};
|
};
|
||||||
let write_addr = |buf: &mut Vec<u8>, val: u64| match offset_size {
|
let write_addr_opt = |buf: &mut Vec<u8>, addr: Option<u64>| match addr {
|
||||||
4 => buf.extend_from_slice(&(val as u32).to_le_bytes()),
|
Some(a) => push_addr(buf, a, offset_size),
|
||||||
_ => buf.extend_from_slice(&val.to_le_bytes()),
|
None => buf.extend(core::iter::repeat_n(0xFF, os)),
|
||||||
|
};
|
||||||
|
let block_prefix = |buf: &mut Vec<u8>, sig: &[u8; 4], block_off: usize| {
|
||||||
|
buf.extend_from_slice(sig);
|
||||||
|
buf.push(0); // version
|
||||||
|
buf.push(client_id);
|
||||||
|
push_addr(buf, ea_base_address, offset_size);
|
||||||
|
buf.extend_from_slice(&(block_off as u64).to_le_bytes()[..arr_off_size]);
|
||||||
|
};
|
||||||
|
// Serialise one data block (paged or not) onto `out`.
|
||||||
|
let write_dblk = |out: &mut Vec<u8>, db: &DataBlock| {
|
||||||
|
let at = out.len();
|
||||||
|
block_prefix(out, b"EADB", db.start);
|
||||||
|
let first = idx_blk + db.start;
|
||||||
|
if db.nelmts > page_nelmts {
|
||||||
|
// Paged: the prefix carries only its own checksum; each page
|
||||||
|
// follows with one of its own.
|
||||||
|
let sum = jenkins_lookup3(&out[at..]);
|
||||||
|
out.extend_from_slice(&sum.to_le_bytes());
|
||||||
|
for p in 0..db.nelmts / page_nelmts {
|
||||||
|
let page_at = out.len();
|
||||||
|
for e in 0..page_nelmts {
|
||||||
|
let i = first + p * page_nelmts + e;
|
||||||
|
push_index_element(out, slot(i), offset_size, chunk_size_bytes);
|
||||||
|
}
|
||||||
|
let sum = jenkins_lookup3(&out[page_at..]);
|
||||||
|
out.extend_from_slice(&sum.to_le_bytes());
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
for i in first..first + db.nelmts {
|
||||||
|
push_index_element(out, slot(i), offset_size, chunk_size_bytes);
|
||||||
|
}
|
||||||
|
let sum = jenkins_lookup3(&out[at..]);
|
||||||
|
out.extend_from_slice(&sum.to_le_bytes());
|
||||||
|
}
|
||||||
|
debug_assert_eq!(out.len() - at, dblk_size(db.nelmts));
|
||||||
};
|
};
|
||||||
|
|
||||||
write_length(&mut aehd, 0);
|
// Header (EAHD). The six statistics are, in order: super blocks, their
|
||||||
write_length(&mut aehd, 0);
|
// bytes, data blocks, their bytes, max index set, elements realised.
|
||||||
write_length(&mut aehd, n_active_dblks);
|
let mut out = Vec::with_capacity((cursor - ea_base_address) as usize);
|
||||||
write_length(&mut aehd, data_blk_total_size);
|
out.extend_from_slice(b"EAHD");
|
||||||
write_length(&mut aehd, num_elements as u64);
|
out.push(0); // version
|
||||||
write_length(&mut aehd, max_idx_set);
|
out.push(client_id);
|
||||||
|
out.push(elem_size as u8);
|
||||||
|
out.push(MAX_NELMTS_BITS);
|
||||||
|
out.push(IDX_BLK_ELMTS);
|
||||||
|
out.push(DATA_BLK_MIN_ELMTS);
|
||||||
|
out.push(SUP_BLK_MIN_DATA_PTRS);
|
||||||
|
out.push(MAX_DBLK_PAGE_NELMTS_BITS);
|
||||||
|
write_length(&mut out, nsuper_blks);
|
||||||
|
write_length(&mut out, super_blk_size);
|
||||||
|
write_length(&mut out, ndata_blks);
|
||||||
|
write_length(&mut out, data_blk_size);
|
||||||
|
write_length(&mut out, max_idx_set as u64);
|
||||||
|
write_length(&mut out, realized);
|
||||||
|
push_addr(&mut out, aeib_address, offset_size);
|
||||||
|
let sum = jenkins_lookup3(&out);
|
||||||
|
out.extend_from_slice(&sum.to_le_bytes());
|
||||||
|
debug_assert_eq!(out.len(), aehd_size);
|
||||||
|
|
||||||
write_addr(&mut aehd, aeib_address);
|
// Index block (EAIB): inline elements, data block and super block
|
||||||
|
// addresses.
|
||||||
let aehd_checksum = jenkins_lookup3(&aehd);
|
let ib_start = out.len();
|
||||||
aehd.extend_from_slice(&aehd_checksum.to_le_bytes());
|
out.extend_from_slice(b"EAIB");
|
||||||
debug_assert_eq!(aehd.len(), aehd_size);
|
out.push(0);
|
||||||
|
out.push(client_id);
|
||||||
// Build AEIB
|
push_addr(&mut out, ea_base_address, offset_size);
|
||||||
let mut aeib = Vec::with_capacity(aeib_size);
|
for i in 0..idx_blk {
|
||||||
aeib.extend_from_slice(b"EAIB");
|
push_index_element(&mut out, slot(i), offset_size, chunk_size_bytes);
|
||||||
aeib.push(0);
|
|
||||||
aeib.push(client_id);
|
|
||||||
|
|
||||||
match offset_size {
|
|
||||||
4 => aeib.extend_from_slice(&(ea_base_address as u32).to_le_bytes()),
|
|
||||||
8 => aeib.extend_from_slice(&ea_base_address.to_le_bytes()),
|
|
||||||
_ => aeib.extend_from_slice(&ea_base_address.to_le_bytes()),
|
|
||||||
}
|
}
|
||||||
|
for db in &direct {
|
||||||
// Inline elements
|
write_addr_opt(&mut out, db.addr);
|
||||||
#[allow(clippy::needless_range_loop)]
|
|
||||||
for i in 0..idx_blk_elmts as usize {
|
|
||||||
if i < n_inline {
|
|
||||||
write_chunk_element(
|
|
||||||
&mut aeib,
|
|
||||||
&chunks[i],
|
|
||||||
offset_size,
|
|
||||||
has_filters,
|
|
||||||
chunk_size_bytes,
|
|
||||||
);
|
|
||||||
} else {
|
|
||||||
write_undefined_element(&mut aeib, offset_size, has_filters, chunk_size_bytes);
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
for (sb_addr, _, _) in &supers {
|
||||||
|
write_addr_opt(&mut out, *sb_addr);
|
||||||
|
}
|
||||||
|
let sum = jenkins_lookup3(&out[ib_start..]);
|
||||||
|
out.extend_from_slice(&sum.to_le_bytes());
|
||||||
|
debug_assert_eq!(out.len() - ib_start, aeib_size);
|
||||||
|
|
||||||
// Data block addresses + build data blocks
|
for db in direct.iter().filter(|d| d.addr.is_some()) {
|
||||||
let mut data_blocks_buf = Vec::new();
|
write_dblk(&mut out, db);
|
||||||
let dblks_base = aeib_address + aeib_size as u64;
|
}
|
||||||
let mut dblk_cursor = dblks_base;
|
for (sb_addr, u, dblks) in &supers {
|
||||||
let mut chunk_idx = n_inline;
|
if sb_addr.is_none() {
|
||||||
|
|
||||||
for &nelmts in &dblk_sizes {
|
|
||||||
if chunk_idx >= num_elements {
|
|
||||||
match offset_size {
|
|
||||||
4 => aeib.extend_from_slice(&u32::MAX.to_le_bytes()),
|
|
||||||
8 => aeib.extend_from_slice(&u64::MAX.to_le_bytes()),
|
|
||||||
_ => aeib.extend_from_slice(&u64::MAX.to_le_bytes()),
|
|
||||||
}
|
|
||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
|
let (ndblks, nelmts, first) = levels[*u];
|
||||||
match offset_size {
|
let sb_start = out.len();
|
||||||
4 => aeib.extend_from_slice(&(dblk_cursor as u32).to_le_bytes()),
|
block_prefix(&mut out, b"EASB", first);
|
||||||
8 => aeib.extend_from_slice(&dblk_cursor.to_le_bytes()),
|
if nelmts > page_nelmts {
|
||||||
_ => aeib.extend_from_slice(&dblk_cursor.to_le_bytes()),
|
// Page-init bits, `npages` per data block, packed MSB-first
|
||||||
}
|
// (`H5VM_bit_set`): every page of an allocated data block is
|
||||||
|
// written.
|
||||||
// Build EADB
|
let npages = nelmts / page_nelmts;
|
||||||
let mut aedb = Vec::new();
|
let mut bitmap = vec![0u8; sblk_bitmap_len(ndblks, nelmts)];
|
||||||
aedb.extend_from_slice(b"EADB");
|
for (k, db) in dblks.iter().enumerate() {
|
||||||
aedb.push(0);
|
if db.addr.is_some() {
|
||||||
aedb.push(client_id);
|
for p in 0..npages {
|
||||||
match offset_size {
|
let bit = k * npages + p;
|
||||||
4 => aedb.extend_from_slice(&(ea_base_address as u32).to_le_bytes()),
|
bitmap[bit / 8] |= 0x80 >> (bit % 8);
|
||||||
8 => aedb.extend_from_slice(&ea_base_address.to_le_bytes()),
|
}
|
||||||
_ => aedb.extend_from_slice(&ea_base_address.to_le_bytes()),
|
}
|
||||||
}
|
|
||||||
|
|
||||||
let blk_off_size = (max_nelmts_bits as usize).div_ceil(8);
|
|
||||||
let blk_off_val = (chunk_idx - n_inline) as u64;
|
|
||||||
aedb.extend_from_slice(&blk_off_val.to_le_bytes()[..blk_off_size]);
|
|
||||||
|
|
||||||
for slot in 0..nelmts {
|
|
||||||
if chunk_idx + slot < num_elements {
|
|
||||||
write_chunk_element(
|
|
||||||
&mut aedb,
|
|
||||||
&chunks[chunk_idx + slot],
|
|
||||||
offset_size,
|
|
||||||
has_filters,
|
|
||||||
chunk_size_bytes,
|
|
||||||
);
|
|
||||||
} else {
|
|
||||||
write_undefined_element(&mut aedb, offset_size, has_filters, chunk_size_bytes);
|
|
||||||
}
|
}
|
||||||
|
out.extend_from_slice(&bitmap);
|
||||||
}
|
}
|
||||||
|
for db in dblks {
|
||||||
let aedb_checksum = jenkins_lookup3(&aedb);
|
write_addr_opt(&mut out, db.addr);
|
||||||
aedb.extend_from_slice(&aedb_checksum.to_le_bytes());
|
}
|
||||||
|
let sum = jenkins_lookup3(&out[sb_start..]);
|
||||||
dblk_cursor += aedb.len() as u64;
|
out.extend_from_slice(&sum.to_le_bytes());
|
||||||
data_blocks_buf.extend_from_slice(&aedb);
|
for db in dblks.iter().filter(|d| d.addr.is_some()) {
|
||||||
chunk_idx += nelmts;
|
write_dblk(&mut out, db);
|
||||||
}
|
|
||||||
|
|
||||||
// Super block addresses (all undefined)
|
|
||||||
for _ in 0..n_sblk_addrs {
|
|
||||||
match offset_size {
|
|
||||||
4 => aeib.extend_from_slice(&u32::MAX.to_le_bytes()),
|
|
||||||
8 => aeib.extend_from_slice(&u64::MAX.to_le_bytes()),
|
|
||||||
_ => aeib.extend_from_slice(&u64::MAX.to_le_bytes()),
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
debug_assert_eq!(out.len() as u64, cursor - ea_base_address);
|
||||||
let aeib_checksum = jenkins_lookup3(&aeib);
|
out
|
||||||
aeib.extend_from_slice(&aeib_checksum.to_le_bytes());
|
|
||||||
debug_assert_eq!(aeib.len(), aeib_size);
|
|
||||||
|
|
||||||
let mut combined = aehd;
|
|
||||||
combined.extend_from_slice(&aeib);
|
|
||||||
combined.extend_from_slice(&data_blocks_buf);
|
|
||||||
combined
|
|
||||||
}
|
|
||||||
|
|
||||||
fn write_chunk_element(
|
|
||||||
buf: &mut Vec<u8>,
|
|
||||||
chunk: &WrittenChunk,
|
|
||||||
offset_size: u8,
|
|
||||||
has_filters: bool,
|
|
||||||
chunk_size_bytes: usize,
|
|
||||||
) {
|
|
||||||
match offset_size {
|
|
||||||
4 => buf.extend_from_slice(&(chunk.address as u32).to_le_bytes()),
|
|
||||||
8 => buf.extend_from_slice(&chunk.address.to_le_bytes()),
|
|
||||||
_ => buf.extend_from_slice(&chunk.address.to_le_bytes()),
|
|
||||||
}
|
|
||||||
if has_filters {
|
|
||||||
let cs_bytes = chunk.compressed_size.to_le_bytes();
|
|
||||||
buf.extend_from_slice(&cs_bytes[..chunk_size_bytes]);
|
|
||||||
buf.extend_from_slice(&chunk.filter_mask.to_le_bytes());
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn write_undefined_element(
|
|
||||||
buf: &mut Vec<u8>,
|
|
||||||
offset_size: u8,
|
|
||||||
has_filters: bool,
|
|
||||||
chunk_size_bytes: usize,
|
|
||||||
) {
|
|
||||||
let os = offset_size as usize;
|
|
||||||
// Use extend with repeat to avoid heap-allocating a temporary Vec on each call.
|
|
||||||
buf.extend(core::iter::repeat_n(0xFF, os));
|
|
||||||
if has_filters {
|
|
||||||
buf.extend(core::iter::repeat_n(0x00, chunk_size_bytes));
|
|
||||||
buf.extend_from_slice(&0u32.to_le_bytes());
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -117,6 +117,9 @@ pub enum FormatError {
|
|||||||
/// A message is marked shared but was parsed without access to the file,
|
/// A message is marked shared but was parsed without access to the file,
|
||||||
/// so the reference to the real message could not be followed.
|
/// so the reference to the real message could not be followed.
|
||||||
UnresolvedSharedMessage,
|
UnresolvedSharedMessage,
|
||||||
|
/// A selection does not fit the dataset it was applied to (wrong rank, or
|
||||||
|
/// it reaches past a dimension's extent).
|
||||||
|
SelectionOutOfBounds(String),
|
||||||
/// The dataset's raw data is stored in external files (External Data
|
/// The dataset's raw data is stored in external files (External Data
|
||||||
/// Files message), which this reader does not follow.
|
/// Files message), which this reader does not follow.
|
||||||
ExternalDataFilesUnsupported,
|
ExternalDataFilesUnsupported,
|
||||||
@@ -333,6 +336,9 @@ impl fmt::Display for FormatError {
|
|||||||
f,
|
f,
|
||||||
"dataset raw data is stored in external file(s), which is not supported"
|
"dataset raw data is stored in external file(s), which is not supported"
|
||||||
),
|
),
|
||||||
|
FormatError::SelectionOutOfBounds(msg) => {
|
||||||
|
write!(f, "selection out of bounds: {msg}")
|
||||||
|
}
|
||||||
FormatError::UnresolvedSharedMessage => write!(
|
FormatError::UnresolvedSharedMessage => write!(
|
||||||
f,
|
f,
|
||||||
"message is shared but no file data was available to resolve it"
|
"message is shared but no file data was available to resolve it"
|
||||||
|
|||||||
File diff suppressed because it is too large
Load Diff
@@ -4,7 +4,7 @@
|
|||||||
//! link messages, contiguous datasets, inline and dense attributes.
|
//! link messages, contiguous datasets, inline and dense attributes.
|
||||||
|
|
||||||
#[cfg(not(feature = "std"))]
|
#[cfg(not(feature = "std"))]
|
||||||
use alloc::{string::String, string::ToString, vec, vec::Vec};
|
use alloc::{format, string::String, string::ToString, vec, vec::Vec};
|
||||||
|
|
||||||
use crate::attribute::AttributeMessage;
|
use crate::attribute::AttributeMessage;
|
||||||
use crate::chunked_write::{
|
use crate::chunked_write::{
|
||||||
@@ -19,7 +19,7 @@ use crate::metadata_index::{DatasetMetadata, MetadataBlock, MetadataIndex};
|
|||||||
use crate::object_header_writer::ObjectHeaderWriter;
|
use crate::object_header_writer::ObjectHeaderWriter;
|
||||||
use crate::superblock::Superblock;
|
use crate::superblock::Superblock;
|
||||||
use crate::type_builders::{
|
use crate::type_builders::{
|
||||||
DatasetBuilder, FillTime, FinishedGroup, GroupBuilder, build_attr_message,
|
DatasetBuilder, FinishedGroup, GroupBuilder, build_attr_message, fill_value_message,
|
||||||
};
|
};
|
||||||
|
|
||||||
// Re-export public types that moved to type_builders for API compatibility.
|
// Re-export public types that moved to type_builders for API compatibility.
|
||||||
@@ -33,6 +33,49 @@ pub(crate) const OFFSET_SIZE: u8 = 8;
|
|||||||
pub(crate) const LENGTH_SIZE: u8 = 8;
|
pub(crate) const LENGTH_SIZE: u8 = 8;
|
||||||
const SUPERBLOCK_SIZE: usize = 48;
|
const SUPERBLOCK_SIZE: usize = 48;
|
||||||
|
|
||||||
|
/// Largest raw data a compact dataset can hold: the layout message (version,
|
||||||
|
/// class, 2-byte size, data) must fit an object header message, whose size
|
||||||
|
/// field is 2 bytes. Bigger "compact" requests fall back to contiguous storage.
|
||||||
|
const MAX_COMPACT_DATA_SIZE: usize = crate::object_header_writer::MAX_MESSAGE_SIZE - 4;
|
||||||
|
|
||||||
|
/// libhdf5's bounds on a file space page size (`H5F_FILE_SPACE_PAGE_SIZE_MIN`
|
||||||
|
/// and `_MAX`).
|
||||||
|
const MIN_FILE_SPACE_PAGE_SIZE: u32 = 512;
|
||||||
|
const MAX_FILE_SPACE_PAGE_SIZE: u32 = 1024 * 1024 * 1024;
|
||||||
|
|
||||||
|
/// Superblock extension object header for a file using the paged file-space
|
||||||
|
/// strategy: a single File Space Info message (0x0017), as libhdf5 writes it
|
||||||
|
/// for `fs_strategy="page"` without persisted free space.
|
||||||
|
fn build_paged_superblock_extension(page_size: u32) -> Result<Vec<u8>, FormatError> {
|
||||||
|
let mut fsinfo = Vec::new();
|
||||||
|
fsinfo.push(1); // version
|
||||||
|
fsinfo.push(1); // strategy: H5F_FSPACE_STRATEGY_PAGE
|
||||||
|
fsinfo.push(0); // persisting free space: no
|
||||||
|
write_length(&mut fsinfo, 1, LENGTH_SIZE); // free-space section threshold
|
||||||
|
write_length(&mut fsinfo, u64::from(page_size), LENGTH_SIZE);
|
||||||
|
fsinfo.extend_from_slice(&0u16.to_le_bytes()); // page end metadata threshold
|
||||||
|
write_undef_offset(&mut fsinfo, OFFSET_SIZE); // EOA before free-space info
|
||||||
|
let mut w = ObjectHeaderWriter::new();
|
||||||
|
// Flags as libhdf5 sets them: bit 2 (never share) and bit 4 (mark if
|
||||||
|
// unknown). Not constant: libhdf5 rewrites the message when it closes a
|
||||||
|
// file it opened for writing.
|
||||||
|
w.add_message_with_flags(MessageType::Unknown(0x0017), fsinfo, 0x14);
|
||||||
|
w.serialize()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// A group or dataset name must be one path component: not empty, not ".",
|
||||||
|
/// and without '/'. `FileWriter` writes a root group plus one level of
|
||||||
|
/// groups, and cannot create intermediate groups for a path.
|
||||||
|
fn check_link_name(name: &str) -> Result<(), FormatError> {
|
||||||
|
if name.is_empty() || name == "." || name.contains('/') {
|
||||||
|
return Err(FormatError::SerializationError(format!(
|
||||||
|
"invalid object name {name:?}: names must be a single path component \
|
||||||
|
(FileWriter does not create nested groups)"
|
||||||
|
)));
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
/// Threshold for switching from compact (inline) to dense attribute storage.
|
/// Threshold for switching from compact (inline) to dense attribute storage.
|
||||||
const DENSE_ATTR_THRESHOLD: usize = 8;
|
const DENSE_ATTR_THRESHOLD: usize = 8;
|
||||||
|
|
||||||
@@ -50,12 +93,12 @@ pub(crate) fn build_chunked_dataset_oh(
|
|||||||
pipeline_message: Option<&[u8]>,
|
pipeline_message: Option<&[u8]>,
|
||||||
attrs: &[AttributeMessage],
|
attrs: &[AttributeMessage],
|
||||||
dense_blob: Option<&DenseAttrBlob>,
|
dense_blob: Option<&DenseAttrBlob>,
|
||||||
fill_time: FillTime,
|
fill_message: &[u8],
|
||||||
) -> Vec<u8> {
|
) -> Result<Vec<u8>, FormatError> {
|
||||||
let mut w = ObjectHeaderWriter::new();
|
let mut w = ObjectHeaderWriter::new();
|
||||||
w.add_message_with_flags(MessageType::Datatype, dt.serialize(), 0x01);
|
w.add_message_with_flags(MessageType::Datatype, dt.serialize(), 0x01);
|
||||||
w.add_message(MessageType::Dataspace, ds.serialize(LENGTH_SIZE));
|
w.add_message(MessageType::Dataspace, ds.serialize(LENGTH_SIZE));
|
||||||
w.add_message_with_flags(MessageType::FillValue, vec![3, fill_time.to_byte()], 0x01);
|
w.add_message_with_flags(MessageType::FillValue, fill_message.to_vec(), 0x01);
|
||||||
w.add_message(MessageType::DataLayout, layout_message.to_vec());
|
w.add_message(MessageType::DataLayout, layout_message.to_vec());
|
||||||
if let Some(pm) = pipeline_message {
|
if let Some(pm) = pipeline_message {
|
||||||
w.add_message(MessageType::FilterPipeline, pm.to_vec());
|
w.add_message(MessageType::FilterPipeline, pm.to_vec());
|
||||||
@@ -77,15 +120,21 @@ pub(crate) fn build_dataset_oh(
|
|||||||
data_size: u64,
|
data_size: u64,
|
||||||
attrs: &[AttributeMessage],
|
attrs: &[AttributeMessage],
|
||||||
dense_blob: Option<&DenseAttrBlob>,
|
dense_blob: Option<&DenseAttrBlob>,
|
||||||
fill_time: FillTime,
|
fill_message: &[u8],
|
||||||
) -> Vec<u8> {
|
) -> Result<Vec<u8>, FormatError> {
|
||||||
let mut w = ObjectHeaderWriter::new();
|
let mut w = ObjectHeaderWriter::new();
|
||||||
w.add_message_with_flags(MessageType::Datatype, dt.serialize(), 0x01);
|
w.add_message_with_flags(MessageType::Datatype, dt.serialize(), 0x01);
|
||||||
w.add_message(MessageType::Dataspace, ds.serialize(LENGTH_SIZE));
|
w.add_message(MessageType::Dataspace, ds.serialize(LENGTH_SIZE));
|
||||||
w.add_message_with_flags(MessageType::FillValue, vec![3, fill_time.to_byte()], 0x01);
|
w.add_message_with_flags(MessageType::FillValue, fill_message.to_vec(), 0x01);
|
||||||
let mut dl = Vec::new();
|
let mut dl = Vec::new();
|
||||||
dl.push(4); // version
|
dl.push(4); // version
|
||||||
dl.push(1); // class = contiguous
|
dl.push(1); // class = contiguous
|
||||||
|
// An empty dataset has no storage: its address must be the undefined
|
||||||
|
// address, as libhdf5 writes it. A real address with size 0 trips
|
||||||
|
// libhdf5's `addr + size <= addr` overflow check, and it refuses the
|
||||||
|
// dataset as "invalid dataset size, likely file corruption" — which made
|
||||||
|
// every store with no sessions or knowledge graph unreadable by h5py.
|
||||||
|
let data_addr = if data_size == 0 { u64::MAX } else { data_addr };
|
||||||
dl.extend_from_slice(&data_addr.to_le_bytes());
|
dl.extend_from_slice(&data_addr.to_le_bytes());
|
||||||
dl.extend_from_slice(&data_size.to_le_bytes());
|
dl.extend_from_slice(&data_size.to_le_bytes());
|
||||||
w.add_message(MessageType::DataLayout, dl);
|
w.add_message(MessageType::DataLayout, dl);
|
||||||
@@ -106,12 +155,12 @@ pub(crate) fn build_compact_dataset_oh(
|
|||||||
data: &[u8],
|
data: &[u8],
|
||||||
attrs: &[AttributeMessage],
|
attrs: &[AttributeMessage],
|
||||||
dense_blob: Option<&DenseAttrBlob>,
|
dense_blob: Option<&DenseAttrBlob>,
|
||||||
fill_time: FillTime,
|
fill_message: &[u8],
|
||||||
) -> Vec<u8> {
|
) -> Result<Vec<u8>, FormatError> {
|
||||||
let mut w = ObjectHeaderWriter::new();
|
let mut w = ObjectHeaderWriter::new();
|
||||||
w.add_message_with_flags(MessageType::Datatype, dt.serialize(), 0x01);
|
w.add_message_with_flags(MessageType::Datatype, dt.serialize(), 0x01);
|
||||||
w.add_message(MessageType::Dataspace, ds.serialize(LENGTH_SIZE));
|
w.add_message(MessageType::Dataspace, ds.serialize(LENGTH_SIZE));
|
||||||
w.add_message_with_flags(MessageType::FillValue, vec![3, fill_time.to_byte()], 0x01);
|
w.add_message_with_flags(MessageType::FillValue, fill_message.to_vec(), 0x01);
|
||||||
// Compact layout message: version=4, class=0, u16 size, inline data
|
// Compact layout message: version=4, class=0, u16 size, inline data
|
||||||
let mut dl = Vec::new();
|
let mut dl = Vec::new();
|
||||||
dl.push(4); // version
|
dl.push(4); // version
|
||||||
@@ -134,7 +183,7 @@ pub(crate) fn build_group_oh(
|
|||||||
dense_link_info: Option<&[u8]>,
|
dense_link_info: Option<&[u8]>,
|
||||||
attrs: &[AttributeMessage],
|
attrs: &[AttributeMessage],
|
||||||
dense_blob: Option<&DenseAttrBlob>,
|
dense_blob: Option<&DenseAttrBlob>,
|
||||||
) -> Vec<u8> {
|
) -> Result<Vec<u8>, FormatError> {
|
||||||
let mut w = ObjectHeaderWriter::new();
|
let mut w = ObjectHeaderWriter::new();
|
||||||
if let Some(li) = dense_link_info {
|
if let Some(li) = dense_link_info {
|
||||||
// Dense link storage: a LinkInfo pointing at the fractal heap + name
|
// Dense link storage: a LinkInfo pointing at the fractal heap + name
|
||||||
@@ -896,12 +945,12 @@ pub(crate) fn build_vds_dataset_oh(
|
|||||||
global_heap_addr: u64,
|
global_heap_addr: u64,
|
||||||
attrs: &[AttributeMessage],
|
attrs: &[AttributeMessage],
|
||||||
dense_blob: Option<&DenseAttrBlob>,
|
dense_blob: Option<&DenseAttrBlob>,
|
||||||
fill_time: FillTime,
|
fill_message: &[u8],
|
||||||
) -> Vec<u8> {
|
) -> Result<Vec<u8>, FormatError> {
|
||||||
let mut w = ObjectHeaderWriter::new();
|
let mut w = ObjectHeaderWriter::new();
|
||||||
w.add_message_with_flags(MessageType::Datatype, dt.serialize(), 0x01);
|
w.add_message_with_flags(MessageType::Datatype, dt.serialize(), 0x01);
|
||||||
w.add_message(MessageType::Dataspace, ds.serialize(LENGTH_SIZE));
|
w.add_message(MessageType::Dataspace, ds.serialize(LENGTH_SIZE));
|
||||||
w.add_message_with_flags(MessageType::FillValue, vec![3, fill_time.to_byte()], 0x01);
|
w.add_message_with_flags(MessageType::FillValue, fill_message.to_vec(), 0x01);
|
||||||
// VDS layout message: version=4, class=3, global_heap_address(8), global_heap_index=1(4)
|
// VDS layout message: version=4, class=3, global_heap_address(8), global_heap_index=1(4)
|
||||||
let mut dl = Vec::new();
|
let mut dl = Vec::new();
|
||||||
dl.push(4u8); // version
|
dl.push(4u8); // version
|
||||||
@@ -950,7 +999,9 @@ pub struct FileWriter {
|
|||||||
alignment_threshold: usize,
|
alignment_threshold: usize,
|
||||||
/// Global alignment boundary in bytes (0 = disabled).
|
/// Global alignment boundary in bytes (0 = disabled).
|
||||||
alignment_bytes: usize,
|
alignment_bytes: usize,
|
||||||
/// Page size for page-buffer mode. When set, a v4 superblock is written.
|
/// File space page size. When set, the file uses libhdf5's paged
|
||||||
|
/// file-space strategy (a File Space Info message in the superblock
|
||||||
|
/// extension).
|
||||||
page_size: Option<u32>,
|
page_size: Option<u32>,
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -982,9 +1033,16 @@ impl FileWriter {
|
|||||||
self
|
self
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Enable page-buffer mode with the given page size. Writing this causes
|
/// Write the file with libhdf5's *paged* file-space strategy and the given
|
||||||
/// the file to be written with a v4 superblock (page_size field) instead
|
/// page size, as `H5Pset_file_space_strategy(H5F_FSPACE_STRATEGY_PAGE)` +
|
||||||
/// of the default v3.
|
/// `H5Pset_file_space_page_size` (h5py: `fs_strategy="page"`,
|
||||||
|
/// `fs_page_size=...`) do: a v3 superblock with an extension holding a
|
||||||
|
/// File Space Info message, and the file padded to a whole number of
|
||||||
|
/// pages. Readers with a page buffer can then fetch metadata page by page.
|
||||||
|
///
|
||||||
|
/// `page_size` must be between 512 bytes and 1 GiB (libhdf5's limits);
|
||||||
|
/// [`Self::finish`] fails otherwise. This used to write a "version 4"
|
||||||
|
/// superblock, which does not exist and no HDF5 library can open.
|
||||||
pub fn with_page_size(&mut self, page_size: u32) -> &mut Self {
|
pub fn with_page_size(&mut self, page_size: u32) -> &mut Self {
|
||||||
self.page_size = Some(page_size);
|
self.page_size = Some(page_size);
|
||||||
self
|
self
|
||||||
@@ -1009,6 +1067,14 @@ impl FileWriter {
|
|||||||
|
|
||||||
pub fn finish(self) -> Result<Vec<u8>, FormatError> {
|
pub fn finish(self) -> Result<Vec<u8>, FormatError> {
|
||||||
let page_size = self.page_size;
|
let page_size = self.page_size;
|
||||||
|
if let Some(ps) = page_size
|
||||||
|
&& !(MIN_FILE_SPACE_PAGE_SIZE..=MAX_FILE_SPACE_PAGE_SIZE).contains(&ps)
|
||||||
|
{
|
||||||
|
return Err(FormatError::SerializationError(format!(
|
||||||
|
"file space page size {ps} is outside libhdf5's \
|
||||||
|
{MIN_FILE_SPACE_PAGE_SIZE}..={MAX_FILE_SPACE_PAGE_SIZE} bytes"
|
||||||
|
)));
|
||||||
|
}
|
||||||
struct DsFlat {
|
struct DsFlat {
|
||||||
name: String,
|
name: String,
|
||||||
dt: Datatype,
|
dt: Datatype,
|
||||||
@@ -1017,7 +1083,8 @@ impl FileWriter {
|
|||||||
attrs: Vec<AttributeMessage>,
|
attrs: Vec<AttributeMessage>,
|
||||||
chunk_options: ChunkOptions,
|
chunk_options: ChunkOptions,
|
||||||
maxshape: Option<Vec<u64>>,
|
maxshape: Option<Vec<u64>>,
|
||||||
fill_time: FillTime,
|
/// Serialized Fill Value message.
|
||||||
|
fill_message: Vec<u8>,
|
||||||
compact: bool,
|
compact: bool,
|
||||||
alignment: usize,
|
alignment: usize,
|
||||||
/// VDS source mappings (set for Virtual datasets).
|
/// VDS source mappings (set for Virtual datasets).
|
||||||
@@ -1067,6 +1134,7 @@ impl FileWriter {
|
|||||||
};
|
};
|
||||||
attrs.extend(p.build_attrs(&raw));
|
attrs.extend(p.build_attrs(&raw));
|
||||||
}
|
}
|
||||||
|
let fill_message = fill_value_message(db.fill_time, db.fill_value.as_deref(), &dt)?;
|
||||||
Ok(DsFlat {
|
Ok(DsFlat {
|
||||||
name: db.name,
|
name: db.name,
|
||||||
dt,
|
dt,
|
||||||
@@ -1075,13 +1143,26 @@ impl FileWriter {
|
|||||||
attrs,
|
attrs,
|
||||||
chunk_options: db.chunk_options,
|
chunk_options: db.chunk_options,
|
||||||
maxshape: db.maxshape,
|
maxshape: db.maxshape,
|
||||||
fill_time: db.fill_time,
|
fill_message,
|
||||||
compact: db.compact,
|
compact: db.compact,
|
||||||
alignment: db.alignment,
|
alignment: db.alignment,
|
||||||
virtual_sources: db.virtual_sources,
|
virtual_sources: db.virtual_sources,
|
||||||
})
|
})
|
||||||
};
|
};
|
||||||
|
|
||||||
|
// Every name becomes a single link in its parent group. The writer
|
||||||
|
// has no nested groups, so a path like "a/b" would be stored as one
|
||||||
|
// link literally named "a/b" — which no HDF5 reader can resolve.
|
||||||
|
let root_names = self.root_datasets.iter().map(|d| d.name.as_str());
|
||||||
|
let group_names = self.groups.iter().flat_map(|g| {
|
||||||
|
core::iter::once(g.name.as_str())
|
||||||
|
.chain(g.datasets.iter().map(|d| d.name.as_str()))
|
||||||
|
.chain(g.external_links.iter().map(|l| l.0.as_str()))
|
||||||
|
});
|
||||||
|
for name in root_names.chain(group_names) {
|
||||||
|
check_link_name(name)?;
|
||||||
|
}
|
||||||
|
|
||||||
let mut all_ds: Vec<DsFlat> = Vec::new();
|
let mut all_ds: Vec<DsFlat> = Vec::new();
|
||||||
let mut groups: Vec<GrpFlat> = Vec::new();
|
let mut groups: Vec<GrpFlat> = Vec::new();
|
||||||
let mut root_ds_indices: Vec<usize> = Vec::new();
|
let mut root_ds_indices: Vec<usize> = Vec::new();
|
||||||
@@ -1114,17 +1195,35 @@ impl FileWriter {
|
|||||||
root_attrs.push(build_attr_message(n, v));
|
root_attrs.push(build_attr_message(n, v));
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Every datatype must have an on-disk encoding before anything is laid
|
||||||
|
// out: `Datatype::serialize` itself cannot report a failure.
|
||||||
|
let group_attrs = groups.iter().flat_map(|g| &g.attrs);
|
||||||
|
let ds_attrs = all_ds.iter().flat_map(|d| &d.attrs);
|
||||||
|
for a in root_attrs.iter().chain(group_attrs).chain(ds_attrs) {
|
||||||
|
a.datatype.check_encodable()?;
|
||||||
|
}
|
||||||
|
for d in &all_ds {
|
||||||
|
d.dt.check_encodable()?;
|
||||||
|
}
|
||||||
|
|
||||||
let is_vds: Vec<bool> = all_ds.iter().map(|d| d.virtual_sources.is_some()).collect();
|
let is_vds: Vec<bool> = all_ds.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()
|
||||||
.map(|(i, d)| !is_vds[i] && (d.chunk_options.is_chunked() || d.maxshape.is_some()))
|
.map(|(i, d)| {
|
||||||
|
// Only a dataset that can grow needs chunks; a maxshape equal
|
||||||
|
// to the shape is as fixed as no maxshape at all.
|
||||||
|
let resizable = d.maxshape.as_ref().is_some_and(|m| *m != d.ds.dimensions);
|
||||||
|
!is_vds[i] && (d.chunk_options.is_chunked() || resizable)
|
||||||
|
})
|
||||||
.collect();
|
.collect();
|
||||||
// Determine which datasets use compact storage
|
// Determine which datasets use compact storage
|
||||||
let is_compact: Vec<bool> = all_ds
|
let is_compact: Vec<bool> = all_ds
|
||||||
.iter()
|
.iter()
|
||||||
.enumerate()
|
.enumerate()
|
||||||
.map(|(i, d)| !is_vds[i] && !is_chunked[i] && d.compact && d.raw.len() <= 65535)
|
.map(|(i, d)| {
|
||||||
|
!is_vds[i] && !is_chunked[i] && d.compact && d.raw.len() <= MAX_COMPACT_DATA_SIZE
|
||||||
|
})
|
||||||
.collect();
|
.collect();
|
||||||
let root_dense = root_attrs.len() > DENSE_ATTR_THRESHOLD;
|
let root_dense = root_attrs.len() > DENSE_ATTR_THRESHOLD;
|
||||||
let group_dense: Vec<bool> = groups
|
let group_dense: Vec<bool> = groups
|
||||||
@@ -1163,9 +1262,9 @@ impl FileWriter {
|
|||||||
}
|
}
|
||||||
let attr_blob = group_dense[gi].then(|| build_dense_attrs(&g.attrs, 0));
|
let attr_blob = group_dense[gi].then(|| build_dense_attrs(&g.attrs, 0));
|
||||||
let dl = group_links_dense[gi].then_some(dummy_link_info.as_slice());
|
let dl = group_links_dense[gi].then_some(dummy_link_info.as_slice());
|
||||||
build_group_oh(&dummy_links, dl, &g.attrs, attr_blob.as_ref()).len()
|
build_group_oh(&dummy_links, dl, &g.attrs, attr_blob.as_ref()).map(|oh| oh.len())
|
||||||
})
|
})
|
||||||
.collect();
|
.collect::<Result<_, _>>()?;
|
||||||
|
|
||||||
let root_dummy_links: Vec<LinkMessage> = {
|
let root_dummy_links: Vec<LinkMessage> = {
|
||||||
let mut links = Vec::new();
|
let mut links = Vec::new();
|
||||||
@@ -1180,7 +1279,7 @@ impl FileWriter {
|
|||||||
let root_oh_size = {
|
let root_oh_size = {
|
||||||
let attr_blob = root_dense.then(|| build_dense_attrs(&root_attrs, 0));
|
let attr_blob = root_dense.then(|| build_dense_attrs(&root_attrs, 0));
|
||||||
let dl = root_links_dense.then_some(dummy_link_info.as_slice());
|
let dl = root_links_dense.then_some(dummy_link_info.as_slice());
|
||||||
build_group_oh(&root_dummy_links, dl, &root_attrs, attr_blob.as_ref()).len()
|
build_group_oh(&root_dummy_links, dl, &root_attrs, attr_blob.as_ref())?.len()
|
||||||
};
|
};
|
||||||
|
|
||||||
struct DataBlob {
|
struct DataBlob {
|
||||||
@@ -1208,8 +1307,8 @@ impl FileWriter {
|
|||||||
0, // dummy address
|
0, // dummy address
|
||||||
&d.attrs,
|
&d.attrs,
|
||||||
dense_blob.as_ref(),
|
dense_blob.as_ref(),
|
||||||
d.fill_time,
|
&d.fill_message,
|
||||||
);
|
)?;
|
||||||
// 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(&[]);
|
||||||
@@ -1221,8 +1320,10 @@ impl FileWriter {
|
|||||||
precompressed: None,
|
precompressed: None,
|
||||||
});
|
});
|
||||||
} else if is_chunked[i] {
|
} else if is_chunked[i] {
|
||||||
let chunk_dims = d.chunk_options.resolve_chunk_dims(&d.ds.dimensions);
|
|
||||||
let elem_size = d.dt.type_size() as usize;
|
let elem_size = d.dt.type_size() as usize;
|
||||||
|
let chunk_dims = d
|
||||||
|
.chunk_options
|
||||||
|
.resolve_chunk_dims_for(&d.ds.dimensions, elem_size);
|
||||||
// Compress once in Pass 1; cache the result so Pass 2 can skip
|
// Compress once in Pass 1; cache the result so Pass 2 can skip
|
||||||
// re-compression and just rebuild the index with real addresses.
|
// re-compression and just rebuild the index with real addresses.
|
||||||
let pre = precompress_chunks(
|
let pre = precompress_chunks(
|
||||||
@@ -1236,7 +1337,7 @@ impl FileWriter {
|
|||||||
&pre,
|
&pre,
|
||||||
dummy_cursor,
|
dummy_cursor,
|
||||||
d.maxshape.as_deref(),
|
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))
|
||||||
@@ -1250,8 +1351,8 @@ impl FileWriter {
|
|||||||
result.pipeline_message.as_deref(),
|
result.pipeline_message.as_deref(),
|
||||||
&d.attrs,
|
&d.attrs,
|
||||||
dense_blob.as_ref(),
|
dense_blob.as_ref(),
|
||||||
d.fill_time,
|
&d.fill_message,
|
||||||
);
|
)?;
|
||||||
dummy_blobs.push(DataBlob {
|
dummy_blobs.push(DataBlob {
|
||||||
data: result.data_bytes,
|
data: result.data_bytes,
|
||||||
oh_bytes: oh,
|
oh_bytes: oh,
|
||||||
@@ -1269,8 +1370,8 @@ impl FileWriter {
|
|||||||
&d.raw,
|
&d.raw,
|
||||||
&d.attrs,
|
&d.attrs,
|
||||||
dense_blob.as_ref(),
|
dense_blob.as_ref(),
|
||||||
d.fill_time,
|
&d.fill_message,
|
||||||
);
|
)?;
|
||||||
dummy_blobs.push(DataBlob {
|
dummy_blobs.push(DataBlob {
|
||||||
data: vec![],
|
data: vec![],
|
||||||
oh_bytes: oh,
|
oh_bytes: oh,
|
||||||
@@ -1289,8 +1390,8 @@ impl FileWriter {
|
|||||||
d.raw.len() as u64,
|
d.raw.len() as u64,
|
||||||
&d.attrs,
|
&d.attrs,
|
||||||
dense_blob.as_ref(),
|
dense_blob.as_ref(),
|
||||||
d.fill_time,
|
&d.fill_message,
|
||||||
);
|
)?;
|
||||||
dummy_blobs.push(DataBlob {
|
dummy_blobs.push(DataBlob {
|
||||||
data: d.raw.clone(),
|
data: d.raw.clone(),
|
||||||
oh_bytes: oh,
|
oh_bytes: oh,
|
||||||
@@ -1302,12 +1403,12 @@ impl FileWriter {
|
|||||||
let actual_ds_oh_sizes: Vec<usize> = dummy_blobs.iter().map(|b| b.oh_bytes.len()).collect();
|
let actual_ds_oh_sizes: Vec<usize> = dummy_blobs.iter().map(|b| b.oh_bytes.len()).collect();
|
||||||
|
|
||||||
// Pass 2: compute real addresses
|
// Pass 2: compute real addresses
|
||||||
// v4 superblocks add a 4-byte page_size field before the checksum.
|
// A paged file carries its File Space Info in a superblock extension
|
||||||
let superblock_size = if page_size.is_some() {
|
// object header, placed right after the superblock.
|
||||||
SUPERBLOCK_SIZE + 4
|
let sb_ext = page_size
|
||||||
} else {
|
.map(build_paged_superblock_extension)
|
||||||
SUPERBLOCK_SIZE
|
.transpose()?;
|
||||||
};
|
let superblock_size = SUPERBLOCK_SIZE + sb_ext.as_ref().map_or(0, Vec::len);
|
||||||
let root_group_addr = superblock_size as u64;
|
let root_group_addr = superblock_size as u64;
|
||||||
let mut cursor2 = superblock_size + root_oh_size;
|
let mut cursor2 = superblock_size + root_oh_size;
|
||||||
|
|
||||||
@@ -1398,8 +1499,8 @@ impl FileWriter {
|
|||||||
heap_addr,
|
heap_addr,
|
||||||
&d.attrs,
|
&d.attrs,
|
||||||
ds_dense_blobs[i].as_ref(),
|
ds_dense_blobs[i].as_ref(),
|
||||||
d.fill_time,
|
&d.fill_message,
|
||||||
);
|
)?;
|
||||||
ds_blobs2.push(DataBlob {
|
ds_blobs2.push(DataBlob {
|
||||||
data: gcol_bytes.clone(),
|
data: gcol_bytes.clone(),
|
||||||
oh_bytes: oh,
|
oh_bytes: oh,
|
||||||
@@ -1416,7 +1517,7 @@ impl FileWriter {
|
|||||||
.expect("chunked dataset missing precompressed cache"),
|
.expect("chunked dataset missing precompressed cache"),
|
||||||
base_address,
|
base_address,
|
||||||
d.maxshape.as_deref(),
|
d.maxshape.as_deref(),
|
||||||
);
|
)?;
|
||||||
cursor2 += result.data_bytes.len();
|
cursor2 += result.data_bytes.len();
|
||||||
let oh = build_chunked_dataset_oh(
|
let oh = build_chunked_dataset_oh(
|
||||||
&d.dt,
|
&d.dt,
|
||||||
@@ -1425,8 +1526,8 @@ impl FileWriter {
|
|||||||
result.pipeline_message.as_deref(),
|
result.pipeline_message.as_deref(),
|
||||||
&d.attrs,
|
&d.attrs,
|
||||||
ds_dense_blobs[i].as_ref(),
|
ds_dense_blobs[i].as_ref(),
|
||||||
d.fill_time,
|
&d.fill_message,
|
||||||
);
|
)?;
|
||||||
ds_blobs2.push(DataBlob {
|
ds_blobs2.push(DataBlob {
|
||||||
data: result.data_bytes,
|
data: result.data_bytes,
|
||||||
oh_bytes: oh,
|
oh_bytes: oh,
|
||||||
@@ -1440,8 +1541,8 @@ impl FileWriter {
|
|||||||
&d.raw,
|
&d.raw,
|
||||||
&d.attrs,
|
&d.attrs,
|
||||||
ds_dense_blobs[i].as_ref(),
|
ds_dense_blobs[i].as_ref(),
|
||||||
d.fill_time,
|
&d.fill_message,
|
||||||
);
|
)?;
|
||||||
ds_blobs2.push(DataBlob {
|
ds_blobs2.push(DataBlob {
|
||||||
data: vec![],
|
data: vec![],
|
||||||
oh_bytes: oh,
|
oh_bytes: oh,
|
||||||
@@ -1465,8 +1566,8 @@ impl FileWriter {
|
|||||||
d.raw.len() as u64,
|
d.raw.len() as u64,
|
||||||
&d.attrs,
|
&d.attrs,
|
||||||
ds_dense_blobs[i].as_ref(),
|
ds_dense_blobs[i].as_ref(),
|
||||||
d.fill_time,
|
&d.fill_message,
|
||||||
);
|
)?;
|
||||||
let mut data = vec![0u8; padding];
|
let mut data = vec![0u8; padding];
|
||||||
data.extend_from_slice(&d.raw);
|
data.extend_from_slice(&d.raw);
|
||||||
cursor2 += d.raw.len();
|
cursor2 += d.raw.len();
|
||||||
@@ -1481,11 +1582,16 @@ impl FileWriter {
|
|||||||
let actual_ds_oh_sizes2: Vec<usize> = ds_blobs2.iter().map(|b| b.oh_bytes.len()).collect();
|
let actual_ds_oh_sizes2: Vec<usize> = ds_blobs2.iter().map(|b| b.oh_bytes.len()).collect();
|
||||||
debug_assert_eq!(actual_ds_oh_sizes, actual_ds_oh_sizes2);
|
debug_assert_eq!(actual_ds_oh_sizes, actual_ds_oh_sizes2);
|
||||||
|
|
||||||
|
// libhdf5 ends a paged file on a page boundary.
|
||||||
|
let data_end = cursor2;
|
||||||
|
if let Some(ps) = page_size {
|
||||||
|
cursor2 = cursor2.next_multiple_of(ps as usize);
|
||||||
|
}
|
||||||
let eof_addr2 = cursor2 as u64;
|
let eof_addr2 = cursor2 as u64;
|
||||||
let mut buf = Vec::with_capacity(cursor2);
|
let mut buf = Vec::with_capacity(cursor2);
|
||||||
|
|
||||||
let sb = Superblock {
|
let sb = Superblock {
|
||||||
version: if page_size.is_some() { 4 } else { 3 },
|
version: 3,
|
||||||
offset_size: OFFSET_SIZE,
|
offset_size: OFFSET_SIZE,
|
||||||
length_size: LENGTH_SIZE,
|
length_size: LENGTH_SIZE,
|
||||||
base_address: 0,
|
base_address: 0,
|
||||||
@@ -1497,11 +1603,18 @@ impl FileWriter {
|
|||||||
free_space_address: None,
|
free_space_address: None,
|
||||||
driver_info_address: None,
|
driver_info_address: None,
|
||||||
consistency_flags: 0,
|
consistency_flags: 0,
|
||||||
superblock_extension_address: Some(u64::MAX),
|
superblock_extension_address: Some(if sb_ext.is_some() {
|
||||||
|
SUPERBLOCK_SIZE as u64
|
||||||
|
} else {
|
||||||
|
u64::MAX
|
||||||
|
}),
|
||||||
checksum: None,
|
checksum: None,
|
||||||
page_size,
|
page_size: None,
|
||||||
};
|
};
|
||||||
buf.extend_from_slice(&sb.serialize());
|
buf.extend_from_slice(&sb.serialize());
|
||||||
|
if let Some(ref ext) = sb_ext {
|
||||||
|
buf.extend_from_slice(ext);
|
||||||
|
}
|
||||||
|
|
||||||
// Root group OH
|
// Root group OH
|
||||||
let mut root_links: Vec<LinkMessage> = Vec::new();
|
let mut root_links: Vec<LinkMessage> = Vec::new();
|
||||||
@@ -1522,7 +1635,7 @@ impl FileWriter {
|
|||||||
root_dl,
|
root_dl,
|
||||||
&root_attrs,
|
&root_attrs,
|
||||||
root_dense_blob.as_ref(),
|
root_dense_blob.as_ref(),
|
||||||
));
|
)?);
|
||||||
if let Some(ref b) = root_link_blob {
|
if let Some(ref b) = root_link_blob {
|
||||||
buf.extend_from_slice(&b.blob);
|
buf.extend_from_slice(&b.blob);
|
||||||
}
|
}
|
||||||
@@ -1547,7 +1660,7 @@ impl FileWriter {
|
|||||||
dl,
|
dl,
|
||||||
&g.attrs,
|
&g.attrs,
|
||||||
group_dense_blobs[gi].as_ref(),
|
group_dense_blobs[gi].as_ref(),
|
||||||
));
|
)?);
|
||||||
if let Some(ref b) = link_blob {
|
if let Some(ref b) = link_blob {
|
||||||
buf.extend_from_slice(&b.blob);
|
buf.extend_from_slice(&b.blob);
|
||||||
}
|
}
|
||||||
@@ -1569,7 +1682,8 @@ impl FileWriter {
|
|||||||
buf.extend_from_slice(&blob.data);
|
buf.extend_from_slice(&blob.data);
|
||||||
}
|
}
|
||||||
|
|
||||||
debug_assert_eq!(buf.len(), cursor2);
|
debug_assert_eq!(buf.len(), data_end);
|
||||||
|
buf.resize(cursor2, 0);
|
||||||
Ok(buf)
|
Ok(buf)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -2143,7 +2257,8 @@ mod tests {
|
|||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn file_writer_v4_superblock() {
|
fn file_writer_paged_file_uses_v3_superblock_and_fsinfo_extension() {
|
||||||
|
// This used to write superblock "version 4", which does not exist.
|
||||||
let mut fw = FileWriter::new();
|
let mut fw = FileWriter::new();
|
||||||
fw.with_page_size(4096);
|
fw.with_page_size(4096);
|
||||||
fw.create_dataset("data").with_f64_data(&[1.0, 2.0]);
|
fw.create_dataset("data").with_f64_data(&[1.0, 2.0]);
|
||||||
@@ -2151,8 +2266,31 @@ 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();
|
||||||
assert_eq!(sb.version, 4, "expected superblock v4");
|
assert_eq!(sb.version, 3);
|
||||||
assert_eq!(sb.page_size, Some(4096));
|
assert_eq!(sb.superblock_extension_address, Some(48));
|
||||||
|
assert_eq!(bytes.len() % 4096, 0);
|
||||||
|
assert_eq!(sb.eof_address, bytes.len() as u64);
|
||||||
|
let ext = ObjectHeader::parse(&bytes, 48, 8, 8).unwrap();
|
||||||
|
let fsinfo = &ext.messages[0];
|
||||||
|
assert_eq!(fsinfo.msg_type, MessageType::Unknown(0x0017));
|
||||||
|
// Byte-for-byte what HDF5 2.0 writes for fs_strategy="page",
|
||||||
|
// fs_page_size=4096.
|
||||||
|
let mut expected = vec![1u8, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0];
|
||||||
|
expected.extend_from_slice(&4096u64.to_le_bytes());
|
||||||
|
expected.extend_from_slice(&[0, 0]);
|
||||||
|
expected.extend_from_slice(&[0xff; 8]);
|
||||||
|
assert_eq!(fsinfo.data, expected);
|
||||||
|
assert_eq!(fsinfo.flags, 0x14);
|
||||||
|
assert_eq!(read_dataset_f64(&bytes, "data"), vec![1.0, 2.0]);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn file_writer_rejects_page_sizes_libhdf5_would() {
|
||||||
|
for ps in [0u32, 511, MAX_FILE_SPACE_PAGE_SIZE + 1] {
|
||||||
|
let mut fw = FileWriter::new();
|
||||||
|
fw.with_page_size(ps);
|
||||||
|
assert!(fw.finish().is_err(), "page size {ps}");
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
|
|||||||
@@ -98,15 +98,50 @@ pub fn parse_fill_value(msg: &HeaderMessage) -> Result<Option<Vec<u8>>, FormatEr
|
|||||||
|
|
||||||
/// The fill value that applies to a dataset given its header messages. The new
|
/// The fill value that applies to a dataset given its header messages. The new
|
||||||
/// message wins over the old one when both are present.
|
/// message wins over the old one when both are present.
|
||||||
|
///
|
||||||
|
/// A *shared* fill value message holds only a reference to the real message,
|
||||||
|
/// which cannot be followed without the file: this returns
|
||||||
|
/// [`FormatError::UnresolvedSharedMessage`] for one (it used to answer "zeros").
|
||||||
|
/// Use [`dataset_fill_value_in`] when the file bytes are at hand.
|
||||||
pub fn dataset_fill_value(messages: &[HeaderMessage]) -> Result<Option<Vec<u8>>, FormatError> {
|
pub fn dataset_fill_value(messages: &[HeaderMessage]) -> Result<Option<Vec<u8>>, FormatError> {
|
||||||
|
fill_value_from(messages, |_| Err(FormatError::UnresolvedSharedMessage))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// [`dataset_fill_value`] for a dataset in `file_data`, following a shared
|
||||||
|
/// fill value message to where it lives: another object header, or the
|
||||||
|
/// file's shared-message (SOHM) heap, as libhdf5 writes it when the file has
|
||||||
|
/// a SOHM index for fill values.
|
||||||
|
pub fn dataset_fill_value_in(
|
||||||
|
file_data: &[u8],
|
||||||
|
messages: &[HeaderMessage],
|
||||||
|
offset_size: u8,
|
||||||
|
length_size: u8,
|
||||||
|
) -> Result<Option<Vec<u8>>, FormatError> {
|
||||||
|
fill_value_from(messages, |msg| {
|
||||||
|
crate::shared_message::message_data_with_sohm(file_data, msg, offset_size, length_size)
|
||||||
|
.map(|data| data.into_owned())
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
fn fill_value_from(
|
||||||
|
messages: &[HeaderMessage],
|
||||||
|
resolve_shared: impl Fn(&HeaderMessage) -> Result<Vec<u8>, FormatError>,
|
||||||
|
) -> Result<Option<Vec<u8>>, FormatError> {
|
||||||
for wanted in [MessageType::FillValue, MessageType::FillValueOld] {
|
for wanted in [MessageType::FillValue, MessageType::FillValueOld] {
|
||||||
if let Some(msg) = messages.iter().find(|m| m.msg_type == wanted) {
|
if let Some(msg) = messages.iter().find(|m| m.msg_type == wanted) {
|
||||||
if crate::shared_message::is_shared(msg.flags) {
|
let value = if crate::shared_message::is_shared(msg.flags) {
|
||||||
// A shared fill value is legal but vanishingly rare; treat it
|
let data = resolve_shared(msg)?;
|
||||||
// as the default rather than misparsing the reference.
|
parse_fill_value(&HeaderMessage {
|
||||||
return Ok(None);
|
msg_type: msg.msg_type,
|
||||||
}
|
size: data.len(),
|
||||||
if let Some(value) = parse_fill_value(msg)? {
|
flags: msg.flags & !0x02,
|
||||||
|
creation_order: msg.creation_order,
|
||||||
|
data,
|
||||||
|
})?
|
||||||
|
} else {
|
||||||
|
parse_fill_value(msg)?
|
||||||
|
};
|
||||||
|
if let Some(value) = value {
|
||||||
return Ok(Some(value));
|
return Ok(Some(value));
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -174,7 +209,7 @@ pub fn read_full_with_fill<E: From<FormatError>>(
|
|||||||
{
|
{
|
||||||
return Err(FormatError::ExternalDataFilesUnsupported.into());
|
return Err(FormatError::ExternalDataFilesUnsupported.into());
|
||||||
}
|
}
|
||||||
let fill = dataset_fill_value(messages)?;
|
let fill = dataset_fill_value_in(file_data, messages, offset_size, length_size)?;
|
||||||
if !has_storage(layout) {
|
if !has_storage(layout) {
|
||||||
return Ok(filled_dataset(dataspace, elem_size, fill.as_deref())?);
|
return Ok(filled_dataset(dataspace, elem_size, fill.as_deref())?);
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -19,8 +19,23 @@ pub const FILTER_SCALEOFFSET: u16 = 6;
|
|||||||
pub const FILTER_LZ4: u16 = 32004;
|
pub const FILTER_LZ4: u16 = 32004;
|
||||||
/// Zstandard compression.
|
/// Zstandard compression.
|
||||||
pub const FILTER_ZSTD: u16 = 32015;
|
pub const FILTER_ZSTD: u16 = 32015;
|
||||||
/// Pcodec lossless numerical codec (clawhdf5 internal; not yet HDF5-registered).
|
/// Pcodec lossless numerical codec — a **private, unregistered** clawhdf5
|
||||||
pub const FILTER_PCODEC: u16 = 32023;
|
/// filter. Pcodec has no ID in the HDF Group's filter registry (checked
|
||||||
|
/// 2026-09-25, `hdf5_plugins/docs/RegisteredFilterPlugins.md`), so it uses an
|
||||||
|
/// ID from the registry's testing/private range (256–511). No libhdf5 plugin
|
||||||
|
/// decodes it: h5py/libhdf5 report the filter as unavailable. Only clawhdf5
|
||||||
|
/// (with the `pcodec` feature) reads these datasets.
|
||||||
|
pub const FILTER_PCODEC: u16 = 480;
|
||||||
|
/// Filter name written with [`FILTER_PCODEC`].
|
||||||
|
pub const FILTER_PCODEC_NAME: &str = "pcodec (clawhdf5 private)";
|
||||||
|
/// The ID clawhdf5 up to 2.7.0 wrote pcodec under. It is registered to
|
||||||
|
/// Granular BitRound (GBR), whose decode is a pass-through, so libhdf5 with
|
||||||
|
/// that plugin would have returned the compressed bytes as data. Read as
|
||||||
|
/// pcodec only when the filter is named exactly [`FILTER_PCODEC_LEGACY_NAME`],
|
||||||
|
/// the name those versions wrote; never written.
|
||||||
|
pub const FILTER_PCODEC_LEGACY: u16 = 32023;
|
||||||
|
/// The filter name clawhdf5 up to 2.7.0 wrote with [`FILTER_PCODEC_LEGACY`].
|
||||||
|
pub const FILTER_PCODEC_LEGACY_NAME: &str = "pcodec";
|
||||||
|
|
||||||
/// Description of a single filter in a pipeline.
|
/// Description of a single filter in a pipeline.
|
||||||
#[derive(Debug, Clone, PartialEq)]
|
#[derive(Debug, Clone, PartialEq)]
|
||||||
|
|||||||
File diff suppressed because it is too large
Load Diff
@@ -1,19 +1,39 @@
|
|||||||
//! SZIP (libaec Adaptive Entropy Coding) decompression.
|
//! SZIP (libaec Adaptive Entropy Coding) decompression.
|
||||||
//!
|
//!
|
||||||
//! Gated by the `szip` feature which links against the system libaec library.
|
//! Gated by the `szip` feature which links against the system libaec library.
|
||||||
|
//!
|
||||||
|
//! libhdf5's SZIP filter (`H5Zszip.c`) prefixes each chunk with its
|
||||||
|
//! uncompressed size and hands the rest to szlib's `SZ_BufftoBuffDecompress`.
|
||||||
|
//! libaec implements that call (`sz_compat.c`) on top of `aec_buffer_decode`
|
||||||
|
//! with some reshaping — 32/64-bit samples are coded as byte planes of 8-bit
|
||||||
|
//! samples, and scanlines that are not a whole number of blocks are padded —
|
||||||
|
//! which [`szip_decompress`] reproduces so its output matches libhdf5's.
|
||||||
|
|
||||||
#[cfg(not(feature = "std"))]
|
#[cfg(not(feature = "std"))]
|
||||||
use alloc::vec::Vec;
|
use alloc::vec::Vec;
|
||||||
|
|
||||||
use crate::error::FormatError;
|
use crate::error::FormatError;
|
||||||
|
|
||||||
/// Decompress SZIP-compressed data using libaec.
|
/// `SZ_MSB_OPTION_MASK`: samples are big-endian.
|
||||||
|
#[cfg(feature = "szip")]
|
||||||
|
const SZ_MSB_OPTION_MASK: u32 = 16;
|
||||||
|
/// `SZ_NN_OPTION_MASK`: nearest-neighbour preprocessing.
|
||||||
|
#[cfg(feature = "szip")]
|
||||||
|
const SZ_NN_OPTION_MASK: u32 = 32;
|
||||||
|
|
||||||
|
/// Decompress one SZIP-filtered chunk.
|
||||||
///
|
///
|
||||||
/// `cd` is the HDF5 SZIP filter client data (matches `H5Z_SZIP_PARM_*` indices):
|
/// `cd` is the HDF5 SZIP filter client data (`H5Z_SZIP_PARM_*` indices):
|
||||||
/// cd[0] = options mask (`H5_SZIP_NN_OPTION_MASK = 0x20` enables NN preprocessing)
|
/// cd[0] = options mask (`SZ_*_OPTION_MASK`: 16 = MSB byte order,
|
||||||
/// cd[1] = pixels per block (H5Z_SZIP_PARM_PPB; 8, 10, 16, or 32)
|
/// 32 = nearest-neighbour preprocessing; K13/EC/LSB/RAW bits carry
|
||||||
/// cd[2] = bits per sample (H5Z_SZIP_PARM_BPP; element bit width)
|
/// no decoding information for libaec)
|
||||||
/// cd[3] = pixels per scan line (H5Z_SZIP_PARM_PPS; informational only)
|
/// cd[1] = pixels per block
|
||||||
|
/// cd[2] = bits per pixel (sample precision, rounded up to 32 or 64 above
|
||||||
|
/// 24 by libhdf5)
|
||||||
|
/// cd[3] = pixels per scanline
|
||||||
|
///
|
||||||
|
/// The chunk is a 4-byte little-endian uncompressed size followed by the
|
||||||
|
/// szlib stream.
|
||||||
pub(crate) fn szip_decompress(
|
pub(crate) fn szip_decompress(
|
||||||
_data: &[u8],
|
_data: &[u8],
|
||||||
_cd: &[u32],
|
_cd: &[u32],
|
||||||
@@ -33,62 +53,174 @@ pub(crate) fn szip_decompress(
|
|||||||
|
|
||||||
#[cfg(feature = "szip")]
|
#[cfg(feature = "szip")]
|
||||||
fn szip_decode_impl(data: &[u8], cd: &[u32], chunk_size: usize) -> Result<Vec<u8>, FormatError> {
|
fn szip_decode_impl(data: &[u8], cd: &[u32], chunk_size: usize) -> Result<Vec<u8>, FormatError> {
|
||||||
if cd.len() < 3 {
|
let err = |m: &str| FormatError::ChunkedReadError(format!("szip: {m}"));
|
||||||
return Err(FormatError::ChunkedReadError(
|
if cd.len() < 4 {
|
||||||
"szip: missing client data".into(),
|
return Err(err("missing client data"));
|
||||||
));
|
|
||||||
}
|
}
|
||||||
let options = cd[0];
|
let options = cd[0];
|
||||||
let pixels_per_block = cd[1];
|
let pixels_per_block = cd[1] as usize;
|
||||||
let bits_per_sample = cd[2]; // H5Z_SZIP_PARM_BPP
|
let bits_per_pixel = cd[2];
|
||||||
if bits_per_sample == 0 || bits_per_sample > 32 {
|
let pixels_per_scanline = cd[3] as usize;
|
||||||
return Err(FormatError::ChunkedReadError(
|
if !(1..=32).contains(&bits_per_pixel) && bits_per_pixel != 64 {
|
||||||
"szip: invalid bits per sample".into(),
|
return Err(err("invalid bits per sample"));
|
||||||
));
|
|
||||||
}
|
}
|
||||||
if chunk_size == 0 {
|
if pixels_per_block == 0 || pixels_per_scanline == 0 {
|
||||||
return Err(FormatError::ChunkedReadError(
|
return Err(err("invalid block or scanline size"));
|
||||||
"szip: unknown output size".into(),
|
|
||||||
));
|
|
||||||
}
|
}
|
||||||
if data.is_empty() {
|
if data.len() < 4 {
|
||||||
return Err(FormatError::ChunkedReadError("szip: empty input".into()));
|
return Err(err("chunk too short"));
|
||||||
}
|
}
|
||||||
|
// H5Zszip.c: UINT32DECODE of the uncompressed size, then the stream.
|
||||||
|
let dest_len = u32::from_le_bytes([data[0], data[1], data[2], data[3]]) as usize;
|
||||||
|
let limit = if chunk_size != 0 {
|
||||||
|
chunk_size
|
||||||
|
} else {
|
||||||
|
crate::filters::MAX_DECOMPRESS_SIZE
|
||||||
|
};
|
||||||
|
if dest_len > limit {
|
||||||
|
return Err(err("declared size exceeds chunk size"));
|
||||||
|
}
|
||||||
|
let stream = &data[4..];
|
||||||
|
|
||||||
// Map HDF5 option mask to libaec flags.
|
// --- libaec sz_compat.c: SZ_BufftoBuffDecompress ---
|
||||||
// HDF5 always stores SZIP data in MSB order, so AEC_DATA_MSB is unconditional.
|
let rsi = pixels_per_scanline.div_ceil(pixels_per_block);
|
||||||
// H5_SZIP_NN_OPTION_MASK (0x20): NN differential preprocessing.
|
let mut flags = 0;
|
||||||
let mut flags: u32 = libaec_sys::AEC_DATA_MSB;
|
if options & SZ_MSB_OPTION_MASK != 0 {
|
||||||
if options & 0x20 != 0 {
|
flags |= libaec_sys::AEC_DATA_MSB;
|
||||||
|
}
|
||||||
|
if options & SZ_NN_OPTION_MASK != 0 {
|
||||||
flags |= libaec_sys::AEC_DATA_PREPROCESS;
|
flags |= libaec_sys::AEC_DATA_PREPROCESS;
|
||||||
}
|
}
|
||||||
|
let pad_scanline = !pixels_per_scanline.is_multiple_of(pixels_per_block);
|
||||||
|
let deinterleave = bits_per_pixel == 32 || bits_per_pixel == 64;
|
||||||
|
let bits_per_sample = if deinterleave { 8 } else { bits_per_pixel };
|
||||||
|
let pixel_size = match bits_per_sample {
|
||||||
|
17.. => 4,
|
||||||
|
9.. => 2,
|
||||||
|
_ => 1,
|
||||||
|
};
|
||||||
|
let scanlines = (dest_len / pixel_size).div_ceil(pixels_per_scanline);
|
||||||
|
let buf_size = if pad_scanline {
|
||||||
|
rsi.checked_mul(pixels_per_block)
|
||||||
|
.and_then(|n| n.checked_mul(pixel_size))
|
||||||
|
.and_then(|n| n.checked_mul(scanlines))
|
||||||
|
.filter(|&n| n <= crate::filters::MAX_DECOMPRESS_SIZE.max(limit))
|
||||||
|
.ok_or_else(|| err("scanline padding too large"))?
|
||||||
|
} else {
|
||||||
|
dest_len
|
||||||
|
};
|
||||||
|
|
||||||
let mut out = vec![0u8; chunk_size];
|
let mut buf = vec![0u8; buf_size];
|
||||||
let mut strm = libaec_sys::AecStream::zeroed();
|
let mut strm = libaec_sys::AecStream::zeroed();
|
||||||
strm.next_in = data.as_ptr();
|
strm.next_in = stream.as_ptr();
|
||||||
strm.avail_in = data.len();
|
strm.avail_in = stream.len();
|
||||||
strm.next_out = out.as_mut_ptr();
|
strm.next_out = buf.as_mut_ptr();
|
||||||
strm.avail_out = chunk_size;
|
strm.avail_out = buf_size;
|
||||||
strm.bits_per_sample = bits_per_sample;
|
strm.bits_per_sample = bits_per_sample;
|
||||||
strm.block_size = pixels_per_block;
|
strm.block_size = pixels_per_block as u32;
|
||||||
strm.rsi = 128; // HDF5 default: 128 blocks per reference sample interval
|
strm.rsi = rsi as u32;
|
||||||
strm.flags = flags;
|
strm.flags = flags;
|
||||||
|
// SAFETY: next_in/avail_in and next_out/avail_out describe live buffers
|
||||||
|
// (`stream` and `buf`) that outlive the call.
|
||||||
let result = unsafe { libaec_sys::aec_buffer_decode(&mut strm) };
|
let result = unsafe { libaec_sys::aec_buffer_decode(&mut strm) };
|
||||||
if result != 0 {
|
if result != 0 {
|
||||||
return Err(FormatError::DecompressionError(format!(
|
return Err(FormatError::DecompressionError(format!(
|
||||||
"szip: libaec error {result}"
|
"szip: libaec error {result}"
|
||||||
)));
|
)));
|
||||||
}
|
}
|
||||||
let decoded_len = chunk_size - strm.avail_out;
|
let mut total_out = strm.total_out;
|
||||||
out.truncate(decoded_len);
|
if pad_scanline {
|
||||||
Ok(out)
|
let line = pixels_per_scanline * pixel_size;
|
||||||
|
let padded_line = rsi * pixels_per_block * pixel_size;
|
||||||
|
// remove_padding: compact each padded line down to `line` bytes.
|
||||||
|
let mut i = line;
|
||||||
|
let mut j = padded_line;
|
||||||
|
while j < total_out {
|
||||||
|
let end = (j + line).min(buf.len());
|
||||||
|
buf.copy_within(j..end, i);
|
||||||
|
i += line;
|
||||||
|
j += padded_line;
|
||||||
|
}
|
||||||
|
total_out = scanlines * line;
|
||||||
|
}
|
||||||
|
if total_out < dest_len {
|
||||||
|
return Err(err("stream decoded to fewer bytes than declared"));
|
||||||
|
}
|
||||||
|
buf.truncate(dest_len);
|
||||||
|
if deinterleave {
|
||||||
|
// deinterleave_buffer: byte planes back into words.
|
||||||
|
let w = (bits_per_pixel / 8) as usize;
|
||||||
|
let n = dest_len / w;
|
||||||
|
let mut out = vec![0u8; dest_len];
|
||||||
|
for i in 0..n {
|
||||||
|
for j in 0..w {
|
||||||
|
out[i * w + j] = buf[j * n + i];
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(out)
|
||||||
|
} else {
|
||||||
|
Ok(buf)
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use super::*;
|
use super::*;
|
||||||
|
|
||||||
|
#[cfg(feature = "szip")]
|
||||||
|
fn unhex(s: &str) -> Vec<u8> {
|
||||||
|
(0..s.len())
|
||||||
|
.step_by(2)
|
||||||
|
.map(|i| u8::from_str_radix(&s[i..i + 2], 16).unwrap())
|
||||||
|
.collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// SZIP chunks written by libhdf5, decoded exactly as libhdf5 decodes
|
||||||
|
/// them. Each case: fixture, chunk byte offset and size (from h5py's
|
||||||
|
/// `get_chunk_info`), the filter's cd_values, and the chunk's values as
|
||||||
|
/// h5py reads them (file byte order, hex). Before the fix every one of
|
||||||
|
/// these came back as garbage or zeros (or "invalid bits per sample" for
|
||||||
|
/// 64-bit): the 4-byte size prefix was fed to libaec, 32/64-bit samples
|
||||||
|
/// were not de-interleaved from byte planes, the reference sample
|
||||||
|
/// interval was fixed at 128 instead of derived from the scanline, padded
|
||||||
|
/// scanlines were not unpadded, and LE data was decoded as MSB.
|
||||||
|
#[cfg(feature = "szip")]
|
||||||
|
#[test]
|
||||||
|
fn szip_decodes_libhdf5_chunks_exactly() {
|
||||||
|
/// (name, file, chunk offset, chunk size, cd_values, decoded hex)
|
||||||
|
type Case<'a> = (&'a str, &'a [u8], usize, usize, [u32; 4], &'a str);
|
||||||
|
let noencoder: &[u8] = include_bytes!("../tests/fixtures/filters/noencoder.h5");
|
||||||
|
let le_data: &[u8] = include_bytes!("../tests/fixtures/filters/le_data.h5");
|
||||||
|
let h5py: &[u8] = include_bytes!("../tests/fixtures/filters/szip_h5py.h5");
|
||||||
|
#[rustfmt::skip]
|
||||||
|
let cases: &[Case] = &[
|
||||||
|
// <i4, 10 px/scanline over 4 px/block: padded scanlines + byte planes.
|
||||||
|
("noencoder /noencoder_szip_dset.h5", noencoder, 6040, 16, [168, 4, 32, 10],
|
||||||
|
"00000000010000000200000003000000040000000500000006000000070000000800000009000000"),
|
||||||
|
// <f4, LSB + NN.
|
||||||
|
("le_data /Szip_float_data_le", le_data, 55224, 48, [169, 4, 32, 12],
|
||||||
|
"abaaaa3eabaa2a3f0000803fabaa2a3f0000803fabaaaa3f0000803fabaaaa3f5555d53fabaaaa3f5555d53f00000040"),
|
||||||
|
// >f4, MSB + NN.
|
||||||
|
("le_data /Szip_float_data_be", le_data, 55396, 48, [177, 4, 32, 12],
|
||||||
|
"3eaaaaab3f2aaaab3f8000003f2aaaab3f8000003faaaaab3f8000003faaaaab3fd555553faaaaab3fd5555540000000"),
|
||||||
|
// <f8 (64-bit), NN.
|
||||||
|
("szip_h5py /f8", h5py, 4016, 100, [169, 8, 64, 10],
|
||||||
|
"00000000000008c000000000000008c000000000000008c000000000000008c000000000000004c000000000000004c000000000000004c000000000000004c000000000000000c000000000000000c000000000000000c000000000000000c0000000000000f8bf000000000000f8bf000000000000f8bf000000000000f8bf000000000000f0bf000000000000f0bf000000000000f0bf000000000000f0bf000000000000e0bf000000000000e0bf000000000000e0bf000000000000e0bf0000000000000000000000000000000000000000000000000000000000000000000000000000e03f000000000000e03f000000000000e03f000000000000e03f000000000000f03f000000000000f03f000000000000f03f000000000000f03f000000000000f83f000000000000f83f000000000000f83f000000000000f83f"),
|
||||||
|
// <i8 (64-bit), entropy coding without NN.
|
||||||
|
("szip_h5py /i8", h5py, 4188, 53, [141, 4, 64, 10],
|
||||||
|
"000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000300000000000000030000000000000003000000000000000300000000000000030000000000000003000000000000000300000000000000030000000000000006000000000000000600000000000000060000000000000006000000000000000600000000000000060000000000000006000000000000000600000000000000090000000000000009000000000000000900000000000000090000000000000009000000000000000900000000000000090000000000000009000000000000000c000000000000000c000000000000000c000000000000000c000000000000000c000000000000000c000000000000000c000000000000000c00000000000000"),
|
||||||
|
// <u2, 35 px/scanline over 8 px/block: padded scanlines, 16-bit samples.
|
||||||
|
("szip_h5py /u2", h5py, 4308, 43, [169, 8, 16, 35],
|
||||||
|
"00000000000000006100610061006100c200c200c200c20023012301230123018401840184018401e501e501e501e5014602460246024602a702a702a702a702080308030803"),
|
||||||
|
];
|
||||||
|
for (name, file, off, len, cd, want) in cases {
|
||||||
|
let want = unhex(want);
|
||||||
|
let got = szip_decompress(&file[*off..off + len], cd, want.len())
|
||||||
|
.unwrap_or_else(|e| panic!("{name}: {e:?}"));
|
||||||
|
assert_eq!(got, want, "{name}");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn szip_disabled_returns_unsupported() {
|
fn szip_disabled_returns_unsupported() {
|
||||||
#[cfg(not(feature = "szip"))]
|
#[cfg(not(feature = "szip"))]
|
||||||
@@ -132,6 +264,8 @@ mod tests {
|
|||||||
assert_eq!(rc, 0, "aec_buffer_encode failed: {rc}");
|
assert_eq!(rc, 0, "aec_buffer_encode failed: {rc}");
|
||||||
let enc_len = encoded.len() - enc.avail_out;
|
let enc_len = encoded.len() - enc.avail_out;
|
||||||
encoded.truncate(enc_len);
|
encoded.truncate(enc_len);
|
||||||
|
// H5Zszip.c prefixes the stream with the uncompressed size.
|
||||||
|
encoded.splice(0..0, (original.len() as u32).to_le_bytes());
|
||||||
|
|
||||||
// Decode through our public interface.
|
// Decode through our public interface.
|
||||||
// cd[0]=0 (no NN bit 0x20), cd[1]=8 (ppb), cd[2]=8 (bpp), cd[3]=1024 (pps).
|
// cd[0]=0 (no NN bit 0x20), cd[1]=8 (ppb), cd[2]=8 (bpp), cd[3]=1024 (pps).
|
||||||
@@ -163,6 +297,8 @@ mod tests {
|
|||||||
assert_eq!(rc, 0, "aec_buffer_encode with NN failed: {rc}");
|
assert_eq!(rc, 0, "aec_buffer_encode with NN failed: {rc}");
|
||||||
let enc_len = encoded.len() - enc.avail_out;
|
let enc_len = encoded.len() - enc.avail_out;
|
||||||
encoded.truncate(enc_len);
|
encoded.truncate(enc_len);
|
||||||
|
// H5Zszip.c prefixes the stream with the uncompressed size.
|
||||||
|
encoded.splice(0..0, (original.len() as u32).to_le_bytes());
|
||||||
|
|
||||||
// cd[0] = 0x20 (H5_SZIP_NN_OPTION_MASK) → decoder must set AEC_DATA_PREPROCESS.
|
// cd[0] = 0x20 (H5_SZIP_NN_OPTION_MASK) → decoder must set AEC_DATA_PREPROCESS.
|
||||||
let cd = [0x20u32, 8, 8, 1024];
|
let cd = [0x20u32, 8, 8, 1024];
|
||||||
|
|||||||
@@ -6,9 +6,35 @@ extern crate alloc;
|
|||||||
#[cfg(not(feature = "std"))]
|
#[cfg(not(feature = "std"))]
|
||||||
use alloc::{format, vec, vec::Vec};
|
use alloc::{format, vec, vec::Vec};
|
||||||
|
|
||||||
|
use crate::chunk_grid::ChunkGrid;
|
||||||
use crate::chunked_read::ChunkInfo;
|
use crate::chunked_read::ChunkInfo;
|
||||||
use crate::error::FormatError;
|
use crate::error::FormatError;
|
||||||
|
|
||||||
|
/// Verify the Jenkins lookup3 checksum stored immediately after
|
||||||
|
/// `data[start..end]`, as every Fixed Array structure carries one.
|
||||||
|
///
|
||||||
|
/// A corrupt chunk index silently yields addresses pointing at the wrong
|
||||||
|
/// bytes, so a mismatch has to be an error rather than a shrug: without this
|
||||||
|
/// the damage surfaces as plausible-looking data from the wrong chunk.
|
||||||
|
#[cfg(feature = "checksum")]
|
||||||
|
fn verify_checksum(data: &[u8], start: usize, end: usize) -> Result<(), FormatError> {
|
||||||
|
ensure_len(data, end, 4)?;
|
||||||
|
let stored = u32::from_le_bytes([data[end], data[end + 1], data[end + 2], data[end + 3]]);
|
||||||
|
let computed = crate::checksum::jenkins_lookup3(&data[start..end]);
|
||||||
|
if computed != stored {
|
||||||
|
return Err(FormatError::ChecksumMismatch {
|
||||||
|
expected: stored,
|
||||||
|
computed,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(not(feature = "checksum"))]
|
||||||
|
fn verify_checksum(_data: &[u8], _start: usize, _end: usize) -> Result<(), FormatError> {
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
/// Parsed Fixed Array header (FAHD).
|
/// Parsed Fixed Array header (FAHD).
|
||||||
#[derive(Debug, Clone)]
|
#[derive(Debug, Clone)]
|
||||||
pub struct FixedArrayHeader {
|
pub struct FixedArrayHeader {
|
||||||
@@ -103,6 +129,8 @@ impl FixedArrayHeader {
|
|||||||
let num_elements = read_length(d, pos, length_size)?;
|
let num_elements = read_length(d, pos, length_size)?;
|
||||||
pos += length_size as usize;
|
pos += length_size as usize;
|
||||||
let data_block_address = read_offset(d, pos, offset_size)?;
|
let data_block_address = read_offset(d, pos, offset_size)?;
|
||||||
|
pos += offset_size as usize;
|
||||||
|
verify_checksum(file_data, offset, offset + pos)?;
|
||||||
|
|
||||||
Ok(FixedArrayHeader {
|
Ok(FixedArrayHeader {
|
||||||
client_id,
|
client_id,
|
||||||
@@ -124,13 +152,13 @@ pub fn read_fixed_array_chunks(
|
|||||||
file_data: &[u8],
|
file_data: &[u8],
|
||||||
header: &FixedArrayHeader,
|
header: &FixedArrayHeader,
|
||||||
dataset_dims: &[u64],
|
dataset_dims: &[u64],
|
||||||
|
max_dims: Option<&[u64]>,
|
||||||
chunk_dimensions: &[u32],
|
chunk_dimensions: &[u32],
|
||||||
element_size: u32,
|
element_size: u32,
|
||||||
offset_size: u8,
|
offset_size: u8,
|
||||||
_length_size: u8,
|
_length_size: u8,
|
||||||
) -> Result<Vec<ChunkInfo>, FormatError> {
|
) -> Result<Vec<ChunkInfo>, FormatError> {
|
||||||
let db_offset = header.data_block_address as usize;
|
let db_offset = header.data_block_address as usize;
|
||||||
let rank = chunk_dimensions.len();
|
|
||||||
|
|
||||||
// Parse data block header: FADB(4) + version(1) + client_id(1) + header_address(offset_size)
|
// Parse data block header: FADB(4) + version(1) + client_id(1) + header_address(offset_size)
|
||||||
let db_header_size = 4 + 1 + 1 + offset_size as usize;
|
let db_header_size = 4 + 1 + 1 + offset_size as usize;
|
||||||
@@ -171,19 +199,10 @@ pub fn read_fixed_array_chunks(
|
|||||||
))
|
))
|
||||||
};
|
};
|
||||||
|
|
||||||
// Compute chunk offsets based on index.
|
// The index is laid out over the chunk grid of the *maximum* dimensions
|
||||||
// Chunks are stored in row-major order within the dataset space.
|
// (row-major), so a dataset smaller than its maxshape has gaps.
|
||||||
let mut num_chunks_per_dim = Vec::with_capacity(rank);
|
let dims_u64: Vec<u64> = chunk_dimensions.iter().map(|&d| d as u64).collect();
|
||||||
for d_idx in 0..rank {
|
let grid = ChunkGrid::fixed_array(dataset_dims, max_dims, &dims_u64)?;
|
||||||
let ch_dim = chunk_dimensions[d_idx] as u64;
|
|
||||||
if ch_dim == 0 {
|
|
||||||
return Err(FormatError::ChunkedReadError(
|
|
||||||
"chunk dimension is zero".into(),
|
|
||||||
));
|
|
||||||
}
|
|
||||||
let ds_dim = dataset_dims[d_idx];
|
|
||||||
num_chunks_per_dim.push(ds_dim.div_ceil(ch_dim));
|
|
||||||
}
|
|
||||||
|
|
||||||
let chunk_byte_size: u64 =
|
let chunk_byte_size: u64 =
|
||||||
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;
|
||||||
@@ -199,7 +218,11 @@ pub fn read_fixed_array_chunks(
|
|||||||
header.element_size,
|
header.element_size,
|
||||||
chunk_byte_size,
|
chunk_byte_size,
|
||||||
)? {
|
)? {
|
||||||
let offsets = index_to_chunk_offsets(i, &num_chunks_per_dim, chunk_dimensions);
|
// A slot beyond the current extent is ignored, as the
|
||||||
|
// library does.
|
||||||
|
let Some(offsets) = grid.offsets(i as u64) else {
|
||||||
|
return Ok(());
|
||||||
|
};
|
||||||
chunks.push(ChunkInfo {
|
chunks.push(ChunkInfo {
|
||||||
chunk_size,
|
chunk_size,
|
||||||
filter_mask,
|
filter_mask,
|
||||||
@@ -223,7 +246,8 @@ pub fn read_fixed_array_chunks(
|
|||||||
|
|
||||||
if !is_paged {
|
if !is_paged {
|
||||||
// Non-paged: prefix, then `num_elements` elements packed directly,
|
// Non-paged: prefix, then `num_elements` elements packed directly,
|
||||||
// then a trailing checksum (which we don't validate).
|
// then a checksum over both.
|
||||||
|
verify_checksum(file_data, db_offset, elem_at(elements_start, num_elements)?)?;
|
||||||
for i in 0..num_elements {
|
for i in 0..num_elements {
|
||||||
push_element(i, elem_at(elements_start, i)?, &mut chunks)?;
|
push_element(i, elem_at(elements_start, i)?, &mut chunks)?;
|
||||||
}
|
}
|
||||||
@@ -254,6 +278,9 @@ pub fn read_fixed_array_chunks(
|
|||||||
available: file_data.len(),
|
available: file_data.len(),
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
|
// The prefix and page bitmap are covered by their own checksum, and each
|
||||||
|
// initialised page by one of its own.
|
||||||
|
verify_checksum(file_data, db_offset, bitmap_start + bitmap_size)?;
|
||||||
|
|
||||||
for p in 0..npages {
|
for p in 0..npages {
|
||||||
let page_first = p * page_nelmts; // < num_elements, cannot overflow
|
let page_first = p * page_nelmts; // < num_elements, cannot overflow
|
||||||
@@ -270,6 +297,7 @@ pub fn read_fixed_array_chunks(
|
|||||||
.checked_mul(page_stride)
|
.checked_mul(page_stride)
|
||||||
.and_then(|o| pages_start.checked_add(o))
|
.and_then(|o| pages_start.checked_add(o))
|
||||||
.ok_or_else(stride_overflow)?;
|
.ok_or_else(stride_overflow)?;
|
||||||
|
verify_checksum(file_data, page_off, elem_at(page_off, page_count)?)?;
|
||||||
for e in 0..page_count {
|
for e in 0..page_count {
|
||||||
push_element(page_first + e, elem_at(page_off, e)?, &mut chunks)?;
|
push_element(page_first + e, elem_at(page_off, e)?, &mut chunks)?;
|
||||||
}
|
}
|
||||||
@@ -335,27 +363,6 @@ fn parse_fa_element(
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Convert a linear chunk index to N-dimensional chunk offsets in dataset space.
|
|
||||||
fn index_to_chunk_offsets(
|
|
||||||
index: usize,
|
|
||||||
num_chunks_per_dim: &[u64],
|
|
||||||
chunk_dimensions: &[u32],
|
|
||||||
) -> Vec<u64> {
|
|
||||||
let rank = num_chunks_per_dim.len();
|
|
||||||
let mut offsets = vec![0u64; rank];
|
|
||||||
let mut remaining = index as u64;
|
|
||||||
for d in (0..rank).rev() {
|
|
||||||
let nchunks = num_chunks_per_dim[d];
|
|
||||||
if nchunks == 0 {
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
let chunk_idx = remaining % nchunks;
|
|
||||||
remaining /= nchunks;
|
|
||||||
offsets[d] = chunk_idx * chunk_dimensions[d] as u64;
|
|
||||||
}
|
|
||||||
offsets
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Read a variable-length little-endian unsigned integer.
|
/// Read a variable-length little-endian unsigned integer.
|
||||||
fn read_variable_length(data: &[u8], size: usize) -> Result<u64, FormatError> {
|
fn read_variable_length(data: &[u8], size: usize) -> Result<u64, FormatError> {
|
||||||
if size > 8 || data.len() < size {
|
if size > 8 || data.len() < size {
|
||||||
@@ -374,46 +381,31 @@ fn read_variable_length(data: &[u8], size: usize) -> Result<u64, FormatError> {
|
|||||||
mod tests {
|
mod tests {
|
||||||
use super::*;
|
use super::*;
|
||||||
|
|
||||||
|
/// Stamp the Jenkins checksum a real file would carry over
|
||||||
|
/// `data[start..end]`, writing it at `end`. Fixtures built by hand need
|
||||||
|
/// this now that the reader validates it — as every HDF5 writer does.
|
||||||
|
fn stamp_checksum(data: &mut [u8], start: usize, end: usize) {
|
||||||
|
let sum = crate::checksum::jenkins_lookup3(&data[start..end]);
|
||||||
|
data[end..end + 4].copy_from_slice(&sum.to_le_bytes());
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn index_to_offsets_1d() {
|
fn index_to_offsets_1d() {
|
||||||
let num_chunks = vec![5u64];
|
let g = ChunkGrid::fixed_array(&[100], None, &[20]).unwrap();
|
||||||
let chunk_dims = vec![20u32];
|
assert_eq!(g.offsets(0).unwrap(), vec![0]);
|
||||||
assert_eq!(index_to_chunk_offsets(0, &num_chunks, &chunk_dims), vec![0]);
|
assert_eq!(g.offsets(1).unwrap(), vec![20]);
|
||||||
assert_eq!(
|
assert_eq!(g.offsets(4).unwrap(), vec![80]);
|
||||||
index_to_chunk_offsets(1, &num_chunks, &chunk_dims),
|
|
||||||
vec![20]
|
|
||||||
);
|
|
||||||
assert_eq!(
|
|
||||||
index_to_chunk_offsets(4, &num_chunks, &chunk_dims),
|
|
||||||
vec![80]
|
|
||||||
);
|
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn index_to_offsets_2d() {
|
fn index_to_offsets_2d() {
|
||||||
// 10x6 dataset with 4x3 chunks => ceil(10/4)=3, ceil(6/3)=2 => 6 chunks
|
// 10x6 dataset with 4x3 chunks => ceil(10/4)=3, ceil(6/3)=2 => 6 chunks
|
||||||
let num_chunks = vec![3u64, 2];
|
let g = ChunkGrid::fixed_array(&[10, 6], None, &[4, 3]).unwrap();
|
||||||
let chunk_dims = vec![4u32, 3];
|
assert_eq!(g.offsets(0).unwrap(), vec![0, 0]);
|
||||||
assert_eq!(
|
assert_eq!(g.offsets(1).unwrap(), vec![0, 3]);
|
||||||
index_to_chunk_offsets(0, &num_chunks, &chunk_dims),
|
assert_eq!(g.offsets(2).unwrap(), vec![4, 0]);
|
||||||
vec![0, 0]
|
assert_eq!(g.offsets(3).unwrap(), vec![4, 3]);
|
||||||
);
|
assert_eq!(g.offsets(5).unwrap(), vec![8, 3]);
|
||||||
assert_eq!(
|
|
||||||
index_to_chunk_offsets(1, &num_chunks, &chunk_dims),
|
|
||||||
vec![0, 3]
|
|
||||||
);
|
|
||||||
assert_eq!(
|
|
||||||
index_to_chunk_offsets(2, &num_chunks, &chunk_dims),
|
|
||||||
vec![4, 0]
|
|
||||||
);
|
|
||||||
assert_eq!(
|
|
||||||
index_to_chunk_offsets(3, &num_chunks, &chunk_dims),
|
|
||||||
vec![4, 3]
|
|
||||||
);
|
|
||||||
assert_eq!(
|
|
||||||
index_to_chunk_offsets(5, &num_chunks, &chunk_dims),
|
|
||||||
vec![8, 3]
|
|
||||||
);
|
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
@@ -439,7 +431,7 @@ mod tests {
|
|||||||
buf[8..16].copy_from_slice(&5u64.to_le_bytes());
|
buf[8..16].copy_from_slice(&5u64.to_le_bytes());
|
||||||
// data_block_address (offset_size=8)
|
// data_block_address (offset_size=8)
|
||||||
buf[16..24].copy_from_slice(&0x1000u64.to_le_bytes());
|
buf[16..24].copy_from_slice(&0x1000u64.to_le_bytes());
|
||||||
// checksum (4 bytes, we don't validate in parse)
|
stamp_checksum(&mut buf, 0, 24);
|
||||||
|
|
||||||
let header = FixedArrayHeader::parse(&buf, 0, 8, 8).unwrap();
|
let header = FixedArrayHeader::parse(&buf, 0, 8, 8).unwrap();
|
||||||
assert_eq!(header.client_id, 1);
|
assert_eq!(header.client_id, 1);
|
||||||
@@ -449,6 +441,54 @@ mod tests {
|
|||||||
assert_eq!(header.data_block_address, 0x1000);
|
assert_eq!(header.data_block_address, 0x1000);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Corruption anywhere in the index must be an error, not a wrong
|
||||||
|
/// address. Every structure carries a checksum; flipping a bit in each in
|
||||||
|
/// turn must be caught, because the alternative is reading a chunk from
|
||||||
|
/// the wrong offset and returning it as data.
|
||||||
|
#[test]
|
||||||
|
fn corrupting_any_fixed_array_structure_is_detected() {
|
||||||
|
let build = || -> (Vec<u8>, usize) {
|
||||||
|
let (os, fahd, db) = (8usize, 0x100usize, 0x200usize);
|
||||||
|
let mut f = vec![0u8; 0x3000];
|
||||||
|
f[fahd..fahd + 4].copy_from_slice(b"FAHD");
|
||||||
|
f[fahd + 6] = os as u8;
|
||||||
|
f[fahd + 7] = 10;
|
||||||
|
f[fahd + 8..fahd + 16].copy_from_slice(&3u64.to_le_bytes());
|
||||||
|
f[fahd + 16..fahd + 24].copy_from_slice(&(db as u64).to_le_bytes());
|
||||||
|
stamp_checksum(&mut f, fahd, fahd + 24);
|
||||||
|
f[db..db + 4].copy_from_slice(b"FADB");
|
||||||
|
f[db + 6..db + 14].copy_from_slice(&(fahd as u64).to_le_bytes());
|
||||||
|
let elems = db + 6 + os;
|
||||||
|
for i in 0..3usize {
|
||||||
|
let addr = 0x1000u64 + i as u64 * 0x100;
|
||||||
|
f[elems + i * os..elems + (i + 1) * os].copy_from_slice(&addr.to_le_bytes());
|
||||||
|
}
|
||||||
|
stamp_checksum(&mut f, db, elems + 3 * os);
|
||||||
|
(f, fahd)
|
||||||
|
};
|
||||||
|
|
||||||
|
let read = |f: &[u8], fahd: usize| -> Result<Vec<ChunkInfo>, FormatError> {
|
||||||
|
let h = FixedArrayHeader::parse(f, fahd, 8, 8)?;
|
||||||
|
read_fixed_array_chunks(f, &h, &[60], None, &[20], 8, 8, 8)
|
||||||
|
};
|
||||||
|
|
||||||
|
let (clean, fahd) = build();
|
||||||
|
assert!(read(&clean, fahd).is_ok(), "the intact fixture must read");
|
||||||
|
|
||||||
|
// A byte inside the header, and one inside a data block element.
|
||||||
|
for &at in &[0x108usize, 0x210usize] {
|
||||||
|
let (mut damaged, fahd) = build();
|
||||||
|
damaged[at] ^= 0x01;
|
||||||
|
assert!(
|
||||||
|
matches!(
|
||||||
|
read(&damaged, fahd),
|
||||||
|
Err(FormatError::ChecksumMismatch { .. })
|
||||||
|
),
|
||||||
|
"corruption at {at:#x} went undetected"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn parse_fixed_array_header_invalid_signature() {
|
fn parse_fixed_array_header_invalid_signature() {
|
||||||
let mut buf = vec![0u8; 256];
|
let mut buf = vec![0u8; 256];
|
||||||
@@ -469,11 +509,12 @@ mod tests {
|
|||||||
buf[fahd + 7] = 200; // max_nelmts_bits — absurd, would overflow a shift
|
buf[fahd + 7] = 200; // max_nelmts_bits — absurd, would overflow a shift
|
||||||
buf[fahd + 8..fahd + 16].copy_from_slice(&3u64.to_le_bytes()); // num_elements
|
buf[fahd + 8..fahd + 16].copy_from_slice(&3u64.to_le_bytes()); // num_elements
|
||||||
buf[fahd + 16..fahd + 24].copy_from_slice(&0x100u64.to_le_bytes());
|
buf[fahd + 16..fahd + 24].copy_from_slice(&0x100u64.to_le_bytes());
|
||||||
|
stamp_checksum(&mut buf, fahd, fahd + 24);
|
||||||
// FADB so parsing reaches the paged check
|
// FADB so parsing reaches the paged check
|
||||||
let db = 0x100usize;
|
let db = 0x100usize;
|
||||||
buf[db..db + 4].copy_from_slice(b"FADB");
|
buf[db..db + 4].copy_from_slice(b"FADB");
|
||||||
let header = FixedArrayHeader::parse(&buf, fahd, 8, 8).unwrap();
|
let header = FixedArrayHeader::parse(&buf, fahd, 8, 8).unwrap();
|
||||||
let r = read_fixed_array_chunks(&buf, &header, &[100], &[20], 8, 8, 8);
|
let r = read_fixed_array_chunks(&buf, &header, &[100], None, &[20], 8, 8, 8);
|
||||||
assert!(r.is_err());
|
assert!(r.is_err());
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -486,9 +527,11 @@ mod tests {
|
|||||||
buf[fahd + 7] = 10;
|
buf[fahd + 7] = 10;
|
||||||
buf[fahd + 8..fahd + 16].copy_from_slice(&u64::MAX.to_le_bytes()); // absurd count
|
buf[fahd + 8..fahd + 16].copy_from_slice(&u64::MAX.to_le_bytes()); // absurd count
|
||||||
buf[fahd + 16..fahd + 24].copy_from_slice(&0x80u64.to_le_bytes());
|
buf[fahd + 16..fahd + 24].copy_from_slice(&0x80u64.to_le_bytes());
|
||||||
|
// Valid checksum, so it is the element count that must be rejected.
|
||||||
|
stamp_checksum(&mut buf, fahd, fahd + 24);
|
||||||
buf[0x80..0x84].copy_from_slice(b"FADB");
|
buf[0x80..0x84].copy_from_slice(b"FADB");
|
||||||
let header = FixedArrayHeader::parse(&buf, fahd, 8, 8).unwrap();
|
let header = FixedArrayHeader::parse(&buf, fahd, 8, 8).unwrap();
|
||||||
let r = read_fixed_array_chunks(&buf, &header, &[100], &[20], 8, 8, 8);
|
let r = read_fixed_array_chunks(&buf, &header, &[100], None, &[20], 8, 8, 8);
|
||||||
assert!(r.is_err());
|
assert!(r.is_err());
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -511,7 +554,7 @@ mod tests {
|
|||||||
data_block_address: (usize::MAX - 4) as u64,
|
data_block_address: (usize::MAX - 4) as u64,
|
||||||
};
|
};
|
||||||
let buf = vec![0u8; 64];
|
let buf = vec![0u8; 64];
|
||||||
let r = read_fixed_array_chunks(&buf, &header, &[100], &[20], 8, 8, 8);
|
let r = read_fixed_array_chunks(&buf, &header, &[100], None, &[20], 8, 8, 8);
|
||||||
assert!(r.is_err());
|
assert!(r.is_err());
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -545,6 +588,7 @@ mod tests {
|
|||||||
file_data[fahd_offset + 8..fahd_offset + 16].copy_from_slice(&num_chunks.to_le_bytes());
|
file_data[fahd_offset + 8..fahd_offset + 16].copy_from_slice(&num_chunks.to_le_bytes());
|
||||||
file_data[fahd_offset + 16..fahd_offset + 24]
|
file_data[fahd_offset + 16..fahd_offset + 24]
|
||||||
.copy_from_slice(&(db_offset as u64).to_le_bytes());
|
.copy_from_slice(&(db_offset as u64).to_le_bytes());
|
||||||
|
stamp_checksum(&mut file_data, fahd_offset, fahd_offset + 24);
|
||||||
|
|
||||||
// Build FADB at db_offset
|
// Build FADB at db_offset
|
||||||
file_data[db_offset..db_offset + 4].copy_from_slice(b"FADB");
|
file_data[db_offset..db_offset + 4].copy_from_slice(b"FADB");
|
||||||
@@ -562,6 +606,7 @@ mod tests {
|
|||||||
let pos = elem_start + i * os;
|
let pos = elem_start + i * os;
|
||||||
file_data[pos..pos + os].copy_from_slice(&addr.to_le_bytes());
|
file_data[pos..pos + os].copy_from_slice(&addr.to_le_bytes());
|
||||||
}
|
}
|
||||||
|
stamp_checksum(&mut file_data, db_offset, elem_start + 5 * os);
|
||||||
|
|
||||||
let header =
|
let header =
|
||||||
FixedArrayHeader::parse(&file_data, fahd_offset, offset_size, length_size).unwrap();
|
FixedArrayHeader::parse(&file_data, fahd_offset, offset_size, length_size).unwrap();
|
||||||
@@ -571,6 +616,7 @@ mod tests {
|
|||||||
&file_data,
|
&file_data,
|
||||||
&header,
|
&header,
|
||||||
&ds_dims,
|
&ds_dims,
|
||||||
|
None,
|
||||||
&chunk_dims,
|
&chunk_dims,
|
||||||
8,
|
8,
|
||||||
offset_size,
|
offset_size,
|
||||||
@@ -611,6 +657,7 @@ mod tests {
|
|||||||
file_data[fahd_offset + 8..fahd_offset + 16].copy_from_slice(&num_chunks.to_le_bytes());
|
file_data[fahd_offset + 8..fahd_offset + 16].copy_from_slice(&num_chunks.to_le_bytes());
|
||||||
file_data[fahd_offset + 16..fahd_offset + 24]
|
file_data[fahd_offset + 16..fahd_offset + 24]
|
||||||
.copy_from_slice(&(db_offset as u64).to_le_bytes());
|
.copy_from_slice(&(db_offset as u64).to_le_bytes());
|
||||||
|
stamp_checksum(&mut file_data, fahd_offset, fahd_offset + 24);
|
||||||
|
|
||||||
file_data[db_offset..db_offset + 4].copy_from_slice(b"FADB");
|
file_data[db_offset..db_offset + 4].copy_from_slice(b"FADB");
|
||||||
file_data[db_offset + 4] = 0;
|
file_data[db_offset + 4] = 0;
|
||||||
@@ -632,6 +679,11 @@ mod tests {
|
|||||||
file_data[pos + os..pos + os + 4].copy_from_slice(&csize.to_le_bytes());
|
file_data[pos + os..pos + os + 4].copy_from_slice(&csize.to_le_bytes());
|
||||||
file_data[pos + os + 4..pos + os + 8].copy_from_slice(&fmask.to_le_bytes());
|
file_data[pos + os + 4..pos + os + 8].copy_from_slice(&fmask.to_le_bytes());
|
||||||
}
|
}
|
||||||
|
stamp_checksum(
|
||||||
|
&mut file_data,
|
||||||
|
db_offset,
|
||||||
|
elem_start + test_chunks.len() * elem_size,
|
||||||
|
);
|
||||||
|
|
||||||
let header =
|
let header =
|
||||||
FixedArrayHeader::parse(&file_data, fahd_offset, offset_size, length_size).unwrap();
|
FixedArrayHeader::parse(&file_data, fahd_offset, offset_size, length_size).unwrap();
|
||||||
@@ -641,6 +693,7 @@ mod tests {
|
|||||||
&file_data,
|
&file_data,
|
||||||
&header,
|
&header,
|
||||||
&ds_dims,
|
&ds_dims,
|
||||||
|
None,
|
||||||
&chunk_dims,
|
&chunk_dims,
|
||||||
8,
|
8,
|
||||||
offset_size,
|
offset_size,
|
||||||
@@ -696,6 +749,7 @@ mod tests {
|
|||||||
file_data[fahd_offset + 8..fahd_offset + 16].copy_from_slice(&num_elements.to_le_bytes());
|
file_data[fahd_offset + 8..fahd_offset + 16].copy_from_slice(&num_elements.to_le_bytes());
|
||||||
file_data[fahd_offset + 16..fahd_offset + 24]
|
file_data[fahd_offset + 16..fahd_offset + 24]
|
||||||
.copy_from_slice(&(db_offset as u64).to_le_bytes());
|
.copy_from_slice(&(db_offset as u64).to_le_bytes());
|
||||||
|
stamp_checksum(&mut file_data, fahd_offset, fahd_offset + 24);
|
||||||
|
|
||||||
// FADB prefix
|
// FADB prefix
|
||||||
file_data[db_offset..db_offset + 4].copy_from_slice(b"FADB");
|
file_data[db_offset..db_offset + 4].copy_from_slice(b"FADB");
|
||||||
@@ -715,6 +769,9 @@ mod tests {
|
|||||||
let base_addr = 0x1000u64;
|
let base_addr = 0x1000u64;
|
||||||
// Page 0 (elements 0..4) and page 2 (elements 8..11) carry addresses;
|
// Page 0 (elements 0..4) and page 2 (elements 8..11) carry addresses;
|
||||||
// page 1's slot is left zero-filled and must be skipped.
|
// page 1's slot is left zero-filled and must be skipped.
|
||||||
|
// The prefix and bitmap carry one checksum, each initialised page
|
||||||
|
// another — as a real file does.
|
||||||
|
stamp_checksum(&mut file_data, db_offset, bitmap_off + bitmap_size);
|
||||||
for &p in &[0usize, 2usize] {
|
for &p in &[0usize, 2usize] {
|
||||||
let page_off = pages_start + p * page_total;
|
let page_off = pages_start + p * page_total;
|
||||||
let count = core::cmp::min(page_nelmts, num_elements as usize - p * page_nelmts);
|
let count = core::cmp::min(page_nelmts, num_elements as usize - p * page_nelmts);
|
||||||
@@ -724,6 +781,7 @@ mod tests {
|
|||||||
let pos = page_off + e * os;
|
let pos = page_off + e * os;
|
||||||
file_data[pos..pos + os].copy_from_slice(&addr.to_le_bytes());
|
file_data[pos..pos + os].copy_from_slice(&addr.to_le_bytes());
|
||||||
}
|
}
|
||||||
|
stamp_checksum(&mut file_data, page_off, page_off + count * os);
|
||||||
}
|
}
|
||||||
|
|
||||||
let header =
|
let header =
|
||||||
@@ -736,6 +794,7 @@ mod tests {
|
|||||||
&file_data,
|
&file_data,
|
||||||
&header,
|
&header,
|
||||||
&ds_dims,
|
&ds_dims,
|
||||||
|
None,
|
||||||
&chunk_dims,
|
&chunk_dims,
|
||||||
8,
|
8,
|
||||||
offset_size,
|
offset_size,
|
||||||
|
|||||||
@@ -0,0 +1,155 @@
|
|||||||
|
//! IEEE-754 half precision (binary16) conversions.
|
||||||
|
//!
|
||||||
|
//! Pure integer bit manipulation, so it works under `no_std` and needs no
|
||||||
|
//! `libm`. The writer ([`crate::type_builders::DatasetBuilder::with_f16_data`]),
|
||||||
|
//! the reader and `clawhdf5-agent`'s half-precision embedding store all use
|
||||||
|
//! these two functions, so a value rounded in memory is bit-for-bit the value
|
||||||
|
//! that reads back from the file.
|
||||||
|
|
||||||
|
/// Largest finite half-precision value. Anything larger in magnitude rounds
|
||||||
|
/// to infinity.
|
||||||
|
pub const F16_MAX: f32 = 65504.0;
|
||||||
|
|
||||||
|
/// Convert an `f32` to the bit pattern of the nearest half-precision value,
|
||||||
|
/// rounding ties to even (the IEEE default, and what numpy and the `half`
|
||||||
|
/// crate do).
|
||||||
|
///
|
||||||
|
/// Values beyond ±[`F16_MAX`] become ±infinity, values too small for a
|
||||||
|
/// subnormal become signed zero, and NaN stays NaN (quiet, payload
|
||||||
|
/// truncated).
|
||||||
|
pub fn f32_to_f16_bits(value: f32) -> u16 {
|
||||||
|
let x = value.to_bits();
|
||||||
|
let sign = (x >> 16) & 0x8000;
|
||||||
|
let exp = x & 0x7F80_0000;
|
||||||
|
let man = x & 0x007F_FFFF;
|
||||||
|
|
||||||
|
// Infinity and NaN.
|
||||||
|
if exp == 0x7F80_0000 {
|
||||||
|
let quiet_nan = if man == 0 { 0 } else { 0x0200 };
|
||||||
|
return (sign | 0x7C00 | quiet_nan | (man >> 13)) as u16;
|
||||||
|
}
|
||||||
|
|
||||||
|
let half_exp = ((exp >> 23) as i32) - 127 + 15;
|
||||||
|
|
||||||
|
// Too large: infinity.
|
||||||
|
if half_exp >= 0x1F {
|
||||||
|
return (sign | 0x7C00) as u16;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Subnormal half, or zero.
|
||||||
|
if half_exp <= 0 {
|
||||||
|
if 14 - half_exp > 24 {
|
||||||
|
return sign as u16;
|
||||||
|
}
|
||||||
|
let man = man | 0x0080_0000; // implicit leading bit
|
||||||
|
let shift = (14 - half_exp) as u32;
|
||||||
|
let mut half_man = man >> shift;
|
||||||
|
let round_bit = 1u32 << (shift - 1);
|
||||||
|
// Round half to even: up if above half, or exactly half and odd.
|
||||||
|
if (man & round_bit) != 0 && (man & (3 * round_bit - 1)) != 0 {
|
||||||
|
half_man += 1;
|
||||||
|
}
|
||||||
|
return (sign | half_man) as u16;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Normal half. A mantissa carry correctly rolls into the exponent (and
|
||||||
|
// from the largest finite value into infinity).
|
||||||
|
let half = sign | ((half_exp as u32) << 10) | (man >> 13);
|
||||||
|
let round_bit = 0x0000_1000;
|
||||||
|
if (man & round_bit) != 0 && (man & (3 * round_bit - 1)) != 0 {
|
||||||
|
(half + 1) as u16
|
||||||
|
} else {
|
||||||
|
half as u16
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Convert the bit pattern of a half-precision value to `f32` (exact: every
|
||||||
|
/// half value is representable as an `f32`).
|
||||||
|
pub fn f16_bits_to_f32(h: u16) -> f32 {
|
||||||
|
let h = h as u32;
|
||||||
|
let sign = (h & 0x8000) << 16;
|
||||||
|
let exp = (h >> 10) & 0x1f;
|
||||||
|
let mant = h & 0x3ff;
|
||||||
|
let bits = if exp == 0 {
|
||||||
|
if mant == 0 {
|
||||||
|
sign // signed zero
|
||||||
|
} else {
|
||||||
|
// Subnormal: normalize into an f32 normal.
|
||||||
|
let mut e: i32 = -1;
|
||||||
|
let mut m = mant;
|
||||||
|
loop {
|
||||||
|
e += 1;
|
||||||
|
m <<= 1;
|
||||||
|
if m & 0x400 != 0 {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
let m = m & 0x3ff;
|
||||||
|
sign | (((127 - 15 - e) as u32) << 23) | (m << 13)
|
||||||
|
}
|
||||||
|
} else if exp == 0x1f {
|
||||||
|
sign | 0x7f80_0000 | (mant << 13) // inf / NaN
|
||||||
|
} else {
|
||||||
|
sign | ((exp + 127 - 15) << 23) | (mant << 13)
|
||||||
|
};
|
||||||
|
f32::from_bits(bits)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Round an `f32` to the nearest half-precision value, returned as `f32`.
|
||||||
|
pub fn round_to_f16(value: f32) -> f32 {
|
||||||
|
f16_bits_to_f32(f32_to_f16_bits(value))
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn every_half_value_round_trips() {
|
||||||
|
for bits in 0..=u16::MAX {
|
||||||
|
let v = f16_bits_to_f32(bits);
|
||||||
|
if v.is_nan() {
|
||||||
|
assert!(f16_bits_to_f32(f32_to_f16_bits(v)).is_nan(), "{bits:#06x}");
|
||||||
|
} else {
|
||||||
|
assert_eq!(f32_to_f16_bits(v), bits, "{bits:#06x} -> {v}");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn matches_the_half_crate() {
|
||||||
|
// Every 257th f32 bit pattern (~16.7M values) covers every exponent,
|
||||||
|
// the subnormal range, both signs, ties and the overflow boundary.
|
||||||
|
let mut bits: u32 = 0;
|
||||||
|
loop {
|
||||||
|
let v = f32::from_bits(bits);
|
||||||
|
let ours = f32_to_f16_bits(v);
|
||||||
|
let theirs = half::f16::from_f32(v);
|
||||||
|
if v.is_nan() {
|
||||||
|
assert!(theirs.is_nan() && f16_bits_to_f32(ours).is_nan());
|
||||||
|
} else {
|
||||||
|
assert_eq!(ours, theirs.to_bits(), "{bits:#010x} ({v:e})");
|
||||||
|
assert_eq!(f16_bits_to_f32(ours).to_bits(), theirs.to_f32().to_bits());
|
||||||
|
}
|
||||||
|
match bits.checked_add(257) {
|
||||||
|
Some(b) => bits = b,
|
||||||
|
None => break,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn rounds_ties_to_even_and_saturates_to_infinity() {
|
||||||
|
// 1 + 2^-11 is exactly halfway between 1.0 and the next half (1 + 2^-10).
|
||||||
|
assert_eq!(round_to_f16(1.0 + 2f32.powi(-11)), 1.0);
|
||||||
|
assert_eq!(
|
||||||
|
round_to_f16(1.0 + 3.0 * 2f32.powi(-11)),
|
||||||
|
1.0 + 2.0 * 2f32.powi(-10)
|
||||||
|
);
|
||||||
|
assert_eq!(round_to_f16(F16_MAX), F16_MAX);
|
||||||
|
assert_eq!(round_to_f16(65520.0), f32::INFINITY); // halfway to 2^16 rounds up
|
||||||
|
assert_eq!(round_to_f16(-1e9), f32::NEG_INFINITY);
|
||||||
|
assert_eq!(round_to_f16(1e-9).to_bits(), 0);
|
||||||
|
assert_eq!(round_to_f16(-1e-9).to_bits(), (-0.0f32).to_bits());
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -54,6 +54,7 @@ pub mod btree_v1;
|
|||||||
pub mod btree_v2;
|
pub mod btree_v2;
|
||||||
pub mod checksum;
|
pub mod checksum;
|
||||||
pub mod chunk_cache;
|
pub mod chunk_cache;
|
||||||
|
mod chunk_grid;
|
||||||
pub mod chunk_index;
|
pub mod chunk_index;
|
||||||
pub mod chunked_read;
|
pub mod chunked_read;
|
||||||
pub mod chunked_write;
|
pub mod chunked_write;
|
||||||
@@ -72,6 +73,7 @@ pub mod filter_pipeline;
|
|||||||
pub mod filters;
|
pub mod filters;
|
||||||
mod filters_szip;
|
mod filters_szip;
|
||||||
pub mod fixed_array;
|
pub mod fixed_array;
|
||||||
|
pub mod float16;
|
||||||
pub mod fractal_heap;
|
pub mod fractal_heap;
|
||||||
pub mod global_heap;
|
pub mod global_heap;
|
||||||
pub mod group_info;
|
pub mod group_info;
|
||||||
@@ -89,6 +91,7 @@ pub mod object_header;
|
|||||||
pub mod object_header_writer;
|
pub mod object_header_writer;
|
||||||
#[cfg(feature = "parallel")]
|
#[cfg(feature = "parallel")]
|
||||||
pub mod parallel_read;
|
pub mod parallel_read;
|
||||||
|
pub mod partial_read;
|
||||||
pub mod profiling;
|
pub mod profiling;
|
||||||
pub mod property_list;
|
pub mod property_list;
|
||||||
pub mod selection;
|
pub mod selection;
|
||||||
|
|||||||
@@ -146,12 +146,7 @@ impl ObjectHeader {
|
|||||||
ensure_len(data, pos, msg_data_size)?;
|
ensure_len(data, pos, msg_data_size)?;
|
||||||
let msg_type = MessageType::from_u16(msg_type_raw);
|
let msg_type = MessageType::from_u16(msg_type_raw);
|
||||||
|
|
||||||
// Check if unknown + must-understand (bit 3 of msg_flags)
|
check_unknown_message(msg_type, msg_flags)?;
|
||||||
if let MessageType::Unknown(id) = msg_type
|
|
||||||
&& msg_flags & 0x08 != 0
|
|
||||||
{
|
|
||||||
return Err(FormatError::UnsupportedMessage(id));
|
|
||||||
}
|
|
||||||
|
|
||||||
if msg_type != MessageType::Nil {
|
if msg_type != MessageType::Nil {
|
||||||
messages.push(HeaderMessage {
|
messages.push(HeaderMessage {
|
||||||
@@ -229,11 +224,7 @@ impl ObjectHeader {
|
|||||||
|
|
||||||
let msg_type = MessageType::from_u16(msg_type_raw);
|
let msg_type = MessageType::from_u16(msg_type_raw);
|
||||||
|
|
||||||
if let MessageType::Unknown(id) = msg_type
|
check_unknown_message(msg_type, msg_flags)?;
|
||||||
&& msg_flags & 0x08 != 0
|
|
||||||
{
|
|
||||||
return Err(FormatError::UnsupportedMessage(id));
|
|
||||||
}
|
|
||||||
|
|
||||||
if msg_type != MessageType::Nil {
|
if msg_type != MessageType::Nil {
|
||||||
messages.push(HeaderMessage {
|
messages.push(HeaderMessage {
|
||||||
@@ -424,11 +415,7 @@ impl ObjectHeader {
|
|||||||
|
|
||||||
let msg_type = MessageType::from_u16(msg_type_raw);
|
let msg_type = MessageType::from_u16(msg_type_raw);
|
||||||
|
|
||||||
if let MessageType::Unknown(id) = msg_type
|
check_unknown_message(msg_type, msg_flags)?;
|
||||||
&& msg_flags & 0x08 != 0
|
|
||||||
{
|
|
||||||
return Err(FormatError::UnsupportedMessage(id));
|
|
||||||
}
|
|
||||||
|
|
||||||
let msg_data = data[pos..pos + msg_data_size].to_vec();
|
let msg_data = data[pos..pos + msg_data_size].to_vec();
|
||||||
|
|
||||||
@@ -509,6 +496,24 @@ impl ObjectHeader {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Header message flag bit 7: fail if the message is unknown, always.
|
||||||
|
const MSG_FLAG_FAIL_IF_UNKNOWN_ALWAYS: u8 = 0x80;
|
||||||
|
|
||||||
|
/// Refuse an unknown message the file says no reader may skip.
|
||||||
|
///
|
||||||
|
/// The parser only ever reads, so bit 3 (fail only when opened for writing)
|
||||||
|
/// is ignored, as libhdf5 ignores it for a read-only open; bit 7 fails
|
||||||
|
/// regardless of access mode. This had the two the wrong way round, failing
|
||||||
|
/// objects libhdf5 reads and reading ones it refuses (`tbogus.h5`).
|
||||||
|
fn check_unknown_message(msg_type: MessageType, msg_flags: u8) -> Result<(), FormatError> {
|
||||||
|
match msg_type {
|
||||||
|
MessageType::Unknown(id) if msg_flags & MSG_FLAG_FAIL_IF_UNKNOWN_ALWAYS != 0 => {
|
||||||
|
Err(FormatError::UnsupportedMessage(id))
|
||||||
|
}
|
||||||
|
_ => Ok(()),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use super::*;
|
use super::*;
|
||||||
@@ -632,14 +637,38 @@ mod tests {
|
|||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn parse_v1_unknown_must_understand_errors() {
|
fn parse_v1_unknown_fail_always_errors() {
|
||||||
// Bit 3 of msg_flags = must understand
|
// Bit 7 of msg_flags = fail if unknown, whatever the access mode.
|
||||||
let messages = [(0x00FFu16, &[0xAA][..], 0x08u8)];
|
let messages = [(0x00FFu16, &[0xAA][..], 0x80u8)];
|
||||||
let data = build_v1_header(&messages, 8, 8);
|
let data = build_v1_header(&messages, 8, 8);
|
||||||
let err = ObjectHeader::parse(&data, 0, 8, 8).unwrap_err();
|
let err = ObjectHeader::parse(&data, 0, 8, 8).unwrap_err();
|
||||||
assert_eq!(err, FormatError::UnsupportedMessage(0x00FF));
|
assert_eq!(err, FormatError::UnsupportedMessage(0x00FF));
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn parse_v1_unknown_fail_on_write_is_ignored_when_reading() {
|
||||||
|
// Bit 3 = fail if unknown *and the file is opened for writing*. This
|
||||||
|
// parser only reads, so libhdf5 (read-only) opens such an object and
|
||||||
|
// so must we. Bits 4/5 (mark if unknown / was unknown) never fail.
|
||||||
|
for flags in [0x08u8, 0x10, 0x20, 0x38] {
|
||||||
|
let messages = [(0x00FFu16, &[0xAA][..], flags)];
|
||||||
|
let data = build_v1_header(&messages, 8, 8);
|
||||||
|
let hdr = ObjectHeader::parse(&data, 0, 8, 8).unwrap();
|
||||||
|
assert_eq!(hdr.messages[0].msg_type, MessageType::Unknown(0x00FF));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn parse_v2_unknown_message_flags() {
|
||||||
|
let data = build_v2_header(0x00, &[(0xF0, &[1, 2], 0x08)], None);
|
||||||
|
assert!(ObjectHeader::parse(&data, 0, 8, 8).is_ok());
|
||||||
|
let data = build_v2_header(0x00, &[(0xF0, &[1, 2], 0x80)], None);
|
||||||
|
assert_eq!(
|
||||||
|
ObjectHeader::parse(&data, 0, 8, 8).unwrap_err(),
|
||||||
|
FormatError::UnsupportedMessage(0xF0)
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn parse_v2_no_timestamps_one_message() {
|
fn parse_v2_no_timestamps_one_message() {
|
||||||
let data = build_v2_header(0x00, &[(0x01, &[10, 20], 0)], None);
|
let data = build_v2_header(0x00, &[(0x01, &[10, 20], 0)], None);
|
||||||
|
|||||||
@@ -1,11 +1,17 @@
|
|||||||
//! Object header writer for v2 format.
|
//! Object header writer for v2 format.
|
||||||
|
|
||||||
#[cfg(not(feature = "std"))]
|
#[cfg(not(feature = "std"))]
|
||||||
use alloc::vec::Vec;
|
use alloc::{format, vec::Vec};
|
||||||
|
|
||||||
use crate::checksum::jenkins_lookup3;
|
use crate::checksum::jenkins_lookup3;
|
||||||
|
use crate::error::FormatError;
|
||||||
use crate::message_type::MessageType;
|
use crate::message_type::MessageType;
|
||||||
|
|
||||||
|
/// Largest message payload a v2 object header can describe: the per-message
|
||||||
|
/// size field is 2 bytes. A bigger message cannot be encoded at all — writing
|
||||||
|
/// its size truncated to 16 bits produced files libhdf5 refuses.
|
||||||
|
pub const MAX_MESSAGE_SIZE: usize = u16::MAX as usize;
|
||||||
|
|
||||||
/// Writer for v2 object headers with proper checksums.
|
/// Writer for v2 object headers with proper checksums.
|
||||||
pub struct ObjectHeaderWriter {
|
pub struct ObjectHeaderWriter {
|
||||||
messages: Vec<(MessageType, Vec<u8>, u8)>, // (type, data, msg_flags)
|
messages: Vec<(MessageType, Vec<u8>, u8)>, // (type, data, msg_flags)
|
||||||
@@ -30,7 +36,22 @@ impl ObjectHeaderWriter {
|
|||||||
}
|
}
|
||||||
|
|
||||||
/// Serialize the complete v2 object header (OHDR + messages + checksum).
|
/// Serialize the complete v2 object header (OHDR + messages + checksum).
|
||||||
pub fn serialize(&self) -> Vec<u8> {
|
///
|
||||||
|
/// Fails with [`FormatError::SerializationError`] when a message is larger
|
||||||
|
/// than [`MAX_MESSAGE_SIZE`] (e.g. an attribute over ~64 KiB, which would
|
||||||
|
/// need dense attribute storage), rather than writing a corrupt header.
|
||||||
|
pub fn serialize(&self) -> Result<Vec<u8>, FormatError> {
|
||||||
|
if let Some((msg_type, data, _)) = self
|
||||||
|
.messages
|
||||||
|
.iter()
|
||||||
|
.find(|(_, data, _)| data.len() > MAX_MESSAGE_SIZE)
|
||||||
|
{
|
||||||
|
return Err(FormatError::SerializationError(format!(
|
||||||
|
"{msg_type:?} message is {} bytes; an object header message holds at most \
|
||||||
|
{MAX_MESSAGE_SIZE} bytes",
|
||||||
|
data.len()
|
||||||
|
)));
|
||||||
|
}
|
||||||
// Calculate total message bytes: each message has type(1) + size(2) + flags(1) + data
|
// Calculate total message bytes: each message has type(1) + size(2) + flags(1) + data
|
||||||
let msg_bytes_total: usize = self
|
let msg_bytes_total: usize = self
|
||||||
.messages
|
.messages
|
||||||
@@ -80,7 +101,7 @@ impl ObjectHeaderWriter {
|
|||||||
let checksum = jenkins_lookup3(&buf);
|
let checksum = jenkins_lookup3(&buf);
|
||||||
buf.extend_from_slice(&checksum.to_le_bytes());
|
buf.extend_from_slice(&checksum.to_le_bytes());
|
||||||
|
|
||||||
buf
|
Ok(buf)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -125,15 +146,22 @@ impl BatchObjectHeaderWriter {
|
|||||||
|
|
||||||
/// Compute the serialized size of each header without actually serializing.
|
/// Compute the serialized size of each header without actually serializing.
|
||||||
/// Returns sizes in the same order as headers were added.
|
/// Returns sizes in the same order as headers were added.
|
||||||
pub fn compute_sizes(&self) -> Vec<usize> {
|
pub fn compute_sizes(&self) -> Result<Vec<usize>, FormatError> {
|
||||||
self.headers.iter().map(|h| h.serialize().len()).collect()
|
self.headers
|
||||||
|
.iter()
|
||||||
|
.map(|h| h.serialize().map(|b| b.len()))
|
||||||
|
.collect()
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Serialize all headers into a single contiguous buffer.
|
/// Serialize all headers into a single contiguous buffer.
|
||||||
/// Returns `(combined_bytes, offsets)` where `offsets[i]` is the byte
|
/// Returns `(combined_bytes, offsets)` where `offsets[i]` is the byte
|
||||||
/// offset of header `i` within the combined buffer.
|
/// offset of header `i` within the combined buffer.
|
||||||
pub fn serialize_all(&self) -> (Vec<u8>, Vec<usize>) {
|
pub fn serialize_all(&self) -> Result<(Vec<u8>, Vec<usize>), FormatError> {
|
||||||
let serialized: Vec<Vec<u8>> = self.headers.iter().map(|h| h.serialize()).collect();
|
let serialized: Vec<Vec<u8>> = self
|
||||||
|
.headers
|
||||||
|
.iter()
|
||||||
|
.map(|h| h.serialize())
|
||||||
|
.collect::<Result<_, _>>()?;
|
||||||
let total: usize = serialized.iter().map(|s| s.len()).sum();
|
let total: usize = serialized.iter().map(|s| s.len()).sum();
|
||||||
let mut buf = Vec::with_capacity(total);
|
let mut buf = Vec::with_capacity(total);
|
||||||
let mut offsets = Vec::with_capacity(serialized.len());
|
let mut offsets = Vec::with_capacity(serialized.len());
|
||||||
@@ -141,7 +169,7 @@ impl BatchObjectHeaderWriter {
|
|||||||
offsets.push(buf.len());
|
offsets.push(buf.len());
|
||||||
buf.extend_from_slice(s);
|
buf.extend_from_slice(s);
|
||||||
}
|
}
|
||||||
(buf, offsets)
|
Ok((buf, offsets))
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -159,7 +187,7 @@ mod tests {
|
|||||||
#[test]
|
#[test]
|
||||||
fn empty_header_roundtrip() {
|
fn empty_header_roundtrip() {
|
||||||
let writer = ObjectHeaderWriter::new();
|
let writer = ObjectHeaderWriter::new();
|
||||||
let bytes = writer.serialize();
|
let bytes = writer.serialize().unwrap();
|
||||||
let hdr = ObjectHeader::parse(&bytes, 0, 8, 8).unwrap();
|
let hdr = ObjectHeader::parse(&bytes, 0, 8, 8).unwrap();
|
||||||
assert_eq!(hdr.version, 2);
|
assert_eq!(hdr.version, 2);
|
||||||
assert_eq!(hdr.messages.len(), 0);
|
assert_eq!(hdr.messages.len(), 0);
|
||||||
@@ -170,7 +198,7 @@ mod tests {
|
|||||||
let mut writer = ObjectHeaderWriter::new();
|
let mut writer = ObjectHeaderWriter::new();
|
||||||
writer.add_message(MessageType::Dataspace, vec![1, 2, 3, 4]);
|
writer.add_message(MessageType::Dataspace, vec![1, 2, 3, 4]);
|
||||||
writer.add_message(MessageType::Datatype, vec![5, 6]);
|
writer.add_message(MessageType::Datatype, vec![5, 6]);
|
||||||
let bytes = writer.serialize();
|
let bytes = writer.serialize().unwrap();
|
||||||
let hdr = ObjectHeader::parse(&bytes, 0, 8, 8).unwrap();
|
let hdr = ObjectHeader::parse(&bytes, 0, 8, 8).unwrap();
|
||||||
assert_eq!(hdr.messages.len(), 2);
|
assert_eq!(hdr.messages.len(), 2);
|
||||||
assert_eq!(hdr.messages[0].msg_type, MessageType::Dataspace);
|
assert_eq!(hdr.messages[0].msg_type, MessageType::Dataspace);
|
||||||
@@ -184,12 +212,30 @@ mod tests {
|
|||||||
let mut writer = ObjectHeaderWriter::new();
|
let mut writer = ObjectHeaderWriter::new();
|
||||||
// Add a message with >255 bytes of payload
|
// Add a message with >255 bytes of payload
|
||||||
writer.add_message(MessageType::Datatype, vec![0xAA; 300]);
|
writer.add_message(MessageType::Datatype, vec![0xAA; 300]);
|
||||||
let bytes = writer.serialize();
|
let bytes = writer.serialize().unwrap();
|
||||||
let hdr = ObjectHeader::parse(&bytes, 0, 8, 8).unwrap();
|
let hdr = ObjectHeader::parse(&bytes, 0, 8, 8).unwrap();
|
||||||
assert_eq!(hdr.messages.len(), 1);
|
assert_eq!(hdr.messages.len(), 1);
|
||||||
assert_eq!(hdr.messages[0].data.len(), 300);
|
assert_eq!(hdr.messages[0].data.len(), 300);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn oversized_message_is_an_error_not_a_truncated_size() {
|
||||||
|
// 65535 bytes is the largest encodable payload.
|
||||||
|
let mut writer = ObjectHeaderWriter::new();
|
||||||
|
writer.add_message(MessageType::Attribute, vec![0; MAX_MESSAGE_SIZE]);
|
||||||
|
let bytes = writer.serialize().unwrap();
|
||||||
|
let hdr = ObjectHeader::parse(&bytes, 0, 8, 8).unwrap();
|
||||||
|
assert_eq!(hdr.messages[0].data.len(), MAX_MESSAGE_SIZE);
|
||||||
|
|
||||||
|
// One byte more used to be written with its size wrapped to 0.
|
||||||
|
let mut writer = ObjectHeaderWriter::new();
|
||||||
|
writer.add_message(MessageType::Attribute, vec![0; MAX_MESSAGE_SIZE + 1]);
|
||||||
|
assert!(matches!(
|
||||||
|
writer.serialize(),
|
||||||
|
Err(FormatError::SerializationError(_))
|
||||||
|
));
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn batch_writer_serialize_all() {
|
fn batch_writer_serialize_all() {
|
||||||
let mut batch = BatchObjectHeaderWriter::new();
|
let mut batch = BatchObjectHeaderWriter::new();
|
||||||
@@ -204,7 +250,7 @@ mod tests {
|
|||||||
batch.add(w2);
|
batch.add(w2);
|
||||||
assert_eq!(batch.len(), 2);
|
assert_eq!(batch.len(), 2);
|
||||||
|
|
||||||
let (buf, offsets) = batch.serialize_all();
|
let (buf, offsets) = batch.serialize_all().unwrap();
|
||||||
assert_eq!(offsets.len(), 2);
|
assert_eq!(offsets.len(), 2);
|
||||||
assert_eq!(offsets[0], 0);
|
assert_eq!(offsets[0], 0);
|
||||||
|
|
||||||
@@ -222,7 +268,7 @@ mod tests {
|
|||||||
fn batch_writer_empty() {
|
fn batch_writer_empty() {
|
||||||
let batch = BatchObjectHeaderWriter::new();
|
let batch = BatchObjectHeaderWriter::new();
|
||||||
assert!(batch.is_empty());
|
assert!(batch.is_empty());
|
||||||
let (buf, offsets) = batch.serialize_all();
|
let (buf, offsets) = batch.serialize_all().unwrap();
|
||||||
assert!(buf.is_empty());
|
assert!(buf.is_empty());
|
||||||
assert!(offsets.is_empty());
|
assert!(offsets.is_empty());
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -10,7 +10,7 @@
|
|||||||
use crate::chunked_read::ChunkInfo;
|
use crate::chunked_read::ChunkInfo;
|
||||||
use crate::error::FormatError;
|
use crate::error::FormatError;
|
||||||
use crate::filter_pipeline::FilterPipeline;
|
use crate::filter_pipeline::FilterPipeline;
|
||||||
use crate::filters::decompress_chunk;
|
use crate::filters::decompress_chunk_masked;
|
||||||
use crate::lane_partition::{self, LaneStats, PartitionStats};
|
use crate::lane_partition::{self, LaneStats, PartitionStats};
|
||||||
|
|
||||||
/// Threshold: only use parallel decompression when chunk count exceeds this.
|
/// Threshold: only use parallel decompression when chunk count exceeds this.
|
||||||
@@ -84,11 +84,13 @@ pub fn decompress_chunks_lane_partitioned(
|
|||||||
}
|
}
|
||||||
let raw_chunk = &file_data[c_addr..c_addr + size];
|
let raw_chunk = &file_data[c_addr..c_addr + size];
|
||||||
|
|
||||||
let decompressed = if chunk_info.filter_mask == 0 {
|
let decompressed = decompress_chunk_masked(
|
||||||
decompress_chunk(raw_chunk, pipeline, chunk_total_bytes, element_size)?
|
raw_chunk,
|
||||||
} else {
|
pipeline,
|
||||||
raw_chunk.to_vec()
|
chunk_total_bytes,
|
||||||
};
|
element_size,
|
||||||
|
chunk_info.filter_mask,
|
||||||
|
)?;
|
||||||
|
|
||||||
stats.chunks_processed += 1;
|
stats.chunks_processed += 1;
|
||||||
stats.compressed_bytes += size as u64;
|
stats.compressed_bytes += size as u64;
|
||||||
@@ -158,11 +160,13 @@ pub fn decompress_chunks_parallel(
|
|||||||
}
|
}
|
||||||
let raw_chunk = &file_data[c_addr..c_addr + size];
|
let raw_chunk = &file_data[c_addr..c_addr + size];
|
||||||
|
|
||||||
let decompressed = if chunk_info.filter_mask == 0 {
|
let decompressed = decompress_chunk_masked(
|
||||||
decompress_chunk(raw_chunk, pipeline, chunk_total_bytes, element_size)?
|
raw_chunk,
|
||||||
} else {
|
pipeline,
|
||||||
raw_chunk.to_vec()
|
chunk_total_bytes,
|
||||||
};
|
element_size,
|
||||||
|
chunk_info.filter_mask,
|
||||||
|
)?;
|
||||||
|
|
||||||
Ok(DecompressedChunk {
|
Ok(DecompressedChunk {
|
||||||
index,
|
index,
|
||||||
@@ -200,11 +204,13 @@ pub fn decompress_chunks_sequential(
|
|||||||
let raw_chunk = &file_data[c_addr..c_addr + size];
|
let raw_chunk = &file_data[c_addr..c_addr + size];
|
||||||
|
|
||||||
let decompressed = if let Some(pl) = pipeline {
|
let decompressed = if let Some(pl) = pipeline {
|
||||||
if chunk_info.filter_mask == 0 {
|
decompress_chunk_masked(
|
||||||
decompress_chunk(raw_chunk, pl, chunk_total_bytes, element_size)?
|
raw_chunk,
|
||||||
} else {
|
pl,
|
||||||
raw_chunk.to_vec()
|
chunk_total_bytes,
|
||||||
}
|
element_size,
|
||||||
|
chunk_info.filter_mask,
|
||||||
|
)?
|
||||||
} else {
|
} else {
|
||||||
raw_chunk.to_vec()
|
raw_chunk.to_vec()
|
||||||
};
|
};
|
||||||
|
|||||||
@@ -0,0 +1,365 @@
|
|||||||
|
//! Selection reads that cost what the selection costs, not what the dataset
|
||||||
|
//! costs.
|
||||||
|
//!
|
||||||
|
//! [`crate::data_read::read_raw_data_selection`] used to decode the *entire*
|
||||||
|
//! dataset and then pick elements out of it, so reading a 64x64 window of a
|
||||||
|
//! large dataset took about as long as reading all of it. Here the selection's
|
||||||
|
//! bounding box is materialised instead — only the rows of a contiguous
|
||||||
|
//! dataset, or only the chunks, that overlap it — and the existing extractor
|
||||||
|
//! runs over that small buffer with the selection translated to the box's
|
||||||
|
//! origin. Extraction semantics are therefore exactly the full-read ones.
|
||||||
|
|
||||||
|
#[cfg(not(feature = "std"))]
|
||||||
|
use alloc::string as alloc_or_std;
|
||||||
|
#[cfg(not(feature = "std"))]
|
||||||
|
use alloc::{format, vec, vec::Vec};
|
||||||
|
#[cfg(feature = "std")]
|
||||||
|
use std::string as alloc_or_std;
|
||||||
|
|
||||||
|
use crate::chunked_read::{alloc_output, checked_byte_len, list_chunks};
|
||||||
|
use crate::data_layout::DataLayout;
|
||||||
|
use crate::data_read::extract_selection_from_buffer;
|
||||||
|
use crate::dataspace::Dataspace;
|
||||||
|
use crate::error::FormatError;
|
||||||
|
use crate::filter_pipeline::FilterPipeline;
|
||||||
|
use crate::filters::{all_filters_skipped, decompress_chunk_masked};
|
||||||
|
use crate::selection::Selection;
|
||||||
|
|
||||||
|
/// The smallest axis-aligned box containing every selected element, as
|
||||||
|
/// `(start, extent)` per dimension. `None` when there is nothing to gain or
|
||||||
|
/// the selection is not valid for `dims` (the caller's full path then reports
|
||||||
|
/// the error exactly as before).
|
||||||
|
fn bounding_box(selection: &Selection, dims: &[u64]) -> Option<(Vec<u64>, Vec<u64>)> {
|
||||||
|
match selection {
|
||||||
|
Selection::Hyperslab {
|
||||||
|
start,
|
||||||
|
stride,
|
||||||
|
count,
|
||||||
|
block,
|
||||||
|
} => {
|
||||||
|
let rank = dims.len();
|
||||||
|
if [start.len(), stride.len(), count.len(), block.len()] != [rank; 4] {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
let mut extent = Vec::with_capacity(rank);
|
||||||
|
for d in 0..rank {
|
||||||
|
if count[d] == 0 || block[d] == 0 {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
// Last selected index + 1, relative to start.
|
||||||
|
let span = (count[d] - 1)
|
||||||
|
.checked_mul(stride[d])?
|
||||||
|
.checked_add(block[d])?;
|
||||||
|
if start[d].checked_add(span)? > dims[d] {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
extent.push(span);
|
||||||
|
}
|
||||||
|
Some((start.clone(), extent))
|
||||||
|
}
|
||||||
|
Selection::Points(points) => {
|
||||||
|
let rank = dims.len();
|
||||||
|
let first = points.first()?;
|
||||||
|
if first.len() != rank {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
let (mut lo, mut hi) = (first.clone(), first.clone());
|
||||||
|
for p in points {
|
||||||
|
if p.len() != rank {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
for d in 0..rank {
|
||||||
|
if p[d] >= dims[d] {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
lo[d] = lo[d].min(p[d]);
|
||||||
|
hi[d] = hi[d].max(p[d]);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
let extent = lo.iter().zip(&hi).map(|(l, h)| h - l + 1).collect();
|
||||||
|
Some((lo, extent))
|
||||||
|
}
|
||||||
|
Selection::All | Selection::None => None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Check that `selection` addresses only elements that exist in a dataset of
|
||||||
|
/// shape `dims`. Without this an out-of-range selection read *something*: a
|
||||||
|
/// hyperslab past the edge came back padded with zeros, and a point whose
|
||||||
|
/// column was out of range wrapped into the next row.
|
||||||
|
pub fn validate(selection: &Selection, dims: &[u64]) -> Result<(), FormatError> {
|
||||||
|
let rank = dims.len();
|
||||||
|
let bad = |msg: alloc_or_std::String| Err(FormatError::SelectionOutOfBounds(msg));
|
||||||
|
match selection {
|
||||||
|
Selection::All | Selection::None => Ok(()),
|
||||||
|
Selection::Hyperslab {
|
||||||
|
start,
|
||||||
|
stride,
|
||||||
|
count,
|
||||||
|
block,
|
||||||
|
} => {
|
||||||
|
if [start.len(), stride.len(), count.len(), block.len()] != [rank; 4] {
|
||||||
|
return bad(format!("hyperslab rank does not match dataset rank {rank}"));
|
||||||
|
}
|
||||||
|
for d in 0..rank {
|
||||||
|
if count[d] == 0 || block[d] == 0 {
|
||||||
|
continue; // selects nothing along this dimension
|
||||||
|
}
|
||||||
|
let end = (count[d] - 1)
|
||||||
|
.checked_mul(stride[d])
|
||||||
|
.and_then(|v| v.checked_add(block[d]))
|
||||||
|
.and_then(|v| v.checked_add(start[d]));
|
||||||
|
if !end.is_some_and(|end| end <= dims[d]) {
|
||||||
|
return bad(format!(
|
||||||
|
"dimension {d}: start {} stride {} count {} block {} exceeds extent {}",
|
||||||
|
start[d], stride[d], count[d], block[d], dims[d]
|
||||||
|
));
|
||||||
|
}
|
||||||
|
if block[d] > stride[d] && count[d] > 1 {
|
||||||
|
return bad(format!(
|
||||||
|
"dimension {d}: block {} larger than stride {} (overlapping blocks)",
|
||||||
|
block[d], stride[d]
|
||||||
|
));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
Selection::Points(points) => {
|
||||||
|
for p in points {
|
||||||
|
if p.len() != rank {
|
||||||
|
return bad(format!("point {p:?} does not match dataset rank {rank}"));
|
||||||
|
}
|
||||||
|
if let Some(d) = (0..rank).find(|&d| p[d] >= dims[d]) {
|
||||||
|
return bad(format!(
|
||||||
|
"point {p:?}: coordinate {} exceeds extent {} of dimension {d}",
|
||||||
|
p[d], dims[d]
|
||||||
|
));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// The same selection expressed relative to `origin`.
|
||||||
|
fn translate(selection: &Selection, origin: &[u64]) -> Selection {
|
||||||
|
match selection {
|
||||||
|
Selection::Hyperslab {
|
||||||
|
start,
|
||||||
|
stride,
|
||||||
|
count,
|
||||||
|
block,
|
||||||
|
} => Selection::Hyperslab {
|
||||||
|
start: start.iter().zip(origin).map(|(s, o)| s - o).collect(),
|
||||||
|
stride: stride.clone(),
|
||||||
|
count: count.clone(),
|
||||||
|
block: block.clone(),
|
||||||
|
},
|
||||||
|
Selection::Points(points) => Selection::Points(
|
||||||
|
points
|
||||||
|
.iter()
|
||||||
|
.map(|p| p.iter().zip(origin).map(|(c, o)| c - o).collect())
|
||||||
|
.collect(),
|
||||||
|
),
|
||||||
|
other => other.clone(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Copy the part of a source region that overlaps the box into `out` (which
|
||||||
|
/// is the box, row-major).
|
||||||
|
///
|
||||||
|
/// The source region starts at `src_origin` in dataset coordinates, has shape
|
||||||
|
/// `src_shape`, and its elements are in `src` row-major. One `memcpy` per
|
||||||
|
/// overlapping row of the last dimension.
|
||||||
|
#[allow(clippy::too_many_arguments)]
|
||||||
|
fn copy_overlap(
|
||||||
|
src: &[u8],
|
||||||
|
src_origin: &[u64],
|
||||||
|
src_shape: &[u64],
|
||||||
|
out: &mut [u8],
|
||||||
|
box_start: &[u64],
|
||||||
|
box_extent: &[u64],
|
||||||
|
elem_size: usize,
|
||||||
|
) {
|
||||||
|
let rank = box_start.len();
|
||||||
|
// Overlap in dataset coordinates.
|
||||||
|
let mut lo = vec![0u64; rank];
|
||||||
|
let mut hi = vec![0u64; rank];
|
||||||
|
for d in 0..rank {
|
||||||
|
lo[d] = src_origin[d].max(box_start[d]);
|
||||||
|
hi[d] = (src_origin[d] + src_shape[d]).min(box_start[d] + box_extent[d]);
|
||||||
|
if lo[d] >= hi[d] {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
let strides = |shape: &[u64]| {
|
||||||
|
let mut s = vec![1u64; rank];
|
||||||
|
for d in (0..rank.saturating_sub(1)).rev() {
|
||||||
|
s[d] = s[d + 1] * shape[d + 1];
|
||||||
|
}
|
||||||
|
s
|
||||||
|
};
|
||||||
|
let (src_strides, out_strides) = (strides(src_shape), strides(box_extent));
|
||||||
|
let last = rank - 1;
|
||||||
|
let run = ((hi[last] - lo[last]) as usize) * elem_size;
|
||||||
|
|
||||||
|
let mut idx = lo.clone();
|
||||||
|
loop {
|
||||||
|
let src_at: u64 = (0..rank)
|
||||||
|
.map(|d| (idx[d] - src_origin[d]) * src_strides[d])
|
||||||
|
.sum();
|
||||||
|
let out_at: u64 = (0..rank)
|
||||||
|
.map(|d| (idx[d] - box_start[d]) * out_strides[d])
|
||||||
|
.sum();
|
||||||
|
let (s, o) = (src_at as usize * elem_size, out_at as usize * elem_size);
|
||||||
|
if let (Some(from), Some(to)) = (src.get(s..s + run), out.get_mut(o..o + run)) {
|
||||||
|
to.copy_from_slice(from);
|
||||||
|
}
|
||||||
|
// Advance over every dimension but the last.
|
||||||
|
let mut d = last;
|
||||||
|
loop {
|
||||||
|
if d == 0 {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
d -= 1;
|
||||||
|
idx[d] += 1;
|
||||||
|
if idx[d] < hi[d] {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
idx[d] = lo[d];
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Read `selection` without materialising the whole dataset, when that is
|
||||||
|
/// possible and worthwhile. `Ok(None)` means "use the full-read path": an
|
||||||
|
/// `All`/`None`/invalid selection, a layout this doesn't handle (compact,
|
||||||
|
/// virtual, storage-less), or a bounding box covering most of the dataset.
|
||||||
|
#[allow(clippy::too_many_arguments)]
|
||||||
|
pub fn read_selection(
|
||||||
|
file_data: &[u8],
|
||||||
|
layout: &DataLayout,
|
||||||
|
dataspace: &Dataspace,
|
||||||
|
elem_size: usize,
|
||||||
|
pipeline: Option<&FilterPipeline>,
|
||||||
|
offset_size: u8,
|
||||||
|
length_size: u8,
|
||||||
|
selection: &Selection,
|
||||||
|
) -> Result<Option<Vec<u8>>, FormatError> {
|
||||||
|
let dims = &dataspace.dimensions;
|
||||||
|
if dims.is_empty() || elem_size == 0 {
|
||||||
|
return Ok(None);
|
||||||
|
}
|
||||||
|
let Some((box_start, box_extent)) = bounding_box(selection, dims) else {
|
||||||
|
return Ok(None);
|
||||||
|
};
|
||||||
|
let total = dataspace.checked_num_elements()?;
|
||||||
|
let box_elements = box_extent
|
||||||
|
.iter()
|
||||||
|
.try_fold(1u64, |acc, &e| acc.checked_mul(e))
|
||||||
|
.ok_or_else(|| FormatError::Overflow("selection bounding box overflows".into()))?;
|
||||||
|
// A box covering most of the dataset gains nothing over the full path.
|
||||||
|
if box_elements.saturating_mul(2) > total {
|
||||||
|
return Ok(None);
|
||||||
|
}
|
||||||
|
let mut boxed = alloc_output(checked_byte_len(box_elements, elem_size)?)?;
|
||||||
|
|
||||||
|
match layout {
|
||||||
|
DataLayout::Contiguous {
|
||||||
|
address: Some(address),
|
||||||
|
..
|
||||||
|
} => {
|
||||||
|
let base = usize::try_from(*address)
|
||||||
|
.map_err(|_| FormatError::Overflow("data address exceeds usize".into()))?;
|
||||||
|
let data = file_data
|
||||||
|
.get(base..)
|
||||||
|
.and_then(|d| d.get(..checked_byte_len(total, elem_size).ok()?))
|
||||||
|
.ok_or(FormatError::UnexpectedEof {
|
||||||
|
expected: base,
|
||||||
|
available: file_data.len(),
|
||||||
|
})?;
|
||||||
|
let origin = vec![0u64; dims.len()];
|
||||||
|
copy_overlap(
|
||||||
|
data,
|
||||||
|
&origin,
|
||||||
|
dims,
|
||||||
|
&mut boxed,
|
||||||
|
&box_start,
|
||||||
|
&box_extent,
|
||||||
|
elem_size,
|
||||||
|
);
|
||||||
|
}
|
||||||
|
DataLayout::Chunked {
|
||||||
|
btree_address: Some(_),
|
||||||
|
..
|
||||||
|
} => {
|
||||||
|
let (chunks, chunk_dims) = list_chunks(
|
||||||
|
file_data,
|
||||||
|
layout,
|
||||||
|
dataspace,
|
||||||
|
elem_size,
|
||||||
|
offset_size,
|
||||||
|
length_size,
|
||||||
|
)?;
|
||||||
|
let rank = dims.len();
|
||||||
|
let chunk_shape: Vec<u64> = chunk_dims.iter().map(|&d| d as u64).collect();
|
||||||
|
let chunk_bytes = crate::chunked_read::checked_chunk_byte_len(&chunk_dims, elem_size)?;
|
||||||
|
for chunk in &chunks {
|
||||||
|
if chunk.offsets.len() < rank || chunk.address == u64::MAX {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
let origin = &chunk.offsets[..rank];
|
||||||
|
let overlaps = (0..rank).all(|d| {
|
||||||
|
origin[d] < box_start[d] + box_extent[d]
|
||||||
|
&& origin[d].saturating_add(chunk_shape[d]) > box_start[d]
|
||||||
|
});
|
||||||
|
if !overlaps {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
let at = usize::try_from(chunk.address)
|
||||||
|
.map_err(|_| FormatError::Overflow("chunk address exceeds usize".into()))?;
|
||||||
|
let raw = at
|
||||||
|
.checked_add(chunk.chunk_size as usize)
|
||||||
|
.and_then(|end| file_data.get(at..end))
|
||||||
|
.ok_or(FormatError::UnexpectedEof {
|
||||||
|
expected: at.saturating_add(chunk.chunk_size as usize),
|
||||||
|
available: file_data.len(),
|
||||||
|
})?;
|
||||||
|
// Mirrors the full-read path: filter-mask bit i set means
|
||||||
|
// filter i was not applied to this chunk.
|
||||||
|
let decoded;
|
||||||
|
let data: &[u8] = match pipeline {
|
||||||
|
Some(pl) if !all_filters_skipped(pl, chunk.filter_mask) => {
|
||||||
|
decoded = decompress_chunk_masked(
|
||||||
|
raw,
|
||||||
|
pl,
|
||||||
|
chunk_bytes,
|
||||||
|
elem_size as u32,
|
||||||
|
chunk.filter_mask,
|
||||||
|
)?;
|
||||||
|
&decoded
|
||||||
|
}
|
||||||
|
_ => raw,
|
||||||
|
};
|
||||||
|
copy_overlap(
|
||||||
|
data,
|
||||||
|
origin,
|
||||||
|
&chunk_shape,
|
||||||
|
&mut boxed,
|
||||||
|
&box_start,
|
||||||
|
&box_extent,
|
||||||
|
elem_size,
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
_ => return Ok(None),
|
||||||
|
}
|
||||||
|
|
||||||
|
extract_selection_from_buffer(
|
||||||
|
&boxed,
|
||||||
|
&box_extent,
|
||||||
|
elem_size,
|
||||||
|
&translate(selection, &box_start),
|
||||||
|
)
|
||||||
|
.map(Some)
|
||||||
|
}
|
||||||
@@ -43,7 +43,7 @@ impl Default for DatasetCreateProps {
|
|||||||
fletcher32: false,
|
fletcher32: false,
|
||||||
lz4: false,
|
lz4: false,
|
||||||
zstd_level: None,
|
zstd_level: None,
|
||||||
fill_time: FillTime::Alloc,
|
fill_time: FillTime::IfSet,
|
||||||
compact: false,
|
compact: false,
|
||||||
alignment: 0,
|
alignment: 0,
|
||||||
}
|
}
|
||||||
@@ -335,7 +335,7 @@ mod tests {
|
|||||||
fn dcpl_defaults() {
|
fn dcpl_defaults() {
|
||||||
let dcpl = DatasetCreateProps::new();
|
let dcpl = DatasetCreateProps::new();
|
||||||
assert!(dcpl.chunk_dims.is_none());
|
assert!(dcpl.chunk_dims.is_none());
|
||||||
assert_eq!(dcpl.fill_time, FillTime::Alloc);
|
assert_eq!(dcpl.fill_time, FillTime::IfSet);
|
||||||
assert!(!dcpl.compact);
|
assert!(!dcpl.compact);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -225,9 +225,12 @@ pub fn parse_sohm_table_message(
|
|||||||
|
|
||||||
/// Parse the SOHM table structure (signature "SMTB") from the file.
|
/// Parse the SOHM table structure (signature "SMTB") from the file.
|
||||||
///
|
///
|
||||||
/// Each index entry: index_type(1) + mesg_types(2) + min_mesg_size(4) +
|
/// Each index entry: version(1) + index_type(1) + mesg_types(2) +
|
||||||
/// list_max(2) + btree_min(2) + num_messages(2) + index_addr(offset_size) +
|
/// min_mesg_size(4) + list_max(2) + btree_min(2) + num_messages(2) +
|
||||||
/// heap_addr(offset_size)
|
/// index_addr(offset_size) + heap_addr(offset_size)
|
||||||
|
///
|
||||||
|
/// The leading per-index version byte (0) was missing here, so every field
|
||||||
|
/// after it was read one byte off — verified against an HDF5 2.0 file.
|
||||||
pub fn parse_sohm_table(
|
pub fn parse_sohm_table(
|
||||||
file_data: &[u8],
|
file_data: &[u8],
|
||||||
table_addr: usize,
|
table_addr: usize,
|
||||||
@@ -240,11 +243,16 @@ pub fn parse_sohm_table(
|
|||||||
}
|
}
|
||||||
let mut pos = table_addr + 4;
|
let mut pos = table_addr + 4;
|
||||||
let os = offset_size as usize;
|
let os = offset_size as usize;
|
||||||
let entry_size = 1 + 2 + 4 + 2 + 2 + 2 + os + os; // 13 + 2*offset_size
|
let entry_size = 1 + 1 + 2 + 4 + 2 + 2 + 2 + os + os; // 14 + 2*offset_size
|
||||||
|
|
||||||
let mut indexes = Vec::with_capacity(nindexes as usize);
|
let mut indexes = Vec::with_capacity(nindexes as usize);
|
||||||
for _ in 0..nindexes {
|
for _ in 0..nindexes {
|
||||||
ensure_len(file_data, pos, entry_size)?;
|
ensure_len(file_data, pos, entry_size)?;
|
||||||
|
let version = file_data[pos];
|
||||||
|
if version != 0 {
|
||||||
|
return Err(FormatError::InvalidSohmTableVersion(version));
|
||||||
|
}
|
||||||
|
pos += 1;
|
||||||
let index_type = file_data[pos];
|
let index_type = file_data[pos];
|
||||||
pos += 1;
|
pos += 1;
|
||||||
let mesg_types = u16::from_le_bytes([file_data[pos], file_data[pos + 1]]);
|
let mesg_types = u16::from_le_bytes([file_data[pos], file_data[pos + 1]]);
|
||||||
@@ -381,6 +389,68 @@ pub fn parse_sohm_btree_entries(
|
|||||||
// ---- SOHM resolution ----
|
// ---- SOHM resolution ----
|
||||||
|
|
||||||
/// Find the SOHM index that handles the given message type.
|
/// Find the SOHM index that handles the given message type.
|
||||||
|
/// Load a file's SOHM table: superblock → superblock extension → Shared
|
||||||
|
/// Message Table message → SMTB. `Ok(None)` when the file has no superblock
|
||||||
|
/// extension or no shared-message table.
|
||||||
|
pub fn load_sohm_table(
|
||||||
|
file_data: &[u8],
|
||||||
|
offset_size: u8,
|
||||||
|
length_size: u8,
|
||||||
|
) -> Result<Option<SohmTable>, FormatError> {
|
||||||
|
let sig = crate::signature::find_signature(file_data)?;
|
||||||
|
let sb = crate::superblock::Superblock::parse(file_data, sig)?;
|
||||||
|
let Some(ext_addr) = sb
|
||||||
|
.superblock_extension_address
|
||||||
|
.filter(|&a| !is_undefined(a, offset_size))
|
||||||
|
else {
|
||||||
|
return Ok(None);
|
||||||
|
};
|
||||||
|
let ext = ObjectHeader::parse(file_data, ext_addr as usize, offset_size, length_size)?;
|
||||||
|
let Some(msg) = ext
|
||||||
|
.messages
|
||||||
|
.iter()
|
||||||
|
.find(|m| m.msg_type == MessageType::SharedMessageTable)
|
||||||
|
else {
|
||||||
|
return Ok(None);
|
||||||
|
};
|
||||||
|
let table_msg = parse_sohm_table_message(&msg.data, offset_size)?;
|
||||||
|
parse_sohm_table(
|
||||||
|
file_data,
|
||||||
|
table_msg.table_address as usize,
|
||||||
|
table_msg.nindexes,
|
||||||
|
offset_size,
|
||||||
|
)
|
||||||
|
.map(Some)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Like [`message_data`], but also follows references into the file's SOHM
|
||||||
|
/// heap (shared object header messages), loading the SOHM table on demand.
|
||||||
|
pub fn message_data_with_sohm<'a>(
|
||||||
|
file_data: &[u8],
|
||||||
|
msg: &'a crate::object_header::HeaderMessage,
|
||||||
|
offset_size: u8,
|
||||||
|
length_size: u8,
|
||||||
|
) -> Result<Cow<'a, [u8]>, FormatError> {
|
||||||
|
if !is_shared(msg.flags) {
|
||||||
|
return Ok(Cow::Borrowed(&msg.data));
|
||||||
|
}
|
||||||
|
let shared_ref = parse_shared_ref(&msg.data, offset_size)?;
|
||||||
|
let table = if shared_ref.heap_id.is_some() {
|
||||||
|
load_sohm_table(file_data, offset_size, length_size)?
|
||||||
|
} else {
|
||||||
|
None
|
||||||
|
};
|
||||||
|
resolve_shared_message_with_sohm(
|
||||||
|
file_data,
|
||||||
|
&shared_ref,
|
||||||
|
msg.msg_type,
|
||||||
|
offset_size,
|
||||||
|
length_size,
|
||||||
|
table.as_ref(),
|
||||||
|
)
|
||||||
|
.map(Cow::Owned)
|
||||||
|
}
|
||||||
|
|
||||||
fn find_index_for_msg_type(table: &SohmTable, msg_type: MessageType) -> Option<&SohmIndex> {
|
fn find_index_for_msg_type(table: &SohmTable, msg_type: MessageType) -> Option<&SohmIndex> {
|
||||||
let type_bit = 1u16 << msg_type.to_u16();
|
let type_bit = 1u16 << msg_type.to_u16();
|
||||||
table
|
table
|
||||||
@@ -707,6 +777,7 @@ mod tests {
|
|||||||
let mut buf = Vec::new();
|
let mut buf = Vec::new();
|
||||||
buf.extend_from_slice(b"SMTB");
|
buf.extend_from_slice(b"SMTB");
|
||||||
for idx in indexes {
|
for idx in indexes {
|
||||||
|
buf.push(0); // version
|
||||||
buf.push(idx.index_type);
|
buf.push(idx.index_type);
|
||||||
buf.extend_from_slice(&idx.mesg_types.to_le_bytes());
|
buf.extend_from_slice(&idx.mesg_types.to_le_bytes());
|
||||||
buf.extend_from_slice(&idx.min_mesg_size.to_le_bytes());
|
buf.extend_from_slice(&idx.min_mesg_size.to_le_bytes());
|
||||||
|
|||||||
@@ -39,7 +39,13 @@ pub struct Superblock {
|
|||||||
pub superblock_extension_address: Option<u64>,
|
pub superblock_extension_address: Option<u64>,
|
||||||
/// CRC32C checksum (v2/v3 only).
|
/// CRC32C checksum (v2/v3 only).
|
||||||
pub checksum: Option<u32>,
|
pub checksum: Option<u32>,
|
||||||
/// Page size for page-buffer mode (v4 only). `None` for v0–v3.
|
/// Page size of the non-standard "version 4" superblock layout (v4 only).
|
||||||
|
/// `None` for v0–v3.
|
||||||
|
///
|
||||||
|
/// HDF5 has no superblock version 4 — libhdf5 refuses it. A real paged
|
||||||
|
/// file is a v2/v3 superblock whose extension holds a File Space Info
|
||||||
|
/// message (what `FileWriter::with_page_size` writes). This field is kept
|
||||||
|
/// only so such files written by older clawhdf5 versions still parse.
|
||||||
pub page_size: Option<u32>,
|
pub page_size: Option<u32>,
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -127,8 +133,9 @@ impl Superblock {
|
|||||||
|
|
||||||
/// Serialize this superblock to bytes.
|
/// Serialize this superblock to bytes.
|
||||||
///
|
///
|
||||||
/// Writes v2/v3 format, or v4 (with `page_size`) when `self.version == 4`.
|
/// Writes v2/v3 format, or the non-standard v4 (with `page_size`) when
|
||||||
/// Computes and appends Jenkins lookup3 checksum.
|
/// `self.version == 4` — which no HDF5 library opens; see
|
||||||
|
/// [`Self::page_size`]. Computes and appends Jenkins lookup3 checksum.
|
||||||
pub fn serialize(&self) -> Vec<u8> {
|
pub fn serialize(&self) -> Vec<u8> {
|
||||||
let mut buf = Vec::with_capacity(48);
|
let mut buf = Vec::with_capacity(48);
|
||||||
buf.extend_from_slice(&HDF5_SIGNATURE);
|
buf.extend_from_slice(&HDF5_SIGNATURE);
|
||||||
|
|||||||
@@ -15,29 +15,81 @@ use crate::datatype::{
|
|||||||
|
|
||||||
/// Controls when fill values are written to dataset storage.
|
/// Controls when fill values are written to dataset storage.
|
||||||
///
|
///
|
||||||
/// Corresponds to the HDF5 fill value message's "fill time" field.
|
/// Corresponds to the HDF5 fill value message's "fill time" field
|
||||||
|
/// (`H5D_fill_time_t`).
|
||||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
|
||||||
pub enum FillTime {
|
pub enum FillTime {
|
||||||
/// Never write fill values (0x02). Avoids initialization overhead
|
/// Never write fill values (`H5D_FILL_TIME_NEVER`). Avoids
|
||||||
/// for datasets that will be fully written before any read.
|
/// initialization overhead for datasets that will be fully written
|
||||||
|
/// before any read.
|
||||||
Never,
|
Never,
|
||||||
/// Write fill values at allocation time (0x0a). This is the default
|
/// Write fill values when storage is allocated (`H5D_FILL_TIME_ALLOC`).
|
||||||
/// and matches the HDF5 C library's behavior.
|
|
||||||
#[default]
|
|
||||||
Alloc,
|
Alloc,
|
||||||
/// Write fill values only when the fill value has been explicitly set (0x06).
|
/// Write fill values at allocation only if one was set explicitly
|
||||||
|
/// (`H5D_FILL_TIME_IFSET`). The default, as in the HDF5 C library.
|
||||||
|
#[default]
|
||||||
IfSet,
|
IfSet,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Space allocation time written with every fill value message: late
|
||||||
|
/// (`H5D_ALLOC_TIME_LATE`), bits 0-1 of the flags byte.
|
||||||
|
const ALLOC_TIME_LATE: u8 = 2;
|
||||||
|
|
||||||
impl FillTime {
|
impl FillTime {
|
||||||
/// Serialize to the byte used in the fill value message (version 3).
|
/// Serialize to the flags byte of a version 3 fill value message: the
|
||||||
|
/// space allocation time (late) in bits 0-1 and the fill time in bits
|
||||||
|
/// 2-3 (`H5D_FILL_TIME_ALLOC` = 0, `NEVER` = 1, `IFSET` = 2).
|
||||||
|
///
|
||||||
|
/// This used to put `Never` in the ALLOC slot, `Alloc` in IFSET and
|
||||||
|
/// `IfSet` in NEVER, so libhdf5 saw every choice as a different one.
|
||||||
pub fn to_byte(self) -> u8 {
|
pub fn to_byte(self) -> u8 {
|
||||||
match self {
|
ALLOC_TIME_LATE | (self.code() << 2)
|
||||||
FillTime::Never => 0x02,
|
}
|
||||||
FillTime::Alloc => 0x0a,
|
|
||||||
FillTime::IfSet => 0x06,
|
/// Decode the fill time from a version 3 fill value message's flags.
|
||||||
|
pub fn from_byte(flags: u8) -> Option<FillTime> {
|
||||||
|
match (flags >> 2) & 0x03 {
|
||||||
|
0 => Some(FillTime::Alloc),
|
||||||
|
1 => Some(FillTime::Never),
|
||||||
|
2 => Some(FillTime::IfSet),
|
||||||
|
_ => None,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn code(self) -> u8 {
|
||||||
|
match self {
|
||||||
|
FillTime::Alloc => 0,
|
||||||
|
FillTime::Never => 1,
|
||||||
|
FillTime::IfSet => 2,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Serialize a version 3 Fill Value message for a dataset of `dt`: the fill
|
||||||
|
/// time, and the user-defined fill value if there is one (bit 5).
|
||||||
|
pub(crate) fn fill_value_message(
|
||||||
|
fill_time: FillTime,
|
||||||
|
value: Option<&[u8]>,
|
||||||
|
dt: &Datatype,
|
||||||
|
) -> Result<Vec<u8>, crate::error::FormatError> {
|
||||||
|
let mut msg = vec![3, fill_time.to_byte()];
|
||||||
|
if let Some(value) = value {
|
||||||
|
if matches!(dt, Datatype::VariableLength { .. }) {
|
||||||
|
return Err(crate::error::FormatError::SerializationError(
|
||||||
|
"a fill value for a variable-length datatype is not supported".into(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
if value.len() != dt.type_size() as usize {
|
||||||
|
return Err(crate::error::FormatError::DataSizeMismatch {
|
||||||
|
expected: dt.type_size() as usize,
|
||||||
|
actual: value.len(),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
msg[1] |= 0x20; // fill value defined
|
||||||
|
msg.extend_from_slice(&(value.len() as u32).to_le_bytes());
|
||||||
|
msg.extend_from_slice(value);
|
||||||
|
}
|
||||||
|
Ok(msg)
|
||||||
}
|
}
|
||||||
|
|
||||||
// ---- Datatype constructors ----
|
// ---- Datatype constructors ----
|
||||||
@@ -56,6 +108,21 @@ pub fn make_f64_type() -> Datatype {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// IEEE-754 half precision (binary16), little-endian — numpy's `float16`.
|
||||||
|
pub fn make_f16_type() -> Datatype {
|
||||||
|
Datatype::FloatingPoint {
|
||||||
|
size: 2,
|
||||||
|
byte_order: DatatypeByteOrder::LittleEndian,
|
||||||
|
bit_offset: 0,
|
||||||
|
bit_precision: 16,
|
||||||
|
exponent_location: 10,
|
||||||
|
exponent_size: 5,
|
||||||
|
mantissa_location: 0,
|
||||||
|
mantissa_size: 10,
|
||||||
|
exponent_bias: 15,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
pub fn make_f32_type() -> Datatype {
|
pub fn make_f32_type() -> Datatype {
|
||||||
Datatype::FloatingPoint {
|
Datatype::FloatingPoint {
|
||||||
size: 4,
|
size: 4,
|
||||||
@@ -317,7 +384,11 @@ pub(crate) fn build_attr_message(name: &str, value: &AttrValue) -> AttributeMess
|
|||||||
raw_data: data.clone(),
|
raw_data: data.clone(),
|
||||||
},
|
},
|
||||||
AttrValue::String(s) => {
|
AttrValue::String(s) => {
|
||||||
let bytes = s.as_bytes();
|
// A fixed-length string type must be at least 1 byte: libhdf5
|
||||||
|
// rejects size 0 ("invalid datatype size") and with it every
|
||||||
|
// attribute on the object. h5py stores "" as one NUL byte.
|
||||||
|
let mut bytes = s.as_bytes().to_vec();
|
||||||
|
bytes.resize(bytes.len().max(1), 0);
|
||||||
AttributeMessage {
|
AttributeMessage {
|
||||||
name: name.to_string(),
|
name: name.to_string(),
|
||||||
datatype: Datatype::String {
|
datatype: Datatype::String {
|
||||||
@@ -326,11 +397,12 @@ pub(crate) fn build_attr_message(name: &str, value: &AttrValue) -> AttributeMess
|
|||||||
charset: CharacterSet::Utf8,
|
charset: CharacterSet::Utf8,
|
||||||
},
|
},
|
||||||
dataspace: scalar_ds(),
|
dataspace: scalar_ds(),
|
||||||
raw_data: bytes.to_vec(),
|
raw_data: bytes,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
AttrValue::StringArray(arr) => {
|
AttrValue::StringArray(arr) => {
|
||||||
let max_len = arr.iter().map(|s| s.len()).max().unwrap_or(0);
|
// At least 1 byte per element, as for a single string.
|
||||||
|
let max_len = arr.iter().map(|s| s.len()).max().unwrap_or(0).max(1);
|
||||||
let mut raw = Vec::new();
|
let mut raw = Vec::new();
|
||||||
for s in arr {
|
for s in arr {
|
||||||
let mut b = s.as_bytes().to_vec();
|
let mut b = s.as_bytes().to_vec();
|
||||||
@@ -416,8 +488,10 @@ pub struct DatasetBuilder {
|
|||||||
pub(crate) data: Option<Vec<u8>>,
|
pub(crate) data: Option<Vec<u8>>,
|
||||||
pub(crate) attrs: Vec<(String, AttrValue)>,
|
pub(crate) attrs: Vec<(String, AttrValue)>,
|
||||||
pub(crate) chunk_options: ChunkOptions,
|
pub(crate) chunk_options: ChunkOptions,
|
||||||
/// Controls when fill values are written. Default is `FillTime::Alloc`.
|
/// Controls when fill values are written. Default is `FillTime::IfSet`.
|
||||||
pub(crate) fill_time: FillTime,
|
pub(crate) fill_time: FillTime,
|
||||||
|
/// User-defined fill value: one element's bytes, as stored.
|
||||||
|
pub(crate) fill_value: Option<Vec<u8>>,
|
||||||
/// Use compact (inline) storage: data is stored in the object header.
|
/// Use compact (inline) storage: data is stored in the object header.
|
||||||
/// Only valid when raw data is <= 65536 bytes and dataset is not chunked.
|
/// Only valid when raw data is <= 65536 bytes and dataset is not chunked.
|
||||||
pub(crate) compact: bool,
|
pub(crate) compact: bool,
|
||||||
@@ -444,6 +518,7 @@ impl DatasetBuilder {
|
|||||||
attrs: Vec::new(),
|
attrs: Vec::new(),
|
||||||
chunk_options: ChunkOptions::default(),
|
chunk_options: ChunkOptions::default(),
|
||||||
fill_time: FillTime::default(),
|
fill_time: FillTime::default(),
|
||||||
|
fill_value: None,
|
||||||
compact: false,
|
compact: false,
|
||||||
alignment: 0,
|
alignment: 0,
|
||||||
virtual_sources: None,
|
virtual_sources: None,
|
||||||
@@ -478,6 +553,24 @@ impl DatasetBuilder {
|
|||||||
self
|
self
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Store `data` as IEEE half precision (numpy `float16`), rounding each
|
||||||
|
/// value to the nearest half ([`crate::float16::f32_to_f16_bits`]).
|
||||||
|
/// Half the bytes of [`Self::with_f32_data`], at about three significant
|
||||||
|
/// decimal digits; values beyond ±65504 become ±infinity. Reading it back
|
||||||
|
/// with `read_f32` yields the rounded values exactly.
|
||||||
|
pub fn with_f16_data(&mut self, data: &[f32]) -> &mut Self {
|
||||||
|
self.datatype = Some(make_f16_type());
|
||||||
|
let mut b = Vec::with_capacity(data.len() * 2);
|
||||||
|
for &v in data {
|
||||||
|
b.extend_from_slice(&crate::float16::f32_to_f16_bits(v).to_le_bytes());
|
||||||
|
}
|
||||||
|
self.data = Some(b);
|
||||||
|
if self.shape.is_none() {
|
||||||
|
self.shape = Some(vec![data.len() as u64]);
|
||||||
|
}
|
||||||
|
self
|
||||||
|
}
|
||||||
|
|
||||||
pub fn with_i32_data(&mut self, data: &[i32]) -> &mut Self {
|
pub fn with_i32_data(&mut self, data: &[i32]) -> &mut Self {
|
||||||
self.datatype = Some(make_i32_type());
|
self.datatype = Some(make_i32_type());
|
||||||
let mut b = Vec::with_capacity(data.len() * 4);
|
let mut b = Vec::with_capacity(data.len() * 4);
|
||||||
@@ -638,7 +731,12 @@ impl DatasetBuilder {
|
|||||||
self
|
self
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Enable Pcodec lossless numerical compression (clawhdf5 filter ID 32023).
|
/// Enable Pcodec lossless numerical compression (private clawhdf5 filter
|
||||||
|
/// ID 480).
|
||||||
|
///
|
||||||
|
/// **Not interoperable:** pcodec has no registered HDF5 filter ID and no
|
||||||
|
/// libhdf5 plugin, so h5py and other HDF5 readers cannot read the
|
||||||
|
/// dataset — only clawhdf5 built with the `pcodec` feature can.
|
||||||
///
|
///
|
||||||
/// Pcodec achieves 30–94% better compression ratio than Zstd for f32/f64
|
/// Pcodec achieves 30–94% better compression ratio than Zstd for f32/f64
|
||||||
/// columns at 1–5 GiB/s decompression speed (arXiv:2502.06112). Requires
|
/// columns at 1–5 GiB/s decompression speed (arXiv:2502.06112). Requires
|
||||||
@@ -682,10 +780,20 @@ impl DatasetBuilder {
|
|||||||
self
|
self
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Set the dataset's fill value: what readers return for storage that
|
||||||
|
/// was never written (e.g. after the dataset is extended). `value` is one
|
||||||
|
/// element's bytes as stored — the dataset datatype's size and byte order
|
||||||
|
/// (`(-1i32).to_le_bytes()` for an `i32` dataset). A size mismatch, or a
|
||||||
|
/// variable-length datatype, makes `finish` fail.
|
||||||
|
pub fn with_fill_value(&mut self, value: &[u8]) -> &mut Self {
|
||||||
|
self.fill_value = Some(value.to_vec());
|
||||||
|
self
|
||||||
|
}
|
||||||
|
|
||||||
/// Use compact (inline) storage for this dataset.
|
/// Use compact (inline) storage for this dataset.
|
||||||
///
|
///
|
||||||
/// The raw data is stored directly in the dataset's object header rather
|
/// The raw data is stored directly in the dataset's object header rather
|
||||||
/// than as a separate data blob. Only effective when raw data <= 65536 bytes
|
/// than as a separate data blob. Only effective when raw data <= 65531 bytes
|
||||||
/// and the dataset is not chunked.
|
/// and the dataset is not chunked.
|
||||||
pub fn compact(&mut self) -> &mut Self {
|
pub fn compact(&mut self) -> &mut Self {
|
||||||
self.compact = true;
|
self.compact = true;
|
||||||
|
|||||||
@@ -148,7 +148,12 @@ pub fn read_vl_strings(
|
|||||||
Ok(result)
|
Ok(result)
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Resolve VL byte sequences from raw data.
|
/// Resolve VL sequences from raw data, returning each element's bytes.
|
||||||
|
///
|
||||||
|
/// Each element is the sequence's full encoding — element count × base type
|
||||||
|
/// size bytes, in the base type's byte order — so a sequence of `i32` yields
|
||||||
|
/// four bytes per value. Decode it with the base type (e.g.
|
||||||
|
/// [`crate::data_read::read_as_i64`]).
|
||||||
pub fn read_vl_bytes(
|
pub fn read_vl_bytes(
|
||||||
file_data: &[u8],
|
file_data: &[u8],
|
||||||
raw_data: &[u8],
|
raw_data: &[u8],
|
||||||
@@ -177,8 +182,10 @@ pub fn read_vl_bytes(
|
|||||||
},
|
},
|
||||||
)?;
|
)?;
|
||||||
|
|
||||||
let len = (vl.length as usize).min(obj.data.len());
|
// The heap object holds the whole sequence. `vl.length` counts
|
||||||
result.push(obj.data[..len].to_vec());
|
// elements, not bytes, so it is only the byte length when the base
|
||||||
|
// type is one byte wide.
|
||||||
|
result.push(obj.data.clone());
|
||||||
}
|
}
|
||||||
|
|
||||||
Ok(result)
|
Ok(result)
|
||||||
|
|||||||
@@ -0,0 +1,13 @@
|
|||||||
|
# Filter conformance fixtures
|
||||||
|
|
||||||
|
Files written by libhdf5 (and its registered filter plugins), used by the
|
||||||
|
filter regression tests in `src/filters.rs` to compare our decoders against
|
||||||
|
the values h5py/libhdf5 read from the same bytes. Chunk byte ranges quoted in
|
||||||
|
the tests come from h5py's `DatasetID.get_chunk_info`.
|
||||||
|
|
||||||
|
| File | Origin | Licence |
|
||||||
|
|------|--------|---------|
|
||||||
|
| `h5ex_d_lz4.h5` | HDF Group `HDF5Examples/C/H5FLT/tfiles/h5ex_d_lz4.h5` (hdf5 repository) | HDF5 licence (BSD-3-Clause style) |
|
||||||
|
| `noencoder.h5` | HDF Group `test/testfiles/noencoder.h5` (hdf5 repository) | HDF5 licence (BSD-3-Clause style) |
|
||||||
|
| `le_data.h5` | HDF Group `test/testfiles/le_data.h5` (hdf5 repository) | HDF5 licence (BSD-3-Clause style) |
|
||||||
|
| `szip_h5py.h5` | Written for these tests with h5py 3 / libhdf5 2.0.0 (libaec szip): `f8` (8x10, chunks 4x10, `('nn', 8)`), `i8` (8x10, chunks 4x10, `('ec', 4)`), `u2` (70, chunks 35, `('nn', 8)`) | Same as this repository |
|
||||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,49 @@
|
|||||||
|
"""Generate shared_fill_value.h5: datasets whose Fill Value message is
|
||||||
|
*shared*, in the two ways libhdf5 can share one.
|
||||||
|
|
||||||
|
- /sohm_a, /sohm_b: the file has a shared-object-header-message (SOHM) index
|
||||||
|
for fill values, so libhdf5 stores the fill value (-7, int32) in the SOHM
|
||||||
|
heap and /sohm_b's header holds only a reference to it. Chunked, with only
|
||||||
|
the first chunk written, so the rest reads as the fill value.
|
||||||
|
- /unwritten_a, /unwritten_b: the same, never written: no storage at all,
|
||||||
|
read entirely as the fill value.
|
||||||
|
|
||||||
|
h5py has no API for SOHM indexes, so the file creation property list is
|
||||||
|
configured by calling the libhdf5 bundled in the h5py wheel through ctypes.
|
||||||
|
Written with h5py 3.16.0 / HDF5 2.0.0. Re-run only to regenerate:
|
||||||
|
|
||||||
|
python gen_shared_fill.py shared_fill_value.h5
|
||||||
|
"""
|
||||||
|
import ctypes
|
||||||
|
import glob
|
||||||
|
import os
|
||||||
|
import sys
|
||||||
|
|
||||||
|
import h5py
|
||||||
|
import numpy as np
|
||||||
|
|
||||||
|
libdir = os.path.join(os.path.dirname(os.path.dirname(h5py.__file__)), "h5py.libs")
|
||||||
|
libs = [p for p in glob.glob(os.path.join(libdir, "libhdf5*.so*")) if "_hl" not in os.path.basename(p)]
|
||||||
|
lib = ctypes.CDLL(libs[0])
|
||||||
|
lib.H5open()
|
||||||
|
|
||||||
|
H5O_SHMESG_FILL_FLAG = 1 << 0x0005
|
||||||
|
|
||||||
|
fcpl = h5py.h5p.create(h5py.h5p.FILE_CREATE)
|
||||||
|
lib.H5Pset_shared_mesg_nindexes.argtypes = [ctypes.c_int64, ctypes.c_uint]
|
||||||
|
lib.H5Pset_shared_mesg_index.argtypes = [ctypes.c_int64, ctypes.c_uint, ctypes.c_uint, ctypes.c_uint]
|
||||||
|
assert lib.H5Pset_shared_mesg_nindexes(fcpl.id, 1) >= 0
|
||||||
|
assert lib.H5Pset_shared_mesg_index(fcpl.id, 0, H5O_SHMESG_FILL_FLAG, 0) >= 0
|
||||||
|
|
||||||
|
fapl = h5py.h5p.create(h5py.h5p.FILE_ACCESS)
|
||||||
|
fapl.set_libver_bounds(h5py.h5f.LIBVER_LATEST, h5py.h5f.LIBVER_LATEST)
|
||||||
|
fid = h5py.h5f.create(sys.argv[1].encode(), h5py.h5f.ACC_TRUNC, fcpl=fcpl, fapl=fapl)
|
||||||
|
with h5py.File(fid) as f:
|
||||||
|
# Chunked, with only the first chunk written: the rest reads as fill.
|
||||||
|
# libhdf5 keeps the first copy of a message in its own header; the second
|
||||||
|
# identical one (the `_b` datasets) is the SOHM reference.
|
||||||
|
for name in ("sohm_a", "sohm_b"):
|
||||||
|
d = f.create_dataset(name, shape=(8,), chunks=(4,), dtype="<i4", fillvalue=-7)
|
||||||
|
d[:4] = np.arange(4)
|
||||||
|
for name in ("unwritten_a", "unwritten_b"):
|
||||||
|
f.create_dataset(name, shape=(3,), dtype="<i4", fillvalue=-7)
|
||||||
@@ -0,0 +1,44 @@
|
|||||||
|
"""Generate std_ref_hdf5_2_0.h5: a dataset of H5T_STD_REF (the reference
|
||||||
|
datatype introduced in HDF5 1.12, datatype message version 4) holding two
|
||||||
|
object references — to /target (a dataset) and /grp (a group).
|
||||||
|
|
||||||
|
h5py has no API for this type, so the file is written by calling the libhdf5
|
||||||
|
bundled in the h5py wheel directly through ctypes. Written with h5py 3.16.0 /
|
||||||
|
HDF5 2.0.0. Re-run only if the fixture ever needs regenerating:
|
||||||
|
|
||||||
|
python gen_std_ref.py std_ref_hdf5_2_0.h5
|
||||||
|
"""
|
||||||
|
import ctypes
|
||||||
|
import glob
|
||||||
|
import os
|
||||||
|
import sys
|
||||||
|
|
||||||
|
import h5py
|
||||||
|
import numpy as np
|
||||||
|
|
||||||
|
libdir = os.path.join(os.path.dirname(os.path.dirname(h5py.__file__)), "h5py.libs")
|
||||||
|
libs = [p for p in glob.glob(os.path.join(libdir, "libhdf5*.so*")) if "_hl" not in os.path.basename(p)]
|
||||||
|
lib = ctypes.CDLL(libs[0])
|
||||||
|
lib.H5open()
|
||||||
|
hid = ctypes.c_int64
|
||||||
|
std_ref = hid.in_dll(lib, "H5T_STD_REF_g").value
|
||||||
|
|
||||||
|
lib.H5Screate_simple.restype = hid
|
||||||
|
lib.H5Screate_simple.argtypes = [ctypes.c_int, ctypes.POINTER(ctypes.c_uint64), ctypes.POINTER(ctypes.c_uint64)]
|
||||||
|
lib.H5Dcreate2.restype = hid
|
||||||
|
lib.H5Dcreate2.argtypes = [hid, ctypes.c_char_p, hid, hid, hid, hid, hid]
|
||||||
|
lib.H5Rcreate_object.argtypes = [hid, ctypes.c_char_p, hid, ctypes.c_void_p]
|
||||||
|
lib.H5Dwrite.argtypes = [hid, hid, hid, hid, hid, ctypes.c_void_p]
|
||||||
|
lib.H5Dclose.argtypes = [hid]
|
||||||
|
|
||||||
|
with h5py.File(sys.argv[1], "w", libver="latest") as f:
|
||||||
|
f.create_dataset("target", data=np.arange(5, dtype="<i4"))
|
||||||
|
f.create_group("grp")
|
||||||
|
fid = f.id.id
|
||||||
|
sid = lib.H5Screate_simple(1, (ctypes.c_uint64 * 1)(2), None)
|
||||||
|
did = lib.H5Dcreate2(fid, b"refs", std_ref, sid, 0, 0, 0)
|
||||||
|
refs = ((ctypes.c_ubyte * 64) * 2)() # H5R_ref_t is a 64-byte buffer
|
||||||
|
assert lib.H5Rcreate_object(fid, b"/target", 0, ctypes.byref(refs[0])) == 0
|
||||||
|
assert lib.H5Rcreate_object(fid, b"/grp", 0, ctypes.byref(refs[1])) == 0
|
||||||
|
assert lib.H5Dwrite(did, std_ref, 0, 0, 0, ctypes.byref(refs)) == 0
|
||||||
|
lib.H5Dclose(did)
|
||||||
Binary file not shown.
Binary file not shown.
BIN
Binary file not shown.
@@ -2,6 +2,15 @@
|
|||||||
|
|
||||||
use clawhdf5_format::data_read::{read_object_references, read_region_references};
|
use clawhdf5_format::data_read::{read_object_references, read_region_references};
|
||||||
use clawhdf5_format::datatype::{Datatype, ReferenceType};
|
use clawhdf5_format::datatype::{Datatype, ReferenceType};
|
||||||
|
/// The Python interpreter to drive interop checks with.
|
||||||
|
///
|
||||||
|
/// `CLAWHDF5_PYTHON` lets these run against a virtualenv holding h5py, which
|
||||||
|
/// on a PEP 668 "externally managed" system is the only place it can be
|
||||||
|
/// installed. Without it the suite silently skips, and a silent skip here is
|
||||||
|
/// how a datatype bug once reached a release.
|
||||||
|
fn python() -> String {
|
||||||
|
std::env::var("CLAWHDF5_PYTHON").unwrap_or_else(|_| "python3".to_string())
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn object_ref_single_valid() {
|
fn object_ref_single_valid() {
|
||||||
@@ -173,7 +182,7 @@ print('ok')
|
|||||||
"#,
|
"#,
|
||||||
path.display()
|
path.display()
|
||||||
);
|
);
|
||||||
let output = std::process::Command::new("python3")
|
let output = std::process::Command::new(python())
|
||||||
.args(["-c", &script])
|
.args(["-c", &script])
|
||||||
.output();
|
.output();
|
||||||
|
|
||||||
@@ -316,3 +325,97 @@ print('ok')
|
|||||||
// Clean up
|
// Clean up
|
||||||
let _ = std::fs::remove_file(&path);
|
let _ = std::fs::remove_file(&path);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// H5T_STD_REF (HDF5 1.12+ references, datatype message version 4)
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
/// `fixtures/std_ref_hdf5_2_0.h5` (see `gen_std_ref.py`) holds a dataset of
|
||||||
|
/// `H5T_STD_REF` with two object references, written by HDF5 2.0 itself. The
|
||||||
|
/// datatype used to be rejected with `InvalidReferenceType(2)`.
|
||||||
|
#[test]
|
||||||
|
fn std_ref_object_references_from_hdf5_2_0() {
|
||||||
|
use clawhdf5_format::data_layout::DataLayout;
|
||||||
|
use clawhdf5_format::dataspace::Dataspace;
|
||||||
|
use clawhdf5_format::group_v2::resolve_path_any;
|
||||||
|
use clawhdf5_format::message_type::MessageType;
|
||||||
|
use clawhdf5_format::object_header::ObjectHeader;
|
||||||
|
use clawhdf5_format::signature::find_signature;
|
||||||
|
use clawhdf5_format::superblock::Superblock;
|
||||||
|
|
||||||
|
let bytes: &[u8] = include_bytes!("fixtures/std_ref_hdf5_2_0.h5");
|
||||||
|
let sb = Superblock::parse(bytes, find_signature(bytes).unwrap()).unwrap();
|
||||||
|
let (os, ls) = (sb.offset_size, sb.length_size);
|
||||||
|
|
||||||
|
let refs_addr = resolve_path_any(bytes, &sb, "refs").unwrap();
|
||||||
|
let header = ObjectHeader::parse(bytes, refs_addr as usize, os, ls).unwrap();
|
||||||
|
let message = |t: MessageType| {
|
||||||
|
&header
|
||||||
|
.messages
|
||||||
|
.iter()
|
||||||
|
.find(|m| m.msg_type == t)
|
||||||
|
.unwrap()
|
||||||
|
.data
|
||||||
|
};
|
||||||
|
|
||||||
|
let (datatype, _) = Datatype::parse(message(MessageType::Datatype)).unwrap();
|
||||||
|
assert_eq!(
|
||||||
|
datatype,
|
||||||
|
Datatype::Reference {
|
||||||
|
size: 18,
|
||||||
|
ref_type: ReferenceType::Object2
|
||||||
|
}
|
||||||
|
);
|
||||||
|
let dataspace = Dataspace::parse(message(MessageType::Dataspace), ls).unwrap();
|
||||||
|
let layout = DataLayout::parse(message(MessageType::DataLayout), os, ls).unwrap();
|
||||||
|
let raw =
|
||||||
|
clawhdf5_format::data_read::read_raw_data(bytes, &layout, &dataspace, &datatype).unwrap();
|
||||||
|
assert_eq!(raw.len(), 2 * 18);
|
||||||
|
|
||||||
|
// The references point at the objects they were created from.
|
||||||
|
let refs = read_object_references(&raw, &datatype, os).unwrap();
|
||||||
|
let addresses: Vec<u64> = refs.iter().map(|r| r.address).collect();
|
||||||
|
assert_eq!(
|
||||||
|
addresses,
|
||||||
|
[
|
||||||
|
resolve_path_any(bytes, &sb, "target").unwrap(),
|
||||||
|
resolve_path_any(bytes, &sb, "grp").unwrap(),
|
||||||
|
]
|
||||||
|
);
|
||||||
|
// And what they point at is a real object header.
|
||||||
|
for address in addresses {
|
||||||
|
ObjectHeader::parse(bytes, address as usize, os, ls).unwrap();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn std_ref_decoding_rejects_malformed_elements() {
|
||||||
|
let dt = Datatype::Reference {
|
||||||
|
size: 18,
|
||||||
|
ref_type: ReferenceType::Object2,
|
||||||
|
};
|
||||||
|
let mut good = vec![0u8; 18];
|
||||||
|
good[..4].copy_from_slice(&[2, 0, 8, 0xb3]);
|
||||||
|
assert_eq!(
|
||||||
|
read_object_references(&good, &dt, 8).unwrap()[0].address,
|
||||||
|
0xb3
|
||||||
|
);
|
||||||
|
|
||||||
|
// Null reference.
|
||||||
|
assert_eq!(
|
||||||
|
read_object_references(&[0u8; 18], &dt, 8).unwrap()[0].address,
|
||||||
|
u64::MAX
|
||||||
|
);
|
||||||
|
for (what, patch) in [
|
||||||
|
("wrong reference type", (0usize, 3u8)),
|
||||||
|
("external flag", (1, 1)),
|
||||||
|
("token longer than the element", (2, 200)),
|
||||||
|
("zero-length token", (2, 0)),
|
||||||
|
] {
|
||||||
|
let mut bad = good.clone();
|
||||||
|
bad[patch.0] = patch.1;
|
||||||
|
assert!(read_object_references(&bad, &dt, 8).is_err(), "{what}");
|
||||||
|
}
|
||||||
|
// Not a whole number of elements.
|
||||||
|
assert!(read_object_references(&good[..17], &dt, 8).is_err());
|
||||||
|
}
|
||||||
|
|||||||
@@ -312,3 +312,49 @@ fn provenance_mismatch_on_corruption() {
|
|||||||
"corrupted data should produce hash mismatch"
|
"corrupted data should produce hash mismatch"
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// Fuzzer finds, kept as regression tests
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
/// `fuzz_btree_v2` crash input from 2026-09-20 (82 bytes): a B-tree v2 header
|
||||||
|
/// followed by internal nodes that point back into themselves. It predates the
|
||||||
|
/// depth cap and record budget added to B-tree v2 traversal that day and no
|
||||||
|
/// longer crashes; this replays the fuzz target's exact code path on it so a
|
||||||
|
/// regression fails CI rather than waiting for a fuzz run.
|
||||||
|
#[test]
|
||||||
|
fn fuzz_btree_v2_crash_f98c19dc_is_a_clean_result() {
|
||||||
|
use clawhdf5_format::btree_v2::{BTreeV2Header, collect_btree_v2_records};
|
||||||
|
let data: &[u8] = &[
|
||||||
|
0x42, 0x54, 0x48, 0x44, 0x00, 0x06, 0x00, 0xed, 0xef, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00,
|
||||||
|
0x00, 0x03, 0x40, 0x14, 0x93, 0x42, 0x54, 0x49, 0x4e, 0x42, 0x00, 0x00, 0x00, 0x00, 0x00,
|
||||||
|
0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
|
||||||
|
0x00, 0x00, 0x00, 0x00, 0x42, 0x54, 0x48, 0x44, 0x00, 0x00, 0x00, 0x13, 0x05, 0x00, 0x00,
|
||||||
|
0x00, 0x00, 0x00, 0x00, 0x80, 0x00, 0x00, 0x00, 0x40, 0x14, 0x93, 0x42, 0x54, 0x00, 0x49,
|
||||||
|
0x00, 0x01, 0x4e, 0x42, 0x42, 0x54, 0xbe,
|
||||||
|
];
|
||||||
|
assert_eq!(data.len(), 82);
|
||||||
|
for offset_size in [4u8, 8] {
|
||||||
|
for length_size in [4u8, 8] {
|
||||||
|
if let Ok(header) = BTreeV2Header::parse(data, 0, offset_size, length_size) {
|
||||||
|
let _ = collect_btree_v2_records(data, &header, offset_size, length_size);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
let (fields, file) = data.split_first_chunk::<20>().unwrap();
|
||||||
|
let header = BTreeV2Header {
|
||||||
|
tree_type: fields[0],
|
||||||
|
node_size: u32::from_le_bytes([fields[1], fields[2], fields[3], fields[4]]),
|
||||||
|
record_size: u16::from_le_bytes([fields[5], fields[6]]),
|
||||||
|
depth: u16::from_le_bytes([fields[7], fields[8]]),
|
||||||
|
root_node_address: u64::from(u32::from_le_bytes([
|
||||||
|
fields[9], fields[10], fields[11], fields[12],
|
||||||
|
])),
|
||||||
|
num_records_in_root: u16::from_le_bytes([fields[13], fields[14]]),
|
||||||
|
total_records: u64::from(u32::from_le_bytes([
|
||||||
|
fields[15], fields[16], fields[17], fields[18],
|
||||||
|
])),
|
||||||
|
};
|
||||||
|
let offset_size = if fields[19] & 1 == 0 { 4 } else { 8 };
|
||||||
|
let _ = collect_btree_v2_records(file, &header, offset_size, 8);
|
||||||
|
}
|
||||||
|
|||||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user