Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
5889b378e9 | ||
|
|
18dc35f7e5 | ||
|
|
e17ab0ceef | ||
|
|
105cf13347 | ||
|
|
0529f72a2c | ||
|
|
a29c1b224b | ||
|
|
c0a9206703 | ||
|
|
57756e69ec | ||
|
|
6ad8ceb426 | ||
|
|
2e7e0456c1 | ||
|
|
dc0113d015 | ||
|
|
8ea455bbcb | ||
|
|
1e18ff5a86 | ||
|
|
306a35347c | ||
|
|
4c60398b30 | ||
|
|
7155409202 | ||
|
|
64d9c5f171 | ||
|
|
84a39ef3c5 | ||
|
|
55ed87d2e8 | ||
|
|
6531158d9f | ||
|
|
aa92fef7bb | ||
|
|
29baabbed2 | ||
|
|
b23946e62d | ||
|
|
52cfcf20b2 | ||
|
|
05c665a898 | ||
|
|
b36c6ec2af | ||
|
|
f3d63dbdcd | ||
|
|
0addf328bc | ||
|
|
d668e45ab5 | ||
|
|
c6a7bbfc67 | ||
|
|
3027380979 | ||
|
|
f507803ec1 | ||
|
|
8803d0754b | ||
|
|
42c3872ec9 | ||
|
|
c19199f3eb | ||
|
|
41db450c92 | ||
|
|
db4a067fe8 | ||
|
|
4aa3c5a1ca | ||
|
|
26e06cc5fd | ||
|
|
390a2e3836 | ||
|
|
f15bf2eb22 | ||
|
|
09480747aa | ||
|
|
39bf2bebf4 | ||
|
|
0ee698accd | ||
|
|
2bfbb7fb4b | ||
|
|
61424d1418 | ||
|
|
65d219c409 | ||
|
|
eb99de1020 |
@@ -33,8 +33,14 @@ jobs:
|
||||
/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
|
||||
|
||||
@@ -4,3 +4,4 @@ benchmarks/longmemeval/*.json
|
||||
|
||||
# Local model weights (MiniLM etc.) — large, not committed
|
||||
weights/
|
||||
.venv
|
||||
|
||||
+494
-1
@@ -28,6 +28,402 @@
|
||||
|
||||
---
|
||||
|
||||
## 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. What it costs is
|
||||
throughput — about 13% of QPS and 16% of build time at 100 000 x 384. So the
|
||||
setting trades ~13% of query speed for ~36% of the process's memory at equal
|
||||
recall. It is **off by default**: the right side of that trade depends on
|
||||
whether the deployment is short of memory or short of CPU.
|
||||
|
||||
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.
|
||||
|
||||
## Read harness
|
||||
|
||||
Produced by `cargo run --release -p clawhdf5-bench --bin read_harness`: a 4096 x
|
||||
2048 `f64` dataset (64 MB) written three ways, read in full and through four
|
||||
hyperslab selections, each from a fresh file handle. The last column is the
|
||||
point: does a selection cost what the *selection* costs?
|
||||
|
||||
### Baseline (v2.4.0): every selection decodes the whole dataset
|
||||
|
||||
4096 x 2048 f64 (64 MB per dataset), chunks 256 x 256, file 129 MB
|
||||
|
||||
| layout | read | selected | time ms | MB/s of selection | vs full read |
|
||||
|---|---|---:|---:|---:|---:|
|
||||
| chunked + deflate | full (first) | 64 MB | 181.8 | 352 | |
|
||||
| chunked + deflate | full (repeat) | 64 MB | 162.1 | 395 | 1.00x |
|
||||
| chunked + deflate | 64 x 64 window (1 chunk) | 0.03 MB | 104.89 | 0 | 0.577x |
|
||||
| chunked + deflate | 512 x 512 window (4-9 chunks) | 2.00 MB | 110.26 | 18 | 0.606x |
|
||||
| chunked + deflate | one row | 0.02 MB | 105.81 | 0 | 0.582x |
|
||||
| chunked + deflate | one column | 0.03 MB | 108.37 | 0 | 0.596x |
|
||||
| chunked | full (first) | 64 MB | 97.4 | 657 | |
|
||||
| chunked | full (repeat) | 64 MB | 86.7 | 738 | 1.00x |
|
||||
| chunked | 64 x 64 window (1 chunk) | 0.03 MB | 40.97 | 1 | 0.420x |
|
||||
| chunked | 512 x 512 window (4-9 chunks) | 2.00 MB | 44.66 | 45 | 0.458x |
|
||||
| chunked | one row | 0.02 MB | 30.88 | 1 | 0.317x |
|
||||
| chunked | one column | 0.03 MB | 30.27 | 1 | 0.311x |
|
||||
| contiguous | full (first) | 64 MB | 57.6 | 1112 | |
|
||||
| contiguous | full (repeat) | 64 MB | 53.5 | 1195 | 1.00x |
|
||||
| contiguous | 64 x 64 window (1 chunk) | 0.03 MB | 30.97 | 1 | 0.538x |
|
||||
| contiguous | 512 x 512 window (4-9 chunks) | 2.00 MB | 31.64 | 63 | 0.550x |
|
||||
| contiguous | one row | 0.02 MB | 31.90 | 0 | 0.554x |
|
||||
| contiguous | one column | 0.03 MB | 29.36 | 1 | 0.510x |
|
||||
|
||||
### After: partial reads
|
||||
|
||||
Only the rows of a contiguous dataset, or the chunks, that overlap the
|
||||
selection's bounding box are read/decoded. A 64 x 64 window of the compressed
|
||||
dataset: **105 -> 0.39 ms**; one row: **106 -> 2.7 ms**; one column:
|
||||
**108 -> 5.2 ms**. (Absolute full-read times differ between the two runs
|
||||
because the machine's speed drifted; compare the *vs full read* column.)
|
||||
|
||||
4096 x 2048 f64 (64 MB per dataset), chunks 256 x 256, file 129 MB
|
||||
|
||||
| layout | read | selected | time ms | MB/s of selection | vs full read |
|
||||
|---|---|---:|---:|---:|---:|
|
||||
| chunked + deflate | full (first) | 64 MB | 112.5 | 569 | |
|
||||
| chunked + deflate | full (repeat) | 64 MB | 104.5 | 612 | 1.00x |
|
||||
| chunked + deflate | 64 x 64 window (1 chunk) | 0.03 MB | 0.39 | 81 | 0.003x |
|
||||
| chunked + deflate | 512 x 512 window (4-9 chunks) | 2.00 MB | 4.85 | 412 | 0.043x |
|
||||
| chunked + deflate | one row | 0.02 MB | 2.69 | 6 | 0.024x |
|
||||
| chunked + deflate | one column | 0.03 MB | 5.23 | 6 | 0.046x |
|
||||
| chunked | full (first) | 64 MB | 70.0 | 915 | |
|
||||
| chunked | full (repeat) | 64 MB | 61.7 | 1037 | 1.00x |
|
||||
| chunked | 64 x 64 window (1 chunk) | 0.03 MB | 0.06 | 541 | 0.001x |
|
||||
| chunked | 512 x 512 window (4-9 chunks) | 2.00 MB | 1.99 | 1005 | 0.028x |
|
||||
| chunked | one row | 0.02 MB | 0.05 | 285 | 0.001x |
|
||||
| chunked | one column | 0.03 MB | 0.45 | 69 | 0.006x |
|
||||
| contiguous | full (first) | 64 MB | 60.3 | 1062 | |
|
||||
| contiguous | full (repeat) | 64 MB | 56.4 | 1134 | 1.00x |
|
||||
| contiguous | 64 x 64 window (1 chunk) | 0.03 MB | 0.08 | 396 | 0.001x |
|
||||
| contiguous | 512 x 512 window (4-9 chunks) | 2.00 MB | 2.12 | 944 | 0.035x |
|
||||
| contiguous | one row | 0.02 MB | 0.03 | 576 | 0.000x |
|
||||
| contiguous | one column | 0.03 MB | 2.55 | 12 | 0.042x |
|
||||
|
||||
### After: parallel cached decode, fewer copies (full reads)
|
||||
|
||||
Full-read times, old and new binaries run alternately at the same moment (this
|
||||
machine's absolute speed drifts over a long session, so only same-moment
|
||||
comparisons mean anything):
|
||||
|
||||
| layout (64 MB `f64`) | before | after |
|
||||
|---|---:|---:|
|
||||
| chunked + deflate | 110 ms | 69 ms |
|
||||
| chunked | 72 ms | 60 ms |
|
||||
| contiguous | 56 ms | 30 ms |
|
||||
|
||||
What changed: the facade's cached read path decompressed chunks one at a time
|
||||
(only the uncached reader was parallel) and pushed every chunk through a 16 MiB
|
||||
cache that a 64 MB read simply churns; it now decodes cache misses in parallel
|
||||
batches and caches only datasets that fit. Unfiltered chunks are copied
|
||||
straight from the file bytes instead of via two intermediate buffers. A
|
||||
contiguous dataset is converted straight from the file bytes (one copy instead
|
||||
of two), and the native-endian conversions no longer zero a buffer they are
|
||||
about to overwrite.
|
||||
|
||||
## Search harness baseline (v2.3.0)
|
||||
|
||||
Produced by `cargo run --release -p clawhdf5-bench --bin search_harness -- --full`
|
||||
on deterministic **clustered** synthetic data (384-dim, unit-normalised; points =
|
||||
cluster centre + noise — uniform random vectors are nearly equidistant in high
|
||||
dimension and say nothing about embeddings). Recall is measured against an exact
|
||||
brute-force scan, 200 queries. This is the *before* picture for the search
|
||||
hot-path work; every change to that path should be justified by a re-run.
|
||||
|
||||
Two things stand out:
|
||||
|
||||
* **HNSW recall does not respond to `ef`** and degrades sharply with size
|
||||
(0.87 → 0.67 → 0.31 recall@10 at 1K / 10K / 100K). Latency plateaus at the same
|
||||
point, i.e. the search exhausts the nodes it can reach: on clustered data the
|
||||
graph is poorly connected. The index selects neighbours by plain top-M
|
||||
distance rather than the HNSW paper's diversity heuristic.
|
||||
* **End-to-end `hybrid_search` is ~1000x slower than its vector stage** (49 ms
|
||||
vs ~0.03 ms at 10K; 884 ms at 100K). Each query rebuilds the BM25 index from
|
||||
scratch and rewrites the whole `.h5` file. The first query after `open()`
|
||||
additionally rebuilds the HNSW index (10.5 s at 100K).
|
||||
|
||||
### HNSW, N = 1000, dim = 384, M = 16, ef_construction = 64
|
||||
|
||||
build: 72.4 ms (13818 vectors/s) · exact scan: 3854 QPS, p50 258 µs
|
||||
|
||||
| ef | recall@10 | QPS | p50 µs | p99 µs |
|
||||
|---:|---:|---:|---:|---:|
|
||||
| 16 | 0.8710 | 59484 | 16 | 31 |
|
||||
| 32 | 0.8730 | 46302 | 21 | 25 |
|
||||
| 64 | 0.8730 | 31683 | 31 | 44 |
|
||||
| 128 | 0.8730 | 24715 | 40 | 49 |
|
||||
| 256 | 0.8730 | 24788 | 40 | 50 |
|
||||
|
||||
### HNSW, N = 10000, dim = 384, M = 16, ef_construction = 64
|
||||
|
||||
build: 802.5 ms (12461 vectors/s) · exact scan: 418 QPS, p50 2363 µs
|
||||
|
||||
| ef | recall@10 | QPS | p50 µs | p99 µs |
|
||||
|---:|---:|---:|---:|---:|
|
||||
| 16 | 0.6695 | 44031 | 19 | 51 |
|
||||
| 32 | 0.6705 | 45066 | 22 | 30 |
|
||||
| 64 | 0.6705 | 32746 | 30 | 41 |
|
||||
| 128 | 0.6705 | 27542 | 36 | 51 |
|
||||
| 256 | 0.6705 | 27754 | 36 | 49 |
|
||||
|
||||
### HNSW, N = 100000, dim = 384, M = 16, ef_construction = 64
|
||||
|
||||
build: 9752.6 ms (10254 vectors/s) · exact scan: 40 QPS, p50 24648 µs
|
||||
|
||||
| ef | recall@10 | QPS | p50 µs | p99 µs |
|
||||
|---:|---:|---:|---:|---:|
|
||||
| 16 | 0.3085 | 18046 | 57 | 84 |
|
||||
| 32 | 0.3110 | 21621 | 43 | 75 |
|
||||
| 64 | 0.3130 | 20015 | 49 | 70 |
|
||||
| 128 | 0.3135 | 15822 | 63 | 99 |
|
||||
| 256 | 0.3135 | 15308 | 66 | 124 |
|
||||
|
||||
### End to end: `HDF5Memory::hybrid_search` (k = 10, weights 0.7 / 0.3)
|
||||
|
||||
| N | ingest ms | checkpoint ms | open ms | first query ms | p50 ms | p99 ms | QPS |
|
||||
|---:|---:|---:|---:|---:|---:|---:|---:|
|
||||
| 1000 | 11 | 3.9 | 0.9 | 68.1 | 5.48 | 5.57 | 182.5 |
|
||||
| 10000 | 114 | 32.2 | 10.9 | 845.0 | 48.56 | 78.65 | 19.8 |
|
||||
| 100000 | 1486 | 713.0 | 354.5 | 10486.5 | 883.51 | 975.23 | 1.1 |
|
||||
wrote /tmp/claude-1000/-home-osobh-projects-clawhdf5/422f755e-dd25-4c35-8613-5439087e3aaa/scratchpad/baseline_full.json
|
||||
|
||||
### After: HNSW neighbour-selection heuristic
|
||||
|
||||
Same harness, same data, after replacing closest-M neighbour selection with the
|
||||
HNSW paper's diversity heuristic (Algorithm 4, keeping pruned connections) for
|
||||
both new links and back-link pruning. Recall@10 at `ef = 64`: **0.87 → 1.00**
|
||||
(1K), **0.67 → 1.00** (10K), **0.31 → 0.98** (100K), and it now rises with
|
||||
`ef` as it should. The cost is a slower build (extra distance evaluations per
|
||||
insert: ~3.5x at 10K); the distance-kernel work that follows targets that.
|
||||
|
||||
### HNSW, N = 1000, dim = 384, M = 16, ef_construction = 64
|
||||
|
||||
build: 221.3 ms (4519 vectors/s) · exact scan: 3851 QPS, p50 258 µs
|
||||
|
||||
| ef | recall@10 | QPS | p50 µs | p99 µs |
|
||||
|---:|---:|---:|---:|---:|
|
||||
| 16 | 0.9990 | 54760 | 18 | 29 |
|
||||
| 32 | 1.0000 | 40422 | 24 | 44 |
|
||||
| 64 | 1.0000 | 27744 | 36 | 51 |
|
||||
| 128 | 1.0000 | 13164 | 74 | 106 |
|
||||
| 256 | 1.0000 | 6879 | 144 | 175 |
|
||||
|
||||
### HNSW, N = 10000, dim = 384, M = 16, ef_construction = 64
|
||||
|
||||
build: 2733.5 ms (3658 vectors/s) · exact scan: 423 QPS, p50 2362 µs
|
||||
|
||||
| ef | recall@10 | QPS | p50 µs | p99 µs |
|
||||
|---:|---:|---:|---:|---:|
|
||||
| 16 | 0.9975 | 31321 | 27 | 61 |
|
||||
| 32 | 1.0000 | 32427 | 29 | 48 |
|
||||
| 64 | 1.0000 | 22738 | 42 | 62 |
|
||||
| 128 | 1.0000 | 10055 | 99 | 129 |
|
||||
| 256 | 1.0000 | 4649 | 214 | 266 |
|
||||
|
||||
### HNSW, N = 100000, dim = 384, M = 16, ef_construction = 64
|
||||
|
||||
build: 36472.8 ms (2742 vectors/s) · exact scan: 40 QPS, p50 24644 µs
|
||||
|
||||
| ef | recall@10 | QPS | p50 µs | p99 µs |
|
||||
|---:|---:|---:|---:|---:|
|
||||
| 16 | 0.9235 | 11394 | 82 | 194 |
|
||||
| 32 | 0.9675 | 12788 | 73 | 161 |
|
||||
| 64 | 0.9840 | 10406 | 91 | 186 |
|
||||
| 128 | 0.9990 | 7633 | 126 | 248 |
|
||||
| 256 | 0.9990 | 2823 | 352 | 510 |
|
||||
|
||||
### After: persistent keyword index, no store rewrite per query
|
||||
|
||||
`hybrid_search` used to rebuild the BM25 index from scratch (re-tokenising every
|
||||
record) and rewrite the whole `.h5` file on **every query**. The index is now
|
||||
kept for the life of the store and updated incrementally, and activation boosts
|
||||
are persisted by the next checkpoint instead of inside the query. Steady-state
|
||||
p50: **5.5 → 0.24 ms** (1K), **49 → 2.1 ms** (10K), **884 → 23 ms** (100K).
|
||||
|
||||
The first query after `open()` is slower than before (it pays for the better —
|
||||
slower — HNSW build plus the one-off keyword index build); persisting the HNSW
|
||||
index removes that.
|
||||
|
||||
### End to end: `HDF5Memory::hybrid_search` (k = 10, weights 0.7 / 0.3)
|
||||
|
||||
| N | ingest ms | checkpoint ms | open ms | first query ms | p50 ms | p99 ms | QPS |
|
||||
|---:|---:|---:|---:|---:|---:|---:|---:|
|
||||
| 1000 | 11 | 3.8 | 0.9 | 195.9 | 0.24 | 0.27 | 4130.4 |
|
||||
| 10000 | 104 | 31.1 | 10.9 | 2627.1 | 2.09 | 2.11 | 479.5 |
|
||||
| 100000 | 1436 | 684.7 | 278.0 | 36308.1 | 22.90 | 25.46 | 43.5 |
|
||||
|
||||
### After: vector index persisted with the checkpoint
|
||||
|
||||
The HNSW graph (not the vectors, which the store already holds) is saved to
|
||||
`<store>.h5.ann` at each checkpoint and reloaded by `open()`, tied to that
|
||||
checkpoint by a generation id. The index is now built once per store (the *cold
|
||||
index build* column — the first query ever), not once per session. First query
|
||||
after `open()`: **196 → 1.7 ms** (1K), **2627 → 15 ms** (10K),
|
||||
**36308 → 159 ms** (100K); what remains is the one-off keyword index build.
|
||||
Batch saves no longer force a full rebuild either: appended records join the
|
||||
index incrementally.
|
||||
|
||||
| N | ingest ms | cold index build ms | checkpoint ms | open ms | first query after open ms | p50 ms | p99 ms | QPS |
|
||||
|---:|---:|---:|---:|---:|---:|---:|---:|---:|
|
||||
| 1000 | 14 | 220 | 6.1 | 1.2 | 1.7 | 0.24 | 0.27 | 4049.9 |
|
||||
| 10000 | 120 | 2916 | 33.1 | 14.0 | 15.4 | 2.15 | 3.30 | 421.2 |
|
||||
| 100000 | 1591 | 40515 | 747.3 | 324.7 | 158.9 | 23.07 | 30.42 | 41.3 |
|
||||
|
||||
### After: unit-vector dot product, reusable visited set
|
||||
|
||||
Cosine distance recomputed both vector norms on every evaluation; the index now
|
||||
stores unit vectors and uses a plain dot product. The per-call `HashSet` of
|
||||
visited nodes became a reusable epoch-stamped array. Recall is unchanged.
|
||||
Build: **2.75 -> 1.89 s** (10K), **~38 -> 21 s** (100K). QPS at `ef = 64`:
|
||||
**22.7K -> 39K** (10K), **10.4K -> 14K** (100K).
|
||||
|
||||
### HNSW, N = 1000, dim = 384, M = 16, ef_construction = 64
|
||||
|
||||
build: 113.4 ms (8821 vectors/s) · exact scan: 4375 QPS, p50 225 µs
|
||||
|
||||
| ef | recall@10 | QPS | p50 µs | p99 µs |
|
||||
|---:|---:|---:|---:|---:|
|
||||
| 16 | 0.9990 | 144379 | 7 | 15 |
|
||||
| 32 | 1.0000 | 110654 | 9 | 17 |
|
||||
| 64 | 1.0000 | 80446 | 12 | 25 |
|
||||
| 128 | 1.0000 | 38220 | 26 | 36 |
|
||||
| 256 | 1.0000 | 20041 | 50 | 62 |
|
||||
|
||||
### HNSW, N = 10000, dim = 384, M = 16, ef_construction = 64
|
||||
|
||||
build: 1519.4 ms (6581 vectors/s) · exact scan: 422 QPS, p50 2368 µs
|
||||
|
||||
| ef | recall@10 | QPS | p50 µs | p99 µs |
|
||||
|---:|---:|---:|---:|---:|
|
||||
| 16 | 0.9975 | 54608 | 15 | 45 |
|
||||
| 32 | 1.0000 | 66009 | 14 | 24 |
|
||||
| 64 | 1.0000 | 49854 | 19 | 31 |
|
||||
| 128 | 1.0000 | 22403 | 45 | 57 |
|
||||
| 256 | 1.0000 | 10096 | 100 | 120 |
|
||||
|
||||
### HNSW, N = 100000, dim = 384, M = 16, ef_construction = 64
|
||||
|
||||
build: 21084.6 ms (4743 vectors/s) · exact scan: 39 QPS, p50 24739 µs
|
||||
|
||||
| ef | recall@10 | QPS | p50 µs | p99 µs |
|
||||
|---:|---:|---:|---:|---:|
|
||||
| 16 | 0.9235 | 15139 | 61 | 154 |
|
||||
| 32 | 0.9675 | 18181 | 53 | 121 |
|
||||
| 64 | 0.9840 | 13980 | 70 | 139 |
|
||||
| 128 | 0.9990 | 10959 | 86 | 174 |
|
||||
| 256 | 0.9990 | 3731 | 254 | 697 |
|
||||
|
||||
|
||||
### After: unranked keyword scores, top-k merge (rankings unchanged)
|
||||
|
||||
A fusion study (`search_harness --fusion-study`) showed that capping the
|
||||
keyword candidate pool is **not** a safe optimisation: against the current
|
||||
full-corpus normalisation the final top-10 overlap is only 0.83-0.92 and the
|
||||
first result changes for 10-35% of queries, for only a 2x saving. So the fusion
|
||||
semantics were left alone and the same answer made cheaper: fusion needs every
|
||||
keyword score but not their ranking, so BM25 now returns them unsorted from a
|
||||
dense accumulator (it hashed every posting and then sorted every match), and
|
||||
the merge selects its top k instead of sorting every candidate. Steady-state
|
||||
p50: **0.24 -> 0.07 ms** (1K), **2.1 -> 0.49 ms** (10K), **23 -> 4.65 ms**
|
||||
(100K) — **79x / 100x / 190x** faster than the v2.3.0 baseline, with identical
|
||||
results.
|
||||
|
||||
### End to end: `HDF5Memory::hybrid_search` (k = 10, weights 0.7 / 0.3)
|
||||
|
||||
| N | ingest ms | cold index build ms | checkpoint ms | open ms | first query after open ms | p50 ms | p99 ms | QPS |
|
||||
|---:|---:|---:|---:|---:|---:|---:|---:|---:|
|
||||
| 1000 | 11 | 112 | 4.0 | 1.1 | 1.4 | 0.07 | 0.08 | 14077.9 |
|
||||
| 10000 | 104 | 1487 | 33.8 | 13.7 | 13.9 | 0.49 | 0.51 | 2020.9 |
|
||||
| 100000 | 1376 | 20285 | 728.9 | 353.1 | 142.2 | 4.65 | 4.78 | 214.7 |
|
||||
|
||||
### After: batched bulk build (optionally parallel); deletions handled in search
|
||||
|
||||
Profiling showed **90% of a build's distance evaluations are in back-link
|
||||
pruning**. The bulk build now inserts in batches: plan each node's neighbours
|
||||
against the graph as it stood at the start of the batch, link, then prune every
|
||||
overflowing list once. That is less work even single-threaded (a node gaining
|
||||
several back-links in a batch is pruned once), and with the `parallel` feature
|
||||
planning and pruning run on a thread pool. The graph is deterministic and the
|
||||
same with or without the feature. Parallelising *within* one insert was tried
|
||||
first and gave only 1.45x on 16 cores (tasks too small).
|
||||
|
||||
| build | 1K | 10K | 100K |
|
||||
|---|---:|---:|---:|
|
||||
| v2.4.0 | 116 ms | 1676 ms | ~21 s |
|
||||
| batched | 83 ms | 1074 ms | 19.2 s |
|
||||
| batched + `parallel` (16 cores) | 34 ms | 388 ms | 5.9 s |
|
||||
|
||||
Recall on clustered data is unchanged or slightly better (100K, `ef = 64`:
|
||||
0.984 -> 0.9945). On uniform random data it dips slightly (10K, `ef = 64`:
|
||||
0.474 -> 0.444), the cost of batch members not seeing each other while
|
||||
planning; batches are capped at 1/16 of the graph and 512 nodes.
|
||||
|
||||
## Vector Search Latency
|
||||
|
||||
Brute-force cosine similarity over 384-dimensional embeddings (OpenAI text-embedding-3-small size).
|
||||
@@ -276,6 +672,102 @@ Session-level:
|
||||
| Vector only | 85.4% | 94.2% | 96.6% | 0.8901 |
|
||||
| Hybrid | **88.2%** | **95.8%** | **97.8%** | **0.9158** |
|
||||
|
||||
### Fusion method — weighted vs. RRF, full haystack, n=500
|
||||
|
||||
Reciprocal rank fusion has been in the codebase since early on but was only
|
||||
reachable as a free function over a linear scan, so it had never been compared
|
||||
with the weighted sum on equal terms. `HDF5Memory::hybrid_search_with` now
|
||||
takes a `Fusion`, and both run over the same HNSW + BM25 candidates:
|
||||
|
||||
| Mode | turn Hit@1 | Hit@5 | Hit@10 | MRR | session Hit@1 | session MRR |
|
||||
|---|---|---|---|---|---|---|
|
||||
| BM25 only | **53.8%** | 75.0% | 81.6% | 0.6320 | 86.2% | 0.8948 |
|
||||
| Vector only | 36.0% | 71.8% | 81.6% | 0.5031 | 85.4% | 0.8901 |
|
||||
| **Weighted 0.4 / 0.6** | 51.6% | **81.4%** | **87.8%** | **0.6430** | **91.0%** | **0.9347** |
|
||||
| RRF (k=60) | 45.0% | 78.8% | 87.6% | 0.5967 | 89.6% | 0.9253 |
|
||||
|
||||
**RRF loses to the tuned weighted sum** — 6.6pp of turn Hit@1 and 0.046 of MRR
|
||||
— and lands almost exactly where the old `0.7/0.3` weighting did (44.2% /
|
||||
0.5856). That is not a coincidence: RRF combines the two stages by rank with
|
||||
*equal* influence, and on this corpus the stages are not equally good. BM25
|
||||
alone beats the vector stage by 17.8pp at Hit@1, so any scheme that treats them
|
||||
as peers gives up rank-1 accuracy, and RRF discards the score magnitudes that
|
||||
would say which stage to believe.
|
||||
|
||||
This is a property of the corpus, not a defect in RRF: its selling point is
|
||||
robustness when the two stages' scores are not comparable and there is no
|
||||
labelled data to tune against. Here there is, so the weighted sum is kept as
|
||||
the default. `Fusion::Rrf` remains available for callers whose stages are more
|
||||
evenly matched.
|
||||
|
||||
### Keyword tokenizer — stemming, full haystack, n=500
|
||||
|
||||
The keyword stage lowercases and splits on non-alphanumerics, with no stemming,
|
||||
so "training" and "trains" are unrelated terms. `TokenFilter::Stemmed` strips
|
||||
common English inflections (plurals, `-ing`/`-ed`, with consonant un-doubling)
|
||||
from documents and queries alike. Turn-level:
|
||||
|
||||
| Mode | Hit@1 | Hit@5 | Hit@10 | MRR | session Hit@1 |
|
||||
|---|---|---|---|---|---|
|
||||
| BM25 only | **53.8%** | 75.0% | 81.6% | 0.6320 | 86.2% |
|
||||
| BM25 only, stemmed | 52.0% | 77.8% | 84.0% | 0.6320 | 88.0% |
|
||||
| Hybrid 0.4/0.6 | 51.6% | **81.4%** | 87.8% | **0.6430** | 91.0% |
|
||||
| Hybrid 0.4/0.6, stemmed | 50.2% | **81.4%** | **88.2%** | 0.6394 | **91.4%** |
|
||||
|
||||
**Stemming is a trade, not a win, and the default stays off.** It reliably buys
|
||||
depth and costs the top rank: on BM25 alone, +2.8pp Hit@5 and +2.4pp Hit@10 for
|
||||
−1.8pp Hit@1, with MRR unchanged to four decimal places — the gains deeper down
|
||||
exactly offset the loss at rank 1. That is what conflation does: merging
|
||||
"train"/"training"/"trains" surfaces documents an exact-match query would never
|
||||
reach, and also lets a near-miss outrank the exact hit.
|
||||
|
||||
On the configuration that actually ships (hybrid 0.4/0.6) the trade is
|
||||
narrower still — Hit@5 identical, Hit@10 +0.4pp, Hit@1 −1.4pp, MRR −0.004 —
|
||||
because the vector stage already supplies much of the recall stemming would
|
||||
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
|
||||
@@ -318,7 +810,8 @@ BM25 at Hit@1. Both dominate `0.7/0.3`.
|
||||
|
||||
The rows below are kept at the three original settings because they are what the
|
||||
mode ablation measured — read them as "the shape of each stage in isolation",
|
||||
and take the operating point from the sweep.
|
||||
and take the operating point from the sweep. `0.4/0.6` is now the shipped
|
||||
default (`hybrid::DEFAULT_FUSION`).
|
||||
|
||||
The same pattern shows up independently in omni-cortex's four-signal RRF ablation,
|
||||
where adding BM25 to a dense retriever raised nDCG@5 while lowering Hit@1 and MRR.
|
||||
|
||||
+271
@@ -1,5 +1,276 @@
|
||||
# Changelog
|
||||
|
||||
## v2.6.0 (2026-09-20)
|
||||
|
||||
### Upgrade Notes
|
||||
- **Re-ranked results change, substantially for the better.** `RerankInput`
|
||||
and `ReRankConfig` gained fields (`relevance`, `relevance_weight`), so
|
||||
literal constructions need updating; `..Default::default()` does not. Any
|
||||
caller that re-ranked was previously getting results ordered by age with the
|
||||
retrieval score discarded — see below.
|
||||
- **Breaking:** `MemoryCache::embeddings` is a `cache::Embeddings` rather than
|
||||
a `Vec<Vec<f32>>` (indexing still yields a `&[f32]` row); `embeddings_flat`
|
||||
is gone, replaced by `flat_embeddings()`; `rebuild_flat()` is a deprecated
|
||||
no-op.
|
||||
- `MemoryConfig` gained `quantized_index` (default `false`, so behaviour is
|
||||
unchanged unless you opt in); literal constructions need the field.
|
||||
|
||||
### Retrieval quality
|
||||
- `clawhdf5-agent`: **re-ranking discarded the retrieval score.**
|
||||
`reranker::rerank` built its combined score from temporal decay, source
|
||||
authority and Hebbian activation only — `RerankInput` had no relevance field
|
||||
— so re-ranking a candidate pool reordered it by age and threw the
|
||||
retriever's ordering away. The OpenClaw backend re-ranked every search, so
|
||||
this was its shipping behaviour: measured over the full LongMemEval haystack
|
||||
it cost **40.6pp of Hit@1** (11.0% vs 51.6%) and two thirds of MRR (0.183 vs
|
||||
0.643). `RerankInput::relevance` and `ReRankConfig::relevance_weight` (1.0 by
|
||||
default) fix it: relevance leads and the metadata signals break near-ties,
|
||||
which restores retrieval (Hit@1 +0.4pp vs no re-ranking) and improves
|
||||
recency discrimination by 6–7pp. **Breaking:** `RerankInput` and
|
||||
`ReRankConfig` gained fields, so literal constructions need updating;
|
||||
`..Default::default()` does not.
|
||||
- `clawhdf5-bench`: the LongMemEval harness feeds the dataset's real session
|
||||
dates to the store instead of a synthetic counter (decay needs true
|
||||
intervals, not just the right order), and reports `newest_gold_first` — on a
|
||||
`knowledge-update` question, did the newest gold session outrank the stale
|
||||
one it supersedes? Plain recall cannot see this, because both are labelled
|
||||
gold. New `--rerank-sweep`.
|
||||
|
||||
### Memory
|
||||
- `clawhdf5-agent`: **`MemoryConfig::quantized_index`** stores the vector
|
||||
index's own copy of the embeddings as `i8` rather than `f32`, which at 100k
|
||||
384-dim entries takes the index from 266 to 123 MiB and the whole reopened
|
||||
store from 399 to 256 MiB (2.72x -> **1.74x** the raw vectors). Quantised
|
||||
distances are approximate and `ef` cannot compensate — recall@10 tops out at
|
||||
0.967 against f32's 0.9995 — so the query path re-scores the candidate pool
|
||||
against the exact embeddings the store already holds, which restores recall
|
||||
(0.9940 vs 0.9945 at ef=64) for about 13% of QPS. **Off by default**: it
|
||||
trades query speed for memory, and which side is worth more depends on the
|
||||
deployment. The setting is persisted, so a reopened store does not silently
|
||||
revert to four times the index memory.
|
||||
- `clawhdf5-ann`: `Storage::Int8` and the `build_with` / `new_with` /
|
||||
`from_graph_bytes_with` constructors that select it. The scale is per row,
|
||||
not global — a fixed `[-1, 1]` scale spends fewer than 12 of the 255 levels
|
||||
on a unit-length 128-dim vector and is unusable (0.35 top-10 overlap against
|
||||
an exact ranking, versus 0.99 per row). `compact()` keeps the storage it was
|
||||
given; serialized indexes still carry f32 vectors, so a quantised index is
|
||||
rebuilt rather than loaded.
|
||||
- `clawhdf5-agent`: **a loaded store holds ~30% less memory** (100k 384-dim
|
||||
entries: 505 -> 357 MiB, 3.44x -> 2.43x the raw vectors). The cache kept
|
||||
every embedding twice — a `Vec<Vec<f32>>` and a flattened copy for the
|
||||
batched kernels, maintained in lock-step — so it now stores only the flat
|
||||
buffer and indexes into it. Recall and query latency are unchanged.
|
||||
**Breaking:** `MemoryCache::embeddings` is a `cache::Embeddings` rather than
|
||||
a `Vec<Vec<f32>>` (indexing still yields a `&[f32]` row); `embeddings_flat`
|
||||
is gone, replaced by `flat_embeddings()`; `rebuild_flat()` is a deprecated
|
||||
no-op. Rows are now always exactly `dim` long — shorter ones are
|
||||
zero-padded — which makes the ragged-row case that used to silently
|
||||
misalign the flattened copy unrepresentable.
|
||||
- `clawhdf5-bench`: `search_harness --footprint` reports live heap use per
|
||||
stage, measured with a counting allocator (RSS cannot see a structure freed
|
||||
into the allocator's own pool).
|
||||
|
||||
### Testing
|
||||
- The Python interop suites honour **`CLAWHDF5_PYTHON`**, and `ci-test.sh`
|
||||
picks up a `.venv/bin/python` automatically. On a PEP 668 "externally
|
||||
managed" system h5py cannot be installed into the system interpreter at all,
|
||||
so every interop suite — the h5py writer round-trips, the facade, netCDF4
|
||||
and the reference files — was skipping silently. A silent skip here is
|
||||
exactly how the v5 compound-datatype bug reached a release.
|
||||
`CLAWHDF5_REQUIRE_INTEROP=1` still turns a skip into a failure.
|
||||
|
||||
## v2.5.0 (2026-09-19)
|
||||
|
||||
### Upgrade Notes
|
||||
- **Retrieval rankings change, for the better.** The default fusion weights
|
||||
move from `0.7/0.3` to `0.4/0.6` (`hybrid::DEFAULT_FUSION`), measured over the
|
||||
full LongMemEval haystack: turn-level Hit@1 51.6% vs 44.2%, MRR 0.643 vs
|
||||
0.586. `unified_search` and the OpenClaw backend pick this up automatically;
|
||||
callers passing weights to `hybrid_search` explicitly are unaffected.
|
||||
- **Out-of-range selections are now errors.** `read_*_selection` used to return
|
||||
data for a selection that ran past a dataset edge — a hyperslab came back
|
||||
zero-padded, and a point with an out-of-range coordinate wrapped into the
|
||||
next row. Both are now `FormatError::SelectionOutOfBounds`. Code relying on
|
||||
the old (wrong) values will start seeing errors.
|
||||
- **Large compressed datasets written without explicit chunk dimensions get a
|
||||
different layout.** They used to be stored as one chunk; they are now split
|
||||
to ~1 MiB chunks. The files stay standard and h5py-readable, and explicit
|
||||
`with_chunks` is unaffected.
|
||||
- `rayon` is now a default dependency of `clawhdf5-agent` (the parallel index
|
||||
build). Opt out with `--no-default-features --features float16,hnsw`.
|
||||
- `clawhdf5-ann` search results no longer shrink when records near the query
|
||||
have been deleted, so a search that previously returned fewer than `k`
|
||||
results now returns `k`.
|
||||
|
||||
### Retrieval quality
|
||||
- `clawhdf5-agent`: optional keyword stemming — `bm25::TokenFilter::Stemmed`
|
||||
and `HDF5Memory::set_token_filter`, so "training" and "trains" match. **Off
|
||||
by default**, on measurement rather than principle: over the full LongMemEval
|
||||
haystack it buys depth and costs the top rank (BM25 alone: Hit@5 +2.8pp,
|
||||
Hit@10 +2.4pp, Hit@1 −1.8pp, MRR unchanged), and on the shipping hybrid
|
||||
configuration the trade is narrower still. See `BENCHMARKS.md`.
|
||||
- `clawhdf5-agent`: **`QueryExpander::expand` panicked on ordinary non-ASCII
|
||||
input** — `"İ AI"` was enough. It searched a lowercased copy of the query and
|
||||
then sliced the *original* with those offsets, which only works while
|
||||
lowercasing preserves byte length (Turkish `İ` is 2 bytes and lowercases to
|
||||
3). Depending on where the offsets drifted it either corrupted the output
|
||||
("İstanbul AI trip" lost a character) or panicked. Matching now walks the
|
||||
original string.
|
||||
- `clawhdf5-agent`: query expansion no longer rewrites text inside words.
|
||||
`replace_word_case_insensitive` did a plain substring replace despite its
|
||||
name, so "training" became "trArtificial Intelligencening" and "programming"
|
||||
became "Pull Requestogramming" — every acronym expansion of ordinary prose
|
||||
was corrupt. Matches now require word boundaries; genuine acronyms
|
||||
(`API`, `database`) still expand.
|
||||
- `clawhdf5-agent`: **the default fusion weights are now the measured ones.**
|
||||
A sweep of every 0.1 step over the full LongMemEval haystack (500 questions,
|
||||
real MiniLM embeddings) shows the long-standing `0.7/0.3` default is
|
||||
*strictly dominated* by `0.4/0.6` — turn-level Hit@1 51.6% vs 44.2%, Hit@5
|
||||
81.4% vs 79.2%, Hit@10 87.8% vs 85.8%, MRR 0.643 vs 0.586, and better at
|
||||
session level too. The finding was recorded in `BENCHMARKS.md` but had never
|
||||
been applied: `unified_search` and the OpenClaw backend both hardcoded
|
||||
`0.7/0.3`. They now use `hybrid::DEFAULT_FUSION`. **Callers passing weights
|
||||
to `hybrid_search` explicitly are unaffected** — pass `0.4`/`0.6` (or use
|
||||
`hybrid_search_with`) to get the tuned behaviour.
|
||||
- `clawhdf5-agent`: fusion is now selectable. New `hybrid::Fusion`
|
||||
(`Weighted { vector, keyword }` or `Rrf { k }`), `hybrid::fuse`,
|
||||
`hybrid::hybrid_search_fused` and `HDF5Memory::hybrid_search_with`.
|
||||
Reciprocal rank fusion existed but was unreachable from the store, so it had
|
||||
never been measured against the weighted sum; the LongMemEval bench now has
|
||||
an `RRF` mode.
|
||||
|
||||
### HDF5 Read Path
|
||||
- **Selection reads cost what the selection costs.** `read_*_selection` decoded
|
||||
the *entire* dataset and then picked elements out, so a 64 x 64 window of a
|
||||
64 MB compressed dataset took 105 ms - about as long as reading all of it.
|
||||
Now only the rows (contiguous) or chunks that overlap the selection's
|
||||
bounding box are read and decompressed: that window takes 0.39 ms, one row
|
||||
2.7 ms, one column 5.2 ms. Results are identical to the full-read path
|
||||
(equivalence-tested over random hyperslabs and point lists, ranks 1-3,
|
||||
contiguous / chunked / deflate). New `read_harness` bench binary.
|
||||
- **Faster full reads** (same-moment A/B, 64 MB `f64`): chunked + deflate
|
||||
110 -> 69 ms, chunked 72 -> 60 ms, contiguous 56 -> 30 ms. The facade's
|
||||
cached read path now decompresses cache misses in parallel batches (it was
|
||||
sequential; only the uncached reader was parallel) and caches only datasets
|
||||
that fit the chunk cache; unfiltered chunks are copied straight from the file
|
||||
bytes; a contiguous dataset is converted straight from the file bytes; and
|
||||
the native-endian conversions no longer zero a buffer before overwriting it.
|
||||
- **Datasets indexed by a version-2 B-tree now read** (layout v4, chunk index
|
||||
type 5 — what `libver='latest'` uses for two or more unlimited dimensions;
|
||||
previously "unsupported chunked layout"). The four copies of the chunk-index
|
||||
dispatch are now one shared function, so every read path gets it.
|
||||
- **`H5T_STD_REF` references** (HDF5 1.12+, datatype message version 4) parse:
|
||||
`ReferenceType` gains `Object2`, `DatasetRegion2` and `Attribute`, and
|
||||
`read_object_references` decodes the new object references. Previously any
|
||||
dataset of this type failed with `InvalidReferenceType(2)`. Tested against a
|
||||
file written by HDF5 2.0 itself (fixture + generator script committed).
|
||||
- **Automatic chunk sizes.** Asking for compression (or any filter) without
|
||||
`with_chunks` used to store the whole dataset as one chunk, so any read had
|
||||
to decompress everything and nothing could be decoded in parallel. Datasets up
|
||||
to 1 MiB stay a single chunk, as before; larger ones are split by halving the
|
||||
dimensions in turn until a chunk is at most 1 MiB (the approach h5py takes).
|
||||
**Behaviour change:** large compressed datasets written without explicit
|
||||
chunk dimensions get a different (standard, h5py-readable) layout. Explicit
|
||||
`with_chunks` is unaffected.
|
||||
- **Out-of-range selections are errors.** They used to return data: a hyperslab
|
||||
past an edge came back padded with zeros, and a point whose column was out of
|
||||
range wrapped into the next row and returned that element. Now
|
||||
`FormatError::SelectionOutOfBounds` (also for a rank mismatch or overlapping
|
||||
blocks).
|
||||
|
||||
### Search
|
||||
- `clawhdf5-ann`: **faster index builds.** Back-link pruning is 90% of a
|
||||
build's distance evaluations; the bulk build now inserts in batches and
|
||||
prunes each overflowing neighbour list once per batch (10K: 1676 -> 1074 ms).
|
||||
With the `parallel` feature, planning and pruning run on a thread pool (10K:
|
||||
388 ms, 100K: ~21 s -> 5.9 s on 16 cores). The graph is deterministic and
|
||||
identical with or without the feature. `clawhdf5-agent`'s `parallel` feature
|
||||
enables it for the agent's index and is now **on by default** (adds `rayon`
|
||||
to the default dependency set; build with `--no-default-features --features
|
||||
float16,hnsw` to opt out).
|
||||
- `clawhdf5-ann`: `HnswIndex::search` returned fewer than `k` results — often
|
||||
none — when the records nearest the query had been deleted: it collected `ef`
|
||||
candidates, *then* dropped the deleted ones, *then* took `k`. Deleted nodes
|
||||
are now traversed as waypoints but never occupy a result slot, so a search
|
||||
returns the `k` nearest live records. Matters for any store that deletes or
|
||||
supersedes memories without compacting straight away.
|
||||
|
||||
## v2.4.0 (2026-09-19)
|
||||
|
||||
### Upgrade Notes
|
||||
- **Search results improve on upgrade.** The HNSW index now reaches true
|
||||
neighbours it previously could not (recall@10 0.31 -> 0.98 at 100K records on
|
||||
clustered data), so `hybrid_search` rankings change for the better. The agent
|
||||
rebuilds its index from the store automatically; a standalone `HnswIndex`
|
||||
persisted with `to_hdf5_bytes` keeps its old graph until rebuilt.
|
||||
- **`hybrid_search` no longer writes the store.** Hebbian activation boosts are
|
||||
persisted by the next checkpoint (any flushing write, `flush_wal`, or when
|
||||
the `HDF5Memory` is dropped) instead of inside every query; a crash before
|
||||
then forgets only the boosts since the last checkpoint. Activation weights
|
||||
are now capped at 16.
|
||||
- A new sidecar file, `<store>.h5.ann`, holds the vector index graph. It is
|
||||
derived data: safe to delete (the index is rebuilt), copied by `snapshot()`,
|
||||
and worth including when copying a store by hand to avoid a rebuild.
|
||||
- `BM25Index` no longer caches IDF and gained `add_document`,
|
||||
`remove_document`, `pad_to`, `scores`, `len` and `is_empty`; results are now
|
||||
deterministic (ties break by record id).
|
||||
|
||||
### Search
|
||||
- `clawhdf5-ann`: **HNSW recall fix.** Neighbours were chosen as the plain
|
||||
closest-M, which on clustered data (what embeddings look like) turns each
|
||||
cluster into an island: recall@10 was 0.87 / 0.67 / 0.31 at 1K / 10K / 100K
|
||||
vectors and did not improve with `ef`. The index now uses the HNSW paper's
|
||||
diversity heuristic (Algorithm 4 with kept pruned connections) when linking a
|
||||
new node and when pruning back-links: recall@10 at `ef = 64` is 1.00 / 1.00 /
|
||||
0.98 and responds to `ef`. Builds are slower (~3.5x at 10K). Existing
|
||||
persisted indexes keep their old graph until rebuilt; the agent rebuilds its
|
||||
index from the cache, so stores pick this up automatically.
|
||||
- `clawhdf5-agent`: **`hybrid_search` is 23-39x faster in steady state** (p50
|
||||
5.5 -> 0.24 ms at 1K records, 49 -> 2.1 ms at 10K, 884 -> 23 ms at 100K).
|
||||
Every query used to rebuild the BM25 index from scratch and rewrite the whole
|
||||
`.h5` file. The keyword index now lives for the life of the store and is
|
||||
updated incrementally (add / remove / in-place update, exactly equivalent to
|
||||
a fresh build - property-tested), and a query no longer writes the store.
|
||||
**Behaviour change:** Hebbian activation boosts are persisted by the next
|
||||
checkpoint (any flushing write, `flush_wal`, or drop) rather than
|
||||
immediately; a crash in between forgets only the boosts since the last
|
||||
checkpoint. Activation weights are now capped (16.0) - they grew without
|
||||
bound.
|
||||
- `clawhdf5-agent`: **the vector index is persisted**, so `open()` no longer
|
||||
rebuilds it on the first search (first query after open: 2627 -> 15 ms at 10K
|
||||
records, 36 s -> 159 ms at 100K). The HNSW graph — not the vectors, which the
|
||||
store already holds — is written to `<store>.h5.ann` at each checkpoint and
|
||||
tied to it by a generation id in `/meta`; a missing, stale, damaged or
|
||||
structurally invalid sidecar is ignored and the index rebuilt. Records
|
||||
replayed from the WAL join the loaded index incrementally; a replayed update
|
||||
or delete invalidates it. `snapshot()` copies it. Batch saves no longer force
|
||||
a full index rebuild.
|
||||
- `clawhdf5-ann`: faster HNSW build and search with identical recall. The
|
||||
cosine metric stores unit vectors and compares them with a plain dot product
|
||||
(it re-derived both norms on every distance evaluation), and the per-call
|
||||
`HashSet` of visited nodes is a reusable epoch-stamped array. Build 2.75 ->
|
||||
1.89 s at 10K and ~38 -> 21 s at 100K; QPS at `ef = 64` 22.7K -> 39K at 10K.
|
||||
Distances returned by `search` are unchanged (1 - cosine). Indexes loaded
|
||||
from older HDF5 files are normalised on load.
|
||||
- `clawhdf5-accel`: the SIMD backend is detected once per process instead of
|
||||
on every kernel call.
|
||||
- `clawhdf5-ann`: `HnswIndex::graph_to_bytes` / `from_graph_bytes` — graph-only
|
||||
serialization (checksummed, every neighbour id and level validated on load).
|
||||
- `clawhdf5-agent`: a further 4-5x on `hybrid_search` with **identical
|
||||
rankings** (p50 now 0.07 / 0.49 / 4.65 ms at 1K / 10K / 100K — 79x / 100x /
|
||||
190x faster than v2.3.0). Fusion needs every keyword score but not their
|
||||
ranking: new `BM25Index::scores` returns them unsorted from a dense
|
||||
accumulator (it hashed every posting, then sorted every match), and
|
||||
`merge_vector_keyword` selects its top k instead of sorting every candidate.
|
||||
Capping the keyword candidate pool was measured and rejected: it changes the
|
||||
top-10 for most queries (`search_harness --fusion-study`).
|
||||
- `clawhdf5-agent`: BM25 results are deterministic (ties break by record id),
|
||||
top-k uses a bounded heap, and the "WAND early termination" that computed a
|
||||
bound and then ignored it is gone. IDF is computed per query.
|
||||
- `clawhdf5-bench`: new `search_harness` binary — HNSW recall@10 / QPS / latency
|
||||
per `ef` against an exact scan, and end-to-end `hybrid_search` timings, on
|
||||
deterministic clustered (or `--uniform`) data. Baseline in `BENCHMARKS.md`.
|
||||
|
||||
## v2.3.0 (2026-09-19)
|
||||
|
||||
### Upgrade Notes
|
||||
|
||||
@@ -33,6 +33,24 @@ Cargo workspace with 16 crates under `crates/` (plus `libaec-sys`, an internal F
|
||||
the approximate `clawhdf5-ann` index for the vector stage (the index mirrors
|
||||
the cache and self-heals on drift). Build the agent with
|
||||
`--no-default-features --features float16` to force the exact linear cosine scan.
|
||||
The agent's `parallel` feature (also default) builds the index on a thread
|
||||
pool; the graph is identical with or without it.
|
||||
The index uses the HNSW paper's diversity heuristic for neighbour selection
|
||||
(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). `MemoryConfig::quantized_index` (off by
|
||||
default, persisted) 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 and costs
|
||||
~13% of QPS. `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
|
||||
`cargo run --release -p clawhdf5-bench --bin search_harness` (baselines in
|
||||
`BENCHMARKS.md`).
|
||||
- WAL (write-ahead log) for crash-safe persistence, with a chained CRC32
|
||||
trailer per entry (each entry's CRC folds in the previous entry's CRC) so a
|
||||
corrupted, reordered, duplicated, or spliced entry stops replay cleanly
|
||||
|
||||
+1
-1
@@ -21,7 +21,7 @@ members = [
|
||||
resolver = "2"
|
||||
|
||||
[workspace.package]
|
||||
version = "2.3.0"
|
||||
version = "2.6.0"
|
||||
edition = "2024"
|
||||
license = "MIT"
|
||||
repository = "https://git.redclaw.dev/quantumclaw/clawhdf5"
|
||||
|
||||
@@ -432,6 +432,13 @@ ClawhDF5's agent memory design draws from 15+ recent papers:
|
||||
| `agent` | no | Full agent memory layer |
|
||||
| `float16` | **yes** | Half-precision embedding storage (2× compression) |
|
||||
| `hnsw` | **yes** | HNSW approximate vector index for `hybrid_search` (via `clawhdf5-ann`); disable for an exact linear scan |
|
||||
|
||||
`MemoryConfig::quantized_index` (off by default) stores 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 — recall matches the `f32` index, at about
|
||||
13% fewer queries per second. See `BENCHMARKS.md`, "Quantising the index copy".
|
||||
| `parallel` | no | Rayon parallel search |
|
||||
| `fast-math` | no | BLAS matrix-vector multiply |
|
||||
| `accelerate` | no | Apple Accelerate / AMX (macOS) |
|
||||
@@ -492,6 +499,14 @@ cargo build -p clawhdf5-agent --features "agent,float16,accelerate,parallel,gpu"
|
||||
# Tests
|
||||
cargo test --workspace # all 1,650+ 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
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
[package]
|
||||
name = "clawhdf5-accel"
|
||||
version = "2.3.0"
|
||||
version = "2.6.0"
|
||||
edition = "2024"
|
||||
description = "SIMD-accelerated operations for rustyhdf5"
|
||||
license = "MIT"
|
||||
|
||||
@@ -61,8 +61,14 @@ pub enum Backend {
|
||||
Scalar,
|
||||
}
|
||||
|
||||
/// Detect the best available SIMD backend at runtime.
|
||||
/// The best available SIMD backend, detected once per process. Every kernel
|
||||
/// dispatches through this, so it sits in the innermost loop of every search.
|
||||
pub fn detect_backend() -> Backend {
|
||||
static BACKEND: std::sync::OnceLock<Backend> = std::sync::OnceLock::new();
|
||||
*BACKEND.get_or_init(detect_backend_uncached)
|
||||
}
|
||||
|
||||
fn detect_backend_uncached() -> Backend {
|
||||
#[cfg(target_arch = "aarch64")]
|
||||
{
|
||||
return Backend::Neon; // Always available on aarch64
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
[package]
|
||||
name = "clawhdf5-agent"
|
||||
version = "2.3.0"
|
||||
version = "2.6.0"
|
||||
edition = "2024"
|
||||
description = "HDF5-backed persistent memory store for on-device AI agents"
|
||||
license = "MIT"
|
||||
@@ -10,12 +10,12 @@ keywords = ["agent", "memory", "hdf5", "vector-search", "embedding"]
|
||||
categories = ["database", "science", "algorithms"]
|
||||
|
||||
[dependencies]
|
||||
clawhdf5-format = { path = "../clawhdf5-format", version = "2.3.0", features = ["parallel", "fast-checksum"] }
|
||||
clawhdf5 = { path = "../clawhdf5", version = "2.3.0" }
|
||||
clawhdf5-io = { path = "../clawhdf5-io", version = "2.3.0", features = ["mmap"] }
|
||||
clawhdf5-accel = { path = "../clawhdf5-accel", version = "2.3.0" }
|
||||
clawhdf5-ann = { path = "../clawhdf5-ann", version = "2.3.0", optional = true }
|
||||
clawhdf5-gpu = { path = "../clawhdf5-gpu", version = "2.3.0", optional = true, default-features = false }
|
||||
clawhdf5-format = { path = "../clawhdf5-format", version = "2.6.0", features = ["parallel", "fast-checksum"] }
|
||||
clawhdf5 = { path = "../clawhdf5", version = "2.6.0" }
|
||||
clawhdf5-io = { path = "../clawhdf5-io", version = "2.6.0", features = ["mmap"] }
|
||||
clawhdf5-accel = { path = "../clawhdf5-accel", version = "2.6.0" }
|
||||
clawhdf5-ann = { path = "../clawhdf5-ann", version = "2.6.0", optional = true }
|
||||
clawhdf5-gpu = { path = "../clawhdf5-gpu", version = "2.6.0", optional = true, default-features = false }
|
||||
serde = { workspace = true }
|
||||
byteorder = "1"
|
||||
half = { workspace = true, optional = true }
|
||||
@@ -45,9 +45,11 @@ name = "memory_bench"
|
||||
harness = false
|
||||
|
||||
[features]
|
||||
default = ["float16", "hnsw"]
|
||||
default = ["float16", "hnsw", "parallel"]
|
||||
float16 = ["half"]
|
||||
parallel = ["rayon"]
|
||||
# Rayon-parallel brute-force search strategies, and a parallel bulk build of
|
||||
# the HNSW index (same graph, several times faster on a multi-core machine).
|
||||
parallel = ["rayon", "clawhdf5-ann?/parallel"]
|
||||
# Compress embeddings with Zstd instead of deflate when
|
||||
# `MemoryConfig::compression` is on. Off by default: it links libzstd (C).
|
||||
zstd = ["clawhdf5/zstd"]
|
||||
|
||||
@@ -118,6 +118,7 @@ mod tests {
|
||||
created_at: "2025-01-01T00:00:00Z".to_string(),
|
||||
wal_enabled: false,
|
||||
wal_max_entries: 500,
|
||||
quantized_index: false,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -37,7 +37,7 @@
|
||||
//! let mem = AsyncHDF5Memory::open_with(path, config).await?;
|
||||
//! mem.save(entry).await?; // buffered → background writer
|
||||
//! mem.save_batch(entries).await?; // also buffered
|
||||
//! let results = mem.hybrid_search(emb, "query".into(), 0.7, 0.3, 5).await;
|
||||
//! let results = mem.hybrid_search(emb, "query".into(), 0.4, 0.6, 5).await;
|
||||
//! mem.shutdown().await?; // final flush + stop
|
||||
//! ```
|
||||
|
||||
|
||||
+388
-114
@@ -3,10 +3,15 @@
|
||||
//! Provides a standard BM25 (Okapi BM25) implementation with an in-memory
|
||||
//! inverted index. Tombstoned documents are excluded from indexing and search.
|
||||
//!
|
||||
//! Optimizations:
|
||||
//! - Cached IDF scores (don't recompute per query)
|
||||
//! - Sorted posting lists by doc_id for cache-friendly access
|
||||
//! - Block-Max WAND early termination
|
||||
//! The index is **incremental**: [`BM25Index::add_document`] and
|
||||
//! [`BM25Index::remove_document`] keep it exactly equivalent to one built from
|
||||
//! scratch over the same live documents, so a store can maintain one index for
|
||||
//! its lifetime instead of re-tokenising the whole corpus per query. To make
|
||||
//! that possible IDF is computed at query time (it depends on the live
|
||||
//! document count) rather than cached at build time.
|
||||
//!
|
||||
//! - Posting lists sorted by doc id
|
||||
//! - Bounded-heap top-k; results ordered by score, then doc id (deterministic)
|
||||
|
||||
use std::cmp::Reverse;
|
||||
use std::collections::{BinaryHeap, HashMap};
|
||||
@@ -41,10 +46,11 @@ const DEFAULT_B: f32 = 0.75;
|
||||
pub struct BM25Index {
|
||||
/// Inverted index: token -> sorted list of (doc_id, term_frequency).
|
||||
inverted: HashMap<String, Vec<(usize, u32)>>,
|
||||
/// Cached IDF scores per token.
|
||||
idf_cache: HashMap<String, f32>,
|
||||
/// Number of tokens in each document (0 for tombstoned docs).
|
||||
doc_lengths: Vec<u32>,
|
||||
/// Sum of `doc_lengths` over live documents (keeps `avg_dl` exact under
|
||||
/// incremental updates).
|
||||
total_length: u64,
|
||||
/// Average document length across non-tombstoned docs.
|
||||
avg_dl: f32,
|
||||
/// Number of non-tombstoned documents.
|
||||
@@ -53,19 +59,27 @@ pub struct BM25Index {
|
||||
k1: f32,
|
||||
/// BM25 b parameter.
|
||||
b: f32,
|
||||
/// Applied to every document and query token, so the two always agree.
|
||||
filter: TokenFilter,
|
||||
}
|
||||
|
||||
impl BM25Index {
|
||||
/// Build a BM25 index from a set of documents, excluding tombstoned entries.
|
||||
pub fn build(documents: &[String], tombstones: &[u8]) -> Self {
|
||||
Self::build_with(documents, tombstones, TokenFilter::default())
|
||||
}
|
||||
|
||||
/// [`BM25Index::build`] with the token filter chosen explicitly.
|
||||
pub fn build_with(documents: &[String], tombstones: &[u8], filter: TokenFilter) -> Self {
|
||||
let mut index = Self {
|
||||
inverted: HashMap::new(),
|
||||
idf_cache: HashMap::new(),
|
||||
doc_lengths: vec![0; documents.len()],
|
||||
total_length: 0,
|
||||
avg_dl: 0.0,
|
||||
num_docs: 0,
|
||||
k1: DEFAULT_K1,
|
||||
b: DEFAULT_B,
|
||||
filter,
|
||||
};
|
||||
index.index_documents(documents, tombstones);
|
||||
index
|
||||
@@ -77,107 +91,165 @@ impl BM25Index {
|
||||
/// Uses Block-Max WAND for early termination when remaining documents
|
||||
/// cannot beat the current top-k threshold.
|
||||
pub fn search(&self, query: &str, k: usize) -> Vec<(usize, f32)> {
|
||||
if self.num_docs == 0 || k == 0 {
|
||||
if k == 0 {
|
||||
return Vec::new();
|
||||
}
|
||||
|
||||
let tokens = tokenize(query);
|
||||
if tokens.is_empty() {
|
||||
return Vec::new();
|
||||
}
|
||||
|
||||
// Collect posting lists and cached IDF scores for query tokens
|
||||
type QueryTerm<'a> = (&'a str, f32, &'a [(usize, u32)]);
|
||||
let mut query_terms: Vec<QueryTerm<'_>> = Vec::new();
|
||||
for token in &tokens {
|
||||
if let (Some(postings), Some(&idf)) = (
|
||||
self.inverted.get(token.as_str()),
|
||||
self.idf_cache.get(token.as_str()),
|
||||
) {
|
||||
query_terms.push((token, idf, postings));
|
||||
// Top-k with a bounded min-heap: O(matches * log k) instead of sorting
|
||||
// every match. Ties break towards the lower doc id so results are
|
||||
// deterministic.
|
||||
let mut heap: BinaryHeap<Reverse<(HeapScore, Reverse<usize>)>> =
|
||||
BinaryHeap::with_capacity(k.min(1024) + 1);
|
||||
for (doc_id, score) in self.scores(query) {
|
||||
heap.push(Reverse((HeapScore(score), Reverse(doc_id))));
|
||||
if heap.len() > k {
|
||||
heap.pop();
|
||||
}
|
||||
}
|
||||
let mut results: Vec<(usize, f32)> = heap
|
||||
.into_iter()
|
||||
.map(|Reverse((HeapScore(score), Reverse(doc_id)))| (doc_id, score))
|
||||
.collect();
|
||||
results.sort_by(|a, b| b.1.total_cmp(&a.1).then(a.0.cmp(&b.0)));
|
||||
results
|
||||
}
|
||||
|
||||
if query_terms.is_empty() {
|
||||
/// The BM25 score of **every** matching document, in doc-id order, unsorted
|
||||
/// by score. Score fusion normalises over the whole matching set, so it
|
||||
/// needs all of these but not their ranking; producing a ranked list of
|
||||
/// every match (`search(query, corpus_len)`) spent most of its time sorting.
|
||||
pub fn scores(&self, query: &str) -> Vec<(usize, f32)> {
|
||||
if self.num_docs == 0 {
|
||||
return Vec::new();
|
||||
}
|
||||
|
||||
// Accumulate BM25 scores per document using WAND-style scoring
|
||||
let mut scores: HashMap<usize, f32> = HashMap::new();
|
||||
|
||||
// Compute maximum possible contribution per term for WAND
|
||||
let max_tf_score: Vec<f32> = query_terms
|
||||
.iter()
|
||||
.map(|(_, idf, _)| {
|
||||
// Upper bound: max TF contribution when tf is high and dl is short
|
||||
let max_tf_num = 10.0 * (self.k1 + 1.0);
|
||||
let max_tf_den = 10.0 + self.k1 * (1.0 - self.b);
|
||||
idf * max_tf_num / max_tf_den
|
||||
})
|
||||
.collect();
|
||||
|
||||
let total_max_contribution: f32 = max_tf_score.iter().sum();
|
||||
|
||||
// Threshold for WAND early termination. `top_k_heap` is a min-heap of
|
||||
// size k (worst-of-the-top-k at the head) so it can be maintained in
|
||||
// O(log k) per update instead of re-sorting the whole buffer.
|
||||
let mut threshold = 0.0f32;
|
||||
let mut top_k_heap: BinaryHeap<Reverse<HeapScore>> = BinaryHeap::with_capacity(k);
|
||||
|
||||
for (term_idx, (_, idf, postings)) in query_terms.iter().enumerate() {
|
||||
for &(doc_id, freq) in *postings {
|
||||
// Term-at-a-time accumulation into a dense array: a common term has a
|
||||
// posting per document, and hashing each one dominated query time.
|
||||
// IDF is computed here rather than cached at build time: it depends on
|
||||
// the live document count, which changes with every incremental
|
||||
// add/remove, and costs one `ln` per query term.
|
||||
let mut acc = vec![0.0f32; self.doc_lengths.len()];
|
||||
let mut matched = false;
|
||||
for token in tokenize_with(query, self.filter) {
|
||||
let Some(postings) = self.inverted.get(token.as_str()) else {
|
||||
continue;
|
||||
};
|
||||
matched = true;
|
||||
let df = postings.len() as f32;
|
||||
let idf = ((self.num_docs as f32 - df + 0.5) / (df + 0.5) + 1.0).ln();
|
||||
for &(doc_id, freq) in postings {
|
||||
let dl = self.doc_lengths[doc_id] as f32;
|
||||
let freq_f = freq as f32;
|
||||
let tf = (freq_f * (self.k1 + 1.0))
|
||||
/ (freq_f + self.k1 * (1.0 - self.b + self.b * dl / self.avg_dl));
|
||||
let contribution = idf * tf;
|
||||
|
||||
let entry = scores.entry(doc_id).or_insert(0.0);
|
||||
*entry += contribution;
|
||||
|
||||
// WAND check: if this doc's current partial score + remaining
|
||||
// max terms can't beat threshold, we can skip (but we still
|
||||
// accumulate since we process term-at-a-time)
|
||||
if term_idx == query_terms.len() - 1 {
|
||||
// Last term: check if this doc beats threshold
|
||||
let final_score = *entry;
|
||||
if top_k_heap.len() >= k {
|
||||
if final_score > threshold {
|
||||
// Replace the current worst-of-top-k.
|
||||
top_k_heap.pop();
|
||||
top_k_heap.push(Reverse(HeapScore(final_score)));
|
||||
threshold = top_k_heap.peek().map(|Reverse(s)| s.0).unwrap_or(0.0);
|
||||
}
|
||||
} else {
|
||||
top_k_heap.push(Reverse(HeapScore(final_score)));
|
||||
if top_k_heap.len() == k {
|
||||
threshold = top_k_heap.peek().map(|Reverse(s)| s.0).unwrap_or(0.0);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
// After processing each term, check if remaining terms can
|
||||
// possibly produce results above threshold
|
||||
let remaining_max: f32 = max_tf_score[term_idx + 1..].iter().sum();
|
||||
if remaining_max < threshold && total_max_contribution > 0.0 {
|
||||
// Early termination: remaining terms can't produce new top-k
|
||||
// entries on their own. But existing partial scores may still
|
||||
// be updated, so we continue (WAND is approximate here).
|
||||
let _ = remaining_max; // hint to compiler
|
||||
acc[doc_id] += idf * tf;
|
||||
}
|
||||
}
|
||||
if !matched {
|
||||
return Vec::new();
|
||||
}
|
||||
// Every contribution is strictly positive (idf = ln(1 + x), x > 0), so
|
||||
// a zero entry is a document no query term touched.
|
||||
acc.into_iter()
|
||||
.enumerate()
|
||||
.filter(|&(_, score)| score > 0.0)
|
||||
.collect()
|
||||
}
|
||||
|
||||
let mut results: Vec<(usize, f32)> = scores.into_iter().collect();
|
||||
results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
|
||||
results.truncate(k);
|
||||
results
|
||||
/// The token filter this index was built with.
|
||||
pub fn token_filter(&self) -> TokenFilter {
|
||||
self.filter
|
||||
}
|
||||
|
||||
/// Number of document slots (live or not) the index covers. Ids are
|
||||
/// positions in the document list it mirrors.
|
||||
pub fn len(&self) -> usize {
|
||||
self.doc_lengths.len()
|
||||
}
|
||||
|
||||
/// `true` when the index covers no document slots.
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.doc_lengths.is_empty()
|
||||
}
|
||||
|
||||
/// Index `text` as document `doc_id`, which must be the next free id
|
||||
/// (`self.len()`) or an existing slot that is currently empty (removed or
|
||||
/// tombstoned). After any sequence of `add_document` / `remove_document`
|
||||
/// calls the index scores exactly as one freshly built from the same live
|
||||
/// documents.
|
||||
pub fn add_document(&mut self, doc_id: usize, text: &str) {
|
||||
if doc_id >= self.doc_lengths.len() {
|
||||
self.doc_lengths.resize(doc_id + 1, 0);
|
||||
}
|
||||
debug_assert_eq!(self.doc_lengths[doc_id], 0, "slot {doc_id} is occupied");
|
||||
|
||||
let tokens = tokenize_with(text, self.filter);
|
||||
let mut term_freqs: HashMap<&str, u32> = HashMap::new();
|
||||
for token in &tokens {
|
||||
*term_freqs.entry(token).or_insert(0) += 1;
|
||||
}
|
||||
for (token, freq) in term_freqs {
|
||||
let postings = self.inverted.entry(token.to_string()).or_default();
|
||||
// Posting lists stay sorted by doc id; appends are the common case.
|
||||
match postings.last() {
|
||||
Some(&(last, _)) if last >= doc_id => {
|
||||
let at = postings.partition_point(|&(id, _)| id < doc_id);
|
||||
postings.insert(at, (doc_id, freq));
|
||||
}
|
||||
_ => postings.push((doc_id, freq)),
|
||||
}
|
||||
}
|
||||
self.doc_lengths[doc_id] = tokens.len() as u32;
|
||||
self.total_length += tokens.len() as u64;
|
||||
self.num_docs += 1;
|
||||
self.refresh_avg_dl();
|
||||
}
|
||||
|
||||
/// Extend the index to cover `len` document slots, leaving new ones empty.
|
||||
/// Used for slots that hold no live document (tombstoned records).
|
||||
pub fn pad_to(&mut self, len: usize) {
|
||||
if len > self.doc_lengths.len() {
|
||||
self.doc_lengths.resize(len, 0);
|
||||
}
|
||||
}
|
||||
|
||||
/// Remove document `doc_id`, whose indexed text was `text`. The text is
|
||||
/// needed to find its postings; pass exactly what was added.
|
||||
pub fn remove_document(&mut self, doc_id: usize, text: &str) {
|
||||
let tokens = tokenize_with(text, self.filter);
|
||||
let mut seen: std::collections::HashSet<&str> = std::collections::HashSet::new();
|
||||
for token in &tokens {
|
||||
if !seen.insert(token) {
|
||||
continue;
|
||||
}
|
||||
if let Some(postings) = self.inverted.get_mut(token.as_str()) {
|
||||
if let Ok(at) = postings.binary_search_by_key(&doc_id, |&(id, _)| id) {
|
||||
postings.remove(at);
|
||||
}
|
||||
if postings.is_empty() {
|
||||
self.inverted.remove(token.as_str());
|
||||
}
|
||||
}
|
||||
}
|
||||
if let Some(len) = self.doc_lengths.get_mut(doc_id) {
|
||||
self.total_length = self.total_length.saturating_sub(u64::from(*len));
|
||||
*len = 0;
|
||||
}
|
||||
self.num_docs = self.num_docs.saturating_sub(1);
|
||||
self.refresh_avg_dl();
|
||||
}
|
||||
|
||||
fn refresh_avg_dl(&mut self) {
|
||||
self.avg_dl = if self.num_docs > 0 {
|
||||
self.total_length as f32 / self.num_docs as f32
|
||||
} else {
|
||||
0.0
|
||||
};
|
||||
}
|
||||
|
||||
/// Rebuild the index from scratch (e.g., after compaction).
|
||||
pub fn rebuild(&mut self, documents: &[String], tombstones: &[u8]) {
|
||||
self.inverted.clear();
|
||||
self.idf_cache.clear();
|
||||
self.doc_lengths = vec![0; documents.len()];
|
||||
self.total_length = 0;
|
||||
self.avg_dl = 0.0;
|
||||
self.num_docs = 0;
|
||||
self.index_documents(documents, tombstones);
|
||||
@@ -193,7 +265,7 @@ impl BM25Index {
|
||||
continue;
|
||||
}
|
||||
|
||||
let tokens = tokenize(doc);
|
||||
let tokens = tokenize_with(doc, self.filter);
|
||||
let doc_len = tokens.len() as u32;
|
||||
self.doc_lengths[i] = doc_len;
|
||||
total_length += doc_len as u64;
|
||||
@@ -214,33 +286,98 @@ impl BM25Index {
|
||||
}
|
||||
|
||||
self.num_docs = count;
|
||||
self.avg_dl = if count > 0 {
|
||||
total_length as f32 / count as f32
|
||||
} else {
|
||||
0.0
|
||||
};
|
||||
self.total_length = total_length;
|
||||
self.refresh_avg_dl();
|
||||
|
||||
// Sort posting lists by doc_id for cache-friendly access
|
||||
for postings in self.inverted.values_mut() {
|
||||
postings.sort_by_key(|&(doc_id, _)| doc_id);
|
||||
}
|
||||
|
||||
// Pre-compute and cache IDF scores
|
||||
for (token, postings) in &self.inverted {
|
||||
let df = postings.len() as f32;
|
||||
let idf = ((self.num_docs as f32 - df + 0.5) / (df + 0.5) + 1.0).ln();
|
||||
self.idf_cache.insert(token.clone(), idf);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Tokenize a string: lowercase, split on non-alphanumeric characters,
|
||||
/// filter empty tokens.
|
||||
/// What [`tokenize_with`] does to each token after splitting.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
|
||||
pub enum TokenFilter {
|
||||
/// Lowercase and split only — the original behaviour.
|
||||
#[default]
|
||||
Plain,
|
||||
/// Also strip common English inflections, so "running" and "runs" match
|
||||
/// "run". Conservative on purpose: only plural and past/continuous verb
|
||||
/// endings, and only on tokens long enough that stripping leaves a real
|
||||
/// stem. A stemmer earns its keep by conflating *related* words; an
|
||||
/// aggressive one also conflates unrelated ones ("universe"/"university"),
|
||||
/// which costs precision.
|
||||
Stemmed,
|
||||
}
|
||||
|
||||
/// Strip common English inflections from an already-lowercased token.
|
||||
///
|
||||
/// Applied identically to documents and queries, so the pair only has to agree
|
||||
/// with itself — the stem need not be a real word.
|
||||
fn stem(token: &str) -> &str {
|
||||
// Below this, stripping does more harm than good ("bed" -> "b").
|
||||
const MIN_STEM: usize = 4;
|
||||
let strip = |suffix: &str, min_len: usize| -> Option<&str> {
|
||||
let stem = token.strip_suffix(suffix)?;
|
||||
(stem.len() >= min_len).then_some(stem)
|
||||
};
|
||||
|
||||
// Plurals first: "studies" -> "studi", "classes" -> "class", "cats" -> "cat".
|
||||
// "ies" keeps its "i" so the result meets "-ied" ("studied" -> "studi").
|
||||
if let Some(stem) = strip("ies", 2) {
|
||||
return &token[..stem.len() + 1];
|
||||
}
|
||||
for suffix in ["sses", "shes", "ches", "xes", "zes"] {
|
||||
if let Some(stem) = strip(suffix, MIN_STEM - 1) {
|
||||
// Keep the sibilant: "classes" -> "class", not "clas".
|
||||
return &token[..stem.len() + 2];
|
||||
}
|
||||
}
|
||||
// Verb endings before the bare plural, so "raced" doesn't become "raced".
|
||||
if let Some(stem) = strip("ing", MIN_STEM - 1).or_else(|| strip("ed", MIN_STEM - 1)) {
|
||||
return undouble(stem);
|
||||
}
|
||||
if !token.ends_with("ss")
|
||||
&& !token.ends_with("us")
|
||||
&& !token.ends_with("is")
|
||||
&& let Some(stem) = strip("s", MIN_STEM - 1)
|
||||
{
|
||||
return stem;
|
||||
}
|
||||
token
|
||||
}
|
||||
|
||||
/// "runn" -> "run": undo the consonant doubling that "-ing"/"-ed" introduce.
|
||||
fn undouble(stem: &str) -> &str {
|
||||
let mut chars = stem.chars().rev();
|
||||
let (Some(last), Some(prev)) = (chars.next(), chars.next()) else {
|
||||
return stem;
|
||||
};
|
||||
let doubled = last == prev && !"aeiou".contains(last) && last.is_ascii_alphabetic();
|
||||
if doubled && stem.len() > 3 {
|
||||
&stem[..stem.len() - 1]
|
||||
} else {
|
||||
stem
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
fn tokenize(text: &str) -> Vec<String> {
|
||||
tokenize_with(text, TokenFilter::Plain)
|
||||
}
|
||||
|
||||
/// Split `text` into scoring tokens under `filter`.
|
||||
pub fn tokenize_with(text: &str, filter: TokenFilter) -> Vec<String> {
|
||||
text.to_lowercase()
|
||||
.split(|c: char| !c.is_alphanumeric())
|
||||
.filter(|s| !s.is_empty())
|
||||
.map(|s| s.to_string())
|
||||
.map(|token| match filter {
|
||||
TokenFilter::Plain => token.to_string(),
|
||||
TokenFilter::Stemmed => stem(token).to_string(),
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
@@ -386,24 +523,21 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cached_idf_consistent_with_computed() {
|
||||
fn score_matches_the_bm25_formula() {
|
||||
let docs = vec![
|
||||
"rust programming".to_string(),
|
||||
"rust systems".to_string(),
|
||||
"python scripting".to_string(),
|
||||
];
|
||||
let tombstones = vec![0, 0, 0];
|
||||
let index = BM25Index::build(&docs, &tombstones);
|
||||
let index = BM25Index::build(&docs, &[0, 0, 0]);
|
||||
|
||||
// IDF for "rust" (appears in 2 of 3 docs)
|
||||
let idf_rust = index.idf_cache.get("rust").unwrap();
|
||||
let expected_idf = ((3.0f32 - 2.0 + 0.5) / (2.0 + 0.5) + 1.0).ln();
|
||||
assert!(
|
||||
(idf_rust - expected_idf).abs() < 1e-6,
|
||||
"cached IDF mismatch: {} vs {}",
|
||||
idf_rust,
|
||||
expected_idf
|
||||
);
|
||||
// "python": df = 1 of N = 3. Every doc has the average length (2) and
|
||||
// tf = 1, so the tf factor is exactly 1 and the score is the IDF.
|
||||
let results = index.search("python", 3);
|
||||
let expected_idf = ((3.0f32 - 1.0 + 0.5) / (1.0 + 0.5) + 1.0).ln();
|
||||
assert_eq!(results.len(), 1);
|
||||
assert_eq!(results[0].0, 2);
|
||||
assert!((results[0].1 - expected_idf).abs() < 1e-6, "{results:?}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -467,4 +601,144 @@ mod tests {
|
||||
);
|
||||
}
|
||||
}
|
||||
/// Documents drawn from a small vocabulary so terms collide heavily.
|
||||
fn random_doc(state: &mut u64) -> String {
|
||||
const VOCAB: &[&str] = &[
|
||||
"alpha", "beta", "gamma", "delta", "eps", "zeta", "eta", "x1",
|
||||
];
|
||||
let mut next = || {
|
||||
*state = state
|
||||
.wrapping_mul(6364136223846793005)
|
||||
.wrapping_add(1442695040888963407);
|
||||
(*state >> 33) as usize
|
||||
};
|
||||
let len = 1 + next() % 9;
|
||||
(0..len)
|
||||
.map(|_| VOCAB[next() % VOCAB.len()])
|
||||
.collect::<Vec<_>>()
|
||||
.join(" ")
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn incremental_updates_match_a_fresh_build_exactly() {
|
||||
for seed in 0..60u64 {
|
||||
let mut state = seed.wrapping_mul(0x9E37_79B9_7F4A_7C15) | 1;
|
||||
let mut docs: Vec<String> = Vec::new();
|
||||
let mut tombstones: Vec<u8> = Vec::new();
|
||||
let mut index = BM25Index::build(&docs, &tombstones);
|
||||
|
||||
for step in 0..80 {
|
||||
state = state.wrapping_mul(6364136223846793005).wrapping_add(1);
|
||||
let live: Vec<usize> = (0..docs.len()).filter(|&i| tombstones[i] == 0).collect();
|
||||
match (state >> 40) % 4 {
|
||||
0 if !live.is_empty() => {
|
||||
// delete
|
||||
let id = live[(state >> 20) as usize % live.len()];
|
||||
index.remove_document(id, &docs[id]);
|
||||
tombstones[id] = 1;
|
||||
}
|
||||
1 if !live.is_empty() => {
|
||||
// update in place
|
||||
let id = live[(state >> 20) as usize % live.len()];
|
||||
let new_text = random_doc(&mut state);
|
||||
index.remove_document(id, &docs[id]);
|
||||
index.add_document(id, &new_text);
|
||||
docs[id] = new_text;
|
||||
}
|
||||
_ => {
|
||||
let text = random_doc(&mut state);
|
||||
index.add_document(docs.len(), &text);
|
||||
docs.push(text);
|
||||
tombstones.push(0);
|
||||
}
|
||||
}
|
||||
|
||||
let fresh = BM25Index::build(&docs, &tombstones);
|
||||
for query in ["alpha", "beta gamma", "x1 zeta alpha delta", "missing"] {
|
||||
let got = index.search(query, 5);
|
||||
let want = fresh.search(query, 5);
|
||||
assert_eq!(got.len(), want.len(), "seed {seed} step {step} {query:?}");
|
||||
for (g, w) in got.iter().zip(&want) {
|
||||
assert_eq!(
|
||||
g.0, w.0,
|
||||
"seed {seed} step {step} {query:?}: {got:?} vs {want:?}"
|
||||
);
|
||||
assert!(
|
||||
(g.1 - w.1).abs() < 1e-5,
|
||||
"seed {seed} step {step} {query:?}"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn scores_is_the_unranked_form_of_a_full_search() {
|
||||
let mut state = 99u64;
|
||||
let docs: Vec<String> = (0..200).map(|_| random_doc(&mut state)).collect();
|
||||
let tombstones: Vec<u8> = (0..200).map(|i| u8::from(i % 7 == 0)).collect();
|
||||
let index = BM25Index::build(&docs, &tombstones);
|
||||
for query in ["alpha", "beta gamma x1", "missing", ""] {
|
||||
let mut all = index.scores(query);
|
||||
all.sort_by(|a, b| b.1.total_cmp(&a.1).then(a.0.cmp(&b.0)));
|
||||
assert_eq!(all, index.search(query, docs.len()), "{query:?}");
|
||||
assert!(all.iter().all(|(id, _)| tombstones[*id] == 0));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stemming_conflates_inflections_of_the_same_word() {
|
||||
let stem_of = |w: &str| tokenize_with(w, TokenFilter::Stemmed).pop().unwrap();
|
||||
// Pairs that should meet.
|
||||
for (a, b) in [
|
||||
("running", "runs"),
|
||||
("trained", "training"),
|
||||
("miles", "mile"),
|
||||
("studies", "studied"),
|
||||
("mentioned", "mentioning"),
|
||||
("classes", "class"),
|
||||
("planned", "planning"),
|
||||
] {
|
||||
assert_eq!(stem_of(a), stem_of(b), "{a} / {b} should share a stem");
|
||||
}
|
||||
// Pairs that must stay apart. Note which pairs are deliberately absent:
|
||||
// "bed"/"bedding" and "gas"/"gassed" both collapse to one stem, which
|
||||
// is what Porter does too and is right — they are related words.
|
||||
for (a, b) in [
|
||||
("universe", "university"),
|
||||
("business", "busy"),
|
||||
("this", "thing"),
|
||||
] {
|
||||
assert_ne!(stem_of(a), stem_of(b), "{a} / {b} must not be conflated");
|
||||
}
|
||||
// Short words and non-inflections are left alone.
|
||||
for word in ["run", "bus", "is", "his", "data", "gas"] {
|
||||
assert_eq!(stem_of(word), word, "{word} should be untouched");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stemming_is_off_by_default_and_applied_consistently() {
|
||||
assert_eq!(tokenize("Running miles"), ["running", "miles"]);
|
||||
assert_eq!(
|
||||
tokenize_with("Running miles", TokenFilter::Stemmed),
|
||||
["run", "mile"]
|
||||
);
|
||||
|
||||
// A query inflected differently from the document still matches.
|
||||
let docs = vec!["I ran while training for the marathon".to_string()];
|
||||
let plain = BM25Index::build_with(&docs, &[0], TokenFilter::Plain);
|
||||
let stemmed = BM25Index::build_with(&docs, &[0], TokenFilter::Stemmed);
|
||||
assert!(plain.search("trains", 1).is_empty());
|
||||
assert_eq!(stemmed.search("trains", 1).len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ties_break_towards_the_lower_doc_id() {
|
||||
let docs: Vec<String> = (0..6).map(|_| "same text".to_string()).collect();
|
||||
let index = BM25Index::build(&docs, &[0; 6]);
|
||||
let ids: Vec<usize> = index.search("same", 3).into_iter().map(|r| r.0).collect();
|
||||
assert_eq!(ids, [0, 1, 2]);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2,16 +2,143 @@
|
||||
|
||||
use crate::vector_search;
|
||||
|
||||
/// 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>,
|
||||
@@ -28,8 +155,7 @@ 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(),
|
||||
@@ -41,15 +167,14 @@ impl MemoryCache {
|
||||
}
|
||||
}
|
||||
|
||||
/// 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);
|
||||
}
|
||||
/// 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).
|
||||
@@ -79,8 +204,7 @@ impl MemoryCache {
|
||||
let idx = self.chunks.len();
|
||||
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);
|
||||
@@ -118,20 +242,7 @@ impl MemoryCache {
|
||||
if idx < self.chunks.len() {
|
||||
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 +294,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 +307,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 +316,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 +334,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 +357,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 +392,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 +428,15 @@ 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]);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -28,13 +28,40 @@ use crate::vector_search;
|
||||
pub fn hybrid_search(
|
||||
query_embedding: &[f32],
|
||||
query_text: &str,
|
||||
vectors: &[Vec<f32>],
|
||||
_chunks: &[String],
|
||||
vectors: &(impl crate::vector_search::VectorSet + Sync + ?Sized),
|
||||
chunks: &[String],
|
||||
tombstones: &[u8],
|
||||
bm25_index: &BM25Index,
|
||||
vector_weight: f32,
|
||||
keyword_weight: f32,
|
||||
k: usize,
|
||||
) -> Vec<(usize, f32)> {
|
||||
hybrid_search_fused(
|
||||
query_embedding,
|
||||
query_text,
|
||||
vectors,
|
||||
chunks,
|
||||
tombstones,
|
||||
bm25_index,
|
||||
Fusion::Weighted {
|
||||
vector: vector_weight,
|
||||
keyword: keyword_weight,
|
||||
},
|
||||
k,
|
||||
)
|
||||
}
|
||||
|
||||
/// [`hybrid_search`] with the fusion method chosen explicitly.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn hybrid_search_fused(
|
||||
query_embedding: &[f32],
|
||||
query_text: &str,
|
||||
vectors: &(impl crate::vector_search::VectorSet + Sync + ?Sized),
|
||||
_chunks: &[String],
|
||||
tombstones: &[u8],
|
||||
bm25_index: &BM25Index,
|
||||
fusion: Fusion,
|
||||
k: usize,
|
||||
) -> Vec<(usize, f32)> {
|
||||
// Get raw scores from both systems. Request all results so normalization
|
||||
// covers the full distribution.
|
||||
@@ -42,12 +69,12 @@ pub fn hybrid_search(
|
||||
let 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)
|
||||
@@ -58,9 +85,9 @@ pub fn hybrid_search(
|
||||
vector_search::cosine_similarity_batch(query_embedding, vectors, tombstones)
|
||||
}
|
||||
};
|
||||
let kw_scores = bm25_index.search(query_text, vectors.len());
|
||||
let kw_scores = bm25_index.scores(query_text);
|
||||
|
||||
merge_vector_keyword(vec_scores, kw_scores, vector_weight, keyword_weight, k)
|
||||
fuse(vec_scores, kw_scores, fusion, k)
|
||||
}
|
||||
|
||||
/// Merge pre-computed vector-similarity and keyword scores into a single ranking.
|
||||
@@ -76,29 +103,112 @@ pub fn merge_vector_keyword(
|
||||
keyword_weight: f32,
|
||||
k: usize,
|
||||
) -> Vec<(usize, f32)> {
|
||||
// Normalize each set to [0, 1].
|
||||
let vec_normalized = normalize_scores(&vec_scores);
|
||||
let kw_normalized = normalize_scores(&kw_scores);
|
||||
fuse(
|
||||
vec_scores,
|
||||
kw_scores,
|
||||
Fusion::Weighted {
|
||||
vector: vector_weight,
|
||||
keyword: keyword_weight,
|
||||
},
|
||||
k,
|
||||
)
|
||||
}
|
||||
|
||||
// Merge scores with weights.
|
||||
let mut merged: HashMap<usize, f32> = HashMap::new();
|
||||
/// How the vector and keyword stages are combined into one ranking.
|
||||
#[derive(Debug, Clone, Copy, PartialEq)]
|
||||
pub enum Fusion {
|
||||
/// Min-max normalise each stage over its own candidates, then take a
|
||||
/// weighted sum. Uses the *scores*, so a stage that separates its
|
||||
/// candidates sharply keeps that separation — and a stage whose candidates
|
||||
/// are all near-identical contributes little.
|
||||
Weighted {
|
||||
/// Weight on the vector stage.
|
||||
vector: f32,
|
||||
/// Weight on the keyword stage.
|
||||
keyword: f32,
|
||||
},
|
||||
/// Reciprocal rank fusion: each stage contributes `1 / (k + rank)`,
|
||||
/// ignoring score magnitudes entirely. Robust when the two stages'
|
||||
/// scores aren't comparable, at the cost of discarding confidence.
|
||||
Rrf {
|
||||
/// The rank-damping constant; 60 is the value from the original paper.
|
||||
k: f32,
|
||||
},
|
||||
}
|
||||
|
||||
for (idx, score) in &vec_normalized {
|
||||
*merged.entry(*idx).or_insert(0.0) += vector_weight * score;
|
||||
impl Default for Fusion {
|
||||
fn default() -> Self {
|
||||
DEFAULT_FUSION
|
||||
}
|
||||
for (idx, score) in &kw_normalized {
|
||||
*merged.entry(*idx).or_insert(0.0) += keyword_weight * score;
|
||||
}
|
||||
|
||||
/// The fusion `hybrid_search` uses unless told otherwise.
|
||||
///
|
||||
/// The weights are not a guess: a sweep of every 0.1 step over the full
|
||||
/// LongMemEval haystack (500 questions, real MiniLM embeddings) found the
|
||||
/// long-standing 0.7/0.3 default *strictly dominated* — 0.4/0.6 is better at
|
||||
/// Hit@1, Hit@5, Hit@10 and MRR, at both turn and session granularity. See
|
||||
/// `BENCHMARKS.md`, "Weight sweep".
|
||||
pub const DEFAULT_FUSION: Fusion = Fusion::Weighted {
|
||||
vector: 0.4,
|
||||
keyword: 0.6,
|
||||
};
|
||||
|
||||
/// Combine one ranked candidate list from each stage into a single top-`k`.
|
||||
///
|
||||
/// Neither list need be sorted; both are consumed.
|
||||
pub fn fuse(
|
||||
vec_scores: Vec<(usize, f32)>,
|
||||
kw_scores: Vec<(usize, f32)>,
|
||||
fusion: Fusion,
|
||||
k: usize,
|
||||
) -> Vec<(usize, f32)> {
|
||||
let mut merged: HashMap<usize, f32> = HashMap::new();
|
||||
match fusion {
|
||||
Fusion::Weighted { vector, keyword } => {
|
||||
// Normalize each set to [0, 1].
|
||||
for (idx, score) in &normalize_scores(&vec_scores) {
|
||||
*merged.entry(*idx).or_insert(0.0) += vector * score;
|
||||
}
|
||||
for (idx, score) in &normalize_scores(&kw_scores) {
|
||||
*merged.entry(*idx).or_insert(0.0) += keyword * score;
|
||||
}
|
||||
}
|
||||
Fusion::Rrf { k: damping } => {
|
||||
for mut stage in [vec_scores, kw_scores] {
|
||||
// Rank 1 is the best score. Ties break by index so a stage's
|
||||
// contribution doesn't depend on the candidate order it
|
||||
// happened to be produced in.
|
||||
stage.sort_by(|a, b| {
|
||||
b.1.partial_cmp(&a.1)
|
||||
.unwrap_or(std::cmp::Ordering::Equal)
|
||||
.then(a.0.cmp(&b.0))
|
||||
});
|
||||
for (rank, (idx, _)) in stage.iter().enumerate() {
|
||||
*merged.entry(*idx).or_insert(0.0) += 1.0 / (damping + (rank + 1) as f32);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let mut results: Vec<(usize, f32)> = merged.into_iter().collect();
|
||||
// Index tie-break: `merged` is a HashMap, so without it the ties that
|
||||
// survive `truncate` differ from run to run.
|
||||
results.sort_by(|a, b| {
|
||||
// survive differ from run to run.
|
||||
let by_score_then_id = |a: &(usize, f32), b: &(usize, f32)| {
|
||||
b.1.partial_cmp(&a.1)
|
||||
.unwrap_or(std::cmp::Ordering::Equal)
|
||||
.then(a.0.cmp(&b.0))
|
||||
});
|
||||
results.truncate(k);
|
||||
};
|
||||
// Only the top k are wanted: partition them out, then order just those,
|
||||
// instead of sorting every candidate (the keyword side can be the corpus).
|
||||
if k == 0 {
|
||||
return Vec::new();
|
||||
}
|
||||
if results.len() > k {
|
||||
results.select_nth_unstable_by(k - 1, by_score_then_id);
|
||||
results.truncate(k);
|
||||
}
|
||||
results.sort_by(by_score_then_id);
|
||||
results
|
||||
}
|
||||
|
||||
@@ -160,7 +270,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,
|
||||
@@ -172,12 +282,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)
|
||||
@@ -188,7 +298,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));
|
||||
@@ -344,6 +454,68 @@ mod tests {
|
||||
assert_eq!(result[0].1, 1.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn default_fusion_is_the_tuned_operating_point() {
|
||||
// A sweep over the full LongMemEval haystack found 0.7/0.3 strictly
|
||||
// dominated by 0.4/0.6 (BENCHMARKS.md). This guards the finding
|
||||
// against being quietly undone.
|
||||
assert_eq!(
|
||||
DEFAULT_FUSION,
|
||||
Fusion::Weighted {
|
||||
vector: 0.4,
|
||||
keyword: 0.6
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rrf_rewards_agreement_between_the_stages_and_ignores_magnitudes() {
|
||||
// Doc 1 is second-best in both stages; doc 0 is best in one and absent
|
||||
// from the other. RRF prefers the doc both stages liked.
|
||||
let vec_scores = vec![(0, 100.0), (1, 0.9)];
|
||||
let kw_scores = vec![(2, 5.0), (1, 4.9)];
|
||||
let ranked = fuse(vec_scores, kw_scores, Fusion::Rrf { k: 60.0 }, 3);
|
||||
assert_eq!(ranked[0].0, 1, "{ranked:?}");
|
||||
|
||||
// Scaling one stage's scores cannot change an RRF ranking, only the
|
||||
// order within that stage can.
|
||||
let a = fuse(
|
||||
vec![(0, 1.0), (1, 0.5)],
|
||||
vec![(1, 2.0), (0, 1.0)],
|
||||
Fusion::Rrf { k: 60.0 },
|
||||
2,
|
||||
);
|
||||
let b = fuse(
|
||||
vec![(0, 1e6), (1, -3.0)],
|
||||
vec![(1, 0.002), (0, 0.001)],
|
||||
Fusion::Rrf { k: 60.0 },
|
||||
2,
|
||||
);
|
||||
assert_eq!(
|
||||
a.iter().map(|r| r.0).collect::<Vec<_>>(),
|
||||
b.iter().map(|r| r.0).collect::<Vec<_>>()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn merge_top_k_matches_a_full_sort() {
|
||||
// Many ties (scores repeat) so the index tie-break is exercised.
|
||||
let vec_scores: Vec<(usize, f32)> = (0..300).map(|i| (i, ((i * 7) % 13) as f32)).collect();
|
||||
let kw_scores: Vec<(usize, f32)> = (100..500).map(|i| (i, ((i * 5) % 11) as f32)).collect();
|
||||
let everything =
|
||||
merge_vector_keyword(vec_scores.clone(), kw_scores.clone(), 0.7, 0.3, 10_000);
|
||||
assert_eq!(everything.len(), 500);
|
||||
assert!(
|
||||
everything
|
||||
.windows(2)
|
||||
.all(|w| { w[0].1 > w[1].1 || (w[0].1 == w[1].1 && w[0].0 < w[1].0) })
|
||||
);
|
||||
for k in [0, 1, 7, 50, 499, 500, 501] {
|
||||
let top = merge_vector_keyword(vec_scores.clone(), kw_scores.clone(), 0.7, 0.3, k);
|
||||
assert_eq!(top, everything[..k.min(500)], "k = {k}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn normalize_scores_all_equal() {
|
||||
let matched = normalize_scores(&[(0, 0.4), (1, 0.4)]);
|
||||
|
||||
@@ -62,7 +62,7 @@ use std::path::{Path, PathBuf};
|
||||
|
||||
use cache::MemoryCache;
|
||||
#[cfg(feature = "hnsw")]
|
||||
use clawhdf5_ann::{DistanceMetric, HnswIndex};
|
||||
use clawhdf5_ann::{DistanceMetric, HnswIndex, Storage};
|
||||
use ephemeral::{EphemeralConfig, EphemeralStore};
|
||||
|
||||
/// HNSW construction parameters used for the agent's vector index. Cosine is the
|
||||
@@ -139,6 +139,17 @@ 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.
|
||||
///
|
||||
/// 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
|
||||
/// restores recall; what it costs is throughput — roughly 13% of queries
|
||||
/// per second and 16% of build time at 100K x 384. See `BENCHMARKS.md`.
|
||||
///
|
||||
/// Has no effect without the `hnsw` feature.
|
||||
pub quantized_index: bool,
|
||||
}
|
||||
|
||||
impl MemoryConfig {
|
||||
@@ -160,6 +171,7 @@ impl MemoryConfig {
|
||||
created_at,
|
||||
wal_enabled: true,
|
||||
wal_max_entries: 500,
|
||||
quantized_index: false,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -205,6 +217,12 @@ pub trait AgentMemory {
|
||||
fn get_session_summary(&self, session_id: &str) -> Result<Option<String>>;
|
||||
}
|
||||
|
||||
/// Ceiling for a record's Hebbian activation weight. Each hit adds
|
||||
/// `hebbian_boost` and the fused score is scaled by `sqrt(weight)`, so without
|
||||
/// a cap a frequently returned record's advantage grows without limit and it
|
||||
/// eventually outranks better matches purely on popularity.
|
||||
pub(crate) const MAX_ACTIVATION_WEIGHT: f32 = 16.0;
|
||||
|
||||
/// Most anomaly alerts kept between `take_anomaly_alerts` calls.
|
||||
const MAX_PENDING_ALERTS: usize = 1024;
|
||||
|
||||
@@ -247,6 +265,17 @@ pub struct HDF5Memory {
|
||||
/// via [`HDF5Memory::take_anomaly_alerts`]. Saves are never blocked on
|
||||
/// these — surfacing is opt-in for callers that want to act on them.
|
||||
anomaly_alerts: Vec<anomaly::AnomalyAlert>,
|
||||
/// Keyword index over `cache.chunks`, kept for the life of the store and
|
||||
/// updated incrementally — it used to be rebuilt from scratch, re-tokenising
|
||||
/// every record, on every single query. Built lazily on first use; see
|
||||
/// [`HDF5Memory::ensure_bm25_fresh`] for how it stays in sync.
|
||||
bm25: Option<bm25::BM25Index>,
|
||||
/// Token filter the keyword index is built with. Changing it drops the
|
||||
/// index; it is not persisted, because the index is not either.
|
||||
bm25_filter: bm25::TokenFilter,
|
||||
/// Activation weights changed since the last checkpoint (searches boost
|
||||
/// the records they return). Cleared by `flush`.
|
||||
activations_dirty: bool,
|
||||
/// Opened with [`HDF5Memory::open_read_only`]: nothing may reach the disk.
|
||||
read_only: bool,
|
||||
/// A WAL that `open()` could not read and moved aside; see
|
||||
@@ -299,6 +328,9 @@ impl HDF5Memory {
|
||||
provenance: provenance::ProvenanceStore::new(),
|
||||
anomaly: anomaly::WriteAnomalyDetector::new(anomaly::AnomalyConfig::default()),
|
||||
anomaly_alerts: Vec::new(),
|
||||
bm25: None,
|
||||
bm25_filter: bm25::TokenFilter::default(),
|
||||
activations_dirty: false,
|
||||
read_only: false,
|
||||
quarantined_wal: None,
|
||||
_lock: Some(lock),
|
||||
@@ -374,8 +406,13 @@ impl HDF5Memory {
|
||||
} else {
|
||||
Some(store_lock::StoreLock::acquire(path)?)
|
||||
};
|
||||
let ((config, mut cache, sessions, knowledge), wal_applied) =
|
||||
storage::read_from_disk_with_mark(path)?;
|
||||
let ((config, mut cache, sessions, knowledge), checkpoint) =
|
||||
storage::read_from_disk_with_meta(path)?;
|
||||
let wal_applied = checkpoint.wal_applied;
|
||||
let n_checkpoint = cache.len();
|
||||
// Set if WAL replay did anything other than append records; the saved
|
||||
// vector index then no longer describes the first `n_checkpoint` ones.
|
||||
let mut replay_only_appended = true;
|
||||
|
||||
// Replay WAL if present
|
||||
let wal_path = path.with_extension("h5.wal");
|
||||
@@ -392,6 +429,9 @@ impl HDF5Memory {
|
||||
&& let Ok(entries) =
|
||||
wal::WalFile::read_entries_for_migration(&wal_path, wal_applied)
|
||||
{
|
||||
replay_only_appended &= entries
|
||||
.iter()
|
||||
.all(|e| e.entry_type == wal::WalEntryType::Save);
|
||||
wal::replay_into_cache(&entries, &mut cache);
|
||||
}
|
||||
None
|
||||
@@ -403,6 +443,9 @@ impl HDF5Memory {
|
||||
// in case the process died between writing the .h5 and
|
||||
// truncating the WAL.
|
||||
let entries = wal::WalFile::read_entries_for_migration(&wal_path, wal_applied)?;
|
||||
replay_only_appended &= entries
|
||||
.iter()
|
||||
.all(|e| e.entry_type == wal::WalEntryType::Save);
|
||||
wal::replay_into_cache(&entries, &mut cache);
|
||||
Some(wal::WalFile::open(&wal_path)?)
|
||||
} else if config.wal_enabled {
|
||||
@@ -411,6 +454,35 @@ impl HDF5Memory {
|
||||
None
|
||||
};
|
||||
|
||||
#[cfg(feature = "hnsw")]
|
||||
let loaded_index = if replay_only_appended {
|
||||
Self::load_vector_index(
|
||||
path,
|
||||
checkpoint.ann_generation,
|
||||
&cache,
|
||||
n_checkpoint,
|
||||
if config.quantized_index {
|
||||
Storage::Int8
|
||||
} else {
|
||||
Storage::Float32
|
||||
},
|
||||
)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
#[cfg(not(feature = "hnsw"))]
|
||||
let _ = (
|
||||
n_checkpoint,
|
||||
replay_only_appended,
|
||||
checkpoint.ann_generation,
|
||||
);
|
||||
#[cfg(feature = "hnsw")]
|
||||
let synced_len = if loaded_index.is_some() {
|
||||
cache.len()
|
||||
} else {
|
||||
0
|
||||
};
|
||||
|
||||
Ok(Self {
|
||||
config,
|
||||
cache,
|
||||
@@ -419,14 +491,14 @@ impl HDF5Memory {
|
||||
wal,
|
||||
strategy: None,
|
||||
ephemeral: None,
|
||||
// Existing data is loaded from disk + WAL replay; mark the index
|
||||
// dirty so it is (re)built from the cache on the first search.
|
||||
// Reuse the vector index saved with the checkpoint when there is
|
||||
// one; otherwise mark it dirty so the first search builds it.
|
||||
#[cfg(feature = "hnsw")]
|
||||
hnsw: None,
|
||||
hnsw_dirty: loaded_index.is_none(),
|
||||
#[cfg(feature = "hnsw")]
|
||||
hnsw_dirty: true,
|
||||
hnsw_synced_len: synced_len,
|
||||
#[cfg(feature = "hnsw")]
|
||||
hnsw_synced_len: 0,
|
||||
hnsw: loaded_index,
|
||||
// No on-disk provenance ledger exists yet (see CLAUDE.md), so
|
||||
// there's no historical hash to verify loaded records against —
|
||||
// the store starts empty and is populated as records are
|
||||
@@ -434,12 +506,170 @@ impl HDF5Memory {
|
||||
provenance: provenance::ProvenanceStore::new(),
|
||||
anomaly: anomaly::WriteAnomalyDetector::new(anomaly::AnomalyConfig::default()),
|
||||
anomaly_alerts: Vec::new(),
|
||||
bm25: None,
|
||||
bm25_filter: bm25::TokenFilter::default(),
|
||||
activations_dirty: false,
|
||||
read_only,
|
||||
quarantined_wal,
|
||||
_lock: lock,
|
||||
})
|
||||
}
|
||||
|
||||
/// Where the vector index graph is kept between sessions.
|
||||
#[cfg_attr(not(feature = "hnsw"), allow(dead_code))]
|
||||
fn vector_index_path(store: &Path) -> PathBuf {
|
||||
store.with_extension("h5.ann")
|
||||
}
|
||||
|
||||
/// Save the vector index graph next to the store, returning the generation
|
||||
/// id the checkpoint must record for it. Only an index that exactly mirrors
|
||||
/// the cache is saved; otherwise any stale sidecar is removed and `None`
|
||||
/// returned, and the next session rebuilds. Failures are not errors — the
|
||||
/// sidecar is a cache of derived data.
|
||||
#[cfg(feature = "hnsw")]
|
||||
fn persist_vector_index(&self) -> Option<u64> {
|
||||
let path = Self::vector_index_path(&self.config.path);
|
||||
let index = match self.hnsw.as_ref() {
|
||||
Some(index)
|
||||
if !self.hnsw_dirty
|
||||
&& self.hnsw_synced_len == self.cache.embeddings.len()
|
||||
&& index.len() == self.cache.embeddings.len() =>
|
||||
{
|
||||
index
|
||||
}
|
||||
_ => {
|
||||
let _ = std::fs::remove_file(&path);
|
||||
return None;
|
||||
}
|
||||
};
|
||||
static COUNTER: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
|
||||
let nanos = std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.map_or(0, |d| d.as_nanos() as u64);
|
||||
let generation = nanos
|
||||
^ (u64::from(std::process::id()) << 32)
|
||||
^ COUNTER.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
|
||||
|
||||
let mut bytes = generation.to_le_bytes().to_vec();
|
||||
bytes.extend_from_slice(&index.graph_to_bytes());
|
||||
let tmp = path.with_extension("ann.tmp");
|
||||
let written =
|
||||
storage::write_synced(&tmp, &bytes).and_then(|()| storage::rename_synced(&tmp, &path));
|
||||
match written {
|
||||
Ok(()) => Some(generation),
|
||||
Err(_) => {
|
||||
let _ = std::fs::remove_file(&tmp);
|
||||
let _ = std::fs::remove_file(&path);
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(not(feature = "hnsw"))]
|
||||
fn persist_vector_index(&self) -> Option<u64> {
|
||||
None
|
||||
}
|
||||
|
||||
/// Load the vector index saved with the checkpoint identified by
|
||||
/// `generation`, covering the first `n_checkpoint` records of `cache`.
|
||||
/// Anything unexpected — no sidecar, another generation, a damaged or
|
||||
/// mismatched graph — yields `None` and the index is rebuilt on demand.
|
||||
#[cfg(feature = "hnsw")]
|
||||
fn load_vector_index(
|
||||
store: &Path,
|
||||
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()?;
|
||||
let (stamp, graph) = bytes.split_at_checked(8)?;
|
||||
if u64::from_le_bytes(stamp.try_into().ok()?) != generation {
|
||||
return None;
|
||||
}
|
||||
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].to_vec()) != id
|
||||
{
|
||||
return None;
|
||||
}
|
||||
}
|
||||
for (id, &t) in cache.tombstones.iter().enumerate() {
|
||||
if t != 0 {
|
||||
index.mark_deleted(id);
|
||||
}
|
||||
}
|
||||
Some(index)
|
||||
}
|
||||
|
||||
/// Bring the keyword index in line with the cache and return it.
|
||||
///
|
||||
/// Appends need no hook: records the index hasn't seen yet (whatever path
|
||||
/// added them) are indexed here, in order. Changes that keep the length the
|
||||
/// same are reported explicitly — [`Self::bm25_on_delete`] and
|
||||
/// [`Self::bm25_on_update`] — and anything that renumbers records
|
||||
/// (compaction) drops the index so it is rebuilt.
|
||||
pub(crate) fn ensure_bm25_fresh(&mut self) -> &bm25::BM25Index {
|
||||
let n = self.cache.chunks.len();
|
||||
let bm25 = match self.bm25.take() {
|
||||
Some(index) if index.len() <= n && index.token_filter() == self.bm25_filter => {
|
||||
let mut index = index;
|
||||
for id in index.len()..n {
|
||||
if self.cache.tombstones[id] == 0 {
|
||||
index.add_document(id, &self.cache.chunks[id]);
|
||||
}
|
||||
}
|
||||
index.pad_to(n);
|
||||
index
|
||||
}
|
||||
_ => bm25::BM25Index::build_with(
|
||||
&self.cache.chunks,
|
||||
&self.cache.tombstones,
|
||||
self.bm25_filter,
|
||||
),
|
||||
};
|
||||
self.bm25.insert(bm25)
|
||||
}
|
||||
|
||||
/// Choose how keyword-search tokens are normalised, rebuilding the index
|
||||
/// on next use. [`bm25::TokenFilter::Stemmed`] matches inflections of the
|
||||
/// same word at some cost in precision; measure before adopting it (see
|
||||
/// `BENCHMARKS.md`).
|
||||
pub fn set_token_filter(&mut self, filter: bm25::TokenFilter) {
|
||||
if filter != self.bm25_filter {
|
||||
self.bm25_filter = filter;
|
||||
self.bm25 = None;
|
||||
}
|
||||
}
|
||||
|
||||
/// Record `id` was tombstoned; its text is still in the cache.
|
||||
fn bm25_on_delete(&mut self, id: usize) {
|
||||
if let Some(index) = self.bm25.as_mut()
|
||||
&& id < index.len()
|
||||
{
|
||||
index.remove_document(id, &self.cache.chunks[id]);
|
||||
}
|
||||
}
|
||||
|
||||
/// Record `id`'s text changed from `old_text` to what the cache holds now.
|
||||
fn bm25_on_update(&mut self, id: usize, old_text: &str) {
|
||||
if let Some(index) = self.bm25.as_mut()
|
||||
&& id < index.len()
|
||||
{
|
||||
index.remove_document(id, old_text);
|
||||
index.add_document(id, &self.cache.chunks[id]);
|
||||
}
|
||||
}
|
||||
|
||||
/// Flush current state to disk and truncate the WAL.
|
||||
///
|
||||
/// Every code path that persists the full cache to the .h5 file must
|
||||
@@ -455,17 +685,24 @@ impl HDF5Memory {
|
||||
// Record which WAL prefix this checkpoint contains, so a crash before
|
||||
// the truncate below can't replay those entries a second time.
|
||||
let wal_applied = self.wal.as_ref().map(|w| w.mark());
|
||||
storage::write_to_disk_with_mark(
|
||||
// Written before the .h5 so a crash in between leaves a sidecar whose
|
||||
// generation matches no checkpoint (ignored), never the reverse.
|
||||
let ann_generation = self.persist_vector_index();
|
||||
storage::write_to_disk_with_meta(
|
||||
&self.config.path,
|
||||
&self.config,
|
||||
&self.cache,
|
||||
&self.sessions,
|
||||
&self.knowledge,
|
||||
wal_applied,
|
||||
&schema::CheckpointMeta {
|
||||
wal_applied,
|
||||
ann_generation,
|
||||
},
|
||||
)?;
|
||||
if let Some(ref mut w) = self.wal {
|
||||
w.truncate()?;
|
||||
}
|
||||
self.activations_dirty = false;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -589,6 +826,16 @@ impl HDF5Memory {
|
||||
// the index length drifts from the cache length (covering any mutation path
|
||||
// that doesn't call a hook, e.g. consolidation pushes).
|
||||
|
||||
/// 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.
|
||||
///
|
||||
@@ -604,11 +851,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,
|
||||
// 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,
|
||||
HNSW_M,
|
||||
HNSW_EF_CONSTRUCTION,
|
||||
DistanceMetric::Cosine,
|
||||
self.index_storage(),
|
||||
);
|
||||
for (i, &t) in self.cache.tombstones.iter().enumerate() {
|
||||
if t != 0 {
|
||||
@@ -623,6 +874,30 @@ impl HDF5Memory {
|
||||
#[cfg(feature = "hnsw")]
|
||||
fn ensure_hnsw_fresh(&mut self) {
|
||||
let n = self.cache.embeddings.len();
|
||||
// Records appended since the index was last in sync (a batch save, or
|
||||
// any path that pushes to the cache without a hook) are inserted
|
||||
// incrementally rather than triggering a rebuild of the whole graph.
|
||||
if !self.hnsw_dirty
|
||||
&& self.hnsw_synced_len < n
|
||||
&& let Some(index) = self.hnsw.as_mut()
|
||||
&& index.len() == self.hnsw_synced_len
|
||||
{
|
||||
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].to_vec()) == id
|
||||
});
|
||||
if appended {
|
||||
for id in self.hnsw_synced_len..n {
|
||||
if self.cache.tombstones[id] != 0 {
|
||||
index.mark_deleted(id);
|
||||
}
|
||||
}
|
||||
self.hnsw_synced_len = n;
|
||||
} else {
|
||||
self.hnsw_dirty = true;
|
||||
}
|
||||
}
|
||||
if self.hnsw_dirty || self.hnsw_synced_len != n {
|
||||
self.hnsw = self.build_hnsw_from_cache();
|
||||
self.hnsw_synced_len = n;
|
||||
@@ -641,7 +916,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 {
|
||||
@@ -777,6 +1052,7 @@ impl HDF5Memory {
|
||||
&entry.session_id,
|
||||
entry.timestamp,
|
||||
);
|
||||
let old_text = std::mem::take(&mut self.cache.chunks[existing_idx]);
|
||||
self.cache.update(
|
||||
existing_idx,
|
||||
entry.chunk,
|
||||
@@ -785,6 +1061,7 @@ impl HDF5Memory {
|
||||
entry.timestamp,
|
||||
entry.session_id,
|
||||
);
|
||||
self.bm25_on_update(existing_idx, &old_text);
|
||||
// In-place embedding change: the index node is stale, force rebuild.
|
||||
self.hnsw_mark_dirty();
|
||||
let needs_flush = self
|
||||
@@ -863,8 +1140,8 @@ impl AgentMemory for HDF5Memory {
|
||||
);
|
||||
indices.push(idx);
|
||||
}
|
||||
// Batch inserts rebuild the index once rather than node-by-node.
|
||||
self.hnsw_mark_dirty();
|
||||
// The vector and keyword indexes pick the new records up
|
||||
// incrementally the next time they are needed.
|
||||
self.flush()?;
|
||||
Ok(indices)
|
||||
}
|
||||
@@ -876,6 +1153,7 @@ impl AgentMemory for HDF5Memory {
|
||||
)));
|
||||
}
|
||||
self.hnsw_on_delete(id);
|
||||
self.bm25_on_delete(id);
|
||||
self.flush()?;
|
||||
|
||||
// Auto-compact if threshold exceeded
|
||||
@@ -893,6 +1171,7 @@ impl AgentMemory for HDF5Memory {
|
||||
if removed > 0 {
|
||||
// Record ids are cache indices, which compaction just renumbered.
|
||||
self.provenance.remap(&index_map);
|
||||
self.bm25 = None;
|
||||
// Compaction renumbers cache indices; rebuild the index to match.
|
||||
self.hnsw_mark_dirty();
|
||||
self.flush()?;
|
||||
@@ -920,6 +1199,13 @@ impl AgentMemory for HDF5Memory {
|
||||
if self.wal.as_ref().is_some_and(|w| !w.is_empty()) && wal_path.exists() {
|
||||
storage::snapshot_file(&wal_path, &snapshot.with_extension("h5.wal"))?;
|
||||
}
|
||||
// The saved vector index belongs to the checkpoint just copied (its
|
||||
// generation id is in that .h5), so it is valid for the snapshot too.
|
||||
// Best effort: without it the snapshot simply rebuilds on first search.
|
||||
let ann_path = Self::vector_index_path(&self.config.path);
|
||||
if ann_path.exists() {
|
||||
let _ = storage::snapshot_file(&ann_path, &Self::vector_index_path(&snapshot));
|
||||
}
|
||||
Ok(snapshot)
|
||||
}
|
||||
|
||||
@@ -1123,7 +1409,8 @@ impl HDF5Memory {
|
||||
k: usize,
|
||||
) -> Vec<SearchResult> {
|
||||
// Persistent tier.
|
||||
let persistent = self.hybrid_search(query_embedding, query_text, 0.7, 0.3, k);
|
||||
let persistent =
|
||||
self.hybrid_search_with(query_embedding, query_text, hybrid::DEFAULT_FUSION, k);
|
||||
const EPHEMERAL_BOOST: f32 = 1.2;
|
||||
let mut results = persistent;
|
||||
|
||||
@@ -1172,6 +1459,18 @@ impl HDF5Memory {
|
||||
|
||||
// --- Tests ---
|
||||
|
||||
impl Drop for HDF5Memory {
|
||||
/// Best-effort checkpoint of activation weights that only searches have
|
||||
/// touched. Everything else is already durable through the WAL or an
|
||||
/// earlier checkpoint; without this a search-only session would forget
|
||||
/// every boost it made.
|
||||
fn drop(&mut self) {
|
||||
if self.activations_dirty && !self.read_only {
|
||||
let _ = self.flush();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
@@ -1595,6 +1894,231 @@ mod tests {
|
||||
assert_eq!(restored.cache.chunks, ["checkpointed", "wal-only"]);
|
||||
}
|
||||
|
||||
/// A store with `n` records spread over a few directions, WAL on.
|
||||
#[cfg(feature = "hnsw")]
|
||||
fn indexed_store(dir: &TempDir, n: usize) -> (HDF5Memory, PathBuf) {
|
||||
let mut config = make_config(dir);
|
||||
config.wal_enabled = true;
|
||||
config.wal_max_entries = 10_000;
|
||||
config.compact_threshold = 0.0;
|
||||
let path = config.path.clone();
|
||||
let mut mem = HDF5Memory::create(config).unwrap();
|
||||
for i in 0..n {
|
||||
let a = i as f32 * 0.37;
|
||||
mem.save(make_entry(
|
||||
&format!("rec{i}"),
|
||||
&[a.cos(), a.sin(), (a * 0.5).cos(), 0.1],
|
||||
))
|
||||
.unwrap();
|
||||
}
|
||||
(mem, path)
|
||||
}
|
||||
|
||||
#[cfg(feature = "hnsw")]
|
||||
fn top_ids(mem: &mut HDF5Memory, q: &[f32]) -> Vec<usize> {
|
||||
mem.hybrid_search(q, "", 1.0, 0.0, 5)
|
||||
.into_iter()
|
||||
.map(|r| r.index)
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[cfg(feature = "hnsw")]
|
||||
#[test]
|
||||
fn vector_index_is_reloaded_not_rebuilt() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let (mut mem, path) = indexed_store(&dir, 60);
|
||||
let q = [0.3f32.cos(), 0.3f32.sin(), 0.9, 0.1];
|
||||
let expected = top_ids(&mut mem, &q); // builds the index
|
||||
mem.flush_wal().unwrap(); // checkpoint + sidecar
|
||||
drop(mem);
|
||||
assert!(HDF5Memory::vector_index_path(&path).exists());
|
||||
|
||||
let mut reopened = HDF5Memory::open(&path).unwrap();
|
||||
assert!(!reopened.hnsw_dirty, "index should come from the sidecar");
|
||||
assert_eq!(reopened.hnsw.as_ref().unwrap().len(), 60);
|
||||
assert_eq!(top_ids(&mut reopened, &q), expected);
|
||||
}
|
||||
|
||||
#[cfg(feature = "hnsw")]
|
||||
#[test]
|
||||
fn records_appended_after_the_checkpoint_join_the_loaded_index() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let (mut mem, path) = indexed_store(&dir, 40);
|
||||
top_ids(&mut mem, &[1.0, 0.0, 0.0, 0.0]);
|
||||
mem.flush_wal().unwrap();
|
||||
// Only in the WAL when the process "dies".
|
||||
mem.save(make_entry("late", &[0.0, 0.0, 0.0, 1.0])).unwrap();
|
||||
drop(mem);
|
||||
|
||||
let mut reopened = HDF5Memory::open(&path).unwrap();
|
||||
assert!(!reopened.hnsw_dirty);
|
||||
assert_eq!(reopened.hnsw.as_ref().unwrap().len(), 41);
|
||||
assert_eq!(top_ids(&mut reopened, &[0.0, 0.0, 0.0, 1.0])[0], 40);
|
||||
}
|
||||
|
||||
#[cfg(feature = "hnsw")]
|
||||
#[test]
|
||||
fn replayed_update_invalidates_the_saved_index() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let (mut mem, path) = indexed_store(&dir, 40);
|
||||
top_ids(&mut mem, &[1.0, 0.0, 0.0, 0.0]);
|
||||
mem.flush_wal().unwrap();
|
||||
// An in-place update after the checkpoint changes record 0's vector;
|
||||
// the saved graph was built over the old one.
|
||||
let mut moved = make_entry("rec0 moved", &[0.0, 0.0, 0.0, 1.0]);
|
||||
moved.tags = mem.cache.tags[0].clone();
|
||||
mem.save_or_update(moved).unwrap();
|
||||
let expected = top_ids(&mut mem, &[0.0, 0.0, 0.0, 1.0]);
|
||||
std::mem::forget(mem); // die without the drop-time checkpoint
|
||||
|
||||
let mut reopened = HDF5Memory::open_read_only(&path).unwrap();
|
||||
assert!(reopened.hnsw_dirty, "saved index must not be reused");
|
||||
assert_eq!(top_ids(&mut reopened, &[0.0, 0.0, 0.0, 1.0]), expected);
|
||||
}
|
||||
|
||||
#[cfg(feature = "hnsw")]
|
||||
#[test]
|
||||
fn stale_or_damaged_index_sidecar_is_ignored() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let (mut mem, path) = indexed_store(&dir, 40);
|
||||
let q = [1.0, 0.0, 0.0, 0.0];
|
||||
top_ids(&mut mem, &q); // builds the index
|
||||
mem.flush_wal().unwrap();
|
||||
let ann = HDF5Memory::vector_index_path(&path);
|
||||
let first_sidecar = std::fs::read(&ann).unwrap();
|
||||
// A second checkpoint gets a new generation.
|
||||
mem.save(make_entry("more", &[0.5, 0.5, 0.0, 0.0])).unwrap();
|
||||
top_ids(&mut mem, &q);
|
||||
mem.flush_wal().unwrap();
|
||||
let expected_after = top_ids(&mut mem, &q);
|
||||
drop(mem);
|
||||
|
||||
// Sidecar from the earlier checkpoint: wrong generation.
|
||||
std::fs::write(&ann, &first_sidecar).unwrap();
|
||||
let mut reopened = HDF5Memory::open_read_only(&path).unwrap();
|
||||
assert!(reopened.hnsw_dirty);
|
||||
assert_eq!(top_ids(&mut reopened, &q), expected_after);
|
||||
drop(reopened);
|
||||
|
||||
// Right generation, damaged graph.
|
||||
let mut mem = HDF5Memory::open(&path).unwrap();
|
||||
top_ids(&mut mem, &q);
|
||||
mem.flush_wal().unwrap();
|
||||
drop(mem);
|
||||
let mut bytes = std::fs::read(&ann).unwrap();
|
||||
let mid = bytes.len() / 2;
|
||||
bytes[mid] ^= 0xFF;
|
||||
std::fs::write(&ann, &bytes).unwrap();
|
||||
let mut reopened = HDF5Memory::open_read_only(&path).unwrap();
|
||||
assert!(reopened.hnsw_dirty);
|
||||
assert_eq!(top_ids(&mut reopened, &q), expected_after);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn set_token_filter_rebuilds_the_keyword_index() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let mut mem = HDF5Memory::create(make_config(&dir)).unwrap();
|
||||
mem.save(make_entry(
|
||||
"I was training for a marathon",
|
||||
&[1.0, 0.0, 0.0, 0.0],
|
||||
))
|
||||
.unwrap();
|
||||
|
||||
// Count only genuine keyword matches: `hybrid_search` also returns
|
||||
// zero-score filler when fewer than k records are relevant.
|
||||
let hits = |mem: &mut HDF5Memory| {
|
||||
mem.hybrid_search(&[0.0, 0.0, 0.0, 0.0], "trains", 0.0, 1.0, 5)
|
||||
.iter()
|
||||
.filter(|r| r.score > 0.0)
|
||||
.count()
|
||||
};
|
||||
assert_eq!(hits(&mut mem), 0);
|
||||
|
||||
mem.set_token_filter(bm25::TokenFilter::Stemmed);
|
||||
assert_eq!(hits(&mut mem), 1, "index should have been rebuilt stemmed");
|
||||
|
||||
// And back, rebuilding again.
|
||||
mem.set_token_filter(bm25::TokenFilter::Plain);
|
||||
assert_eq!(hits(&mut mem), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn keyword_index_stays_in_sync_through_every_mutation() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let mut config = make_config(&dir);
|
||||
config.compact_threshold = 0.0; // compact only when asked
|
||||
let mut mem = HDF5Memory::create(config).unwrap();
|
||||
let check = |mem: &mut HDF5Memory, what: &str| {
|
||||
let fresh = bm25::BM25Index::build(&mem.cache.chunks, &mem.cache.tombstones);
|
||||
let n = mem.cache.len();
|
||||
for query in ["apple", "banana cherry", "date", "nothing"] {
|
||||
let kept = mem.ensure_bm25_fresh().search(query, n);
|
||||
assert_eq!(kept, fresh.search(query, n), "{what}: {query:?}");
|
||||
}
|
||||
};
|
||||
let tagged = |chunk: &str, tag: &str| {
|
||||
let mut e = make_entry(chunk, &[1.0, 0.0, 0.0, 0.0]);
|
||||
e.tags = tag.into();
|
||||
e
|
||||
};
|
||||
|
||||
check(&mut mem, "empty");
|
||||
mem.save(tagged("apple banana", "a")).unwrap();
|
||||
mem.save(tagged("banana cherry cherry", "b")).unwrap();
|
||||
check(&mut mem, "after saves");
|
||||
mem.save_batch(vec![tagged("date apple", "c"), tagged("cherry", "d")])
|
||||
.unwrap();
|
||||
check(&mut mem, "after save_batch");
|
||||
mem.save_or_update(tagged("date date date", "a")).unwrap();
|
||||
check(&mut mem, "after in-place update");
|
||||
mem.delete(1).unwrap();
|
||||
check(&mut mem, "after delete");
|
||||
mem.save(tagged("apple cherry", "e")).unwrap();
|
||||
check(&mut mem, "after save following a delete");
|
||||
mem.compact().unwrap();
|
||||
check(&mut mem, "after compact");
|
||||
mem.hybrid_search(&[1.0, 0.0, 0.0, 0.0], "apple", 0.5, 0.5, 3);
|
||||
check(&mut mem, "after a search");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn search_does_not_write_the_store_but_boosts_persist_on_drop() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let config = make_config(&dir);
|
||||
let path = config.path.clone();
|
||||
let mut mem = HDF5Memory::create(config).unwrap();
|
||||
mem.save(make_entry("findable", &[1.0, 0.0, 0.0, 0.0]))
|
||||
.unwrap();
|
||||
let before = std::fs::read(&path).unwrap();
|
||||
|
||||
for _ in 0..3 {
|
||||
mem.hybrid_search(&[1.0, 0.0, 0.0, 0.0], "findable", 1.0, 0.0, 1);
|
||||
}
|
||||
assert_eq!(
|
||||
std::fs::read(&path).unwrap(),
|
||||
before,
|
||||
"a query must not rewrite the store"
|
||||
);
|
||||
let boosted = mem.cache.activation_weights[0];
|
||||
assert!(boosted > 1.0);
|
||||
drop(mem);
|
||||
|
||||
let reopened = HDF5Memory::open(&path).unwrap();
|
||||
assert_eq!(reopened.cache.activation_weights[0], boosted);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn activation_weight_is_capped() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let mut mem = HDF5Memory::create(make_config(&dir)).unwrap();
|
||||
mem.save(make_entry("popular", &[1.0, 0.0, 0.0, 0.0]))
|
||||
.unwrap();
|
||||
for _ in 0..500 {
|
||||
mem.hybrid_search(&[1.0, 0.0, 0.0, 0.0], "popular", 1.0, 0.0, 1);
|
||||
}
|
||||
assert_eq!(mem.cache.activation_weights[0], MAX_ACTIVATION_WEIGHT);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn store_has_a_single_writer() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
|
||||
@@ -466,7 +466,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,
|
||||
@@ -531,11 +531,14 @@ impl MemoryBackend for ClawhdfBackend {
|
||||
query_embedding: &[f32],
|
||||
k: usize,
|
||||
) -> Vec<MemorySearchResult> {
|
||||
// 1. Hybrid retrieval (RRF-blended vector + BM25).
|
||||
// 1. Hybrid retrieval (vector + BM25, fused by score).
|
||||
let candidates = k.saturating_mul(3).max(10);
|
||||
let raw = self
|
||||
.memory
|
||||
.hybrid_search(query_embedding, query_text, 0.7, 0.3, candidates);
|
||||
let raw = self.memory.hybrid_search_with(
|
||||
query_embedding,
|
||||
query_text,
|
||||
crate::hybrid::DEFAULT_FUSION,
|
||||
candidates,
|
||||
);
|
||||
|
||||
if raw.is_empty() {
|
||||
return Vec::new();
|
||||
@@ -551,6 +554,7 @@ impl MemoryBackend for ClawhdfBackend {
|
||||
timestamp: r.timestamp,
|
||||
source_channel: r.source_channel.clone(),
|
||||
raw_activation: r.activation,
|
||||
relevance: r.score,
|
||||
})
|
||||
.collect();
|
||||
|
||||
@@ -713,11 +717,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)
|
||||
|
||||
@@ -6,6 +6,11 @@
|
||||
//! - Temporal expansion (time-related rewrites)
|
||||
//! - Morphological variants (stemming-like transforms)
|
||||
//! - Knowledge graph expansion (entity aliases and neighbors)
|
||||
//!
|
||||
//! The morphological rules are crude suffix swaps, so some variants are not
|
||||
//! words ("during" -> "dured"). That is tolerable for a BM25 stage, which
|
||||
//! simply finds no postings for a nonsense term, but it means expansion is not
|
||||
//! free: measure before enabling it on a retrieval path.
|
||||
|
||||
use crate::knowledge::KnowledgeCache;
|
||||
|
||||
@@ -340,20 +345,87 @@ fn contains_phrase(text: &str, phrase: &str) -> bool {
|
||||
|
||||
/// Replace a phrase in `text` case-insensitively, preserving surrounding case.
|
||||
fn replace_word_case_insensitive(text: &str, from: &str, to: &str) -> String {
|
||||
case_insensitive_replace(text, from, to)
|
||||
replace_first(text, from, to, MatchKind::WholeWord)
|
||||
}
|
||||
|
||||
fn case_insensitive_replace(text: &str, from: &str, to: &str) -> String {
|
||||
let lower = text.to_lowercase();
|
||||
let lower_from = from.to_lowercase();
|
||||
if let Some(pos) = lower.find(&lower_from) {
|
||||
let end = pos + from.len();
|
||||
format!("{}{}{}", &text[..pos], to, &text[end..])
|
||||
} else {
|
||||
text.to_string()
|
||||
replace_first(text, from, to, MatchKind::Substring)
|
||||
}
|
||||
|
||||
/// Whether a match may fall inside a larger word.
|
||||
#[derive(Clone, Copy, PartialEq)]
|
||||
enum MatchKind {
|
||||
/// Match anywhere, including inside another word.
|
||||
Substring,
|
||||
/// Match only when both ends sit on a word boundary.
|
||||
WholeWord,
|
||||
}
|
||||
|
||||
/// Replace the first case-insensitive match of `from` in `text` with `to`.
|
||||
///
|
||||
/// Matching walks the *original* string rather than a lowercased copy. The
|
||||
/// previous implementation searched `text.to_lowercase()` and then sliced
|
||||
/// `text` with the offsets it found, which only holds while lowercasing
|
||||
/// preserves byte length. It does not: Turkish `İ` (2 bytes) lowercases to
|
||||
/// `i` + U+0307 (3 bytes), so every later offset was wrong — silently
|
||||
/// corrupting the output, or panicking when an offset landed inside a
|
||||
/// character or past the end. `"İ AI"` was enough to panic.
|
||||
fn replace_first(text: &str, from: &str, to: &str, kind: MatchKind) -> String {
|
||||
match find_case_insensitive(text, from, kind) {
|
||||
Some((start, end)) => {
|
||||
let mut out = String::with_capacity(text.len() - (end - start) + to.len());
|
||||
out.push_str(&text[..start]);
|
||||
out.push_str(to);
|
||||
out.push_str(&text[end..]);
|
||||
out
|
||||
}
|
||||
None => text.to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Byte range of the first case-insensitive match of `needle` in `haystack`.
|
||||
fn find_case_insensitive(haystack: &str, needle: &str, kind: MatchKind) -> Option<(usize, usize)> {
|
||||
if needle.is_empty() {
|
||||
return None;
|
||||
}
|
||||
let lowered: Vec<char> = needle.chars().flat_map(char::to_lowercase).collect();
|
||||
let is_word = |c: char| c.is_alphanumeric() || c == '_';
|
||||
|
||||
for (start, _) in haystack.char_indices() {
|
||||
if kind == MatchKind::WholeWord
|
||||
&& haystack[..start].chars().next_back().is_some_and(is_word)
|
||||
{
|
||||
continue; // mid-word: "ai" inside "training"
|
||||
}
|
||||
let mut matched = 0usize;
|
||||
let mut end = start;
|
||||
for (offset, ch) in haystack[start..].char_indices() {
|
||||
if matched == lowered.len() {
|
||||
break;
|
||||
}
|
||||
let mut consumed_all = true;
|
||||
for lc in ch.to_lowercase() {
|
||||
if lowered.get(matched) != Some(&lc) {
|
||||
consumed_all = false;
|
||||
break;
|
||||
}
|
||||
matched += 1;
|
||||
}
|
||||
if !consumed_all {
|
||||
break;
|
||||
}
|
||||
end = start + offset + ch.len_utf8();
|
||||
}
|
||||
if matched == lowered.len()
|
||||
&& !(kind == MatchKind::WholeWord
|
||||
&& haystack[end..].chars().next().is_some_and(is_word))
|
||||
{
|
||||
return Some((start, end));
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
/// Simple whitespace/punctuation tokenizer.
|
||||
fn tokenize(text: &str) -> Vec<String> {
|
||||
text.split(|c: char| !c.is_alphanumeric())
|
||||
@@ -637,4 +709,86 @@ mod tests {
|
||||
expanded.iter().map(|x| &x.text).collect::<Vec<_>>()
|
||||
);
|
||||
}
|
||||
#[test]
|
||||
fn acronyms_only_match_whole_words() {
|
||||
let ex = QueryExpander::new(QueryExpansionConfig::default());
|
||||
// "training" contains "ai", "programming" contains "pr". These used to
|
||||
// be rewritten to "trArtificial Intelligencening" and
|
||||
// "Pull Requestogramming".
|
||||
for query in [
|
||||
"How many miles during my marathon training?",
|
||||
"Which programming language did I pick?",
|
||||
"I updated the maintainer list",
|
||||
] {
|
||||
for expansion in ex.expand(query) {
|
||||
assert!(
|
||||
expansion.expansion_type != "acronym",
|
||||
"{query:?} produced {expansion:?}"
|
||||
);
|
||||
}
|
||||
}
|
||||
// A real acronym still expands, in both directions.
|
||||
let texts: Vec<String> = ex
|
||||
.expand("What about the API and the database?")
|
||||
.into_iter()
|
||||
.filter(|e| e.expansion_type == "acronym")
|
||||
.map(|e| e.text)
|
||||
.collect();
|
||||
assert!(
|
||||
texts
|
||||
.iter()
|
||||
.any(|t| t.contains("Application Programming Interface")),
|
||||
"{texts:?}"
|
||||
);
|
||||
assert!(texts.iter().any(|t| t.contains("DB")), "{texts:?}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn non_ascii_queries_do_not_panic_or_corrupt() {
|
||||
let ex = QueryExpander::new(QueryExpansionConfig::default());
|
||||
// Turkish 'İ' is 2 bytes but lowercases to 3, so offsets taken from a
|
||||
// lowercased copy no longer line up with the original. `"İ AI"` used
|
||||
// to panic; `"İstanbul AI trip"` used to silently eat a character.
|
||||
for query in ["İ AI", "İé AI", "İİ ML", "İstanbul AI trip", "ǰ ML notes"] {
|
||||
for expansion in ex.expand(query) {
|
||||
assert!(
|
||||
expansion.text.contains('İ') || expansion.text.contains('ǰ'),
|
||||
"{query:?} lost its leading character: {expansion:?}"
|
||||
);
|
||||
}
|
||||
}
|
||||
let expanded = ex.expand("İstanbul AI trip");
|
||||
assert!(
|
||||
expanded
|
||||
.iter()
|
||||
.any(|e| e.text == "İstanbul Artificial Intelligence trip"),
|
||||
"{expanded:?}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn whole_word_matching_handles_string_edges_and_case() {
|
||||
assert_eq!(
|
||||
replace_word_case_insensitive("ai tools", "AI", "Artificial Intelligence"),
|
||||
"Artificial Intelligence tools"
|
||||
);
|
||||
assert_eq!(
|
||||
replace_word_case_insensitive("tools for ai", "AI", "Artificial Intelligence"),
|
||||
"tools for Artificial Intelligence"
|
||||
);
|
||||
assert_eq!(
|
||||
replace_word_case_insensitive("the aim", "AI", "Artificial Intelligence"),
|
||||
"the aim",
|
||||
"must not match inside a word"
|
||||
);
|
||||
assert_eq!(
|
||||
replace_word_case_insensitive("no match here", "xyz", "abc"),
|
||||
"no match here"
|
||||
);
|
||||
// Only the first occurrence is replaced, as before.
|
||||
assert_eq!(
|
||||
replace_word_case_insensitive("ai and ai", "ai", "ML"),
|
||||
"ML and ai"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -4,8 +4,10 @@
|
||||
//! into a single composite score for each retrieved result.
|
||||
|
||||
/// 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);
|
||||
|
||||
@@ -22,6 +22,7 @@ pub const ZEROCLAW_VERSION: &str = "0.8.0";
|
||||
/// the checkpoint was taken with an empty WAL.
|
||||
const WAL_APPLIED_LEN_ATTR: &str = "wal_applied_len";
|
||||
const WAL_APPLIED_CRC_ATTR: &str = "wal_applied_crc";
|
||||
const ANN_GENERATION_ATTR: &str = "ann_generation";
|
||||
|
||||
/// Build a complete HDF5 file from the in-memory state.
|
||||
pub fn build_hdf5_file(
|
||||
@@ -43,6 +44,34 @@ pub fn build_hdf5_file_with_mark(
|
||||
knowledge: &KnowledgeCache,
|
||||
wal_applied: Option<WalMark>,
|
||||
) -> Result<Vec<u8>, MemoryError> {
|
||||
let meta = CheckpointMeta {
|
||||
wal_applied,
|
||||
ann_generation: None,
|
||||
};
|
||||
build_hdf5_file_with_meta(config, cache, sessions, knowledge, &meta)
|
||||
}
|
||||
|
||||
/// Bookkeeping a checkpoint records in `/meta` beside the store's contents.
|
||||
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
|
||||
pub struct CheckpointMeta {
|
||||
/// The WAL prefix this checkpoint already contains; see [`WalMark`].
|
||||
pub wal_applied: Option<WalMark>,
|
||||
/// Identifies the vector-index sidecar (`<store>.h5.ann`) written with this
|
||||
/// checkpoint. A sidecar is loaded only if it carries the same value, so
|
||||
/// one left over from another checkpoint can never be attached to records
|
||||
/// it wasn't built from.
|
||||
pub ann_generation: Option<u64>,
|
||||
}
|
||||
|
||||
/// [`build_hdf5_file`] with checkpoint bookkeeping.
|
||||
pub fn build_hdf5_file_with_meta(
|
||||
config: &MemoryConfig,
|
||||
cache: &MemoryCache,
|
||||
sessions: &SessionCache,
|
||||
knowledge: &KnowledgeCache,
|
||||
checkpoint: &CheckpointMeta,
|
||||
) -> Result<Vec<u8>, MemoryError> {
|
||||
let wal_applied = checkpoint.wal_applied;
|
||||
let mut builder = clawhdf5::FileBuilder::new();
|
||||
|
||||
// /meta group with schema attributes
|
||||
@@ -75,6 +104,10 @@ pub fn build_hdf5_file_with_mark(
|
||||
"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(
|
||||
"edgehdf5_version",
|
||||
AttrValue::String(ZEROCLAW_VERSION.into()),
|
||||
@@ -83,6 +116,11 @@ pub fn build_hdf5_file_with_mark(
|
||||
meta.set_attr(WAL_APPLIED_LEN_ATTR, AttrValue::I64(mark.len as i64));
|
||||
meta.set_attr(WAL_APPLIED_CRC_ATTR, AttrValue::I64(i64::from(mark.crc)));
|
||||
}
|
||||
if let Some(generation) = checkpoint.ann_generation {
|
||||
// Stored as the i64 with the same bits; attributes have no u64 scalar
|
||||
// round trip through every reader.
|
||||
meta.set_attr(ANN_GENERATION_ATTR, AttrValue::I64(generation as i64));
|
||||
}
|
||||
// Need at least one dataset in the group for it to be a proper group
|
||||
meta.create_dataset("_marker").with_u8_data(&[1]).compact();
|
||||
let finished_meta = meta.finish();
|
||||
@@ -119,7 +157,7 @@ fn build_memory_group(
|
||||
{
|
||||
let ds = group
|
||||
.create_dataset("embeddings")
|
||||
.with_f32_data(&flat)
|
||||
.with_f32_data(flat)
|
||||
.with_shape(&[n, d]);
|
||||
|
||||
// Chunk size tuning: target ~256KB per chunk for optimal I/O
|
||||
@@ -386,6 +424,22 @@ pub fn read_wal_mark(file: &clawhdf5::File) -> Option<WalMark> {
|
||||
Some(WalMark { len, crc })
|
||||
}
|
||||
|
||||
/// Read the checkpoint bookkeeping from `/meta`.
|
||||
pub fn read_checkpoint_meta(file: &clawhdf5::File) -> CheckpointMeta {
|
||||
let ann_generation = file
|
||||
.group("meta")
|
||||
.ok()
|
||||
.and_then(|g| g.attrs().ok())
|
||||
.and_then(|attrs| match attrs.get(ANN_GENERATION_ATTR) {
|
||||
Some(AttrValue::I64(v)) => Some(*v as u64),
|
||||
_ => None,
|
||||
});
|
||||
CheckpointMeta {
|
||||
wal_applied: read_wal_mark(file),
|
||||
ann_generation,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn validate_and_load(
|
||||
file: &clawhdf5::File,
|
||||
) -> Result<(MemoryConfig, MemoryCache, SessionCache, KnowledgeCache), MemoryError> {
|
||||
@@ -434,6 +488,7 @@ 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),
|
||||
quantized_index: optional_bool_attr(&attrs, "quantized_index", false),
|
||||
};
|
||||
|
||||
// Load /memory group
|
||||
@@ -513,12 +568,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,
|
||||
@@ -526,7 +576,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;
|
||||
@@ -534,7 +584,6 @@ fn load_memory_group(
|
||||
cache.tombstones = tombstones;
|
||||
cache.norms = norms;
|
||||
cache.activation_weights = activation_weights;
|
||||
cache.rebuild_flat();
|
||||
|
||||
Ok(cache)
|
||||
}
|
||||
|
||||
@@ -4,7 +4,7 @@ use std::path::Path;
|
||||
|
||||
use crate::bm25;
|
||||
use crate::hybrid;
|
||||
use crate::{HDF5Memory, MemoryError, Result, SearchResult};
|
||||
use crate::{HDF5Memory, MAX_ACTIVATION_WEIGHT, MemoryError, Result, SearchResult};
|
||||
|
||||
impl HDF5Memory {
|
||||
/// Vector + keyword scoring stage of [`HDF5Memory::hybrid_search`].
|
||||
@@ -20,8 +20,7 @@ impl HDF5Memory {
|
||||
query_embedding: &[f32],
|
||||
query_text: &str,
|
||||
bm25: &bm25::BM25Index,
|
||||
vector_weight: f32,
|
||||
keyword_weight: f32,
|
||||
fusion: hybrid::Fusion,
|
||||
k: usize,
|
||||
) -> Vec<(usize, f32)> {
|
||||
self.ensure_hnsw_fresh();
|
||||
@@ -30,29 +29,40 @@ impl HDF5Memory {
|
||||
// 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 candidates = index.search(query_embedding, pool, pool);
|
||||
// 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))
|
||||
.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();
|
||||
let kw_scores = bm25.search(query_text, self.cache.len());
|
||||
hybrid::merge_vector_keyword(
|
||||
vec_scores,
|
||||
kw_scores,
|
||||
vector_weight,
|
||||
keyword_weight,
|
||||
k,
|
||||
)
|
||||
// Fusion normalises over every keyword match, so it needs all
|
||||
// the scores — but not ranked.
|
||||
let kw_scores = bm25.scores(query_text);
|
||||
hybrid::fuse(vec_scores, kw_scores, fusion, k)
|
||||
}
|
||||
_ => hybrid::hybrid_search(
|
||||
_ => hybrid::hybrid_search_fused(
|
||||
query_embedding,
|
||||
query_text,
|
||||
&self.cache.embeddings,
|
||||
&self.cache.chunks,
|
||||
&self.cache.tombstones,
|
||||
bm25,
|
||||
vector_weight,
|
||||
keyword_weight,
|
||||
fusion,
|
||||
k,
|
||||
),
|
||||
}
|
||||
@@ -64,19 +74,17 @@ impl HDF5Memory {
|
||||
query_embedding: &[f32],
|
||||
query_text: &str,
|
||||
bm25: &bm25::BM25Index,
|
||||
vector_weight: f32,
|
||||
keyword_weight: f32,
|
||||
fusion: hybrid::Fusion,
|
||||
k: usize,
|
||||
) -> Vec<(usize, f32)> {
|
||||
hybrid::hybrid_search(
|
||||
hybrid::hybrid_search_fused(
|
||||
query_embedding,
|
||||
query_text,
|
||||
&self.cache.embeddings,
|
||||
&self.cache.chunks,
|
||||
&self.cache.tombstones,
|
||||
bm25,
|
||||
vector_weight,
|
||||
keyword_weight,
|
||||
fusion,
|
||||
k,
|
||||
)
|
||||
}
|
||||
@@ -90,15 +98,35 @@ impl HDF5Memory {
|
||||
keyword_weight: f32,
|
||||
k: usize,
|
||||
) -> Vec<SearchResult> {
|
||||
let bm25 = bm25::BM25Index::build(&self.cache.chunks, &self.cache.tombstones);
|
||||
let scored = self.vector_keyword_search(
|
||||
self.hybrid_search_with(
|
||||
query_embedding,
|
||||
query_text,
|
||||
&bm25,
|
||||
vector_weight,
|
||||
keyword_weight,
|
||||
hybrid::Fusion::Weighted {
|
||||
vector: vector_weight,
|
||||
keyword: keyword_weight,
|
||||
},
|
||||
k,
|
||||
);
|
||||
)
|
||||
}
|
||||
|
||||
/// [`HDF5Memory::hybrid_search`] with the fusion method chosen explicitly.
|
||||
///
|
||||
/// [`hybrid::DEFAULT_FUSION`] is what the weighted form defaults to;
|
||||
/// [`hybrid::Fusion::Rrf`] combines the two stages by rank instead of by
|
||||
/// score.
|
||||
pub fn hybrid_search_with(
|
||||
&mut self,
|
||||
query_embedding: &[f32],
|
||||
query_text: &str,
|
||||
fusion: hybrid::Fusion,
|
||||
k: usize,
|
||||
) -> Vec<SearchResult> {
|
||||
// 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 mut results: Vec<SearchResult> = scored
|
||||
.into_iter()
|
||||
.map(|(idx, score)| {
|
||||
@@ -132,15 +160,26 @@ impl HDF5Memory {
|
||||
.map(|r| r.index)
|
||||
.collect();
|
||||
self.apply_hebbian_boost(&hit_indices);
|
||||
self.flush().ok();
|
||||
self.bm25 = Some(bm25);
|
||||
|
||||
results
|
||||
}
|
||||
|
||||
/// 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
|
||||
/// `hybrid_search` cost O(store size) in disk I/O. They are a ranking hint,
|
||||
/// not user data: a crash before the next checkpoint only forgets the
|
||||
/// boosts since the last one.
|
||||
fn apply_hebbian_boost(&mut self, hit_indices: &[usize]) {
|
||||
for &idx in hit_indices {
|
||||
self.cache.activation_weights[idx] += self.config.hebbian_boost;
|
||||
if hit_indices.is_empty() || self.config.hebbian_boost == 0.0 {
|
||||
return;
|
||||
}
|
||||
for &idx in hit_indices {
|
||||
let w = &mut self.cache.activation_weights[idx];
|
||||
*w = (*w + self.config.hebbian_boost).min(MAX_ACTIVATION_WEIGHT);
|
||||
}
|
||||
self.activations_dirty = true;
|
||||
}
|
||||
|
||||
/// Get the chunk text for a memory entry by index.
|
||||
|
||||
@@ -34,7 +34,23 @@ pub fn write_to_disk_with_mark(
|
||||
knowledge: &KnowledgeCache,
|
||||
wal_applied: Option<WalMark>,
|
||||
) -> Result<(), MemoryError> {
|
||||
let bytes = schema::build_hdf5_file_with_mark(config, cache, sessions, knowledge, wal_applied)?;
|
||||
let meta = schema::CheckpointMeta {
|
||||
wal_applied,
|
||||
ann_generation: None,
|
||||
};
|
||||
write_to_disk_with_meta(path, config, cache, sessions, knowledge, &meta)
|
||||
}
|
||||
|
||||
/// [`write_to_disk`] with full checkpoint bookkeeping.
|
||||
pub fn write_to_disk_with_meta(
|
||||
path: &Path,
|
||||
config: &MemoryConfig,
|
||||
cache: &MemoryCache,
|
||||
sessions: &SessionCache,
|
||||
knowledge: &KnowledgeCache,
|
||||
checkpoint: &schema::CheckpointMeta,
|
||||
) -> Result<(), MemoryError> {
|
||||
let bytes = schema::build_hdf5_file_with_meta(config, cache, sessions, knowledge, checkpoint)?;
|
||||
|
||||
if bytes.is_empty() {
|
||||
return Err(MemoryError::Hdf5("build_hdf5_file produced 0 bytes".into()));
|
||||
@@ -47,7 +63,7 @@ pub fn write_to_disk_with_mark(
|
||||
}
|
||||
|
||||
/// Write `bytes` to `path` and flush them to stable storage.
|
||||
fn write_synced(path: &Path, bytes: &[u8]) -> Result<(), MemoryError> {
|
||||
pub(crate) fn write_synced(path: &Path, bytes: &[u8]) -> Result<(), MemoryError> {
|
||||
use std::io::Write;
|
||||
let mut f = std::fs::File::create(path).map_err(MemoryError::Io)?;
|
||||
f.write_all(bytes).map_err(MemoryError::Io)?;
|
||||
@@ -62,7 +78,7 @@ fn write_synced(path: &Path, bytes: &[u8]) -> Result<(), MemoryError> {
|
||||
/// This is per-checkpoint/snapshot cost only (each is already a full file
|
||||
/// write). Individual WAL appends are deliberately not synced — see the
|
||||
/// durability notes in the crate docs.
|
||||
fn rename_synced(from: &Path, to: &Path) -> Result<(), MemoryError> {
|
||||
pub(crate) fn rename_synced(from: &Path, to: &Path) -> Result<(), MemoryError> {
|
||||
std::fs::rename(from, to).map_err(MemoryError::Io)?;
|
||||
#[cfg(unix)]
|
||||
if let Some(dir) = to.parent() {
|
||||
@@ -113,6 +129,20 @@ pub fn read_from_disk_with_mark(path: &Path) -> Result<(StoreState, Option<WalMa
|
||||
Ok(((config, cache, sessions, knowledge), wal_applied))
|
||||
}
|
||||
|
||||
/// [`read_from_disk`], plus all checkpoint bookkeeping.
|
||||
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())
|
||||
.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();
|
||||
let meta = schema::read_checkpoint_meta(&file);
|
||||
Ok(((config, cache, sessions, knowledge), meta))
|
||||
}
|
||||
|
||||
/// Copy an HDF5 file atomically to a destination.
|
||||
pub fn snapshot_file(src: &Path, dest: &Path) -> Result<std::path::PathBuf, MemoryError> {
|
||||
let dest_file = if dest.is_dir() {
|
||||
|
||||
@@ -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));
|
||||
|
||||
@@ -165,3 +165,72 @@ 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);
|
||||
}
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
[package]
|
||||
name = "clawhdf5-android"
|
||||
version = "2.3.0"
|
||||
version = "2.6.0"
|
||||
edition = "2024"
|
||||
description = "Android JNI bridge for edgehdf5-memory HDF5 backend"
|
||||
license = "MIT"
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
[package]
|
||||
name = "clawhdf5-ann"
|
||||
version = "2.3.0"
|
||||
version = "2.6.0"
|
||||
edition = "2024"
|
||||
description = "HNSW approximate nearest neighbor index stored as HDF5"
|
||||
license = "MIT"
|
||||
@@ -10,9 +10,9 @@ keywords = ["hdf5", "ann", "hnsw", "nearest-neighbor"]
|
||||
categories = ["algorithms", "science"]
|
||||
|
||||
[dependencies]
|
||||
clawhdf5-format = { path = "../clawhdf5-format", version = "2.3.0" }
|
||||
clawhdf5-io = { path = "../clawhdf5-io", version = "2.3.0" }
|
||||
clawhdf5-accel = { path = "../clawhdf5-accel", version = "2.3.0" }
|
||||
clawhdf5-format = { path = "../clawhdf5-format", version = "2.6.0" }
|
||||
clawhdf5-io = { path = "../clawhdf5-io", version = "2.6.0" }
|
||||
clawhdf5-accel = { path = "../clawhdf5-accel", version = "2.6.0" }
|
||||
rayon = { version = "1", optional = true }
|
||||
|
||||
[features]
|
||||
|
||||
+1131
-136
File diff suppressed because it is too large
Load Diff
@@ -5,4 +5,4 @@
|
||||
|
||||
mod hnsw;
|
||||
|
||||
pub use hnsw::{DistanceMetric, HnswIndex};
|
||||
pub use hnsw::{DistanceMetric, HnswIndex, Storage};
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
[package]
|
||||
name = "clawhdf5-bench"
|
||||
version = "2.3.0"
|
||||
version = "2.6.0"
|
||||
edition = "2024"
|
||||
description = "Benchmark harnesses for clawhdf5-agent (Track 8)"
|
||||
license = "MIT"
|
||||
@@ -13,6 +13,14 @@ path = "src/bin/longmemeval_bench.rs"
|
||||
name = "memory_arena"
|
||||
path = "src/bin/memory_arena.rs"
|
||||
|
||||
[[bin]]
|
||||
name = "read_harness"
|
||||
path = "src/bin/read_harness.rs"
|
||||
|
||||
[[bin]]
|
||||
name = "search_harness"
|
||||
path = "src/bin/search_harness.rs"
|
||||
|
||||
[[bin]]
|
||||
name = "footprint_bench"
|
||||
path = "src/bin/footprint_bench.rs"
|
||||
@@ -48,6 +56,9 @@ harness = false
|
||||
|
||||
[dependencies]
|
||||
clawhdf5-agent = { path = "../clawhdf5-agent" }
|
||||
clawhdf5-ann = { path = "../clawhdf5-ann" }
|
||||
clawhdf5 = { path = "../clawhdf5" }
|
||||
clawhdf5-format = { path = "../clawhdf5-format" }
|
||||
clawhdf5-io = { path = "../clawhdf5-io" }
|
||||
mpi = { version = "0.8", optional = true }
|
||||
serde = { workspace = true }
|
||||
|
||||
@@ -55,44 +55,150 @@ use std::time::{Duration, Instant};
|
||||
#[path = "longmemeval_bench/embedder.rs"]
|
||||
mod embedder;
|
||||
|
||||
use clawhdf5_agent::{AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry};
|
||||
use clawhdf5_agent::bm25::TokenFilter;
|
||||
use clawhdf5_agent::hybrid::Fusion;
|
||||
use clawhdf5_agent::reranker::{ReRankConfig, RerankInput, rerank};
|
||||
use clawhdf5_agent::{AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry, SearchResult};
|
||||
use serde::Deserialize;
|
||||
use tempfile::TempDir;
|
||||
|
||||
const EMBEDDING_DIM: usize = 384;
|
||||
|
||||
/// A mode's fusion, as one short string for the reports.
|
||||
fn describe(mode: Mode) -> String {
|
||||
let fusion = match mode.fusion {
|
||||
Fusion::Weighted { vector, keyword } => format!("vector_{vector:.1}_keyword_{keyword:.1}"),
|
||||
Fusion::Rrf { k } => format!("rrf_k{k:.0}"),
|
||||
};
|
||||
let tokens = match mode.tokens {
|
||||
TokenFilter::Plain => fusion,
|
||||
TokenFilter::Stemmed => format!("{fusion}_stemmed"),
|
||||
};
|
||||
match mode.rerank {
|
||||
None => tokens,
|
||||
Some(cfg) if cfg.relevance_weight == 0.0 => format!("{tokens}_rerank_metadata"),
|
||||
Some(cfg) => format!(
|
||||
"{tokens}_rerank_blended_hl{:.0}d",
|
||||
cfg.temporal_half_life_secs / 86_400.0
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
/// A retrieval configuration: how much of the score comes from each stage.
|
||||
#[derive(Clone, Copy)]
|
||||
struct Mode {
|
||||
label: &'static str,
|
||||
vector_weight: f32,
|
||||
keyword_weight: f32,
|
||||
/// How the two retrieval stages are combined into one ranking.
|
||||
fusion: Fusion,
|
||||
/// How keyword tokens are normalised before indexing and querying.
|
||||
tokens: TokenFilter,
|
||||
/// Re-rank the retrieved candidates with recency and friends, relative to
|
||||
/// the question's own date.
|
||||
rerank: Option<ReRankConfig>,
|
||||
}
|
||||
|
||||
impl Mode {
|
||||
const fn weighted(label: &'static str, vector: f32, keyword: f32) -> Self {
|
||||
Self {
|
||||
label,
|
||||
fusion: Fusion::Weighted { vector, keyword },
|
||||
tokens: TokenFilter::Plain,
|
||||
rerank: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg_attr(not(feature = "embeddings"), allow(dead_code))]
|
||||
fn reranked(mut self, label: &'static str, rerank: ReRankConfig) -> Self {
|
||||
self.label = label;
|
||||
self.rerank = Some(rerank);
|
||||
self
|
||||
}
|
||||
|
||||
const fn stemmed(mut self, label: &'static str) -> Self {
|
||||
self.label = label;
|
||||
self.tokens = TokenFilter::Stemmed;
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
/// The only mode available without real embeddings. Passing zero vectors with
|
||||
/// `vector_weight = 0.0` is what made the vector stage inert.
|
||||
const BM25_ONLY: Mode = Mode {
|
||||
label: "BM25 only (vector stage inert)",
|
||||
vector_weight: 0.0,
|
||||
keyword_weight: 1.0,
|
||||
};
|
||||
const BM25_ONLY: Mode = Mode::weighted("BM25 only (vector stage inert)", 0.0, 1.0);
|
||||
#[cfg(feature = "embeddings")]
|
||||
const VECTOR_ONLY: Mode = Mode {
|
||||
label: "Vector only (MiniLM + HNSW)",
|
||||
vector_weight: 1.0,
|
||||
keyword_weight: 0.0,
|
||||
};
|
||||
const VECTOR_ONLY: Mode = Mode::weighted("Vector only (MiniLM + HNSW)", 1.0, 0.0);
|
||||
/// Tuned by `--sweep` over the full haystack. The former 0.7/0.3 was a
|
||||
/// documented default that had never been searched, and the sweep found it
|
||||
/// strictly dominated: 0.4/0.6 is better on Hit@1, Hit@5, Hit@10 and MRR at
|
||||
/// both granularities.
|
||||
#[cfg(feature = "embeddings")]
|
||||
const HYBRID: Mode = Mode {
|
||||
label: "Hybrid (0.4 vector / 0.6 BM25, tuned)",
|
||||
vector_weight: 0.4,
|
||||
keyword_weight: 0.6,
|
||||
const HYBRID: Mode = Mode::weighted("Hybrid (0.4 vector / 0.6 BM25, tuned)", 0.4, 0.6);
|
||||
|
||||
/// Reciprocal rank fusion, the documented alternative to the weighted sum.
|
||||
/// It ignores score magnitudes, so there is nothing to tune — which is the
|
||||
/// claim being tested.
|
||||
#[cfg(feature = "embeddings")]
|
||||
const RRF: Mode = Mode {
|
||||
label: "Hybrid (reciprocal rank fusion, k=60)",
|
||||
fusion: Fusion::Rrf { k: 60.0 },
|
||||
tokens: TokenFilter::Plain,
|
||||
rerank: None,
|
||||
};
|
||||
|
||||
/// The same two configurations with stemmed keyword tokens, so the tokenizer's
|
||||
/// effect is isolated from everything else.
|
||||
const BM25_STEMMED: Mode = BM25_ONLY.stemmed("BM25 only, stemmed tokens");
|
||||
|
||||
/// Re-ranking as it behaved before `relevance` was an input: the combined
|
||||
/// score was recency + authority + activation only, so the retriever's own
|
||||
/// ordering was discarded.
|
||||
#[cfg(feature = "embeddings")]
|
||||
fn hybrid_rerank_metadata_only() -> Mode {
|
||||
HYBRID.reranked(
|
||||
"Hybrid + rerank (metadata only, pre-fix)",
|
||||
ReRankConfig {
|
||||
relevance_weight: 0.0,
|
||||
..ReRankConfig::default()
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
/// Re-ranking as it behaves now: relevance leads, recency nudges.
|
||||
#[cfg(feature = "embeddings")]
|
||||
fn hybrid_rerank_blended() -> Mode {
|
||||
HYBRID.reranked(
|
||||
"Hybrid + rerank (relevance + recency)",
|
||||
ReRankConfig::default(),
|
||||
)
|
||||
}
|
||||
|
||||
/// The same blend at several half-lives. Decay is `2^(-age / half_life)`, so a
|
||||
/// half-life far shorter than the gaps between memories sends every score to
|
||||
/// zero and the signal vanishes; far longer and everything scores ~1 and it
|
||||
/// vanishes the other way. The right value tracks how far apart the memories
|
||||
/// actually are.
|
||||
#[cfg(feature = "embeddings")]
|
||||
fn hybrid_rerank_half_lives() -> Vec<Mode> {
|
||||
[
|
||||
("1 day", 86_400.0),
|
||||
("7 days", 7.0 * 86_400.0),
|
||||
("30 days", 30.0 * 86_400.0),
|
||||
("90 days", 90.0 * 86_400.0),
|
||||
]
|
||||
.into_iter()
|
||||
.map(|(label, half_life)| {
|
||||
HYBRID.reranked(
|
||||
Box::leak(format!("Hybrid + rerank, half-life {label}").into_boxed_str()),
|
||||
ReRankConfig {
|
||||
temporal_half_life_secs: half_life,
|
||||
..ReRankConfig::default()
|
||||
},
|
||||
)
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
#[cfg(feature = "embeddings")]
|
||||
const HYBRID_STEMMED: Mode = HYBRID.stemmed("Hybrid 0.4/0.6, stemmed tokens");
|
||||
|
||||
/// Every 0.1 step of vector weight, keyword weight taking the remainder.
|
||||
///
|
||||
/// Labels are leaked to `&'static str` because `Mode::label` is a `&'static
|
||||
@@ -104,11 +210,11 @@ fn sweep_modes() -> Vec<Mode> {
|
||||
(0..=10)
|
||||
.map(|i| {
|
||||
let v = i as f32 / 10.0;
|
||||
Mode {
|
||||
label: Box::leak(format!("sweep v={v:.1} / k={:.1}", 1.0 - v).into_boxed_str()),
|
||||
vector_weight: v,
|
||||
keyword_weight: 1.0 - v,
|
||||
}
|
||||
Mode::weighted(
|
||||
Box::leak(format!("sweep v={v:.1} / k={:.1}", 1.0 - v).into_boxed_str()),
|
||||
v,
|
||||
1.0 - v,
|
||||
)
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
@@ -181,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)
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
@@ -199,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
|
||||
}
|
||||
@@ -261,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,
|
||||
}
|
||||
|
||||
@@ -276,19 +433,26 @@ fn evaluate_question(
|
||||
config.compact_threshold = 0.0;
|
||||
|
||||
let mut memory = HDF5Memory::create(config).expect("failed to create HDF5Memory");
|
||||
memory.set_token_filter(mode.tokens);
|
||||
|
||||
// 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),
|
||||
@@ -302,7 +466,6 @@ fn evaluate_question(
|
||||
},
|
||||
});
|
||||
turn_has_answer.push(turn.has_answer);
|
||||
ts += 1.0;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -319,17 +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(
|
||||
&query_emb,
|
||||
&q.question,
|
||||
mode.vector_weight,
|
||||
mode.keyword_weight,
|
||||
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;
|
||||
@@ -384,6 +617,7 @@ fn evaluate_question(
|
||||
hit5_turn,
|
||||
hit10_turn,
|
||||
rr_turn,
|
||||
newest_gold_first,
|
||||
latency,
|
||||
}
|
||||
}
|
||||
@@ -472,10 +706,7 @@ fn print_report(
|
||||
println!(" LongMemEval Benchmark — {}", mode.label);
|
||||
println!("=================================================================");
|
||||
println!();
|
||||
println!(
|
||||
"Mode: vector_weight={:.1} / keyword_weight={:.1}",
|
||||
mode.vector_weight, mode.keyword_weight
|
||||
);
|
||||
println!("Mode: {}", describe(mode));
|
||||
println!();
|
||||
println!("Scoring target: RETRIEVAL RECALL (did the gold memory land in top-k).");
|
||||
println!(" No answer is generated or scored. This is NOT the official");
|
||||
@@ -538,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!(
|
||||
@@ -602,10 +851,7 @@ fn print_report(
|
||||
println!("```json");
|
||||
println!("{{");
|
||||
println!(" \"benchmark\": \"longmemeval\",");
|
||||
println!(
|
||||
" \"mode\": \"vector_{:.1}_keyword_{:.1}\",",
|
||||
mode.vector_weight, mode.keyword_weight
|
||||
);
|
||||
println!(" \"mode\": \"{}\",", describe(mode));
|
||||
println!(" \"dataset_variant\": \"{}\",", profile.variant());
|
||||
println!(" \"scoring_target\": \"retrieval_recall\",");
|
||||
println!(" \"k\": 10,");
|
||||
@@ -654,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}",
|
||||
@@ -676,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() {
|
||||
@@ -684,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"));
|
||||
}
|
||||
@@ -702,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."
|
||||
@@ -767,19 +1036,34 @@ 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, VECTOR_ONLY, HYBRID]
|
||||
vec![
|
||||
BM25_ONLY,
|
||||
VECTOR_ONLY,
|
||||
HYBRID,
|
||||
RRF,
|
||||
BM25_STEMMED,
|
||||
HYBRID_STEMMED,
|
||||
hybrid_rerank_metadata_only(),
|
||||
hybrid_rerank_blended(),
|
||||
]
|
||||
}
|
||||
}
|
||||
#[cfg(not(feature = "embeddings"))]
|
||||
{
|
||||
vec![BM25_ONLY]
|
||||
vec![BM25_ONLY, BM25_STEMMED]
|
||||
}
|
||||
} else {
|
||||
if sweep {
|
||||
eprintln!("warning: --sweep needs --embeddings; running BM25 only");
|
||||
}
|
||||
vec![BM25_ONLY]
|
||||
// Stemming is a property of the keyword stage, so it can be compared
|
||||
// without a model.
|
||||
vec![BM25_ONLY, BM25_STEMMED]
|
||||
};
|
||||
|
||||
for (mode_idx, mode) in modes.iter().enumerate() {
|
||||
@@ -882,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);
|
||||
@@ -893,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");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,176 @@
|
||||
//! HDF5 read-path measurement harness: full reads vs. hyperslab selections on
|
||||
//! a chunked 2-D dataset, compressed and uncompressed, plus a contiguous one.
|
||||
//!
|
||||
//! The question it answers for every read-path change: does the cost of a
|
||||
//! selection scale with the *selection*, or with the whole dataset?
|
||||
//!
|
||||
//! ```text
|
||||
//! cargo run --release -p clawhdf5-bench --bin read_harness
|
||||
//! cargo run --release -p clawhdf5-bench --bin read_harness -- --large # 512 MB
|
||||
//! ```
|
||||
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use clawhdf5::{File, FileBuilder};
|
||||
use clawhdf5_format::selection::Selection;
|
||||
|
||||
const CHUNK: u64 = 256;
|
||||
|
||||
struct Layout {
|
||||
name: &'static str,
|
||||
chunked: bool,
|
||||
deflate: bool,
|
||||
}
|
||||
|
||||
const LAYOUTS: [Layout; 3] = [
|
||||
Layout {
|
||||
name: "chunked + deflate",
|
||||
chunked: true,
|
||||
deflate: true,
|
||||
},
|
||||
Layout {
|
||||
name: "chunked",
|
||||
chunked: true,
|
||||
deflate: false,
|
||||
},
|
||||
Layout {
|
||||
name: "contiguous",
|
||||
chunked: false,
|
||||
deflate: false,
|
||||
},
|
||||
];
|
||||
|
||||
/// Smooth-ish, compressible data whose value encodes its position, so a read
|
||||
/// can be verified exactly.
|
||||
fn value(row: u64, col: u64) -> f64 {
|
||||
(row * 100_003 + col) as f64 * 0.5
|
||||
}
|
||||
|
||||
fn write_file(path: &std::path::Path, rows: u64, cols: u64) {
|
||||
let data: Vec<f64> = (0..rows)
|
||||
.flat_map(|r| (0..cols).map(move |c| value(r, c)))
|
||||
.collect();
|
||||
let mut builder = FileBuilder::new();
|
||||
for (i, layout) in LAYOUTS.iter().enumerate() {
|
||||
let ds = builder.create_dataset(&format!("d{i}"));
|
||||
ds.with_f64_data(&data).with_shape(&[rows, cols]);
|
||||
if layout.chunked {
|
||||
ds.with_chunks(&[CHUNK, CHUNK]);
|
||||
}
|
||||
if layout.deflate {
|
||||
ds.with_deflate(4);
|
||||
}
|
||||
}
|
||||
builder.write(path).unwrap();
|
||||
}
|
||||
|
||||
fn median(mut samples: Vec<Duration>) -> Duration {
|
||||
samples.sort();
|
||||
samples[samples.len() / 2]
|
||||
}
|
||||
|
||||
fn time<T>(reps: usize, mut f: impl FnMut() -> T) -> Duration {
|
||||
median(
|
||||
(0..reps)
|
||||
.map(|_| {
|
||||
let t = Instant::now();
|
||||
std::hint::black_box(f());
|
||||
t.elapsed()
|
||||
})
|
||||
.collect(),
|
||||
)
|
||||
}
|
||||
|
||||
fn slab(start: [u64; 2], count: [u64; 2]) -> Selection {
|
||||
Selection::Hyperslab {
|
||||
start: start.to_vec(),
|
||||
stride: vec![1, 1],
|
||||
count: count.to_vec(),
|
||||
block: vec![1, 1],
|
||||
}
|
||||
}
|
||||
|
||||
fn main() {
|
||||
let large = std::env::args().any(|a| a == "--large");
|
||||
let (rows, cols) = if large { (8192, 8192) } else { (4096, 2048) };
|
||||
let total_mb = (rows * cols * 8) as f64 / (1 << 20) as f64;
|
||||
if cfg!(debug_assertions) {
|
||||
eprintln!("warning: debug build — numbers are meaningless. Use --release.");
|
||||
}
|
||||
|
||||
let dir = tempfile::TempDir::new().unwrap();
|
||||
let path = dir.path().join("read_harness.h5");
|
||||
write_file(&path, rows, cols);
|
||||
let file_mb = std::fs::metadata(&path).unwrap().len() as f64 / (1 << 20) as f64;
|
||||
|
||||
println!("## Read harness");
|
||||
println!(
|
||||
"\n{rows} x {cols} f64 ({total_mb:.0} MB per dataset), chunks {CHUNK} x {CHUNK}, file {file_mb:.0} MB\n"
|
||||
);
|
||||
|
||||
// (label, selection, elements selected)
|
||||
let selections: Vec<(&str, Selection, u64)> = vec![
|
||||
(
|
||||
"64 x 64 window (1 chunk)",
|
||||
slab([300, 300], [64, 64]),
|
||||
64 * 64,
|
||||
),
|
||||
(
|
||||
"512 x 512 window (4-9 chunks)",
|
||||
slab([1000, 700], [512, 512]),
|
||||
512 * 512,
|
||||
),
|
||||
("one row", slab([rows / 2, 0], [1, cols]), cols),
|
||||
("one column", slab([0, cols / 2], [rows, 1]), rows),
|
||||
];
|
||||
|
||||
println!("| layout | read | selected | time ms | MB/s of selection | vs full read |");
|
||||
println!("|---|---|---:|---:|---:|---:|");
|
||||
for (i, layout) in LAYOUTS.iter().enumerate() {
|
||||
// Fresh handle per layout so one dataset's cached chunks don't help
|
||||
// (or evict) another's.
|
||||
let file = File::open(&path).unwrap();
|
||||
let ds = file.dataset(&format!("d{i}")).unwrap();
|
||||
|
||||
let full_cold = time(1, || ds.read_f64().unwrap());
|
||||
let full = time(3, || ds.read_f64().unwrap());
|
||||
println!(
|
||||
"| {} | full (first) | {total_mb:.0} MB | {:.1} | {:.0} | |",
|
||||
layout.name,
|
||||
full_cold.as_secs_f64() * 1e3,
|
||||
total_mb / full_cold.as_secs_f64()
|
||||
);
|
||||
println!(
|
||||
"| {} | full (repeat) | {total_mb:.0} MB | {:.1} | {:.0} | 1.00x |",
|
||||
layout.name,
|
||||
full.as_secs_f64() * 1e3,
|
||||
total_mb / full.as_secs_f64()
|
||||
);
|
||||
|
||||
for (label, selection, elements) in &selections {
|
||||
// A fresh handle again: measure the selection on its own, not
|
||||
// served from chunks the full read just cached.
|
||||
let file = File::open(&path).unwrap();
|
||||
let ds = file.dataset(&format!("d{i}")).unwrap();
|
||||
let got = ds.read_f64_selection(selection).unwrap();
|
||||
assert_eq!(got.len() as u64, *elements, "{label}");
|
||||
if let Selection::Hyperslab { start, .. } = selection {
|
||||
assert_eq!(got[0], value(start[0], start[1]), "{label}: wrong data");
|
||||
}
|
||||
let took = time(5, || {
|
||||
let file = File::open(&path).unwrap();
|
||||
let ds = file.dataset(&format!("d{i}")).unwrap();
|
||||
ds.read_f64_selection(selection).unwrap()
|
||||
});
|
||||
let mb = (*elements * 8) as f64 / (1 << 20) as f64;
|
||||
println!(
|
||||
"| {} | {label} | {:.2} MB | {:.2} | {:.0} | {:.3}x |",
|
||||
layout.name,
|
||||
mb,
|
||||
took.as_secs_f64() * 1e3,
|
||||
mb / took.as_secs_f64(),
|
||||
took.as_secs_f64() / full_cold.as_secs_f64()
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,679 @@
|
||||
//! Search measurement harness: recall vs. speed for the HNSW index, and
|
||||
//! end-to-end `hybrid_search` latency as the store grows.
|
||||
//!
|
||||
//! Every search-path change should be justified by a before/after run of this
|
||||
//! binary. It reports, for deterministic synthetic data:
|
||||
//!
|
||||
//! * **ANN** — index build time, and for each `ef`: recall@10 against an exact
|
||||
//! brute-force scan, queries/second, and p50/p99 latency.
|
||||
//! * **End to end** — `HDF5Memory`: ingest time, checkpoint time, `open()`
|
||||
//! time, the one-off cold index build (first query ever), the first query
|
||||
//! after a reopen, and steady-state `hybrid_search` p50/p99 at each size.
|
||||
//!
|
||||
//! Data is *clustered* (points = cluster centre + noise, unit-normalised), not
|
||||
//! uniform: uniform random high-dimensional vectors are nearly equidistant,
|
||||
//! which makes recall numbers meaningless and is nothing like embeddings.
|
||||
//!
|
||||
//! ```text
|
||||
//! cargo run --release -p clawhdf5-bench --bin search_harness # 1K, 10K
|
||||
//! 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
|
||||
//! ```
|
||||
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use clawhdf5_agent::{AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry};
|
||||
use clawhdf5_ann::{DistanceMetric, HnswIndex, Storage};
|
||||
|
||||
const DIM: usize = 384;
|
||||
const K: usize = 10;
|
||||
const N_QUERIES: usize = 200;
|
||||
const HNSW_M: usize = 16;
|
||||
const HNSW_EF_CONSTRUCTION: usize = 64;
|
||||
const EF_VALUES: [usize; 5] = [16, 32, 64, 128, 256];
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Deterministic data
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
struct Rng(u64);
|
||||
|
||||
impl Rng {
|
||||
fn next_u64(&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)
|
||||
}
|
||||
|
||||
/// Uniform in [0, 1).
|
||||
fn unit(&mut self) -> f32 {
|
||||
(self.next_u64() >> 40) as f32 / (1u64 << 24) as f32
|
||||
}
|
||||
|
||||
/// Approximately standard normal (sum of uniforms).
|
||||
fn gauss(&mut self) -> f32 {
|
||||
let sum: f32 = (0..6).map(|_| self.unit()).sum();
|
||||
(sum - 3.0) * std::f32::consts::SQRT_2
|
||||
}
|
||||
|
||||
fn below(&mut self, n: usize) -> usize {
|
||||
(self.next_u64() % n as u64) as usize
|
||||
}
|
||||
}
|
||||
|
||||
fn normalize(v: &mut [f32]) {
|
||||
let norm = v.iter().map(|x| x * x).sum::<f32>().sqrt();
|
||||
if norm > 0.0 {
|
||||
v.iter_mut().for_each(|x| *x /= norm);
|
||||
}
|
||||
}
|
||||
|
||||
struct Dataset {
|
||||
vectors: Vec<Vec<f32>>,
|
||||
queries: Vec<Vec<f32>>,
|
||||
/// Cluster id of each vector (used to give records topical text).
|
||||
cluster_of: Vec<usize>,
|
||||
query_cluster: Vec<usize>,
|
||||
}
|
||||
|
||||
/// `--uniform`: isotropic random unit vectors instead of clusters. Not a
|
||||
/// realistic workload, but a useful second distribution — a recall problem
|
||||
/// 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);
|
||||
|
||||
/// `--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) {
|
||||
let random_unit = |rng: &mut Rng| {
|
||||
let mut v: Vec<f32> = (0..DIM).map(|_| rng.gauss()).collect();
|
||||
normalize(&mut v);
|
||||
v
|
||||
};
|
||||
return Dataset {
|
||||
vectors: (0..n).map(|_| random_unit(&mut rng)).collect(),
|
||||
queries: (0..N_QUERIES).map(|_| random_unit(&mut rng)).collect(),
|
||||
cluster_of: vec![0; n],
|
||||
query_cluster: vec![0; N_QUERIES],
|
||||
};
|
||||
}
|
||||
let n_clusters = (n / 100).clamp(8, 512);
|
||||
let centres: Vec<Vec<f32>> = (0..n_clusters)
|
||||
.map(|_| {
|
||||
let mut c: Vec<f32> = (0..DIM).map(|_| rng.gauss()).collect();
|
||||
normalize(&mut c);
|
||||
c
|
||||
})
|
||||
.collect();
|
||||
let point = |rng: &mut Rng, cluster: usize| {
|
||||
// Noise comparable to the centre's per-dimension magnitude, so
|
||||
// clusters overlap and the nearest neighbours are non-trivial.
|
||||
let scale = 0.6 / (DIM as f32).sqrt();
|
||||
let mut v: Vec<f32> = centres[cluster]
|
||||
.iter()
|
||||
.map(|c| c + rng.gauss() * scale)
|
||||
.collect();
|
||||
normalize(&mut v);
|
||||
v
|
||||
};
|
||||
let mut vectors = Vec::with_capacity(n);
|
||||
let mut cluster_of = Vec::with_capacity(n);
|
||||
for _ in 0..n {
|
||||
let c = rng.below(n_clusters);
|
||||
vectors.push(point(&mut rng, c));
|
||||
cluster_of.push(c);
|
||||
}
|
||||
let mut queries = Vec::with_capacity(N_QUERIES);
|
||||
let mut query_cluster = Vec::with_capacity(N_QUERIES);
|
||||
for _ in 0..N_QUERIES {
|
||||
let c = rng.below(n_clusters);
|
||||
queries.push(point(&mut rng, c));
|
||||
query_cluster.push(c);
|
||||
}
|
||||
Dataset {
|
||||
vectors,
|
||||
queries,
|
||||
cluster_of,
|
||||
query_cluster,
|
||||
}
|
||||
}
|
||||
|
||||
const WORDS: &[&str] = &[
|
||||
"deploy", "latency", "cache", "schema", "index", "vector", "memory", "agent", "kernel",
|
||||
"buffer", "socket", "thread", "tensor", "gradient", "ledger", "invoice", "meeting", "roadmap",
|
||||
"customer", "contract", "sensor", "orbit", "protein", "genome", "harbor", "bridge", "engine",
|
||||
"battery", "harvest", "weather", "museum", "recipe",
|
||||
];
|
||||
|
||||
/// Text whose vocabulary is biased by cluster, so keyword and vector signals
|
||||
/// agree the way they do for real embedded text.
|
||||
fn text_for(cluster: usize, i: usize, rng: &mut Rng) -> String {
|
||||
let topic = [
|
||||
WORDS[cluster % WORDS.len()],
|
||||
WORDS[(cluster / 7 + 3) % WORDS.len()],
|
||||
];
|
||||
let mut words = Vec::with_capacity(14);
|
||||
for j in 0..14 {
|
||||
if j % 3 == 0 {
|
||||
words.push(topic[j / 3 % 2]);
|
||||
} else {
|
||||
words.push(WORDS[rng.below(WORDS.len())]);
|
||||
}
|
||||
}
|
||||
format!("record {i}: {}", words.join(" "))
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// 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
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, v)| (i, v.iter().zip(query).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.truncate(k);
|
||||
scored.into_iter().map(|(i, _)| i).collect()
|
||||
}
|
||||
|
||||
struct Latency {
|
||||
p50: Duration,
|
||||
p99: Duration,
|
||||
qps: f64,
|
||||
}
|
||||
|
||||
fn summarize(mut samples: Vec<Duration>) -> Latency {
|
||||
samples.sort();
|
||||
let total: Duration = samples.iter().sum();
|
||||
let at = |q: f64| samples[((samples.len() - 1) as f64 * q).round() as usize];
|
||||
Latency {
|
||||
p50: at(0.50),
|
||||
p99: at(0.99),
|
||||
qps: samples.len() as f64 / total.as_secs_f64(),
|
||||
}
|
||||
}
|
||||
|
||||
/// 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);
|
||||
|
||||
// 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() {
|
||||
LIVE_BYTES.fetch_add(layout.size() as i64, std::sync::atomic::Ordering::Relaxed);
|
||||
}
|
||||
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() {
|
||||
LIVE_BYTES.fetch_add(
|
||||
new_size as i64 - layout.size() as i64,
|
||||
std::sync::atomic::Ordering::Relaxed,
|
||||
);
|
||||
}
|
||||
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
|
||||
}
|
||||
|
||||
fn mib(bytes: u64) -> f64 {
|
||||
bytes as f64 / (1 << 20) as f64
|
||||
}
|
||||
|
||||
fn micros(d: Duration) -> f64 {
|
||||
d.as_secs_f64() * 1e6
|
||||
}
|
||||
|
||||
fn millis(d: Duration) -> f64 {
|
||||
d.as_secs_f64() * 1e3
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// ANN: recall vs speed
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
fn bench_ann(n: usize, json: &mut Vec<serde_json::Value>) {
|
||||
let data = make_dataset(n, 0xA11CE ^ n as u64);
|
||||
let truth: Vec<Vec<usize>> = data
|
||||
.queries
|
||||
.iter()
|
||||
.map(|q| exact_top_k(&data.vectors, q, K))
|
||||
.collect();
|
||||
|
||||
let started = Instant::now();
|
||||
let index = HnswIndex::build_with(
|
||||
&data.vectors,
|
||||
HNSW_M,
|
||||
HNSW_EF_CONSTRUCTION,
|
||||
DistanceMetric::Cosine,
|
||||
storage(),
|
||||
);
|
||||
let build = started.elapsed();
|
||||
|
||||
// Exact scan baseline, for scale.
|
||||
let exact = summarize(
|
||||
data.queries
|
||||
.iter()
|
||||
.map(|q| {
|
||||
let t = Instant::now();
|
||||
std::hint::black_box(exact_top_k(&data.vectors, q, K));
|
||||
t.elapsed()
|
||||
})
|
||||
.collect(),
|
||||
);
|
||||
|
||||
println!(
|
||||
"\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",
|
||||
millis(build),
|
||||
n as f64 / build.as_secs_f64(),
|
||||
exact.qps,
|
||||
micros(exact.p50)
|
||||
);
|
||||
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 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();
|
||||
}
|
||||
let recall = hits as f64 / (K * data.queries.len()) as f64;
|
||||
let lat = summarize(samples);
|
||||
println!(
|
||||
"| {ef} | {recall:.4} | {:.0} | {:.0} | {:.0} |",
|
||||
lat.qps,
|
||||
micros(lat.p50),
|
||||
micros(lat.p99)
|
||||
);
|
||||
json.push(serde_json::json!({
|
||||
"bench": "hnsw", "n": n, "ef": ef, "recall_at_10": recall,
|
||||
"qps": lat.qps, "p50_us": micros(lat.p50), "p99_us": micros(lat.p99),
|
||||
"build_ms": millis(build),
|
||||
}));
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// End to end: HDF5Memory::hybrid_search
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
fn bench_end_to_end(n: usize, json: &mut Vec<serde_json::Value>) {
|
||||
let data = make_dataset(n, 0xE2E ^ n as u64);
|
||||
let dir = tempfile::TempDir::new().unwrap();
|
||||
let path = dir.path().join("store.h5");
|
||||
let mut rng = Rng(7);
|
||||
|
||||
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 query_texts: Vec<String> = data
|
||||
.query_cluster
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, c)| text_for(*c, i, &mut rng))
|
||||
.collect();
|
||||
|
||||
let mut mem = HDF5Memory::create(MemoryConfig::new(path.clone(), "bench", DIM)).unwrap();
|
||||
let t = Instant::now();
|
||||
mem.save_batch(entries).unwrap();
|
||||
let ingest = t.elapsed();
|
||||
// The very first query builds the vector and keyword indexes from
|
||||
// scratch. It happens once per store, not once per session: the checkpoint
|
||||
// below saves the vector index, so a later `open()` reloads it.
|
||||
let t = Instant::now();
|
||||
std::hint::black_box(mem.hybrid_search(&data.queries[1], &query_texts[1], 0.7, 0.3, K));
|
||||
let cold_build = t.elapsed();
|
||||
|
||||
let t = Instant::now();
|
||||
mem.flush_wal().unwrap();
|
||||
let checkpoint = t.elapsed();
|
||||
drop(mem);
|
||||
|
||||
let t = Instant::now();
|
||||
let mut mem = HDF5Memory::open(&path).unwrap();
|
||||
let open = t.elapsed();
|
||||
|
||||
// The first query after open pays for whatever is rebuilt lazily.
|
||||
let t = Instant::now();
|
||||
std::hint::black_box(mem.hybrid_search(&data.queries[0], &query_texts[0], 0.7, 0.3, K));
|
||||
let first_query = t.elapsed();
|
||||
|
||||
// Fewer steady-state samples at large N: each query is currently O(N).
|
||||
let samples_wanted = if n >= 100_000 { 20 } else { N_QUERIES.min(100) };
|
||||
let steady = summarize(
|
||||
(0..samples_wanted)
|
||||
.map(|i| {
|
||||
let t = Instant::now();
|
||||
std::hint::black_box(mem.hybrid_search(
|
||||
&data.queries[i % N_QUERIES],
|
||||
&query_texts[i % N_QUERIES],
|
||||
0.7,
|
||||
0.3,
|
||||
K,
|
||||
));
|
||||
t.elapsed()
|
||||
})
|
||||
.collect(),
|
||||
);
|
||||
|
||||
println!(
|
||||
"| {n} | {:.0} | {:.0} | {:.1} | {:.1} | {:.1} | {:.2} | {:.2} | {:.1} |",
|
||||
millis(ingest),
|
||||
millis(cold_build),
|
||||
millis(checkpoint),
|
||||
millis(open),
|
||||
millis(first_query),
|
||||
millis(steady.p50),
|
||||
millis(steady.p99),
|
||||
steady.qps
|
||||
);
|
||||
json.push(serde_json::json!({
|
||||
"bench": "hybrid_search", "n": n,
|
||||
"ingest_ms": millis(ingest), "cold_index_build_ms": millis(cold_build),
|
||||
"checkpoint_ms": millis(checkpoint),
|
||||
"open_ms": millis(open), "first_query_ms": millis(first_query),
|
||||
"p50_ms": millis(steady.p50), "p99_ms": millis(steady.p99), "qps": steady.qps,
|
||||
}));
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Fusion study: does capping the keyword candidate pool change the ranking?
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// `hybrid_search` min-max normalises each signal over the candidates it is
|
||||
/// given. The vector stage supplies a pool of `max(8k, 64)`; the keyword stage
|
||||
/// supplies *every* matching record, which is what now dominates query time.
|
||||
/// This compares the current fusion with one whose keyword stage is capped to
|
||||
/// a pool, reporting how often the final top-k agree and what each costs.
|
||||
fn fusion_study(n: usize) {
|
||||
use clawhdf5_agent::bm25::BM25Index;
|
||||
use clawhdf5_agent::hybrid::merge_vector_keyword;
|
||||
|
||||
let data = make_dataset(n, 0xE2E ^ n as u64);
|
||||
let mut rng = Rng(7);
|
||||
let texts: Vec<String> = (0..n)
|
||||
.map(|i| text_for(data.cluster_of[i], i, &mut rng))
|
||||
.collect();
|
||||
let query_texts: Vec<String> = data
|
||||
.query_cluster
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, c)| text_for(*c, i, &mut rng))
|
||||
.collect();
|
||||
let bm25 = BM25Index::build(&texts, &vec![0u8; n]);
|
||||
let index = HnswIndex::build_with(
|
||||
&data.vectors,
|
||||
HNSW_M,
|
||||
HNSW_EF_CONSTRUCTION,
|
||||
DistanceMetric::Cosine,
|
||||
storage(),
|
||||
);
|
||||
|
||||
let vec_pool = (K * 8).max(64);
|
||||
println!("\n### Fusion study, N = {n} (k = {K}, weights 0.7 / 0.3, vector pool {vec_pool})\n");
|
||||
println!(
|
||||
"| keyword pool | top-{K} overlap vs full | identical top-{K} | same #1 | keyword+merge µs |"
|
||||
);
|
||||
println!("|---:|---:|---:|---:|---:|");
|
||||
|
||||
let fuse = |q: usize, kw_pool: usize| -> (Vec<usize>, Duration) {
|
||||
let vec_scores: Vec<(usize, f32)> = index
|
||||
.search(&data.queries[q], vec_pool, vec_pool)
|
||||
.into_iter()
|
||||
.map(|(id, d)| (id, 1.0 - d))
|
||||
.collect();
|
||||
let t = Instant::now();
|
||||
let kw = bm25.search(&query_texts[q], kw_pool);
|
||||
let merged = merge_vector_keyword(vec_scores, kw, 0.7, 0.3, K);
|
||||
let took = t.elapsed();
|
||||
(merged.into_iter().map(|(id, _)| id).collect(), took)
|
||||
};
|
||||
|
||||
let full: Vec<(Vec<usize>, Duration)> = (0..N_QUERIES).map(|q| fuse(q, n)).collect();
|
||||
let full_time: Duration = full.iter().map(|f| f.1).sum();
|
||||
println!(
|
||||
"| all ({n}) | 1.0000 | 100.0% | 100.0% | {:.0} |",
|
||||
micros(full_time) / N_QUERIES as f64
|
||||
);
|
||||
for pool in [vec_pool, vec_pool * 4, 1000] {
|
||||
if pool >= n {
|
||||
continue;
|
||||
}
|
||||
let (mut overlap, mut identical, mut same_first) = (0usize, 0usize, 0usize);
|
||||
let mut time = Duration::ZERO;
|
||||
for (q, (want, _)) in full.iter().enumerate() {
|
||||
let (got, took) = fuse(q, pool);
|
||||
time += took;
|
||||
overlap += got.iter().filter(|id| want.contains(id)).count();
|
||||
identical += usize::from(&got == want);
|
||||
same_first += usize::from(got.first() == want.first());
|
||||
}
|
||||
println!(
|
||||
"| {pool} | {:.4} | {:.1}% | {:.1}% | {:.0} |",
|
||||
overlap as f64 / (K * N_QUERIES) as f64,
|
||||
100.0 * identical as f64 / N_QUERIES as f64,
|
||||
100.0 * same_first as f64 / N_QUERIES as f64,
|
||||
micros(time) / N_QUERIES as f64
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// 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();
|
||||
let reopened = HDF5Memory::open(&path).unwrap();
|
||||
let after_open = heap_bytes();
|
||||
let loaded = after_open.saturating_sub(before_open);
|
||||
drop(reopened);
|
||||
|
||||
let raw = (n * DIM * 4) as u64;
|
||||
println!(
|
||||
"| {n} | {:.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),
|
||||
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");
|
||||
let ann_only = args.iter().any(|a| a == "--ann-only");
|
||||
if args.iter().any(|a| a == "--fusion-study") {
|
||||
for &n in if full {
|
||||
&[10_000, 100_000][..]
|
||||
} else {
|
||||
&[10_000][..]
|
||||
} {
|
||||
fusion_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)");
|
||||
}
|
||||
let json_path = args
|
||||
.iter()
|
||||
.position(|a| a == "--json")
|
||||
.and_then(|i| args.get(i + 1))
|
||||
.cloned();
|
||||
let sizes: &[usize] = if full {
|
||||
&[1_000, 10_000, 100_000]
|
||||
} else {
|
||||
&[1_000, 10_000]
|
||||
};
|
||||
|
||||
if cfg!(debug_assertions) {
|
||||
eprintln!("warning: debug build — numbers are meaningless. Use --release.");
|
||||
}
|
||||
|
||||
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 | reopened / raw |"
|
||||
);
|
||||
println!("|---:|---:|---:|---:|---:|---:|---:|");
|
||||
for &n in sizes {
|
||||
bench_footprint(n);
|
||||
}
|
||||
return;
|
||||
}
|
||||
// `--e2e-only` skips the index benchmarks, so the end-to-end section runs
|
||||
// in a process that has not already spun up a thread pool.
|
||||
if !args.iter().any(|a| a == "--e2e-only") {
|
||||
for &n in sizes {
|
||||
bench_ann(n, &mut json);
|
||||
}
|
||||
}
|
||||
|
||||
if ann_only {
|
||||
return;
|
||||
}
|
||||
println!("\n### End to end: `HDF5Memory::hybrid_search` (k = {K}, weights 0.7 / 0.3)\n");
|
||||
println!(
|
||||
"| N | ingest ms | cold index build ms | checkpoint ms | open ms | first query after open ms | p50 ms | p99 ms | QPS |"
|
||||
);
|
||||
println!("|---:|---:|---:|---:|---:|---:|---:|---:|---:|");
|
||||
for &n in sizes {
|
||||
bench_end_to_end(n, &mut json);
|
||||
}
|
||||
|
||||
if let Some(path) = json_path {
|
||||
std::fs::write(&path, serde_json::to_string_pretty(&json).unwrap()).unwrap();
|
||||
eprintln!("wrote {path}");
|
||||
}
|
||||
}
|
||||
@@ -1,6 +1,6 @@
|
||||
[package]
|
||||
name = "clawhdf5-cli"
|
||||
version = "2.3.0"
|
||||
version = "2.6.0"
|
||||
edition = "2024"
|
||||
license = "MIT"
|
||||
description = "CLI for clawhdf5 agent memory — create, save, search, recall, stats"
|
||||
@@ -14,7 +14,7 @@ name = "clawhdf5"
|
||||
path = "src/main.rs"
|
||||
|
||||
[dependencies]
|
||||
clawhdf5-agent = { path = "../clawhdf5-agent", version = "2.3.0" }
|
||||
clawhdf5-agent = { path = "../clawhdf5-agent", version = "2.6.0" }
|
||||
clap = { version = "4", features = ["derive", "env"] }
|
||||
serde_json = "1"
|
||||
serde = { workspace = true }
|
||||
|
||||
@@ -28,6 +28,10 @@ enum Commands {
|
||||
/// Enable write-ahead log
|
||||
#[arg(long)]
|
||||
wal: bool,
|
||||
/// Store the vector index's copy of the embeddings as int8, roughly
|
||||
/// halving a loaded store's memory at about 13% fewer queries/second
|
||||
#[arg(long)]
|
||||
quantized_index: bool,
|
||||
},
|
||||
/// Save a memory entry (reads JSON from stdin or --json)
|
||||
Save {
|
||||
@@ -88,9 +92,15 @@ 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,
|
||||
quantized_index,
|
||||
} => {
|
||||
let mut config = MemoryConfig::new(cli.path.clone(), &agent_id, dim);
|
||||
config.wal_enabled = wal;
|
||||
config.quantized_index = quantized_index;
|
||||
let mem = HDF5Memory::create(config)?;
|
||||
let j = serde_json::json!({
|
||||
"status": "created",
|
||||
@@ -98,6 +108,7 @@ fn run(cli: Cli) -> Result<(), Box<dyn std::error::Error>> {
|
||||
"agent_id": agent_id,
|
||||
"embedding_dim": dim,
|
||||
"wal_enabled": wal,
|
||||
"quantized_index": quantized_index,
|
||||
"count": mem.count(),
|
||||
});
|
||||
println!("{}", serde_json::to_string_pretty(&j)?);
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
[package]
|
||||
name = "clawhdf5-derive"
|
||||
version = "2.3.0"
|
||||
version = "2.6.0"
|
||||
edition = "2024"
|
||||
description = "Derive macros for rustyhdf5 HDF5 traits"
|
||||
license = "MIT"
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
[package]
|
||||
name = "clawhdf5-filters"
|
||||
version = "2.3.0"
|
||||
version = "2.6.0"
|
||||
edition = "2024"
|
||||
description = "Filter and compression pipeline for clawhdf5"
|
||||
license = "MIT"
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
[package]
|
||||
name = "clawhdf5-format"
|
||||
version = "2.3.0"
|
||||
version = "2.6.0"
|
||||
edition = "2024"
|
||||
description = "Pure-Rust HDF5 binary format parsing and writing — no C dependencies"
|
||||
license = "MIT"
|
||||
@@ -25,7 +25,7 @@ pco = { version = "1.0", optional = true }
|
||||
[dev-dependencies]
|
||||
serde_json = "1"
|
||||
criterion = { workspace = true }
|
||||
clawhdf5-derive = { path = "../clawhdf5-derive", version = "2.3.0" }
|
||||
clawhdf5-derive = { path = "../clawhdf5-derive", version = "2.6.0" }
|
||||
|
||||
[[bench]]
|
||||
name = "bench"
|
||||
|
||||
@@ -374,6 +374,11 @@ impl ChunkCache {
|
||||
|
||||
// ----- Index operations -----
|
||||
|
||||
/// The most decompressed bytes this cache will hold.
|
||||
pub fn max_bytes(&self) -> usize {
|
||||
self.inner.lock().map(|g| g.max_bytes).unwrap_or(0)
|
||||
}
|
||||
|
||||
/// Bind the cache to the dataset at chunk-index address `addr`.
|
||||
///
|
||||
/// The cache is shared per file across all of its datasets. If the cache
|
||||
|
||||
@@ -165,12 +165,29 @@ pub(crate) fn checked_chunk_byte_len(
|
||||
/// process when the allocation fails; a size taken from the file must surface
|
||||
/// as an error instead.
|
||||
pub(crate) fn alloc_output(len: usize) -> Result<Vec<u8>, FormatError> {
|
||||
let mut out = Vec::new();
|
||||
out.try_reserve_exact(len).map_err(|_| {
|
||||
FormatError::Overflow(format!("cannot allocate {len} bytes for dataset output"))
|
||||
})?;
|
||||
out.resize(len, 0);
|
||||
Ok(out)
|
||||
if len == 0 {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let failed =
|
||||
|| FormatError::Overflow(format!("cannot allocate {len} bytes for dataset output"));
|
||||
let layout = core::alloc::Layout::array::<u8>(len).map_err(|_| failed())?;
|
||||
// Ask the allocator for zeroed memory instead of reserving and then
|
||||
// writing zeros: for a large buffer the OS hands out already-zero pages
|
||||
// lazily, where an explicit fill touches every page up front — and most of
|
||||
// the buffer is about to be overwritten with chunk data anyway.
|
||||
//
|
||||
// SAFETY (both arms): `layout` has non-zero size (len > 0) and alignment 1.
|
||||
#[cfg(feature = "std")]
|
||||
let ptr = unsafe { std::alloc::alloc_zeroed(layout) };
|
||||
#[cfg(not(feature = "std"))]
|
||||
let ptr = unsafe { alloc::alloc::alloc_zeroed(layout) };
|
||||
if ptr.is_null() {
|
||||
return Err(failed());
|
||||
}
|
||||
// SAFETY: `ptr` came from the global allocator with the layout of
|
||||
// `[u8; len]`, which is exactly what `Vec<u8>` with capacity `len` frees;
|
||||
// all `len` bytes are initialised (zero).
|
||||
Ok(unsafe { Vec::from_raw_parts(ptr, len, len) })
|
||||
}
|
||||
|
||||
fn read_offset(data: &[u8], pos: usize, size: u8) -> Result<u64, FormatError> {
|
||||
@@ -362,6 +379,116 @@ pub fn generate_implicit_chunks(
|
||||
}
|
||||
|
||||
/// Read a chunked dataset, decompressing chunks as needed.
|
||||
/// Chunks decompressed together before being copied out, bounding the extra
|
||||
/// memory a parallel full read holds at once.
|
||||
const DECODE_BATCH: usize = 128;
|
||||
|
||||
/// B-tree v2 record types used for chunk indexing.
|
||||
const BT2_CHUNK_UNFILTERED: u8 = 10;
|
||||
const BT2_CHUNK_FILTERED: u8 = 11;
|
||||
|
||||
/// Chunks indexed by a version-2 B-tree (layout v4, index type 5).
|
||||
///
|
||||
/// Record layouts (all little endian):
|
||||
/// * type 10, unfiltered: address, then one 8-byte *scaled* offset per
|
||||
/// dimension (offset / chunk dimension);
|
||||
/// * type 11, filtered: address, stored chunk size (a variable number of
|
||||
/// bytes), 4-byte filter mask, then the scaled offsets.
|
||||
///
|
||||
/// The width of the stored-size field depends on the largest possible chunk;
|
||||
/// rather than re-derive the library's formula it is taken from the record
|
||||
/// size the tree header declares, which is what actually governs the bytes.
|
||||
fn read_btree_v2_chunks(
|
||||
file_data: &[u8],
|
||||
addr: u64,
|
||||
chunk_dims: &[usize],
|
||||
elem_size: usize,
|
||||
offset_size: u8,
|
||||
length_size: u8,
|
||||
) -> Result<Vec<ChunkInfo>, FormatError> {
|
||||
use crate::btree_v2::{BTreeV2Header, collect_btree_v2_records};
|
||||
|
||||
let bad = |what: &str| FormatError::ChunkedReadError(format!("B-tree v2 chunk index: {what}"));
|
||||
let header = BTreeV2Header::parse(file_data, addr as usize, offset_size, length_size)?;
|
||||
let rank = chunk_dims.len();
|
||||
let os = offset_size as usize;
|
||||
let record_size = header.record_size as usize;
|
||||
let size_len = match header.tree_type {
|
||||
BT2_CHUNK_UNFILTERED => {
|
||||
if record_size != os + 8 * rank {
|
||||
return Err(bad("unexpected record size for unfiltered chunks"));
|
||||
}
|
||||
0
|
||||
}
|
||||
BT2_CHUNK_FILTERED => {
|
||||
let fixed = os + 4 + 8 * rank;
|
||||
let size_len = record_size
|
||||
.checked_sub(fixed)
|
||||
.ok_or_else(|| bad("record too small"))?;
|
||||
if !(1..=8).contains(&size_len) {
|
||||
return Err(bad("implausible chunk-size field width"));
|
||||
}
|
||||
size_len
|
||||
}
|
||||
_ => return Err(bad("tree is not a chunk index")),
|
||||
};
|
||||
let unfiltered_bytes = checked_chunk_byte_len(chunk_dims, elem_size)?;
|
||||
let unfiltered_bytes =
|
||||
u32::try_from(unfiltered_bytes).map_err(|_| bad("chunk larger than 4 GiB"))?;
|
||||
|
||||
let records = collect_btree_v2_records(file_data, &header, offset_size, length_size)?;
|
||||
let mut chunks = Vec::with_capacity(records.len());
|
||||
for record in &records {
|
||||
let data = record.data.as_slice();
|
||||
if data.len() < record_size {
|
||||
return Err(bad("truncated record"));
|
||||
}
|
||||
let address = read_offset(data, 0, offset_size)?;
|
||||
let mut pos = os;
|
||||
let (chunk_size, filter_mask) = if size_len == 0 {
|
||||
(unfiltered_bytes, 0)
|
||||
} else {
|
||||
let mut size = 0u64;
|
||||
for (i, &b) in data[pos..pos + size_len].iter().enumerate() {
|
||||
size |= u64::from(b) << (8 * i);
|
||||
}
|
||||
pos += size_len;
|
||||
let mask = u32::from_le_bytes([data[pos], data[pos + 1], data[pos + 2], data[pos + 3]]);
|
||||
pos += 4;
|
||||
(
|
||||
u32::try_from(size).map_err(|_| bad("stored chunk larger than 4 GiB"))?,
|
||||
mask,
|
||||
)
|
||||
};
|
||||
let mut offsets = Vec::with_capacity(rank);
|
||||
for &dim in chunk_dims {
|
||||
let scaled = u64::from_le_bytes([
|
||||
data[pos],
|
||||
data[pos + 1],
|
||||
data[pos + 2],
|
||||
data[pos + 3],
|
||||
data[pos + 4],
|
||||
data[pos + 5],
|
||||
data[pos + 6],
|
||||
data[pos + 7],
|
||||
]);
|
||||
pos += 8;
|
||||
offsets.push(
|
||||
scaled
|
||||
.checked_mul(dim as u64)
|
||||
.ok_or_else(|| bad("chunk offset overflows"))?,
|
||||
);
|
||||
}
|
||||
chunks.push(ChunkInfo {
|
||||
chunk_size,
|
||||
filter_mask,
|
||||
offsets,
|
||||
address,
|
||||
});
|
||||
}
|
||||
Ok(chunks)
|
||||
}
|
||||
|
||||
/// Every allocated chunk of a chunked dataset, for any supported chunk index,
|
||||
/// plus the spatial chunk dimensions. Chunks the file never allocated (sparse
|
||||
/// datasets) are simply absent from the list.
|
||||
@@ -487,6 +614,18 @@ pub fn list_chunks(
|
||||
length_size,
|
||||
)?
|
||||
}
|
||||
(4, Some(5)) => {
|
||||
// Version-2 B-tree: what the library uses for a dataset with two
|
||||
// or more unlimited dimensions.
|
||||
read_btree_v2_chunks(
|
||||
file_data,
|
||||
addr,
|
||||
&chunk_dims,
|
||||
elem_size,
|
||||
offset_size,
|
||||
length_size,
|
||||
)?
|
||||
}
|
||||
(v, idx) => {
|
||||
return Err(FormatError::ChunkedReadError(format!(
|
||||
"unsupported chunked layout version={v}, index_type={idx:?}"
|
||||
@@ -629,29 +768,12 @@ pub fn read_chunked_data_cached(
|
||||
length_size: u8,
|
||||
cache: &ChunkCache,
|
||||
) -> Result<Vec<u8>, FormatError> {
|
||||
let (
|
||||
chunk_dimensions,
|
||||
version,
|
||||
chunk_index_type,
|
||||
addr_opt,
|
||||
single_filtered_size,
|
||||
single_filter_mask,
|
||||
) = match layout {
|
||||
let (chunk_dimensions, addr_opt) = match layout {
|
||||
DataLayout::Chunked {
|
||||
chunk_dimensions,
|
||||
btree_address,
|
||||
version,
|
||||
chunk_index_type,
|
||||
single_chunk_filtered_size,
|
||||
single_chunk_filter_mask,
|
||||
} => (
|
||||
chunk_dimensions,
|
||||
*version,
|
||||
*chunk_index_type,
|
||||
*btree_address,
|
||||
*single_chunk_filtered_size,
|
||||
*single_chunk_filter_mask,
|
||||
),
|
||||
..
|
||||
} => (chunk_dimensions, *btree_address),
|
||||
_ => {
|
||||
return Err(FormatError::ChunkedReadError(
|
||||
"expected chunked layout".into(),
|
||||
@@ -688,69 +810,14 @@ pub fn read_chunked_data_cached(
|
||||
|
||||
// Populate chunk index on first access
|
||||
if !cache.has_index() {
|
||||
let chunks = match (version, chunk_index_type) {
|
||||
(3, _) => collect_chunk_info(file_data, addr, ndims, offset_size, length_size)?,
|
||||
(4, Some(1)) => {
|
||||
let chunk_byte_size = checked_chunk_byte_len(&chunk_dims, elem_size)?;
|
||||
let (csize, fmask) = if let Some(fs) = single_filtered_size {
|
||||
(fs as u32, single_filter_mask.unwrap_or(0))
|
||||
} else {
|
||||
(chunk_byte_size as u32, 0)
|
||||
};
|
||||
vec![ChunkInfo {
|
||||
chunk_size: csize,
|
||||
filter_mask: fmask,
|
||||
offsets: vec![0u64; rank],
|
||||
address: addr,
|
||||
}]
|
||||
}
|
||||
(4, Some(2)) => {
|
||||
let spatial_chunk_dims: &[u32] = &chunk_dimensions[..rank];
|
||||
generate_implicit_chunks(
|
||||
addr,
|
||||
&dataspace.dimensions,
|
||||
spatial_chunk_dims,
|
||||
elem_size as u32,
|
||||
)
|
||||
}
|
||||
(4, Some(3)) => {
|
||||
let spatial_chunk_dims: &[u32] = &chunk_dimensions[..rank];
|
||||
let header =
|
||||
FixedArrayHeader::parse(file_data, addr as usize, offset_size, length_size)?;
|
||||
read_fixed_array_chunks(
|
||||
file_data,
|
||||
&header,
|
||||
&dataspace.dimensions,
|
||||
spatial_chunk_dims,
|
||||
elem_size as u32,
|
||||
offset_size,
|
||||
length_size,
|
||||
)?
|
||||
}
|
||||
(4, Some(4)) => {
|
||||
let spatial_chunk_dims: &[u32] = &chunk_dimensions[..rank];
|
||||
let header = ExtensibleArrayHeader::parse(
|
||||
file_data,
|
||||
addr as usize,
|
||||
offset_size,
|
||||
length_size,
|
||||
)?;
|
||||
read_extensible_array_chunks(
|
||||
file_data,
|
||||
&header,
|
||||
&dataspace.dimensions,
|
||||
spatial_chunk_dims,
|
||||
elem_size as u32,
|
||||
offset_size,
|
||||
length_size,
|
||||
)?
|
||||
}
|
||||
(v, idx) => {
|
||||
return Err(FormatError::ChunkedReadError(format!(
|
||||
"unsupported chunked layout version={v}, index_type={idx:?}"
|
||||
)));
|
||||
}
|
||||
};
|
||||
let (chunks, _) = list_chunks(
|
||||
file_data,
|
||||
layout,
|
||||
dataspace,
|
||||
elem_size,
|
||||
offset_size,
|
||||
length_size,
|
||||
)?;
|
||||
cache.populate_index(&chunks, rank);
|
||||
}
|
||||
|
||||
@@ -777,52 +844,86 @@ pub fn read_chunked_data_cached(
|
||||
|
||||
let chunk_total_bytes = checked_chunk_byte_len(&chunk_dims, elem_size)?;
|
||||
|
||||
for chunk_info in &chunks {
|
||||
let coord: Vec<u64> = chunk_info.offsets.iter().take(rank).copied().collect();
|
||||
|
||||
// Try decompressed cache first
|
||||
let decompressed = if let Some(cached) = cache.get_decompressed_aligned(&coord) {
|
||||
cached
|
||||
} else {
|
||||
// Decompress from file
|
||||
let c_addr = chunk_info.address as usize;
|
||||
let size = chunk_info.chunk_size as usize;
|
||||
ensure_len(file_data, c_addr, size)?;
|
||||
let raw_chunk = &file_data[c_addr..c_addr + size];
|
||||
let dec = if let Some(pl) = pipeline {
|
||||
if chunk_info.filter_mask == 0 {
|
||||
decompress_chunk(raw_chunk, pl, chunk_total_bytes, elem_size as u32)?
|
||||
} else {
|
||||
raw_chunk.to_vec()
|
||||
}
|
||||
} else {
|
||||
raw_chunk.to_vec()
|
||||
};
|
||||
cache.put_decompressed(coord, dec)
|
||||
};
|
||||
|
||||
let mut place = |data: &[u8], chunk_info: &ChunkInfo| {
|
||||
if rank == 0 {
|
||||
let copy_len = data.len().min(output.len());
|
||||
output[..copy_len].copy_from_slice(&data[..copy_len]);
|
||||
return;
|
||||
}
|
||||
let chunk_offsets: Vec<usize> = chunk_info
|
||||
.offsets
|
||||
.iter()
|
||||
.take(rank)
|
||||
.map(|&o| o as usize)
|
||||
.collect();
|
||||
copy_chunk_to_output(
|
||||
data,
|
||||
&mut output,
|
||||
&chunk_offsets,
|
||||
&chunk_dims,
|
||||
&ds_dims,
|
||||
&ds_strides,
|
||||
&chunk_strides,
|
||||
elem_size,
|
||||
rank,
|
||||
);
|
||||
};
|
||||
let raw_bytes = |chunk_info: &ChunkInfo| -> Result<&[u8], FormatError> {
|
||||
let c_addr = chunk_info.address as usize;
|
||||
let size = chunk_info.chunk_size as usize;
|
||||
ensure_len(file_data, c_addr, size)?;
|
||||
Ok(&file_data[c_addr..c_addr + size])
|
||||
};
|
||||
|
||||
if rank == 0 {
|
||||
let copy_len = decompressed.len().min(output.len());
|
||||
output[..copy_len].copy_from_slice(&decompressed[..copy_len]);
|
||||
} else {
|
||||
copy_chunk_to_output(
|
||||
&decompressed,
|
||||
&mut output,
|
||||
&chunk_offsets,
|
||||
&chunk_dims,
|
||||
&ds_dims,
|
||||
&ds_strides,
|
||||
&chunk_strides,
|
||||
elem_size,
|
||||
rank,
|
||||
);
|
||||
// Chunks stored as-is (no pipeline, or the filter mask says this chunk
|
||||
// skipped it) are copied straight from the file bytes: they are already in
|
||||
// memory, so routing them through a Vec and then an aligned cache buffer
|
||||
// was two extra copies of the whole dataset for nothing.
|
||||
let stored_raw = |c: &ChunkInfo| pipeline.is_none() || c.filter_mask != 0;
|
||||
let mut misses: Vec<&ChunkInfo> = Vec::new();
|
||||
for chunk_info in &chunks {
|
||||
if stored_raw(chunk_info) {
|
||||
place(raw_bytes(chunk_info)?, chunk_info);
|
||||
continue;
|
||||
}
|
||||
let coord: Vec<u64> = chunk_info.offsets.iter().take(rank).copied().collect();
|
||||
match cache.get_decompressed_aligned(&coord) {
|
||||
Some(cached) => place(&cached, chunk_info),
|
||||
None => misses.push(chunk_info),
|
||||
}
|
||||
}
|
||||
|
||||
// Decompress what the cache didn't have, a bounded batch at a time — in
|
||||
// parallel with the `parallel` feature (this path, the one the facade
|
||||
// uses, was sequential; only the uncached reader was parallel). Chunks are
|
||||
// cached only when the whole dataset fits: pushing a larger dataset
|
||||
// through the cache just evicts each chunk moments after inserting it.
|
||||
let cache_them = total_bytes <= cache.max_bytes();
|
||||
if let Some(pl) = pipeline {
|
||||
let decode = |c: &&ChunkInfo| -> Result<Vec<u8>, FormatError> {
|
||||
decompress_chunk(raw_bytes(c)?, pl, chunk_total_bytes, elem_size as u32)
|
||||
};
|
||||
for batch in misses.chunks(DECODE_BATCH) {
|
||||
#[cfg(feature = "parallel")]
|
||||
let decoded: Vec<Result<Vec<u8>, FormatError>> = if batch.len() >= 4 {
|
||||
use rayon::prelude::*;
|
||||
batch.par_iter().map(decode).collect()
|
||||
} else {
|
||||
batch.iter().map(decode).collect()
|
||||
};
|
||||
#[cfg(not(feature = "parallel"))]
|
||||
let decoded: Vec<Result<Vec<u8>, FormatError>> = batch.iter().map(decode).collect();
|
||||
|
||||
for (chunk_info, data) in batch.iter().zip(decoded) {
|
||||
let data = data?;
|
||||
if cache_them {
|
||||
let coord: Vec<u64> = chunk_info.offsets.iter().take(rank).copied().collect();
|
||||
let cached = cache.put_decompressed(coord, data);
|
||||
place(&cached, chunk_info);
|
||||
} else {
|
||||
place(&data, chunk_info);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -985,29 +1086,12 @@ pub fn read_chunked_data_sweep(
|
||||
cache: &ChunkCache,
|
||||
sweep: &mut SweepContext,
|
||||
) -> Result<Vec<u8>, FormatError> {
|
||||
let (
|
||||
chunk_dimensions,
|
||||
version,
|
||||
chunk_index_type,
|
||||
addr_opt,
|
||||
single_filtered_size,
|
||||
single_filter_mask,
|
||||
) = match layout {
|
||||
let (chunk_dimensions, addr_opt) = match layout {
|
||||
DataLayout::Chunked {
|
||||
chunk_dimensions,
|
||||
btree_address,
|
||||
version,
|
||||
chunk_index_type,
|
||||
single_chunk_filtered_size,
|
||||
single_chunk_filter_mask,
|
||||
} => (
|
||||
chunk_dimensions,
|
||||
*version,
|
||||
*chunk_index_type,
|
||||
*btree_address,
|
||||
*single_chunk_filtered_size,
|
||||
*single_chunk_filter_mask,
|
||||
),
|
||||
..
|
||||
} => (chunk_dimensions, *btree_address),
|
||||
_ => {
|
||||
return Err(FormatError::ChunkedReadError(
|
||||
"expected chunked layout".into(),
|
||||
@@ -1044,69 +1128,14 @@ pub fn read_chunked_data_sweep(
|
||||
|
||||
// Populate chunk index on first access
|
||||
if !cache.has_index() {
|
||||
let chunks = match (version, chunk_index_type) {
|
||||
(3, _) => collect_chunk_info(file_data, addr, ndims, offset_size, length_size)?,
|
||||
(4, Some(1)) => {
|
||||
let chunk_byte_size = checked_chunk_byte_len(&chunk_dims, elem_size)?;
|
||||
let (csize, fmask) = if let Some(fs) = single_filtered_size {
|
||||
(fs as u32, single_filter_mask.unwrap_or(0))
|
||||
} else {
|
||||
(chunk_byte_size as u32, 0)
|
||||
};
|
||||
vec![ChunkInfo {
|
||||
chunk_size: csize,
|
||||
filter_mask: fmask,
|
||||
offsets: vec![0u64; rank],
|
||||
address: addr,
|
||||
}]
|
||||
}
|
||||
(4, Some(2)) => {
|
||||
let spatial_chunk_dims: &[u32] = &chunk_dimensions[..rank];
|
||||
generate_implicit_chunks(
|
||||
addr,
|
||||
&dataspace.dimensions,
|
||||
spatial_chunk_dims,
|
||||
elem_size as u32,
|
||||
)
|
||||
}
|
||||
(4, Some(3)) => {
|
||||
let spatial_chunk_dims: &[u32] = &chunk_dimensions[..rank];
|
||||
let header =
|
||||
FixedArrayHeader::parse(file_data, addr as usize, offset_size, length_size)?;
|
||||
read_fixed_array_chunks(
|
||||
file_data,
|
||||
&header,
|
||||
&dataspace.dimensions,
|
||||
spatial_chunk_dims,
|
||||
elem_size as u32,
|
||||
offset_size,
|
||||
length_size,
|
||||
)?
|
||||
}
|
||||
(4, Some(4)) => {
|
||||
let spatial_chunk_dims: &[u32] = &chunk_dimensions[..rank];
|
||||
let header = ExtensibleArrayHeader::parse(
|
||||
file_data,
|
||||
addr as usize,
|
||||
offset_size,
|
||||
length_size,
|
||||
)?;
|
||||
read_extensible_array_chunks(
|
||||
file_data,
|
||||
&header,
|
||||
&dataspace.dimensions,
|
||||
spatial_chunk_dims,
|
||||
elem_size as u32,
|
||||
offset_size,
|
||||
length_size,
|
||||
)?
|
||||
}
|
||||
(v, idx) => {
|
||||
return Err(FormatError::ChunkedReadError(format!(
|
||||
"unsupported chunked layout version={v}, index_type={idx:?}"
|
||||
)));
|
||||
}
|
||||
};
|
||||
let (chunks, _) = list_chunks(
|
||||
file_data,
|
||||
layout,
|
||||
dataspace,
|
||||
elem_size,
|
||||
offset_size,
|
||||
length_size,
|
||||
)?;
|
||||
cache.populate_index(&chunks, rank);
|
||||
}
|
||||
|
||||
@@ -1211,29 +1240,12 @@ pub fn read_chunked_data_indexed(
|
||||
length_size: u8,
|
||||
cache: &ChunkCache,
|
||||
) -> Result<Vec<u8>, FormatError> {
|
||||
let (
|
||||
chunk_dimensions,
|
||||
version,
|
||||
chunk_index_type,
|
||||
addr_opt,
|
||||
single_filtered_size,
|
||||
single_filter_mask,
|
||||
) = match layout {
|
||||
let (chunk_dimensions, addr_opt) = match layout {
|
||||
DataLayout::Chunked {
|
||||
chunk_dimensions,
|
||||
btree_address,
|
||||
version,
|
||||
chunk_index_type,
|
||||
single_chunk_filtered_size,
|
||||
single_chunk_filter_mask,
|
||||
} => (
|
||||
chunk_dimensions,
|
||||
*version,
|
||||
*chunk_index_type,
|
||||
*btree_address,
|
||||
*single_chunk_filtered_size,
|
||||
*single_chunk_filter_mask,
|
||||
),
|
||||
..
|
||||
} => (chunk_dimensions, *btree_address),
|
||||
_ => {
|
||||
return Err(FormatError::ChunkedReadError(
|
||||
"expected chunked layout".into(),
|
||||
@@ -1270,69 +1282,14 @@ pub fn read_chunked_data_indexed(
|
||||
|
||||
// Build chunk index on first access
|
||||
if !cache.has_chunk_index() {
|
||||
let chunks = match (version, chunk_index_type) {
|
||||
(3, _) => collect_chunk_info(file_data, addr, ndims, offset_size, length_size)?,
|
||||
(4, Some(1)) => {
|
||||
let chunk_byte_size = checked_chunk_byte_len(&chunk_dims, elem_size)?;
|
||||
let (csize, fmask) = if let Some(fs) = single_filtered_size {
|
||||
(fs as u32, single_filter_mask.unwrap_or(0))
|
||||
} else {
|
||||
(chunk_byte_size as u32, 0)
|
||||
};
|
||||
vec![ChunkInfo {
|
||||
chunk_size: csize,
|
||||
filter_mask: fmask,
|
||||
offsets: vec![0u64; rank],
|
||||
address: addr,
|
||||
}]
|
||||
}
|
||||
(4, Some(2)) => {
|
||||
let spatial_chunk_dims: &[u32] = &chunk_dimensions[..rank];
|
||||
generate_implicit_chunks(
|
||||
addr,
|
||||
&dataspace.dimensions,
|
||||
spatial_chunk_dims,
|
||||
elem_size as u32,
|
||||
)
|
||||
}
|
||||
(4, Some(3)) => {
|
||||
let spatial_chunk_dims: &[u32] = &chunk_dimensions[..rank];
|
||||
let header =
|
||||
FixedArrayHeader::parse(file_data, addr as usize, offset_size, length_size)?;
|
||||
read_fixed_array_chunks(
|
||||
file_data,
|
||||
&header,
|
||||
&dataspace.dimensions,
|
||||
spatial_chunk_dims,
|
||||
elem_size as u32,
|
||||
offset_size,
|
||||
length_size,
|
||||
)?
|
||||
}
|
||||
(4, Some(4)) => {
|
||||
let spatial_chunk_dims: &[u32] = &chunk_dimensions[..rank];
|
||||
let header = ExtensibleArrayHeader::parse(
|
||||
file_data,
|
||||
addr as usize,
|
||||
offset_size,
|
||||
length_size,
|
||||
)?;
|
||||
read_extensible_array_chunks(
|
||||
file_data,
|
||||
&header,
|
||||
&dataspace.dimensions,
|
||||
spatial_chunk_dims,
|
||||
elem_size as u32,
|
||||
offset_size,
|
||||
length_size,
|
||||
)?
|
||||
}
|
||||
(v, idx) => {
|
||||
return Err(FormatError::ChunkedReadError(format!(
|
||||
"unsupported chunked layout version={v}, index_type={idx:?}"
|
||||
)));
|
||||
}
|
||||
};
|
||||
let (chunks, _) = list_chunks(
|
||||
file_data,
|
||||
layout,
|
||||
dataspace,
|
||||
elem_size,
|
||||
offset_size,
|
||||
length_size,
|
||||
)?;
|
||||
cache.populate_chunk_index(&chunks, rank);
|
||||
// Also populate the legacy index for compatibility
|
||||
if !cache.has_index() {
|
||||
@@ -2284,21 +2241,23 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cached_read_second_call_uses_cache() {
|
||||
fn cached_read_second_call_reuses_the_index() {
|
||||
let values: Vec<f64> = (0..20).map(|i| i as f64).collect();
|
||||
let (file_data, layout, dataspace) = build_1d_chunked_file(&values, 10);
|
||||
let datatype = make_f64_type();
|
||||
let cache = ChunkCache::new();
|
||||
|
||||
// First read — populates index + decompressed cache
|
||||
// First read — populates the chunk index. These chunks are stored
|
||||
// unfiltered, so they are copied straight from the file bytes and the
|
||||
// decompressed-chunk cache is (deliberately) not involved.
|
||||
let raw1 = read_chunked_data_cached(
|
||||
&file_data, &layout, &dataspace, &datatype, None, 8, 8, &cache,
|
||||
)
|
||||
.unwrap();
|
||||
assert!(cache.has_index());
|
||||
assert!(cache.cached_chunk_count() > 0);
|
||||
assert_eq!(cache.cached_chunk_count(), 0);
|
||||
|
||||
// Second read — should hit the decompressed cache
|
||||
// Second read — reuses the cached index
|
||||
let raw2 = read_chunked_data_cached(
|
||||
&file_data, &layout, &dataspace, &datatype, None, 8, 8, &cache,
|
||||
)
|
||||
|
||||
@@ -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
|
||||
@@ -49,6 +48,38 @@ pub struct ChunkOptions {
|
||||
pub pcodec: bool,
|
||||
}
|
||||
|
||||
/// Largest chunk the automatic choice produces, in bytes.
|
||||
const AUTO_CHUNK_TARGET_BYTES: u64 = 1 << 20;
|
||||
|
||||
/// Extent assumed for a dimension that is currently empty (an unlimited
|
||||
/// dimension not yet written to) — the same stand-in h5py uses.
|
||||
const AUTO_CHUNK_EMPTY_DIM: u64 = 1024;
|
||||
|
||||
/// Choose chunk dimensions for a dataset nobody specified them for.
|
||||
///
|
||||
/// Asking for compression (or any filter) without chunk dimensions used to
|
||||
/// make the whole dataset one chunk. That defeats the point of chunking: any
|
||||
/// read — even a single row — must decompress everything, and a large dataset
|
||||
/// cannot be decompressed in parallel. Datasets up to the target size stay a
|
||||
/// single chunk, exactly as before; larger ones are split by halving the
|
||||
/// dimensions in turn (so chunks keep roughly the dataset's proportions, the
|
||||
/// approach h5py takes) until a chunk fits the target.
|
||||
pub fn auto_chunk_dims(shape: &[u64], elem_size: usize) -> Vec<u64> {
|
||||
let mut dims: Vec<u64> = shape
|
||||
.iter()
|
||||
.map(|&d| if d == 0 { AUTO_CHUNK_EMPTY_DIM } else { d })
|
||||
.collect();
|
||||
let elem = elem_size.max(1) as u64;
|
||||
let bytes = |dims: &[u64]| dims.iter().fold(elem, |acc, &d| acc.saturating_mul(d));
|
||||
let mut axis = 0;
|
||||
while bytes(&dims) > AUTO_CHUNK_TARGET_BYTES && dims.iter().any(|&d| d > 1) {
|
||||
let i = axis % dims.len();
|
||||
dims[i] = dims[i].div_ceil(2);
|
||||
axis += 1;
|
||||
}
|
||||
dims
|
||||
}
|
||||
|
||||
impl ChunkOptions {
|
||||
/// Whether any chunking option is enabled.
|
||||
pub fn is_chunked(&self) -> bool {
|
||||
@@ -135,11 +166,17 @@ impl ChunkOptions {
|
||||
|
||||
/// Determine chunk dimensions, using user-specified or auto-computing.
|
||||
pub fn resolve_chunk_dims(&self, shape: &[u64]) -> Vec<u64> {
|
||||
if let Some(ref dims) = self.chunk_dims {
|
||||
dims.clone()
|
||||
} else {
|
||||
// Auto chunk: use the full dataset shape (single chunk)
|
||||
shape.to_vec()
|
||||
// Without the element size, assume 8 bytes (the widest common scalar);
|
||||
// the writer uses `resolve_chunk_dims_for`.
|
||||
self.resolve_chunk_dims_for(shape, 8)
|
||||
}
|
||||
|
||||
/// Chunk dimensions for a dataset of `shape` whose elements are `elem_size`
|
||||
/// bytes: the caller's if given, otherwise chosen automatically.
|
||||
pub fn resolve_chunk_dims_for(&self, shape: &[u64], elem_size: usize) -> Vec<u64> {
|
||||
match self.chunk_dims {
|
||||
Some(ref dims) => dims.clone(),
|
||||
None => auto_chunk_dims(shape, elem_size),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -890,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;
|
||||
@@ -1143,6 +1181,45 @@ mod tests {
|
||||
assert_eq!(dims, vec![100, 50]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn auto_chunking_splits_only_large_datasets() {
|
||||
let bytes = |dims: &[u64], elem: u64| dims.iter().product::<u64>() * elem;
|
||||
// Up to the target: one chunk, as before.
|
||||
assert_eq!(auto_chunk_dims(&[100, 50], 8), [100, 50]);
|
||||
assert_eq!(auto_chunk_dims(&[131_072], 8), [131_072]); // exactly 1 MiB
|
||||
// Larger: split, keeping proportions, never above the target.
|
||||
let big = auto_chunk_dims(&[4096, 2048], 8);
|
||||
assert!(bytes(&big, 8) <= AUTO_CHUNK_TARGET_BYTES, "{big:?}");
|
||||
assert!(bytes(&big, 8) > AUTO_CHUNK_TARGET_BYTES / 4, "{big:?}");
|
||||
assert_eq!(big[0] / big[1], 2, "proportions kept: {big:?}");
|
||||
// Every dimension stays within the dataset and at least 1.
|
||||
for shape in [
|
||||
vec![10_000_000u64],
|
||||
vec![3, 5_000_000],
|
||||
vec![1, 1, 9_000_000],
|
||||
vec![7; 9],
|
||||
] {
|
||||
let dims = auto_chunk_dims(&shape, 4);
|
||||
assert!(
|
||||
dims.iter().zip(&shape).all(|(c, s)| *c >= 1 && c <= s),
|
||||
"{shape:?} -> {dims:?}"
|
||||
);
|
||||
assert!(
|
||||
bytes(&dims, 4) <= AUTO_CHUNK_TARGET_BYTES,
|
||||
"{shape:?} -> {dims:?}"
|
||||
);
|
||||
}
|
||||
// An empty (unlimited, unwritten) dimension still gets a usable chunk.
|
||||
let growable = auto_chunk_dims(&[0, 128], 8);
|
||||
assert!(growable[0] >= 1 && bytes(&growable, 8) <= AUTO_CHUNK_TARGET_BYTES);
|
||||
// Explicit dimensions always win.
|
||||
let explicit = ChunkOptions {
|
||||
chunk_dims: Some(vec![10, 10]),
|
||||
..Default::default()
|
||||
};
|
||||
assert_eq!(explicit.resolve_chunk_dims_for(&[4096, 2048], 8), [10, 10]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn chunk_options_pipeline_deflate() {
|
||||
// Auto-shuffle is applied before compression by default (matches h5py).
|
||||
@@ -1435,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())
|
||||
@@ -1449,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));
|
||||
}
|
||||
|
||||
@@ -307,6 +307,24 @@ pub fn read_raw_data_selection(
|
||||
) -> Result<Vec<u8>, FormatError> {
|
||||
use crate::selection::Selection;
|
||||
|
||||
crate::partial_read::validate(selection, &dataspace.dimensions)?;
|
||||
|
||||
// Read only what the selection's bounding box touches when that is
|
||||
// possible; everything below is the decode-everything-then-pick path,
|
||||
// kept for the cases `partial_read` declines.
|
||||
if let Some(selected) = crate::partial_read::read_selection(
|
||||
file_data,
|
||||
layout,
|
||||
dataspace,
|
||||
datatype.type_size() as usize,
|
||||
pipeline,
|
||||
offset_size,
|
||||
length_size,
|
||||
selection,
|
||||
)? {
|
||||
return Ok(selected);
|
||||
}
|
||||
|
||||
match selection {
|
||||
Selection::All => {
|
||||
return read_raw_data_full(
|
||||
@@ -858,6 +876,30 @@ fn get_size(dt: &Datatype) -> usize {
|
||||
dt.type_size() as usize
|
||||
}
|
||||
|
||||
/// Reinterpret little-endian bytes as `count` native values of `T` on a
|
||||
/// little-endian target, in one copy.
|
||||
///
|
||||
/// The buffer is allocated uninitialised and filled by the copy. It used to be
|
||||
/// `vec![0; count]` first, which for a large dataset meant writing every page
|
||||
/// twice (zero it, then overwrite it) — about as expensive as the copy itself.
|
||||
#[cfg(target_endian = "little")]
|
||||
fn native_le_to_vec<T: Copy>(raw: &[u8], count: usize) -> Vec<T> {
|
||||
let bytes = count * core::mem::size_of::<T>();
|
||||
debug_assert!(bytes <= raw.len());
|
||||
let mut result: Vec<T> = Vec::with_capacity(count);
|
||||
// SAFETY: `result` has capacity for `count` values of `T`, i.e. `bytes`
|
||||
// bytes; `raw` holds at least `bytes` bytes (callers derive `count` from
|
||||
// `raw.len() / size_of::<T>()`); the regions cannot overlap because
|
||||
// `result` was just allocated. Every `T` used here (f32/f64/i32/i64) is
|
||||
// valid for any bit pattern, so after the copy all `count` values are
|
||||
// initialised and `set_len` is sound.
|
||||
unsafe {
|
||||
core::ptr::copy_nonoverlapping(raw.as_ptr(), result.as_mut_ptr().cast::<u8>(), bytes);
|
||||
result.set_len(count);
|
||||
}
|
||||
result
|
||||
}
|
||||
|
||||
/// Convert raw bytes to `f64` values.
|
||||
pub fn read_as_f64(raw: &[u8], datatype: &Datatype) -> Result<Vec<f64>, FormatError> {
|
||||
// Array datatypes (e.g. an array-typed compound member) are read as a flat
|
||||
@@ -885,14 +927,7 @@ pub fn read_as_f64(raw: &[u8], datatype: &Datatype) -> Result<Vec<f64>, FormatEr
|
||||
..
|
||||
}
|
||||
) {
|
||||
let mut result = vec![0.0f64; count];
|
||||
// SAFETY: On LE platforms, f64 in-memory representation matches LE bytes.
|
||||
// We copy raw bytes directly into the f64 buffer.
|
||||
// SAFETY: The byte slice is properly aligned for this type and the length is divisible by size_of::<T>().
|
||||
unsafe {
|
||||
core::ptr::copy_nonoverlapping(raw.as_ptr(), result.as_mut_ptr() as *mut u8, raw.len());
|
||||
}
|
||||
return Ok(result);
|
||||
return Ok(native_le_to_vec::<f64>(raw, count));
|
||||
}
|
||||
|
||||
let order = get_byte_order(datatype);
|
||||
@@ -975,12 +1010,7 @@ pub fn read_as_i64(raw: &[u8], datatype: &Datatype) -> Result<Vec<i64>, FormatEr
|
||||
}
|
||||
)
|
||||
{
|
||||
let mut result = vec![0i64; count];
|
||||
// SAFETY: The byte slice is properly aligned for this type and the length is divisible by size_of::<T>().
|
||||
unsafe {
|
||||
core::ptr::copy_nonoverlapping(raw.as_ptr(), result.as_mut_ptr() as *mut u8, raw.len());
|
||||
}
|
||||
return Ok(result);
|
||||
return Ok(native_le_to_vec::<i64>(raw, count));
|
||||
}
|
||||
|
||||
let order = get_byte_order(datatype);
|
||||
@@ -1044,12 +1074,7 @@ pub fn read_as_f32(raw: &[u8], datatype: &Datatype) -> Result<Vec<f32>, FormatEr
|
||||
..
|
||||
}
|
||||
) {
|
||||
let mut result = vec![0.0f32; count];
|
||||
// SAFETY: The byte slice is properly aligned for this type and the length is divisible by size_of::<T>().
|
||||
unsafe {
|
||||
core::ptr::copy_nonoverlapping(raw.as_ptr(), result.as_mut_ptr() as *mut u8, raw.len());
|
||||
}
|
||||
return Ok(result);
|
||||
return Ok(native_le_to_vec::<f32>(raw, count));
|
||||
}
|
||||
|
||||
let order = get_byte_order(datatype);
|
||||
@@ -1126,12 +1151,7 @@ pub fn read_as_i32(raw: &[u8], datatype: &Datatype) -> Result<Vec<i32>, FormatEr
|
||||
}
|
||||
)
|
||||
{
|
||||
let mut result = vec![0i32; count];
|
||||
// SAFETY: The byte slice is properly aligned for this type and the length is divisible by size_of::<T>().
|
||||
unsafe {
|
||||
core::ptr::copy_nonoverlapping(raw.as_ptr(), result.as_mut_ptr() as *mut u8, raw.len());
|
||||
}
|
||||
return Ok(result);
|
||||
return Ok(native_le_to_vec::<i32>(raw, count));
|
||||
}
|
||||
|
||||
let order = get_byte_order(datatype);
|
||||
@@ -1407,6 +1427,26 @@ pub fn read_object_references(
|
||||
}
|
||||
Ok(result)
|
||||
}
|
||||
Datatype::Reference {
|
||||
ref_type: crate::datatype::ReferenceType::Object2,
|
||||
size,
|
||||
} => {
|
||||
let elem_size = *size as usize;
|
||||
if elem_size == 0 {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
if !raw.len().is_multiple_of(elem_size) {
|
||||
return Err(FormatError::DataSizeMismatch {
|
||||
expected: 0,
|
||||
actual: raw.len(),
|
||||
});
|
||||
}
|
||||
raw.chunks_exact(elem_size)
|
||||
.map(|element| {
|
||||
decode_std_object_ref(element).map(|address| ObjectReference { address })
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
_ => Err(FormatError::TypeMismatch {
|
||||
expected: "Reference(Object)",
|
||||
actual: datatype_name(datatype),
|
||||
@@ -1414,6 +1454,46 @@ pub fn read_object_references(
|
||||
}
|
||||
}
|
||||
|
||||
/// Decode one `H5T_STD_REF` object reference as stored in a dataset:
|
||||
/// `type(1) flags(1) token_size(1) token(token_size)`, zero-padded to the
|
||||
/// element size. For a reference within the same file the token is the target
|
||||
/// object's header address. An all-zero element is a null reference and
|
||||
/// decodes to the undefined address (`u64::MAX`).
|
||||
fn decode_std_object_ref(element: &[u8]) -> Result<u64, FormatError> {
|
||||
const STD_REF_OBJECT: u8 = 2;
|
||||
const FLAG_EXTERNAL: u8 = 0x01;
|
||||
if element.iter().all(|&b| b == 0) {
|
||||
return Ok(u64::MAX);
|
||||
}
|
||||
let [ref_type, flags, token_size, token @ ..] = element else {
|
||||
return Err(FormatError::UnexpectedEof {
|
||||
expected: 3,
|
||||
available: element.len(),
|
||||
});
|
||||
};
|
||||
if *ref_type != STD_REF_OBJECT {
|
||||
return Err(FormatError::InvalidReferenceType(*ref_type));
|
||||
}
|
||||
if flags & FLAG_EXTERNAL != 0 {
|
||||
// Carries a file name as well; nothing here follows those.
|
||||
return Err(FormatError::TypeMismatch {
|
||||
expected: "object reference within this file",
|
||||
actual: "external object reference",
|
||||
});
|
||||
}
|
||||
let n = *token_size as usize;
|
||||
if n == 0 || n > 8 || n > token.len() {
|
||||
return Err(FormatError::UnexpectedEof {
|
||||
expected: 3 + n,
|
||||
available: element.len(),
|
||||
});
|
||||
}
|
||||
Ok(token[..n]
|
||||
.iter()
|
||||
.rev()
|
||||
.fold(0u64, |addr, &byte| (addr << 8) | u64::from(byte)))
|
||||
}
|
||||
|
||||
/// Read region references from raw bytes.
|
||||
///
|
||||
/// Region references encode a dataset selection (hyperslab, point list, etc.)
|
||||
|
||||
@@ -36,8 +36,18 @@ pub enum CharacterSet {
|
||||
/// Reference type.
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub enum ReferenceType {
|
||||
/// Legacy object reference: the target's object header address.
|
||||
Object,
|
||||
/// Legacy dataset region reference.
|
||||
DatasetRegion,
|
||||
/// `H5T_STD_REF` object reference (HDF5 1.12+, datatype message version
|
||||
/// 4): a small header followed by an object token. Decoded by
|
||||
/// `data_read::read_object_references`.
|
||||
Object2,
|
||||
/// `H5T_STD_REF` dataset region reference.
|
||||
DatasetRegion2,
|
||||
/// `H5T_STD_REF` attribute reference.
|
||||
Attribute,
|
||||
}
|
||||
|
||||
/// A member of a compound datatype.
|
||||
@@ -424,9 +434,15 @@ impl Datatype {
|
||||
7 => {
|
||||
// Reference
|
||||
let ref_type_val = bf0 & 0x0F;
|
||||
let ref_type = match ref_type_val {
|
||||
0 => ReferenceType::Object,
|
||||
1 => ReferenceType::DatasetRegion,
|
||||
// Datatype message version 4 (HDF5 1.12) revised this class:
|
||||
// types 2-4 are the new `H5T_STD_REF` references, and the high
|
||||
// nibble of the first flag byte carries their encoding version.
|
||||
let ref_type = match (ref_type_val, version) {
|
||||
(0, _) => ReferenceType::Object,
|
||||
(1, _) => ReferenceType::DatasetRegion,
|
||||
(2, 4..) => ReferenceType::Object2,
|
||||
(3, 4..) => ReferenceType::DatasetRegion2,
|
||||
(4, 4..) => ReferenceType::Attribute,
|
||||
_ => return Err(FormatError::InvalidReferenceType(ref_type_val)),
|
||||
};
|
||||
Ok((Datatype::Reference { size, ref_type }, pos))
|
||||
@@ -1563,6 +1579,28 @@ mod tests {
|
||||
assert_eq!(err, FormatError::InvalidCharacterSet(2));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_reference_v4_std_ref_from_hdf5_2_0() {
|
||||
// Datatype message of an H5T_STD_REF dataset written by HDF5 2.0:
|
||||
// class 7, version 4, type 2 (object), encoding version 1, 18 bytes.
|
||||
let bytes = [0x47, 0x12, 0x00, 0x00, 0x12, 0x00, 0x00, 0x00];
|
||||
let (dt, consumed) = Datatype::parse(&bytes).unwrap();
|
||||
assert_eq!(consumed, 8);
|
||||
assert_eq!(
|
||||
dt,
|
||||
Datatype::Reference {
|
||||
size: 18,
|
||||
ref_type: ReferenceType::Object2
|
||||
}
|
||||
);
|
||||
// The new types are only valid from datatype version 4.
|
||||
let old_version = [0x37, 0x12, 0x00, 0x00, 0x12, 0x00, 0x00, 0x00];
|
||||
assert_eq!(
|
||||
Datatype::parse(&old_version).unwrap_err(),
|
||||
FormatError::InvalidReferenceType(2)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_error_invalid_reference_type() {
|
||||
let buf = build_dt_header(7, 1, [5, 0, 0], 8);
|
||||
|
||||
@@ -117,6 +117,9 @@ pub enum FormatError {
|
||||
/// A message is marked shared but was parsed without access to the file,
|
||||
/// so the reference to the real message could not be followed.
|
||||
UnresolvedSharedMessage,
|
||||
/// A selection does not fit the dataset it was applied to (wrong rank, or
|
||||
/// it reaches past a dimension's extent).
|
||||
SelectionOutOfBounds(String),
|
||||
/// The dataset's raw data is stored in external files (External Data
|
||||
/// Files message), which this reader does not follow.
|
||||
ExternalDataFilesUnsupported,
|
||||
@@ -333,6 +336,9 @@ impl fmt::Display for FormatError {
|
||||
f,
|
||||
"dataset raw data is stored in external file(s), which is not supported"
|
||||
),
|
||||
FormatError::SelectionOutOfBounds(msg) => {
|
||||
write!(f, "selection out of bounds: {msg}")
|
||||
}
|
||||
FormatError::UnresolvedSharedMessage => write!(
|
||||
f,
|
||||
"message is shared but no file data was available to resolve it"
|
||||
|
||||
@@ -1221,8 +1221,10 @@ impl FileWriter {
|
||||
precompressed: None,
|
||||
});
|
||||
} else if is_chunked[i] {
|
||||
let chunk_dims = d.chunk_options.resolve_chunk_dims(&d.ds.dimensions);
|
||||
let elem_size = d.dt.type_size() as usize;
|
||||
let chunk_dims = d
|
||||
.chunk_options
|
||||
.resolve_chunk_dims_for(&d.ds.dimensions, elem_size);
|
||||
// Compress once in Pass 1; cache the result so Pass 2 can skip
|
||||
// re-compression and just rebuild the index with real addresses.
|
||||
let pre = precompress_chunks(
|
||||
|
||||
@@ -845,9 +845,33 @@ fn shuffle_decompress(data: &[u8], element_size: usize) -> Result<Vec<u8>, Forma
|
||||
let num_elements = data.len() / element_size;
|
||||
let mut result = vec![0u8; data.len()];
|
||||
|
||||
for i in 0..num_elements {
|
||||
for j in 0..element_size {
|
||||
result[i * element_size + j] = data[j * num_elements + i];
|
||||
// The shuffled stream is `element_size` byte planes of `num_elements`
|
||||
// bytes each; un-shuffling interleaves them. This is on the read path of
|
||||
// every compressed dataset (shuffle is applied automatically before
|
||||
// compression). The naive `result[i * es + j] = data[j * n + i]` form does
|
||||
// a multiply and two bounds checks per byte and defeats vectorisation;
|
||||
// fixed-width plane arrays sliced to a common length let the compiler
|
||||
// hoist the checks and emit interleaves for the common 4- and 8-byte
|
||||
// element sizes.
|
||||
fn interleave<const W: usize>(data: &[u8], n: usize, out: &mut [u8]) {
|
||||
let planes: [&[u8]; W] = core::array::from_fn(|j| &data[j * n..(j + 1) * n]);
|
||||
for (i, element) in out.as_chunks_mut::<W>().0.iter_mut().enumerate() {
|
||||
for (byte, plane) in element.iter_mut().zip(&planes) {
|
||||
*byte = plane[i];
|
||||
}
|
||||
}
|
||||
}
|
||||
match element_size {
|
||||
2 => interleave::<2>(data, num_elements, &mut result),
|
||||
4 => interleave::<4>(data, num_elements, &mut result),
|
||||
8 => interleave::<8>(data, num_elements, &mut result),
|
||||
16 => interleave::<16>(data, num_elements, &mut result),
|
||||
_ => {
|
||||
for (i, element) in result.chunks_exact_mut(element_size).enumerate() {
|
||||
for (j, byte) in element.iter_mut().enumerate() {
|
||||
*byte = data[j * num_elements + i];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1848,4 +1872,20 @@ mod tests {
|
||||
};
|
||||
assert!(decompress_chunk(&data, &pipeline, 16, 1).is_err());
|
||||
}
|
||||
#[test]
|
||||
fn unshuffle_inverts_shuffle_for_every_element_size() {
|
||||
for element_size in [1usize, 2, 3, 4, 5, 8, 12, 16, 24] {
|
||||
for elements in [0usize, 1, 2, 7, 64, 1000] {
|
||||
let original: Vec<u8> = (0..element_size * elements)
|
||||
.map(|i| (i * 31 + 7) as u8)
|
||||
.collect();
|
||||
let shuffled = shuffle_compress(&original, element_size).unwrap();
|
||||
assert_eq!(
|
||||
shuffle_decompress(&shuffled, element_size).unwrap(),
|
||||
original,
|
||||
"element_size {element_size}, {elements} elements"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -89,6 +89,7 @@ pub mod object_header;
|
||||
pub mod object_header_writer;
|
||||
#[cfg(feature = "parallel")]
|
||||
pub mod parallel_read;
|
||||
pub mod partial_read;
|
||||
pub mod profiling;
|
||||
pub mod property_list;
|
||||
pub mod selection;
|
||||
|
||||
@@ -0,0 +1,359 @@
|
||||
//! Selection reads that cost what the selection costs, not what the dataset
|
||||
//! costs.
|
||||
//!
|
||||
//! [`crate::data_read::read_raw_data_selection`] used to decode the *entire*
|
||||
//! dataset and then pick elements out of it, so reading a 64x64 window of a
|
||||
//! large dataset took about as long as reading all of it. Here the selection's
|
||||
//! bounding box is materialised instead — only the rows of a contiguous
|
||||
//! dataset, or only the chunks, that overlap it — and the existing extractor
|
||||
//! runs over that small buffer with the selection translated to the box's
|
||||
//! origin. Extraction semantics are therefore exactly the full-read ones.
|
||||
|
||||
#[cfg(not(feature = "std"))]
|
||||
use alloc::string as alloc_or_std;
|
||||
#[cfg(not(feature = "std"))]
|
||||
use alloc::{format, vec, vec::Vec};
|
||||
#[cfg(feature = "std")]
|
||||
use std::string as alloc_or_std;
|
||||
|
||||
use crate::chunked_read::{alloc_output, checked_byte_len, list_chunks};
|
||||
use crate::data_layout::DataLayout;
|
||||
use crate::data_read::extract_selection_from_buffer;
|
||||
use crate::dataspace::Dataspace;
|
||||
use crate::error::FormatError;
|
||||
use crate::filter_pipeline::FilterPipeline;
|
||||
use crate::filters::decompress_chunk;
|
||||
use crate::selection::Selection;
|
||||
|
||||
/// The smallest axis-aligned box containing every selected element, as
|
||||
/// `(start, extent)` per dimension. `None` when there is nothing to gain or
|
||||
/// the selection is not valid for `dims` (the caller's full path then reports
|
||||
/// the error exactly as before).
|
||||
fn bounding_box(selection: &Selection, dims: &[u64]) -> Option<(Vec<u64>, Vec<u64>)> {
|
||||
match selection {
|
||||
Selection::Hyperslab {
|
||||
start,
|
||||
stride,
|
||||
count,
|
||||
block,
|
||||
} => {
|
||||
let rank = dims.len();
|
||||
if [start.len(), stride.len(), count.len(), block.len()] != [rank; 4] {
|
||||
return None;
|
||||
}
|
||||
let mut extent = Vec::with_capacity(rank);
|
||||
for d in 0..rank {
|
||||
if count[d] == 0 || block[d] == 0 {
|
||||
return None;
|
||||
}
|
||||
// Last selected index + 1, relative to start.
|
||||
let span = (count[d] - 1)
|
||||
.checked_mul(stride[d])?
|
||||
.checked_add(block[d])?;
|
||||
if start[d].checked_add(span)? > dims[d] {
|
||||
return None;
|
||||
}
|
||||
extent.push(span);
|
||||
}
|
||||
Some((start.clone(), extent))
|
||||
}
|
||||
Selection::Points(points) => {
|
||||
let rank = dims.len();
|
||||
let first = points.first()?;
|
||||
if first.len() != rank {
|
||||
return None;
|
||||
}
|
||||
let (mut lo, mut hi) = (first.clone(), first.clone());
|
||||
for p in points {
|
||||
if p.len() != rank {
|
||||
return None;
|
||||
}
|
||||
for d in 0..rank {
|
||||
if p[d] >= dims[d] {
|
||||
return None;
|
||||
}
|
||||
lo[d] = lo[d].min(p[d]);
|
||||
hi[d] = hi[d].max(p[d]);
|
||||
}
|
||||
}
|
||||
let extent = lo.iter().zip(&hi).map(|(l, h)| h - l + 1).collect();
|
||||
Some((lo, extent))
|
||||
}
|
||||
Selection::All | Selection::None => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Check that `selection` addresses only elements that exist in a dataset of
|
||||
/// shape `dims`. Without this an out-of-range selection read *something*: a
|
||||
/// hyperslab past the edge came back padded with zeros, and a point whose
|
||||
/// column was out of range wrapped into the next row.
|
||||
pub fn validate(selection: &Selection, dims: &[u64]) -> Result<(), FormatError> {
|
||||
let rank = dims.len();
|
||||
let bad = |msg: alloc_or_std::String| Err(FormatError::SelectionOutOfBounds(msg));
|
||||
match selection {
|
||||
Selection::All | Selection::None => Ok(()),
|
||||
Selection::Hyperslab {
|
||||
start,
|
||||
stride,
|
||||
count,
|
||||
block,
|
||||
} => {
|
||||
if [start.len(), stride.len(), count.len(), block.len()] != [rank; 4] {
|
||||
return bad(format!("hyperslab rank does not match dataset rank {rank}"));
|
||||
}
|
||||
for d in 0..rank {
|
||||
if count[d] == 0 || block[d] == 0 {
|
||||
continue; // selects nothing along this dimension
|
||||
}
|
||||
let end = (count[d] - 1)
|
||||
.checked_mul(stride[d])
|
||||
.and_then(|v| v.checked_add(block[d]))
|
||||
.and_then(|v| v.checked_add(start[d]));
|
||||
if !end.is_some_and(|end| end <= dims[d]) {
|
||||
return bad(format!(
|
||||
"dimension {d}: start {} stride {} count {} block {} exceeds extent {}",
|
||||
start[d], stride[d], count[d], block[d], dims[d]
|
||||
));
|
||||
}
|
||||
if block[d] > stride[d] && count[d] > 1 {
|
||||
return bad(format!(
|
||||
"dimension {d}: block {} larger than stride {} (overlapping blocks)",
|
||||
block[d], stride[d]
|
||||
));
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
Selection::Points(points) => {
|
||||
for p in points {
|
||||
if p.len() != rank {
|
||||
return bad(format!("point {p:?} does not match dataset rank {rank}"));
|
||||
}
|
||||
if let Some(d) = (0..rank).find(|&d| p[d] >= dims[d]) {
|
||||
return bad(format!(
|
||||
"point {p:?}: coordinate {} exceeds extent {} of dimension {d}",
|
||||
p[d], dims[d]
|
||||
));
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// The same selection expressed relative to `origin`.
|
||||
fn translate(selection: &Selection, origin: &[u64]) -> Selection {
|
||||
match selection {
|
||||
Selection::Hyperslab {
|
||||
start,
|
||||
stride,
|
||||
count,
|
||||
block,
|
||||
} => Selection::Hyperslab {
|
||||
start: start.iter().zip(origin).map(|(s, o)| s - o).collect(),
|
||||
stride: stride.clone(),
|
||||
count: count.clone(),
|
||||
block: block.clone(),
|
||||
},
|
||||
Selection::Points(points) => Selection::Points(
|
||||
points
|
||||
.iter()
|
||||
.map(|p| p.iter().zip(origin).map(|(c, o)| c - o).collect())
|
||||
.collect(),
|
||||
),
|
||||
other => other.clone(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Copy the part of a source region that overlaps the box into `out` (which
|
||||
/// is the box, row-major).
|
||||
///
|
||||
/// The source region starts at `src_origin` in dataset coordinates, has shape
|
||||
/// `src_shape`, and its elements are in `src` row-major. One `memcpy` per
|
||||
/// overlapping row of the last dimension.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn copy_overlap(
|
||||
src: &[u8],
|
||||
src_origin: &[u64],
|
||||
src_shape: &[u64],
|
||||
out: &mut [u8],
|
||||
box_start: &[u64],
|
||||
box_extent: &[u64],
|
||||
elem_size: usize,
|
||||
) {
|
||||
let rank = box_start.len();
|
||||
// Overlap in dataset coordinates.
|
||||
let mut lo = vec![0u64; rank];
|
||||
let mut hi = vec![0u64; rank];
|
||||
for d in 0..rank {
|
||||
lo[d] = src_origin[d].max(box_start[d]);
|
||||
hi[d] = (src_origin[d] + src_shape[d]).min(box_start[d] + box_extent[d]);
|
||||
if lo[d] >= hi[d] {
|
||||
return;
|
||||
}
|
||||
}
|
||||
let strides = |shape: &[u64]| {
|
||||
let mut s = vec![1u64; rank];
|
||||
for d in (0..rank.saturating_sub(1)).rev() {
|
||||
s[d] = s[d + 1] * shape[d + 1];
|
||||
}
|
||||
s
|
||||
};
|
||||
let (src_strides, out_strides) = (strides(src_shape), strides(box_extent));
|
||||
let last = rank - 1;
|
||||
let run = ((hi[last] - lo[last]) as usize) * elem_size;
|
||||
|
||||
let mut idx = lo.clone();
|
||||
loop {
|
||||
let src_at: u64 = (0..rank)
|
||||
.map(|d| (idx[d] - src_origin[d]) * src_strides[d])
|
||||
.sum();
|
||||
let out_at: u64 = (0..rank)
|
||||
.map(|d| (idx[d] - box_start[d]) * out_strides[d])
|
||||
.sum();
|
||||
let (s, o) = (src_at as usize * elem_size, out_at as usize * elem_size);
|
||||
if let (Some(from), Some(to)) = (src.get(s..s + run), out.get_mut(o..o + run)) {
|
||||
to.copy_from_slice(from);
|
||||
}
|
||||
// Advance over every dimension but the last.
|
||||
let mut d = last;
|
||||
loop {
|
||||
if d == 0 {
|
||||
return;
|
||||
}
|
||||
d -= 1;
|
||||
idx[d] += 1;
|
||||
if idx[d] < hi[d] {
|
||||
break;
|
||||
}
|
||||
idx[d] = lo[d];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Read `selection` without materialising the whole dataset, when that is
|
||||
/// possible and worthwhile. `Ok(None)` means "use the full-read path": an
|
||||
/// `All`/`None`/invalid selection, a layout this doesn't handle (compact,
|
||||
/// virtual, storage-less), or a bounding box covering most of the dataset.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn read_selection(
|
||||
file_data: &[u8],
|
||||
layout: &DataLayout,
|
||||
dataspace: &Dataspace,
|
||||
elem_size: usize,
|
||||
pipeline: Option<&FilterPipeline>,
|
||||
offset_size: u8,
|
||||
length_size: u8,
|
||||
selection: &Selection,
|
||||
) -> Result<Option<Vec<u8>>, FormatError> {
|
||||
let dims = &dataspace.dimensions;
|
||||
if dims.is_empty() || elem_size == 0 {
|
||||
return Ok(None);
|
||||
}
|
||||
let Some((box_start, box_extent)) = bounding_box(selection, dims) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let total = dataspace.checked_num_elements()?;
|
||||
let box_elements = box_extent
|
||||
.iter()
|
||||
.try_fold(1u64, |acc, &e| acc.checked_mul(e))
|
||||
.ok_or_else(|| FormatError::Overflow("selection bounding box overflows".into()))?;
|
||||
// A box covering most of the dataset gains nothing over the full path.
|
||||
if box_elements.saturating_mul(2) > total {
|
||||
return Ok(None);
|
||||
}
|
||||
let mut boxed = alloc_output(checked_byte_len(box_elements, elem_size)?)?;
|
||||
|
||||
match layout {
|
||||
DataLayout::Contiguous {
|
||||
address: Some(address),
|
||||
..
|
||||
} => {
|
||||
let base = usize::try_from(*address)
|
||||
.map_err(|_| FormatError::Overflow("data address exceeds usize".into()))?;
|
||||
let data = file_data
|
||||
.get(base..)
|
||||
.and_then(|d| d.get(..checked_byte_len(total, elem_size).ok()?))
|
||||
.ok_or(FormatError::UnexpectedEof {
|
||||
expected: base,
|
||||
available: file_data.len(),
|
||||
})?;
|
||||
let origin = vec![0u64; dims.len()];
|
||||
copy_overlap(
|
||||
data,
|
||||
&origin,
|
||||
dims,
|
||||
&mut boxed,
|
||||
&box_start,
|
||||
&box_extent,
|
||||
elem_size,
|
||||
);
|
||||
}
|
||||
DataLayout::Chunked {
|
||||
btree_address: Some(_),
|
||||
..
|
||||
} => {
|
||||
let (chunks, chunk_dims) = list_chunks(
|
||||
file_data,
|
||||
layout,
|
||||
dataspace,
|
||||
elem_size,
|
||||
offset_size,
|
||||
length_size,
|
||||
)?;
|
||||
let rank = dims.len();
|
||||
let chunk_shape: Vec<u64> = chunk_dims.iter().map(|&d| d as u64).collect();
|
||||
let chunk_bytes = crate::chunked_read::checked_chunk_byte_len(&chunk_dims, elem_size)?;
|
||||
for chunk in &chunks {
|
||||
if chunk.offsets.len() < rank || chunk.address == u64::MAX {
|
||||
continue;
|
||||
}
|
||||
let origin = &chunk.offsets[..rank];
|
||||
let overlaps = (0..rank).all(|d| {
|
||||
origin[d] < box_start[d] + box_extent[d]
|
||||
&& origin[d].saturating_add(chunk_shape[d]) > box_start[d]
|
||||
});
|
||||
if !overlaps {
|
||||
continue;
|
||||
}
|
||||
let at = usize::try_from(chunk.address)
|
||||
.map_err(|_| FormatError::Overflow("chunk address exceeds usize".into()))?;
|
||||
let raw = at
|
||||
.checked_add(chunk.chunk_size as usize)
|
||||
.and_then(|end| file_data.get(at..end))
|
||||
.ok_or(FormatError::UnexpectedEof {
|
||||
expected: at.saturating_add(chunk.chunk_size as usize),
|
||||
available: file_data.len(),
|
||||
})?;
|
||||
// Mirrors the full-read path: a non-zero filter mask means the
|
||||
// chunk was stored unfiltered.
|
||||
let decoded;
|
||||
let data: &[u8] = match pipeline {
|
||||
Some(pl) if chunk.filter_mask == 0 => {
|
||||
decoded = decompress_chunk(raw, pl, chunk_bytes, elem_size as u32)?;
|
||||
&decoded
|
||||
}
|
||||
_ => raw,
|
||||
};
|
||||
copy_overlap(
|
||||
data,
|
||||
origin,
|
||||
&chunk_shape,
|
||||
&mut boxed,
|
||||
&box_start,
|
||||
&box_extent,
|
||||
elem_size,
|
||||
);
|
||||
}
|
||||
}
|
||||
_ => return Ok(None),
|
||||
}
|
||||
|
||||
extract_selection_from_buffer(
|
||||
&boxed,
|
||||
&box_extent,
|
||||
elem_size,
|
||||
&translate(selection, &box_start),
|
||||
)
|
||||
.map(Some)
|
||||
}
|
||||
@@ -0,0 +1,44 @@
|
||||
"""Generate std_ref_hdf5_2_0.h5: a dataset of H5T_STD_REF (the reference
|
||||
datatype introduced in HDF5 1.12, datatype message version 4) holding two
|
||||
object references — to /target (a dataset) and /grp (a group).
|
||||
|
||||
h5py has no API for this type, so the file is written by calling the libhdf5
|
||||
bundled in the h5py wheel directly through ctypes. Written with h5py 3.16.0 /
|
||||
HDF5 2.0.0. Re-run only if the fixture ever needs regenerating:
|
||||
|
||||
python gen_std_ref.py std_ref_hdf5_2_0.h5
|
||||
"""
|
||||
import ctypes
|
||||
import glob
|
||||
import os
|
||||
import sys
|
||||
|
||||
import h5py
|
||||
import numpy as np
|
||||
|
||||
libdir = os.path.join(os.path.dirname(os.path.dirname(h5py.__file__)), "h5py.libs")
|
||||
libs = [p for p in glob.glob(os.path.join(libdir, "libhdf5*.so*")) if "_hl" not in os.path.basename(p)]
|
||||
lib = ctypes.CDLL(libs[0])
|
||||
lib.H5open()
|
||||
hid = ctypes.c_int64
|
||||
std_ref = hid.in_dll(lib, "H5T_STD_REF_g").value
|
||||
|
||||
lib.H5Screate_simple.restype = hid
|
||||
lib.H5Screate_simple.argtypes = [ctypes.c_int, ctypes.POINTER(ctypes.c_uint64), ctypes.POINTER(ctypes.c_uint64)]
|
||||
lib.H5Dcreate2.restype = hid
|
||||
lib.H5Dcreate2.argtypes = [hid, ctypes.c_char_p, hid, hid, hid, hid, hid]
|
||||
lib.H5Rcreate_object.argtypes = [hid, ctypes.c_char_p, hid, ctypes.c_void_p]
|
||||
lib.H5Dwrite.argtypes = [hid, hid, hid, hid, hid, ctypes.c_void_p]
|
||||
lib.H5Dclose.argtypes = [hid]
|
||||
|
||||
with h5py.File(sys.argv[1], "w", libver="latest") as f:
|
||||
f.create_dataset("target", data=np.arange(5, dtype="<i4"))
|
||||
f.create_group("grp")
|
||||
fid = f.id.id
|
||||
sid = lib.H5Screate_simple(1, (ctypes.c_uint64 * 1)(2), None)
|
||||
did = lib.H5Dcreate2(fid, b"refs", std_ref, sid, 0, 0, 0)
|
||||
refs = ((ctypes.c_ubyte * 64) * 2)() # H5R_ref_t is a 64-byte buffer
|
||||
assert lib.H5Rcreate_object(fid, b"/target", 0, ctypes.byref(refs[0])) == 0
|
||||
assert lib.H5Rcreate_object(fid, b"/grp", 0, ctypes.byref(refs[1])) == 0
|
||||
assert lib.H5Dwrite(did, std_ref, 0, 0, 0, ctypes.byref(refs)) == 0
|
||||
lib.H5Dclose(did)
|
||||
Binary file not shown.
@@ -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();
|
||||
|
||||
@@ -316,3 +325,97 @@ print('ok')
|
||||
// Clean up
|
||||
let _ = std::fs::remove_file(&path);
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// H5T_STD_REF (HDF5 1.12+ references, datatype message version 4)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// `fixtures/std_ref_hdf5_2_0.h5` (see `gen_std_ref.py`) holds a dataset of
|
||||
/// `H5T_STD_REF` with two object references, written by HDF5 2.0 itself. The
|
||||
/// datatype used to be rejected with `InvalidReferenceType(2)`.
|
||||
#[test]
|
||||
fn std_ref_object_references_from_hdf5_2_0() {
|
||||
use clawhdf5_format::data_layout::DataLayout;
|
||||
use clawhdf5_format::dataspace::Dataspace;
|
||||
use clawhdf5_format::group_v2::resolve_path_any;
|
||||
use clawhdf5_format::message_type::MessageType;
|
||||
use clawhdf5_format::object_header::ObjectHeader;
|
||||
use clawhdf5_format::signature::find_signature;
|
||||
use clawhdf5_format::superblock::Superblock;
|
||||
|
||||
let bytes: &[u8] = include_bytes!("fixtures/std_ref_hdf5_2_0.h5");
|
||||
let sb = Superblock::parse(bytes, find_signature(bytes).unwrap()).unwrap();
|
||||
let (os, ls) = (sb.offset_size, sb.length_size);
|
||||
|
||||
let refs_addr = resolve_path_any(bytes, &sb, "refs").unwrap();
|
||||
let header = ObjectHeader::parse(bytes, refs_addr as usize, os, ls).unwrap();
|
||||
let message = |t: MessageType| {
|
||||
&header
|
||||
.messages
|
||||
.iter()
|
||||
.find(|m| m.msg_type == t)
|
||||
.unwrap()
|
||||
.data
|
||||
};
|
||||
|
||||
let (datatype, _) = Datatype::parse(message(MessageType::Datatype)).unwrap();
|
||||
assert_eq!(
|
||||
datatype,
|
||||
Datatype::Reference {
|
||||
size: 18,
|
||||
ref_type: ReferenceType::Object2
|
||||
}
|
||||
);
|
||||
let dataspace = Dataspace::parse(message(MessageType::Dataspace), ls).unwrap();
|
||||
let layout = DataLayout::parse(message(MessageType::DataLayout), os, ls).unwrap();
|
||||
let raw =
|
||||
clawhdf5_format::data_read::read_raw_data(bytes, &layout, &dataspace, &datatype).unwrap();
|
||||
assert_eq!(raw.len(), 2 * 18);
|
||||
|
||||
// The references point at the objects they were created from.
|
||||
let refs = read_object_references(&raw, &datatype, os).unwrap();
|
||||
let addresses: Vec<u64> = refs.iter().map(|r| r.address).collect();
|
||||
assert_eq!(
|
||||
addresses,
|
||||
[
|
||||
resolve_path_any(bytes, &sb, "target").unwrap(),
|
||||
resolve_path_any(bytes, &sb, "grp").unwrap(),
|
||||
]
|
||||
);
|
||||
// And what they point at is a real object header.
|
||||
for address in addresses {
|
||||
ObjectHeader::parse(bytes, address as usize, os, ls).unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn std_ref_decoding_rejects_malformed_elements() {
|
||||
let dt = Datatype::Reference {
|
||||
size: 18,
|
||||
ref_type: ReferenceType::Object2,
|
||||
};
|
||||
let mut good = vec![0u8; 18];
|
||||
good[..4].copy_from_slice(&[2, 0, 8, 0xb3]);
|
||||
assert_eq!(
|
||||
read_object_references(&good, &dt, 8).unwrap()[0].address,
|
||||
0xb3
|
||||
);
|
||||
|
||||
// Null reference.
|
||||
assert_eq!(
|
||||
read_object_references(&[0u8; 18], &dt, 8).unwrap()[0].address,
|
||||
u64::MAX
|
||||
);
|
||||
for (what, patch) in [
|
||||
("wrong reference type", (0usize, 3u8)),
|
||||
("external flag", (1, 1)),
|
||||
("token longer than the element", (2, 200)),
|
||||
("zero-length token", (2, 0)),
|
||||
] {
|
||||
let mut bad = good.clone();
|
||||
bad[patch.0] = patch.1;
|
||||
assert!(read_object_references(&bad, &dt, 8).is_err(), "{what}");
|
||||
}
|
||||
// Not a whole number of elements.
|
||||
assert!(read_object_references(&good[..17], &dt, 8).is_err());
|
||||
}
|
||||
|
||||
@@ -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,6 +1,6 @@
|
||||
[package]
|
||||
name = "clawhdf5-gpu"
|
||||
version = "2.3.0"
|
||||
version = "2.6.0"
|
||||
edition = "2024"
|
||||
description = "GPU-accelerated vector operations for rustyhdf5 using wgpu compute shaders"
|
||||
license = "MIT"
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
[package]
|
||||
name = "clawhdf5-io"
|
||||
version = "2.3.0"
|
||||
version = "2.6.0"
|
||||
edition = "2024"
|
||||
description = "I/O abstraction layer for rustyhdf5"
|
||||
license = "MIT"
|
||||
@@ -10,7 +10,7 @@ keywords = ["hdf5", "io", "science", "data"]
|
||||
categories = ["filesystem", "science"]
|
||||
|
||||
[dependencies]
|
||||
clawhdf5-format = { path = "../clawhdf5-format", version = "2.3.0" }
|
||||
clawhdf5-format = { path = "../clawhdf5-format", version = "2.6.0" }
|
||||
memmap2 = { version = "0.9", optional = true }
|
||||
libc = { version = "0.2", optional = true }
|
||||
tokio = { version = "1", features = ["fs", "io-util"], optional = true }
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
[package]
|
||||
name = "clawhdf5-migrate"
|
||||
version = "2.3.0"
|
||||
version = "2.6.0"
|
||||
edition = "2024"
|
||||
description = "CLI to migrate SQLite agent memory databases to HDF5 format"
|
||||
license = "MIT"
|
||||
@@ -14,9 +14,9 @@ name = "clawhdf5-migrate"
|
||||
path = "src/main.rs"
|
||||
|
||||
[dependencies]
|
||||
clawhdf5-agent = { path = "../clawhdf5-agent", version = "2.3.0" }
|
||||
clawhdf5-format = { path = "../clawhdf5-format", version = "2.3.0" }
|
||||
clawhdf5 = { path = "../clawhdf5", version = "2.3.0" }
|
||||
clawhdf5-agent = { path = "../clawhdf5-agent", version = "2.6.0" }
|
||||
clawhdf5-format = { path = "../clawhdf5-format", version = "2.6.0" }
|
||||
clawhdf5 = { path = "../clawhdf5", version = "2.6.0" }
|
||||
rusqlite = { version = "0.31", features = ["bundled"] }
|
||||
clap = { version = "4", features = ["derive"] }
|
||||
half = { workspace = true }
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
[package]
|
||||
name = "clawhdf5-napi"
|
||||
version = "2.3.0"
|
||||
version = "2.6.0"
|
||||
edition = "2024"
|
||||
description = "Node.js native addon (napi-rs) exposing clawhdf5-agent to TypeScript/JavaScript"
|
||||
license = "MIT"
|
||||
@@ -10,7 +10,7 @@ repository = "https://git.redclaw.dev/quantumclaw/clawhdf5"
|
||||
crate-type = ["cdylib"]
|
||||
|
||||
[dependencies]
|
||||
clawhdf5-agent = { path = "../clawhdf5-agent", version = "2.3.0" }
|
||||
clawhdf5-agent = { path = "../clawhdf5-agent", version = "2.6.0" }
|
||||
napi = { version = "2", default-features = false, features = ["napi9"] }
|
||||
napi-derive = "2"
|
||||
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
[package]
|
||||
name = "clawhdf5-netcdf4"
|
||||
version = "2.3.0"
|
||||
version = "2.6.0"
|
||||
edition = "2024"
|
||||
description = "NetCDF-4 read support built on rustyhdf5 — pure Rust, no C dependencies"
|
||||
license = "MIT"
|
||||
@@ -10,8 +10,8 @@ keywords = ["netcdf", "netcdf4", "hdf5", "science", "climate"]
|
||||
categories = ["parser-implementations", "science"]
|
||||
|
||||
[dependencies]
|
||||
clawhdf5 = { path = "../clawhdf5", version = "2.3.0" }
|
||||
clawhdf5-format = { path = "../clawhdf5-format", version = "2.3.0" }
|
||||
clawhdf5 = { path = "../clawhdf5", version = "2.6.0" }
|
||||
clawhdf5-format = { path = "../clawhdf5-format", version = "2.6.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,6 +1,6 @@
|
||||
[package]
|
||||
name = "clawhdf5-py"
|
||||
version = "2.3.0"
|
||||
version = "2.6.0"
|
||||
edition = "2024"
|
||||
description = "Python bindings for rustyhdf5 — a pure-Rust HDF5 library"
|
||||
license = "MIT"
|
||||
@@ -14,8 +14,8 @@ name = "clawhdf5"
|
||||
crate-type = ["cdylib", "rlib"]
|
||||
|
||||
[dependencies]
|
||||
clawhdf5_rs = { path = "../clawhdf5", version = "2.3.0", package = "clawhdf5" }
|
||||
clawhdf5-format = { path = "../clawhdf5-format", version = "2.3.0" }
|
||||
clawhdf5_rs = { path = "../clawhdf5", version = "2.6.0", package = "clawhdf5" }
|
||||
clawhdf5-format = { path = "../clawhdf5-format", version = "2.6.0" }
|
||||
pyo3 = "0.29"
|
||||
numpy = "0.29"
|
||||
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "maturin"
|
||||
|
||||
[project]
|
||||
name = "rustyhdf5"
|
||||
version = "2.3.0"
|
||||
version = "2.6.0"
|
||||
description = "Python bindings for rustyhdf5 — a pure-Rust HDF5 library"
|
||||
requires-python = ">=3.8"
|
||||
license = { text = "MIT" }
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
[package]
|
||||
name = "clawhdf5"
|
||||
version = "2.3.0"
|
||||
version = "2.6.0"
|
||||
edition = "2024"
|
||||
description = "Pure-Rust HDF5 reader/writer — no C dependencies"
|
||||
license = "MIT"
|
||||
@@ -10,16 +10,16 @@ keywords = ["hdf5", "science", "data", "binary"]
|
||||
categories = ["parser-implementations", "science", "encoding"]
|
||||
|
||||
[dependencies]
|
||||
clawhdf5-format = { path = "../clawhdf5-format", version = "2.3.0" }
|
||||
clawhdf5-io = { path = "../clawhdf5-io", version = "2.3.0" }
|
||||
clawhdf5-format = { path = "../clawhdf5-format", version = "2.6.0" }
|
||||
clawhdf5-io = { path = "../clawhdf5-io", version = "2.6.0" }
|
||||
rayon = { version = "1", optional = true }
|
||||
|
||||
[dev-dependencies]
|
||||
tempfile = { workspace = true }
|
||||
criterion = { workspace = true }
|
||||
clawhdf5-io = { path = "../clawhdf5-io", version = "2.3.0", features = ["mmap"] }
|
||||
clawhdf5-format = { path = "../clawhdf5-format", version = "2.3.0", features = ["parallel", "fast-checksum"] }
|
||||
clawhdf5-filters = { path = "../clawhdf5-filters", version = "2.3.0" }
|
||||
clawhdf5-io = { path = "../clawhdf5-io", version = "2.6.0", features = ["mmap"] }
|
||||
clawhdf5-format = { path = "../clawhdf5-format", version = "2.6.0", features = ["parallel", "fast-checksum"] }
|
||||
clawhdf5-filters = { path = "../clawhdf5-filters", version = "2.6.0" }
|
||||
|
||||
[[bench]]
|
||||
name = "mmap_bench"
|
||||
|
||||
@@ -381,8 +381,13 @@ impl<'f> Dataset<'f> {
|
||||
|
||||
/// Read all data as `f64` values.
|
||||
pub fn read_f64(&self) -> Result<Vec<f64>, Error> {
|
||||
let raw = self.read_raw()?;
|
||||
let dt = self.datatype()?;
|
||||
// A contiguous dataset is converted straight from the file bytes; going
|
||||
// through `read_raw` first copied the whole dataset an extra time.
|
||||
if let Ok(Some(bytes)) = self.read_raw_ref() {
|
||||
return Ok(data_read::read_as_f64(bytes, &dt)?);
|
||||
}
|
||||
let raw = self.read_raw()?;
|
||||
Ok(data_read::read_as_f64(&raw, &dt)?)
|
||||
}
|
||||
|
||||
@@ -393,29 +398,49 @@ impl<'f> Dataset<'f> {
|
||||
///
|
||||
/// Read all data as `f32` values.
|
||||
pub fn read_f32(&self) -> Result<Vec<f32>, Error> {
|
||||
let raw = self.read_raw()?;
|
||||
let dt = self.datatype()?;
|
||||
// A contiguous dataset is converted straight from the file bytes; going
|
||||
// through `read_raw` first copied the whole dataset an extra time.
|
||||
if let Ok(Some(bytes)) = self.read_raw_ref() {
|
||||
return Ok(data_read::read_as_f32(bytes, &dt)?);
|
||||
}
|
||||
let raw = self.read_raw()?;
|
||||
Ok(data_read::read_as_f32(&raw, &dt)?)
|
||||
}
|
||||
|
||||
/// Read all data as `i32` values.
|
||||
pub fn read_i32(&self) -> Result<Vec<i32>, Error> {
|
||||
let raw = self.read_raw()?;
|
||||
let dt = self.datatype()?;
|
||||
// A contiguous dataset is converted straight from the file bytes; going
|
||||
// through `read_raw` first copied the whole dataset an extra time.
|
||||
if let Ok(Some(bytes)) = self.read_raw_ref() {
|
||||
return Ok(data_read::read_as_i32(bytes, &dt)?);
|
||||
}
|
||||
let raw = self.read_raw()?;
|
||||
Ok(data_read::read_as_i32(&raw, &dt)?)
|
||||
}
|
||||
|
||||
/// Read all data as `i64` values.
|
||||
pub fn read_i64(&self) -> Result<Vec<i64>, Error> {
|
||||
let raw = self.read_raw()?;
|
||||
let dt = self.datatype()?;
|
||||
// A contiguous dataset is converted straight from the file bytes; going
|
||||
// through `read_raw` first copied the whole dataset an extra time.
|
||||
if let Ok(Some(bytes)) = self.read_raw_ref() {
|
||||
return Ok(data_read::read_as_i64(bytes, &dt)?);
|
||||
}
|
||||
let raw = self.read_raw()?;
|
||||
Ok(data_read::read_as_i64(&raw, &dt)?)
|
||||
}
|
||||
|
||||
/// Read all data as `u64` values.
|
||||
pub fn read_u64(&self) -> Result<Vec<u64>, Error> {
|
||||
let raw = self.read_raw()?;
|
||||
let dt = self.datatype()?;
|
||||
// A contiguous dataset is converted straight from the file bytes; going
|
||||
// through `read_raw` first copied the whole dataset an extra time.
|
||||
if let Ok(Some(bytes)) = self.read_raw_ref() {
|
||||
return Ok(data_read::read_as_u64(bytes, &dt)?);
|
||||
}
|
||||
let raw = self.read_raw()?;
|
||||
Ok(data_read::read_as_u64(&raw, &dt)?)
|
||||
}
|
||||
|
||||
@@ -458,6 +483,7 @@ impl<'f> Dataset<'f> {
|
||||
|| (matches!(dl, DataLayout::Chunked { .. })
|
||||
&& !clawhdf5_format::fill_value::is_default(fill.as_deref()));
|
||||
if fill_matters {
|
||||
clawhdf5_format::partial_read::validate(selection, &ds.dimensions)?;
|
||||
let full = self.read_raw()?;
|
||||
return Ok(data_read::extract_selection_from_buffer(
|
||||
&full,
|
||||
|
||||
@@ -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");
|
||||
@@ -913,3 +922,128 @@ with h5py.File("{dst_str}", "r") as f:
|
||||
"[(1, 2.5), (3, 4.5)] ('a', 'b') [18446744073709551615, 0, 9223372036854775808] uint64"
|
||||
);
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// h5py writes datasets indexed by a version-2 B-tree -> clawhdf5 reads
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// With `libver='latest'`, a chunked dataset with two or more unlimited
|
||||
/// dimensions indexes its chunks with a version-2 B-tree (layout v4, index
|
||||
/// type 5). These used to fail with "unsupported chunked layout".
|
||||
#[test]
|
||||
fn h5py_btree_v2_chunk_index_clawhdf5_reads() {
|
||||
skip_if_no_python!();
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = dir.path().join("bt2.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:
|
||||
a = np.arange(60 * 45, dtype="<i4").reshape(60, 45)
|
||||
f.create_dataset("plain", data=a, chunks=(7, 8), maxshape=(None, None))
|
||||
f.create_dataset("gz", data=a, chunks=(7, 8), maxshape=(None, None), compression="gzip", shuffle=True)
|
||||
# Enough chunks (2500) that the tree has internal nodes.
|
||||
big = np.arange(200 * 200, dtype="<i4").reshape(200, 200)
|
||||
f.create_dataset("deep", data=big, chunks=(4, 4), maxshape=(None, None))
|
||||
s = f.create_dataset("sparse", shape=(30, 30), dtype="<i4", chunks=(5, 5), maxshape=(None, None), fillvalue=-9)
|
||||
s[10:15, 20:25] = 4
|
||||
s[29, 29] = 1
|
||||
with h5py.File("{path_str}", "r") as f:
|
||||
print("sparse", f["sparse"][...].ravel().tolist())
|
||||
print("slab", f["deep"][37:141:13, 5:190:31].ravel().tolist())
|
||||
"#
|
||||
);
|
||||
let out = run_python_output(&script);
|
||||
let expected: std::collections::HashMap<&str, Vec<i32>> = out
|
||||
.lines()
|
||||
.map(|l| {
|
||||
let (name, list) = l.split_once(' ').unwrap();
|
||||
(name, parse_int_list(list))
|
||||
})
|
||||
.collect();
|
||||
|
||||
let file = File::open(&path).unwrap();
|
||||
let small: Vec<i32> = (0..60 * 45).collect();
|
||||
assert_eq!(file.dataset("plain").unwrap().read_i32().unwrap(), small);
|
||||
assert_eq!(file.dataset("gz").unwrap().read_i32().unwrap(), small);
|
||||
let deep: Vec<i32> = (0..200 * 200).collect();
|
||||
assert_eq!(file.dataset("deep").unwrap().read_i32().unwrap(), deep);
|
||||
assert_eq!(
|
||||
file.dataset("sparse").unwrap().read_i32().unwrap(),
|
||||
expected["sparse"]
|
||||
);
|
||||
// Partial read through the same index: rows 37,50,..,128 x cols 5,36,..,160.
|
||||
let slab = clawhdf5_format::selection::Selection::Hyperslab {
|
||||
start: vec![37, 5],
|
||||
stride: vec![13, 31],
|
||||
count: vec![8, 6],
|
||||
block: vec![1, 1],
|
||||
};
|
||||
assert_eq!(
|
||||
file.dataset("deep")
|
||||
.unwrap()
|
||||
.read_i32_selection(&slab)
|
||||
.unwrap(),
|
||||
expected["slab"]
|
||||
);
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// clawhdf5 auto-chunks a large compressed dataset -> h5py reads
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Compression without explicit chunk dimensions used to store the whole
|
||||
/// dataset as a single chunk. Large datasets are now split automatically;
|
||||
/// h5py must read the result and see sensibly sized chunks.
|
||||
#[test]
|
||||
fn clawhdf5_auto_chunked_dataset_h5py_reads() {
|
||||
skip_if_no_python!();
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = dir.path().join("auto_chunk.h5");
|
||||
let path_str = path.display().to_string();
|
||||
|
||||
let (rows, cols) = (1500u64, 1100u64); // 13.2 MB of f64
|
||||
let data: Vec<f64> = (0..rows * cols).map(|i| (i % 9973) as f64 * 0.25).collect();
|
||||
let mut builder = FileBuilder::new();
|
||||
builder
|
||||
.create_dataset("big")
|
||||
.with_f64_data(&data)
|
||||
.with_shape(&[rows, cols])
|
||||
.with_deflate(4);
|
||||
builder
|
||||
.create_dataset("small")
|
||||
.with_f64_data(&data[..600])
|
||||
.with_shape(&[20, 30])
|
||||
.with_deflate(4);
|
||||
builder.write(&path).unwrap();
|
||||
|
||||
let out = run_python_output(&format!(
|
||||
r#"
|
||||
import h5py, numpy as np
|
||||
with h5py.File("{path_str}", "r") as f:
|
||||
big, small = f["big"], f["small"]
|
||||
expect = (np.arange(1500 * 1100) % 9973) * 0.25
|
||||
ok = bool(np.array_equal(big[...].ravel(), expect)) and bool(np.array_equal(small[...].ravel(), expect[:600]))
|
||||
chunk_bytes = int(np.prod(big.chunks)) * 8
|
||||
print(ok, chunk_bytes <= 1 << 20, chunk_bytes >= 1 << 17, small.chunks == (20, 30), big.compression)
|
||||
"#
|
||||
));
|
||||
assert_eq!(out.trim(), "True True True True gzip");
|
||||
|
||||
// And it reads back here, in full and partially.
|
||||
let file = File::open(&path).unwrap();
|
||||
let ds = file.dataset("big").unwrap();
|
||||
assert_eq!(ds.read_f64().unwrap(), data);
|
||||
let row = clawhdf5_format::selection::Selection::Hyperslab {
|
||||
start: vec![777, 0],
|
||||
stride: vec![1, 1],
|
||||
count: vec![1, cols],
|
||||
block: vec![1, 1],
|
||||
};
|
||||
let start = (777 * cols) as usize;
|
||||
assert_eq!(
|
||||
ds.read_f64_selection(&row).unwrap(),
|
||||
data[start..start + cols as usize]
|
||||
);
|
||||
}
|
||||
|
||||
@@ -0,0 +1,197 @@
|
||||
//! Selection reads must return exactly what a full read followed by element
|
||||
//! extraction returns — for every layout, rank and selection shape — while
|
||||
//! touching only what the selection needs.
|
||||
|
||||
use clawhdf5::{File, FileBuilder};
|
||||
use clawhdf5_format::selection::Selection;
|
||||
|
||||
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 below(&mut self, n: u64) -> u64 {
|
||||
self.next() % n.max(1)
|
||||
}
|
||||
}
|
||||
|
||||
/// Row-major reference extraction from a full read.
|
||||
fn reference(full: &[i32], dims: &[u64], selection: &Selection) -> Vec<i32> {
|
||||
let strides: Vec<u64> = (0..dims.len())
|
||||
.map(|d| dims[d + 1..].iter().product())
|
||||
.collect();
|
||||
let at =
|
||||
|coord: &[u64]| full[coord.iter().zip(&strides).map(|(c, s)| c * s).sum::<u64>() as usize];
|
||||
match selection {
|
||||
Selection::Points(points) => points.iter().map(|p| at(p)).collect(),
|
||||
Selection::Hyperslab {
|
||||
start,
|
||||
stride,
|
||||
count,
|
||||
block,
|
||||
} => {
|
||||
// Selected indices per dimension, then their cartesian product.
|
||||
let per_dim: Vec<Vec<u64>> = (0..dims.len())
|
||||
.map(|d| {
|
||||
(0..count[d])
|
||||
.flat_map(|c| (0..block[d]).map(move |b| (c, b)))
|
||||
.map(|(c, b)| start[d] + c * stride[d] + b)
|
||||
.collect()
|
||||
})
|
||||
.collect();
|
||||
let mut out = Vec::new();
|
||||
let mut idx = vec![0usize; dims.len()];
|
||||
loop {
|
||||
let coord: Vec<u64> = idx
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(d, &i)| per_dim[d][i])
|
||||
.collect();
|
||||
out.push(at(&coord));
|
||||
let mut d = dims.len();
|
||||
loop {
|
||||
if d == 0 {
|
||||
return out;
|
||||
}
|
||||
d -= 1;
|
||||
idx[d] += 1;
|
||||
if idx[d] < per_dim[d].len() {
|
||||
break;
|
||||
}
|
||||
idx[d] = 0;
|
||||
}
|
||||
}
|
||||
}
|
||||
_ => unreachable!(),
|
||||
}
|
||||
}
|
||||
|
||||
fn random_hyperslab(rng: &mut Rng, dims: &[u64]) -> Selection {
|
||||
let mut start = Vec::new();
|
||||
let mut stride = Vec::new();
|
||||
let mut count = Vec::new();
|
||||
let mut block = Vec::new();
|
||||
for &dim in dims {
|
||||
let b = 1 + rng.below(3);
|
||||
let st = b + rng.below(4); // stride >= block: no overlap
|
||||
let s = rng.below(dim - b + 1);
|
||||
let max_count = (dim - s - b) / st + 1;
|
||||
let c = 1 + rng.below(max_count.min(6));
|
||||
start.push(s);
|
||||
stride.push(st);
|
||||
count.push(c);
|
||||
block.push(b);
|
||||
}
|
||||
Selection::Hyperslab {
|
||||
start,
|
||||
stride,
|
||||
count,
|
||||
block,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn selection_reads_match_full_reads_for_every_layout() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let mut rng = Rng(7);
|
||||
// (dims, chunk dims)
|
||||
let shapes: [(&[u64], &[u64]); 3] = [
|
||||
(&[97], &[10]),
|
||||
(&[41, 53], &[8, 9]),
|
||||
(&[11, 13, 17], &[4, 5, 6]),
|
||||
];
|
||||
for (dims, chunks) in shapes {
|
||||
let n: u64 = dims.iter().product();
|
||||
let data: Vec<i32> = (0..n as i32).map(|v| v * 3 - 7).collect();
|
||||
|
||||
let path = dir.path().join(format!("r{}.h5", dims.len()));
|
||||
let mut builder = FileBuilder::new();
|
||||
builder
|
||||
.create_dataset("contiguous")
|
||||
.with_i32_data(&data)
|
||||
.with_shape(dims);
|
||||
builder
|
||||
.create_dataset("chunked")
|
||||
.with_i32_data(&data)
|
||||
.with_shape(dims)
|
||||
.with_chunks(chunks);
|
||||
builder
|
||||
.create_dataset("deflated")
|
||||
.with_i32_data(&data)
|
||||
.with_shape(dims)
|
||||
.with_chunks(chunks)
|
||||
.with_deflate(3);
|
||||
builder.write(&path).unwrap();
|
||||
|
||||
let file = File::open(&path).unwrap();
|
||||
for name in ["contiguous", "chunked", "deflated"] {
|
||||
let ds = file.dataset(name).unwrap();
|
||||
let full = ds.read_i32().unwrap();
|
||||
assert_eq!(full, data, "{name} full read");
|
||||
|
||||
for case in 0..60 {
|
||||
let selection = if case % 5 == 4 {
|
||||
let points = (0..1 + rng.below(12))
|
||||
.map(|_| dims.iter().map(|&d| rng.below(d)).collect())
|
||||
.collect();
|
||||
Selection::Points(points)
|
||||
} else {
|
||||
random_hyperslab(&mut rng, dims)
|
||||
};
|
||||
assert_eq!(
|
||||
ds.read_i32_selection(&selection).unwrap(),
|
||||
reference(&full, dims, &selection),
|
||||
"{name} rank {} case {case}: {selection:?}",
|
||||
dims.len()
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn out_of_bounds_selections_are_errors() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = dir.path().join("oob.h5");
|
||||
let mut builder = FileBuilder::new();
|
||||
builder
|
||||
.create_dataset("d")
|
||||
.with_i32_data(&(0..100).collect::<Vec<i32>>())
|
||||
.with_shape(&[10, 10])
|
||||
.with_chunks(&[4, 4]);
|
||||
builder.write(&path).unwrap();
|
||||
let file = File::open(&path).unwrap();
|
||||
let ds = file.dataset("d").unwrap();
|
||||
let beyond = Selection::Hyperslab {
|
||||
start: vec![8, 8],
|
||||
stride: vec![1, 1],
|
||||
count: vec![5, 5],
|
||||
block: vec![1, 1],
|
||||
};
|
||||
use clawhdf5::Error;
|
||||
use clawhdf5_format::error::FormatError;
|
||||
let is_oob = |s: &Selection| {
|
||||
matches!(
|
||||
ds.read_i32_selection(s),
|
||||
Err(Error::Format(FormatError::SelectionOutOfBounds(_)))
|
||||
)
|
||||
};
|
||||
// Used to come back padded with zeros.
|
||||
assert!(is_oob(&beyond));
|
||||
// Row out of range.
|
||||
assert!(is_oob(&Selection::Points(vec![vec![10, 0]])));
|
||||
// Column out of range: used to wrap into the next row and return its value.
|
||||
assert!(is_oob(&Selection::Points(vec![vec![0, 12]])));
|
||||
// Wrong rank.
|
||||
assert!(is_oob(&Selection::Points(vec![vec![3]])));
|
||||
// In range is fine.
|
||||
assert_eq!(
|
||||
ds.read_i32_selection(&Selection::Points(vec![vec![9, 9]]))
|
||||
.unwrap(),
|
||||
[99]
|
||||
);
|
||||
}
|
||||
@@ -364,6 +364,11 @@ cargo install --path crates/clawhdf5-cli
|
||||
clawhdf5 --path agent.h5 create --agent-id my-agent --dim 384 --wal
|
||||
```
|
||||
|
||||
Add `--quantized-index` to store the vector index's copy of the embeddings as
|
||||
int8. That roughly halves a loaded store's memory at about 13% fewer queries
|
||||
per second, with recall unchanged — the query path re-scores candidates
|
||||
against the exact embeddings. The setting is recorded in the file.
|
||||
|
||||
Output:
|
||||
```json
|
||||
{
|
||||
|
||||
+45
-13
@@ -62,13 +62,21 @@ only files using the native type through the C API / h5py low-level API hit this
|
||||
|
||||
## Revised reference datatype (class 7, version 4) is not parsed
|
||||
|
||||
**Status:** open, unconfirmed against a real file.
|
||||
**Status:** fixed 2026-09-19 for object references; region and attribute
|
||||
references are recognised but not decoded.
|
||||
|
||||
**Summary:** HDF5 1.12+ `H5T_STD_REF` references use datatype version 4 with
|
||||
reference types 2–4 (object2 / region2 / attribute), which `Datatype::parse`
|
||||
rejects with `InvalidReferenceType`. h5py still writes the legacy v1
|
||||
object/region references, which read correctly, so no reproducing file has been
|
||||
generated yet; one written with the C API (`H5T_STD_REF`) is needed.
|
||||
**Summary:** HDF5 1.12+ `H5T_STD_REF` references use datatype message version 4
|
||||
with reference types 2-4 (object / region / attribute), which `Datatype::parse`
|
||||
rejected with `InvalidReferenceType`. h5py still writes the legacy references,
|
||||
so no file had been available to test against.
|
||||
|
||||
**Fix:** a real file was produced by driving the libhdf5 bundled in the h5py
|
||||
wheel through ctypes (`tests/fixtures/gen_std_ref.py` ->
|
||||
`std_ref_hdf5_2_0.h5`). The three new types parse as
|
||||
`ReferenceType::{Object2, DatasetRegion2, Attribute}`, and
|
||||
`read_object_references` decodes `Object2` elements (type, flags, token size,
|
||||
token = target object header address). External references (flag bit 0) and
|
||||
the region/attribute payloads are errors rather than misreads.
|
||||
|
||||
## `clawhdf5-gpu` `gpu_tests` can hang under the default parallel test runner
|
||||
|
||||
@@ -116,16 +124,17 @@ not reported).
|
||||
|
||||
## B-tree v2 chunk index (layout v4, index type 5) is not supported
|
||||
|
||||
**Status:** open.
|
||||
**Status:** fixed 2026-09-19.
|
||||
|
||||
**Summary:** a chunked dataset with **two or more unlimited dimensions** written
|
||||
with `libver='latest'` indexes its chunks with a version-2 B-tree. Reading it
|
||||
fails with `ChunkedReadError("unsupported chunked layout version=4,
|
||||
index_type=Some(5)")`. Single-chunk, implicit, fixed-array and
|
||||
extensible-array indexes (and the v3 B-tree v1) are supported.
|
||||
with `libver='latest'` indexes its chunks with a version-2 B-tree, and reading it
|
||||
failed with `unsupported chunked layout version=4, index_type=Some(5)`.
|
||||
|
||||
**Repro:** `f.create_dataset("d", shape=(5, 7), chunks=(2, 3), maxshape=(None, None))`
|
||||
with `h5py.File(..., libver='latest')`.
|
||||
**Fix:** record types 10 (unfiltered) and 11 (filtered) are decoded — address,
|
||||
stored size, filter mask, scaled offsets — through the shared chunk-listing
|
||||
function, so full reads, cached reads, partial reads and fill-value handling
|
||||
all work. Covered by an h5py interop test (plain, gzip+shuffle, a 2500-chunk
|
||||
tree with internal nodes, a sparse dataset with a fill value, a hyperslab).
|
||||
|
||||
## External links and external raw data are not followed
|
||||
|
||||
@@ -137,3 +146,26 @@ 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.
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@redclaw/clawhdf5",
|
||||
"version": "2.3.0",
|
||||
"version": "2.6.0",
|
||||
"description": "Node.js bindings for clawhdf5 — HDF5-backed agent memory with hippocampal consolidation",
|
||||
"main": "index.js",
|
||||
"types": "index.d.ts",
|
||||
|
||||
+26
-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=()
|
||||
@@ -63,6 +70,13 @@ run_step "cargo clippy (format feature matrix)" cargo clippy \
|
||||
--features parallel,lz4,zstd,pcodec,fast-checksum \
|
||||
-- -D warnings
|
||||
|
||||
# The HNSW index's parallel bulk build is feature-gated too.
|
||||
run_step "cargo clippy (ann parallel)" cargo clippy \
|
||||
-p clawhdf5-ann \
|
||||
--all-targets \
|
||||
--features parallel \
|
||||
-- -D warnings
|
||||
|
||||
# 4. Tests (exclude clawhdf5-py)
|
||||
run_step "cargo test" cargo test \
|
||||
--workspace \
|
||||
@@ -72,14 +86,24 @@ run_step "cargo test (format feature matrix)" cargo test \
|
||||
-p clawhdf5-format \
|
||||
--features parallel,lz4,zstd,pcodec,fast-checksum
|
||||
|
||||
run_step "cargo test (ann parallel)" cargo test \
|
||||
-p clawhdf5-ann \
|
||||
--features parallel
|
||||
|
||||
# 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