Compare commits
48
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
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 |
+57
-11
@@ -9,15 +9,15 @@ jobs:
|
||||
runs-on: ubuntu-latest
|
||||
container: rust:latest
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- name: Cache cargo registry/target
|
||||
uses: actions/cache@v4
|
||||
with:
|
||||
path: |
|
||||
~/.cargo/registry
|
||||
~/.cargo/git
|
||||
target
|
||||
key: ${{ runner.os }}-cargo-${{ hashFiles('**/Cargo.lock') }}
|
||||
# Plain git rather than actions/checkout: that is a JavaScript action,
|
||||
# and rust:latest has no `node`, so it failed with exit 127 before any
|
||||
# code was built — on every push. actions/cache went for the same reason.
|
||||
- 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: Install rustfmt & clippy components
|
||||
run: rustup component add rustfmt clippy
|
||||
- name: Install thumbv7em-none-eabihf target
|
||||
@@ -28,13 +28,59 @@ jobs:
|
||||
# dependency a failure (CLAWHDF5_REQUIRE_INTEROP below).
|
||||
run: |
|
||||
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
|
||||
/opt/interop/bin/pip install --no-cache-dir h5py numpy netCDF4 xarray
|
||||
echo "/opt/interop/bin" >> "$GITHUB_PATH"
|
||||
- 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; print('h5py', h5py.__version__, 'HDF5', h5py.version.hdf5_version, 'netCDF4', netCDF4.__version__)"
|
||||
- name: Run CI script
|
||||
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"
|
||||
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
|
||||
weights/
|
||||
.venv
|
||||
|
||||
+326
@@ -28,6 +28,229 @@
|
||||
|
||||
---
|
||||
|
||||
## Memory footprint
|
||||
|
||||
`cargo run --release -p clawhdf5-bench --bin search_harness -- --footprint --full`,
|
||||
384-dim `f32`. The figure that matters is **reopened**: a store loaded from
|
||||
disk, which is what a long-lived process holds.
|
||||
|
||||
Measured with a counting global allocator, not RSS. RSS cannot see 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 at all. Measured that way a store holding the corpus twice and one
|
||||
holding it once came out *identical* (1.00x both), which is how the first
|
||||
attempt at this measurement went.
|
||||
|
||||
| N | vectors (raw) | reopened, before | reopened, after |
|
||||
|---:|---:|---:|---:|
|
||||
| 1 000 | 1 MiB | 5 MiB (3.41x) | 4 MiB (2.39x) |
|
||||
| 10 000 | 15 MiB | 50 MiB (3.43x) | 35 MiB (2.42x) |
|
||||
| 100 000 | 146 MiB | 505 MiB (3.44x) | **357 MiB (2.43x)** |
|
||||
|
||||
The cache stored every embedding twice — once as a `Vec<Vec<f32>>` and once
|
||||
flattened for the batched kernels, kept in lock-step on every push, update and
|
||||
compaction. Storing only the flat buffer and indexing into it gives back
|
||||
almost exactly one copy of the corpus (148 MiB at 100k) and one heap
|
||||
allocation per entry. Recall and query latency are unchanged.
|
||||
|
||||
What remains at 2.43x: the flat vectors (1.0x), the HNSW index's own copy of
|
||||
them (1.0x), and text, ids and graph (~0.4x). The index copy is the next
|
||||
target — it is what a quantised or borrowed representation would address.
|
||||
|
||||
### Quantising the index copy (`quantized_index`)
|
||||
|
||||
`MemoryConfig::quantized_index` stores the index's copy as `i8` instead of
|
||||
`f32`. Same harness, same binary, `--footprint --full` with and without
|
||||
`--int8`:
|
||||
|
||||
| N | vectors (raw) | indexes, f32 | indexes, int8 | reopened, f32 | reopened, int8 |
|
||||
|---:|---:|---:|---:|---:|---:|
|
||||
| 1 000 | 1 MiB | 2 MiB | 1 MiB | 4 MiB (2.40x) | 2 MiB (1.64x) |
|
||||
| 10 000 | 15 MiB | 32 MiB | 14 MiB | 44 MiB (3.03x) | 27 MiB (1.81x) |
|
||||
| 100 000 | 146 MiB | 266 MiB | **123 MiB** | 399 MiB (2.72x) | **256 MiB (1.74x)** |
|
||||
|
||||
The scale is **per row**, not global. A unit-length row in `d` dimensions has
|
||||
components around `1/sqrt(d)`, so a fixed `[-1, 1]` scale spends fewer than 12
|
||||
of the 255 levels on a 128-dimensional vector: measured against an exact
|
||||
ranking that gives 0.35 top-10 overlap — unusable. Scaling each row by its own
|
||||
largest component brings the same measurement to 0.99.
|
||||
|
||||
Quantised distances still cost recall on their own, and **`ef` does not buy it
|
||||
back**, because the loss is in the distances rather than in the graph
|
||||
(`--ann-only --full`, N = 100 000):
|
||||
|
||||
| ef | recall@10, f32 | recall@10, int8 | recall@10, int8 + re-score |
|
||||
|---:|---:|---:|---:|
|
||||
| 32 | 0.9775 | 0.9415 | 0.9785 |
|
||||
| 64 | 0.9945 | 0.9625 | 0.9940 |
|
||||
| 128 | 0.9995 | 0.9670 | 0.9990 |
|
||||
| 256 | 0.9995 | 0.9670 (ceiling) | 0.9990 |
|
||||
|
||||
Re-scoring closes the gap: the store already holds the exact embeddings, so
|
||||
the query path re-scores the candidate pool against them before fusion. That
|
||||
is done automatically whenever the index is quantised.
|
||||
|
||||
**On AVX2 this costs nothing — it pays.** The first measurement of this put
|
||||
the cost at ~13% of QPS and ~16% of build time, but that compared a scalar
|
||||
int8 loop against `clawhdf5-accel`'s hand-written AVX2 kernels for `f32`:
|
||||
the gap was a missing kernel, not a property of int8. With
|
||||
`clawhdf5_accel::dot_i8` (AVX2: sign-extend to `i16`, then `madd_epi16`),
|
||||
medians of three alternating runs at N = 100 000, same binary:
|
||||
|
||||
| | f32 | int8 | int8 + re-score |
|
||||
|---|---:|---:|---:|
|
||||
| build | 3197 ms | **1778 ms** | 1826 ms |
|
||||
| QPS at ef = 64 | 13 399 | 29 195 | **21 848** |
|
||||
| recall@10 at ef = 64 | 0.9945 | 0.9625 | **0.9940** |
|
||||
|
||||
So at equal recall the quantised index answers **1.63x as many queries per
|
||||
second**, builds **1.8x faster**, and holds a quarter of the vectors. (Compare
|
||||
only at equal `ef`: with re-scoring the harness raises `ef` to at least the
|
||||
candidate pool, so the `ef = 16` and `ef = 32` rows are not like-for-like.)
|
||||
|
||||
#### On ARM (Raspberry Pi 5, Cortex-A76)
|
||||
|
||||
`dot_i8` has two aarch64 kernels: `SDOT` for CPUs with the ARMv8.2
|
||||
dot-product extension (Cortex-A76 and later, Neoverse-N1, all Apple Silicon)
|
||||
and plain NEON (`vmull_s8` + `vpadalq_s16`) otherwise. Medians of three runs
|
||||
at N = 100 000, ef = 64, recall@10 0.9940 in every int8 row against f32's
|
||||
0.9945:
|
||||
|
||||
| int8 kernel | build | QPS | vs f32 |
|
||||
|---|---:|---:|---:|
|
||||
| *(f32 baseline)* | 33 413 ms | 6 164 | 1.00x |
|
||||
| scalar (what v2.7.0 shipped) | 18 950 ms | ~6 190 | 1.00x |
|
||||
| plain NEON | ~17 000 ms | 6 640 | 1.08x |
|
||||
| **SDOT** | **14 464 ms** | **7 267** | **1.18x** |
|
||||
|
||||
These are Pi 5 numbers, not "ARM" numbers: a Pi has far less memory bandwidth
|
||||
and cache than an Apple M-series or a flagship phone, so the ratios will move
|
||||
on other hardware. The plain-NEON row is that code on an A76 with `SDOT`
|
||||
disabled, not a measurement of a pre-A76 core.
|
||||
|
||||
**A correction.** Until this was measured, this section said aarch64 "falls
|
||||
back to the scalar loop, where the original trade still applies" — that is,
|
||||
that quantised search was ~13% slower than f32 on ARM. That was extrapolated
|
||||
from x86 and it was wrong. On x86-64 the portable baseline is SSE2 while the
|
||||
f32 kernels are hand-written AVX2, so scalar int8 lost; on aarch64 NEON *is*
|
||||
the baseline, the compiler vectorises the scalar loop well, and scalar int8
|
||||
already matched f32 for search while building 1.76x faster.
|
||||
|
||||
So on every configuration measured — x86-64 AVX2, and Pi 5 with each of the
|
||||
three int8 kernels — the quantised index is at least as fast as f32 at equal
|
||||
recall, builds faster, and holds a quarter of the vectors.
|
||||
|
||||
A measurement trap worth recording: the synthetic `clustered` generator in the
|
||||
`clawhdf5-ann` tests draws clusters far tighter than any real embedding, so
|
||||
neighbours there sit closer together than the quantisation error and top-10
|
||||
*identity* is noise. Scored on that fixture int8 looks catastrophic (0.57
|
||||
overlap) — a fact about the fixture, not the storage. The tests use random
|
||||
vectors, and recall is measured against brute-force ground truth rather than
|
||||
against the f32 index, whose own approximation errors a re-scored search is
|
||||
entitled to get right.
|
||||
|
||||
### Search options: source filters, re-ranking, confidence
|
||||
|
||||
Measured 2026-09-24 on tank (AMD Ryzen 7 7800X3D). `HDF5Memory::search`
|
||||
with `SearchOptions`, clustered 384-dim data, k = 10, Hebbian boosting off.
|
||||
Filters keep 50%, 10% or 1% of the store at random, or two whole clusters
|
||||
chosen *away* from each query's own — the case an ANN index handles worst,
|
||||
because nothing it finds near the query is allowed. Recall is vector-only
|
||||
against an exact scan of the allowed records; latency is full hybrid search
|
||||
(vector + BM25). 200 queries, medians of three runs (recall was identical in
|
||||
every run).
|
||||
|
||||
```bash
|
||||
cargo run --release -p clawhdf5-bench --bin search_harness -- --options-study --full
|
||||
```
|
||||
|
||||
| N | options | filtered recall@10 | p50 ms | p99 ms |
|
||||
|---:|---|---:|---:|---:|
|
||||
| 10 000 | no filter | 1.0000 | 0.488 | 0.518 |
|
||||
| 10 000 | random 50% | 1.0000 | 0.520 | 0.563 |
|
||||
| 10 000 | random 10% | 1.0000 | 0.342 | 0.367 |
|
||||
| 10 000 | random 1% | 1.0000 | 0.220 | 0.239 |
|
||||
| 10 000 | 2 clusters away from the query | 1.0000 | 0.252 | 0.278 |
|
||||
| 10 000 | re-rank | — | 0.562 | 0.674 |
|
||||
| 10 000 | re-rank + confidence | — | 0.554 | 0.575 |
|
||||
| 100 000 | no filter | 0.9995 | 4.600 | 5.245 |
|
||||
| 100 000 | random 50% | 1.0000 | 4.770 | 5.614 |
|
||||
| 100 000 | random 10% | 1.0000 | 3.244 | 4.008 |
|
||||
| 100 000 | random 1% | 1.0000 | 2.288 | 2.943 |
|
||||
| 100 000 | 2 clusters away from the query | 1.0000 | 2.298 | 3.048 |
|
||||
| 100 000 | re-rank | — | 4.747 | 5.411 |
|
||||
| 100 000 | re-rank + confidence | — | 4.748 | 5.526 |
|
||||
|
||||
A filtered search finds the exact filtered top 10 and is never slower than an
|
||||
unfiltered one. The filter applies before ranking — over-fetching the index in
|
||||
proportion to what the filter removes, and scanning the allowed records
|
||||
exactly whenever that costs fewer distance evaluations than the index would
|
||||
(roughly `pool × M`). The first version compared the over-fetch to the store
|
||||
size instead, and that measured badly at 100K: the 1% filter took 5.9 ms at
|
||||
recall 0.9965 and the away-from-query filter 12.3 ms at 0.976, both through an
|
||||
index asked for ~16 000 candidates, where scanning the few hundred or thousand
|
||||
allowed records is exact and cheap. Re-ranking a 3k candidate pool and
|
||||
confidence rejection add about 3%.
|
||||
|
||||
### float16 embedding storage (`MemoryConfig::float16`)
|
||||
|
||||
Measured 2026-09-23 on tank (AMD Ryzen 7 7800X3D). The same clustered
|
||||
384-dim data in an `f32` store and a `float16` store, both with the default
|
||||
int8 index and Hebbian boosting off (so every query sees the same store).
|
||||
Recall is vector-only `hybrid_search` against an exact scan of the original
|
||||
`f32` vectors, 200 queries. Six runs, three with each store going first;
|
||||
medians. Nothing depended on the order.
|
||||
|
||||
```bash
|
||||
cargo run --release -p clawhdf5-bench --bin search_harness -- --float16-study --full
|
||||
cargo run --release -p clawhdf5-bench --bin search_harness -- --float16-study --full --f16-first
|
||||
```
|
||||
|
||||
| N | embeddings | file MiB | checkpoint ms | open ms | recall@10 | top-10 overlap | hybrid p50 ms |
|
||||
|---:|---|---:|---:|---:|---:|---:|---:|
|
||||
| 1 000 | f32 | 1.6 | 8 | 1.6 | 1.0000 | | 0.072 |
|
||||
| 1 000 | float16 | 0.8 | 6 | 1.9 | 0.9980 | 0.9980 | 0.072 |
|
||||
| 10 000 | f32 | 15.4 | 66 | 15.2 | 1.0000 | | 0.495 |
|
||||
| 10 000 | float16 | 8.2 | 45 | 18.1 | 1.0000 | 1.0000 | 0.494 |
|
||||
| 100 000 | f32 | 154.0 | 752 | 299.8 | 0.9940 | | 4.676 |
|
||||
| 100 000 | float16 | **80.8** | **512** | **252.1** | 0.9990 | 0.9940 | 4.654 |
|
||||
|
||||
The file is 48% smaller, checkpoints write less and open reads less. At
|
||||
small N opening is slightly slower (widening halves costs more than the I/O it
|
||||
saves: +3 ms at 10K). Recall does not move: half precision keeps about three
|
||||
significant digits, far finer than the gaps between neighbours on unit-length
|
||||
embeddings. Recall was identical in every run; the 0.999 against 0.994 at
|
||||
100K is two slightly different HNSW graphs, not an improvement to claim.
|
||||
|
||||
The in-memory cache holds the half-rounded values, so the store searches the
|
||||
same before and after a reopen; RAM use is unchanged (the cache is still
|
||||
`f32`). What `float16` saves is disk, and the I/O that goes with it.
|
||||
|
||||
### Opening a store (`read_from_disk`)
|
||||
|
||||
`HDF5Memory::open` memory-mapped the file, copied the whole mapping into a
|
||||
`Vec`, and handed that to `File::from_bytes` — while `File::open` memory-maps
|
||||
the file itself. Dropping the copy takes **store open from 455 ms to 327 ms**
|
||||
at 100 000 x 384 (`--e2e-only --full`; two runs after the change, 326.8 and
|
||||
328.1 ms).
|
||||
|
||||
It does **not** lower the process's peak memory, which is worth stating
|
||||
precisely because it is the obvious thing to assume. The harness now reports a
|
||||
high-water mark alongside the retained figure:
|
||||
|
||||
| N | reopened MiB | peak during open MiB |
|
||||
|---:|---:|---:|
|
||||
| 1 000 | 4 | 5 |
|
||||
| 10 000 | 44 | 61 |
|
||||
| 100 000 | 399 | 562 |
|
||||
|
||||
The peak is set *after* the parse, by the index build, so a buffer allocated
|
||||
and freed during the parse never reaches the high-water mark. Holding a
|
||||
deliberate extra copy of the file across the whole parse leaves the peak
|
||||
unmoved, which is how this was confirmed rather than assumed. What the change
|
||||
saves is the copy itself: a full-file memcpy on every open, and the transient
|
||||
that goes with it.
|
||||
|
||||
## Read harness
|
||||
|
||||
Produced by `cargo run --release -p clawhdf5-bench --bin read_harness`: a 4096 x
|
||||
@@ -653,6 +876,46 @@ add. There is no case here for changing the default; `TokenFilter::Stemmed`
|
||||
is available via `HDF5Memory::set_token_filter` for callers who want Hit@5/@10
|
||||
over rank-1 precision.
|
||||
|
||||
### Re-ranking and recency — full haystack, n=500
|
||||
|
||||
`reranker::rerank` combines temporal decay, source authority and Hebbian
|
||||
activation. Until now its combined score contained **no relevance term at
|
||||
all** — `RerankInput` did not carry the retrieval score — so a caller that
|
||||
re-ranked its candidates threw the retriever's ordering away and returned them
|
||||
ordered by age. The OpenClaw backend did exactly that on every search.
|
||||
|
||||
Measuring that is unambiguous. "Recency" below is the share of
|
||||
`knowledge-update` questions where the newest gold session outranked the stale
|
||||
one (see `newest_gold_first`); ~45% is chance.
|
||||
|
||||
| Mode | Hit@1 | Hit@5 | Hit@10 | MRR | recency |
|
||||
|---|---|---|---|---|---|
|
||||
| Hybrid 0.4/0.6, no re-rank | 51.6% | **81.4%** | 87.8% | 0.6430 | 45.0% |
|
||||
| + re-rank, **metadata only** (pre-fix) | 11.0% | 24.8% | 43.8% | 0.1829 | **87.5%** |
|
||||
| + re-rank, relevance-led, half-life 1 day | **52.0%** | 79.8% | 87.8% | 0.6403 | 51.7% |
|
||||
| + re-rank, relevance-led, half-life 7 days | 51.8% | 80.8% | 87.6% | **0.6437** | **52.2%** |
|
||||
| + re-rank, relevance-led, half-life 30 days | 51.8% | 81.0% | 87.8% | 0.6427 | 51.4% |
|
||||
| + re-rank, relevance-led, half-life 90 days | **52.0%** | 80.4% | 87.8% | 0.6425 | 50.8% |
|
||||
|
||||
**The pre-fix row is the finding.** Ordering candidates by recency alone costs
|
||||
40.6pp of Hit@1 and two thirds of MRR: the results are the newest memories in
|
||||
the pool rather than the ones that answer the question. It does ace the recency
|
||||
metric, which is exactly what makes that metric worth having — a number that
|
||||
only goes up when a change is good would not have caught this.
|
||||
|
||||
With relevance leading, retrieval is preserved (Hit@1 +0.4pp, MRR −0.003
|
||||
against no re-ranking) and recency discrimination gains 6–7pp. That is a real
|
||||
improvement but not a solved problem: recency only breaks near-ties, so it
|
||||
cannot reach the 87.5% the degenerate ordering gets. Those two rows are the
|
||||
ends of a trade-off, and the default sits deliberately near the relevance end.
|
||||
|
||||
**Half-life is not a sensitive knob.** Across 1, 7, 30 and 90 days recency
|
||||
moves 1.4pp and MRR 0.003 — inside the noise of a 500-question run — because
|
||||
the temporal term is capped by its weight (0.3) while relevance differences
|
||||
between candidates are larger. The 24-hour default is kept; there is no
|
||||
measured reason to change it, and a corpus-matched value is not the lever it
|
||||
looks like.
|
||||
|
||||
### Weight sweep — full haystack, n=500
|
||||
|
||||
`0.7/0.3` was a documented default, never a searched one. Sweeping
|
||||
@@ -999,6 +1262,69 @@ cargo run --release --bin ephemeral_perf
|
||||
|
||||
---
|
||||
|
||||
## Deflate backend: zlib-rs vs zlib-ng
|
||||
|
||||
Measured 2026-09-23 on tank (AMD Ryzen 7 7800X3D, 8C/16T). The default
|
||||
deflate backend is now **zlib-rs**, a pure-Rust port of zlib-ng; zlib-ng (C,
|
||||
built with cmake) was the default before and is still available as
|
||||
`fast-deflate`. Both builds were compiled once into separate target
|
||||
directories and run **alternately, three rounds each**; figures are medians.
|
||||
|
||||
```bash
|
||||
# zlib-rs (default)
|
||||
cargo bench -p clawhdf5-filters --bench deflate_bench
|
||||
cargo bench -p clawhdf5-bench --bench h5bench_write --features libhdf5-compare -- '^write_2d_chunked/'
|
||||
cargo run --release -p clawhdf5-bench --bin read_harness
|
||||
# zlib-ng: add --features fast-deflate (filters) or clawhdf5-format/fast-deflate (bench)
|
||||
```
|
||||
|
||||
| Workload | zlib-rs | zlib-ng | rs / ng |
|
||||
|---|---:|---:|---:|
|
||||
| HDF5 chunked write, deflate-6, 512×512 f32 | 1.458 ms | 1.484 ms | 0.98 |
|
||||
| HDF5 chunked write, deflate-6, 128×128 f32 | 157.6 µs | 152.8 µs | 1.03 |
|
||||
| HDF5 chunked write, deflate-6, 32×32 f32 | 62.9 µs | 60.2 µs | 1.05 |
|
||||
| HDF5 read, 64 MB chunked + deflate, full | 64.4 ms | 65.2 ms | 0.99 |
|
||||
| HDF5 read, 64×64 window (1 chunk) | 0.18 ms | 0.17 ms | 1.06 |
|
||||
| HDF5 read, 512×512 window (4–9 chunks) | 4.10 ms | 4.10 ms | 1.00 |
|
||||
| HDF5 read, one row / one column | 1.00 / 2.01 ms | 0.95 / 1.95 ms | 1.05 / 1.03 |
|
||||
| Raw inflate, 8 MB f64 | 5.92 ms | 6.06 ms | 0.98 |
|
||||
| Raw inflate, 1 MB sine | 82.8 µs | 68.6 µs | 1.21 |
|
||||
| Raw deflate-6, 8 MB f64 / 1 MB sine | 92.2 / 2.01 ms | 83.1 / 1.84 ms | 1.11 / 1.09 |
|
||||
|
||||
Compressed output is **byte-identical** between the two at levels 1, 6 and 9
|
||||
on all three inputs, so files do not change size. libhdf5 1.14.6 took 51.4 ms
|
||||
for the 512×512 write in the same session (35× the zlib-rs figure).
|
||||
|
||||
On the HDF5 paths zlib-rs is within 6% of zlib-ng everywhere, and ahead on the
|
||||
largest write. The raw codec loops show zlib-ng still slightly faster at
|
||||
compression (~10%), which chunked writes do not expose because encoding runs
|
||||
in parallel across chunks.
|
||||
|
||||
**Two findings along the way.** The first measurement had zlib-rs 1.2–1.9×
|
||||
slower on single-chunk reads and 3.7× slower on a 1 MB inflate — slower even
|
||||
than miniz_oxide. Neither was zlib-rs's fault:
|
||||
|
||||
1. **Runtime CPU detection was off.** zlib-rs needs its `std` feature to
|
||||
detect and use SIMD at runtime; flate2 turns it on through its default
|
||||
`runtime_detection` feature, which our `default-features = false` flate2
|
||||
dependency was disabling. With it, a 1 MB inflate goes 282 → 83 µs.
|
||||
`clawhdf5-format/zlib-rs` and `clawhdf5-filters/zlib-rs` now enable it.
|
||||
2. **The codec was fed through a 32 KiB buffer.** Both deflate paths used
|
||||
flate2's streaming `read::ZlibDecoder` / `write::ZlibEncoder`. A chunk's
|
||||
decompressed size is known, so they now hand the codec the whole input in
|
||||
one call, into an output buffer sized up front. Worth ~5% on chunked
|
||||
writes and ~10% on zlib-ng's 1 MB inflate. It also closed a hole: the
|
||||
streaming reader returned a truncated stream's bytes without an error, so
|
||||
a truncated chunk read back short; it is now an error.
|
||||
|
||||
| 1 MB inflate, same build otherwise | zlib-rs | zlib-ng |
|
||||
|---|---:|---:|
|
||||
| streaming reader, no runtime detection | 284.1 µs | 76.7 µs |
|
||||
| one-shot, no runtime detection | 282.3 µs | 68.5 µs |
|
||||
| one-shot + runtime detection (shipped) | **82.8 µs** | **68.6 µs** |
|
||||
|
||||
---
|
||||
|
||||
## h5bench-Equivalent I/O Benchmarks
|
||||
|
||||
Criterion harness mirroring h5bench serial workloads. clawhdf5 benchmarks dated 2026-07-01;
|
||||
|
||||
+356
@@ -1,5 +1,361 @@
|
||||
# Changelog
|
||||
|
||||
## Unreleased
|
||||
|
||||
### Upgrade Notes
|
||||
- **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`.
|
||||
- **`MemoryConfig::float16` now does what it says.** It was persisted and
|
||||
otherwise ignored; embeddings were always stored as `f32`. A store created
|
||||
with it on now writes half-precision embeddings (48% smaller files) and
|
||||
rounds embeddings to half precision as they are saved. A store that already
|
||||
had `float16 = true` rounds its embeddings when next opened and writes them
|
||||
as `float16` at its next checkpoint. Off by default.
|
||||
- **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.
|
||||
|
||||
### 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.
|
||||
|
||||
### 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.
|
||||
|
||||
### 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. 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`: **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::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-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
|
||||
|
||||
@@ -19,7 +19,7 @@ Cargo workspace with 16 crates under `crates/` (plus `libaec-sys`, an internal F
|
||||
| `clawhdf5-agent` | Agent memory, session history, knowledge graph storage |
|
||||
| `clawhdf5-gpu` | GPU-accelerated I/O via wgpu (hand-written WGSL compute shaders) |
|
||||
| `clawhdf5-accel` | CPU SIMD acceleration path |
|
||||
| `clawhdf5-migrate` | Schema migration engine |
|
||||
| `clawhdf5-migrate` | SQLite → HDF5 agent-memory migration |
|
||||
| `clawhdf5-android` | Android JNI bindings |
|
||||
| `clawhdf5-cli` | Command-line interface |
|
||||
| `clawhdf5-napi` | Node.js native addon bindings |
|
||||
@@ -27,7 +27,12 @@ Cargo workspace with 16 crates under `crates/` (plus `libaec-sys`, an internal F
|
||||
| `clawhdf5-bench` | Benchmark suite |
|
||||
|
||||
## 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
|
||||
`clawhdf5-agent` `hnsw` feature is **on by default**, so `hybrid_search` uses
|
||||
the approximate `clawhdf5-ann` index for the vector stage (the index mirrors
|
||||
@@ -39,7 +44,20 @@ Cargo workspace with 16 crates under `crates/` (plus `libaec-sys`, an internal F
|
||||
(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()`
|
||||
(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
|
||||
activation boosts are persisted by the next checkpoint (or on drop), not per
|
||||
query. Measure any search-path change with
|
||||
@@ -69,8 +87,26 @@ 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
|
||||
`<store>.h5.wal.corrupt-<ts>` rather than blocking `open()`; a WAL with an
|
||||
unknown *newer* version still fails and is left untouched.
|
||||
- `MemoryConfig::compression` uses deflate by default; enable the agent's
|
||||
`zstd` feature to compress embeddings with Zstd instead (links libzstd).
|
||||
- `MemoryConfig::float16` (off by default, persisted; CLI `create --float16`)
|
||||
writes `/memory/embeddings` as IEEE half precision (48% smaller file at
|
||||
100K, same recall). `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; the OpenClaw
|
||||
backend is `search` with re-rank + confidence on. 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).
|
||||
- `Dataset::verify_provenance()` (clawhdf5 facade, `provenance` feature, on by
|
||||
default) recomputes a dataset's SHA-256 and compares it against the
|
||||
`_provenance_sha256` attribute written automatically on save when
|
||||
@@ -103,6 +139,24 @@ cargo build --release
|
||||
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
|
||||
```bash
|
||||
cargo run -p clawhdf5-cli -- --help
|
||||
|
||||
+4
-1
@@ -21,8 +21,11 @@ members = [
|
||||
resolver = "2"
|
||||
|
||||
[workspace.package]
|
||||
version = "2.5.0"
|
||||
version = "2.7.0"
|
||||
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"
|
||||
repository = "https://git.redclaw.dev/quantumclaw/clawhdf5"
|
||||
|
||||
|
||||
@@ -3,24 +3,101 @@
|
||||
**The memory layer AI agents deserve. One file. Pure Rust. Zero C dependencies.**
|
||||
|
||||
[](LICENSE)
|
||||
[](https://www.rust-lang.org)
|
||||
[](#performance)
|
||||
[](BENCHMARKS.md#longmemeval-results)
|
||||
[](BENCHMARKS.md#memory-footprint)
|
||||
[](https://www.rust-lang.org)
|
||||
[](#building)
|
||||
[](BENCHMARKS.md#longmemeval-results)
|
||||
[](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, integrity-checked memory — all stored in a single portable file.
|
||||
|
||||
> **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.
|
||||
> - **An agent memory layer built on top of it** — vector search, knowledge graph, hippocampal-style consolidation, in `clawhdf5-agent`.
|
||||
|
||||
```
|
||||
cargo add clawhdf5 # core HDF5 read/write, no agent layer
|
||||
cargo add clawhdf5-agent --features agent # + agent memory layer
|
||||
The crates are not on crates.io yet, so depend on them from git:
|
||||
|
||||
```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)**
|
||||
|
||||
## 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 OpenClaw 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): 48% smaller files at the same recall.
|
||||
- `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 OpenClaw-only.
|
||||
|
||||
**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?
|
||||
@@ -35,7 +112,7 @@ Every AI agent needs memory. Today that means scattered Markdown files, SQLite d
|
||||
| Memory consolidation | Manual pruning | Hippocampal-inspired automatic tiers |
|
||||
| Temporal queries | Custom code | Native temporal index (716ns) |
|
||||
| Multi-modal | Multiple stores | Unified cross-modal search |
|
||||
| Security | Hope for the best | Provenance tracking + anomaly detection |
|
||||
| Integrity | Hope for the best | Chained-CRC WAL, checksummed chunk indexes, write-anomaly alerts, opt-in SHA-256 dataset provenance |
|
||||
| Portability | Config + DB + files | **One `.h5` file. Copy it anywhere.** |
|
||||
|
||||
---
|
||||
@@ -58,8 +135,28 @@ 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 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
|
||||
|
||||
**HNSW (the default backend for `hybrid_search`)** — `search_harness`, clustered
|
||||
384-dim data, M = 16, ef_construction = 64, recall measured against an exact scan.
|
||||
See [BENCHMARKS.md § Search harness](BENCHMARKS.md#search-harness-baseline-v230)
|
||||
and [§ Quantising the index copy](BENCHMARKS.md#quantising-the-index-copy-quantized_index):
|
||||
|
||||
| N = 100K, ef = 64 | recall@10 | QPS | build |
|
||||
|---|---:|---:|---:|
|
||||
| `f32` index | 0.9945 | 13 399 | 3.2 s |
|
||||
| `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.
|
||||
|
||||
**Brute-force and IVF paths** (Criterion, i7-12650H):
|
||||
|
||||
| Scale | Flat | IVF (nprobe=10) | IVF-PQ | vs MemX¹ |
|
||||
|-------|------|-----------------|--------|----------|
|
||||
| 1K | **54 µs** | — | — | — |
|
||||
@@ -75,7 +172,7 @@ Figures below are from an independent reproduction run on a second machine (AMD
|
||||
|
||||
| Operation | Latency | Scale |
|
||||
|-----------|---------|-------|
|
||||
| Hybrid search (RRF) | **222 µs** | 1K records |
|
||||
| Hybrid search (`HDF5Memory::hybrid_search`, p50) | **70 µs** / 0.49 ms / 4.65 ms | 1K / 10K / 100K records |
|
||||
| BM25 keyword search | **67 µs** | 1K records |
|
||||
| Knowledge graph BFS | **24 µs** | 1K entities |
|
||||
| Spreading activation | **17 µs** | 100 entities |
|
||||
@@ -115,13 +212,17 @@ declaration:
|
||||
|
||||
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
|
||||
0.0 to 1.0 found the long-standing `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`, or `0.3/0.7` if rank-1 precision matters most. See
|
||||
[BENCHMARKS.md § Weight sweep](BENCHMARKS.md#longmemeval-results).
|
||||
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. Since
|
||||
v2.5.0 `0.4/0.6` is the default (`hybrid::DEFAULT_FUSION`, used by
|
||||
`unified_search`, `hybrid_search_with` and the OpenClaw backend); 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
|
||||
inert and only the BM25 row is produced, which is what every previously published
|
||||
The benchmark's vector stage requires `clawhdf5-bench`'s `embeddings` feature
|
||||
(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.
|
||||
|
||||
On the easier `longmemeval_oracle` variant (evidence sessions only) the same
|
||||
@@ -146,19 +247,40 @@ retrieval recall reported as QA accuracy typically overstates by 20–30 points.
|
||||
|
||||
### Memory Footprint
|
||||
|
||||
| Records | File Size | Bytes/Record | With Compression |
|
||||
|---------|-----------|--------------|------------------|
|
||||
| 1K | ~6.5 MB | ~6.5 KB | ~2.1 MB (3.1x) |
|
||||
| 10K | ~65 MB | ~6.5 KB | ~21 MB (3.1x) |
|
||||
| 100K | ~645 MB | ~6.5 KB | ~208 MB (3.1x) |
|
||||
**On disk** — 384-dim embeddings, 200-char text
|
||||
([BENCHMARKS.md § Memory Footprint](BENCHMARKS.md#memory-footprint-1)):
|
||||
|
||||
| Records | File Size | Bytes/Record | Gzip-6 compressed |
|
||||
|---------|-----------|--------------|-------------------|
|
||||
| 1K | 1.7 MB | 1.8 KB | 277 KB (6.1x) |
|
||||
| 10K | 17.0 MB | 1.7 KB | 2.7 MB (6.2x) |
|
||||
| 100K | 169.8 MB | 1.7 KB | 26.9 MB (6.2x) |
|
||||
|
||||
With `MemoryConfig::float16` the embeddings take half the space: an agent
|
||||
store of 100K × 384 records is 80.8 MiB instead of 154.0.
|
||||
|
||||
**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.
|
||||
|
||||
### Consolidation Efficiency
|
||||
|
||||
1,000 records (10 signal + 990 noise), `working_capacity = 100`
|
||||
([BENCHMARKS.md § Consolidation Efficiency](BENCHMARKS.md#consolidation-efficiency)):
|
||||
|
||||
| Metric | Before | After | Delta |
|
||||
|--------|--------|-------|-------|
|
||||
| Records in store | 1,000 | ~110 | −89% |
|
||||
| Hit@1 recall | ~60% | ~90% | +30% |
|
||||
| Search latency | ~2.8 ms | ~0.3 ms | **9x faster** |
|
||||
| Records in store | 1,000 | 100 | −90% |
|
||||
| Hit@1 recall (signal records) | 100% | 100% | no loss |
|
||||
| Search latency | 2.75 ms | 0.31 ms | **8.8x faster** |
|
||||
|
||||
**Full benchmark details: [BENCHMARKS.md](BENCHMARKS.md)**
|
||||
|
||||
@@ -166,74 +288,73 @@ retrieval recall reported as QA accuracy typically overstates by 20–30 points.
|
||||
|
||||
## 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 │
|
||||
└────────┬────────┘
|
||||
│
|
||||
┌────────────▼────────────┐
|
||||
│ Hybrid Retrieval │
|
||||
│ Vector + BM25 + RRF │
|
||||
└────────────┬────────────┘
|
||||
│
|
||||
┌──────────────────▼──────────────────┐
|
||||
│ Multi-Factor Re-Ranking │
|
||||
│ temporal · authority · activation │
|
||||
└──────────────────┬──────────────────┘
|
||||
│
|
||||
┌────────────▼────────────┐
|
||||
│ Confidence Rejection │
|
||||
│ (suppress bad matches) │
|
||||
└────────────┬────────────┘
|
||||
│
|
||||
┌────────────────────────▼────────────────────────┐
|
||||
│ Memory Store (HDF5) │
|
||||
│ │
|
||||
│ ┌───────────┐ ┌───────────┐ ┌───────────────┐ │
|
||||
│ │ Working │→│ Episodic │→│ Semantic │ │
|
||||
│ │ (bounded) │ │ (bounded) │ │ (long-term) │ │
|
||||
│ └───────────┘ └───────────┘ └───────────────┘ │
|
||||
│ │
|
||||
│ ┌──────────┐ ┌──────────┐ ┌────────────────┐ │
|
||||
│ │Knowledge │ │Temporal │ │ Multi-Modal │ │
|
||||
│ │ Graph │ │ Index │ │ Embeddings │ │
|
||||
│ └──────────┘ └──────────┘ └────────────────┘ │
|
||||
│ │
|
||||
│ ┌──────────┐ ┌──────────┐ ┌────────────────┐ │
|
||||
│ │Provenance│ │ Anomaly │ │ Source │ │
|
||||
│ │ Tracking │ │Detection │ │ Isolation │ │
|
||||
│ └──────────┘ └──────────┘ └────────────────┘ │
|
||||
└─────────────────────────────────────────────────┘
|
||||
│
|
||||
┌────────┴────────┐
|
||||
│ agent_memory.h5 │
|
||||
│ single file │
|
||||
└─────────────────┘
|
||||
┌─────────────────┐
|
||||
│ Agent Query │
|
||||
└────────┬────────┘
|
||||
│
|
||||
┌─────────────────▼──────────────────┐
|
||||
│ HDF5Memory::search │
|
||||
│ optional source-channel filter │
|
||||
│ HNSW vector + BM25 keyword │
|
||||
│ weighted fusion (0.4 / 0.6) │
|
||||
│ × √(Hebbian activation) │
|
||||
└─────────────────┬──────────────────┘
|
||||
│ opt-in (SearchOptions);
|
||||
│ the OpenClaw backend turns both on
|
||||
┌─────────────────▼──────────────────┐
|
||||
│ Multi-factor re-ranking │
|
||||
│ relevance · recency · authority · │
|
||||
│ activation │
|
||||
├────────────────────────────────────┤
|
||||
│ Confidence rejection │
|
||||
│ (suppress bad matches) │
|
||||
└─────────────────┬──────────────────┘
|
||||
│
|
||||
┌────────────────────────────▼────────────────────────────┐
|
||||
│ In memory │
|
||||
│ cache (flat f32 embeddings) · BM25 index · HNSW index │
|
||||
│ provenance ledger + anomaly alerts (session-scoped) │
|
||||
└────────────────────────────┬────────────────────────────┘
|
||||
│ WAL append; checkpoint
|
||||
┌────────────────────────────▼────────────────────────────┐
|
||||
│ 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) │
|
||||
│ agent_memory.h5.lock single-writer lock │
|
||||
└─────────────────────────────────────────────────────────┘
|
||||
```
|
||||
|
||||
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 | What It Does |
|
||||
|--------|-------------|
|
||||
| **`knowledge`** | Entity/relation graph with BFS traversal, spreading activation, fuzzy entity resolution |
|
||||
| **`consolidation`** | Three-tier memory (Working → Episodic → Semantic) with importance scoring 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 |
|
||||
| **`reranker`** | Multi-factor re-ranking: temporal recency, source authority, activation weight |
|
||||
| **`confidence`** | Low-confidence rejection — suppresses spurious recalls when nothing matches |
|
||||
| **`knowledge`** | Entity/relation graph with BFS traversal, spreading activation, fuzzy (Levenshtein) entity resolution |
|
||||
| **`consolidation`** | Three-tier memory (Working → Episodic → Semantic) with importance scoring, novelty, and time-decay |
|
||||
| **`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: retrieval relevance (leads, weight 1.0), temporal recency, source authority, activation weight. Opt-in via `SearchOptions::with_rerank`; on in the OpenClaw backend |
|
||||
| **`confidence`** | Low-confidence rejection — suppresses spurious recalls when nothing matches. Opt-in via `SearchOptions::with_confidence`; on in the OpenClaw backend |
|
||||
| **`temporal`** | Sorted timestamp index, session DAG, entity timeline, temporal query hints |
|
||||
| **`multimodal`** | Cross-modal search across text/image/audio/video embeddings |
|
||||
| **`provenance`** | Source attribution, FNV-1a content hashing, integrity verification |
|
||||
| **`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) |
|
||||
| **`anomaly`** | Write rate limiting, 15 injection-pattern detectors, source-distribution analysis. Alerts never block a save; drain them with `take_anomaly_alerts` |
|
||||
| **`openclaw`** | OpenClaw integration: MemoryBackend trait, Markdown ↔ HDF5 conversion |
|
||||
| **`vector_search`** | Flat cosine, pre-normed, SIMD, BLAS, GPU, parallel search paths |
|
||||
| **`ivf` / `pq`** | IVF-PQ approximate nearest neighbor for billion-scale search |
|
||||
| **`bm25`** | BM25 keyword index with TF-IDF scoring |
|
||||
| **`ivf` / `pq`** | Standalone IVF and IVF-PQ indexes (benchmarked to 100K vectors); not used by `HDF5Memory`, whose ANN index is HNSW |
|
||||
| **`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 |
|
||||
| **`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 |
|
||||
| **`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) |
|
||||
|
||||
---
|
||||
@@ -265,7 +386,7 @@ assert_eq!(values, vec![22.5, 23.1, 21.8]);
|
||||
use clawhdf5_agent::{HDF5Memory, MemoryConfig, MemoryEntry, AgentMemory};
|
||||
|
||||
// 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)?;
|
||||
|
||||
// Save a memory
|
||||
@@ -278,13 +399,38 @@ memory.save(MemoryEntry {
|
||||
tags: "preference".into(),
|
||||
})?;
|
||||
|
||||
// Search
|
||||
let results = memory.search(&query_embedding, 5)?;
|
||||
// Hybrid search: vector + BM25, weighted 0.4 / 0.6 (the measured default)
|
||||
let results = memory.hybrid_search(&query_embedding, "user preferences", 0.4, 0.6, 5);
|
||||
for result in results {
|
||||
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 the OpenClaw backend runs.
|
||||
let careful = memory.search(
|
||||
&query_embedding,
|
||||
"user preferences",
|
||||
&SearchOptions::new(5)
|
||||
.with_rerank(ReRankConfig::default())
|
||||
.with_confidence(ConfidenceConfig::default()),
|
||||
);
|
||||
```
|
||||
|
||||
### Knowledge Graph
|
||||
|
||||
```rust
|
||||
@@ -309,8 +455,8 @@ let neighbors = kg.bfs_neighbors(alice, 2); // 2-hop neighborhood
|
||||
let activated = kg.spreading_activation(&[alice], 0.5, 0.01, 5);
|
||||
|
||||
// Entity resolution — fuzzy matching
|
||||
let resolved = kg.resolve_or_create("alice", "person", -1, 2);
|
||||
// Returns existing Alice entity (Levenshtein distance ≤ 2)
|
||||
let (id, created) = kg.resolve_or_create("alice", "person", -1, 2);
|
||||
// id == alice, created == false: matched the existing entity (Levenshtein distance ≤ 2)
|
||||
```
|
||||
|
||||
### Memory Consolidation
|
||||
@@ -321,15 +467,19 @@ use clawhdf5_agent::consolidation::*;
|
||||
let config = ConsolidationConfig::default();
|
||||
let mut engine = ConsolidationEngine::new(config);
|
||||
|
||||
// Add memories — automatically scored for importance
|
||||
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);
|
||||
let now = 1_700_000_000.0; // seconds since the epoch
|
||||
|
||||
// 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)
|
||||
engine.access_memory(0);
|
||||
engine.access_memory(id, now);
|
||||
|
||||
// Run consolidation cycle
|
||||
let stats = engine.consolidate();
|
||||
engine.consolidate(now);
|
||||
let stats = engine.get_stats();
|
||||
// Working memories promote to Episodic (if important enough)
|
||||
// Episodic memories promote to Semantic (if accessed enough)
|
||||
// Low-decay memories get evicted when tiers are full
|
||||
@@ -357,13 +507,13 @@ let recent = index.latest(10);
|
||||
use clawhdf5_agent::openclaw::*;
|
||||
|
||||
// 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
|
||||
let md = std::fs::read_to_string("MEMORY.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);
|
||||
|
||||
// Export back to Markdown
|
||||
@@ -375,22 +525,23 @@ let exported = backend.export_markdown("MEMORY.md")?;
|
||||
## Crate Map
|
||||
|
||||
```
|
||||
clawhdf5 workspace (16 crates, ~92K lines of Rust; plus libaec-sys, an
|
||||
internal FFI bindings crate for the optional szip feature)
|
||||
clawhdf5 workspace (16 crates, ~86K lines of Rust in src/, ~104K with tests
|
||||
and benches; plus libaec-sys, an internal FFI bindings
|
||||
crate for the optional szip feature)
|
||||
│
|
||||
├── Core HDF5
|
||||
│ ├── clawhdf5-format — Binary parser/writer (no_std), shared type definitions
|
||||
│ ├── clawhdf5-io — I/O abstraction (buffered, mmap, async)
|
||||
│ ├── clawhdf5-format — Binary parser/writer (no_std-capable), shared type definitions
|
||||
│ ├── 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-derive — Proc macros
|
||||
│ ├── clawhdf5 — High-level API
|
||||
│ ├── 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)
|
||||
│
|
||||
├── Agent Memory
|
||||
│ ├── clawhdf5-agent — Memory engine (20.9K lines, 32 modules; WAL is CRC32-checked per entry)
|
||||
│ ├── clawhdf5-ann — HNSW approximate nearest neighbor (default backend; optional `parallel` feature)
|
||||
│ ├── clawhdf5-agent — Memory engine (24.7K lines, 32 modules; chained-CRC WAL)
|
||||
│ ├── clawhdf5-ann — HNSW approximate nearest neighbor (default backend; f32 or int8 storage; `parallel` build)
|
||||
│ ├── clawhdf5-migrate — SQLite → HDF5 migration
|
||||
│ ├── clawhdf5-android — Android JNI bridge
|
||||
│ └── clawhdf5-cli — CLI tool
|
||||
@@ -411,10 +562,10 @@ ClawhDF5's agent memory design draws from 15+ recent papers:
|
||||
|
||||
| Paper | Key Insight | ClawhDF5 Module |
|
||||
|-------|-------------|-----------------|
|
||||
| **MemX** (2026) | RRF + multi-factor re-ranking | `hybrid`, `reranker` |
|
||||
| **Graph-Native Cognitive Memory** (2026) | Graph-structured belief revision | `knowledge` |
|
||||
| **MemX** (2026) | Hybrid fusion + multi-factor re-ranking | `hybrid`, `reranker` |
|
||||
| **Graph-Native Cognitive Memory** (2026) | Graph-structured memory (weighted, timestamped relations; entity timelines) | `knowledge`, `temporal` |
|
||||
| **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` |
|
||||
| **RAGdb** (2025) | Zero-dependency edge RAG | Architecture |
|
||||
| **MemoryGraft** (2025) | Memory poisoning attacks | `anomaly`, `provenance` |
|
||||
@@ -429,15 +580,43 @@ ClawhDF5's agent memory design draws from 15+ recent papers:
|
||||
|
||||
| Flag | Default | Description |
|
||||
|------|---------|-------------|
|
||||
| `agent` | no | Full agent memory layer |
|
||||
| `float16` | **yes** | Half-precision embedding storage (2× compression) |
|
||||
| `float16` | **yes** | Half-precision cosine kernel (`cosine_similarity_f16`). Half-precision *storage* is the `MemoryConfig::float16` setting below, and needs no feature |
|
||||
| `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 |
|
||||
| `accelerate` | no | Apple Accelerate / AMX (macOS) |
|
||||
| `openblas` | no | OpenBLAS (Linux) |
|
||||
| `gpu` | no | GPU search via wgpu |
|
||||
| `async` | no | Tokio async with background flush |
|
||||
| `agent` | no | Reserved; currently enables nothing (the agent layer is always built) |
|
||||
|
||||
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` (off by default; CLI `create --float16`) 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 vector
|
||||
recall and search latency do not change. Embeddings are rounded as they are
|
||||
saved, so the store searches the same before and after a reopen; values must
|
||||
lie within ±65504. See
|
||||
[BENCHMARKS.md § float16 embedding storage](BENCHMARKS.md#float16-embedding-storage-memoryconfigfloat16).
|
||||
|
||||
### `clawhdf5-format`
|
||||
|
||||
@@ -447,26 +626,31 @@ ClawhDF5's agent memory design draws from 15+ recent papers:
|
||||
| `deflate` | yes | Deflate compression |
|
||||
| `checksum` | yes | Jenkins lookup3 verification |
|
||||
| `provenance` | yes | SHA-256 provenance attributes |
|
||||
| `fast-deflate` | **yes** | zlib-ng backend for faster deflate |
|
||||
| `system-zlib-decompress` | **yes** | Use the system zlib for decompression where available |
|
||||
| `zlib-rs` | **yes** | Pure-Rust deflate backend ([zlib-rs](https://github.com/trifectatechfoundation/zlib-rs)) |
|
||||
| `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) |
|
||||
| `fast-checksum` | no | crc32fast-accelerated checksums |
|
||||
| `lz4` | no | LZ4 block compression filter (id 32004) |
|
||||
| `zstd` | no | Zstandard compression filter (id 32015) |
|
||||
| `pcodec` | no | Pcodec lossless numerical codec (id 32023, via `pco` crate) |
|
||||
| `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 |
|
||||
| `szip` | no | SZIP filter (id 4) via libaec (C, through the internal `libaec-sys` crate) |
|
||||
|
||||
### `clawhdf5-ann`
|
||||
|
||||
| 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`
|
||||
|
||||
| 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 |
|
||||
|
||||
> **Parallel I/O (MPI) limitation:** `mpi-io`'s read path is a root-rank read
|
||||
@@ -480,18 +664,26 @@ ClawhDF5's agent memory design draws from 15+ recent papers:
|
||||
## Building
|
||||
|
||||
```bash
|
||||
# Default
|
||||
# Default (pure Rust: no cmake or C compiler needed)
|
||||
cargo build --workspace
|
||||
|
||||
# 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)
|
||||
cargo build -p clawhdf5-agent --features "agent,float16,accelerate,parallel,gpu"
|
||||
cargo build -p clawhdf5-agent --features "accelerate,gpu"
|
||||
|
||||
# Tests
|
||||
cargo test --workspace # all 1,650+ tests
|
||||
cargo test --workspace # all 1,850+ 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
|
||||
cargo bench -p clawhdf5-agent # agent memory suite
|
||||
@@ -504,25 +696,42 @@ cargo bench -p clawhdf5-bench # h5bench-equivalent I/O suite
|
||||
|
||||
```
|
||||
agent_memory.h5
|
||||
├── /meta
|
||||
│ ├── schema_version: "1.0"
|
||||
│ ├── agent_id, embedder, embedding_dim
|
||||
│ └── created_at
|
||||
├── /meta (attributes)
|
||||
│ ├── schema_version: "1.0", edgehdf5_version
|
||||
│ ├── agent_id, embedder, embedding_dim, chunk_size, overlap, 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
|
||||
│ ├── chunks: string[N]
|
||||
│ ├── embeddings: f32[N × D] (or f16 with float16 flag)
|
||||
│ ├── tombstones: u8[N]
|
||||
│ └── norms: f32[N] (pre-computed L2)
|
||||
│ ├── chunks: string[N]
|
||||
│ ├── embeddings: f32[N × D], or f16 for a `float16` store
|
||||
│ │ (chunked; deflate, or Zstd with the `zstd`
|
||||
│ │ 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
|
||||
│ ├── ids: string[S]
|
||||
│ └── summaries: string[S]
|
||||
│ ├── ids, channels, summaries: string[S]
|
||||
│ ├── start_idxs, end_idxs: i64[S]
|
||||
│ └── timestamps: f64[S]
|
||||
└── /knowledge_graph
|
||||
├── entity_names: string[E]
|
||||
├── relation_srcs: i64[R]
|
||||
├── relation_tgts: i64[R]
|
||||
└── relation_types: string[R]
|
||||
├── entity_ids, entity_emb_idxs: i64[E]; entity_names, entity_types: string[E]
|
||||
├── relation_srcs, relation_tgts: i64[R]; relation_types: string[R]
|
||||
├── relation_weights: f32[R]; relation_ts: f64[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
|
||||
@@ -577,6 +786,6 @@ MIT
|
||||
---
|
||||
|
||||
<p align="center">
|
||||
<em>Built by <a href="https://github.com/redclawsystems">RedClaw Systems</a></em><br>
|
||||
<em>~92,000 lines of Rust. Zero C dependencies. One file to remember everything.</em>
|
||||
<em>Built by <a href="https://git.redclaw.dev/quantumclaw">RedClaw Systems</a></em><br>
|
||||
<em>~86,000 lines of Rust. Zero C dependencies. One file to remember everything.</em>
|
||||
</p>
|
||||
|
||||
@@ -1,7 +1,8 @@
|
||||
[package]
|
||||
name = "clawhdf5-accel"
|
||||
version = "2.5.0"
|
||||
version = "2.7.0"
|
||||
edition = "2024"
|
||||
rust-version.workspace = true
|
||||
description = "SIMD-accelerated operations for rustyhdf5"
|
||||
license = "MIT"
|
||||
repository = "https://git.redclaw.dev/quantumclaw/clawhdf5"
|
||||
|
||||
@@ -25,6 +25,55 @@ unsafe fn hsum_256(v: __m256) -> f32 {
|
||||
_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.
|
||||
///
|
||||
/// # 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.
|
||||
pub fn vector_norm(v: &[f32]) -> f32 {
|
||||
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
|
||||
}
|
||||
|
||||
/// 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)
|
||||
}
|
||||
|
||||
/// 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]
|
||||
name = "clawhdf5-agent"
|
||||
version = "2.5.0"
|
||||
version = "2.7.0"
|
||||
edition = "2024"
|
||||
rust-version.workspace = true
|
||||
description = "HDF5-backed persistent memory store for on-device AI agents"
|
||||
license = "MIT"
|
||||
repository = "https://git.redclaw.dev/quantumclaw/clawhdf5"
|
||||
@@ -10,12 +11,12 @@ keywords = ["agent", "memory", "hdf5", "vector-search", "embedding"]
|
||||
categories = ["database", "science", "algorithms"]
|
||||
|
||||
[dependencies]
|
||||
clawhdf5-format = { path = "../clawhdf5-format", version = "2.5.0", features = ["parallel", "fast-checksum"] }
|
||||
clawhdf5 = { path = "../clawhdf5", version = "2.5.0" }
|
||||
clawhdf5-io = { path = "../clawhdf5-io", version = "2.5.0", features = ["mmap"] }
|
||||
clawhdf5-accel = { path = "../clawhdf5-accel", version = "2.5.0" }
|
||||
clawhdf5-ann = { path = "../clawhdf5-ann", version = "2.5.0", optional = true }
|
||||
clawhdf5-gpu = { path = "../clawhdf5-gpu", version = "2.5.0", optional = true, default-features = false }
|
||||
clawhdf5-format = { path = "../clawhdf5-format", version = "2.7.0", features = ["parallel", "fast-checksum"] }
|
||||
clawhdf5 = { path = "../clawhdf5", version = "2.7.0" }
|
||||
clawhdf5-io = { path = "../clawhdf5-io", version = "2.7.0", features = ["mmap"] }
|
||||
clawhdf5-accel = { path = "../clawhdf5-accel", version = "2.7.0" }
|
||||
clawhdf5-ann = { path = "../clawhdf5-ann", version = "2.7.0", optional = true }
|
||||
clawhdf5-gpu = { path = "../clawhdf5-gpu", version = "2.7.0", optional = true, default-features = false }
|
||||
serde = { workspace = true }
|
||||
byteorder = "1"
|
||||
half = { workspace = true, optional = true }
|
||||
|
||||
@@ -118,6 +118,10 @@ mod tests {
|
||||
created_at: "2025-01-01T00:00:00Z".to_string(),
|
||||
wal_enabled: false,
|
||||
wal_max_entries: 500,
|
||||
quantized_index: false,
|
||||
hnsw_m: 16,
|
||||
hnsw_ef_construction: 64,
|
||||
hnsw_ef_search: 0,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -88,8 +88,11 @@ impl BM25Index {
|
||||
/// Search the index for a query, returning the top `k` results
|
||||
/// as `(doc_id, score)` pairs sorted by score descending.
|
||||
///
|
||||
/// Uses Block-Max WAND for early termination when remaining documents
|
||||
/// cannot beat the current top-k threshold.
|
||||
/// Scores every matching document exhaustively, then keeps the top `k`.
|
||||
/// 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)> {
|
||||
if k == 0 {
|
||||
return Vec::new();
|
||||
@@ -561,8 +564,9 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn wand_returns_same_results_as_exhaustive() {
|
||||
// WAND-style search should produce same scores as exhaustive
|
||||
fn top_k_search_matches_ranking_every_score() {
|
||||
// `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)
|
||||
.map(|i| {
|
||||
if i % 3 == 0 {
|
||||
|
||||
@@ -1,17 +1,145 @@
|
||||
//! In-memory cache for memory entries, sessions, and knowledge graph.
|
||||
|
||||
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.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct MemoryCache {
|
||||
pub chunks: Vec<String>,
|
||||
pub embeddings: Vec<Vec<f32>>,
|
||||
/// `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 embeddings: Embeddings,
|
||||
pub source_channels: Vec<String>,
|
||||
pub timestamps: Vec<f64>,
|
||||
pub session_ids: Vec<String>,
|
||||
@@ -22,14 +150,18 @@ pub struct MemoryCache {
|
||||
pub norms: Vec<f32>,
|
||||
/// Hebbian activation weights (default 1.0 per entry).
|
||||
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 {
|
||||
pub fn new(embedding_dim: usize) -> Self {
|
||||
Self {
|
||||
chunks: Vec::new(),
|
||||
embeddings: Vec::new(),
|
||||
embeddings_flat: Vec::new(),
|
||||
embeddings: Embeddings::new(embedding_dim),
|
||||
source_channels: Vec::new(),
|
||||
timestamps: Vec::new(),
|
||||
session_ids: Vec::new(),
|
||||
@@ -38,18 +170,52 @@ impl MemoryCache {
|
||||
embedding_dim,
|
||||
norms: Vec::new(),
|
||||
activation_weights: Vec::new(),
|
||||
half_precision: false,
|
||||
}
|
||||
}
|
||||
|
||||
/// Rebuild `embeddings_flat` from `embeddings` from scratch. Callers that
|
||||
/// populate `embeddings` directly (bulk loads) must call this afterward.
|
||||
pub fn rebuild_flat(&mut self) {
|
||||
self.embeddings_flat.clear();
|
||||
self.embeddings_flat
|
||||
.reserve(self.embeddings.len() * self.embedding_dim);
|
||||
for emb in &self.embeddings {
|
||||
self.embeddings_flat.extend_from_slice(emb);
|
||||
/// Switch half-precision rounding on or off. Turning it on rounds every
|
||||
/// embedding already held (and recomputes norms where one changed) —
|
||||
/// e.g. a `float16` store whose last checkpoint predates half-precision
|
||||
/// storage and so is still `f32` on disk.
|
||||
pub fn set_half_precision(&mut self, on: bool) {
|
||||
self.half_precision = on;
|
||||
if !on {
|
||||
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).
|
||||
@@ -77,10 +243,10 @@ impl MemoryCache {
|
||||
tags: String,
|
||||
) -> usize {
|
||||
let idx = self.chunks.len();
|
||||
let embedding = self.stored_form(embedding);
|
||||
let norm = vector_search::compute_norm(&embedding);
|
||||
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.timestamps.push(timestamp);
|
||||
self.session_ids.push(session_id);
|
||||
@@ -116,22 +282,10 @@ impl MemoryCache {
|
||||
session_id: String,
|
||||
) {
|
||||
if idx < self.chunks.len() {
|
||||
let embedding = self.stored_form(embedding);
|
||||
let norm = vector_search::compute_norm(&embedding);
|
||||
self.chunks[idx] = chunk;
|
||||
let dim = self.embedding_dim;
|
||||
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.embeddings.set(idx, &embedding);
|
||||
self.source_channels[idx] = source_channel;
|
||||
self.timestamps[idx] = timestamp;
|
||||
self.session_ids[idx] = session_id;
|
||||
@@ -183,7 +337,7 @@ impl MemoryCache {
|
||||
new_idx += 1;
|
||||
let norm = vector_search::compute_norm(&self.embeddings[i]);
|
||||
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_timestamps.push(self.timestamps[i]);
|
||||
new_session_ids.push(self.session_ids[i].clone());
|
||||
@@ -196,7 +350,8 @@ impl MemoryCache {
|
||||
|
||||
let removed = old_len - new_chunks.len();
|
||||
self.chunks = new_chunks;
|
||||
self.embeddings = new_embeddings;
|
||||
self.embeddings
|
||||
.reset_from(self.embedding_dim, new_embeddings);
|
||||
self.source_channels = new_source_channels;
|
||||
self.timestamps = new_timestamps;
|
||||
self.session_ids = new_session_ids;
|
||||
@@ -204,16 +359,14 @@ impl MemoryCache {
|
||||
self.tombstones = new_tombstones;
|
||||
self.norms = new_norms;
|
||||
self.activation_weights = new_activation_weights;
|
||||
self.rebuild_flat();
|
||||
|
||||
(removed, index_map)
|
||||
}
|
||||
|
||||
/// Flatten all embeddings into a single Vec<f32> for HDF5 storage.
|
||||
/// `embeddings_flat` is already maintained incrementally, so this just
|
||||
/// clones it — kept as a method for callers that want an owned copy.
|
||||
pub fn flat_embeddings(&self) -> Vec<f32> {
|
||||
self.embeddings_flat.clone()
|
||||
/// All embeddings as one owned `[N x dim]` buffer, for HDF5 storage.
|
||||
/// Prefer [`MemoryCache::flat_embeddings`] where a borrow will do.
|
||||
pub fn flat_embeddings_owned(&self) -> Vec<f32> {
|
||||
self.embeddings.as_flat().to_vec()
|
||||
}
|
||||
}
|
||||
|
||||
@@ -224,7 +377,7 @@ mod tests {
|
||||
/// `embeddings_flat` must always equal a from-scratch flatten of `embeddings`.
|
||||
fn assert_flat_in_sync(cache: &MemoryCache) {
|
||||
let expected: Vec<f32> = cache.embeddings.iter().flatten().copied().collect();
|
||||
assert_eq!(cache.embeddings_flat, expected);
|
||||
assert_eq!(cache.embeddings.as_flat(), expected);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -247,7 +400,10 @@ mod tests {
|
||||
String::new(),
|
||||
);
|
||||
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]
|
||||
@@ -279,7 +435,7 @@ mod tests {
|
||||
);
|
||||
assert_flat_in_sync(&cache);
|
||||
assert_eq!(
|
||||
cache.embeddings_flat,
|
||||
cache.embeddings.as_flat(),
|
||||
vec![7.0, 8.0, 9.0, 4.0, 5.0, 6.0],
|
||||
"update must overwrite the correct flat slice, not just append"
|
||||
);
|
||||
@@ -315,14 +471,71 @@ mod tests {
|
||||
cache.mark_deleted(1);
|
||||
cache.compact();
|
||||
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]
|
||||
fn rebuild_flat_matches_manual_flatten() {
|
||||
let mut cache = MemoryCache::new(2);
|
||||
cache.embeddings = vec![vec![1.0, 2.0], vec![3.0, 4.0]];
|
||||
cache.rebuild_flat();
|
||||
assert_eq!(cache.embeddings_flat, vec![1.0, 2.0, 3.0, 4.0]);
|
||||
cache
|
||||
.embeddings
|
||||
.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));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -28,7 +28,7 @@ use crate::vector_search;
|
||||
pub fn hybrid_search(
|
||||
query_embedding: &[f32],
|
||||
query_text: &str,
|
||||
vectors: &[Vec<f32>],
|
||||
vectors: &(impl crate::vector_search::VectorSet + Sync + ?Sized),
|
||||
chunks: &[String],
|
||||
tombstones: &[u8],
|
||||
bm25_index: &BM25Index,
|
||||
@@ -56,7 +56,7 @@ pub fn hybrid_search(
|
||||
pub fn hybrid_search_fused(
|
||||
query_embedding: &[f32],
|
||||
query_text: &str,
|
||||
vectors: &[Vec<f32>],
|
||||
vectors: &(impl crate::vector_search::VectorSet + Sync + ?Sized),
|
||||
_chunks: &[String],
|
||||
tombstones: &[u8],
|
||||
bm25_index: &BM25Index,
|
||||
@@ -65,31 +65,34 @@ pub fn hybrid_search_fused(
|
||||
) -> Vec<(usize, f32)> {
|
||||
// Get raw scores from both systems. Request all results so normalization
|
||||
// covers the full distribution.
|
||||
// Use parallel search when rayon feature is enabled and vector count > 10K.
|
||||
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 vec_scores = exact_vector_scores(query_embedding, vectors, tombstones);
|
||||
let kw_scores = bm25_index.scores(query_text);
|
||||
|
||||
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.
|
||||
///
|
||||
/// Both score sets are independently min-max normalized to [0, 1] and combined
|
||||
@@ -270,7 +273,7 @@ fn normalize_scores(scores: &[(usize, f32)]) -> Vec<(usize, f32)> {
|
||||
pub fn rrf_hybrid_search(
|
||||
query_embedding: &[f32],
|
||||
query_text: &str,
|
||||
vectors: &[Vec<f32>],
|
||||
vectors: &(impl crate::vector_search::VectorSet + Sync + ?Sized),
|
||||
_chunks: &[String],
|
||||
tombstones: &[u8],
|
||||
bm25_index: &BM25Index,
|
||||
@@ -282,12 +285,12 @@ pub fn rrf_hybrid_search(
|
||||
let mut vec_scores = {
|
||||
#[cfg(feature = "parallel")]
|
||||
{
|
||||
if vectors.len() > 10_000 {
|
||||
if vectors.count() > 10_000 {
|
||||
vector_search::parallel_cosine_batch(
|
||||
query_embedding,
|
||||
vectors,
|
||||
tombstones,
|
||||
vectors.len(),
|
||||
vectors.count(),
|
||||
)
|
||||
} else {
|
||||
vector_search::cosine_similarity_batch(query_embedding, vectors, tombstones)
|
||||
@@ -298,7 +301,7 @@ pub fn rrf_hybrid_search(
|
||||
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.
|
||||
vec_scores.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
|
||||
|
||||
@@ -62,21 +62,17 @@ use std::path::{Path, PathBuf};
|
||||
|
||||
use cache::MemoryCache;
|
||||
#[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};
|
||||
|
||||
/// 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
|
||||
// the `ephemeral` module; they are not needed directly in lib.rs internals.
|
||||
#[allow(unused_imports)]
|
||||
pub use ephemeral::{EphemeralEntry, EphemeralStats};
|
||||
use knowledge::KnowledgeCache;
|
||||
use memory_strategy::{Exchange, MemoryStrategy, StrategyOutput};
|
||||
pub use search::SearchOptions;
|
||||
use session::SessionCache;
|
||||
|
||||
// --- Error type ---
|
||||
@@ -89,6 +85,9 @@ pub enum MemoryError {
|
||||
NotFound(String),
|
||||
/// Another `HDF5Memory` (in this or another process) has the store open.
|
||||
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),
|
||||
}
|
||||
|
||||
impl std::fmt::Display for MemoryError {
|
||||
@@ -99,6 +98,7 @@ impl std::fmt::Display for MemoryError {
|
||||
MemoryError::Schema(e) => write!(f, "schema error: {e}"),
|
||||
MemoryError::NotFound(e) => write!(f, "not found: {e}"),
|
||||
MemoryError::Locked(e) => write!(f, "store is locked: {e}"),
|
||||
MemoryError::InvalidEntry(e) => write!(f, "invalid entry: {e}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -130,6 +130,12 @@ pub struct MemoryConfig {
|
||||
pub embedding_dim: usize,
|
||||
pub chunk_size: 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`).
|
||||
pub float16: bool,
|
||||
pub compression: bool,
|
||||
pub compression_level: u32,
|
||||
@@ -139,6 +145,40 @@ pub struct MemoryConfig {
|
||||
pub created_at: String,
|
||||
pub wal_enabled: bool,
|
||||
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 {
|
||||
@@ -160,6 +200,10 @@ impl MemoryConfig {
|
||||
created_at,
|
||||
wal_enabled: true,
|
||||
wal_max_entries: 500,
|
||||
quantized_index: true,
|
||||
hnsw_m: 16,
|
||||
hnsw_ef_construction: 64,
|
||||
hnsw_ef_search: 0,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -285,7 +329,8 @@ impl HDF5Memory {
|
||||
/// Create a new HDF5 memory file with the given configuration.
|
||||
pub fn create(config: MemoryConfig) -> Result<Self> {
|
||||
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 knowledge = KnowledgeCache::new();
|
||||
|
||||
@@ -444,7 +489,17 @@ impl HDF5Memory {
|
||||
|
||||
#[cfg(feature = "hnsw")]
|
||||
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 {
|
||||
None
|
||||
};
|
||||
@@ -558,6 +613,7 @@ impl HDF5Memory {
|
||||
generation: Option<u64>,
|
||||
cache: &MemoryCache,
|
||||
n_checkpoint: usize,
|
||||
storage: Storage,
|
||||
) -> Option<HnswIndex> {
|
||||
let generation = generation?;
|
||||
let bytes = std::fs::read(Self::vector_index_path(store)).ok()?;
|
||||
@@ -565,15 +621,17 @@ impl HDF5Memory {
|
||||
if u64::from_le_bytes(stamp.try_into().ok()?) != generation {
|
||||
return None;
|
||||
}
|
||||
let vectors = cache.embeddings.get(..n_checkpoint)?.to_vec();
|
||||
let mut index = HnswIndex::from_graph_bytes(graph, vectors).ok()?;
|
||||
let vectors: Vec<Vec<f32>> = (0..n_checkpoint)
|
||||
.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 {
|
||||
return None;
|
||||
}
|
||||
// Records appended since (replayed from the WAL) join incrementally.
|
||||
for id in n_checkpoint..cache.embeddings.len() {
|
||||
if cache.embeddings[id].len() != index.dimension()
|
||||
|| index.insert(cache.embeddings[id].clone()) != id
|
||||
|| index.insert(cache.embeddings[id].to_vec()) != id
|
||||
{
|
||||
return None;
|
||||
}
|
||||
@@ -801,6 +859,42 @@ impl HDF5Memory {
|
||||
// the index length drifts from the cache length (covering any mutation path
|
||||
// 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
|
||||
/// soft-deletions so node ids stay aligned with cache indices.
|
||||
///
|
||||
@@ -816,11 +910,15 @@ impl HDF5Memory {
|
||||
if self.cache.embeddings.iter().any(|e| e.len() != dim) {
|
||||
return None;
|
||||
}
|
||||
let mut index = HnswIndex::build_with_metric(
|
||||
&self.cache.embeddings,
|
||||
HNSW_M,
|
||||
HNSW_EF_CONSTRUCTION,
|
||||
// The index owns its vectors, so it needs rows rather than the cache's
|
||||
// flat buffer. This copy is the index's own; the cache keeps one.
|
||||
let rows: Vec<Vec<f32>> = self.cache.embeddings.iter().map(<[f32]>::to_vec).collect();
|
||||
let mut index = HnswIndex::build_with(
|
||||
&rows,
|
||||
self.hnsw_m(),
|
||||
self.hnsw_ef_construction(),
|
||||
DistanceMetric::Cosine,
|
||||
self.index_storage(),
|
||||
);
|
||||
for (i, &t) in self.cache.tombstones.iter().enumerate() {
|
||||
if t != 0 {
|
||||
@@ -846,7 +944,7 @@ impl HDF5Memory {
|
||||
let dim = index.dimension();
|
||||
let appended = (self.hnsw_synced_len..n).all(|id| {
|
||||
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 {
|
||||
for id in self.hnsw_synced_len..n {
|
||||
@@ -877,7 +975,7 @@ impl HDF5Memory {
|
||||
let emb_len = self.cache.embeddings[idx].len();
|
||||
match self.hnsw.as_mut() {
|
||||
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 {
|
||||
self.hnsw_synced_len = self.cache.embeddings.len();
|
||||
} else {
|
||||
@@ -985,7 +1083,29 @@ impl HDF5Memory {
|
||||
/// 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
|
||||
/// 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> {
|
||||
self.check_embedding(&entry.embedding)?;
|
||||
if let Some(existing_idx) = self.cache.find_by_tags(&entry.tags) {
|
||||
if let Some(ref mut w) = self.wal {
|
||||
let wal_entry = wal::WalEntry {
|
||||
@@ -1041,6 +1161,7 @@ impl HDF5Memory {
|
||||
|
||||
impl AgentMemory for HDF5Memory {
|
||||
fn save(&mut self, entry: MemoryEntry) -> Result<usize> {
|
||||
self.check_embedding(&entry.embedding)?;
|
||||
if let Some(ref mut w) = self.wal {
|
||||
let wal_entry = wal::WalEntry {
|
||||
entry_type: wal::WalEntryType::Save,
|
||||
@@ -1082,6 +1203,10 @@ impl AgentMemory for HDF5Memory {
|
||||
}
|
||||
|
||||
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());
|
||||
for entry in entries {
|
||||
let idx = self.cache.push(
|
||||
@@ -1242,6 +1367,9 @@ impl HDF5Memory {
|
||||
})?;
|
||||
let view = memory_strategy::CacheStoreView::new(&self.cache, &self.knowledge);
|
||||
let output = strat.evaluate(&exchange, &view);
|
||||
for e in &output.entries {
|
||||
self.check_embedding(&e.embedding)?;
|
||||
}
|
||||
for e in &output.entries {
|
||||
self.cache.push(
|
||||
e.chunk.clone(),
|
||||
@@ -1326,6 +1454,16 @@ impl HDF5Memory {
|
||||
let mut promoted = 0;
|
||||
|
||||
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
|
||||
.ephemeral
|
||||
.as_mut()
|
||||
|
||||
@@ -13,9 +13,8 @@ use std::path::{Path, PathBuf};
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use crate::{
|
||||
AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry,
|
||||
confidence::{ConfidenceConfig, ScoredResult, reject_low_confidence},
|
||||
reranker::{ReRankConfig, RerankInput, rerank},
|
||||
AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry, SearchOptions,
|
||||
confidence::ConfidenceConfig, reranker::ReRankConfig,
|
||||
};
|
||||
|
||||
// ─────────────────────────────────────────────────────────────────────────────
|
||||
@@ -466,7 +465,7 @@ impl ClawhdfBackend {
|
||||
let record = MemoryRecord {
|
||||
id: i as u64,
|
||||
chunk: cache.chunks[i].clone(),
|
||||
embedding: cache.embeddings[i].clone(),
|
||||
embedding: cache.embeddings[i].to_vec(),
|
||||
tier: MemoryTier::Working,
|
||||
importance: cache.activation_weights[i],
|
||||
access_count: 0,
|
||||
@@ -524,70 +523,27 @@ impl ClawhdfBackend {
|
||||
|
||||
impl MemoryBackend for ClawhdfBackend {
|
||||
/// Search using hybrid vector + BM25 retrieval, then re-rank and
|
||||
/// confidence-filter.
|
||||
/// confidence-filter — [`HDF5Memory::search`] with both stages on.
|
||||
fn search(
|
||||
&mut self,
|
||||
query_text: &str,
|
||||
query_embedding: &[f32],
|
||||
k: usize,
|
||||
) -> Vec<MemorySearchResult> {
|
||||
// 1. Hybrid retrieval (vector + BM25, fused by score).
|
||||
let candidates = k.saturating_mul(3).max(10);
|
||||
let raw = self.memory.hybrid_search_with(
|
||||
query_embedding,
|
||||
query_text,
|
||||
crate::hybrid::DEFAULT_FUSION,
|
||||
candidates,
|
||||
);
|
||||
|
||||
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
|
||||
let options = SearchOptions::new(k)
|
||||
.with_rerank(self.rerank_config)
|
||||
.with_confidence(self.confidence_config.clone())
|
||||
.at_time(Self::now_secs());
|
||||
self.memory
|
||||
.search(query_embedding, query_text, &options)
|
||||
.into_iter()
|
||||
.take(k)
|
||||
.filter_map(|sr| {
|
||||
let r = raw_by_idx.get(&sr.index)?;
|
||||
let path = r.source_channel.clone();
|
||||
Some(MemorySearchResult {
|
||||
text: r.chunk.clone(),
|
||||
score: sr.score,
|
||||
path: path.clone(),
|
||||
line_range: None,
|
||||
timestamp: Some(r.timestamp),
|
||||
source: path,
|
||||
})
|
||||
.map(|r| MemorySearchResult {
|
||||
text: r.chunk,
|
||||
score: r.score,
|
||||
path: r.source_channel.clone(),
|
||||
line_range: None,
|
||||
timestamp: Some(r.timestamp),
|
||||
source: r.source_channel,
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
@@ -716,11 +672,13 @@ impl MemoryBackend for ClawhdfBackend {
|
||||
|
||||
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
|
||||
.embeddings
|
||||
.norms
|
||||
.iter()
|
||||
.enumerate()
|
||||
.filter(|(i, emb)| cache.tombstones[*i] == 0 && !emb.is_empty())
|
||||
.filter(|(i, norm)| cache.tombstones[*i] == 0 && **norm > 0.0)
|
||||
.count();
|
||||
|
||||
let file_size_bytes = std::fs::metadata(&self.hdf5_path)
|
||||
|
||||
@@ -4,8 +4,10 @@
|
||||
//! into a single composite score for each retrieved result.
|
||||
|
||||
/// Configuration for the multi-factor re-ranker.
|
||||
#[derive(Debug, Clone)]
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
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).
|
||||
pub temporal_weight: f32,
|
||||
/// Weight applied to the source authority score (0.0–1.0).
|
||||
@@ -20,6 +22,9 @@ pub struct ReRankConfig {
|
||||
impl Default for ReRankConfig {
|
||||
fn default() -> 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,
|
||||
authority_weight: 0.2,
|
||||
activation_weight: 0.5,
|
||||
@@ -41,6 +46,8 @@ pub struct ReRankResult {
|
||||
pub authority_score: f32,
|
||||
/// Normalised Hebbian activation score in [0, 1].
|
||||
pub activation_score: f32,
|
||||
/// The retrieval score carried through from the input.
|
||||
pub relevance_score: f32,
|
||||
}
|
||||
|
||||
/// Compute an exponential decay temporal score.
|
||||
@@ -105,6 +112,15 @@ pub struct RerankInput {
|
||||
pub source_channel: String,
|
||||
/// Raw Hebbian activation weight for this entry.
|
||||
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.
|
||||
@@ -138,7 +154,8 @@ pub fn rerank(
|
||||
let auth = source_authority_score(&inp.source_channel);
|
||||
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.activation_weight * act;
|
||||
|
||||
@@ -148,6 +165,7 @@ pub fn rerank(
|
||||
temporal_score: ts,
|
||||
authority_score: auth,
|
||||
activation_score: act,
|
||||
relevance_score: inp.relevance,
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
@@ -253,22 +271,51 @@ mod tests {
|
||||
timestamp: 0.0, // very old
|
||||
source_channel: "other".to_string(),
|
||||
raw_activation: 0.1,
|
||||
relevance: 0.0,
|
||||
},
|
||||
RerankInput {
|
||||
index: 1,
|
||||
timestamp: 86_400.0, // one day ago
|
||||
source_channel: "conversation".to_string(),
|
||||
raw_activation: 0.5,
|
||||
relevance: 0.0,
|
||||
},
|
||||
RerankInput {
|
||||
index: 2,
|
||||
timestamp: 172_800.0, // "now"
|
||||
source_channel: "user_correction".to_string(),
|
||||
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]
|
||||
fn rerank_returns_all_entries() {
|
||||
let inputs = make_inputs();
|
||||
@@ -302,6 +349,7 @@ mod tests {
|
||||
#[test]
|
||||
fn rerank_score_breakdown_matches_manual_calculation() {
|
||||
let config = ReRankConfig {
|
||||
relevance_weight: 0.0,
|
||||
temporal_weight: 1.0,
|
||||
authority_weight: 0.0,
|
||||
activation_weight: 0.0,
|
||||
@@ -312,6 +360,7 @@ mod tests {
|
||||
timestamp: 0.0,
|
||||
source_channel: "other".to_string(),
|
||||
raw_activation: 0.5,
|
||||
relevance: 0.0,
|
||||
}];
|
||||
let now = 3600.0_f64; // exactly one half-life later
|
||||
let results = rerank(&inputs, &config, now);
|
||||
|
||||
@@ -104,6 +104,19 @@ pub fn build_hdf5_file_with_meta(
|
||||
"wal_max_entries",
|
||||
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(
|
||||
"edgehdf5_version",
|
||||
AttrValue::String(ZEROCLAW_VERSION.into()),
|
||||
@@ -146,20 +159,27 @@ fn build_memory_group(
|
||||
// chunks: fixed-length string array
|
||||
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 d = cache.embedding_dim as u64;
|
||||
let flat = cache.flat_embeddings();
|
||||
{
|
||||
let ds = group
|
||||
.create_dataset("embeddings")
|
||||
.with_f32_data(&flat)
|
||||
.with_shape(&[n, d]);
|
||||
let ds = group.create_dataset("embeddings");
|
||||
let elem_bytes: u64 = if config.float16 {
|
||||
ds.with_f16_data(flat);
|
||||
2
|
||||
} else {
|
||||
ds.with_f32_data(flat);
|
||||
4
|
||||
};
|
||||
ds.with_shape(&[n, d]);
|
||||
|
||||
// Chunk size tuning: target ~256KB per chunk for optimal I/O
|
||||
if n > 0 && d > 0 {
|
||||
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]);
|
||||
|
||||
// Compression. Shuffle is applied automatically (auto-shuffle
|
||||
@@ -484,10 +504,32 @@ pub fn validate_and_load(
|
||||
wal_max_entries: optional_i64_attr(&attrs, "wal_max_entries")
|
||||
.and_then(|v| usize::try_from(v).ok())
|
||||
.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
|
||||
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
|
||||
let session_cache = load_sessions_group(file)?;
|
||||
@@ -563,12 +605,7 @@ fn load_memory_group(
|
||||
.collect(),
|
||||
};
|
||||
|
||||
// Unflatten embeddings
|
||||
let embeddings: Vec<Vec<f32>> = flat_embeddings
|
||||
.chunks(embedding_dim)
|
||||
.map(|c| c.to_vec())
|
||||
.collect();
|
||||
|
||||
// No unflattening: the cache stores the buffer as it is on disk.
|
||||
// Read activation_weights if present, default to vec![1.0; N] for backward compat
|
||||
let activation_weights = match read_f32_dataset(&group, "activation_weights") {
|
||||
Ok(w) if w.len() == n => w,
|
||||
@@ -576,7 +613,7 @@ fn load_memory_group(
|
||||
};
|
||||
|
||||
cache.chunks = chunks;
|
||||
cache.embeddings = embeddings;
|
||||
cache.embeddings.set_flat(embedding_dim, flat_embeddings);
|
||||
cache.source_channels = source_channels;
|
||||
cache.timestamps = timestamps;
|
||||
cache.session_ids = session_ids;
|
||||
@@ -584,7 +621,6 @@ fn load_memory_group(
|
||||
cache.tombstones = tombstones;
|
||||
cache.norms = norms;
|
||||
cache.activation_weights = activation_weights;
|
||||
cache.rebuild_flat();
|
||||
|
||||
Ok(cache)
|
||||
}
|
||||
@@ -736,6 +772,13 @@ fn read_string_dataset_from_group(
|
||||
.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> {
|
||||
let ds = group
|
||||
.dataset(name)
|
||||
|
||||
@@ -2,18 +2,107 @@
|
||||
|
||||
use std::path::Path;
|
||||
|
||||
use std::collections::HashSet;
|
||||
|
||||
use crate::bm25;
|
||||
use crate::confidence::{ConfidenceConfig, ScoredResult, reject_low_confidence};
|
||||
use crate::hybrid;
|
||||
use crate::reranker::{ReRankConfig, RerankInput, rerank};
|
||||
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 {
|
||||
/// 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
|
||||
/// previous behaviour, also used as the correctness oracle in tests). With
|
||||
/// `hnsw` enabled and an index available, the vector candidates come from an
|
||||
/// approximate-nearest-neighbour search over an over-fetched pool, then merge
|
||||
/// 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")]
|
||||
fn vector_keyword_search(
|
||||
&mut self,
|
||||
@@ -22,24 +111,103 @@ impl HDF5Memory {
|
||||
bm25: &bm25::BM25Index,
|
||||
fusion: hybrid::Fusion,
|
||||
k: usize,
|
||||
exclude: Option<&[u8]>,
|
||||
) -> Vec<(usize, f32)> {
|
||||
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() {
|
||||
Some(index) if !index.is_empty() && index.dimension() == query_embedding.len() => {
|
||||
// Over-fetch so the merge sees a useful vector pool; cosine
|
||||
// distance from the index converts back to similarity (1 - d).
|
||||
let pool = (k * 8).max(64);
|
||||
let vec_scores: Vec<(usize, f32)> = index
|
||||
.search(query_embedding, pool, pool)
|
||||
let ef = self.hnsw_ef_search(k).max(pool);
|
||||
let candidates = index.search(query_embedding, pool, ef);
|
||||
// A quantised index returns approximate distances, and no
|
||||
// amount of `ef` fixes that — the loss is in the distances,
|
||||
// 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()
|
||||
.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();
|
||||
// Fusion normalises over every keyword match, so it needs all
|
||||
// the scores — but not ranked.
|
||||
let kw_scores = bm25.scores(query_text);
|
||||
let mut kw_scores = bm25.scores(query_text);
|
||||
if let Some(ex) = exclude {
|
||||
if vec_scores.len() < k.min(allowed) {
|
||||
// The allowed records are not where the index looked.
|
||||
return self.exact_masked_search(
|
||||
query_embedding,
|
||||
query_text,
|
||||
bm25,
|
||||
fusion,
|
||||
k,
|
||||
ex,
|
||||
);
|
||||
}
|
||||
kw_scores.retain(|(id, _)| ex[*id] == 0);
|
||||
}
|
||||
hybrid::fuse(vec_scores, kw_scores, fusion, k)
|
||||
}
|
||||
_ => hybrid::hybrid_search_fused(
|
||||
_ => match exclude {
|
||||
Some(ex) => {
|
||||
self.exact_masked_search(query_embedding, query_text, bm25, fusion, k, ex)
|
||||
}
|
||||
None => hybrid::hybrid_search_fused(
|
||||
query_embedding,
|
||||
query_text,
|
||||
&self.cache.embeddings,
|
||||
&self.cache.chunks,
|
||||
&self.cache.tombstones,
|
||||
bm25,
|
||||
fusion,
|
||||
k,
|
||||
),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(not(feature = "hnsw"))]
|
||||
fn vector_keyword_search(
|
||||
&mut self,
|
||||
query_embedding: &[f32],
|
||||
query_text: &str,
|
||||
bm25: &bm25::BM25Index,
|
||||
fusion: hybrid::Fusion,
|
||||
k: usize,
|
||||
exclude: Option<&[u8]>,
|
||||
) -> Vec<(usize, f32)> {
|
||||
match exclude {
|
||||
Some(ex) => self.exact_masked_search(query_embedding, query_text, bm25, fusion, k, ex),
|
||||
None => hybrid::hybrid_search_fused(
|
||||
query_embedding,
|
||||
query_text,
|
||||
&self.cache.embeddings,
|
||||
@@ -52,25 +220,33 @@ impl HDF5Memory {
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(not(feature = "hnsw"))]
|
||||
fn vector_keyword_search(
|
||||
&mut self,
|
||||
/// 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)> {
|
||||
hybrid::hybrid_search_fused(
|
||||
query_embedding,
|
||||
query_text,
|
||||
&self.cache.embeddings,
|
||||
&self.cache.chunks,
|
||||
&self.cache.tombstones,
|
||||
bm25,
|
||||
fusion,
|
||||
k,
|
||||
)
|
||||
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.
|
||||
@@ -105,12 +281,53 @@ impl HDF5Memory {
|
||||
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
|
||||
// incrementally. Take it out for the duration of the call so the
|
||||
// vector stage can borrow `self` mutably, then put it back.
|
||||
self.ensure_bm25_fresh();
|
||||
let bm25 = self.bm25.take().expect("ensure_bm25_fresh leaves an index");
|
||||
let scored = self.vector_keyword_search(query_embedding, query_text, &bm25, fusion, k);
|
||||
let scored = self.vector_keyword_search(
|
||||
query_embedding,
|
||||
query_text,
|
||||
&bm25,
|
||||
options.fusion,
|
||||
fetch,
|
||||
exclude.as_deref(),
|
||||
);
|
||||
self.bm25 = Some(bm25);
|
||||
|
||||
let mut results: Vec<SearchResult> = scored
|
||||
.into_iter()
|
||||
.map(|(idx, score)| {
|
||||
@@ -134,6 +351,25 @@ impl HDF5Memory {
|
||||
.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`
|
||||
// records are relevant, the rest of the list is zero-score filler;
|
||||
// boosting it would teach the store that arbitrary records are
|
||||
@@ -144,11 +380,45 @@ impl HDF5Memory {
|
||||
.map(|r| r.index)
|
||||
.collect();
|
||||
self.apply_hebbian_boost(&hit_indices);
|
||||
self.bm25 = Some(bm25);
|
||||
|
||||
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
|
||||
/// the next checkpoint (any write that flushes, `flush_wal`, or drop) — not
|
||||
/// by rewriting the whole store inside the query, which is what made
|
||||
|
||||
@@ -113,13 +113,11 @@ pub type StoreState = (MemoryConfig, MemoryCache, SessionCache, KnowledgeCache);
|
||||
/// [`read_from_disk`], plus the checkpoint's [`WalMark`] (if any) so the
|
||||
/// caller can skip WAL entries this file already contains.
|
||||
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)?;
|
||||
|
||||
// Advise the OS we'll need the whole file for parsing
|
||||
mmap.advise_willneed(0, mmap.len());
|
||||
|
||||
// Parse the HDF5 file from the mmap'd bytes
|
||||
let file = clawhdf5::File::from_bytes(mmap.as_bytes().to_vec())
|
||||
// `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()`
|
||||
// did the same work and then copied the whole store — a second full copy
|
||||
// of the file, live for the whole parse, on top of the mapping.
|
||||
let file = clawhdf5::File::open(path)
|
||||
.map_err(|e| MemoryError::Hdf5(format!("cannot open {}: {e}", path.display())))?;
|
||||
|
||||
let (mut config, cache, sessions, knowledge) = schema::validate_and_load(&file)?;
|
||||
@@ -133,9 +131,7 @@ pub fn read_from_disk_with_mark(path: &Path) -> Result<(StoreState, Option<WalMa
|
||||
pub fn read_from_disk_with_meta(
|
||||
path: &Path,
|
||||
) -> Result<(StoreState, schema::CheckpointMeta), MemoryError> {
|
||||
let mmap = clawhdf5_io::MmapReader::open(path).map_err(MemoryError::Io)?;
|
||||
mmap.advise_willneed(0, mmap.len());
|
||||
let file = clawhdf5::File::from_bytes(mmap.as_bytes().to_vec())
|
||||
let file = clawhdf5::File::open(path)
|
||||
.map_err(|e| MemoryError::Hdf5(format!("cannot open {}: {e}", path.display())))?;
|
||||
let (mut config, cache, sessions, knowledge) = schema::validate_and_load(&file)?;
|
||||
config.path = path.to_path_buf();
|
||||
|
||||
@@ -4,6 +4,44 @@
|
||||
//! `clawhdf5_accel`, with optional float16 support via the `half` crate.
|
||||
//! 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.
|
||||
///
|
||||
/// 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.
|
||||
pub fn cosine_similarity_batch(
|
||||
query: &[f32],
|
||||
vectors: &[Vec<f32>],
|
||||
vectors: &(impl VectorSet + ?Sized),
|
||||
tombstones: &[u8],
|
||||
) -> Vec<(usize, f32)> {
|
||||
let query_norm = clawhdf5_accel::vector_norm(query);
|
||||
@@ -30,7 +68,7 @@ pub fn cosine_similarity_batch(
|
||||
return Vec::new();
|
||||
}
|
||||
|
||||
let n = vectors.len();
|
||||
let n = vectors.count();
|
||||
let mut results: Vec<(usize, f32)> = Vec::with_capacity(n);
|
||||
|
||||
// Process 4 vectors at a time where possible
|
||||
@@ -42,8 +80,9 @@ pub fn cosine_similarity_batch(
|
||||
if i < tombstones.len() && tombstones[i] != 0 {
|
||||
continue;
|
||||
}
|
||||
let vec_norm = clawhdf5_accel::vector_norm(&vectors[i]);
|
||||
let score = crate::cosine_similarity_prenorm(query, query_norm, &vectors[i], vec_norm);
|
||||
let vec_norm = clawhdf5_accel::vector_norm(vectors.row(i));
|
||||
let score =
|
||||
crate::cosine_similarity_prenorm(query, query_norm, vectors.row(i), vec_norm);
|
||||
results.push((i, score));
|
||||
}
|
||||
}
|
||||
@@ -53,8 +92,8 @@ pub fn cosine_similarity_batch(
|
||||
if i < tombstones.len() && tombstones[i] != 0 {
|
||||
continue;
|
||||
}
|
||||
let vec_norm = clawhdf5_accel::vector_norm(&vectors[i]);
|
||||
let score = crate::cosine_similarity_prenorm(query, query_norm, &vectors[i], vec_norm);
|
||||
let vec_norm = clawhdf5_accel::vector_norm(vectors.row(i));
|
||||
let score = crate::cosine_similarity_prenorm(query, query_norm, vectors.row(i), vec_norm);
|
||||
results.push((i, score));
|
||||
}
|
||||
|
||||
@@ -68,7 +107,7 @@ pub fn cosine_similarity_batch(
|
||||
/// collections. Uses `score = dot(query, vec) / (query_norm * stored_norm)`.
|
||||
pub fn cosine_similarity_batch_prenorm(
|
||||
query: &[f32],
|
||||
vectors: &[Vec<f32>],
|
||||
vectors: &(impl VectorSet + ?Sized),
|
||||
norms: &[f32],
|
||||
tombstones: &[u8],
|
||||
) -> Vec<(usize, f32)> {
|
||||
@@ -77,7 +116,7 @@ pub fn cosine_similarity_batch_prenorm(
|
||||
return Vec::new();
|
||||
}
|
||||
|
||||
let n = vectors.len();
|
||||
let n = vectors.count();
|
||||
let mut results: Vec<(usize, f32)> = Vec::with_capacity(n);
|
||||
|
||||
for i in 0..n {
|
||||
@@ -85,7 +124,7 @@ pub fn cosine_similarity_batch_prenorm(
|
||||
continue;
|
||||
}
|
||||
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));
|
||||
}
|
||||
|
||||
@@ -162,7 +201,7 @@ pub fn cosine_similarity_f16(
|
||||
#[cfg(feature = "parallel")]
|
||||
pub fn parallel_cosine_batch(
|
||||
query: &[f32],
|
||||
vectors: &[Vec<f32>],
|
||||
vectors: &(impl VectorSet + Sync + ?Sized),
|
||||
tombstones: &[u8],
|
||||
k: usize,
|
||||
) -> Vec<(usize, f32)> {
|
||||
@@ -174,24 +213,27 @@ pub fn parallel_cosine_batch(
|
||||
}
|
||||
|
||||
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 {
|
||||
return Vec::new();
|
||||
}
|
||||
|
||||
let mut all_results: Vec<(usize, f32)> = vectors
|
||||
.par_chunks(chunk_size)
|
||||
.enumerate()
|
||||
.flat_map(|(chunk_idx, chunk)| {
|
||||
// Chunk over index ranges: the corpus may be one flat buffer rather than
|
||||
// a slice of rows, so there is nothing to `par_chunks` over.
|
||||
let n = vectors.count();
|
||||
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 mut local: Vec<(usize, f32)> = Vec::with_capacity(chunk.len());
|
||||
for (j, vec) in chunk.iter().enumerate() {
|
||||
let i = base + j;
|
||||
let end = (base + chunk_size).min(n);
|
||||
let mut local: Vec<(usize, f32)> = Vec::with_capacity(end - base);
|
||||
for i in base..end {
|
||||
if i < tombstones.len() && tombstones[i] != 0 {
|
||||
continue;
|
||||
}
|
||||
let vec_norm = clawhdf5_accel::vector_norm(vec);
|
||||
let score = crate::cosine_similarity_prenorm(query, query_norm, vec, vec_norm);
|
||||
let vec_norm = clawhdf5_accel::vector_norm(vectors.row(i));
|
||||
let score =
|
||||
crate::cosine_similarity_prenorm(query, query_norm, vectors.row(i), vec_norm);
|
||||
local.push((i, score));
|
||||
}
|
||||
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")]
|
||||
pub fn parallel_cosine_batch_prenorm(
|
||||
query: &[f32],
|
||||
vectors: &[Vec<f32>],
|
||||
vectors: &(impl VectorSet + Sync + ?Sized),
|
||||
norms: &[f32],
|
||||
tombstones: &[u8],
|
||||
k: usize,
|
||||
@@ -222,23 +264,26 @@ pub fn parallel_cosine_batch_prenorm(
|
||||
}
|
||||
|
||||
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 {
|
||||
return Vec::new();
|
||||
}
|
||||
|
||||
let mut all_results: Vec<(usize, f32)> = vectors
|
||||
.par_chunks(chunk_size)
|
||||
.enumerate()
|
||||
.flat_map(|(chunk_idx, chunk)| {
|
||||
// Chunk over index ranges: the corpus may be one flat buffer rather than
|
||||
// a slice of rows, so there is nothing to `par_chunks` over.
|
||||
let n = vectors.count();
|
||||
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 mut local: Vec<(usize, f32)> = Vec::with_capacity(chunk.len());
|
||||
for (j, vec) in chunk.iter().enumerate() {
|
||||
let i = base + j;
|
||||
let end = (base + chunk_size).min(n);
|
||||
let mut local: Vec<(usize, f32)> = Vec::with_capacity(end - base);
|
||||
for i in base..end {
|
||||
if i < tombstones.len() && tombstones[i] != 0 {
|
||||
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.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
|
||||
|
||||
Binary file not shown.
@@ -0,0 +1,208 @@
|
||||
//! `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);
|
||||
}
|
||||
@@ -0,0 +1,94 @@
|
||||
//! 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");
|
||||
}
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
@@ -1,7 +1,8 @@
|
||||
[package]
|
||||
name = "clawhdf5-android"
|
||||
version = "2.5.0"
|
||||
version = "2.7.0"
|
||||
edition = "2024"
|
||||
rust-version.workspace = true
|
||||
description = "Android JNI bridge for edgehdf5-memory HDF5 backend"
|
||||
license = "MIT"
|
||||
|
||||
|
||||
@@ -1,7 +1,8 @@
|
||||
[package]
|
||||
name = "clawhdf5-ann"
|
||||
version = "2.5.0"
|
||||
version = "2.7.0"
|
||||
edition = "2024"
|
||||
rust-version.workspace = true
|
||||
description = "HNSW approximate nearest neighbor index stored as HDF5"
|
||||
license = "MIT"
|
||||
repository = "https://git.redclaw.dev/quantumclaw/clawhdf5"
|
||||
@@ -10,9 +11,9 @@ keywords = ["hdf5", "ann", "hnsw", "nearest-neighbor"]
|
||||
categories = ["algorithms", "science"]
|
||||
|
||||
[dependencies]
|
||||
clawhdf5-format = { path = "../clawhdf5-format", version = "2.5.0" }
|
||||
clawhdf5-io = { path = "../clawhdf5-io", version = "2.5.0" }
|
||||
clawhdf5-accel = { path = "../clawhdf5-accel", version = "2.5.0" }
|
||||
clawhdf5-format = { path = "../clawhdf5-format", version = "2.7.0" }
|
||||
clawhdf5-io = { path = "../clawhdf5-io", version = "2.7.0" }
|
||||
clawhdf5-accel = { path = "../clawhdf5-accel", version = "2.7.0" }
|
||||
rayon = { version = "1", optional = true }
|
||||
|
||||
[features]
|
||||
|
||||
+412
-62
@@ -154,6 +154,218 @@ impl Ord for FarCandidate {
|
||||
}
|
||||
}
|
||||
|
||||
/// How the index keeps its copy of the vectors.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
|
||||
pub enum Storage {
|
||||
/// Exactly as given: `dim * 4` bytes per vector.
|
||||
#[default]
|
||||
Float32,
|
||||
/// Each component scaled to an `i8`: `dim` bytes per vector, a quarter of
|
||||
/// the space, at some cost in precision.
|
||||
///
|
||||
/// Only meaningful for [`DistanceMetric::Cosine`]: rows are stored
|
||||
/// unit-length, so a quantised dot product reconstructs the similarity
|
||||
/// directly. Requesting it for `L2` keeps `Float32`, because an L2
|
||||
/// distance cannot be recovered from a dot product alone.
|
||||
Int8,
|
||||
}
|
||||
|
||||
/// The index's copy of the vectors, flat and row-major.
|
||||
#[derive(Debug, Clone)]
|
||||
enum Vectors {
|
||||
F32 {
|
||||
dim: usize,
|
||||
flat: Vec<f32>,
|
||||
},
|
||||
/// `flat[i * dim + j]` is component `j` of vector `i` divided by
|
||||
/// `scales[i]`; multiplying back recovers it.
|
||||
///
|
||||
/// The scale is per row rather than global. A unit-length row in `d`
|
||||
/// dimensions has components around `1/sqrt(d)`, so a fixed `[-1, 1]`
|
||||
/// scale spends fewer than 12 of the 255 levels on a 128-dimensional
|
||||
/// vector and the reconstruction error swamps the gaps between near
|
||||
/// neighbours — measured at 0.35 top-10 overlap with the exact ranking.
|
||||
/// Scaling each row by its own largest component uses the full range.
|
||||
Int8 {
|
||||
dim: usize,
|
||||
flat: Vec<i8>,
|
||||
scales: Vec<f32>,
|
||||
},
|
||||
}
|
||||
|
||||
/// Levels either side of zero. 127, not 128, so the range is symmetric.
|
||||
const INT8_LEVELS: f32 = 127.0;
|
||||
|
||||
/// Quantise one row, returning the codes and the scale that inverts them.
|
||||
fn quantise_row(v: &[f32], out: &mut Vec<i8>) -> f32 {
|
||||
let max_abs = v.iter().fold(0.0f32, |m, x| m.max(x.abs()));
|
||||
if max_abs <= f32::MIN_POSITIVE {
|
||||
out.extend(core::iter::repeat_n(0i8, v.len()));
|
||||
return 0.0;
|
||||
}
|
||||
let inv = INT8_LEVELS / max_abs;
|
||||
out.extend(
|
||||
v.iter()
|
||||
.map(|x| (x * inv).round().clamp(-INT8_LEVELS, INT8_LEVELS) as i8),
|
||||
);
|
||||
max_abs / INT8_LEVELS
|
||||
}
|
||||
|
||||
impl Vectors {
|
||||
fn new(dim: usize, storage: Storage, metric: DistanceMetric) -> Self {
|
||||
match storage {
|
||||
Storage::Int8 if metric == DistanceMetric::Cosine => Vectors::Int8 {
|
||||
dim,
|
||||
flat: Vec::new(),
|
||||
scales: Vec::new(),
|
||||
},
|
||||
_ => Vectors::F32 {
|
||||
dim,
|
||||
flat: Vec::new(),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
fn dim(&self) -> usize {
|
||||
match self {
|
||||
Vectors::F32 { dim, .. } | Vectors::Int8 { dim, .. } => *dim,
|
||||
}
|
||||
}
|
||||
|
||||
fn storage(&self) -> Storage {
|
||||
match self {
|
||||
Vectors::F32 { .. } => Storage::Float32,
|
||||
Vectors::Int8 { .. } => Storage::Int8,
|
||||
}
|
||||
}
|
||||
|
||||
fn len(&self) -> usize {
|
||||
let dim = self.dim();
|
||||
if dim == 0 {
|
||||
return 0;
|
||||
}
|
||||
match self {
|
||||
Vectors::F32 { flat, .. } => flat.len() / dim,
|
||||
Vectors::Int8 { flat, .. } => flat.len() / dim,
|
||||
}
|
||||
}
|
||||
|
||||
/// Set the row width, for a store seeded empty by `new`.
|
||||
fn set_dim(&mut self, new_dim: usize) {
|
||||
match self {
|
||||
Vectors::F32 { dim, .. } | Vectors::Int8 { dim, .. } => *dim = new_dim,
|
||||
}
|
||||
}
|
||||
|
||||
fn push(&mut self, vector: &[f32]) {
|
||||
match self {
|
||||
Vectors::F32 { flat, .. } => flat.extend_from_slice(vector),
|
||||
Vectors::Int8 { flat, scales, .. } => scales.push(quantise_row(vector, flat)),
|
||||
}
|
||||
}
|
||||
|
||||
/// Row `i` as `f32`, for callers that need the values back (serialization,
|
||||
/// and the f32 fast paths). Quantised rows are reconstructed, so this is
|
||||
/// lossy in exactly the way the storage is.
|
||||
fn row(&self, i: usize) -> Vec<f32> {
|
||||
let dim = self.dim();
|
||||
let start = i * dim;
|
||||
match self {
|
||||
Vectors::F32 { flat, .. } => flat[start..start + dim].to_vec(),
|
||||
Vectors::Int8 { flat, scales, .. } => flat[start..start + dim]
|
||||
.iter()
|
||||
.map(|&q| f32::from(q) * scales[i])
|
||||
.collect(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Distance between two stored vectors.
|
||||
fn dist(&self, a: usize, b: usize, metric: DistanceMetric) -> f32 {
|
||||
let dim = self.dim();
|
||||
match self {
|
||||
Vectors::F32 { flat, .. } => {
|
||||
let (x, y) = (a * dim, b * dim);
|
||||
compute_distance(&flat[x..x + dim], &flat[y..y + dim], metric)
|
||||
}
|
||||
Vectors::Int8 { flat, scales, .. } => {
|
||||
let (x, y) = (a * dim, b * dim);
|
||||
let dot = dot_i8(&flat[x..x + dim], &flat[y..y + dim]);
|
||||
1.0 - dot as f32 * scales[a] * scales[b]
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Distance from a prepared query to stored vector `i`.
|
||||
fn dist_query(&self, query: &Query, i: usize, metric: DistanceMetric) -> f32 {
|
||||
let dim = self.dim();
|
||||
let start = i * dim;
|
||||
match (self, query) {
|
||||
(Vectors::F32 { flat, .. }, Query::F32(q)) => {
|
||||
compute_distance(q, &flat[start..start + dim], metric)
|
||||
}
|
||||
(Vectors::Int8 { flat, scales, .. }, Query::Int8(q, q_scale)) => {
|
||||
let dot = dot_i8(q, &flat[start..start + dim]);
|
||||
1.0 - dot as f32 * q_scale * scales[i]
|
||||
}
|
||||
// Mixed forms cannot occur: `Query` is built from the same storage.
|
||||
_ => f32::MAX,
|
||||
}
|
||||
}
|
||||
|
||||
/// Build a store from prepared rows.
|
||||
fn from_rows(rows: &[Vec<f32>], storage: Storage, metric: DistanceMetric) -> Self {
|
||||
let dim = rows.first().map_or(0, Vec::len);
|
||||
let mut out = Vectors::new(dim, storage, metric);
|
||||
for row in rows {
|
||||
out.push(row);
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
/// Prepare `query` for comparison against this store.
|
||||
fn query(&self, query: Vec<f32>) -> Query {
|
||||
match self {
|
||||
Vectors::F32 { .. } => Query::F32(query),
|
||||
Vectors::Int8 { .. } => {
|
||||
let mut codes = Vec::with_capacity(query.len());
|
||||
let scale = quantise_row(&query, &mut codes);
|
||||
Query::Int8(codes, scale)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// What a layer search is measuring distance *to*: an incoming query, or a
|
||||
/// node already in the index (which is what insertion compares against).
|
||||
enum Target<'a> {
|
||||
Query(&'a Query),
|
||||
Node(usize),
|
||||
}
|
||||
|
||||
impl Vectors {
|
||||
fn dist_to(&self, target: &Target<'_>, i: usize, metric: DistanceMetric) -> f32 {
|
||||
match target {
|
||||
Target::Query(q) => self.dist_query(q, i, metric),
|
||||
Target::Node(n) => self.dist(*n, i, metric),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// A search query in whichever form the store compares against.
|
||||
enum Query {
|
||||
F32(Vec<f32>),
|
||||
/// Codes and the scale that inverts them, as in [`Vectors::Int8`].
|
||||
Int8(Vec<i8>, f32),
|
||||
}
|
||||
|
||||
/// Sum of products, widened so it cannot overflow. Runtime-dispatched to the
|
||||
/// same SIMD backend as the f32 kernels, so the two storages are compared on
|
||||
/// equal terms.
|
||||
#[inline]
|
||||
fn dot_i8(a: &[i8], b: &[i8]) -> i32 {
|
||||
clawhdf5_accel::dot_i8(a, b)
|
||||
}
|
||||
|
||||
/// Magic for [`HnswIndex::graph_to_bytes`].
|
||||
const GRAPH_MAGIC: &[u8; 4] = b"CHG1";
|
||||
|
||||
@@ -174,8 +386,8 @@ pub const HNSW_FORMAT_VERSION: i64 = 2;
|
||||
/// HDF5 format.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct HnswIndex {
|
||||
/// All vectors in the index.
|
||||
vectors: Vec<Vec<f32>>,
|
||||
/// All vectors in the index, flat and row-major.
|
||||
vectors: Vectors,
|
||||
/// Adjacency lists per layer. `graph[layer][node]` = list of neighbor IDs.
|
||||
graph: Vec<Vec<Vec<usize>>>,
|
||||
/// Soft-deletion flags, one per node. Deleted nodes remain in the graph for
|
||||
@@ -214,6 +426,20 @@ impl HnswIndex {
|
||||
m: usize,
|
||||
ef_construction: usize,
|
||||
metric: DistanceMetric,
|
||||
) -> Self {
|
||||
Self::build_with(vectors, m, ef_construction, metric, Storage::default())
|
||||
}
|
||||
|
||||
/// Build an index, choosing how the vectors are stored.
|
||||
///
|
||||
/// [`Storage::Int8`] keeps them at a quarter of the size; see its docs for
|
||||
/// what that costs and when it applies.
|
||||
pub fn build_with(
|
||||
vectors: &[Vec<f32>],
|
||||
m: usize,
|
||||
ef_construction: usize,
|
||||
metric: DistanceMetric,
|
||||
storage: Storage,
|
||||
) -> Self {
|
||||
assert!(!vectors.is_empty(), "cannot build index from empty vectors");
|
||||
assert!(m >= 2, "m must be at least 2");
|
||||
@@ -224,8 +450,11 @@ impl HnswIndex {
|
||||
|
||||
let m_max0 = m * 2;
|
||||
let n = vectors.len();
|
||||
let prepared: Vec<Vec<f32>> = vectors.iter().map(|v| prepare(v.clone(), metric)).collect();
|
||||
let vectors: &[Vec<f32>] = &prepared;
|
||||
let mut prepared = Vectors::new(dim, storage, metric);
|
||||
for v in vectors {
|
||||
prepared.push(&prepare(v.clone(), metric));
|
||||
}
|
||||
let vectors = &prepared;
|
||||
|
||||
// Assign levels to all nodes
|
||||
let mut node_levels = Vec::with_capacity(n);
|
||||
@@ -322,9 +551,20 @@ impl HnswIndex {
|
||||
/// point for incremental [`HnswIndex::insert`] and as the result of
|
||||
/// [`HnswIndex::compact`] when every vector has been deleted.
|
||||
pub fn new(m: usize, ef_construction: usize, metric: DistanceMetric) -> Self {
|
||||
Self::new_with(m, ef_construction, metric, Storage::default())
|
||||
}
|
||||
|
||||
/// [`HnswIndex::new`], choosing how the vectors are stored.
|
||||
pub fn new_with(
|
||||
m: usize,
|
||||
ef_construction: usize,
|
||||
metric: DistanceMetric,
|
||||
storage: Storage,
|
||||
) -> Self {
|
||||
assert!(m >= 2, "m must be at least 2");
|
||||
Self {
|
||||
vectors: Vec::new(),
|
||||
// The dimension is set by the first insert.
|
||||
vectors: Vectors::new(0, storage, metric),
|
||||
graph: Vec::new(),
|
||||
deleted: Vec::new(),
|
||||
entry_point: 0,
|
||||
@@ -351,7 +591,8 @@ impl HnswIndex {
|
||||
// Seed an empty index.
|
||||
if id == 0 {
|
||||
let node_level = assign_level(0, self.m);
|
||||
self.vectors.push(vector);
|
||||
self.vectors.set_dim(vector.len());
|
||||
self.vectors.push(&vector);
|
||||
self.deleted.push(false);
|
||||
self.node_levels.push(node_level);
|
||||
self.graph = (0..=node_level).map(|_| vec![Vec::new(); 1]).collect();
|
||||
@@ -361,12 +602,12 @@ impl HnswIndex {
|
||||
|
||||
assert_eq!(
|
||||
vector.len(),
|
||||
self.vectors[0].len(),
|
||||
self.vectors.dim(),
|
||||
"insert dimension mismatch"
|
||||
);
|
||||
|
||||
let node_level = assign_level(id, self.m);
|
||||
self.vectors.push(vector);
|
||||
self.vectors.push(&vector);
|
||||
self.deleted.push(false);
|
||||
self.node_levels.push(node_level);
|
||||
|
||||
@@ -387,7 +628,7 @@ impl HnswIndex {
|
||||
ep = greedy_closest(
|
||||
&self.vectors,
|
||||
&self.graph[layer],
|
||||
&self.vectors[id],
|
||||
&Target::Node(id),
|
||||
ep,
|
||||
self.metric,
|
||||
);
|
||||
@@ -400,7 +641,7 @@ impl HnswIndex {
|
||||
let neighbors = search_layer(
|
||||
&self.vectors,
|
||||
&self.graph[layer],
|
||||
&self.vectors[id],
|
||||
&Target::Node(id),
|
||||
ep,
|
||||
self.ef_construction,
|
||||
self.metric,
|
||||
@@ -466,16 +707,25 @@ impl HnswIndex {
|
||||
pub fn compact(&mut self) -> Vec<Option<usize>> {
|
||||
let mut mapping = vec![None; self.vectors.len()];
|
||||
let mut surviving: Vec<Vec<f32>> = Vec::with_capacity(self.active_len());
|
||||
for (old, v) in self.vectors.iter().enumerate() {
|
||||
for (old, slot) in mapping.iter_mut().enumerate() {
|
||||
if !self.deleted[old] {
|
||||
mapping[old] = Some(surviving.len());
|
||||
surviving.push(v.clone());
|
||||
*slot = Some(surviving.len());
|
||||
surviving.push(self.vectors.row(old));
|
||||
}
|
||||
}
|
||||
// Rebuilding must keep the storage the caller chose; a compaction is
|
||||
// not the place to silently quadruple the index's memory.
|
||||
let storage = self.vectors.storage();
|
||||
*self = if surviving.is_empty() {
|
||||
Self::new(self.m, self.ef_construction, self.metric)
|
||||
Self::new_with(self.m, self.ef_construction, self.metric, storage)
|
||||
} else {
|
||||
Self::build_with_metric(&surviving, self.m, self.ef_construction, self.metric)
|
||||
Self::build_with(
|
||||
&surviving,
|
||||
self.m,
|
||||
self.ef_construction,
|
||||
self.metric,
|
||||
storage,
|
||||
)
|
||||
};
|
||||
mapping
|
||||
}
|
||||
@@ -490,24 +740,22 @@ impl HnswIndex {
|
||||
/// # Returns
|
||||
/// A vector of `(id, distance)` pairs sorted by distance (closest first).
|
||||
pub fn search(&self, query: &[f32], k: usize, ef: usize) -> Vec<(usize, f32)> {
|
||||
if self.vectors.is_empty() {
|
||||
if self.vectors.len() == 0 {
|
||||
return Vec::new();
|
||||
}
|
||||
assert_eq!(
|
||||
query.len(),
|
||||
self.vectors[0].len(),
|
||||
"query dimension mismatch"
|
||||
);
|
||||
assert_eq!(query.len(), self.vectors.dim(), "query dimension mismatch");
|
||||
let ef = ef.max(k);
|
||||
let prepared_query = prepare(query.to_vec(), self.metric);
|
||||
let query = prepared_query.as_slice();
|
||||
// Prepared and, for a quantised store, quantised once per search
|
||||
// rather than once per comparison.
|
||||
let prepared = self.vectors.query(prepare(query.to_vec(), self.metric));
|
||||
let target = Target::Query(&prepared);
|
||||
|
||||
let mut ep = self.entry_point;
|
||||
let top_layer = self.graph.len().saturating_sub(1);
|
||||
|
||||
// Greedy search from top layer down to layer 1
|
||||
for layer in (1..=top_layer).rev() {
|
||||
ep = greedy_closest(&self.vectors, &self.graph[layer], query, ep, self.metric);
|
||||
ep = greedy_closest(&self.vectors, &self.graph[layer], &target, ep, self.metric);
|
||||
}
|
||||
|
||||
// Search layer 0 for the ef nearest *live* nodes. Deleted nodes are
|
||||
@@ -516,7 +764,7 @@ impl HnswIndex {
|
||||
let candidates = search_layer(
|
||||
&self.vectors,
|
||||
&self.graph[0],
|
||||
query,
|
||||
&target,
|
||||
ep,
|
||||
ef,
|
||||
self.metric,
|
||||
@@ -543,14 +791,13 @@ impl HnswIndex {
|
||||
pub fn to_hdf5_bytes(&self) -> Result<Vec<u8>, FormatError> {
|
||||
let mut fw = FmtWriter::new();
|
||||
let n = self.vectors.len();
|
||||
let dim = if n > 0 { self.vectors[0].len() } else { 0 };
|
||||
let dim = self.vectors.dim();
|
||||
|
||||
// Flatten vectors into a 1D array for storage
|
||||
let flat_vectors: Vec<f32> = self
|
||||
.vectors
|
||||
.iter()
|
||||
.flat_map(|v| v.iter().copied())
|
||||
.collect();
|
||||
let mut flat_vectors: Vec<f32> = Vec::with_capacity(n * dim);
|
||||
for i in 0..n {
|
||||
flat_vectors.extend_from_slice(&self.vectors.row(i));
|
||||
}
|
||||
|
||||
let mut group = fw.create_group("ann");
|
||||
|
||||
@@ -709,7 +956,9 @@ impl HnswIndex {
|
||||
};
|
||||
|
||||
Ok(Self {
|
||||
vectors,
|
||||
// Serialized files carry f32 vectors and no storage tag: a
|
||||
// quantised index is rebuilt, not loaded.
|
||||
vectors: Vectors::from_rows(&vectors, Storage::Float32, metric),
|
||||
graph,
|
||||
deleted,
|
||||
entry_point,
|
||||
@@ -773,6 +1022,16 @@ impl HnswIndex {
|
||||
/// `bytes` is validated — a corrupt or mismatched graph is an error, never
|
||||
/// an index that panics or walks out of bounds during a search.
|
||||
pub fn from_graph_bytes(bytes: &[u8], vectors: Vec<Vec<f32>>) -> Result<Self, FormatError> {
|
||||
Self::from_graph_bytes_with(bytes, vectors, Storage::default())
|
||||
}
|
||||
|
||||
/// As [`from_graph_bytes`](Self::from_graph_bytes), choosing how the
|
||||
/// rehydrated vectors are stored.
|
||||
pub fn from_graph_bytes_with(
|
||||
bytes: &[u8],
|
||||
vectors: Vec<Vec<f32>>,
|
||||
storage: Storage,
|
||||
) -> Result<Self, FormatError> {
|
||||
let bad = |what: &str| FormatError::SerializationError(format!("HNSW graph: {what}"));
|
||||
let body_len = bytes
|
||||
.len()
|
||||
@@ -863,7 +1122,14 @@ impl HnswIndex {
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
vectors: vectors.into_iter().map(|v| prepare(v, metric)).collect(),
|
||||
vectors: Vectors::from_rows(
|
||||
&vectors
|
||||
.into_iter()
|
||||
.map(|v| prepare(v, metric))
|
||||
.collect::<Vec<_>>(),
|
||||
storage,
|
||||
metric,
|
||||
),
|
||||
graph,
|
||||
deleted,
|
||||
entry_point,
|
||||
@@ -882,16 +1148,17 @@ impl HnswIndex {
|
||||
|
||||
/// Returns true if the index is empty.
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.vectors.is_empty()
|
||||
self.vectors.len() == 0
|
||||
}
|
||||
|
||||
/// How this index stores its copy of the vectors.
|
||||
pub fn storage(&self) -> Storage {
|
||||
self.vectors.storage()
|
||||
}
|
||||
|
||||
/// Returns the dimension of vectors in the index.
|
||||
pub fn dimension(&self) -> usize {
|
||||
if self.vectors.is_empty() {
|
||||
0
|
||||
} else {
|
||||
self.vectors[0].len()
|
||||
}
|
||||
self.vectors.dim()
|
||||
}
|
||||
|
||||
/// Returns the number of layers in the graph.
|
||||
@@ -916,17 +1183,17 @@ impl HnswIndex {
|
||||
|
||||
/// Greedy search: find the single closest node to `query` starting from `ep`.
|
||||
fn greedy_closest(
|
||||
vectors: &[Vec<f32>],
|
||||
vectors: &Vectors,
|
||||
layer: &[Vec<usize>],
|
||||
query: &[f32],
|
||||
target: &Target<'_>,
|
||||
mut ep: usize,
|
||||
metric: DistanceMetric,
|
||||
) -> usize {
|
||||
let mut best_dist = compute_distance(query, &vectors[ep], metric);
|
||||
let mut best_dist = vectors.dist_to(target, ep, metric);
|
||||
loop {
|
||||
let mut changed = false;
|
||||
for &neighbor in &layer[ep] {
|
||||
let d = compute_distance(query, &vectors[neighbor], metric);
|
||||
let d = vectors.dist_to(target, neighbor, metric);
|
||||
if d < best_dist {
|
||||
best_dist = d;
|
||||
ep = neighbor;
|
||||
@@ -950,15 +1217,15 @@ fn greedy_closest(
|
||||
/// instead meant a query whose neighbourhood had been deleted got back fewer
|
||||
/// than `k` results, or none, however many live records were nearby.
|
||||
fn search_layer(
|
||||
vectors: &[Vec<f32>],
|
||||
vectors: &Vectors,
|
||||
layer: &[Vec<usize>],
|
||||
query: &[f32],
|
||||
target: &Target<'_>,
|
||||
ep: usize,
|
||||
ef: usize,
|
||||
metric: DistanceMetric,
|
||||
skip: Option<&[bool]>,
|
||||
) -> Vec<Candidate> {
|
||||
let ep_dist = compute_distance(query, &vectors[ep], metric);
|
||||
let ep_dist = vectors.dist_to(target, ep, metric);
|
||||
|
||||
// Min-heap of candidates to explore
|
||||
let mut candidates = BinaryHeap::new();
|
||||
@@ -980,7 +1247,7 @@ fn search_layer(
|
||||
visited.begin(vectors.len());
|
||||
visited.insert(ep);
|
||||
search_layer_visit(
|
||||
vectors, layer, query, ef, metric, skip, visited, candidates, results,
|
||||
vectors, layer, target, ef, metric, skip, visited, candidates, results,
|
||||
)
|
||||
})
|
||||
}
|
||||
@@ -1023,9 +1290,9 @@ thread_local! {
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn search_layer_visit(
|
||||
vectors: &[Vec<f32>],
|
||||
vectors: &Vectors,
|
||||
layer: &[Vec<usize>],
|
||||
query: &[f32],
|
||||
target: &Target<'_>,
|
||||
ef: usize,
|
||||
metric: DistanceMetric,
|
||||
skip: Option<&[bool]>,
|
||||
@@ -1044,7 +1311,7 @@ fn search_layer_visit(
|
||||
continue;
|
||||
}
|
||||
|
||||
let d = compute_distance(query, &vectors[neighbor], metric);
|
||||
let d = vectors.dist_to(target, neighbor, metric);
|
||||
let furthest_dist = results.peek().map_or(f32::MAX, |f| f.distance);
|
||||
|
||||
if d < furthest_dist || results.len() < ef {
|
||||
@@ -1095,7 +1362,7 @@ fn search_layer_visit(
|
||||
/// remaining slots are then filled with the closest rejected candidates, so a
|
||||
/// node is never left under-connected.
|
||||
fn select_neighbors(
|
||||
vectors: &[Vec<f32>],
|
||||
vectors: &Vectors,
|
||||
candidates: &[(usize, f32)],
|
||||
max_conn: usize,
|
||||
metric: DistanceMetric,
|
||||
@@ -1111,7 +1378,7 @@ fn select_neighbors(
|
||||
}
|
||||
let diverse = selected
|
||||
.iter()
|
||||
.all(|&s| compute_distance(&vectors[id], &vectors[s], metric) > dist_to_node);
|
||||
.all(|&s| vectors.dist(id, s, metric) > dist_to_node);
|
||||
if diverse {
|
||||
selected.push(id);
|
||||
} else {
|
||||
@@ -1136,7 +1403,7 @@ fn batch_len(linked: usize) -> usize {
|
||||
/// layers, found by searching the graph as it currently stands.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn plan_batch(
|
||||
vectors: &[Vec<f32>],
|
||||
vectors: &Vectors,
|
||||
graph: &[Vec<Vec<usize>>],
|
||||
node_levels: &[usize],
|
||||
batch: std::ops::Range<usize>,
|
||||
@@ -1150,7 +1417,7 @@ fn plan_batch(
|
||||
let mut ep = entry_point;
|
||||
// Phase 1: greedy descent from the top layer down to node_level + 1.
|
||||
for layer in (node_level + 1..=ep_level).rev() {
|
||||
ep = greedy_closest(vectors, &graph[layer], &vectors[i], ep, metric);
|
||||
ep = greedy_closest(vectors, &graph[layer], &Target::Node(i), ep, metric);
|
||||
}
|
||||
// Phase 2: search and select on every layer the node lives on.
|
||||
let mut plan = Vec::with_capacity(node_level.min(ep_level) + 1);
|
||||
@@ -1159,7 +1426,7 @@ fn plan_batch(
|
||||
let neighbors = search_layer(
|
||||
vectors,
|
||||
&graph[layer],
|
||||
&vectors[i],
|
||||
&Target::Node(i),
|
||||
ep,
|
||||
ef_construction,
|
||||
metric,
|
||||
@@ -1186,7 +1453,7 @@ fn plan_batch(
|
||||
/// Prune every `(layer, node)` neighbour list in `overflowed` back to its
|
||||
/// limit. Each list belongs to a different node, so they are independent.
|
||||
fn prune_overflowed(
|
||||
vectors: &[Vec<f32>],
|
||||
vectors: &Vectors,
|
||||
graph: &mut [Vec<Vec<usize>>],
|
||||
overflowed: Vec<(usize, usize)>,
|
||||
(m, m_max0): (usize, usize),
|
||||
@@ -1226,7 +1493,7 @@ const PARALLEL_MIN: usize = 8;
|
||||
/// prunes is too fine-grained to parallelise profitably — measured 1.45x on 16
|
||||
/// cores; bulk builds batch their pruning instead, see `prune_overflowed`.)
|
||||
fn link_back(
|
||||
vectors: &[Vec<f32>],
|
||||
vectors: &Vectors,
|
||||
layer: &mut [Vec<usize>],
|
||||
new_id: usize,
|
||||
selected: &[usize],
|
||||
@@ -1248,7 +1515,7 @@ fn link_back(
|
||||
|
||||
/// Trim `node`'s neighbour list back to `max_conn` with [`select_neighbors`].
|
||||
fn prune_connections(
|
||||
vectors: &[Vec<f32>],
|
||||
vectors: &Vectors,
|
||||
neighbors: &mut Vec<usize>,
|
||||
node: usize,
|
||||
max_conn: usize,
|
||||
@@ -1259,7 +1526,7 @@ fn prune_connections(
|
||||
}
|
||||
let mut scored: Vec<(usize, f32)> = neighbors
|
||||
.iter()
|
||||
.map(|&n| (n, compute_distance(&vectors[node], &vectors[n], metric)))
|
||||
.map(|&n| (n, vectors.dist(node, n, metric)))
|
||||
.collect();
|
||||
scored.sort_by(|a, b| a.1.total_cmp(&b.1).then(a.0.cmp(&b.0)));
|
||||
*neighbors = select_neighbors(vectors, &scored, max_conn, metric);
|
||||
@@ -1519,6 +1786,88 @@ mod tests {
|
||||
assert!(recall >= 0.95, "incremental recall@10 = {recall}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn int8_storage_needs_an_exact_re_score_to_match_f32() {
|
||||
// Cosine only: rows are unit-length, so a quantised dot product
|
||||
// reconstructs the similarity directly.
|
||||
//
|
||||
// Not the `clustered` generator: its clusters are far tighter than any
|
||||
// real embedding, so neighbours sit closer together than the
|
||||
// quantisation error and top-10 identity there is noise — that would
|
||||
// measure the fixture, not the storage.
|
||||
let mut vectors = make_random_vectors(3060, 128, 5);
|
||||
let queries = vectors.split_off(3000);
|
||||
let f32_index =
|
||||
HnswIndex::build_with(&vectors, 8, 40, DistanceMetric::Cosine, Storage::Float32);
|
||||
let quantised =
|
||||
HnswIndex::build_with(&vectors, 8, 40, DistanceMetric::Cosine, Storage::Int8);
|
||||
assert_eq!(quantised.storage(), Storage::Int8);
|
||||
|
||||
// Ground truth, not the f32 index's answers: re-scoring can beat that
|
||||
// index, and measuring against it would score being right as drift.
|
||||
let truth: Vec<Vec<usize>> = queries
|
||||
.iter()
|
||||
.map(|q| {
|
||||
let mut d: Vec<(usize, f32)> = vectors
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, v)| (i, compute_distance(q, v, DistanceMetric::Cosine)))
|
||||
.collect();
|
||||
d.sort_by(|a, b| a.1.total_cmp(&b.1));
|
||||
d[..10].iter().map(|x| x.0).collect()
|
||||
})
|
||||
.collect();
|
||||
|
||||
let recall = |got: &dyn Fn(&[f32]) -> Vec<usize>| -> f64 {
|
||||
let mut hits = 0;
|
||||
for (q, want) in queries.iter().zip(&truth) {
|
||||
hits += got(q).iter().filter(|id| want.contains(id)).count();
|
||||
}
|
||||
hits as f64 / (10 * queries.len()) as f64
|
||||
};
|
||||
|
||||
let exact_recall = recall(&|q| f32_index.search(q, 10, 64).iter().map(|r| r.0).collect());
|
||||
let raw_recall = recall(&|q| quantised.search(q, 10, 64).iter().map(|r| r.0).collect());
|
||||
// Quantised distances alone cost recall, and `ef` cannot buy it back:
|
||||
// the loss is in the distances, not in the graph.
|
||||
assert!(
|
||||
raw_recall < exact_recall,
|
||||
"int8 alone should cost recall: {raw_recall} vs {exact_recall}"
|
||||
);
|
||||
|
||||
// Re-scoring a wider candidate pool against the exact vectors — what a
|
||||
// caller holding them (the agent's embedding cache) does — puts it
|
||||
// back, because only the *ordering* was approximate.
|
||||
let rescored_recall = recall(&|q| {
|
||||
let mut pool: Vec<(usize, f32)> = quantised
|
||||
.search(q, 40, 64)
|
||||
.into_iter()
|
||||
.map(|(id, _)| {
|
||||
(
|
||||
id,
|
||||
compute_distance(q, &vectors[id], DistanceMetric::Cosine),
|
||||
)
|
||||
})
|
||||
.collect();
|
||||
pool.sort_by(|a, b| a.1.total_cmp(&b.1));
|
||||
pool.truncate(10);
|
||||
pool.into_iter().map(|p| p.0).collect()
|
||||
});
|
||||
assert!(
|
||||
rescored_recall >= exact_recall - 0.01,
|
||||
"int8 + exact re-score should match f32: {rescored_recall} vs {exact_recall} (raw {raw_recall})"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn int8_storage_falls_back_to_f32_for_non_cosine_metrics() {
|
||||
// L2 distance is not recoverable from a quantised dot product, so the
|
||||
// store silently stays f32 rather than returning wrong distances.
|
||||
let vectors = clustered(100, 8, 5, 3);
|
||||
let index = HnswIndex::build_with(&vectors, 8, 40, DistanceMetric::L2, Storage::Int8);
|
||||
assert_eq!(index.storage(), Storage::Float32);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn deletions_near_the_query_do_not_shrink_or_degrade_results() {
|
||||
let mut vectors = clustered(2040, 16, 20, 11);
|
||||
@@ -1650,6 +1999,7 @@ mod tests {
|
||||
vec![1.2, 0.0], // 3
|
||||
vec![-2.0, 0.0], // 4
|
||||
];
|
||||
let store = Vectors::from_rows(&vectors, Storage::Float32, DistanceMetric::L2);
|
||||
let scored: Vec<(usize, f32)> = (1..5)
|
||||
.map(|i| {
|
||||
(
|
||||
@@ -1659,12 +2009,12 @@ mod tests {
|
||||
})
|
||||
.collect();
|
||||
assert_eq!(
|
||||
select_neighbors(&vectors, &scored, 2, DistanceMetric::L2),
|
||||
select_neighbors(&store, &scored, 2, DistanceMetric::L2),
|
||||
[1, 4]
|
||||
);
|
||||
// Spare capacity is filled with the closest rejected candidates.
|
||||
assert_eq!(
|
||||
select_neighbors(&vectors, &scored, 3, DistanceMetric::L2),
|
||||
select_neighbors(&store, &scored, 3, DistanceMetric::L2),
|
||||
[1, 4, 2]
|
||||
);
|
||||
}
|
||||
@@ -1790,7 +2140,7 @@ mod tests {
|
||||
|
||||
// Verify vectors match
|
||||
for i in 0..loaded.len() {
|
||||
assert_eq!(loaded.vectors[i], index.vectors[i]);
|
||||
assert_eq!(loaded.vectors.row(i), index.vectors.row(i));
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -5,4 +5,4 @@
|
||||
|
||||
mod hnsw;
|
||||
|
||||
pub use hnsw::{DistanceMetric, HnswIndex};
|
||||
pub use hnsw::{DistanceMetric, HnswIndex, Storage};
|
||||
|
||||
@@ -1,7 +1,8 @@
|
||||
[package]
|
||||
name = "clawhdf5-bench"
|
||||
version = "2.5.0"
|
||||
version = "2.7.0"
|
||||
edition = "2024"
|
||||
rust-version.workspace = true
|
||||
description = "Benchmark harnesses for clawhdf5-agent (Track 8)"
|
||||
license = "MIT"
|
||||
|
||||
|
||||
@@ -57,7 +57,8 @@ mod embedder;
|
||||
|
||||
use clawhdf5_agent::bm25::TokenFilter;
|
||||
use clawhdf5_agent::hybrid::Fusion;
|
||||
use clawhdf5_agent::{AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry};
|
||||
use clawhdf5_agent::reranker::{ReRankConfig, RerankInput, rerank};
|
||||
use clawhdf5_agent::{AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry, SearchResult};
|
||||
use serde::Deserialize;
|
||||
use tempfile::TempDir;
|
||||
|
||||
@@ -69,9 +70,17 @@ fn describe(mode: Mode) -> String {
|
||||
Fusion::Weighted { vector, keyword } => format!("vector_{vector:.1}_keyword_{keyword:.1}"),
|
||||
Fusion::Rrf { k } => format!("rrf_k{k:.0}"),
|
||||
};
|
||||
match mode.tokens {
|
||||
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
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -83,6 +92,9 @@ struct Mode {
|
||||
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 {
|
||||
@@ -91,9 +103,17 @@ impl Mode {
|
||||
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;
|
||||
@@ -121,11 +141,61 @@ 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");
|
||||
|
||||
@@ -217,6 +287,37 @@ struct Question {
|
||||
haystack_session_ids: Vec<String>,
|
||||
haystack_sessions: Vec<Vec<Turn>>,
|
||||
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)
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
@@ -235,11 +336,21 @@ struct Metrics {
|
||||
rr_turn: f64,
|
||||
abstention_correct: 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>,
|
||||
count: u32,
|
||||
}
|
||||
|
||||
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 {
|
||||
self.hit1_session as f64 / self.count.max(1) as f64 * 100.0
|
||||
}
|
||||
@@ -297,6 +408,16 @@ struct EvalResult {
|
||||
hit5_turn: bool,
|
||||
hit10_turn: bool,
|
||||
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,
|
||||
}
|
||||
|
||||
@@ -317,15 +438,21 @@ fn evaluate_question(
|
||||
// Build MemoryEntry list from all haystack sessions
|
||||
let mut entries: Vec<MemoryEntry> = 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() {
|
||||
let sess_id = q
|
||||
.haystack_session_ids
|
||||
.get(sess_idx)
|
||||
.map(String::as_str)
|
||||
.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 {
|
||||
chunk: turn.content.clone(),
|
||||
embedding: embedding_for(embeddings, &turn.content),
|
||||
@@ -339,7 +466,6 @@ fn evaluate_question(
|
||||
},
|
||||
});
|
||||
turn_has_answer.push(turn.has_answer);
|
||||
ts += 1.0;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -356,11 +482,87 @@ fn evaluate_question(
|
||||
// Set of session IDs that contain the answer
|
||||
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 t0 = Instant::now();
|
||||
let results = memory.hybrid_search_with(&query_emb, &q.question, mode.fusion, top_k);
|
||||
// Re-ranking only reorders; it needs a candidate pool larger than `top_k`
|
||||
// to have anything to promote.
|
||||
let pool = if mode.rerank.is_some() {
|
||||
top_k * 4
|
||||
} else {
|
||||
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();
|
||||
|
||||
// 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
|
||||
let mut hit1_session = false;
|
||||
let mut hit5_session = false;
|
||||
@@ -415,6 +617,7 @@ fn evaluate_question(
|
||||
hit5_turn,
|
||||
hit10_turn,
|
||||
rr_turn,
|
||||
newest_gold_first,
|
||||
latency,
|
||||
}
|
||||
}
|
||||
@@ -566,6 +769,24 @@ fn print_report(
|
||||
);
|
||||
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 {
|
||||
println!("## Abstention Accuracy");
|
||||
println!(
|
||||
@@ -679,6 +900,14 @@ fn print_report(
|
||||
} else {
|
||||
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!(
|
||||
" \"avg\": {:.1}, \"p50\": {:.1}, \"p95\": {:.1}, \"p99\": {:.1}",
|
||||
@@ -701,6 +930,8 @@ fn main() {
|
||||
let mut limit: Option<usize> = None;
|
||||
let mut weights_dir: Option<String> = None;
|
||||
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);
|
||||
while let Some(arg) = args.next() {
|
||||
match arg.as_str() {
|
||||
@@ -709,6 +940,16 @@ fn main() {
|
||||
limit = Some(v.parse().expect("--limit must be a positive integer"));
|
||||
}
|
||||
"--sweep" => sweep = true,
|
||||
"--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" => {
|
||||
weights_dir = Some(args.next().expect("--embeddings needs a directory"));
|
||||
}
|
||||
@@ -727,6 +968,9 @@ fn main() {
|
||||
BM25-only, vector-only, and hybrid separately. Requires\n\
|
||||
--features embeddings; without it the vector stage is\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\
|
||||
--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\
|
||||
never searched; this is what searches it."
|
||||
@@ -792,6 +1036,10 @@ fn main() {
|
||||
{
|
||||
if sweep {
|
||||
sweep_modes()
|
||||
} else if rerank_sweep {
|
||||
let mut modes = vec![HYBRID, hybrid_rerank_metadata_only()];
|
||||
modes.extend(hybrid_rerank_half_lives());
|
||||
modes
|
||||
} else {
|
||||
vec![
|
||||
BM25_ONLY,
|
||||
@@ -800,6 +1048,8 @@ fn main() {
|
||||
RRF,
|
||||
BM25_STEMMED,
|
||||
HYBRID_STEMMED,
|
||||
hybrid_rerank_metadata_only(),
|
||||
hybrid_rerank_blended(),
|
||||
]
|
||||
}
|
||||
}
|
||||
@@ -916,6 +1166,14 @@ fn run_mode(
|
||||
entry.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;
|
||||
entry.latency_ns.push(ns);
|
||||
@@ -927,3 +1185,30 @@ fn run_mode(
|
||||
eprintln!();
|
||||
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");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -19,12 +19,14 @@
|
||||
//! 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 -- --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
|
||||
//! ```
|
||||
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use clawhdf5_agent::{AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry};
|
||||
use clawhdf5_ann::{DistanceMetric, HnswIndex};
|
||||
use clawhdf5_ann::{DistanceMetric, HnswIndex, Storage};
|
||||
|
||||
const DIM: usize = 384;
|
||||
const K: usize = 10;
|
||||
@@ -84,6 +86,25 @@ struct Dataset {
|
||||
/// that appears only on clustered data points at graph connectivity.
|
||||
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 {
|
||||
let mut rng = Rng(seed);
|
||||
if UNIFORM.load(std::sync::atomic::Ordering::Relaxed) {
|
||||
@@ -169,6 +190,11 @@ fn text_for(cluster: usize, i: usize, rng: &mut Rng) -> String {
|
||||
// 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> {
|
||||
// Vectors are unit length, so cosine order == dot-product order.
|
||||
let mut scored: Vec<(usize, f32)> = vectors
|
||||
@@ -198,6 +224,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 {
|
||||
d.as_secs_f64() * 1e6
|
||||
}
|
||||
@@ -219,11 +322,12 @@ fn bench_ann(n: usize, json: &mut Vec<serde_json::Value>) {
|
||||
.collect();
|
||||
|
||||
let started = Instant::now();
|
||||
let index = HnswIndex::build_with_metric(
|
||||
let index = HnswIndex::build_with(
|
||||
&data.vectors,
|
||||
HNSW_M,
|
||||
HNSW_EF_CONSTRUCTION,
|
||||
DistanceMetric::Cosine,
|
||||
storage(),
|
||||
);
|
||||
let build = started.elapsed();
|
||||
|
||||
@@ -240,7 +344,8 @@ fn bench_ann(n: usize, json: &mut Vec<serde_json::Value>) {
|
||||
);
|
||||
|
||||
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!(
|
||||
"build: {:.1} ms ({:.0} vectors/s) · exact scan: {:.0} QPS, p50 {:.0} µs\n",
|
||||
@@ -251,12 +356,26 @@ fn bench_ann(n: usize, json: &mut Vec<serde_json::Value>) {
|
||||
);
|
||||
println!("| ef | recall@{K} | QPS | p50 µs | p99 µs |");
|
||||
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 {
|
||||
let mut hits = 0usize;
|
||||
let mut samples = Vec::with_capacity(data.queries.len());
|
||||
for (q, want) in data.queries.iter().zip(&truth) {
|
||||
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());
|
||||
hits += got.iter().filter(|(id, _)| want.contains(id)).count();
|
||||
}
|
||||
@@ -369,6 +488,319 @@ fn bench_end_to_end(n: usize, json: &mut Vec<serde_json::Value>) {
|
||||
}));
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// 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?
|
||||
// ---------------------------------------------------------------------------
|
||||
@@ -394,11 +826,12 @@ fn fusion_study(n: usize) {
|
||||
.map(|(i, c)| text_for(*c, i, &mut rng))
|
||||
.collect();
|
||||
let bm25 = BM25Index::build(&texts, &vec![0u8; n]);
|
||||
let index = HnswIndex::build_with_metric(
|
||||
let index = HnswIndex::build_with(
|
||||
&data.vectors,
|
||||
HNSW_M,
|
||||
HNSW_EF_CONSTRUCTION,
|
||||
DistanceMetric::Cosine,
|
||||
storage(),
|
||||
);
|
||||
|
||||
let vec_pool = (K * 8).max(64);
|
||||
@@ -450,6 +883,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() {
|
||||
let args: Vec<String> = std::env::args().skip(1).collect();
|
||||
let full = args.iter().any(|a| a == "--full");
|
||||
@@ -464,6 +960,45 @@ fn main() {
|
||||
}
|
||||
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") {
|
||||
UNIFORM.store(true, std::sync::atomic::Ordering::Relaxed);
|
||||
println!("(uniform random data)");
|
||||
@@ -485,6 +1020,18 @@ fn main() {
|
||||
|
||||
let mut json = Vec::new();
|
||||
println!("## Search harness");
|
||||
|
||||
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") {
|
||||
|
||||
@@ -1,7 +1,8 @@
|
||||
[package]
|
||||
name = "clawhdf5-cli"
|
||||
version = "2.5.0"
|
||||
version = "2.7.0"
|
||||
edition = "2024"
|
||||
rust-version.workspace = true
|
||||
license = "MIT"
|
||||
description = "CLI for clawhdf5 agent memory — create, save, search, recall, stats"
|
||||
repository = "https://git.redclaw.dev/quantumclaw/clawhdf5"
|
||||
@@ -14,7 +15,7 @@ name = "clawhdf5"
|
||||
path = "src/main.rs"
|
||||
|
||||
[dependencies]
|
||||
clawhdf5-agent = { path = "../clawhdf5-agent", version = "2.5.0" }
|
||||
clawhdf5-agent = { path = "../clawhdf5-agent", version = "2.7.0" }
|
||||
clap = { version = "4", features = ["derive", "env"] }
|
||||
serde_json = "1"
|
||||
serde = { workspace = true }
|
||||
|
||||
@@ -28,6 +28,19 @@ enum Commands {
|
||||
/// Enable write-ahead log
|
||||
#[arg(long)]
|
||||
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 on disk as IEEE half precision (float16): half
|
||||
/// the bytes, about three significant digits; values must lie within
|
||||
/// ±65504
|
||||
#[arg(long)]
|
||||
float16: bool,
|
||||
},
|
||||
/// Save a memory entry (reads JSON from stdin or --json)
|
||||
Save {
|
||||
@@ -88,9 +101,24 @@ fn main() {
|
||||
|
||||
fn run(cli: Cli) -> Result<(), Box<dyn std::error::Error>> {
|
||||
match cli.command {
|
||||
Commands::Create { agent_id, dim, wal } => {
|
||||
Commands::Create {
|
||||
agent_id,
|
||||
dim,
|
||||
wal,
|
||||
f32_index,
|
||||
quantized_index: _,
|
||||
float16,
|
||||
} => {
|
||||
let mut config = MemoryConfig::new(cli.path.clone(), &agent_id, dim);
|
||||
config.wal_enabled = wal;
|
||||
config.float16 = 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 mem = HDF5Memory::create(config)?;
|
||||
let j = serde_json::json!({
|
||||
"status": "created",
|
||||
@@ -98,6 +126,8 @@ fn run(cli: Cli) -> Result<(), Box<dyn std::error::Error>> {
|
||||
"agent_id": agent_id,
|
||||
"embedding_dim": dim,
|
||||
"wal_enabled": wal,
|
||||
"quantized_index": config_quantized,
|
||||
"float16": float16,
|
||||
"count": mem.count(),
|
||||
});
|
||||
println!("{}", serde_json::to_string_pretty(&j)?);
|
||||
|
||||
@@ -1,7 +1,8 @@
|
||||
[package]
|
||||
name = "clawhdf5-derive"
|
||||
version = "2.5.0"
|
||||
version = "2.7.0"
|
||||
edition = "2024"
|
||||
rust-version.workspace = true
|
||||
description = "Derive macros for rustyhdf5 HDF5 traits"
|
||||
license = "MIT"
|
||||
repository = "https://git.redclaw.dev/quantumclaw/clawhdf5"
|
||||
|
||||
@@ -1,7 +1,8 @@
|
||||
[package]
|
||||
name = "clawhdf5-filters"
|
||||
version = "2.5.0"
|
||||
version = "2.7.0"
|
||||
edition = "2024"
|
||||
rust-version.workspace = true
|
||||
description = "Filter and compression pipeline for clawhdf5"
|
||||
license = "MIT"
|
||||
repository = "https://git.redclaw.dev/quantumclaw/clawhdf5"
|
||||
@@ -25,8 +26,12 @@ name = "compression_bench"
|
||||
harness = false
|
||||
|
||||
[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"]
|
||||
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 = []
|
||||
|
||||
@@ -8,16 +8,18 @@ Filter and compression pipeline for clawhdf5.
|
||||
## Features
|
||||
|
||||
- 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)
|
||||
|
||||
## Usage
|
||||
|
||||
```rust
|
||||
use clawhdf5_filters::{deflate_decode, deflate_encode};
|
||||
use clawhdf5_filters::{deflate_compress, deflate_decompress};
|
||||
|
||||
let compressed = deflate_encode(&data, 6).unwrap();
|
||||
let decompressed = deflate_decode(&compressed).unwrap();
|
||||
let compressed = deflate_compress(&data, 6).unwrap();
|
||||
// The second argument bounds the output: the expected decompressed size.
|
||||
let decompressed = deflate_decompress(&compressed, data.len()).unwrap();
|
||||
```
|
||||
|
||||
## 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):
|
||||
//! 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
|
||||
//! and is typically the fastest option on macOS. zlib-ng is the fastest portable
|
||||
//! option and what C HDF5 uses internally.
|
||||
//! and is typically the fastest option on macOS. zlib-rs is a pure-Rust port of
|
||||
//! zlib-ng; see `BENCHMARKS.md` for how the two compare.
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// 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.
|
||||
///
|
||||
/// When the output size is known (typical for HDF5 chunks), this avoids
|
||||
/// dynamic reallocation by writing directly into a pre-sized buffer.
|
||||
/// Decompress into a buffer pre-sized to `output_size`, the expected
|
||||
/// decompressed length (known for HDF5 chunks). Output longer than that is an
|
||||
/// error, as is a stream that ends early.
|
||||
pub(crate) fn flate2_decompress_preallocated(
|
||||
data: &[u8],
|
||||
output_size: usize,
|
||||
) -> Result<Vec<u8>, String> {
|
||||
use std::io::Read;
|
||||
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)
|
||||
inflate_bounded(data, output_size, output_size)
|
||||
}
|
||||
|
||||
/// Absolute ceiling on decompressed output when the caller has no size hint,
|
||||
/// preventing unbounded allocation from a hostile/corrupted zlib stream.
|
||||
const MAX_DECOMPRESS_SIZE: usize = 256 * 1024 * 1024;
|
||||
|
||||
/// Streaming decompress with dynamic sizing (when output size is unknown).
|
||||
///
|
||||
/// Bounded by [`MAX_DECOMPRESS_SIZE`] since there is no chunk-size hint to
|
||||
/// validate against here — an unbounded `read_to_end` would let a hostile
|
||||
/// zlib stream force arbitrarily large allocation (a "zlib bomb").
|
||||
/// Decompress with no size hint, bounded by [`MAX_DECOMPRESS_SIZE`] so a
|
||||
/// hostile zlib stream cannot force arbitrarily large allocation (a "zlib
|
||||
/// bomb").
|
||||
pub(crate) fn flate2_decompress_streaming(data: &[u8]) -> Result<Vec<u8>, String> {
|
||||
use std::io::Read;
|
||||
let decoder = flate2::read::ZlibDecoder::new(data);
|
||||
let mut result = Vec::new();
|
||||
decoder
|
||||
.take(MAX_DECOMPRESS_SIZE as u64 + 1)
|
||||
.read_to_end(&mut result)
|
||||
.map_err(|e| e.to_string())?;
|
||||
if result.len() > MAX_DECOMPRESS_SIZE {
|
||||
return Err(format!(
|
||||
"decompressed output exceeds {} MiB limit",
|
||||
MAX_DECOMPRESS_SIZE / 1024 / 1024
|
||||
));
|
||||
}
|
||||
Ok(result)
|
||||
let hint = data.len().saturating_mul(4).min(1 << 20);
|
||||
inflate_bounded(data, hint, MAX_DECOMPRESS_SIZE).map_err(|e| {
|
||||
if e.ends_with("exceeds size limit") {
|
||||
format!(
|
||||
"decompressed output exceeds {} MiB limit",
|
||||
MAX_DECOMPRESS_SIZE / 1024 / 1024
|
||||
)
|
||||
} else {
|
||||
e
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
/// 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> {
|
||||
use std::io::Write;
|
||||
let mut encoder = flate2::write::ZlibEncoder::new(Vec::new(), flate2::Compression::new(level));
|
||||
encoder.write_all(data).map_err(|e| e.to_string())?;
|
||||
encoder.finish().map_err(|e| e.to_string())
|
||||
use flate2::{Compress, Compression, FlushCompress, Status};
|
||||
|
||||
// zlib's compressBound, plus the zlib header and trailer.
|
||||
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:
|
||||
/// 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
|
||||
/// decompression (avoids reallocation).
|
||||
@@ -344,7 +397,7 @@ pub fn decompress(data: &[u8], output_hint: usize) -> Result<Vec<u8>, String> {
|
||||
///
|
||||
/// Selection order:
|
||||
/// 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> {
|
||||
#[cfg(all(target_os = "macos", feature = "apple-compression"))]
|
||||
{
|
||||
@@ -377,9 +430,19 @@ pub fn active_backend() -> &'static str {
|
||||
{
|
||||
"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(
|
||||
all(target_os = "macos", feature = "apple-compression"),
|
||||
feature = "fast-deflate"
|
||||
feature = "fast-deflate",
|
||||
feature = "zlib-rs"
|
||||
)))]
|
||||
{
|
||||
"miniz_oxide"
|
||||
@@ -436,7 +499,7 @@ mod tests {
|
||||
fn backend_name_is_set() {
|
||||
let name = active_backend();
|
||||
assert!(
|
||||
["miniz_oxide", "zlib-ng", "apple-compression"].contains(&name),
|
||||
["miniz_oxide", "zlib-rs", "zlib-ng", "apple-compression"].contains(&name),
|
||||
"unexpected backend: {name}"
|
||||
);
|
||||
}
|
||||
|
||||
@@ -2,12 +2,14 @@
|
||||
//!
|
||||
//! Provides deflate (zlib) decompression/compression with multiple backend options:
|
||||
//!
|
||||
//! - **Default**: `miniz_oxide` (pure Rust, no C dependencies)
|
||||
//! - **`fast-deflate` feature**: `zlib-ng` via flate2 (~2-3x faster, matches C HDF5)
|
||||
//! - **Default (`zlib-rs` feature)**: `zlib-rs` via flate2 (pure Rust, no C
|
||||
//! dependencies)
|
||||
//! - **`fast-deflate` feature**: `zlib-ng` via flate2 (C, built with cmake)
|
||||
//! - **`apple-compression` feature**: Apple Compression Framework on macOS
|
||||
//! (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;
|
||||
|
||||
@@ -115,7 +117,7 @@ mod tests {
|
||||
fn backend_reports_name() {
|
||||
let name = deflate_backend();
|
||||
assert!(
|
||||
["miniz_oxide", "zlib-ng", "apple-compression"].contains(&name),
|
||||
["miniz_oxide", "zlib-rs", "zlib-ng", "apple-compression"].contains(&name),
|
||||
"unexpected backend: {name}"
|
||||
);
|
||||
}
|
||||
|
||||
@@ -1,7 +1,8 @@
|
||||
[package]
|
||||
name = "clawhdf5-format"
|
||||
version = "2.5.0"
|
||||
version = "2.7.0"
|
||||
edition = "2024"
|
||||
rust-version.workspace = true
|
||||
description = "Pure-Rust HDF5 binary format parsing and writing — no C dependencies"
|
||||
license = "MIT"
|
||||
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 }
|
||||
|
||||
[dev-dependencies]
|
||||
half = { workspace = true }
|
||||
serde_json = "1"
|
||||
criterion = { workspace = true }
|
||||
clawhdf5-derive = { path = "../clawhdf5-derive", version = "2.5.0" }
|
||||
clawhdf5-derive = { path = "../clawhdf5-derive", version = "2.7.0" }
|
||||
|
||||
[[bench]]
|
||||
name = "bench"
|
||||
harness = false
|
||||
|
||||
[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 = []
|
||||
checksum = []
|
||||
deflate = ["flate2"]
|
||||
@@ -42,7 +47,10 @@ fast-checksum = ["crc32fast"]
|
||||
fast-deflate = ["flate2/zlib-ng"]
|
||||
system-zlib = ["flate2/zlib-default"]
|
||||
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"]
|
||||
zstd = ["dep:zstd"]
|
||||
blake3_hash = ["blake3"]
|
||||
|
||||
@@ -1 +1,4 @@
|
||||
target/
|
||||
corpus/
|
||||
artifacts/
|
||||
coverage/
|
||||
|
||||
@@ -1,15 +1,36 @@
|
||||
#![no_main]
|
||||
use clawhdf5_format::btree_v2::{BTreeV2Header, collect_btree_v2_records};
|
||||
use libfuzzer_sys::fuzz_target;
|
||||
|
||||
fuzz_target!(|data: &[u8]| {
|
||||
for &offset_size in &[4u8, 8] {
|
||||
for &length_size in &[4u8, 8] {
|
||||
let _ = clawhdf5_format::btree_v2::BTreeV2Header::parse(
|
||||
data,
|
||||
0,
|
||||
offset_size,
|
||||
length_size,
|
||||
);
|
||||
if let Ok(header) = BTreeV2Header::parse(data, 0, offset_size, length_size) {
|
||||
let _ = collect_btree_v2_records(data, &header, 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
|
||||
}
|
||||
|
||||
/// 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.
|
||||
pub fn collect_btree_v2_records(
|
||||
file_data: &[u8],
|
||||
@@ -182,6 +193,22 @@ pub fn collect_btree_v2_records(
|
||||
if header.total_records == 0 || header.num_records_in_root == 0 {
|
||||
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);
|
||||
|
||||
@@ -206,6 +233,7 @@ pub fn collect_btree_v2_records(
|
||||
offset_size,
|
||||
length_size,
|
||||
max_leaf_nrec,
|
||||
&mut budget,
|
||||
&mut records,
|
||||
)?;
|
||||
Ok(records)
|
||||
@@ -273,6 +301,7 @@ fn collect_internal_records(
|
||||
offset_size: u8,
|
||||
length_size: u8,
|
||||
max_leaf_nrec: u64,
|
||||
budget: &mut usize,
|
||||
out: &mut Vec<BTreeV2Record>,
|
||||
) -> Result<(), FormatError> {
|
||||
// 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.
|
||||
for (i, &(child_addr, child_nrec)) in children.iter().enumerate() {
|
||||
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 =
|
||||
parse_leaf_records(file_data, child_addr as usize, child_nrec, record_size)?;
|
||||
out.extend(leaf_recs);
|
||||
@@ -364,6 +395,7 @@ fn collect_internal_records(
|
||||
offset_size,
|
||||
length_size,
|
||||
max_leaf_nrec,
|
||||
budget,
|
||||
out,
|
||||
)?;
|
||||
}
|
||||
@@ -393,6 +425,7 @@ fn collect_internal_records(
|
||||
available: file_data.len(),
|
||||
});
|
||||
}
|
||||
spend(budget, 1)?;
|
||||
out.push(BTreeV2Record {
|
||||
data: file_data[rec_start..rec_end].to_vec(),
|
||||
});
|
||||
@@ -466,6 +499,124 @@ mod tests {
|
||||
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]
|
||||
fn parse_header() {
|
||||
let data = build_btree_v2_header(5, 512, 11, 0, 0x1000, 3, 3, 8, 8);
|
||||
|
||||
@@ -15,7 +15,6 @@ use crate::filter_pipeline::{
|
||||
FilterDescription, FilterPipeline,
|
||||
};
|
||||
use crate::filters::compress_chunk;
|
||||
|
||||
/// 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
|
||||
@@ -928,6 +927,7 @@ pub fn write_selection_to_buffer(
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
|
||||
use super::*;
|
||||
use crate::chunked_read::read_chunked_data;
|
||||
use crate::data_layout::DataLayout;
|
||||
@@ -1512,9 +1512,20 @@ mod tests {
|
||||
|
||||
// ---- 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")]
|
||||
fn h5py_available() -> bool {
|
||||
std::process::Command::new("python3")
|
||||
std::process::Command::new(python())
|
||||
.args(["-c", "import h5py"])
|
||||
.output()
|
||||
.map(|o| o.status.success())
|
||||
@@ -1526,10 +1537,10 @@ mod tests {
|
||||
if !h5py_available() {
|
||||
panic!("h5py not installed — skipping interop test");
|
||||
}
|
||||
let o = std::process::Command::new("python3")
|
||||
let o = std::process::Command::new(python())
|
||||
.args(["-c", script])
|
||||
.output()
|
||||
.expect("python3");
|
||||
.expect("python interpreter");
|
||||
if !o.status.success() {
|
||||
panic!("h5py: {}", String::from_utf8_lossy(&o.stderr));
|
||||
}
|
||||
|
||||
@@ -1076,6 +1076,21 @@ pub fn read_as_f32(raw: &[u8], datatype: &Datatype) -> Result<Vec<f32>, FormatEr
|
||||
) {
|
||||
return Ok(native_le_to_vec::<f32>(raw, count));
|
||||
}
|
||||
// Little-endian half precision (numpy float16): widen directly.
|
||||
if matches!(
|
||||
datatype,
|
||||
Datatype::FloatingPoint {
|
||||
size: 2,
|
||||
byte_order: DatatypeByteOrder::LittleEndian,
|
||||
..
|
||||
}
|
||||
) {
|
||||
let (halves, _) = raw[..count * 2].as_chunks::<2>();
|
||||
return Ok(halves
|
||||
.iter()
|
||||
.map(|&b| f16_bits_to_f32(u16::from_le_bytes(b)))
|
||||
.collect());
|
||||
}
|
||||
|
||||
let order = get_byte_order(datatype);
|
||||
let mut result = Vec::with_capacity(count);
|
||||
@@ -1622,36 +1637,7 @@ fn read_f16_bytes(bytes: &[u8], order: &DatatypeByteOrder) -> f32 {
|
||||
f16_bits_to_f32(u16::from_le_bytes(buf))
|
||||
}
|
||||
|
||||
/// Convert the bit pattern of an IEEE-754 half (binary16) to an `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)
|
||||
}
|
||||
use crate::float16::f16_bits_to_f32;
|
||||
|
||||
fn read_f32_bytes(bytes: &[u8], order: &DatatypeByteOrder) -> f32 {
|
||||
let mut buf = [0u8; 4];
|
||||
|
||||
@@ -640,7 +640,8 @@ impl Datatype {
|
||||
mantissa_size,
|
||||
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 {
|
||||
DatatypeByteOrder::BigEndian => {
|
||||
bf0 |= 0x01;
|
||||
@@ -650,9 +651,14 @@ impl Datatype {
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
// bf[1] bits 0-1: mantissa normalization = 2 (MSB not stored, IEEE 754)
|
||||
let bf1 = 0x3fu8; // matching what h5py generates
|
||||
let mut buf = Self::build_header(1, 1, [bf0, bf1, 0], *size);
|
||||
// Bits 8-15: the sign bit's position, the top bit of the value.
|
||||
// This was hard-coded to 63, which is right only for f64: the
|
||||
// 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_precision.to_le_bytes());
|
||||
buf.push(*exponent_location);
|
||||
@@ -818,6 +824,24 @@ fn build_dt_header(class: u8, version: u8, bf: [u8; 3], size: u32) -> Vec<u8> {
|
||||
mod tests {
|
||||
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
|
||||
fn build_fixed_point(
|
||||
size: u32,
|
||||
|
||||
@@ -12,6 +12,31 @@ use alloc::{format, vec, vec::Vec};
|
||||
use crate::chunked_read::ChunkInfo;
|
||||
use crate::error::FormatError;
|
||||
|
||||
/// Verify the Jenkins lookup3 checksum stored immediately after
|
||||
/// `data[start..end]`, as every Extensible Array structure carries one.
|
||||
///
|
||||
/// A corrupt chunk index yields addresses pointing at the wrong bytes, so a
|
||||
/// mismatch is an error: otherwise the damage surfaces as plausible data read
|
||||
/// 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 Extensible Array header (AEHD).
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ExtensibleArrayHeader {
|
||||
@@ -145,6 +170,8 @@ impl ExtensibleArrayHeader {
|
||||
pos += ls; // skip nelmts
|
||||
pos += ls; // skip max_idx_set (6th stats field)
|
||||
let index_block_address = read_offset(d, pos, offset_size)?;
|
||||
pos += offset_size as usize;
|
||||
verify_checksum(file_data, offset, offset + pos)?;
|
||||
|
||||
Ok(ExtensibleArrayHeader {
|
||||
client_id,
|
||||
@@ -270,6 +297,40 @@ fn index_to_chunk_offsets(
|
||||
|
||||
/// Collect elements from a data block at the given offset.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
/// Layout of super block `u`, per the HDF5 spec: the number of data blocks it
|
||||
/// owns and how many elements each of them holds.
|
||||
///
|
||||
/// `ndblks` and `dblk_nelmts` each double every *other* level, a half-step
|
||||
/// apart, so the blocks grow as 1x16, 1x32, 2x32, 2x64, 4x64 ... for a
|
||||
/// 16-element minimum. Treating either as doubling every level (the previous
|
||||
/// implementation) puts every element after the first data block at the wrong
|
||||
/// index.
|
||||
fn sblk_info(u: usize, data_blk_min_elmts: usize) -> Option<(usize, usize)> {
|
||||
let ndblks = 1usize.checked_shl((u / 2) as u32)?;
|
||||
let dblk_nelmts = 1usize
|
||||
.checked_shl(u.div_ceil(2) as u32)?
|
||||
.checked_mul(data_blk_min_elmts)?;
|
||||
Some((ndblks, dblk_nelmts))
|
||||
}
|
||||
|
||||
/// Width of the "offset of the block in the array" field carried by super and
|
||||
/// data blocks (`hdr->arr_off_size`).
|
||||
fn arr_off_size(header: &ExtensibleArrayHeader) -> usize {
|
||||
(header.max_nelmts_bits as usize).div_ceil(8)
|
||||
}
|
||||
|
||||
/// Elements per data block page, once a data block is large enough to be paged.
|
||||
fn page_nelmts(header: &ExtensibleArrayHeader) -> Option<usize> {
|
||||
1usize.checked_shl(u32::from(header.max_dblk_nelmts_bits))
|
||||
}
|
||||
|
||||
/// Read the elements of one data block (EADB).
|
||||
///
|
||||
/// `page_init` is the owning super block's page-init bitmap and `first_page`
|
||||
/// this block's first bit in it; both are only consulted when the block is
|
||||
/// paged. The bitmap lives in the super block, not here — a paged data block
|
||||
/// stores only its prefix, then one slot per page.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn read_data_block_elements(
|
||||
file_data: &[u8],
|
||||
db_offset: usize,
|
||||
@@ -280,117 +341,101 @@ fn read_data_block_elements(
|
||||
start_index: usize,
|
||||
num_chunks_per_dim: &[u64],
|
||||
chunk_dimensions: &[u32],
|
||||
page_init: &[u8],
|
||||
first_page: usize,
|
||||
) -> Result<Vec<ChunkInfo>, FormatError> {
|
||||
// AEDB: signature(4) + version(1) + client_id(1) + header_address(offset_size)
|
||||
let db_header_size = 4 + 1 + 1 + offset_size as usize;
|
||||
// EADB: signature(4) + version(1) + client_id(1) + header_address(offset_size)
|
||||
// + block offset(arr_off_size)
|
||||
let db_header_size = 4 + 1 + 1 + offset_size as usize + arr_off_size(header);
|
||||
ensure_len(file_data, db_offset, db_header_size)?;
|
||||
|
||||
let d = &file_data[db_offset..];
|
||||
if &d[0..4] != b"EADB" {
|
||||
if &file_data[db_offset..db_offset + 4] != b"EADB" {
|
||||
return Err(FormatError::ChunkedReadError(
|
||||
"invalid Extensible Array data block signature".into(),
|
||||
));
|
||||
}
|
||||
// Skip version(1) + client_id(1) + header_address(offset_size) + block_offset
|
||||
// Block offset is encoded in ceil(max_nelmts_bits/8) bytes
|
||||
let blk_off_size = (header.max_nelmts_bits as usize).div_ceil(8);
|
||||
let mut pos = db_offset + db_header_size + blk_off_size;
|
||||
|
||||
// Check if paged
|
||||
if header.max_nelmts_bits >= usize::BITS as u8 {
|
||||
return Err(FormatError::Overflow(
|
||||
"max_nelmts_bits exceeds usize bit width".into(),
|
||||
));
|
||||
}
|
||||
let page_nelmts = 1usize << header.max_nelmts_bits;
|
||||
let is_paged = nelmts > page_nelmts;
|
||||
let mut pos = db_offset + db_header_size;
|
||||
let page = page_nelmts(header).ok_or_else(|| {
|
||||
FormatError::Overflow("Extensible Array page element count overflows usize".into())
|
||||
})?;
|
||||
|
||||
let mut chunks = Vec::new();
|
||||
|
||||
if !is_paged {
|
||||
for i in 0..nelmts {
|
||||
let read_run = |from: usize,
|
||||
count: usize,
|
||||
first_index: usize,
|
||||
chunks: &mut Vec<ChunkInfo>|
|
||||
-> Result<usize, FormatError> {
|
||||
let mut p = from;
|
||||
for i in 0..count {
|
||||
let (info, consumed) = read_element(
|
||||
file_data,
|
||||
pos,
|
||||
p,
|
||||
header.client_id,
|
||||
header.element_size,
|
||||
offset_size,
|
||||
chunk_byte_size,
|
||||
start_index + i,
|
||||
first_index + i,
|
||||
num_chunks_per_dim,
|
||||
chunk_dimensions,
|
||||
)?;
|
||||
if let Some(ci) = info {
|
||||
chunks.push(ci);
|
||||
}
|
||||
pos += consumed;
|
||||
p += consumed;
|
||||
}
|
||||
} else {
|
||||
// Paged: elements are split into pages of page_nelmts.
|
||||
// After the data block header comes a page bitmap, then each page
|
||||
// has page_nelmts elements followed by a 4-byte checksum.
|
||||
let npages = nelmts.div_ceil(page_nelmts);
|
||||
// Page bitmap: ceil(npages / 8) bytes
|
||||
let bitmap_size = npages.div_ceil(8);
|
||||
// Read bitmap
|
||||
if pos + bitmap_size > file_data.len() {
|
||||
return Err(FormatError::UnexpectedEof {
|
||||
expected: pos + bitmap_size,
|
||||
available: file_data.len(),
|
||||
});
|
||||
}
|
||||
let bitmap = &file_data[pos..pos + bitmap_size];
|
||||
pos += bitmap_size;
|
||||
Ok(p)
|
||||
};
|
||||
|
||||
if nelmts <= page {
|
||||
// Prefix and elements are covered by one checksum.
|
||||
let elem_bytes = if header.client_id == 0 {
|
||||
offset_size as usize
|
||||
} else {
|
||||
header.element_size as usize
|
||||
};
|
||||
let end = nelmts
|
||||
.checked_mul(elem_bytes)
|
||||
.and_then(|b| pos.checked_add(b))
|
||||
.ok_or_else(|| FormatError::Overflow("Extensible Array data block span".into()))?;
|
||||
verify_checksum(file_data, db_offset, end)?;
|
||||
read_run(pos, nelmts, start_index, &mut chunks)?;
|
||||
return Ok(chunks);
|
||||
}
|
||||
|
||||
let mut global_idx = start_index;
|
||||
for page_idx in 0..npages {
|
||||
let byte_idx = page_idx / 8;
|
||||
let bit_idx = page_idx % 8;
|
||||
let page_has_data = (bitmap[byte_idx] >> bit_idx) & 1 != 0;
|
||||
|
||||
let elems_this_page = if page_idx == npages - 1 {
|
||||
let remainder = nelmts % page_nelmts;
|
||||
if remainder == 0 {
|
||||
page_nelmts
|
||||
} else {
|
||||
remainder
|
||||
}
|
||||
} else {
|
||||
page_nelmts
|
||||
};
|
||||
|
||||
if page_has_data {
|
||||
for i in 0..elems_this_page {
|
||||
let (info, consumed) = read_element(
|
||||
file_data,
|
||||
pos,
|
||||
header.client_id,
|
||||
header.element_size,
|
||||
offset_size,
|
||||
chunk_byte_size,
|
||||
global_idx + i,
|
||||
num_chunks_per_dim,
|
||||
chunk_dimensions,
|
||||
)?;
|
||||
if let Some(ci) = info {
|
||||
chunks.push(ci);
|
||||
}
|
||||
pos += consumed;
|
||||
}
|
||||
// Skip page checksum (4 bytes)
|
||||
pos += 4;
|
||||
} else {
|
||||
// Empty page: skip all elements + checksum
|
||||
pos += elems_this_page * elem_bytes + 4;
|
||||
}
|
||||
global_idx += elems_this_page;
|
||||
// Paged: the prefix ends with its own checksum, then one slot per page,
|
||||
// each holding `page` elements followed by a checksum. Pages whose bit is
|
||||
// clear were never written; their slot still occupies the file, so stride
|
||||
// over it rather than reading zeros as addresses.
|
||||
verify_checksum(file_data, db_offset, pos)?;
|
||||
pos += 4;
|
||||
let elem_bytes = if header.client_id == 0 {
|
||||
offset_size as usize
|
||||
} else {
|
||||
header.element_size as usize
|
||||
};
|
||||
let page_stride = page
|
||||
.checked_mul(elem_bytes)
|
||||
.and_then(|b| b.checked_add(4))
|
||||
.ok_or_else(|| FormatError::Overflow("Extensible Array page stride".into()))?;
|
||||
let npages = nelmts.div_ceil(page);
|
||||
for p in 0..npages {
|
||||
// One bit per page across the whole super block, packed contiguously
|
||||
// and MSB-first within each byte, as H5VM_bit_get reads it.
|
||||
let bit = first_page + p;
|
||||
let initialised = page_init
|
||||
.get(bit / 8)
|
||||
.is_some_and(|byte| byte & (0x80 >> (bit % 8)) != 0);
|
||||
if initialised {
|
||||
let count = core::cmp::min(page, nelmts - p * page);
|
||||
// Each page carries its own checksum, over a full page's worth of
|
||||
// slots even when the last one holds fewer live elements.
|
||||
verify_checksum(file_data, pos, pos + page * elem_bytes)?;
|
||||
read_run(pos, count, start_index + p * page, &mut chunks)?;
|
||||
}
|
||||
pos = pos
|
||||
.checked_add(page_stride)
|
||||
.ok_or_else(|| FormatError::Overflow("Extensible Array page offset".into()))?;
|
||||
}
|
||||
|
||||
Ok(chunks)
|
||||
@@ -427,30 +472,83 @@ pub fn read_extensible_array_chunks(
|
||||
let chunk_byte_size: u64 =
|
||||
chunk_dimensions.iter().map(|&d| d as u64).product::<u64>() * element_size as u64;
|
||||
|
||||
// Parse index block (AEIB)
|
||||
// Parse index block (EAIB): signature(4) + version(1) + client_id(1)
|
||||
// + header address(offset_size), then the inline elements, then the
|
||||
// direct data block addresses, then the super block addresses.
|
||||
let ib_offset = header.index_block_address as usize;
|
||||
let ib_header_size = 4 + 1 + 1 + offset_size as usize; // sig + ver + client + hdr_addr
|
||||
let ib_header_size = 4 + 1 + 1 + os;
|
||||
ensure_len(file_data, ib_offset, ib_header_size)?;
|
||||
|
||||
let ib = &file_data[ib_offset..];
|
||||
if &ib[0..4] != b"EAIB" {
|
||||
if &file_data[ib_offset..ib_offset + 4] != b"EAIB" {
|
||||
return Err(FormatError::ChunkedReadError(
|
||||
"invalid Extensible Array index block signature".into(),
|
||||
));
|
||||
}
|
||||
// Skip version(1) + client_id(1) + header_address(offset_size)
|
||||
let mut pos = ib_offset + ib_header_size;
|
||||
|
||||
let mut chunks = Vec::new();
|
||||
let mut global_index = 0usize;
|
||||
let total_elements = header.num_elements as usize;
|
||||
|
||||
// 1. Read inline elements in index block
|
||||
let n_inline = header.idx_blk_elmts as usize;
|
||||
for i in 0..n_inline {
|
||||
if global_index + i >= total_elements {
|
||||
break;
|
||||
let dmin = header.min_dblk_nelmts as usize;
|
||||
if dmin == 0 || !dmin.is_power_of_two() {
|
||||
return Err(FormatError::ChunkedReadError(
|
||||
"Extensible Array data block minimum is not a power of two".into(),
|
||||
));
|
||||
}
|
||||
// nsblks = 1 + (max_nelmts_bits - log2(data_blk_min_elmts)), and the index
|
||||
// block holds 2 * (sup_blk_min_data_ptrs - 1) data block addresses.
|
||||
let log2_dmin = dmin.trailing_zeros() as usize;
|
||||
let nsblks = 1 + (header.max_nelmts_bits as usize).saturating_sub(log2_dmin);
|
||||
let ndblk_addrs = 2 * (header.super_blk_min_nelmts as usize).saturating_sub(1);
|
||||
|
||||
// The data blocks listed directly in the index block are the first
|
||||
// `ndblk_addrs` in super-block order, each sized by the level it belongs
|
||||
// to; the super block addresses that follow resume at the next level.
|
||||
let mut direct: Vec<usize> = Vec::with_capacity(ndblk_addrs);
|
||||
let mut level = 0usize;
|
||||
while direct.len() < ndblk_addrs {
|
||||
if level >= nsblks {
|
||||
return Err(FormatError::ChunkedReadError(
|
||||
"Extensible Array index block claims more data blocks than the array has".into(),
|
||||
));
|
||||
}
|
||||
let (ndblks, dblk_nelmts) = sblk_info(level, dmin).ok_or_else(|| {
|
||||
FormatError::Overflow("Extensible Array super block layout overflows usize".into())
|
||||
})?;
|
||||
for _ in 0..ndblks {
|
||||
direct.push(dblk_nelmts);
|
||||
}
|
||||
level += 1;
|
||||
}
|
||||
if direct.len() != ndblk_addrs {
|
||||
// A partial level in the index block is not a layout HDF5 produces,
|
||||
// and guessing where the super blocks resume would misplace elements.
|
||||
return Err(FormatError::ChunkedReadError(
|
||||
"Extensible Array index block ends mid super block".into(),
|
||||
));
|
||||
}
|
||||
|
||||
// One checksum covers the prefix, every inline element slot, and every
|
||||
// data block and super block address.
|
||||
let elem_bytes = if header.client_id == 0 {
|
||||
os
|
||||
} else {
|
||||
header.element_size as usize
|
||||
};
|
||||
let ib_end = (header.idx_blk_elmts as usize)
|
||||
.checked_mul(elem_bytes)
|
||||
.and_then(|b| pos.checked_add(b))
|
||||
.and_then(|p| {
|
||||
ndblk_addrs
|
||||
.checked_add(nsblks - level)
|
||||
.and_then(|n| n.checked_mul(os).and_then(|b| p.checked_add(b)))
|
||||
})
|
||||
.ok_or_else(|| FormatError::Overflow("Extensible Array index block span".into()))?;
|
||||
verify_checksum(file_data, ib_offset, ib_end)?;
|
||||
|
||||
// 1. Elements stored inline in the index block.
|
||||
let n_inline = (header.idx_blk_elmts as usize).min(total_elements);
|
||||
for i in 0..n_inline {
|
||||
let (info, consumed) = read_element(
|
||||
file_data,
|
||||
pos,
|
||||
@@ -458,7 +556,7 @@ pub fn read_extensible_array_chunks(
|
||||
header.element_size,
|
||||
offset_size,
|
||||
chunk_byte_size,
|
||||
global_index + i,
|
||||
i,
|
||||
&num_chunks_per_dim,
|
||||
chunk_dimensions,
|
||||
)?;
|
||||
@@ -467,154 +565,90 @@ pub fn read_extensible_array_chunks(
|
||||
}
|
||||
pos += consumed;
|
||||
}
|
||||
global_index += n_inline.min(total_elements);
|
||||
|
||||
// If all elements were inline, we're done
|
||||
let mut global_index = n_inline;
|
||||
if global_index >= total_elements {
|
||||
return Ok(chunks);
|
||||
}
|
||||
|
||||
// Compute data block and super block counts
|
||||
let min_dblk = header.min_dblk_nelmts as usize;
|
||||
let sblk_min = header.super_blk_min_nelmts as usize;
|
||||
|
||||
// The first sblk_min super block levels have their data blocks listed directly
|
||||
// in the index block. Compute their sizes.
|
||||
let mut n_direct_dblks = 0usize;
|
||||
let mut dblk_sizes: Vec<usize> = Vec::new();
|
||||
{
|
||||
let mut nelmts = min_dblk;
|
||||
for sb_level in 0..sblk_min {
|
||||
if sb_level >= usize::BITS as usize {
|
||||
return Err(FormatError::Overflow(
|
||||
"sb_level exceeds usize bit width".into(),
|
||||
// 2. Data blocks listed directly in the index block.
|
||||
for &dblk_nelmts in &direct {
|
||||
if global_index >= total_elements {
|
||||
return Ok(chunks);
|
||||
}
|
||||
ensure_len(file_data, pos, os)?;
|
||||
let addr = read_offset(file_data, pos, offset_size)?;
|
||||
pos += os;
|
||||
if !is_undefined_addr(addr, offset_size) {
|
||||
if dblk_nelmts > page_nelmts(header).unwrap_or(usize::MAX) {
|
||||
// Would need a page-init bitmap, which only a super block
|
||||
// carries. HDF5 never pages these small early blocks.
|
||||
return Err(FormatError::ChunkedReadError(
|
||||
"Extensible Array index block references a paged data block".into(),
|
||||
));
|
||||
}
|
||||
let ndblks = 1usize << sb_level;
|
||||
for _ in 0..ndblks {
|
||||
dblk_sizes.push(nelmts);
|
||||
n_direct_dblks += 1;
|
||||
}
|
||||
if sb_level > 0 {
|
||||
nelmts *= 2;
|
||||
}
|
||||
chunks.extend(read_data_block_elements(
|
||||
file_data,
|
||||
addr as usize,
|
||||
dblk_nelmts,
|
||||
header,
|
||||
offset_size,
|
||||
chunk_byte_size,
|
||||
global_index,
|
||||
&num_chunks_per_dim,
|
||||
chunk_dimensions,
|
||||
&[],
|
||||
0,
|
||||
)?);
|
||||
}
|
||||
global_index += dblk_nelmts;
|
||||
}
|
||||
|
||||
// Read direct data block addresses from index block
|
||||
let mut dblk_addrs: Vec<u64> = Vec::with_capacity(n_direct_dblks);
|
||||
for _ in 0..n_direct_dblks {
|
||||
if pos + os > file_data.len() {
|
||||
// 3. Everything else lives in super blocks, one address per remaining
|
||||
// level, starting at the level after the direct data blocks.
|
||||
for u in level..nsblks {
|
||||
if global_index >= total_elements {
|
||||
break;
|
||||
}
|
||||
let addr = read_offset(file_data, pos, offset_size)?;
|
||||
dblk_addrs.push(addr);
|
||||
ensure_len(file_data, pos, os)?;
|
||||
let sb_addr = read_offset(file_data, pos, offset_size)?;
|
||||
pos += os;
|
||||
}
|
||||
|
||||
// Read elements from direct data blocks
|
||||
for (i, &addr) in dblk_addrs.iter().enumerate() {
|
||||
if i >= dblk_sizes.len() {
|
||||
break;
|
||||
let (ndblks, dblk_nelmts) = sblk_info(u, dmin).ok_or_else(|| {
|
||||
FormatError::Overflow("Extensible Array super block layout overflows usize".into())
|
||||
})?;
|
||||
if !is_undefined_addr(sb_addr, offset_size) {
|
||||
chunks.extend(read_super_block(
|
||||
file_data,
|
||||
sb_addr as usize,
|
||||
ndblks,
|
||||
dblk_nelmts,
|
||||
header,
|
||||
offset_size,
|
||||
chunk_byte_size,
|
||||
global_index,
|
||||
&num_chunks_per_dim,
|
||||
chunk_dimensions,
|
||||
)?);
|
||||
}
|
||||
let nelmts = dblk_sizes[i];
|
||||
if is_undefined_addr(addr, offset_size) {
|
||||
global_index += nelmts;
|
||||
continue;
|
||||
}
|
||||
let block_chunks = read_data_block_elements(
|
||||
file_data,
|
||||
addr as usize,
|
||||
nelmts,
|
||||
header,
|
||||
offset_size,
|
||||
chunk_byte_size,
|
||||
global_index,
|
||||
&num_chunks_per_dim,
|
||||
chunk_dimensions,
|
||||
)?;
|
||||
chunks.extend(block_chunks);
|
||||
global_index += nelmts;
|
||||
}
|
||||
|
||||
// Remaining elements are in super blocks
|
||||
let total_in_ib_and_direct: usize = n_inline + dblk_sizes.iter().sum::<usize>();
|
||||
if total_elements <= total_in_ib_and_direct {
|
||||
return Ok(chunks);
|
||||
}
|
||||
let remaining_elements = total_elements - total_in_ib_and_direct;
|
||||
|
||||
// Compute super block layout
|
||||
let mut sb_addrs: Vec<u64> = Vec::new();
|
||||
let mut sb_infos: Vec<(usize, usize)> = Vec::new();
|
||||
{
|
||||
let mut covered = 0usize;
|
||||
let mut sb_level = sblk_min;
|
||||
let mut nelmts_per_dblk = min_dblk;
|
||||
for lev in 0..sblk_min {
|
||||
if lev > 0 {
|
||||
nelmts_per_dblk *= 2;
|
||||
}
|
||||
}
|
||||
|
||||
while covered < remaining_elements {
|
||||
if sb_level >= usize::BITS as usize {
|
||||
return Err(FormatError::Overflow(
|
||||
"sb_level exceeds usize bit width".into(),
|
||||
));
|
||||
}
|
||||
let ndblks = 1usize << sb_level;
|
||||
nelmts_per_dblk *= 2;
|
||||
let total_in_sb = ndblks * nelmts_per_dblk;
|
||||
sb_infos.push((ndblks, nelmts_per_dblk));
|
||||
covered += total_in_sb;
|
||||
sb_level += 1;
|
||||
}
|
||||
}
|
||||
|
||||
// Read super block addresses from index block
|
||||
for _ in 0..sb_infos.len() {
|
||||
if pos + os > file_data.len() {
|
||||
break;
|
||||
}
|
||||
let addr = read_offset(file_data, pos, offset_size)?;
|
||||
sb_addrs.push(addr);
|
||||
pos += os;
|
||||
}
|
||||
|
||||
// Process each super block
|
||||
for (sb_idx, &sb_addr) in sb_addrs.iter().enumerate() {
|
||||
let (ndblks, nelmts_per_dblk) = sb_infos[sb_idx];
|
||||
if is_undefined_addr(sb_addr, offset_size) {
|
||||
global_index += ndblks * nelmts_per_dblk;
|
||||
continue;
|
||||
}
|
||||
let sb_chunks = read_super_block(
|
||||
file_data,
|
||||
sb_addr as usize,
|
||||
ndblks,
|
||||
nelmts_per_dblk,
|
||||
header,
|
||||
offset_size,
|
||||
chunk_byte_size,
|
||||
global_index,
|
||||
&num_chunks_per_dim,
|
||||
chunk_dimensions,
|
||||
)?;
|
||||
chunks.extend(sb_chunks);
|
||||
global_index += ndblks * nelmts_per_dblk;
|
||||
global_index =
|
||||
global_index.saturating_add(ndblks.checked_mul(dblk_nelmts).ok_or_else(|| {
|
||||
FormatError::Overflow("Extensible Array super block span".into())
|
||||
})?);
|
||||
}
|
||||
|
||||
Ok(chunks)
|
||||
}
|
||||
|
||||
/// Read a super block (AESB) and its data blocks.
|
||||
/// Read a super block (EASB) and the data blocks it owns.
|
||||
///
|
||||
/// On disk: signature(4) + version(1) + client_id(1) + header address
|
||||
/// + block offset + the page-init bitmap for every data block it owns
|
||||
/// + one address per data block + checksum.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn read_super_block(
|
||||
file_data: &[u8],
|
||||
sb_offset: usize,
|
||||
ndblks: usize,
|
||||
nelmts_per_dblk: usize,
|
||||
dblk_nelmts: usize,
|
||||
header: &ExtensibleArrayHeader,
|
||||
offset_size: u8,
|
||||
chunk_byte_size: u64,
|
||||
@@ -623,9 +657,7 @@ fn read_super_block(
|
||||
chunk_dimensions: &[u32],
|
||||
) -> Result<Vec<ChunkInfo>, FormatError> {
|
||||
let os = offset_size as usize;
|
||||
|
||||
// AESB: signature(4) + version(1) + client_id(1) + header_address(offset_size)
|
||||
let sb_header_size = 4 + 1 + 1 + os;
|
||||
let sb_header_size = 4 + 1 + 1 + os + arr_off_size(header);
|
||||
ensure_len(file_data, sb_offset, sb_header_size)?;
|
||||
|
||||
if &file_data[sb_offset..sb_offset + 4] != b"EASB" {
|
||||
@@ -634,43 +666,57 @@ fn read_super_block(
|
||||
));
|
||||
}
|
||||
|
||||
let mut pos = sb_offset + sb_header_size;
|
||||
|
||||
// Read data block addresses
|
||||
let mut dblk_addrs: Vec<u64> = Vec::with_capacity(ndblks);
|
||||
for _ in 0..ndblks {
|
||||
if pos + os > file_data.len() {
|
||||
return Err(FormatError::UnexpectedEof {
|
||||
expected: pos + os,
|
||||
available: file_data.len(),
|
||||
});
|
||||
}
|
||||
let addr = read_offset(file_data, pos, offset_size)?;
|
||||
dblk_addrs.push(addr);
|
||||
pos += os;
|
||||
}
|
||||
// Page-init bitmap: one bit per page, `npages` bits per data block, packed
|
||||
// contiguously. HDF5 sizes the buffer `ndblks * ceil(npages / 8)`, which
|
||||
// is bigger than the bits need when `npages` is not a multiple of eight.
|
||||
// Zero-sized unless this level's data blocks are paged.
|
||||
let page = page_nelmts(header).ok_or_else(|| {
|
||||
FormatError::Overflow("Extensible Array page element count overflows usize".into())
|
||||
})?;
|
||||
let npages = if dblk_nelmts > page {
|
||||
dblk_nelmts / page
|
||||
} else {
|
||||
0
|
||||
};
|
||||
let per_dblk_bitmap = npages.div_ceil(8);
|
||||
let bitmap_bytes = per_dblk_bitmap
|
||||
.checked_mul(ndblks)
|
||||
.ok_or_else(|| FormatError::Overflow("Extensible Array page bitmap size".into()))?;
|
||||
let bitmap_start = sb_offset + sb_header_size;
|
||||
ensure_len(file_data, bitmap_start, bitmap_bytes)?;
|
||||
let bitmap = &file_data[bitmap_start..bitmap_start + bitmap_bytes];
|
||||
|
||||
let mut pos = bitmap_start + bitmap_bytes;
|
||||
let mut chunks = Vec::new();
|
||||
let mut global_idx = start_index;
|
||||
|
||||
for &addr in &dblk_addrs {
|
||||
if is_undefined_addr(addr, offset_size) {
|
||||
global_idx += nelmts_per_dblk;
|
||||
continue;
|
||||
// One checksum covers the prefix, the bitmap and every data block address.
|
||||
let sb_end = ndblks
|
||||
.checked_mul(os)
|
||||
.and_then(|b| pos.checked_add(b))
|
||||
.ok_or_else(|| FormatError::Overflow("Extensible Array super block span".into()))?;
|
||||
verify_checksum(file_data, sb_offset, sb_end)?;
|
||||
|
||||
for i in 0..ndblks {
|
||||
ensure_len(file_data, pos, os)?;
|
||||
let addr = read_offset(file_data, pos, offset_size)?;
|
||||
pos += os;
|
||||
if !is_undefined_addr(addr, offset_size) {
|
||||
chunks.extend(read_data_block_elements(
|
||||
file_data,
|
||||
addr as usize,
|
||||
dblk_nelmts,
|
||||
header,
|
||||
offset_size,
|
||||
chunk_byte_size,
|
||||
global_idx,
|
||||
num_chunks_per_dim,
|
||||
chunk_dimensions,
|
||||
bitmap,
|
||||
i * npages,
|
||||
)?);
|
||||
}
|
||||
let block_chunks = read_data_block_elements(
|
||||
file_data,
|
||||
addr as usize,
|
||||
nelmts_per_dblk,
|
||||
header,
|
||||
offset_size,
|
||||
chunk_byte_size,
|
||||
global_idx,
|
||||
num_chunks_per_dim,
|
||||
chunk_dimensions,
|
||||
)?;
|
||||
chunks.extend(block_chunks);
|
||||
global_idx += nelmts_per_dblk;
|
||||
global_idx += dblk_nelmts;
|
||||
}
|
||||
|
||||
Ok(chunks)
|
||||
@@ -679,6 +725,14 @@ fn read_super_block(
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
/// Stamp the Jenkins checksum a real file would carry over
|
||||
/// `data[start..end]`, writing it at `end`. Hand-built fixtures need this
|
||||
/// now that the reader validates it, exactly as HDF5 writes it.
|
||||
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]
|
||||
fn index_to_offsets_1d() {
|
||||
let num_chunks = vec![5u64];
|
||||
@@ -734,6 +788,7 @@ mod tests {
|
||||
buf[44..52].copy_from_slice(&5u64.to_le_bytes()); // stat[4] = num_elements
|
||||
buf[52..60].copy_from_slice(&0u64.to_le_bytes()); // stat[5]
|
||||
buf[60..68].copy_from_slice(&0x1000u64.to_le_bytes()); // index_block_address
|
||||
stamp_checksum(&mut buf, 0, 68);
|
||||
|
||||
let hdr = ExtensibleArrayHeader::parse(&buf, 0, os, ls).unwrap();
|
||||
assert_eq!(hdr.client_id, 0);
|
||||
@@ -819,6 +874,7 @@ mod tests {
|
||||
.copy_from_slice(&(num_chunks as u64).to_le_bytes());
|
||||
file_data[aehd_offset + 60..aehd_offset + 68]
|
||||
.copy_from_slice(&(aeib_offset as u64).to_le_bytes());
|
||||
stamp_checksum(&mut file_data, aehd_offset, aehd_offset + 68);
|
||||
// checksum (4 bytes at +68) — not validated
|
||||
|
||||
// Build AEIB at aeib_offset
|
||||
@@ -836,6 +892,23 @@ mod tests {
|
||||
let p = elem_start + i * osv;
|
||||
file_data[p..p + osv].copy_from_slice(&addr.to_le_bytes());
|
||||
}
|
||||
// The index block's checksum covers its prefix, every inline element
|
||||
// slot, and every data block and super block address slot:
|
||||
// ndblk_addrs = 2 * (sup_blk_min_data_ptrs - 1), and the super block
|
||||
// pointers make up the rest of nsblks levels.
|
||||
let sup_ptrs = file_data[aehd_offset + 10] as usize;
|
||||
let dmin = file_data[aehd_offset + 9] as usize;
|
||||
let nsblks = 1 + 10 - dmin.trailing_zeros() as usize;
|
||||
let ndblk_addrs = 2 * (sup_ptrs - 1);
|
||||
// Levels consumed by those direct data blocks (1, 1, 2, 2, ... per level).
|
||||
let mut consumed = 0usize;
|
||||
let mut levels = 0usize;
|
||||
while consumed < ndblk_addrs {
|
||||
consumed += 1 << (levels / 2);
|
||||
levels += 1;
|
||||
}
|
||||
let ib_end = elem_start + num_chunks * osv + (ndblk_addrs + nsblks - levels) * osv;
|
||||
stamp_checksum(&mut file_data, aeib_offset, ib_end);
|
||||
|
||||
let header = ExtensibleArrayHeader::parse(&file_data, aehd_offset, os, ls).unwrap();
|
||||
let ds_dims = vec![40u64]; // 2 chunks × 20 elements
|
||||
@@ -885,6 +958,7 @@ mod tests {
|
||||
// idx_blk_addr at offset 12 + 6*8 = 60
|
||||
file_data[aehd_offset + 60..aehd_offset + 68]
|
||||
.copy_from_slice(&(aeib_offset as u64).to_le_bytes());
|
||||
stamp_checksum(&mut file_data, aehd_offset, aehd_offset + 68);
|
||||
|
||||
// AEIB
|
||||
file_data[aeib_offset..aeib_offset + 4].copy_from_slice(b"EAIB");
|
||||
@@ -903,42 +977,48 @@ mod tests {
|
||||
pos += osv;
|
||||
}
|
||||
|
||||
// Direct data block addresses: first sb_level=0 has 1 dblk, sb_level=1 has 1 dblk
|
||||
// Total direct dblks for sblk_min=2: 2^0 + 2^1 = 1 + 2 = 3 (oops)
|
||||
// Actually: sblk_min levels. level 0: 2^0=1 dblk, level 1: 2^1=2 dblks => 3 dblks
|
||||
// But we only have 2 remaining elements.
|
||||
// dblk sizes: level 0: 1 dblk of min_dblk=2; level 1: 2 dblks of 2 each (nelmts doubles at level > 0)
|
||||
// Wait, re-reading the code: at level 0, nelmts=min_dblk=2, 1 dblk.
|
||||
// At level 1, 1 dblk, nelmts still 2 (doubles only at level > 0... but the code says
|
||||
// `if sb_level > 0 { nelmts *= 2 }` after pushing). Let me re-check.
|
||||
// After push at level 0: nelmts=2. Then if 0>0 false, no double. Push 1 dblk of 2.
|
||||
// Level 1: ndblks=2. Push 2 dblks of 2. Then 1>0 true, nelmts=4.
|
||||
// Total: 3 dblks with sizes [2, 2, 2]. Total = 6.
|
||||
// We only need 2 more elements. So only the first dblk has data.
|
||||
let n_direct_dblks = 3;
|
||||
// Direct data block addresses. With sup_blk_min_data_ptrs = 2 the index
|
||||
// block holds 2 * (2 - 1) = 2 of them, which are the data blocks of
|
||||
// super block levels 0 and 1: one of `min_dblk_nelmts` elements, then
|
||||
// one of twice that (ndblks = 2^(u/2), dblk_nelmts = 2^((u+1)/2) * min).
|
||||
// Only the first is allocated here; the rest of the array is empty.
|
||||
let ndblk_addrs = 2 * (sblk_min as usize - 1);
|
||||
file_data[pos..pos + osv].copy_from_slice(&(aedb_offset as u64).to_le_bytes());
|
||||
pos += osv;
|
||||
// 2 more dblk addresses - undefined
|
||||
for _ in 1..n_direct_dblks {
|
||||
for _ in 1..ndblk_addrs {
|
||||
file_data[pos..pos + osv].copy_from_slice(&u64::MAX.to_le_bytes());
|
||||
pos += osv;
|
||||
}
|
||||
// Super block addresses fill the remaining levels; all unallocated.
|
||||
let nsblks = 1 + 10 - (min_dblk_nelmts as usize).trailing_zeros() as usize;
|
||||
let mut consumed = 0usize;
|
||||
let mut levels = 0usize;
|
||||
while consumed < ndblk_addrs {
|
||||
consumed += 1 << (levels / 2);
|
||||
levels += 1;
|
||||
}
|
||||
for _ in 0..(nsblks - levels) {
|
||||
file_data[pos..pos + osv].copy_from_slice(&u64::MAX.to_le_bytes());
|
||||
pos += osv;
|
||||
}
|
||||
stamp_checksum(&mut file_data, aeib_offset, pos);
|
||||
|
||||
// EADB at aedb_offset (min_dblk_nelmts elements)
|
||||
// EADB holding the first data block's `min_dblk_nelmts` elements.
|
||||
file_data[aedb_offset..aedb_offset + 4].copy_from_slice(b"EADB");
|
||||
file_data[aedb_offset + 4] = 0;
|
||||
file_data[aedb_offset + 5] = 0;
|
||||
file_data[aedb_offset + 6..aedb_offset + 14]
|
||||
.copy_from_slice(&(aehd_offset as u64).to_le_bytes());
|
||||
// block_offset: ceil(max_nelmts_bits/8) = ceil(10/8) = 2 bytes
|
||||
// block_offset = 0 for first data block
|
||||
let blk_off_size = (10usize).div_ceil(8); // max_nelmts_bits=10
|
||||
let mut dbpos = aedb_offset + 6 + osv + blk_off_size;
|
||||
// Block offset field: ceil(max_nelmts_bits / 8) bytes, zero here.
|
||||
let blk_off_size = (10usize).div_ceil(8);
|
||||
let db_elems = aedb_offset + 6 + osv + blk_off_size;
|
||||
let mut dbpos = db_elems;
|
||||
for i in 0..min_dblk_nelmts as usize {
|
||||
let addr = base_addr + (idx_blk_elmts as u64 + i as u64) * chunk_byte_size;
|
||||
file_data[dbpos..dbpos + osv].copy_from_slice(&addr.to_le_bytes());
|
||||
dbpos += osv;
|
||||
}
|
||||
stamp_checksum(&mut file_data, aedb_offset, dbpos);
|
||||
|
||||
let header = ExtensibleArrayHeader::parse(&file_data, aehd_offset, os, ls).unwrap();
|
||||
let ds_dims = vec![40u64];
|
||||
|
||||
@@ -86,6 +86,12 @@ pub(crate) fn build_dataset_oh(
|
||||
let mut dl = Vec::new();
|
||||
dl.push(4); // version
|
||||
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_size.to_le_bytes());
|
||||
w.add_message(MessageType::DataLayout, dl);
|
||||
|
||||
@@ -629,21 +629,70 @@ fn deflate_decompress(data: &[u8], expected_bytes: usize) -> Result<Vec<u8>, For
|
||||
// Fall through to flate2 on error
|
||||
}
|
||||
|
||||
use std::io::Read;
|
||||
let decoder = flate2::read::ZlibDecoder::new(data);
|
||||
let mut result = Vec::with_capacity(limit.min(1 << 20));
|
||||
// Read one byte past the limit so an over-size stream is distinguishable
|
||||
// A chunk's decompressed size is known, so allocate it once; without one,
|
||||
// start from a multiple of the input and grow.
|
||||
let size_hint = if expected_bytes != 0 {
|
||||
expected_bytes
|
||||
} else {
|
||||
data.len().saturating_mul(4).min(1 << 20)
|
||||
};
|
||||
inflate_bounded(data, size_hint, limit).map_err(FormatError::DecompressionError)
|
||||
}
|
||||
|
||||
/// Inflate a zlib stream into a buffer sized up front, handing the decoder the
|
||||
/// whole input at once.
|
||||
///
|
||||
/// `flate2::read::ZlibDecoder` feeds its input through a 32 KiB buffer and
|
||||
/// grows the output as it goes; on single chunks that cost zlib-rs up to 3.7x
|
||||
/// against zlib-ng (`BENCHMARKS.md`, "Deflate backend"). Output beyond `limit`
|
||||
/// is an error, as is a stream that ends before its end-of-stream marker (the
|
||||
/// streaming reader returned the bytes it had and no error).
|
||||
#[cfg(feature = "deflate")]
|
||||
pub(crate) 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.
|
||||
decoder
|
||||
.take(limit as u64 + 1)
|
||||
.read_to_end(&mut result)
|
||||
.map_err(|e| FormatError::DecompressionError(e.to_string()))?;
|
||||
if result.len() > limit {
|
||||
return Err(FormatError::DecompressionError(
|
||||
"deflate: output exceeds size limit".into(),
|
||||
));
|
||||
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() => {
|
||||
// Out of room: double, up to the limit.
|
||||
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 => {
|
||||
// Room left, so the decoder stopped for want of input.
|
||||
if inflater.total_in() as usize >= data.len()
|
||||
|| (inflater.total_in(), inflater.total_out()) == (in_before, out_before)
|
||||
{
|
||||
return Err("deflate: truncated stream".into());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
/// Direct FFI to Apple's system libz for fast decompression.
|
||||
@@ -722,14 +771,41 @@ fn deflate_decompress(_data: &[u8], _expected_bytes: usize) -> Result<Vec<u8>, F
|
||||
/// Compress data with zlib.
|
||||
#[cfg(feature = "deflate")]
|
||||
fn deflate_compress(data: &[u8], level: u32) -> Result<Vec<u8>, FormatError> {
|
||||
use std::io::Write;
|
||||
let mut encoder = flate2::write::ZlibEncoder::new(Vec::new(), flate2::Compression::new(level));
|
||||
encoder
|
||||
.write_all(data)
|
||||
.map_err(|e| FormatError::CompressionError(e.to_string()))?;
|
||||
encoder
|
||||
.finish()
|
||||
.map_err(|e| FormatError::CompressionError(e.to_string()))
|
||||
deflate_bounded(data, level).map_err(FormatError::CompressionError)
|
||||
}
|
||||
|
||||
/// Deflate `data` into a zlib stream in one pass, into a buffer sized for the
|
||||
/// worst case up front (the same reasoning as [`inflate_bounded`]).
|
||||
#[cfg(feature = "deflate")]
|
||||
pub(crate) fn deflate_bounded(data: &[u8], level: u32) -> Result<Vec<u8>, String> {
|
||||
use flate2::{Compress, Compression, FlushCompress, Status};
|
||||
|
||||
// zlib's compressBound, plus the zlib header and trailer.
|
||||
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),
|
||||
// The bound should make running out of room unreachable; grow
|
||||
// rather than fail if it happens.
|
||||
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());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(not(feature = "deflate"))]
|
||||
@@ -1833,6 +1909,74 @@ mod tests {
|
||||
assert!(deflate_decompress(&compressed, 64).is_err());
|
||||
}
|
||||
|
||||
#[cfg(feature = "deflate")]
|
||||
fn noisy_bytes(n: usize) -> Vec<u8> {
|
||||
// Compressible but not trivially so.
|
||||
(0..n)
|
||||
.map(|i| ((i as f64 * 0.01).sin() * 127.0 + 128.0) as u8 ^ (i as u8 & 3))
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[cfg(feature = "deflate")]
|
||||
fn deflate_decompress_accepts_output_exactly_at_chunk_size() {
|
||||
let data = noisy_bytes(100_000);
|
||||
let compressed = deflate_compress(&data, 6).unwrap();
|
||||
assert_eq!(deflate_decompress(&compressed, data.len()).unwrap(), data);
|
||||
// One byte short of the real size is over the limit.
|
||||
assert!(deflate_decompress(&compressed, data.len() - 1).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[cfg(feature = "deflate")]
|
||||
fn deflate_decompress_without_size_grows_the_buffer() {
|
||||
// No chunk size: the output starts at 4x the input and has to grow.
|
||||
let data = vec![7u8; 3 * 1024 * 1024];
|
||||
let compressed = deflate_compress(&data, 6).unwrap();
|
||||
assert!(compressed.len() * 4 < data.len());
|
||||
assert_eq!(deflate_decompress(&compressed, 0).unwrap(), data);
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[cfg(feature = "deflate")]
|
||||
fn deflate_decompress_rejects_truncated_stream() {
|
||||
// The streaming reader this replaced returned the bytes it had and no
|
||||
// error, so a truncated chunk read back short.
|
||||
let data = noisy_bytes(100_000);
|
||||
let compressed = deflate_compress(&data, 6).unwrap();
|
||||
for cut in [compressed.len() - 1, compressed.len() / 2, 3] {
|
||||
assert!(
|
||||
deflate_decompress(&compressed[..cut], data.len()).is_err(),
|
||||
"truncated to {cut} of {} bytes",
|
||||
compressed.len()
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[cfg(feature = "deflate")]
|
||||
fn deflate_compress_roundtrips_incompressible_data() {
|
||||
// Random-looking input compresses to slightly more than it started
|
||||
// as; the output must still fit the pre-sized buffer (or grow).
|
||||
let mut x = 0x9E37_79B9_7F4A_7C15u64;
|
||||
let data: Vec<u8> = (0..200_000)
|
||||
.map(|_| {
|
||||
x ^= x << 13;
|
||||
x ^= x >> 7;
|
||||
x ^= x << 17;
|
||||
x as u8
|
||||
})
|
||||
.collect();
|
||||
for level in [0, 1, 6, 9] {
|
||||
let compressed = deflate_compress(&data, level).unwrap();
|
||||
assert_eq!(deflate_decompress(&compressed, data.len()).unwrap(), data);
|
||||
}
|
||||
assert_eq!(
|
||||
deflate_decompress(&deflate_compress(&[], 6).unwrap(), 0).unwrap(),
|
||||
Vec::<u8>::new()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[cfg(feature = "zstd")]
|
||||
fn zstd_decompress_rejects_output_exceeding_chunk_size() {
|
||||
|
||||
@@ -9,6 +9,31 @@ use alloc::{format, vec, vec::Vec};
|
||||
use crate::chunked_read::ChunkInfo;
|
||||
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).
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct FixedArrayHeader {
|
||||
@@ -103,6 +128,8 @@ impl FixedArrayHeader {
|
||||
let num_elements = read_length(d, pos, length_size)?;
|
||||
pos += length_size as usize;
|
||||
let data_block_address = read_offset(d, pos, offset_size)?;
|
||||
pos += offset_size as usize;
|
||||
verify_checksum(file_data, offset, offset + pos)?;
|
||||
|
||||
Ok(FixedArrayHeader {
|
||||
client_id,
|
||||
@@ -223,7 +250,8 @@ pub fn read_fixed_array_chunks(
|
||||
|
||||
if !is_paged {
|
||||
// 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 {
|
||||
push_element(i, elem_at(elements_start, i)?, &mut chunks)?;
|
||||
}
|
||||
@@ -254,6 +282,9 @@ pub fn read_fixed_array_chunks(
|
||||
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 {
|
||||
let page_first = p * page_nelmts; // < num_elements, cannot overflow
|
||||
@@ -270,6 +301,7 @@ pub fn read_fixed_array_chunks(
|
||||
.checked_mul(page_stride)
|
||||
.and_then(|o| pages_start.checked_add(o))
|
||||
.ok_or_else(stride_overflow)?;
|
||||
verify_checksum(file_data, page_off, elem_at(page_off, page_count)?)?;
|
||||
for e in 0..page_count {
|
||||
push_element(page_first + e, elem_at(page_off, e)?, &mut chunks)?;
|
||||
}
|
||||
@@ -374,6 +406,14 @@ fn read_variable_length(data: &[u8], size: usize) -> Result<u64, FormatError> {
|
||||
mod tests {
|
||||
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]
|
||||
fn index_to_offsets_1d() {
|
||||
let num_chunks = vec![5u64];
|
||||
@@ -439,7 +479,7 @@ mod tests {
|
||||
buf[8..16].copy_from_slice(&5u64.to_le_bytes());
|
||||
// data_block_address (offset_size=8)
|
||||
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();
|
||||
assert_eq!(header.client_id, 1);
|
||||
@@ -449,6 +489,54 @@ mod tests {
|
||||
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], &[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]
|
||||
fn parse_fixed_array_header_invalid_signature() {
|
||||
let mut buf = vec![0u8; 256];
|
||||
@@ -469,6 +557,7 @@ mod tests {
|
||||
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 + 16..fahd + 24].copy_from_slice(&0x100u64.to_le_bytes());
|
||||
stamp_checksum(&mut buf, fahd, fahd + 24);
|
||||
// FADB so parsing reaches the paged check
|
||||
let db = 0x100usize;
|
||||
buf[db..db + 4].copy_from_slice(b"FADB");
|
||||
@@ -486,6 +575,8 @@ mod tests {
|
||||
buf[fahd + 7] = 10;
|
||||
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());
|
||||
// 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");
|
||||
let header = FixedArrayHeader::parse(&buf, fahd, 8, 8).unwrap();
|
||||
let r = read_fixed_array_chunks(&buf, &header, &[100], &[20], 8, 8, 8);
|
||||
@@ -545,6 +636,7 @@ mod tests {
|
||||
file_data[fahd_offset + 8..fahd_offset + 16].copy_from_slice(&num_chunks.to_le_bytes());
|
||||
file_data[fahd_offset + 16..fahd_offset + 24]
|
||||
.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
|
||||
file_data[db_offset..db_offset + 4].copy_from_slice(b"FADB");
|
||||
@@ -562,6 +654,7 @@ mod tests {
|
||||
let pos = elem_start + i * os;
|
||||
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 =
|
||||
FixedArrayHeader::parse(&file_data, fahd_offset, offset_size, length_size).unwrap();
|
||||
@@ -611,6 +704,7 @@ mod tests {
|
||||
file_data[fahd_offset + 8..fahd_offset + 16].copy_from_slice(&num_chunks.to_le_bytes());
|
||||
file_data[fahd_offset + 16..fahd_offset + 24]
|
||||
.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 + 4] = 0;
|
||||
@@ -632,6 +726,11 @@ mod tests {
|
||||
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());
|
||||
}
|
||||
stamp_checksum(
|
||||
&mut file_data,
|
||||
db_offset,
|
||||
elem_start + test_chunks.len() * elem_size,
|
||||
);
|
||||
|
||||
let header =
|
||||
FixedArrayHeader::parse(&file_data, fahd_offset, offset_size, length_size).unwrap();
|
||||
@@ -696,6 +795,7 @@ mod tests {
|
||||
file_data[fahd_offset + 8..fahd_offset + 16].copy_from_slice(&num_elements.to_le_bytes());
|
||||
file_data[fahd_offset + 16..fahd_offset + 24]
|
||||
.copy_from_slice(&(db_offset as u64).to_le_bytes());
|
||||
stamp_checksum(&mut file_data, fahd_offset, fahd_offset + 24);
|
||||
|
||||
// FADB prefix
|
||||
file_data[db_offset..db_offset + 4].copy_from_slice(b"FADB");
|
||||
@@ -715,6 +815,9 @@ mod tests {
|
||||
let base_addr = 0x1000u64;
|
||||
// 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.
|
||||
// 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] {
|
||||
let page_off = pages_start + p * page_total;
|
||||
let count = core::cmp::min(page_nelmts, num_elements as usize - p * page_nelmts);
|
||||
@@ -724,6 +827,7 @@ mod tests {
|
||||
let pos = page_off + e * os;
|
||||
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 =
|
||||
|
||||
@@ -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());
|
||||
}
|
||||
}
|
||||
@@ -72,6 +72,7 @@ pub mod filter_pipeline;
|
||||
pub mod filters;
|
||||
mod filters_szip;
|
||||
pub mod fixed_array;
|
||||
pub mod float16;
|
||||
pub mod fractal_heap;
|
||||
pub mod global_heap;
|
||||
pub mod group_info;
|
||||
|
||||
@@ -56,6 +56,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 {
|
||||
Datatype::FloatingPoint {
|
||||
size: 4,
|
||||
@@ -478,6 +493,24 @@ impl DatasetBuilder {
|
||||
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 {
|
||||
self.datatype = Some(make_i32_type());
|
||||
let mut b = Vec::with_capacity(data.len() * 4);
|
||||
|
||||
@@ -2,6 +2,15 @@
|
||||
|
||||
use clawhdf5_format::data_read::{read_object_references, read_region_references};
|
||||
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]
|
||||
fn object_ref_single_valid() {
|
||||
@@ -173,7 +182,7 @@ print('ok')
|
||||
"#,
|
||||
path.display()
|
||||
);
|
||||
let output = std::process::Command::new("python3")
|
||||
let output = std::process::Command::new(python())
|
||||
.args(["-c", &script])
|
||||
.output();
|
||||
|
||||
|
||||
@@ -4,9 +4,18 @@
|
||||
//! (and vice versa). They require python3 + h5py to be installed.
|
||||
|
||||
use clawhdf5_format::file_writer::{AttrValue, CompoundTypeBuilder, EnumTypeBuilder, FileWriter};
|
||||
/// 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())
|
||||
}
|
||||
|
||||
fn h5py_available() -> bool {
|
||||
std::process::Command::new("python3")
|
||||
std::process::Command::new(python())
|
||||
.args(["-c", "import h5py"])
|
||||
.output()
|
||||
.map(|o| o.status.success())
|
||||
@@ -17,10 +26,10 @@ fn h5py_read(_path: &std::path::Path, script: &str) -> String {
|
||||
if !h5py_available() {
|
||||
panic!("h5py not installed — skipping interop test");
|
||||
}
|
||||
let o = std::process::Command::new("python3")
|
||||
let o = std::process::Command::new(python())
|
||||
.args(["-c", script])
|
||||
.output()
|
||||
.expect("python3");
|
||||
.expect("python interpreter");
|
||||
if !o.status.success() {
|
||||
panic!("h5py: {}", String::from_utf8_lossy(&o.stderr));
|
||||
}
|
||||
|
||||
@@ -1,7 +1,8 @@
|
||||
[package]
|
||||
name = "clawhdf5-gpu"
|
||||
version = "2.5.0"
|
||||
version = "2.7.0"
|
||||
edition = "2024"
|
||||
rust-version.workspace = true
|
||||
description = "GPU-accelerated vector operations for rustyhdf5 using wgpu compute shaders"
|
||||
license = "MIT"
|
||||
repository = "https://git.redclaw.dev/quantumclaw/clawhdf5"
|
||||
|
||||
@@ -1,7 +1,8 @@
|
||||
[package]
|
||||
name = "clawhdf5-io"
|
||||
version = "2.5.0"
|
||||
version = "2.7.0"
|
||||
edition = "2024"
|
||||
rust-version.workspace = true
|
||||
description = "I/O abstraction layer for rustyhdf5"
|
||||
license = "MIT"
|
||||
repository = "https://git.redclaw.dev/quantumclaw/clawhdf5"
|
||||
@@ -10,7 +11,7 @@ keywords = ["hdf5", "io", "science", "data"]
|
||||
categories = ["filesystem", "science"]
|
||||
|
||||
[dependencies]
|
||||
clawhdf5-format = { path = "../clawhdf5-format", version = "2.5.0" }
|
||||
clawhdf5-format = { path = "../clawhdf5-format", version = "2.7.0" }
|
||||
memmap2 = { version = "0.9", optional = true }
|
||||
libc = { version = "0.2", optional = true }
|
||||
tokio = { version = "1", features = ["fs", "io-util"], optional = true }
|
||||
|
||||
@@ -1,7 +1,8 @@
|
||||
[package]
|
||||
name = "clawhdf5-migrate"
|
||||
version = "2.5.0"
|
||||
version = "2.7.0"
|
||||
edition = "2024"
|
||||
rust-version.workspace = true
|
||||
description = "CLI to migrate SQLite agent memory databases to HDF5 format"
|
||||
license = "MIT"
|
||||
repository = "https://git.redclaw.dev/quantumclaw/clawhdf5"
|
||||
@@ -14,9 +15,9 @@ name = "clawhdf5-migrate"
|
||||
path = "src/main.rs"
|
||||
|
||||
[dependencies]
|
||||
clawhdf5-agent = { path = "../clawhdf5-agent", version = "2.5.0" }
|
||||
clawhdf5-format = { path = "../clawhdf5-format", version = "2.5.0" }
|
||||
clawhdf5 = { path = "../clawhdf5", version = "2.5.0" }
|
||||
clawhdf5-agent = { path = "../clawhdf5-agent", version = "2.7.0" }
|
||||
clawhdf5-format = { path = "../clawhdf5-format", version = "2.7.0" }
|
||||
clawhdf5 = { path = "../clawhdf5", version = "2.7.0" }
|
||||
rusqlite = { version = "0.31", features = ["bundled"] }
|
||||
clap = { version = "4", features = ["derive"] }
|
||||
half = { workspace = true }
|
||||
|
||||
@@ -1,7 +1,8 @@
|
||||
[package]
|
||||
name = "clawhdf5-napi"
|
||||
version = "2.5.0"
|
||||
version = "2.7.0"
|
||||
edition = "2024"
|
||||
rust-version.workspace = true
|
||||
description = "Node.js native addon (napi-rs) exposing clawhdf5-agent to TypeScript/JavaScript"
|
||||
license = "MIT"
|
||||
repository = "https://git.redclaw.dev/quantumclaw/clawhdf5"
|
||||
@@ -10,7 +11,7 @@ repository = "https://git.redclaw.dev/quantumclaw/clawhdf5"
|
||||
crate-type = ["cdylib"]
|
||||
|
||||
[dependencies]
|
||||
clawhdf5-agent = { path = "../clawhdf5-agent", version = "2.5.0" }
|
||||
clawhdf5-agent = { path = "../clawhdf5-agent", version = "2.7.0" }
|
||||
napi = { version = "2", default-features = false, features = ["napi9"] }
|
||||
napi-derive = "2"
|
||||
|
||||
|
||||
@@ -1,7 +1,8 @@
|
||||
[package]
|
||||
name = "clawhdf5-netcdf4"
|
||||
version = "2.5.0"
|
||||
version = "2.7.0"
|
||||
edition = "2024"
|
||||
rust-version.workspace = true
|
||||
description = "NetCDF-4 read support built on rustyhdf5 — pure Rust, no C dependencies"
|
||||
license = "MIT"
|
||||
repository = "https://git.redclaw.dev/quantumclaw/clawhdf5"
|
||||
@@ -10,8 +11,8 @@ keywords = ["netcdf", "netcdf4", "hdf5", "science", "climate"]
|
||||
categories = ["parser-implementations", "science"]
|
||||
|
||||
[dependencies]
|
||||
clawhdf5 = { path = "../clawhdf5", version = "2.5.0" }
|
||||
clawhdf5-format = { path = "../clawhdf5-format", version = "2.5.0" }
|
||||
clawhdf5 = { path = "../clawhdf5", version = "2.7.0" }
|
||||
clawhdf5-format = { path = "../clawhdf5-format", version = "2.7.0" }
|
||||
|
||||
[dev-dependencies]
|
||||
tempfile = { workspace = true }
|
||||
|
||||
@@ -9,6 +9,15 @@ use clawhdf5_netcdf4::{AttrValue, NetCDF4File};
|
||||
// ---------------------------------------------------------------------------
|
||||
// Helpers
|
||||
// ---------------------------------------------------------------------------
|
||||
/// 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())
|
||||
}
|
||||
|
||||
/// When `CLAWHDF5_REQUIRE_INTEROP=1` (set in CI), a missing Python dependency
|
||||
/// is a test failure instead of a silent skip.
|
||||
@@ -17,7 +26,7 @@ fn interop_required() -> bool {
|
||||
}
|
||||
|
||||
fn netcdf4_python_available() -> bool {
|
||||
Command::new("python3")
|
||||
Command::new(python())
|
||||
.args(["-c", "import netCDF4; print(netCDF4.__version__)"])
|
||||
.output()
|
||||
.map(|o| o.status.success())
|
||||
@@ -25,7 +34,7 @@ fn netcdf4_python_available() -> bool {
|
||||
}
|
||||
|
||||
fn xarray_available() -> bool {
|
||||
Command::new("python3")
|
||||
Command::new(python())
|
||||
.args(["-c", "import xarray; print(xarray.__version__)"])
|
||||
.output()
|
||||
.map(|o| o.status.success())
|
||||
@@ -59,7 +68,7 @@ macro_rules! skip_if_no_xarray {
|
||||
}
|
||||
|
||||
fn run_python(script: &str) {
|
||||
let output = Command::new("python3")
|
||||
let output = Command::new(python())
|
||||
.args(["-c", script])
|
||||
.output()
|
||||
.expect("failed to run python3");
|
||||
|
||||
@@ -1,7 +1,8 @@
|
||||
[package]
|
||||
name = "clawhdf5-py"
|
||||
version = "2.5.0"
|
||||
version = "2.7.0"
|
||||
edition = "2024"
|
||||
rust-version.workspace = true
|
||||
description = "Python bindings for rustyhdf5 — a pure-Rust HDF5 library"
|
||||
license = "MIT"
|
||||
repository = "https://git.redclaw.dev/quantumclaw/clawhdf5"
|
||||
@@ -14,8 +15,8 @@ name = "clawhdf5"
|
||||
crate-type = ["cdylib", "rlib"]
|
||||
|
||||
[dependencies]
|
||||
clawhdf5_rs = { path = "../clawhdf5", version = "2.5.0", package = "clawhdf5" }
|
||||
clawhdf5-format = { path = "../clawhdf5-format", version = "2.5.0" }
|
||||
clawhdf5_rs = { path = "../clawhdf5", version = "2.7.0", package = "clawhdf5" }
|
||||
clawhdf5-format = { path = "../clawhdf5-format", version = "2.7.0" }
|
||||
pyo3 = "0.29"
|
||||
numpy = "0.29"
|
||||
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "maturin"
|
||||
|
||||
[project]
|
||||
name = "rustyhdf5"
|
||||
version = "2.5.0"
|
||||
version = "2.7.0"
|
||||
description = "Python bindings for rustyhdf5 — a pure-Rust HDF5 library"
|
||||
requires-python = ">=3.8"
|
||||
license = { text = "MIT" }
|
||||
|
||||
@@ -1,7 +1,8 @@
|
||||
[package]
|
||||
name = "clawhdf5"
|
||||
version = "2.5.0"
|
||||
version = "2.7.0"
|
||||
edition = "2024"
|
||||
rust-version.workspace = true
|
||||
description = "Pure-Rust HDF5 reader/writer — no C dependencies"
|
||||
license = "MIT"
|
||||
repository = "https://git.redclaw.dev/quantumclaw/clawhdf5"
|
||||
@@ -10,16 +11,16 @@ keywords = ["hdf5", "science", "data", "binary"]
|
||||
categories = ["parser-implementations", "science", "encoding"]
|
||||
|
||||
[dependencies]
|
||||
clawhdf5-format = { path = "../clawhdf5-format", version = "2.5.0" }
|
||||
clawhdf5-io = { path = "../clawhdf5-io", version = "2.5.0" }
|
||||
clawhdf5-format = { path = "../clawhdf5-format", version = "2.7.0" }
|
||||
clawhdf5-io = { path = "../clawhdf5-io", version = "2.7.0" }
|
||||
rayon = { version = "1", optional = true }
|
||||
|
||||
[dev-dependencies]
|
||||
tempfile = { workspace = true }
|
||||
criterion = { workspace = true }
|
||||
clawhdf5-io = { path = "../clawhdf5-io", version = "2.5.0", features = ["mmap"] }
|
||||
clawhdf5-format = { path = "../clawhdf5-format", version = "2.5.0", features = ["parallel", "fast-checksum"] }
|
||||
clawhdf5-filters = { path = "../clawhdf5-filters", version = "2.5.0" }
|
||||
clawhdf5-io = { path = "../clawhdf5-io", version = "2.7.0", features = ["mmap"] }
|
||||
clawhdf5-format = { path = "../clawhdf5-format", version = "2.7.0", features = ["parallel", "fast-checksum"] }
|
||||
clawhdf5-filters = { path = "../clawhdf5-filters", version = "2.7.0" }
|
||||
|
||||
[[bench]]
|
||||
name = "mmap_bench"
|
||||
@@ -30,9 +31,10 @@ name = "parallel_bench"
|
||||
harness = false
|
||||
|
||||
[features]
|
||||
default = ["mmap", "fast-deflate", "provenance"]
|
||||
default = ["mmap", "provenance"]
|
||||
mmap = ["clawhdf5-io/mmap"]
|
||||
parallel = ["clawhdf5-format/parallel", "rayon"]
|
||||
# zlib-ng (C, needs cmake) instead of the default pure-Rust zlib-rs.
|
||||
fast-deflate = ["clawhdf5-format/fast-deflate"]
|
||||
apple-compression = []
|
||||
zstd = ["clawhdf5-format/zstd"]
|
||||
|
||||
@@ -9,6 +9,15 @@ use clawhdf5::{AttrValue, CompoundTypeBuilder, DType, File, FileBuilder};
|
||||
// ---------------------------------------------------------------------------
|
||||
// Helpers
|
||||
// ---------------------------------------------------------------------------
|
||||
/// 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())
|
||||
}
|
||||
|
||||
/// When `CLAWHDF5_REQUIRE_INTEROP=1` (set in CI), a missing Python dependency
|
||||
/// is a test failure instead of a silent skip.
|
||||
@@ -17,7 +26,7 @@ fn interop_required() -> bool {
|
||||
}
|
||||
|
||||
fn python_available() -> bool {
|
||||
Command::new("python3")
|
||||
Command::new(python())
|
||||
.args(["-c", "import h5py; print(h5py.__version__)"])
|
||||
.output()
|
||||
.map(|o| o.status.success())
|
||||
@@ -39,7 +48,7 @@ macro_rules! skip_if_no_python {
|
||||
|
||||
/// Run a Python script and panic if it fails.
|
||||
fn run_python(script: &str) {
|
||||
let output = Command::new("python3")
|
||||
let output = Command::new(python())
|
||||
.args(["-c", script])
|
||||
.output()
|
||||
.expect("failed to run python3");
|
||||
@@ -52,7 +61,7 @@ fn run_python(script: &str) {
|
||||
|
||||
/// Run a Python script and return stdout as a trimmed string.
|
||||
fn run_python_output(script: &str) -> String {
|
||||
let output = Command::new("python3")
|
||||
let output = Command::new(python())
|
||||
.args(["-c", script])
|
||||
.output()
|
||||
.expect("failed to run python3");
|
||||
@@ -1038,3 +1047,506 @@ with h5py.File("{path_str}", "r") as f:
|
||||
data[start..start + cols as usize]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn h5py_deep_btree_v2_chunk_index_clawhdf5_reads() {
|
||||
// Two unlimited dimensions give a B-tree v2 chunk index, and 2x2 chunks
|
||||
// over 400x400 give 40 000 index records — enough for HDF5 to build a
|
||||
// tree of depth 2. Small h5py files only ever produce depth-0 trees, so
|
||||
// this is the one fixture that walks internal nodes: the path where the
|
||||
// traversal's record budget (the guard against crafted shared-subtree
|
||||
// trees) is spent, which must never refuse a real file.
|
||||
skip_if_no_python!();
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = dir.path().join("deep_btree.h5");
|
||||
let path_str = path.display().to_string();
|
||||
|
||||
let script = format!(
|
||||
r#"
|
||||
import h5py, numpy as np
|
||||
with h5py.File("{path_str}", "w", libver="latest") as f:
|
||||
d = f.create_dataset("x", shape=(400, 400), maxshape=(None, None),
|
||||
chunks=(2, 2), dtype="i4")
|
||||
d[...] = np.arange(160000, dtype="i4").reshape(400, 400)
|
||||
"#
|
||||
);
|
||||
run_python(&script);
|
||||
|
||||
// The fixture is only meaningful if HDF5 really built internal nodes.
|
||||
let bytes = std::fs::read(&path).unwrap();
|
||||
let at = bytes
|
||||
.windows(4)
|
||||
.position(|w| w == b"BTHD")
|
||||
.expect("expected a B-tree v2 chunk index");
|
||||
let depth = u16::from_le_bytes([bytes[at + 12], bytes[at + 13]]);
|
||||
assert!(
|
||||
depth >= 1,
|
||||
"fixture tree has depth {depth}; it tests nothing"
|
||||
);
|
||||
|
||||
let file = File::open(&path).unwrap();
|
||||
let values = file.dataset("x").unwrap().read_i32().unwrap();
|
||||
assert_eq!(values.len(), 160_000);
|
||||
for (i, &v) in values.iter().enumerate() {
|
||||
assert_eq!(v, i as i32, "element {i}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn h5py_extensible_array_chunk_index_clawhdf5_reads() {
|
||||
// One unlimited dimension means an Extensible Array chunk index. Only its
|
||||
// first few elements live inline in the index block (4 by default), and
|
||||
// every other fixture here is small enough to stop there — which is how
|
||||
// the data block and super block layouts came to be wrong without a test
|
||||
// noticing. The counts below step over each boundary in turn:
|
||||
// 4 inline elements only
|
||||
// 37 past the first direct data block
|
||||
// 400 into the first super block
|
||||
// 5000 several super block levels
|
||||
// 200000 data blocks large enough to be paged
|
||||
skip_if_no_python!();
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
|
||||
for n in [4usize, 37, 400, 5_000, 200_000] {
|
||||
let path = dir.path().join(format!("ea_{n}.h5"));
|
||||
let path_str = path.display().to_string();
|
||||
run_python(&format!(
|
||||
r#"
|
||||
import h5py, numpy as np
|
||||
with h5py.File("{path_str}", "w", libver="latest") as f:
|
||||
d = f.create_dataset("x", shape=({n},), maxshape=(None,), chunks=(1,), dtype="i4")
|
||||
d[...] = np.arange({n}, dtype="i4")
|
||||
"#
|
||||
));
|
||||
|
||||
let bytes = std::fs::read(&path).unwrap();
|
||||
assert!(
|
||||
bytes.windows(4).any(|w| w == b"EAHD"),
|
||||
"n={n}: fixture is not indexed by an Extensible Array"
|
||||
);
|
||||
|
||||
let file = File::open(&path).unwrap();
|
||||
let values = file.dataset("x").unwrap().read_i32().unwrap();
|
||||
assert_eq!(values.len(), n, "n={n}");
|
||||
let wrong = values
|
||||
.iter()
|
||||
.enumerate()
|
||||
.filter(|&(i, &v)| v != i as i32)
|
||||
.count();
|
||||
assert_eq!(wrong, 0, "n={n}: {wrong} of {n} elements read back wrong");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn h5py_sparse_extensible_array_leaves_pages_uninitialised() {
|
||||
// Writing a scattered subset leaves whole pages of a paged data block
|
||||
// never initialised. Those pages still occupy their slot on disk, so the
|
||||
// reader has to skip them by stride and take the fill value instead —
|
||||
// driven by the page-init bitmap, which is packed one bit per page across
|
||||
// the whole super block, MSB first.
|
||||
skip_if_no_python!();
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = dir.path().join("ea_sparse.h5");
|
||||
let path_str = path.display().to_string();
|
||||
let n = 200_000usize;
|
||||
let step = 997usize;
|
||||
|
||||
run_python(&format!(
|
||||
r#"
|
||||
import h5py
|
||||
with h5py.File("{path_str}", "w", libver="latest") as f:
|
||||
d = f.create_dataset("x", shape=({n},), maxshape=(None,), chunks=(1,),
|
||||
dtype="i4", fillvalue=-1)
|
||||
for i in list(range(0, {n}, {step})) + list(range(0, 40)):
|
||||
d[i] = i
|
||||
"#
|
||||
));
|
||||
|
||||
let file = File::open(&path).unwrap();
|
||||
let values = file.dataset("x").unwrap().read_i32().unwrap();
|
||||
assert_eq!(values.len(), n);
|
||||
let wrong = values
|
||||
.iter()
|
||||
.enumerate()
|
||||
.filter(|&(i, &v)| {
|
||||
let expected = if i % step == 0 || i < 40 {
|
||||
i as i32
|
||||
} else {
|
||||
-1
|
||||
};
|
||||
v != expected
|
||||
})
|
||||
.count();
|
||||
assert_eq!(wrong, 0, "{wrong} of {n} elements read back wrong");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn h5py_filtered_and_2d_extensible_array_clawhdf5_reads() {
|
||||
// Filtered elements carry a size and filter mask beside the address, and
|
||||
// a second (fixed) dimension changes how a linear index maps back to
|
||||
// chunk offsets. Both run through the same traversal.
|
||||
skip_if_no_python!();
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
|
||||
let gz = dir.path().join("ea_gzip.h5");
|
||||
let gz_str = gz.display().to_string();
|
||||
run_python(&format!(
|
||||
r#"
|
||||
import h5py, numpy as np
|
||||
with h5py.File("{gz_str}", "w", libver="latest") as f:
|
||||
d = f.create_dataset("x", shape=(5000,), maxshape=(None,), chunks=(1,),
|
||||
dtype="i4", compression="gzip", compression_opts=4)
|
||||
d[...] = np.arange(5000, dtype="i4")
|
||||
"#
|
||||
));
|
||||
let values = File::open(&gz)
|
||||
.unwrap()
|
||||
.dataset("x")
|
||||
.unwrap()
|
||||
.read_i32()
|
||||
.unwrap();
|
||||
assert_eq!(values.len(), 5000);
|
||||
assert_eq!(
|
||||
values
|
||||
.iter()
|
||||
.enumerate()
|
||||
.filter(|&(i, &v)| v != i as i32)
|
||||
.count(),
|
||||
0
|
||||
);
|
||||
|
||||
let two_d = dir.path().join("ea_2d.h5");
|
||||
let two_d_str = two_d.display().to_string();
|
||||
run_python(&format!(
|
||||
r#"
|
||||
import h5py, numpy as np
|
||||
with h5py.File("{two_d_str}", "w", libver="latest") as f:
|
||||
d = f.create_dataset("x", shape=(3000, 4), maxshape=(None, 4), chunks=(1, 4), dtype="i4")
|
||||
d[...] = np.arange(12000, dtype="i4").reshape(3000, 4)
|
||||
"#
|
||||
));
|
||||
let values = File::open(&two_d)
|
||||
.unwrap()
|
||||
.dataset("x")
|
||||
.unwrap()
|
||||
.read_i32()
|
||||
.unwrap();
|
||||
assert_eq!(values.len(), 12_000);
|
||||
assert_eq!(
|
||||
values
|
||||
.iter()
|
||||
.enumerate()
|
||||
.filter(|&(i, &v)| v != i as i32)
|
||||
.count(),
|
||||
0
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn h5py_fixed_array_chunk_index_clawhdf5_reads() {
|
||||
// Fixed dimensions plus libver='latest' give a Fixed Array chunk index.
|
||||
// Its data blocks are paged above 2^page_bits elements (1024 by default),
|
||||
// and unlike the Extensible Array it keeps the page-init bitmap in the
|
||||
// data block itself — a difference worth pinning down, since assuming
|
||||
// otherwise is exactly what made the Extensible Array reader wrong. The
|
||||
// sparse case leaves whole pages uninitialised so the bitmap is actually
|
||||
// consulted rather than being all ones.
|
||||
skip_if_no_python!();
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
|
||||
for n in [100usize, 5_000, 200_000] {
|
||||
let path = dir.path().join(format!("fa_{n}.h5"));
|
||||
let path_str = path.display().to_string();
|
||||
run_python(&format!(
|
||||
r#"
|
||||
import h5py, numpy as np
|
||||
with h5py.File("{path_str}", "w", libver="latest") as f:
|
||||
d = f.create_dataset("x", shape=({n},), chunks=(1,), dtype="i4")
|
||||
d[...] = np.arange({n}, dtype="i4")
|
||||
"#
|
||||
));
|
||||
let bytes = std::fs::read(&path).unwrap();
|
||||
assert!(
|
||||
bytes.windows(4).any(|w| w == b"FAHD"),
|
||||
"n={n}: fixture is not indexed by a Fixed Array"
|
||||
);
|
||||
let values = File::open(&path)
|
||||
.unwrap()
|
||||
.dataset("x")
|
||||
.unwrap()
|
||||
.read_i32()
|
||||
.unwrap();
|
||||
assert_eq!(values.len(), n, "n={n}");
|
||||
let wrong = values
|
||||
.iter()
|
||||
.enumerate()
|
||||
.filter(|&(i, &v)| v != i as i32)
|
||||
.count();
|
||||
assert_eq!(wrong, 0, "n={n}: {wrong} elements read back wrong");
|
||||
}
|
||||
|
||||
let sparse = dir.path().join("fa_sparse.h5");
|
||||
let sparse_str = sparse.display().to_string();
|
||||
let (n, step) = (200_000usize, 997usize);
|
||||
run_python(&format!(
|
||||
r#"
|
||||
import h5py
|
||||
with h5py.File("{sparse_str}", "w", libver="latest") as f:
|
||||
d = f.create_dataset("x", shape=({n},), chunks=(1,), dtype="i4", fillvalue=-1)
|
||||
for i in list(range(0, {n}, {step})) + list(range(0, 40)):
|
||||
d[i] = i
|
||||
"#
|
||||
));
|
||||
let values = File::open(&sparse)
|
||||
.unwrap()
|
||||
.dataset("x")
|
||||
.unwrap()
|
||||
.read_i32()
|
||||
.unwrap();
|
||||
assert_eq!(values.len(), n);
|
||||
let wrong = values
|
||||
.iter()
|
||||
.enumerate()
|
||||
.filter(|&(i, &v)| {
|
||||
let expected = if i % step == 0 || i < 40 {
|
||||
i as i32
|
||||
} else {
|
||||
-1
|
||||
};
|
||||
v != expected
|
||||
})
|
||||
.count();
|
||||
assert_eq!(wrong, 0, "sparse: {wrong} of {n} elements read back wrong");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn corrupting_a_chunk_index_is_an_error_not_wrong_data() {
|
||||
// Every Fixed/Extensible Array structure carries a Jenkins checksum, and
|
||||
// the reader now verifies it. The point is not the checksum itself but
|
||||
// what it prevents: a damaged index otherwise yields addresses pointing
|
||||
// at the wrong bytes, and the caller receives another chunk's data as if
|
||||
// it were the one asked for.
|
||||
skip_if_no_python!();
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
|
||||
for (name, maxshape) in [("fixed", "None"), ("extensible", "(None,)")] {
|
||||
let path = dir.path().join(format!("{name}.h5"));
|
||||
let path_str = path.display().to_string();
|
||||
let shape_arg = if maxshape == "None" {
|
||||
String::new()
|
||||
} else {
|
||||
format!(", maxshape={maxshape}")
|
||||
};
|
||||
run_python(&format!(
|
||||
r#"
|
||||
import h5py, numpy as np
|
||||
with h5py.File("{path_str}", "w", libver="latest") as f:
|
||||
d = f.create_dataset("x", shape=(400,), chunks=(1,), dtype="i4"{shape_arg})
|
||||
d[...] = np.arange(400, dtype="i4")
|
||||
"#
|
||||
));
|
||||
|
||||
let clean = std::fs::read(&path).unwrap();
|
||||
assert_eq!(
|
||||
File::open(&path)
|
||||
.unwrap()
|
||||
.dataset("x")
|
||||
.unwrap()
|
||||
.read_i32()
|
||||
.unwrap()
|
||||
.len(),
|
||||
400,
|
||||
"{name}: the intact file must read"
|
||||
);
|
||||
|
||||
// Flip a low bit of a chunk address inside a data block. Structurally
|
||||
// everything still parses — the index still has the right shape and
|
||||
// the address still lands inside the file — so nothing but the
|
||||
// checksum can notice. Without it the read succeeds and hands back
|
||||
// whatever bytes now sit at that address.
|
||||
let sig: &[u8] = if name == "fixed" { b"FADB" } else { b"EADB" };
|
||||
let block = clean
|
||||
.windows(4)
|
||||
.position(|w| w == sig)
|
||||
.unwrap_or_else(|| panic!("{name}: no data block in the fixture"));
|
||||
// Past the prefix (signature, version, client id, header address, and
|
||||
// for the Extensible Array a block offset), into the first address.
|
||||
let at = block + 4 + 1 + 1 + 8 + if name == "fixed" { 0 } else { 4 } + 1;
|
||||
let mut damaged = clean.clone();
|
||||
damaged[at] ^= 0x10;
|
||||
let damaged_path = dir.path().join(format!("{name}_damaged.h5"));
|
||||
std::fs::write(&damaged_path, &damaged).unwrap();
|
||||
|
||||
let result = File::open(&damaged_path)
|
||||
.unwrap()
|
||||
.dataset("x")
|
||||
.and_then(|d| d.read_i32());
|
||||
assert!(
|
||||
result.is_err(),
|
||||
"{name}: corruption produced data instead of an error"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Half precision (float16) in both directions
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Values that exercise rounding: ties, subnormals, the overflow boundary and
|
||||
/// ordinary embedding-sized components.
|
||||
fn f16_probe_values() -> Vec<f32> {
|
||||
let mut v = vec![
|
||||
0.0,
|
||||
-0.0,
|
||||
1.0,
|
||||
-1.0,
|
||||
0.5,
|
||||
1.0 + 2f32.powi(-11),
|
||||
1.0 + 3.0 * 2f32.powi(-11),
|
||||
65504.0,
|
||||
65519.0,
|
||||
65520.0,
|
||||
-70000.0,
|
||||
6.0e-8,
|
||||
3.0e-8,
|
||||
1.0e-9,
|
||||
1.0e-5,
|
||||
0.1,
|
||||
0.333_333,
|
||||
1234.567,
|
||||
f32::INFINITY,
|
||||
f32::NEG_INFINITY,
|
||||
];
|
||||
// A deterministic spread of embedding-like values.
|
||||
let mut x = 0x2545_F491u32;
|
||||
for _ in 0..4000 {
|
||||
x ^= x << 13;
|
||||
x ^= x >> 17;
|
||||
x ^= x << 5;
|
||||
v.push((x as f32 / u32::MAX as f32 - 0.5) * 0.4);
|
||||
}
|
||||
v
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn clawhdf5_writes_f16_h5py_reads() {
|
||||
skip_if_no_python!();
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = dir.path().join("ours_f16.h5");
|
||||
let path_str = path.display().to_string();
|
||||
let values = f16_probe_values();
|
||||
|
||||
let mut fb = FileBuilder::new();
|
||||
fb.create_dataset("plain").with_f16_data(&values);
|
||||
fb.create_dataset("chunked")
|
||||
.with_f16_data(&values)
|
||||
.with_shape(&[values.len() as u64])
|
||||
.with_chunks(&[512])
|
||||
.with_deflate(6);
|
||||
fb.write(&path).unwrap();
|
||||
|
||||
// h5py must see a genuine float16 dataset, and our rounding must agree
|
||||
// with numpy's own float32 -> float16 conversion bit for bit.
|
||||
let input = values
|
||||
.iter()
|
||||
.map(|v| format!("{:?}", v.to_bits()))
|
||||
.collect::<Vec<_>>()
|
||||
.join(",");
|
||||
let script = format!(
|
||||
r#"
|
||||
import h5py, numpy as np
|
||||
src = np.array([{input}], dtype=np.uint32).view(np.float32)
|
||||
expected = src.astype(np.float16).view(np.uint16)
|
||||
with h5py.File("{path_str}", "r") as f:
|
||||
for name in ("plain", "chunked"):
|
||||
d = f[name]
|
||||
assert d.dtype == np.float16, (name, d.dtype)
|
||||
got = d[:].view(np.uint16)
|
||||
bad = np.nonzero(got != expected)[0]
|
||||
assert bad.size == 0, (name, bad[:5], got[bad[:5]], expected[bad[:5]])
|
||||
print("ok")
|
||||
"#
|
||||
);
|
||||
assert_eq!(run_python_output(&script), "ok");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn h5py_writes_f16_clawhdf5_reads() {
|
||||
skip_if_no_python!();
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = dir.path().join("h5py_f16.h5");
|
||||
let path_str = path.display().to_string();
|
||||
let values = f16_probe_values();
|
||||
let input = values
|
||||
.iter()
|
||||
.map(|v| format!("{:?}", v.to_bits()))
|
||||
.collect::<Vec<_>>()
|
||||
.join(",");
|
||||
|
||||
let script = format!(
|
||||
r#"
|
||||
import h5py, numpy as np
|
||||
src = np.array([{input}], dtype=np.uint32).view(np.float32).astype(np.float16)
|
||||
with h5py.File("{path_str}", "w") as f:
|
||||
f.create_dataset("plain", data=src)
|
||||
f.create_dataset("chunked", data=src, chunks=(512,), compression="gzip", shuffle=True)
|
||||
f.create_dataset("big_endian", data=src.astype(">f2"))
|
||||
"#
|
||||
);
|
||||
run_python(&script);
|
||||
|
||||
let expected: Vec<u32> = values
|
||||
.iter()
|
||||
.map(|&v| clawhdf5_format::float16::round_to_f16(v).to_bits())
|
||||
.collect();
|
||||
let file = File::open(&path).unwrap();
|
||||
for name in ["plain", "chunked", "big_endian"] {
|
||||
let ds = file.dataset(name).unwrap();
|
||||
assert_eq!(
|
||||
ds.dtype().unwrap(),
|
||||
DType::Other("float16".into()),
|
||||
"{name}"
|
||||
);
|
||||
let got: Vec<u32> = ds.read_f32().unwrap().iter().map(|v| v.to_bits()).collect();
|
||||
assert_eq!(got, expected, "{name}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn clawhdf5_writes_f32_h5py_reads() {
|
||||
// Every f32 dataset used to be unreadable by h5py ("sign bit position out
|
||||
// of bounds"): the float datatype's sign position was hard-coded for f64.
|
||||
skip_if_no_python!();
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = dir.path().join("ours_f32.h5");
|
||||
let path_str = path.display().to_string();
|
||||
let values: Vec<f32> = vec![1.5, -2.25, 3.0e-7, 65536.5, f32::MAX, -0.0];
|
||||
|
||||
let mut fb = FileBuilder::new();
|
||||
fb.create_dataset("plain").with_f32_data(&values);
|
||||
fb.create_dataset("chunked")
|
||||
.with_f32_data(&values)
|
||||
.with_shape(&[values.len() as u64])
|
||||
.with_chunks(&[4])
|
||||
.with_deflate(6);
|
||||
fb.write(&path).unwrap();
|
||||
|
||||
let bits = values
|
||||
.iter()
|
||||
.map(|v| v.to_bits().to_string())
|
||||
.collect::<Vec<_>>()
|
||||
.join(",");
|
||||
let script = format!(
|
||||
r#"
|
||||
import h5py, numpy as np
|
||||
expected = np.array([{bits}], dtype=np.uint32)
|
||||
with h5py.File("{path_str}", "r") as f:
|
||||
for name in ("plain", "chunked"):
|
||||
d = f[name]
|
||||
assert d.dtype == np.float32, (name, d.dtype)
|
||||
assert (d[:].view(np.uint32) == expected).all(), (name, d[:])
|
||||
print("ok")
|
||||
"#
|
||||
);
|
||||
assert_eq!(run_python_output(&script), "ok");
|
||||
}
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
name = "libaec-sys"
|
||||
version = "0.1.0"
|
||||
edition = "2024"
|
||||
rust-version.workspace = true
|
||||
links = "aec"
|
||||
|
||||
[build-dependencies]
|
||||
|
||||
+7
-1
@@ -364,6 +364,12 @@ cargo install --path crates/clawhdf5-cli
|
||||
clawhdf5 --path agent.h5 create --agent-id my-agent --dim 384 --wal
|
||||
```
|
||||
|
||||
New stores hold the vector index's copy of the embeddings as int8, which
|
||||
roughly halves a loaded store's memory and is faster at equal recall — the
|
||||
query path re-scores candidates against the exact embeddings. Pass
|
||||
`--f32-index` to keep an f32 index instead. The setting is recorded in the
|
||||
file, and stores created before it existed keep their f32 index.
|
||||
|
||||
Output:
|
||||
```json
|
||||
{
|
||||
@@ -561,4 +567,4 @@ let final_results = confidence::reject_low_confidence(
|
||||
|
||||
---
|
||||
|
||||
<p align="center"><em>Built by <a href="https://github.com/redclawsystems">RedClaw Systems</a></em></p>
|
||||
<p align="center"><em>Built by <a href="https://git.redclaw.dev/quantumclaw">RedClaw Systems</a></em></p>
|
||||
|
||||
+1
-1
@@ -232,4 +232,4 @@ clawhdf5-agent = { version = "2.0", features = ["agent", "float16", "accelerate"
|
||||
|
||||
---
|
||||
|
||||
<p align="center"><em>Built by <a href="https://github.com/redclawsystems">RedClaw Systems</a></em></p>
|
||||
<p align="center"><em>Built by <a href="https://git.redclaw.dev/quantumclaw">RedClaw Systems</a></em></p>
|
||||
|
||||
@@ -146,3 +146,136 @@ created with `external=[...]` storage returns
|
||||
`FormatError::ExternalDataFilesUnsupported`. Neither is resolved. If support is
|
||||
added, file names must be confined to the opened file's directory, as the
|
||||
virtual-dataset resolver now does.
|
||||
|
||||
---
|
||||
|
||||
## Python interop suites skip silently when no interpreter has h5py
|
||||
|
||||
**Status:** fixed on `main` in `a29c1b2` (2026-09-19).
|
||||
|
||||
On a system where `python3` is a PEP 668 "externally managed" interpreter,
|
||||
h5py cannot be installed into it at all, and every interop suite — the h5py
|
||||
writer round-trips, the facade suite, netCDF4, and the reference files —
|
||||
returned `false` from its availability probe and skipped without failing. CI
|
||||
reported `SKIP` and a green run. This is the same class of gap that let the
|
||||
compound-datatype v5 bug above reach a release.
|
||||
|
||||
The probes now read `CLAWHDF5_PYTHON`, and `scripts/ci-test.sh` picks up
|
||||
`.venv/bin/python` automatically. To restore the coverage on a fresh checkout:
|
||||
|
||||
```bash
|
||||
python3 -m venv .venv && .venv/bin/pip install h5py numpy netCDF4
|
||||
```
|
||||
|
||||
Set `CLAWHDF5_REQUIRE_INTEROP=1` in any automated runner so a missing
|
||||
interpreter is a failure rather than a skip.
|
||||
|
||||
|
||||
---
|
||||
|
||||
## Crafted B-tree v2 structures crash or exhaust the reader
|
||||
|
||||
**Status:** fixed on `main` (2026-09-20), after v2.6.0. **Every release up to
|
||||
and including v2.6.0 is affected.**
|
||||
|
||||
B-tree v2 traversal (`clawhdf5-format`, `btree_v2::collect_btree_v2_records`)
|
||||
recursed one frame per level with the depth taken from the file, and followed
|
||||
child addresses without checking whether they were shared. Two consequences
|
||||
for anyone reading untrusted files:
|
||||
|
||||
- A node that is its own child, under a header claiming 65 535 levels, overflows
|
||||
the stack and aborts the process. The file is under 100 bytes.
|
||||
- Levels whose children all point at one node below make the traversal visit it
|
||||
fan-out^depth times: ~30 million records from ~5 KB, and memory exhaustion one
|
||||
level deeper.
|
||||
|
||||
B-tree v2 backs dense attribute storage, v2 groups, shared object header
|
||||
messages and chunk indexes, so opening an object that uses any of them is
|
||||
enough. Both are now errors: depth is capped at 64, and traversal stops once it
|
||||
has produced more records than the file could physically hold.
|
||||
|
||||
---
|
||||
|
||||
## Extensible Array chunk indexes read back wrong data past the inline elements
|
||||
|
||||
**Status:** fixed on `main` (2026-09-20), after v2.6.0. **Every release up to
|
||||
and including v2.6.0 is affected.**
|
||||
|
||||
A dataset created with exactly one unlimited dimension (`maxshape=(None, ...)`,
|
||||
the usual append-only/resizable case) is indexed by an Extensible Array. Its
|
||||
index block holds the first `idx_blk_elmts` chunk entries inline — 4 by
|
||||
default — and everything after that lives in data blocks and super blocks whose
|
||||
layout `clawhdf5-format` computed incorrectly.
|
||||
|
||||
Consequences, by dataset size (1 chunk per element):
|
||||
|
||||
| chunks | result before the fix |
|
||||
|---|---|
|
||||
| <= 36 | correct (inline, plus two data blocks that happened to line up) |
|
||||
| 37 | 1 element wrong |
|
||||
| 400 | 364 elements wrong |
|
||||
| >= ~1000 | `invalid Extensible Array data block signature` |
|
||||
|
||||
The dangerous case is the middle one: values were returned from the wrong
|
||||
chunks rather than an error being raised. Any reader that accepted the data at
|
||||
face value saw plausible but incorrect numbers.
|
||||
|
||||
The root causes were the super block sizing formulas (`ndblks` and
|
||||
`dblk_nelmts` each double every *other* level, a half-step apart), a missing
|
||||
block-offset field in the super block, and a page-init bitmap read from the
|
||||
wrong structure. All four are fixed and covered by interop tests against
|
||||
HDF5 2.0 at sizes that cross each boundary, including paged data blocks.
|
||||
|
||||
Files written by this crate are unaffected — this was purely a read-path bug.
|
||||
|
||||
## Every `f32` dataset we wrote was unreadable by h5py / libhdf5
|
||||
|
||||
**Status:** fixed 2026-09-23, after v2.7.0. **Every
|
||||
release up to and including v2.7.0 is affected** — the encoder was already
|
||||
wrong in v2.1.0.
|
||||
|
||||
The floating-point datatype message carries the position of the sign bit
|
||||
(bits 8–15 of its class bit field). `clawhdf5-format` wrote 63 for every
|
||||
float, which is correct only for `f64`. libhdf5 validates the field, so opening
|
||||
any `f32` dataset written by this crate failed:
|
||||
|
||||
```
|
||||
KeyError: 'Unable to synchronously open object (sign bit position out of bounds)'
|
||||
```
|
||||
|
||||
That covers every agent store (`/memory/embeddings`, `norms` and
|
||||
`activation_weights` are `f32`). `clawhdf5` itself ignores the field on read,
|
||||
and the interop suites only ever wrote `f64` from our side, so nothing here
|
||||
noticed.
|
||||
|
||||
**Fix:** the sign position is computed from the type (`bit_offset +
|
||||
bit_precision - 1`: 15, 31, 63 for half, single, double). Regression tests:
|
||||
`float_sign_location_is_the_top_bit_of_the_value` (byte level),
|
||||
`clawhdf5_writes_f32_h5py_reads` and the agent's
|
||||
`h5py_reads_every_dataset_of_an_agent_store`.
|
||||
|
||||
**Existing files:** an agent store is rewritten in full at every checkpoint, so
|
||||
it becomes readable by h5py at its next checkpoint with a fixed build. Other
|
||||
files with `f32` datasets need to be rewritten.
|
||||
|
||||
## Empty datasets we wrote were unreadable by h5py / libhdf5
|
||||
|
||||
**Status:** fixed 2026-09-23, after v2.7.0. Every
|
||||
release up to and including v2.7.0 is affected.
|
||||
|
||||
A dataset with no elements was written with a real file address and a storage
|
||||
size of 0. libhdf5 guards contiguous storage with an overflow check
|
||||
(`addr + size <= addr`) that is always true when the size is 0, so it rejected
|
||||
the dataset:
|
||||
|
||||
```
|
||||
KeyError: 'Unable to synchronously open object (invalid dataset size, likely file corruption)'
|
||||
```
|
||||
|
||||
In practice: every agent store without sessions or a knowledge graph — the
|
||||
`/sessions` and `/knowledge_graph` datasets are empty until something is added
|
||||
— could not be read by h5py even once the `f32` bug above was fixed. Found by
|
||||
the same agent-store interop test.
|
||||
|
||||
**Fix:** an empty contiguous dataset gets the undefined address (all `0xff`),
|
||||
which is what libhdf5 itself writes.
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@redclaw/clawhdf5",
|
||||
"version": "2.5.0",
|
||||
"version": "2.7.0",
|
||||
"description": "Node.js bindings for clawhdf5 — HDF5-backed agent memory with hippocampal consolidation",
|
||||
"main": "index.js",
|
||||
"types": "index.d.ts",
|
||||
|
||||
+63
-2
@@ -20,6 +20,13 @@
|
||||
set -uo pipefail
|
||||
|
||||
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
|
||||
|
||||
# Interop suites drive a Python interpreter. On a PEP 668 "externally managed"
|
||||
# system h5py can only live in a virtualenv, so pick one up here — before any
|
||||
# test step, since the non-ignored interop suites read the same variable.
|
||||
if [ -z "${CLAWHDF5_PYTHON:-}" ] && [ -x "$SCRIPT_DIR/../.venv/bin/python" ]; then
|
||||
export CLAWHDF5_PYTHON="$SCRIPT_DIR/../.venv/bin/python"
|
||||
fi
|
||||
PASS=0
|
||||
FAIL=0
|
||||
STEPS=()
|
||||
@@ -70,6 +77,50 @@ run_step "cargo clippy (ann parallel)" cargo clippy \
|
||||
--features parallel \
|
||||
-- -D warnings
|
||||
|
||||
# zlib-ng is opt-in (`fast-deflate`; the default is pure-Rust zlib-rs), so
|
||||
# nothing above builds it. Keep it compiling and passing.
|
||||
run_step "cargo clippy (fast-deflate / zlib-ng)" cargo clippy \
|
||||
-p clawhdf5-format -p clawhdf5-filters -p clawhdf5 \
|
||||
--all-targets \
|
||||
--features clawhdf5-format/fast-deflate,clawhdf5-filters/fast-deflate \
|
||||
-- -D warnings
|
||||
|
||||
# The README promises that the core crates build no C by default. Hold it to
|
||||
# that: fail if a crate that compiles C (a *-sys crate, cc or cmake) enters the
|
||||
# default dependency tree of any of them. clawhdf5-migrate (bundled SQLite),
|
||||
# clawhdf5-napi (Node) and clawhdf5-gpu (graphics drivers) are exempt.
|
||||
no_c_in_default_build() {
|
||||
local crate found=0
|
||||
for crate in clawhdf5-format clawhdf5-io clawhdf5-filters clawhdf5 \
|
||||
clawhdf5-agent clawhdf5-ann clawhdf5-accel clawhdf5-netcdf4 clawhdf5-cli; do
|
||||
local c_deps
|
||||
c_deps=$(cargo tree -q -p "$crate" -e normal,build --prefix none \
|
||||
| grep -E '^([a-z0-9_-]+-sys|cc|cmake) v' | sort -u)
|
||||
if [ -n "$c_deps" ]; then
|
||||
echo "$crate pulls in C by default:"
|
||||
echo "$c_deps" | sed 's/^/ /'
|
||||
found=1
|
||||
fi
|
||||
done
|
||||
return $found
|
||||
}
|
||||
run_step "no C in the default build (core crates)" no_c_in_default_build
|
||||
|
||||
# The workspace declares a minimum Rust version (rust-version in Cargo.toml);
|
||||
# check that it really builds there, so the README badge and the manifests
|
||||
# cannot drift from the truth. Separate target dir: a different toolchain
|
||||
# would otherwise invalidate the main build.
|
||||
msrv_check() {
|
||||
local msrv
|
||||
msrv=$(sed -n 's/^rust-version = "\(.*\)"/\1/p' "$SCRIPT_DIR/../Cargo.toml")
|
||||
[ -n "$msrv" ] || { echo "no rust-version in Cargo.toml"; return 1; }
|
||||
rustup toolchain install "$msrv" --profile minimal >/dev/null || return 1
|
||||
echo "checking with Rust $msrv"
|
||||
CARGO_TARGET_DIR="$SCRIPT_DIR/../target/msrv" cargo "+$msrv" check \
|
||||
--workspace --exclude clawhdf5-py --all-targets
|
||||
}
|
||||
run_step "MSRV check" msrv_check
|
||||
|
||||
# 4. Tests (exclude clawhdf5-py)
|
||||
run_step "cargo test" cargo test \
|
||||
--workspace \
|
||||
@@ -83,14 +134,24 @@ run_step "cargo test (ann parallel)" cargo test \
|
||||
-p clawhdf5-ann \
|
||||
--features parallel
|
||||
|
||||
run_step "cargo test (fast-deflate / zlib-ng)" cargo test \
|
||||
-p clawhdf5-format -p clawhdf5-filters -p clawhdf5 \
|
||||
--features clawhdf5-format/fast-deflate,clawhdf5-filters/fast-deflate
|
||||
|
||||
# 5. Python interop suites. The h5py writer tests are #[ignore]d so a plain
|
||||
# `cargo test` stays hermetic; run them explicitly here.
|
||||
if python3 -c "import h5py" >/dev/null 2>&1 || [ "${CLAWHDF5_REQUIRE_INTEROP:-0}" = "1" ]; then
|
||||
# On a PEP 668 "externally managed" system h5py can only live in a
|
||||
# virtualenv, so honour CLAWHDF5_PYTHON (and a local .venv) rather than
|
||||
# skipping — the tests read the same variable.
|
||||
PYTHON="${CLAWHDF5_PYTHON:-python3}"
|
||||
if "$PYTHON" -c "import h5py" >/dev/null 2>&1 || [ "${CLAWHDF5_REQUIRE_INTEROP:-0}" = "1" ]; then
|
||||
run_step "h5py interop (format, ignored tests)" cargo test \
|
||||
-p clawhdf5-format --test writer_h5py_tests -- --include-ignored
|
||||
else
|
||||
echo ""
|
||||
echo "==> [h5py interop] SKIPPED: python3 with h5py not available"
|
||||
echo "==> [h5py interop] SKIPPED: no h5py in $PYTHON"
|
||||
echo " (set CLAWHDF5_PYTHON=/path/to/venv/bin/python, or create .venv;"
|
||||
echo " CLAWHDF5_REQUIRE_INTEROP=1 makes this a failure instead)"
|
||||
STEPS+=("SKIP: h5py interop (format, ignored tests)")
|
||||
fi
|
||||
|
||||
|
||||
Reference in New Issue
Block a user