Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
4c60398b30 | ||
|
|
7155409202 | ||
|
|
64d9c5f171 | ||
|
|
84a39ef3c5 | ||
|
|
55ed87d2e8 | ||
|
|
6531158d9f | ||
|
|
aa92fef7bb | ||
|
|
29baabbed2 | ||
|
|
b23946e62d | ||
|
|
52cfcf20b2 | ||
|
|
05c665a898 | ||
|
|
b36c6ec2af | ||
|
|
f3d63dbdcd | ||
|
|
0addf328bc | ||
|
|
d668e45ab5 | ||
|
|
c6a7bbfc67 | ||
|
|
3027380979 | ||
|
|
f507803ec1 | ||
|
|
8803d0754b | ||
|
|
42c3872ec9 | ||
|
|
c19199f3eb | ||
|
|
41db450c92 |
+164
-1
@@ -28,6 +28,90 @@
|
|||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
|
## 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)
|
## Search harness baseline (v2.3.0)
|
||||||
|
|
||||||
Produced by `cargo run --release -p clawhdf5-bench --bin search_harness -- --full`
|
Produced by `cargo run --release -p clawhdf5-bench --bin search_harness -- --full`
|
||||||
@@ -243,6 +327,28 @@ results.
|
|||||||
| 10000 | 104 | 1487 | 33.8 | 13.7 | 13.9 | 0.49 | 0.51 | 2020.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 |
|
| 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
|
## Vector Search Latency
|
||||||
|
|
||||||
Brute-force cosine similarity over 384-dimensional embeddings (OpenAI text-embedding-3-small size).
|
Brute-force cosine similarity over 384-dimensional embeddings (OpenAI text-embedding-3-small size).
|
||||||
@@ -491,6 +597,62 @@ Session-level:
|
|||||||
| Vector only | 85.4% | 94.2% | 96.6% | 0.8901 |
|
| Vector only | 85.4% | 94.2% | 96.6% | 0.8901 |
|
||||||
| Hybrid | **88.2%** | **95.8%** | **97.8%** | **0.9158** |
|
| 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.
|
||||||
|
|
||||||
### Weight sweep — full haystack, n=500
|
### Weight sweep — full haystack, n=500
|
||||||
|
|
||||||
`0.7/0.3` was a documented default, never a searched one. Sweeping
|
`0.7/0.3` was a documented default, never a searched one. Sweeping
|
||||||
@@ -533,7 +695,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
|
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",
|
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,
|
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.
|
where adding BM25 to a dense retriever raised nDCG@5 while lowering Hit@1 and MRR.
|
||||||
|
|||||||
+116
@@ -1,5 +1,121 @@
|
|||||||
# Changelog
|
# Changelog
|
||||||
|
|
||||||
|
## v2.5.0 (2026-09-19)
|
||||||
|
|
||||||
|
### Upgrade Notes
|
||||||
|
- **Retrieval rankings change, for the better.** The default fusion weights
|
||||||
|
move from `0.7/0.3` to `0.4/0.6` (`hybrid::DEFAULT_FUSION`), measured over the
|
||||||
|
full LongMemEval haystack: turn-level Hit@1 51.6% vs 44.2%, MRR 0.643 vs
|
||||||
|
0.586. `unified_search` and the OpenClaw backend pick this up automatically;
|
||||||
|
callers passing weights to `hybrid_search` explicitly are unaffected.
|
||||||
|
- **Out-of-range selections are now errors.** `read_*_selection` used to return
|
||||||
|
data for a selection that ran past a dataset edge — a hyperslab came back
|
||||||
|
zero-padded, and a point with an out-of-range coordinate wrapped into the
|
||||||
|
next row. Both are now `FormatError::SelectionOutOfBounds`. Code relying on
|
||||||
|
the old (wrong) values will start seeing errors.
|
||||||
|
- **Large compressed datasets written without explicit chunk dimensions get a
|
||||||
|
different layout.** They used to be stored as one chunk; they are now split
|
||||||
|
to ~1 MiB chunks. The files stay standard and h5py-readable, and explicit
|
||||||
|
`with_chunks` is unaffected.
|
||||||
|
- `rayon` is now a default dependency of `clawhdf5-agent` (the parallel index
|
||||||
|
build). Opt out with `--no-default-features --features float16,hnsw`.
|
||||||
|
- `clawhdf5-ann` search results no longer shrink when records near the query
|
||||||
|
have been deleted, so a search that previously returned fewer than `k`
|
||||||
|
results now returns `k`.
|
||||||
|
|
||||||
|
### Retrieval quality
|
||||||
|
- `clawhdf5-agent`: optional keyword stemming — `bm25::TokenFilter::Stemmed`
|
||||||
|
and `HDF5Memory::set_token_filter`, so "training" and "trains" match. **Off
|
||||||
|
by default**, on measurement rather than principle: over the full LongMemEval
|
||||||
|
haystack it buys depth and costs the top rank (BM25 alone: Hit@5 +2.8pp,
|
||||||
|
Hit@10 +2.4pp, Hit@1 −1.8pp, MRR unchanged), and on the shipping hybrid
|
||||||
|
configuration the trade is narrower still. See `BENCHMARKS.md`.
|
||||||
|
- `clawhdf5-agent`: **`QueryExpander::expand` panicked on ordinary non-ASCII
|
||||||
|
input** — `"İ AI"` was enough. It searched a lowercased copy of the query and
|
||||||
|
then sliced the *original* with those offsets, which only works while
|
||||||
|
lowercasing preserves byte length (Turkish `İ` is 2 bytes and lowercases to
|
||||||
|
3). Depending on where the offsets drifted it either corrupted the output
|
||||||
|
("İstanbul AI trip" lost a character) or panicked. Matching now walks the
|
||||||
|
original string.
|
||||||
|
- `clawhdf5-agent`: query expansion no longer rewrites text inside words.
|
||||||
|
`replace_word_case_insensitive` did a plain substring replace despite its
|
||||||
|
name, so "training" became "trArtificial Intelligencening" and "programming"
|
||||||
|
became "Pull Requestogramming" — every acronym expansion of ordinary prose
|
||||||
|
was corrupt. Matches now require word boundaries; genuine acronyms
|
||||||
|
(`API`, `database`) still expand.
|
||||||
|
- `clawhdf5-agent`: **the default fusion weights are now the measured ones.**
|
||||||
|
A sweep of every 0.1 step over the full LongMemEval haystack (500 questions,
|
||||||
|
real MiniLM embeddings) shows the long-standing `0.7/0.3` default is
|
||||||
|
*strictly dominated* by `0.4/0.6` — turn-level Hit@1 51.6% vs 44.2%, Hit@5
|
||||||
|
81.4% vs 79.2%, Hit@10 87.8% vs 85.8%, MRR 0.643 vs 0.586, and better at
|
||||||
|
session level too. The finding was recorded in `BENCHMARKS.md` but had never
|
||||||
|
been applied: `unified_search` and the OpenClaw backend both hardcoded
|
||||||
|
`0.7/0.3`. They now use `hybrid::DEFAULT_FUSION`. **Callers passing weights
|
||||||
|
to `hybrid_search` explicitly are unaffected** — pass `0.4`/`0.6` (or use
|
||||||
|
`hybrid_search_with`) to get the tuned behaviour.
|
||||||
|
- `clawhdf5-agent`: fusion is now selectable. New `hybrid::Fusion`
|
||||||
|
(`Weighted { vector, keyword }` or `Rrf { k }`), `hybrid::fuse`,
|
||||||
|
`hybrid::hybrid_search_fused` and `HDF5Memory::hybrid_search_with`.
|
||||||
|
Reciprocal rank fusion existed but was unreachable from the store, so it had
|
||||||
|
never been measured against the weighted sum; the LongMemEval bench now has
|
||||||
|
an `RRF` mode.
|
||||||
|
|
||||||
|
### HDF5 Read Path
|
||||||
|
- **Selection reads cost what the selection costs.** `read_*_selection` decoded
|
||||||
|
the *entire* dataset and then picked elements out, so a 64 x 64 window of a
|
||||||
|
64 MB compressed dataset took 105 ms - about as long as reading all of it.
|
||||||
|
Now only the rows (contiguous) or chunks that overlap the selection's
|
||||||
|
bounding box are read and decompressed: that window takes 0.39 ms, one row
|
||||||
|
2.7 ms, one column 5.2 ms. Results are identical to the full-read path
|
||||||
|
(equivalence-tested over random hyperslabs and point lists, ranks 1-3,
|
||||||
|
contiguous / chunked / deflate). New `read_harness` bench binary.
|
||||||
|
- **Faster full reads** (same-moment A/B, 64 MB `f64`): chunked + deflate
|
||||||
|
110 -> 69 ms, chunked 72 -> 60 ms, contiguous 56 -> 30 ms. The facade's
|
||||||
|
cached read path now decompresses cache misses in parallel batches (it was
|
||||||
|
sequential; only the uncached reader was parallel) and caches only datasets
|
||||||
|
that fit the chunk cache; unfiltered chunks are copied straight from the file
|
||||||
|
bytes; a contiguous dataset is converted straight from the file bytes; and
|
||||||
|
the native-endian conversions no longer zero a buffer before overwriting it.
|
||||||
|
- **Datasets indexed by a version-2 B-tree now read** (layout v4, chunk index
|
||||||
|
type 5 — what `libver='latest'` uses for two or more unlimited dimensions;
|
||||||
|
previously "unsupported chunked layout"). The four copies of the chunk-index
|
||||||
|
dispatch are now one shared function, so every read path gets it.
|
||||||
|
- **`H5T_STD_REF` references** (HDF5 1.12+, datatype message version 4) parse:
|
||||||
|
`ReferenceType` gains `Object2`, `DatasetRegion2` and `Attribute`, and
|
||||||
|
`read_object_references` decodes the new object references. Previously any
|
||||||
|
dataset of this type failed with `InvalidReferenceType(2)`. Tested against a
|
||||||
|
file written by HDF5 2.0 itself (fixture + generator script committed).
|
||||||
|
- **Automatic chunk sizes.** Asking for compression (or any filter) without
|
||||||
|
`with_chunks` used to store the whole dataset as one chunk, so any read had
|
||||||
|
to decompress everything and nothing could be decoded in parallel. Datasets up
|
||||||
|
to 1 MiB stay a single chunk, as before; larger ones are split by halving the
|
||||||
|
dimensions in turn until a chunk is at most 1 MiB (the approach h5py takes).
|
||||||
|
**Behaviour change:** large compressed datasets written without explicit
|
||||||
|
chunk dimensions get a different (standard, h5py-readable) layout. Explicit
|
||||||
|
`with_chunks` is unaffected.
|
||||||
|
- **Out-of-range selections are errors.** They used to return data: a hyperslab
|
||||||
|
past an edge came back padded with zeros, and a point whose column was out of
|
||||||
|
range wrapped into the next row and returned that element. Now
|
||||||
|
`FormatError::SelectionOutOfBounds` (also for a rank mismatch or overlapping
|
||||||
|
blocks).
|
||||||
|
|
||||||
|
### Search
|
||||||
|
- `clawhdf5-ann`: **faster index builds.** Back-link pruning is 90% of a
|
||||||
|
build's distance evaluations; the bulk build now inserts in batches and
|
||||||
|
prunes each overflowing neighbour list once per batch (10K: 1676 -> 1074 ms).
|
||||||
|
With the `parallel` feature, planning and pruning run on a thread pool (10K:
|
||||||
|
388 ms, 100K: ~21 s -> 5.9 s on 16 cores). The graph is deterministic and
|
||||||
|
identical with or without the feature. `clawhdf5-agent`'s `parallel` feature
|
||||||
|
enables it for the agent's index and is now **on by default** (adds `rayon`
|
||||||
|
to the default dependency set; build with `--no-default-features --features
|
||||||
|
float16,hnsw` to opt out).
|
||||||
|
- `clawhdf5-ann`: `HnswIndex::search` returned fewer than `k` results — often
|
||||||
|
none — when the records nearest the query had been deleted: it collected `ef`
|
||||||
|
candidates, *then* dropped the deleted ones, *then* took `k`. Deleted nodes
|
||||||
|
are now traversed as waypoints but never occupy a result slot, so a search
|
||||||
|
returns the `k` nearest live records. Matters for any store that deletes or
|
||||||
|
supersedes memories without compacting straight away.
|
||||||
|
|
||||||
## v2.4.0 (2026-09-19)
|
## v2.4.0 (2026-09-19)
|
||||||
|
|
||||||
### Upgrade Notes
|
### Upgrade Notes
|
||||||
|
|||||||
@@ -33,6 +33,8 @@ 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 approximate `clawhdf5-ann` index for the vector stage (the index mirrors
|
||||||
the cache and self-heals on drift). Build the agent with
|
the cache and self-heals on drift). Build the agent with
|
||||||
`--no-default-features --features float16` to force the exact linear cosine scan.
|
`--no-default-features --features float16` to force the exact linear cosine scan.
|
||||||
|
The agent's `parallel` feature (also default) builds the index on a thread
|
||||||
|
pool; the graph is identical with or without it.
|
||||||
The index uses the HNSW paper's diversity heuristic for neighbour selection
|
The index uses the HNSW paper's diversity heuristic for neighbour selection
|
||||||
(plain closest-M capped recall on clustered data: 0.31 recall@10 at 100K). Its
|
(plain closest-M capped recall on clustered data: 0.31 recall@10 at 100K). Its
|
||||||
graph is saved to `<store>.h5.ann` at each checkpoint and reloaded by `open()`
|
graph is saved to `<store>.h5.ann` at each checkpoint and reloaded by `open()`
|
||||||
|
|||||||
+1
-1
@@ -21,7 +21,7 @@ members = [
|
|||||||
resolver = "2"
|
resolver = "2"
|
||||||
|
|
||||||
[workspace.package]
|
[workspace.package]
|
||||||
version = "2.4.0"
|
version = "2.5.0"
|
||||||
edition = "2024"
|
edition = "2024"
|
||||||
license = "MIT"
|
license = "MIT"
|
||||||
repository = "https://git.redclaw.dev/quantumclaw/clawhdf5"
|
repository = "https://git.redclaw.dev/quantumclaw/clawhdf5"
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
[package]
|
[package]
|
||||||
name = "clawhdf5-accel"
|
name = "clawhdf5-accel"
|
||||||
version = "2.4.0"
|
version = "2.5.0"
|
||||||
edition = "2024"
|
edition = "2024"
|
||||||
description = "SIMD-accelerated operations for rustyhdf5"
|
description = "SIMD-accelerated operations for rustyhdf5"
|
||||||
license = "MIT"
|
license = "MIT"
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
[package]
|
[package]
|
||||||
name = "clawhdf5-agent"
|
name = "clawhdf5-agent"
|
||||||
version = "2.4.0"
|
version = "2.5.0"
|
||||||
edition = "2024"
|
edition = "2024"
|
||||||
description = "HDF5-backed persistent memory store for on-device AI agents"
|
description = "HDF5-backed persistent memory store for on-device AI agents"
|
||||||
license = "MIT"
|
license = "MIT"
|
||||||
@@ -10,12 +10,12 @@ keywords = ["agent", "memory", "hdf5", "vector-search", "embedding"]
|
|||||||
categories = ["database", "science", "algorithms"]
|
categories = ["database", "science", "algorithms"]
|
||||||
|
|
||||||
[dependencies]
|
[dependencies]
|
||||||
clawhdf5-format = { path = "../clawhdf5-format", version = "2.4.0", features = ["parallel", "fast-checksum"] }
|
clawhdf5-format = { path = "../clawhdf5-format", version = "2.5.0", features = ["parallel", "fast-checksum"] }
|
||||||
clawhdf5 = { path = "../clawhdf5", version = "2.4.0" }
|
clawhdf5 = { path = "../clawhdf5", version = "2.5.0" }
|
||||||
clawhdf5-io = { path = "../clawhdf5-io", version = "2.4.0", features = ["mmap"] }
|
clawhdf5-io = { path = "../clawhdf5-io", version = "2.5.0", features = ["mmap"] }
|
||||||
clawhdf5-accel = { path = "../clawhdf5-accel", version = "2.4.0" }
|
clawhdf5-accel = { path = "../clawhdf5-accel", version = "2.5.0" }
|
||||||
clawhdf5-ann = { path = "../clawhdf5-ann", version = "2.4.0", optional = true }
|
clawhdf5-ann = { path = "../clawhdf5-ann", version = "2.5.0", optional = true }
|
||||||
clawhdf5-gpu = { path = "../clawhdf5-gpu", version = "2.4.0", optional = true, default-features = false }
|
clawhdf5-gpu = { path = "../clawhdf5-gpu", version = "2.5.0", optional = true, default-features = false }
|
||||||
serde = { workspace = true }
|
serde = { workspace = true }
|
||||||
byteorder = "1"
|
byteorder = "1"
|
||||||
half = { workspace = true, optional = true }
|
half = { workspace = true, optional = true }
|
||||||
@@ -45,9 +45,11 @@ name = "memory_bench"
|
|||||||
harness = false
|
harness = false
|
||||||
|
|
||||||
[features]
|
[features]
|
||||||
default = ["float16", "hnsw"]
|
default = ["float16", "hnsw", "parallel"]
|
||||||
float16 = ["half"]
|
float16 = ["half"]
|
||||||
parallel = ["rayon"]
|
# Rayon-parallel brute-force search strategies, and a parallel bulk build of
|
||||||
|
# the HNSW index (same graph, several times faster on a multi-core machine).
|
||||||
|
parallel = ["rayon", "clawhdf5-ann?/parallel"]
|
||||||
# Compress embeddings with Zstd instead of deflate when
|
# Compress embeddings with Zstd instead of deflate when
|
||||||
# `MemoryConfig::compression` is on. Off by default: it links libzstd (C).
|
# `MemoryConfig::compression` is on. Off by default: it links libzstd (C).
|
||||||
zstd = ["clawhdf5/zstd"]
|
zstd = ["clawhdf5/zstd"]
|
||||||
|
|||||||
@@ -37,7 +37,7 @@
|
|||||||
//! let mem = AsyncHDF5Memory::open_with(path, config).await?;
|
//! let mem = AsyncHDF5Memory::open_with(path, config).await?;
|
||||||
//! mem.save(entry).await?; // buffered → background writer
|
//! mem.save(entry).await?; // buffered → background writer
|
||||||
//! mem.save_batch(entries).await?; // also buffered
|
//! mem.save_batch(entries).await?; // also buffered
|
||||||
//! let results = mem.hybrid_search(emb, "query".into(), 0.7, 0.3, 5).await;
|
//! let results = mem.hybrid_search(emb, "query".into(), 0.4, 0.6, 5).await;
|
||||||
//! mem.shutdown().await?; // final flush + stop
|
//! mem.shutdown().await?; // final flush + stop
|
||||||
//! ```
|
//! ```
|
||||||
|
|
||||||
|
|||||||
@@ -59,11 +59,18 @@ pub struct BM25Index {
|
|||||||
k1: f32,
|
k1: f32,
|
||||||
/// BM25 b parameter.
|
/// BM25 b parameter.
|
||||||
b: f32,
|
b: f32,
|
||||||
|
/// Applied to every document and query token, so the two always agree.
|
||||||
|
filter: TokenFilter,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl BM25Index {
|
impl BM25Index {
|
||||||
/// Build a BM25 index from a set of documents, excluding tombstoned entries.
|
/// Build a BM25 index from a set of documents, excluding tombstoned entries.
|
||||||
pub fn build(documents: &[String], tombstones: &[u8]) -> Self {
|
pub fn build(documents: &[String], tombstones: &[u8]) -> Self {
|
||||||
|
Self::build_with(documents, tombstones, TokenFilter::default())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// [`BM25Index::build`] with the token filter chosen explicitly.
|
||||||
|
pub fn build_with(documents: &[String], tombstones: &[u8], filter: TokenFilter) -> Self {
|
||||||
let mut index = Self {
|
let mut index = Self {
|
||||||
inverted: HashMap::new(),
|
inverted: HashMap::new(),
|
||||||
doc_lengths: vec![0; documents.len()],
|
doc_lengths: vec![0; documents.len()],
|
||||||
@@ -72,6 +79,7 @@ impl BM25Index {
|
|||||||
num_docs: 0,
|
num_docs: 0,
|
||||||
k1: DEFAULT_K1,
|
k1: DEFAULT_K1,
|
||||||
b: DEFAULT_B,
|
b: DEFAULT_B,
|
||||||
|
filter,
|
||||||
};
|
};
|
||||||
index.index_documents(documents, tombstones);
|
index.index_documents(documents, tombstones);
|
||||||
index
|
index
|
||||||
@@ -120,7 +128,7 @@ impl BM25Index {
|
|||||||
// add/remove, and costs one `ln` per query term.
|
// add/remove, and costs one `ln` per query term.
|
||||||
let mut acc = vec![0.0f32; self.doc_lengths.len()];
|
let mut acc = vec![0.0f32; self.doc_lengths.len()];
|
||||||
let mut matched = false;
|
let mut matched = false;
|
||||||
for token in tokenize(query) {
|
for token in tokenize_with(query, self.filter) {
|
||||||
let Some(postings) = self.inverted.get(token.as_str()) else {
|
let Some(postings) = self.inverted.get(token.as_str()) else {
|
||||||
continue;
|
continue;
|
||||||
};
|
};
|
||||||
@@ -146,6 +154,11 @@ impl BM25Index {
|
|||||||
.collect()
|
.collect()
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// The token filter this index was built with.
|
||||||
|
pub fn token_filter(&self) -> TokenFilter {
|
||||||
|
self.filter
|
||||||
|
}
|
||||||
|
|
||||||
/// Number of document slots (live or not) the index covers. Ids are
|
/// Number of document slots (live or not) the index covers. Ids are
|
||||||
/// positions in the document list it mirrors.
|
/// positions in the document list it mirrors.
|
||||||
pub fn len(&self) -> usize {
|
pub fn len(&self) -> usize {
|
||||||
@@ -168,7 +181,7 @@ impl BM25Index {
|
|||||||
}
|
}
|
||||||
debug_assert_eq!(self.doc_lengths[doc_id], 0, "slot {doc_id} is occupied");
|
debug_assert_eq!(self.doc_lengths[doc_id], 0, "slot {doc_id} is occupied");
|
||||||
|
|
||||||
let tokens = tokenize(text);
|
let tokens = tokenize_with(text, self.filter);
|
||||||
let mut term_freqs: HashMap<&str, u32> = HashMap::new();
|
let mut term_freqs: HashMap<&str, u32> = HashMap::new();
|
||||||
for token in &tokens {
|
for token in &tokens {
|
||||||
*term_freqs.entry(token).or_insert(0) += 1;
|
*term_freqs.entry(token).or_insert(0) += 1;
|
||||||
@@ -201,7 +214,7 @@ impl BM25Index {
|
|||||||
/// Remove document `doc_id`, whose indexed text was `text`. The text is
|
/// Remove document `doc_id`, whose indexed text was `text`. The text is
|
||||||
/// needed to find its postings; pass exactly what was added.
|
/// needed to find its postings; pass exactly what was added.
|
||||||
pub fn remove_document(&mut self, doc_id: usize, text: &str) {
|
pub fn remove_document(&mut self, doc_id: usize, text: &str) {
|
||||||
let tokens = tokenize(text);
|
let tokens = tokenize_with(text, self.filter);
|
||||||
let mut seen: std::collections::HashSet<&str> = std::collections::HashSet::new();
|
let mut seen: std::collections::HashSet<&str> = std::collections::HashSet::new();
|
||||||
for token in &tokens {
|
for token in &tokens {
|
||||||
if !seen.insert(token) {
|
if !seen.insert(token) {
|
||||||
@@ -252,7 +265,7 @@ impl BM25Index {
|
|||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
|
|
||||||
let tokens = tokenize(doc);
|
let tokens = tokenize_with(doc, self.filter);
|
||||||
let doc_len = tokens.len() as u32;
|
let doc_len = tokens.len() as u32;
|
||||||
self.doc_lengths[i] = doc_len;
|
self.doc_lengths[i] = doc_len;
|
||||||
total_length += doc_len as u64;
|
total_length += doc_len as u64;
|
||||||
@@ -285,11 +298,86 @@ impl BM25Index {
|
|||||||
|
|
||||||
/// Tokenize a string: lowercase, split on non-alphanumeric characters,
|
/// Tokenize a string: lowercase, split on non-alphanumeric characters,
|
||||||
/// filter empty tokens.
|
/// filter empty tokens.
|
||||||
|
/// What [`tokenize_with`] does to each token after splitting.
|
||||||
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
|
||||||
|
pub enum TokenFilter {
|
||||||
|
/// Lowercase and split only — the original behaviour.
|
||||||
|
#[default]
|
||||||
|
Plain,
|
||||||
|
/// Also strip common English inflections, so "running" and "runs" match
|
||||||
|
/// "run". Conservative on purpose: only plural and past/continuous verb
|
||||||
|
/// endings, and only on tokens long enough that stripping leaves a real
|
||||||
|
/// stem. A stemmer earns its keep by conflating *related* words; an
|
||||||
|
/// aggressive one also conflates unrelated ones ("universe"/"university"),
|
||||||
|
/// which costs precision.
|
||||||
|
Stemmed,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Strip common English inflections from an already-lowercased token.
|
||||||
|
///
|
||||||
|
/// Applied identically to documents and queries, so the pair only has to agree
|
||||||
|
/// with itself — the stem need not be a real word.
|
||||||
|
fn stem(token: &str) -> &str {
|
||||||
|
// Below this, stripping does more harm than good ("bed" -> "b").
|
||||||
|
const MIN_STEM: usize = 4;
|
||||||
|
let strip = |suffix: &str, min_len: usize| -> Option<&str> {
|
||||||
|
let stem = token.strip_suffix(suffix)?;
|
||||||
|
(stem.len() >= min_len).then_some(stem)
|
||||||
|
};
|
||||||
|
|
||||||
|
// Plurals first: "studies" -> "studi", "classes" -> "class", "cats" -> "cat".
|
||||||
|
// "ies" keeps its "i" so the result meets "-ied" ("studied" -> "studi").
|
||||||
|
if let Some(stem) = strip("ies", 2) {
|
||||||
|
return &token[..stem.len() + 1];
|
||||||
|
}
|
||||||
|
for suffix in ["sses", "shes", "ches", "xes", "zes"] {
|
||||||
|
if let Some(stem) = strip(suffix, MIN_STEM - 1) {
|
||||||
|
// Keep the sibilant: "classes" -> "class", not "clas".
|
||||||
|
return &token[..stem.len() + 2];
|
||||||
|
}
|
||||||
|
}
|
||||||
|
// Verb endings before the bare plural, so "raced" doesn't become "raced".
|
||||||
|
if let Some(stem) = strip("ing", MIN_STEM - 1).or_else(|| strip("ed", MIN_STEM - 1)) {
|
||||||
|
return undouble(stem);
|
||||||
|
}
|
||||||
|
if !token.ends_with("ss")
|
||||||
|
&& !token.ends_with("us")
|
||||||
|
&& !token.ends_with("is")
|
||||||
|
&& let Some(stem) = strip("s", MIN_STEM - 1)
|
||||||
|
{
|
||||||
|
return stem;
|
||||||
|
}
|
||||||
|
token
|
||||||
|
}
|
||||||
|
|
||||||
|
/// "runn" -> "run": undo the consonant doubling that "-ing"/"-ed" introduce.
|
||||||
|
fn undouble(stem: &str) -> &str {
|
||||||
|
let mut chars = stem.chars().rev();
|
||||||
|
let (Some(last), Some(prev)) = (chars.next(), chars.next()) else {
|
||||||
|
return stem;
|
||||||
|
};
|
||||||
|
let doubled = last == prev && !"aeiou".contains(last) && last.is_ascii_alphabetic();
|
||||||
|
if doubled && stem.len() > 3 {
|
||||||
|
&stem[..stem.len() - 1]
|
||||||
|
} else {
|
||||||
|
stem
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
fn tokenize(text: &str) -> Vec<String> {
|
fn tokenize(text: &str) -> Vec<String> {
|
||||||
|
tokenize_with(text, TokenFilter::Plain)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Split `text` into scoring tokens under `filter`.
|
||||||
|
pub fn tokenize_with(text: &str, filter: TokenFilter) -> Vec<String> {
|
||||||
text.to_lowercase()
|
text.to_lowercase()
|
||||||
.split(|c: char| !c.is_alphanumeric())
|
.split(|c: char| !c.is_alphanumeric())
|
||||||
.filter(|s| !s.is_empty())
|
.filter(|s| !s.is_empty())
|
||||||
.map(|s| s.to_string())
|
.map(|token| match filter {
|
||||||
|
TokenFilter::Plain => token.to_string(),
|
||||||
|
TokenFilter::Stemmed => stem(token).to_string(),
|
||||||
|
})
|
||||||
.collect()
|
.collect()
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -599,6 +687,53 @@ mod tests {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn stemming_conflates_inflections_of_the_same_word() {
|
||||||
|
let stem_of = |w: &str| tokenize_with(w, TokenFilter::Stemmed).pop().unwrap();
|
||||||
|
// Pairs that should meet.
|
||||||
|
for (a, b) in [
|
||||||
|
("running", "runs"),
|
||||||
|
("trained", "training"),
|
||||||
|
("miles", "mile"),
|
||||||
|
("studies", "studied"),
|
||||||
|
("mentioned", "mentioning"),
|
||||||
|
("classes", "class"),
|
||||||
|
("planned", "planning"),
|
||||||
|
] {
|
||||||
|
assert_eq!(stem_of(a), stem_of(b), "{a} / {b} should share a stem");
|
||||||
|
}
|
||||||
|
// Pairs that must stay apart. Note which pairs are deliberately absent:
|
||||||
|
// "bed"/"bedding" and "gas"/"gassed" both collapse to one stem, which
|
||||||
|
// is what Porter does too and is right — they are related words.
|
||||||
|
for (a, b) in [
|
||||||
|
("universe", "university"),
|
||||||
|
("business", "busy"),
|
||||||
|
("this", "thing"),
|
||||||
|
] {
|
||||||
|
assert_ne!(stem_of(a), stem_of(b), "{a} / {b} must not be conflated");
|
||||||
|
}
|
||||||
|
// Short words and non-inflections are left alone.
|
||||||
|
for word in ["run", "bus", "is", "his", "data", "gas"] {
|
||||||
|
assert_eq!(stem_of(word), word, "{word} should be untouched");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn stemming_is_off_by_default_and_applied_consistently() {
|
||||||
|
assert_eq!(tokenize("Running miles"), ["running", "miles"]);
|
||||||
|
assert_eq!(
|
||||||
|
tokenize_with("Running miles", TokenFilter::Stemmed),
|
||||||
|
["run", "mile"]
|
||||||
|
);
|
||||||
|
|
||||||
|
// A query inflected differently from the document still matches.
|
||||||
|
let docs = vec!["I ran while training for the marathon".to_string()];
|
||||||
|
let plain = BM25Index::build_with(&docs, &[0], TokenFilter::Plain);
|
||||||
|
let stemmed = BM25Index::build_with(&docs, &[0], TokenFilter::Stemmed);
|
||||||
|
assert!(plain.search("trains", 1).is_empty());
|
||||||
|
assert_eq!(stemmed.search("trains", 1).len(), 1);
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn ties_break_towards_the_lower_doc_id() {
|
fn ties_break_towards_the_lower_doc_id() {
|
||||||
let docs: Vec<String> = (0..6).map(|_| "same text".to_string()).collect();
|
let docs: Vec<String> = (0..6).map(|_| "same text".to_string()).collect();
|
||||||
|
|||||||
@@ -29,12 +29,39 @@ pub fn hybrid_search(
|
|||||||
query_embedding: &[f32],
|
query_embedding: &[f32],
|
||||||
query_text: &str,
|
query_text: &str,
|
||||||
vectors: &[Vec<f32>],
|
vectors: &[Vec<f32>],
|
||||||
_chunks: &[String],
|
chunks: &[String],
|
||||||
tombstones: &[u8],
|
tombstones: &[u8],
|
||||||
bm25_index: &BM25Index,
|
bm25_index: &BM25Index,
|
||||||
vector_weight: f32,
|
vector_weight: f32,
|
||||||
keyword_weight: f32,
|
keyword_weight: f32,
|
||||||
k: usize,
|
k: usize,
|
||||||
|
) -> Vec<(usize, f32)> {
|
||||||
|
hybrid_search_fused(
|
||||||
|
query_embedding,
|
||||||
|
query_text,
|
||||||
|
vectors,
|
||||||
|
chunks,
|
||||||
|
tombstones,
|
||||||
|
bm25_index,
|
||||||
|
Fusion::Weighted {
|
||||||
|
vector: vector_weight,
|
||||||
|
keyword: keyword_weight,
|
||||||
|
},
|
||||||
|
k,
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// [`hybrid_search`] with the fusion method chosen explicitly.
|
||||||
|
#[allow(clippy::too_many_arguments)]
|
||||||
|
pub fn hybrid_search_fused(
|
||||||
|
query_embedding: &[f32],
|
||||||
|
query_text: &str,
|
||||||
|
vectors: &[Vec<f32>],
|
||||||
|
_chunks: &[String],
|
||||||
|
tombstones: &[u8],
|
||||||
|
bm25_index: &BM25Index,
|
||||||
|
fusion: Fusion,
|
||||||
|
k: usize,
|
||||||
) -> Vec<(usize, f32)> {
|
) -> Vec<(usize, f32)> {
|
||||||
// Get raw scores from both systems. Request all results so normalization
|
// Get raw scores from both systems. Request all results so normalization
|
||||||
// covers the full distribution.
|
// covers the full distribution.
|
||||||
@@ -60,7 +87,7 @@ pub fn hybrid_search(
|
|||||||
};
|
};
|
||||||
let kw_scores = bm25_index.scores(query_text);
|
let kw_scores = bm25_index.scores(query_text);
|
||||||
|
|
||||||
merge_vector_keyword(vec_scores, kw_scores, vector_weight, keyword_weight, k)
|
fuse(vec_scores, kw_scores, fusion, k)
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Merge pre-computed vector-similarity and keyword scores into a single ranking.
|
/// Merge pre-computed vector-similarity and keyword scores into a single ranking.
|
||||||
@@ -76,18 +103,92 @@ pub fn merge_vector_keyword(
|
|||||||
keyword_weight: f32,
|
keyword_weight: f32,
|
||||||
k: usize,
|
k: usize,
|
||||||
) -> Vec<(usize, f32)> {
|
) -> Vec<(usize, f32)> {
|
||||||
// Normalize each set to [0, 1].
|
fuse(
|
||||||
let vec_normalized = normalize_scores(&vec_scores);
|
vec_scores,
|
||||||
let kw_normalized = normalize_scores(&kw_scores);
|
kw_scores,
|
||||||
|
Fusion::Weighted {
|
||||||
|
vector: vector_weight,
|
||||||
|
keyword: keyword_weight,
|
||||||
|
},
|
||||||
|
k,
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
// Merge scores with weights.
|
/// How the vector and keyword stages are combined into one ranking.
|
||||||
let mut merged: HashMap<usize, f32> = HashMap::new();
|
#[derive(Debug, Clone, Copy, PartialEq)]
|
||||||
|
pub enum Fusion {
|
||||||
|
/// Min-max normalise each stage over its own candidates, then take a
|
||||||
|
/// weighted sum. Uses the *scores*, so a stage that separates its
|
||||||
|
/// candidates sharply keeps that separation — and a stage whose candidates
|
||||||
|
/// are all near-identical contributes little.
|
||||||
|
Weighted {
|
||||||
|
/// Weight on the vector stage.
|
||||||
|
vector: f32,
|
||||||
|
/// Weight on the keyword stage.
|
||||||
|
keyword: f32,
|
||||||
|
},
|
||||||
|
/// Reciprocal rank fusion: each stage contributes `1 / (k + rank)`,
|
||||||
|
/// ignoring score magnitudes entirely. Robust when the two stages'
|
||||||
|
/// scores aren't comparable, at the cost of discarding confidence.
|
||||||
|
Rrf {
|
||||||
|
/// The rank-damping constant; 60 is the value from the original paper.
|
||||||
|
k: f32,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
for (idx, score) in &vec_normalized {
|
impl Default for Fusion {
|
||||||
*merged.entry(*idx).or_insert(0.0) += vector_weight * score;
|
fn default() -> Self {
|
||||||
|
DEFAULT_FUSION
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 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);
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
for (idx, score) in &kw_normalized {
|
|
||||||
*merged.entry(*idx).or_insert(0.0) += keyword_weight * score;
|
|
||||||
}
|
}
|
||||||
|
|
||||||
let mut results: Vec<(usize, f32)> = merged.into_iter().collect();
|
let mut results: Vec<(usize, f32)> = merged.into_iter().collect();
|
||||||
@@ -353,6 +454,49 @@ mod tests {
|
|||||||
assert_eq!(result[0].1, 1.0);
|
assert_eq!(result[0].1, 1.0);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn default_fusion_is_the_tuned_operating_point() {
|
||||||
|
// A sweep over the full LongMemEval haystack found 0.7/0.3 strictly
|
||||||
|
// dominated by 0.4/0.6 (BENCHMARKS.md). This guards the finding
|
||||||
|
// against being quietly undone.
|
||||||
|
assert_eq!(
|
||||||
|
DEFAULT_FUSION,
|
||||||
|
Fusion::Weighted {
|
||||||
|
vector: 0.4,
|
||||||
|
keyword: 0.6
|
||||||
|
}
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn rrf_rewards_agreement_between_the_stages_and_ignores_magnitudes() {
|
||||||
|
// Doc 1 is second-best in both stages; doc 0 is best in one and absent
|
||||||
|
// from the other. RRF prefers the doc both stages liked.
|
||||||
|
let vec_scores = vec![(0, 100.0), (1, 0.9)];
|
||||||
|
let kw_scores = vec![(2, 5.0), (1, 4.9)];
|
||||||
|
let ranked = fuse(vec_scores, kw_scores, Fusion::Rrf { k: 60.0 }, 3);
|
||||||
|
assert_eq!(ranked[0].0, 1, "{ranked:?}");
|
||||||
|
|
||||||
|
// Scaling one stage's scores cannot change an RRF ranking, only the
|
||||||
|
// order within that stage can.
|
||||||
|
let a = fuse(
|
||||||
|
vec![(0, 1.0), (1, 0.5)],
|
||||||
|
vec![(1, 2.0), (0, 1.0)],
|
||||||
|
Fusion::Rrf { k: 60.0 },
|
||||||
|
2,
|
||||||
|
);
|
||||||
|
let b = fuse(
|
||||||
|
vec![(0, 1e6), (1, -3.0)],
|
||||||
|
vec![(1, 0.002), (0, 0.001)],
|
||||||
|
Fusion::Rrf { k: 60.0 },
|
||||||
|
2,
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
a.iter().map(|r| r.0).collect::<Vec<_>>(),
|
||||||
|
b.iter().map(|r| r.0).collect::<Vec<_>>()
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn merge_top_k_matches_a_full_sort() {
|
fn merge_top_k_matches_a_full_sort() {
|
||||||
// Many ties (scores repeat) so the index tie-break is exercised.
|
// Many ties (scores repeat) so the index tie-break is exercised.
|
||||||
|
|||||||
@@ -258,6 +258,9 @@ pub struct HDF5Memory {
|
|||||||
/// every record, on every single query. Built lazily on first use; see
|
/// every record, on every single query. Built lazily on first use; see
|
||||||
/// [`HDF5Memory::ensure_bm25_fresh`] for how it stays in sync.
|
/// [`HDF5Memory::ensure_bm25_fresh`] for how it stays in sync.
|
||||||
bm25: Option<bm25::BM25Index>,
|
bm25: Option<bm25::BM25Index>,
|
||||||
|
/// Token filter the keyword index is built with. Changing it drops the
|
||||||
|
/// index; it is not persisted, because the index is not either.
|
||||||
|
bm25_filter: bm25::TokenFilter,
|
||||||
/// Activation weights changed since the last checkpoint (searches boost
|
/// Activation weights changed since the last checkpoint (searches boost
|
||||||
/// the records they return). Cleared by `flush`.
|
/// the records they return). Cleared by `flush`.
|
||||||
activations_dirty: bool,
|
activations_dirty: bool,
|
||||||
@@ -314,6 +317,7 @@ impl HDF5Memory {
|
|||||||
anomaly: anomaly::WriteAnomalyDetector::new(anomaly::AnomalyConfig::default()),
|
anomaly: anomaly::WriteAnomalyDetector::new(anomaly::AnomalyConfig::default()),
|
||||||
anomaly_alerts: Vec::new(),
|
anomaly_alerts: Vec::new(),
|
||||||
bm25: None,
|
bm25: None,
|
||||||
|
bm25_filter: bm25::TokenFilter::default(),
|
||||||
activations_dirty: false,
|
activations_dirty: false,
|
||||||
read_only: false,
|
read_only: false,
|
||||||
quarantined_wal: None,
|
quarantined_wal: None,
|
||||||
@@ -481,6 +485,7 @@ impl HDF5Memory {
|
|||||||
anomaly: anomaly::WriteAnomalyDetector::new(anomaly::AnomalyConfig::default()),
|
anomaly: anomaly::WriteAnomalyDetector::new(anomaly::AnomalyConfig::default()),
|
||||||
anomaly_alerts: Vec::new(),
|
anomaly_alerts: Vec::new(),
|
||||||
bm25: None,
|
bm25: None,
|
||||||
|
bm25_filter: bm25::TokenFilter::default(),
|
||||||
activations_dirty: false,
|
activations_dirty: false,
|
||||||
read_only,
|
read_only,
|
||||||
quarantined_wal,
|
quarantined_wal,
|
||||||
@@ -591,7 +596,7 @@ impl HDF5Memory {
|
|||||||
pub(crate) fn ensure_bm25_fresh(&mut self) -> &bm25::BM25Index {
|
pub(crate) fn ensure_bm25_fresh(&mut self) -> &bm25::BM25Index {
|
||||||
let n = self.cache.chunks.len();
|
let n = self.cache.chunks.len();
|
||||||
let bm25 = match self.bm25.take() {
|
let bm25 = match self.bm25.take() {
|
||||||
Some(index) if index.len() <= n => {
|
Some(index) if index.len() <= n && index.token_filter() == self.bm25_filter => {
|
||||||
let mut index = index;
|
let mut index = index;
|
||||||
for id in index.len()..n {
|
for id in index.len()..n {
|
||||||
if self.cache.tombstones[id] == 0 {
|
if self.cache.tombstones[id] == 0 {
|
||||||
@@ -601,11 +606,26 @@ impl HDF5Memory {
|
|||||||
index.pad_to(n);
|
index.pad_to(n);
|
||||||
index
|
index
|
||||||
}
|
}
|
||||||
_ => bm25::BM25Index::build(&self.cache.chunks, &self.cache.tombstones),
|
_ => bm25::BM25Index::build_with(
|
||||||
|
&self.cache.chunks,
|
||||||
|
&self.cache.tombstones,
|
||||||
|
self.bm25_filter,
|
||||||
|
),
|
||||||
};
|
};
|
||||||
self.bm25.insert(bm25)
|
self.bm25.insert(bm25)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Choose how keyword-search tokens are normalised, rebuilding the index
|
||||||
|
/// on next use. [`bm25::TokenFilter::Stemmed`] matches inflections of the
|
||||||
|
/// same word at some cost in precision; measure before adopting it (see
|
||||||
|
/// `BENCHMARKS.md`).
|
||||||
|
pub fn set_token_filter(&mut self, filter: bm25::TokenFilter) {
|
||||||
|
if filter != self.bm25_filter {
|
||||||
|
self.bm25_filter = filter;
|
||||||
|
self.bm25 = None;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
/// Record `id` was tombstoned; its text is still in the cache.
|
/// Record `id` was tombstoned; its text is still in the cache.
|
||||||
fn bm25_on_delete(&mut self, id: usize) {
|
fn bm25_on_delete(&mut self, id: usize) {
|
||||||
if let Some(index) = self.bm25.as_mut()
|
if let Some(index) = self.bm25.as_mut()
|
||||||
@@ -1350,7 +1370,8 @@ impl HDF5Memory {
|
|||||||
k: usize,
|
k: usize,
|
||||||
) -> Vec<SearchResult> {
|
) -> Vec<SearchResult> {
|
||||||
// Persistent tier.
|
// Persistent tier.
|
||||||
let persistent = self.hybrid_search(query_embedding, query_text, 0.7, 0.3, k);
|
let persistent =
|
||||||
|
self.hybrid_search_with(query_embedding, query_text, hybrid::DEFAULT_FUSION, k);
|
||||||
const EPHEMERAL_BOOST: f32 = 1.2;
|
const EPHEMERAL_BOOST: f32 = 1.2;
|
||||||
let mut results = persistent;
|
let mut results = persistent;
|
||||||
|
|
||||||
@@ -1954,6 +1975,34 @@ mod tests {
|
|||||||
assert_eq!(top_ids(&mut reopened, &q), expected_after);
|
assert_eq!(top_ids(&mut reopened, &q), expected_after);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn set_token_filter_rebuilds_the_keyword_index() {
|
||||||
|
let dir = TempDir::new().unwrap();
|
||||||
|
let mut mem = HDF5Memory::create(make_config(&dir)).unwrap();
|
||||||
|
mem.save(make_entry(
|
||||||
|
"I was training for a marathon",
|
||||||
|
&[1.0, 0.0, 0.0, 0.0],
|
||||||
|
))
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
// Count only genuine keyword matches: `hybrid_search` also returns
|
||||||
|
// zero-score filler when fewer than k records are relevant.
|
||||||
|
let hits = |mem: &mut HDF5Memory| {
|
||||||
|
mem.hybrid_search(&[0.0, 0.0, 0.0, 0.0], "trains", 0.0, 1.0, 5)
|
||||||
|
.iter()
|
||||||
|
.filter(|r| r.score > 0.0)
|
||||||
|
.count()
|
||||||
|
};
|
||||||
|
assert_eq!(hits(&mut mem), 0);
|
||||||
|
|
||||||
|
mem.set_token_filter(bm25::TokenFilter::Stemmed);
|
||||||
|
assert_eq!(hits(&mut mem), 1, "index should have been rebuilt stemmed");
|
||||||
|
|
||||||
|
// And back, rebuilding again.
|
||||||
|
mem.set_token_filter(bm25::TokenFilter::Plain);
|
||||||
|
assert_eq!(hits(&mut mem), 0);
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn keyword_index_stays_in_sync_through_every_mutation() {
|
fn keyword_index_stays_in_sync_through_every_mutation() {
|
||||||
let dir = TempDir::new().unwrap();
|
let dir = TempDir::new().unwrap();
|
||||||
|
|||||||
@@ -531,11 +531,14 @@ impl MemoryBackend for ClawhdfBackend {
|
|||||||
query_embedding: &[f32],
|
query_embedding: &[f32],
|
||||||
k: usize,
|
k: usize,
|
||||||
) -> Vec<MemorySearchResult> {
|
) -> 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 candidates = k.saturating_mul(3).max(10);
|
||||||
let raw = self
|
let raw = self.memory.hybrid_search_with(
|
||||||
.memory
|
query_embedding,
|
||||||
.hybrid_search(query_embedding, query_text, 0.7, 0.3, candidates);
|
query_text,
|
||||||
|
crate::hybrid::DEFAULT_FUSION,
|
||||||
|
candidates,
|
||||||
|
);
|
||||||
|
|
||||||
if raw.is_empty() {
|
if raw.is_empty() {
|
||||||
return Vec::new();
|
return Vec::new();
|
||||||
|
|||||||
@@ -6,6 +6,11 @@
|
|||||||
//! - Temporal expansion (time-related rewrites)
|
//! - Temporal expansion (time-related rewrites)
|
||||||
//! - Morphological variants (stemming-like transforms)
|
//! - Morphological variants (stemming-like transforms)
|
||||||
//! - Knowledge graph expansion (entity aliases and neighbors)
|
//! - Knowledge graph expansion (entity aliases and neighbors)
|
||||||
|
//!
|
||||||
|
//! The morphological rules are crude suffix swaps, so some variants are not
|
||||||
|
//! words ("during" -> "dured"). That is tolerable for a BM25 stage, which
|
||||||
|
//! simply finds no postings for a nonsense term, but it means expansion is not
|
||||||
|
//! free: measure before enabling it on a retrieval path.
|
||||||
|
|
||||||
use crate::knowledge::KnowledgeCache;
|
use crate::knowledge::KnowledgeCache;
|
||||||
|
|
||||||
@@ -340,18 +345,85 @@ fn contains_phrase(text: &str, phrase: &str) -> bool {
|
|||||||
|
|
||||||
/// Replace a phrase in `text` case-insensitively, preserving surrounding case.
|
/// Replace a phrase in `text` case-insensitively, preserving surrounding case.
|
||||||
fn replace_word_case_insensitive(text: &str, from: &str, to: &str) -> String {
|
fn replace_word_case_insensitive(text: &str, from: &str, to: &str) -> String {
|
||||||
case_insensitive_replace(text, from, to)
|
replace_first(text, from, to, MatchKind::WholeWord)
|
||||||
}
|
}
|
||||||
|
|
||||||
fn case_insensitive_replace(text: &str, from: &str, to: &str) -> String {
|
fn case_insensitive_replace(text: &str, from: &str, to: &str) -> String {
|
||||||
let lower = text.to_lowercase();
|
replace_first(text, from, to, MatchKind::Substring)
|
||||||
let lower_from = from.to_lowercase();
|
}
|
||||||
if let Some(pos) = lower.find(&lower_from) {
|
|
||||||
let end = pos + from.len();
|
/// Whether a match may fall inside a larger word.
|
||||||
format!("{}{}{}", &text[..pos], to, &text[end..])
|
#[derive(Clone, Copy, PartialEq)]
|
||||||
} else {
|
enum MatchKind {
|
||||||
text.to_string()
|
/// Match anywhere, including inside another word.
|
||||||
|
Substring,
|
||||||
|
/// Match only when both ends sit on a word boundary.
|
||||||
|
WholeWord,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Replace the first case-insensitive match of `from` in `text` with `to`.
|
||||||
|
///
|
||||||
|
/// Matching walks the *original* string rather than a lowercased copy. The
|
||||||
|
/// previous implementation searched `text.to_lowercase()` and then sliced
|
||||||
|
/// `text` with the offsets it found, which only holds while lowercasing
|
||||||
|
/// preserves byte length. It does not: Turkish `İ` (2 bytes) lowercases to
|
||||||
|
/// `i` + U+0307 (3 bytes), so every later offset was wrong — silently
|
||||||
|
/// corrupting the output, or panicking when an offset landed inside a
|
||||||
|
/// character or past the end. `"İ AI"` was enough to panic.
|
||||||
|
fn replace_first(text: &str, from: &str, to: &str, kind: MatchKind) -> String {
|
||||||
|
match find_case_insensitive(text, from, kind) {
|
||||||
|
Some((start, end)) => {
|
||||||
|
let mut out = String::with_capacity(text.len() - (end - start) + to.len());
|
||||||
|
out.push_str(&text[..start]);
|
||||||
|
out.push_str(to);
|
||||||
|
out.push_str(&text[end..]);
|
||||||
|
out
|
||||||
}
|
}
|
||||||
|
None => text.to_string(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Byte range of the first case-insensitive match of `needle` in `haystack`.
|
||||||
|
fn find_case_insensitive(haystack: &str, needle: &str, kind: MatchKind) -> Option<(usize, usize)> {
|
||||||
|
if needle.is_empty() {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
let lowered: Vec<char> = needle.chars().flat_map(char::to_lowercase).collect();
|
||||||
|
let is_word = |c: char| c.is_alphanumeric() || c == '_';
|
||||||
|
|
||||||
|
for (start, _) in haystack.char_indices() {
|
||||||
|
if kind == MatchKind::WholeWord
|
||||||
|
&& haystack[..start].chars().next_back().is_some_and(is_word)
|
||||||
|
{
|
||||||
|
continue; // mid-word: "ai" inside "training"
|
||||||
|
}
|
||||||
|
let mut matched = 0usize;
|
||||||
|
let mut end = start;
|
||||||
|
for (offset, ch) in haystack[start..].char_indices() {
|
||||||
|
if matched == lowered.len() {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
let mut consumed_all = true;
|
||||||
|
for lc in ch.to_lowercase() {
|
||||||
|
if lowered.get(matched) != Some(&lc) {
|
||||||
|
consumed_all = false;
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
matched += 1;
|
||||||
|
}
|
||||||
|
if !consumed_all {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
end = start + offset + ch.len_utf8();
|
||||||
|
}
|
||||||
|
if matched == lowered.len()
|
||||||
|
&& !(kind == MatchKind::WholeWord
|
||||||
|
&& haystack[end..].chars().next().is_some_and(is_word))
|
||||||
|
{
|
||||||
|
return Some((start, end));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
None
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Simple whitespace/punctuation tokenizer.
|
/// Simple whitespace/punctuation tokenizer.
|
||||||
@@ -637,4 +709,86 @@ mod tests {
|
|||||||
expanded.iter().map(|x| &x.text).collect::<Vec<_>>()
|
expanded.iter().map(|x| &x.text).collect::<Vec<_>>()
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
#[test]
|
||||||
|
fn acronyms_only_match_whole_words() {
|
||||||
|
let ex = QueryExpander::new(QueryExpansionConfig::default());
|
||||||
|
// "training" contains "ai", "programming" contains "pr". These used to
|
||||||
|
// be rewritten to "trArtificial Intelligencening" and
|
||||||
|
// "Pull Requestogramming".
|
||||||
|
for query in [
|
||||||
|
"How many miles during my marathon training?",
|
||||||
|
"Which programming language did I pick?",
|
||||||
|
"I updated the maintainer list",
|
||||||
|
] {
|
||||||
|
for expansion in ex.expand(query) {
|
||||||
|
assert!(
|
||||||
|
expansion.expansion_type != "acronym",
|
||||||
|
"{query:?} produced {expansion:?}"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
// A real acronym still expands, in both directions.
|
||||||
|
let texts: Vec<String> = ex
|
||||||
|
.expand("What about the API and the database?")
|
||||||
|
.into_iter()
|
||||||
|
.filter(|e| e.expansion_type == "acronym")
|
||||||
|
.map(|e| e.text)
|
||||||
|
.collect();
|
||||||
|
assert!(
|
||||||
|
texts
|
||||||
|
.iter()
|
||||||
|
.any(|t| t.contains("Application Programming Interface")),
|
||||||
|
"{texts:?}"
|
||||||
|
);
|
||||||
|
assert!(texts.iter().any(|t| t.contains("DB")), "{texts:?}");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn non_ascii_queries_do_not_panic_or_corrupt() {
|
||||||
|
let ex = QueryExpander::new(QueryExpansionConfig::default());
|
||||||
|
// Turkish 'İ' is 2 bytes but lowercases to 3, so offsets taken from a
|
||||||
|
// lowercased copy no longer line up with the original. `"İ AI"` used
|
||||||
|
// to panic; `"İstanbul AI trip"` used to silently eat a character.
|
||||||
|
for query in ["İ AI", "İé AI", "İİ ML", "İstanbul AI trip", "ǰ ML notes"] {
|
||||||
|
for expansion in ex.expand(query) {
|
||||||
|
assert!(
|
||||||
|
expansion.text.contains('İ') || expansion.text.contains('ǰ'),
|
||||||
|
"{query:?} lost its leading character: {expansion:?}"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
let expanded = ex.expand("İstanbul AI trip");
|
||||||
|
assert!(
|
||||||
|
expanded
|
||||||
|
.iter()
|
||||||
|
.any(|e| e.text == "İstanbul Artificial Intelligence trip"),
|
||||||
|
"{expanded:?}"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn whole_word_matching_handles_string_edges_and_case() {
|
||||||
|
assert_eq!(
|
||||||
|
replace_word_case_insensitive("ai tools", "AI", "Artificial Intelligence"),
|
||||||
|
"Artificial Intelligence tools"
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
replace_word_case_insensitive("tools for ai", "AI", "Artificial Intelligence"),
|
||||||
|
"tools for Artificial Intelligence"
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
replace_word_case_insensitive("the aim", "AI", "Artificial Intelligence"),
|
||||||
|
"the aim",
|
||||||
|
"must not match inside a word"
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
replace_word_case_insensitive("no match here", "xyz", "abc"),
|
||||||
|
"no match here"
|
||||||
|
);
|
||||||
|
// Only the first occurrence is replaced, as before.
|
||||||
|
assert_eq!(
|
||||||
|
replace_word_case_insensitive("ai and ai", "ai", "ML"),
|
||||||
|
"ML and ai"
|
||||||
|
);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -20,8 +20,7 @@ impl HDF5Memory {
|
|||||||
query_embedding: &[f32],
|
query_embedding: &[f32],
|
||||||
query_text: &str,
|
query_text: &str,
|
||||||
bm25: &bm25::BM25Index,
|
bm25: &bm25::BM25Index,
|
||||||
vector_weight: f32,
|
fusion: hybrid::Fusion,
|
||||||
keyword_weight: f32,
|
|
||||||
k: usize,
|
k: usize,
|
||||||
) -> Vec<(usize, f32)> {
|
) -> Vec<(usize, f32)> {
|
||||||
self.ensure_hnsw_fresh();
|
self.ensure_hnsw_fresh();
|
||||||
@@ -38,23 +37,16 @@ impl HDF5Memory {
|
|||||||
// Fusion normalises over every keyword match, so it needs all
|
// Fusion normalises over every keyword match, so it needs all
|
||||||
// the scores — but not ranked.
|
// the scores — but not ranked.
|
||||||
let kw_scores = bm25.scores(query_text);
|
let kw_scores = bm25.scores(query_text);
|
||||||
hybrid::merge_vector_keyword(
|
hybrid::fuse(vec_scores, kw_scores, fusion, k)
|
||||||
vec_scores,
|
|
||||||
kw_scores,
|
|
||||||
vector_weight,
|
|
||||||
keyword_weight,
|
|
||||||
k,
|
|
||||||
)
|
|
||||||
}
|
}
|
||||||
_ => hybrid::hybrid_search(
|
_ => hybrid::hybrid_search_fused(
|
||||||
query_embedding,
|
query_embedding,
|
||||||
query_text,
|
query_text,
|
||||||
&self.cache.embeddings,
|
&self.cache.embeddings,
|
||||||
&self.cache.chunks,
|
&self.cache.chunks,
|
||||||
&self.cache.tombstones,
|
&self.cache.tombstones,
|
||||||
bm25,
|
bm25,
|
||||||
vector_weight,
|
fusion,
|
||||||
keyword_weight,
|
|
||||||
k,
|
k,
|
||||||
),
|
),
|
||||||
}
|
}
|
||||||
@@ -66,19 +58,17 @@ impl HDF5Memory {
|
|||||||
query_embedding: &[f32],
|
query_embedding: &[f32],
|
||||||
query_text: &str,
|
query_text: &str,
|
||||||
bm25: &bm25::BM25Index,
|
bm25: &bm25::BM25Index,
|
||||||
vector_weight: f32,
|
fusion: hybrid::Fusion,
|
||||||
keyword_weight: f32,
|
|
||||||
k: usize,
|
k: usize,
|
||||||
) -> Vec<(usize, f32)> {
|
) -> Vec<(usize, f32)> {
|
||||||
hybrid::hybrid_search(
|
hybrid::hybrid_search_fused(
|
||||||
query_embedding,
|
query_embedding,
|
||||||
query_text,
|
query_text,
|
||||||
&self.cache.embeddings,
|
&self.cache.embeddings,
|
||||||
&self.cache.chunks,
|
&self.cache.chunks,
|
||||||
&self.cache.tombstones,
|
&self.cache.tombstones,
|
||||||
bm25,
|
bm25,
|
||||||
vector_weight,
|
fusion,
|
||||||
keyword_weight,
|
|
||||||
k,
|
k,
|
||||||
)
|
)
|
||||||
}
|
}
|
||||||
@@ -91,20 +81,36 @@ impl HDF5Memory {
|
|||||||
vector_weight: f32,
|
vector_weight: f32,
|
||||||
keyword_weight: f32,
|
keyword_weight: f32,
|
||||||
k: usize,
|
k: usize,
|
||||||
|
) -> Vec<SearchResult> {
|
||||||
|
self.hybrid_search_with(
|
||||||
|
query_embedding,
|
||||||
|
query_text,
|
||||||
|
hybrid::Fusion::Weighted {
|
||||||
|
vector: vector_weight,
|
||||||
|
keyword: keyword_weight,
|
||||||
|
},
|
||||||
|
k,
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// [`HDF5Memory::hybrid_search`] with the fusion method chosen explicitly.
|
||||||
|
///
|
||||||
|
/// [`hybrid::DEFAULT_FUSION`] is what the weighted form defaults to;
|
||||||
|
/// [`hybrid::Fusion::Rrf`] combines the two stages by rank instead of by
|
||||||
|
/// score.
|
||||||
|
pub fn hybrid_search_with(
|
||||||
|
&mut self,
|
||||||
|
query_embedding: &[f32],
|
||||||
|
query_text: &str,
|
||||||
|
fusion: hybrid::Fusion,
|
||||||
|
k: usize,
|
||||||
) -> Vec<SearchResult> {
|
) -> Vec<SearchResult> {
|
||||||
// The keyword index lives for the life of the store and is updated
|
// The keyword index lives for the life of the store and is updated
|
||||||
// incrementally. Take it out for the duration of the call so the
|
// incrementally. Take it out for the duration of the call so the
|
||||||
// vector stage can borrow `self` mutably, then put it back.
|
// vector stage can borrow `self` mutably, then put it back.
|
||||||
self.ensure_bm25_fresh();
|
self.ensure_bm25_fresh();
|
||||||
let bm25 = self.bm25.take().expect("ensure_bm25_fresh leaves an index");
|
let bm25 = self.bm25.take().expect("ensure_bm25_fresh leaves an index");
|
||||||
let scored = self.vector_keyword_search(
|
let scored = self.vector_keyword_search(query_embedding, query_text, &bm25, fusion, k);
|
||||||
query_embedding,
|
|
||||||
query_text,
|
|
||||||
&bm25,
|
|
||||||
vector_weight,
|
|
||||||
keyword_weight,
|
|
||||||
k,
|
|
||||||
);
|
|
||||||
let mut results: Vec<SearchResult> = scored
|
let mut results: Vec<SearchResult> = scored
|
||||||
.into_iter()
|
.into_iter()
|
||||||
.map(|(idx, score)| {
|
.map(|(idx, score)| {
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
[package]
|
[package]
|
||||||
name = "clawhdf5-android"
|
name = "clawhdf5-android"
|
||||||
version = "2.4.0"
|
version = "2.5.0"
|
||||||
edition = "2024"
|
edition = "2024"
|
||||||
description = "Android JNI bridge for edgehdf5-memory HDF5 backend"
|
description = "Android JNI bridge for edgehdf5-memory HDF5 backend"
|
||||||
license = "MIT"
|
license = "MIT"
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
[package]
|
[package]
|
||||||
name = "clawhdf5-ann"
|
name = "clawhdf5-ann"
|
||||||
version = "2.4.0"
|
version = "2.5.0"
|
||||||
edition = "2024"
|
edition = "2024"
|
||||||
description = "HNSW approximate nearest neighbor index stored as HDF5"
|
description = "HNSW approximate nearest neighbor index stored as HDF5"
|
||||||
license = "MIT"
|
license = "MIT"
|
||||||
@@ -10,9 +10,9 @@ keywords = ["hdf5", "ann", "hnsw", "nearest-neighbor"]
|
|||||||
categories = ["algorithms", "science"]
|
categories = ["algorithms", "science"]
|
||||||
|
|
||||||
[dependencies]
|
[dependencies]
|
||||||
clawhdf5-format = { path = "../clawhdf5-format", version = "2.4.0" }
|
clawhdf5-format = { path = "../clawhdf5-format", version = "2.5.0" }
|
||||||
clawhdf5-io = { path = "../clawhdf5-io", version = "2.4.0" }
|
clawhdf5-io = { path = "../clawhdf5-io", version = "2.5.0" }
|
||||||
clawhdf5-accel = { path = "../clawhdf5-accel", version = "2.4.0" }
|
clawhdf5-accel = { path = "../clawhdf5-accel", version = "2.5.0" }
|
||||||
rayon = { version = "1", optional = true }
|
rayon = { version = "1", optional = true }
|
||||||
|
|
||||||
[features]
|
[features]
|
||||||
|
|||||||
+250
-63
@@ -248,66 +248,62 @@ impl HnswIndex {
|
|||||||
let mut entry_point = 0;
|
let mut entry_point = 0;
|
||||||
let mut ep_level = node_levels[0];
|
let mut ep_level = node_levels[0];
|
||||||
|
|
||||||
// Insert nodes one by one
|
// Insert in batches. Each batch is planned against the graph as it
|
||||||
for i in 1..n {
|
// stood when the batch began (read-only, so the plans are independent
|
||||||
let node_level = node_levels[i];
|
// and run in parallel with the `parallel` feature), then linked, then
|
||||||
let mut ep = entry_point;
|
// every neighbour list that overflowed is pruned once. Pruning is ~90%
|
||||||
|
// of a build's distance evaluations, and a node that gains several
|
||||||
// Phase 1: greedy search from top layer down to node_level + 1
|
// back-links in one batch is pruned once instead of once per link.
|
||||||
let start_layer = ep_level;
|
//
|
||||||
for layer in (node_level + 1..=start_layer).rev() {
|
// Nodes in the same batch cannot see each other while planning, so
|
||||||
ep = greedy_closest(vectors, &graph[layer], &vectors[i], ep, metric);
|
// batches start at one node and grow only as the graph does — a batch
|
||||||
|
// is never more than a small fraction of what is already linked. The
|
||||||
|
// result is deterministic and identical with or without `parallel`.
|
||||||
|
let mut next = 1;
|
||||||
|
while next < n {
|
||||||
|
let mut end = (next + batch_len(next)).min(n);
|
||||||
|
// A node that raises the top layer becomes the new entry point and
|
||||||
|
// changes how every later node descends: give it a batch alone.
|
||||||
|
if let Some(tall) = (next..end).find(|&i| node_levels[i] > ep_level) {
|
||||||
|
end = if tall == next { next + 1 } else { tall };
|
||||||
}
|
}
|
||||||
|
|
||||||
// Phase 2: search and connect at layers node_level down to 0
|
let plans = plan_batch(
|
||||||
let bottom = if node_level < start_layer {
|
vectors,
|
||||||
node_level
|
&graph,
|
||||||
} else {
|
&node_levels,
|
||||||
start_layer
|
next..end,
|
||||||
};
|
entry_point,
|
||||||
for layer in (0..=bottom).rev() {
|
ep_level,
|
||||||
|
(m, m_max0, ef_construction),
|
||||||
|
metric,
|
||||||
|
);
|
||||||
|
|
||||||
|
let mut overflowed: Vec<(usize, usize)> = Vec::new();
|
||||||
|
for (offset, plan) in plans.into_iter().enumerate() {
|
||||||
|
let node = next + offset;
|
||||||
|
for (layer, selected) in plan {
|
||||||
let max_conn = if layer == 0 { m_max0 } else { m };
|
let max_conn = if layer == 0 { m_max0 } else { m };
|
||||||
|
|
||||||
let neighbors = search_layer(
|
|
||||||
vectors,
|
|
||||||
&graph[layer],
|
|
||||||
&vectors[i],
|
|
||||||
ep,
|
|
||||||
ef_construction,
|
|
||||||
metric,
|
|
||||||
);
|
|
||||||
|
|
||||||
let scored: Vec<(usize, f32)> =
|
|
||||||
neighbors.iter().map(|c| (c.id, c.distance)).collect();
|
|
||||||
let selected = select_neighbors(vectors, &scored, max_conn, metric);
|
|
||||||
|
|
||||||
// Add bidirectional connections
|
|
||||||
graph[layer][i] = selected.clone();
|
|
||||||
for &neighbor in &selected {
|
for &neighbor in &selected {
|
||||||
graph[layer][neighbor].push(i);
|
let list = &mut graph[layer][neighbor];
|
||||||
// Prune if over limit
|
list.push(node);
|
||||||
if graph[layer][neighbor].len() > max_conn {
|
if list.len() == max_conn + 1 {
|
||||||
prune_connections(
|
overflowed.push((layer, neighbor));
|
||||||
vectors,
|
|
||||||
&mut graph[layer][neighbor],
|
|
||||||
neighbor,
|
|
||||||
max_conn,
|
|
||||||
metric,
|
|
||||||
);
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
graph[layer][node] = selected;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
prune_overflowed(vectors, &mut graph, overflowed, (m, m_max0), metric);
|
||||||
|
|
||||||
if !selected.is_empty() {
|
for (i, &level) in node_levels.iter().enumerate().take(end).skip(next) {
|
||||||
ep = selected[0];
|
if level > ep_level {
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// Update entry point if this node has a higher level
|
|
||||||
if node_level > ep_level {
|
|
||||||
entry_point = i;
|
entry_point = i;
|
||||||
ep_level = node_level;
|
ep_level = level;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
next = end;
|
||||||
|
}
|
||||||
|
|
||||||
Self {
|
Self {
|
||||||
vectors: prepared,
|
vectors: prepared,
|
||||||
@@ -408,22 +404,19 @@ impl HnswIndex {
|
|||||||
ep,
|
ep,
|
||||||
self.ef_construction,
|
self.ef_construction,
|
||||||
self.metric,
|
self.metric,
|
||||||
|
None,
|
||||||
);
|
);
|
||||||
let scored: Vec<(usize, f32)> = neighbors.iter().map(|c| (c.id, c.distance)).collect();
|
let scored: Vec<(usize, f32)> = neighbors.iter().map(|c| (c.id, c.distance)).collect();
|
||||||
let selected = select_neighbors(&self.vectors, &scored, max_conn, self.metric);
|
let selected = select_neighbors(&self.vectors, &scored, max_conn, self.metric);
|
||||||
self.graph[layer][id] = selected.clone();
|
self.graph[layer][id] = selected.clone();
|
||||||
for &neighbor in &selected {
|
link_back(
|
||||||
self.graph[layer][neighbor].push(id);
|
|
||||||
if self.graph[layer][neighbor].len() > max_conn {
|
|
||||||
prune_connections(
|
|
||||||
&self.vectors,
|
&self.vectors,
|
||||||
&mut self.graph[layer][neighbor],
|
&mut self.graph[layer],
|
||||||
neighbor,
|
id,
|
||||||
|
&selected,
|
||||||
max_conn,
|
max_conn,
|
||||||
self.metric,
|
self.metric,
|
||||||
);
|
);
|
||||||
}
|
|
||||||
}
|
|
||||||
if !selected.is_empty() {
|
if !selected.is_empty() {
|
||||||
ep = selected[0];
|
ep = selected[0];
|
||||||
}
|
}
|
||||||
@@ -517,13 +510,21 @@ impl HnswIndex {
|
|||||||
ep = greedy_closest(&self.vectors, &self.graph[layer], query, ep, self.metric);
|
ep = greedy_closest(&self.vectors, &self.graph[layer], query, ep, self.metric);
|
||||||
}
|
}
|
||||||
|
|
||||||
// Search layer 0 with ef candidates. Deleted nodes are still traversed
|
// Search layer 0 for the ef nearest *live* nodes. Deleted nodes are
|
||||||
// (they remain valid graph waypoints) but are filtered from the result.
|
// still traversed (they remain valid graph waypoints) but take no
|
||||||
let candidates = search_layer(&self.vectors, &self.graph[0], query, ep, ef, self.metric);
|
// result slot, so deletions near the query don't shrink the answer.
|
||||||
|
let candidates = search_layer(
|
||||||
|
&self.vectors,
|
||||||
|
&self.graph[0],
|
||||||
|
query,
|
||||||
|
ep,
|
||||||
|
ef,
|
||||||
|
self.metric,
|
||||||
|
Some(&self.deleted),
|
||||||
|
);
|
||||||
|
|
||||||
candidates
|
candidates
|
||||||
.into_iter()
|
.into_iter()
|
||||||
.filter(|c| !self.deleted[c.id])
|
|
||||||
.take(k)
|
.take(k)
|
||||||
.map(|c| (c.id, c.distance))
|
.map(|c| (c.id, c.distance))
|
||||||
.collect()
|
.collect()
|
||||||
@@ -940,6 +941,14 @@ fn greedy_closest(
|
|||||||
}
|
}
|
||||||
|
|
||||||
/// Search a single layer for the ef closest nodes to `query`.
|
/// Search a single layer for the ef closest nodes to `query`.
|
||||||
|
/// Best-first search of one layer, returning up to `ef` nodes by ascending
|
||||||
|
/// distance.
|
||||||
|
///
|
||||||
|
/// `skip` marks nodes that must not be *returned* (soft-deleted ones). They
|
||||||
|
/// are still traversed — a tombstone is a perfectly good waypoint — but they
|
||||||
|
/// never occupy one of the `ef` result slots. Filtering them out afterwards
|
||||||
|
/// instead meant a query whose neighbourhood had been deleted got back fewer
|
||||||
|
/// than `k` results, or none, however many live records were nearby.
|
||||||
fn search_layer(
|
fn search_layer(
|
||||||
vectors: &[Vec<f32>],
|
vectors: &[Vec<f32>],
|
||||||
layer: &[Vec<usize>],
|
layer: &[Vec<usize>],
|
||||||
@@ -947,6 +956,7 @@ fn search_layer(
|
|||||||
ep: usize,
|
ep: usize,
|
||||||
ef: usize,
|
ef: usize,
|
||||||
metric: DistanceMetric,
|
metric: DistanceMetric,
|
||||||
|
skip: Option<&[bool]>,
|
||||||
) -> Vec<Candidate> {
|
) -> Vec<Candidate> {
|
||||||
let ep_dist = compute_distance(query, &vectors[ep], metric);
|
let ep_dist = compute_distance(query, &vectors[ep], metric);
|
||||||
|
|
||||||
@@ -959,16 +969,18 @@ fn search_layer(
|
|||||||
|
|
||||||
// Max-heap of current results (furthest first)
|
// Max-heap of current results (furthest first)
|
||||||
let mut results = BinaryHeap::new();
|
let mut results = BinaryHeap::new();
|
||||||
|
if !skip.is_some_and(|s| s[ep]) {
|
||||||
results.push(FarCandidate {
|
results.push(FarCandidate {
|
||||||
id: ep,
|
id: ep,
|
||||||
distance: ep_dist,
|
distance: ep_dist,
|
||||||
});
|
});
|
||||||
|
}
|
||||||
|
|
||||||
VISITED.with_borrow_mut(|visited| {
|
VISITED.with_borrow_mut(|visited| {
|
||||||
visited.begin(vectors.len());
|
visited.begin(vectors.len());
|
||||||
visited.insert(ep);
|
visited.insert(ep);
|
||||||
search_layer_visit(
|
search_layer_visit(
|
||||||
vectors, layer, query, ef, metric, visited, candidates, results,
|
vectors, layer, query, ef, metric, skip, visited, candidates, results,
|
||||||
)
|
)
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
@@ -1016,6 +1028,7 @@ fn search_layer_visit(
|
|||||||
query: &[f32],
|
query: &[f32],
|
||||||
ef: usize,
|
ef: usize,
|
||||||
metric: DistanceMetric,
|
metric: DistanceMetric,
|
||||||
|
skip: Option<&[bool]>,
|
||||||
visited: &mut Visited,
|
visited: &mut Visited,
|
||||||
mut candidates: BinaryHeap<Candidate>,
|
mut candidates: BinaryHeap<Candidate>,
|
||||||
mut results: BinaryHeap<FarCandidate>,
|
mut results: BinaryHeap<FarCandidate>,
|
||||||
@@ -1039,6 +1052,9 @@ fn search_layer_visit(
|
|||||||
id: neighbor,
|
id: neighbor,
|
||||||
distance: d,
|
distance: d,
|
||||||
});
|
});
|
||||||
|
if skip.is_some_and(|s| s[neighbor]) {
|
||||||
|
continue; // explore through it, but never return it
|
||||||
|
}
|
||||||
results.push(FarCandidate {
|
results.push(FarCandidate {
|
||||||
id: neighbor,
|
id: neighbor,
|
||||||
distance: d,
|
distance: d,
|
||||||
@@ -1111,6 +1127,125 @@ fn select_neighbors(
|
|||||||
selected
|
selected
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// How many nodes to plan together once `linked` nodes are in the graph.
|
||||||
|
fn batch_len(linked: usize) -> usize {
|
||||||
|
(linked / 16).clamp(1, 512)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// For each node in `batch`: the neighbours to link it to on each of its
|
||||||
|
/// layers, found by searching the graph as it currently stands.
|
||||||
|
#[allow(clippy::too_many_arguments)]
|
||||||
|
fn plan_batch(
|
||||||
|
vectors: &[Vec<f32>],
|
||||||
|
graph: &[Vec<Vec<usize>>],
|
||||||
|
node_levels: &[usize],
|
||||||
|
batch: std::ops::Range<usize>,
|
||||||
|
entry_point: usize,
|
||||||
|
ep_level: usize,
|
||||||
|
(m, m_max0, ef_construction): (usize, usize, usize),
|
||||||
|
metric: DistanceMetric,
|
||||||
|
) -> Vec<Vec<(usize, Vec<usize>)>> {
|
||||||
|
let plan_one = |i: usize| -> Vec<(usize, Vec<usize>)> {
|
||||||
|
let node_level = node_levels[i];
|
||||||
|
let mut ep = entry_point;
|
||||||
|
// Phase 1: greedy descent from the top layer down to node_level + 1.
|
||||||
|
for layer in (node_level + 1..=ep_level).rev() {
|
||||||
|
ep = greedy_closest(vectors, &graph[layer], &vectors[i], ep, metric);
|
||||||
|
}
|
||||||
|
// Phase 2: search and select on every layer the node lives on.
|
||||||
|
let mut plan = Vec::with_capacity(node_level.min(ep_level) + 1);
|
||||||
|
for layer in (0..=node_level.min(ep_level)).rev() {
|
||||||
|
let max_conn = if layer == 0 { m_max0 } else { m };
|
||||||
|
let neighbors = search_layer(
|
||||||
|
vectors,
|
||||||
|
&graph[layer],
|
||||||
|
&vectors[i],
|
||||||
|
ep,
|
||||||
|
ef_construction,
|
||||||
|
metric,
|
||||||
|
None,
|
||||||
|
);
|
||||||
|
let scored: Vec<(usize, f32)> = neighbors.iter().map(|c| (c.id, c.distance)).collect();
|
||||||
|
let selected = select_neighbors(vectors, &scored, max_conn, metric);
|
||||||
|
if let Some(&closest) = selected.first() {
|
||||||
|
ep = closest;
|
||||||
|
}
|
||||||
|
plan.push((layer, selected));
|
||||||
|
}
|
||||||
|
plan
|
||||||
|
};
|
||||||
|
|
||||||
|
#[cfg(feature = "parallel")]
|
||||||
|
if batch.len() >= PARALLEL_MIN {
|
||||||
|
use rayon::prelude::*;
|
||||||
|
return batch.into_par_iter().map(plan_one).collect();
|
||||||
|
}
|
||||||
|
batch.map(plan_one).collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Prune every `(layer, node)` neighbour list in `overflowed` back to its
|
||||||
|
/// limit. Each list belongs to a different node, so they are independent.
|
||||||
|
fn prune_overflowed(
|
||||||
|
vectors: &[Vec<f32>],
|
||||||
|
graph: &mut [Vec<Vec<usize>>],
|
||||||
|
overflowed: Vec<(usize, usize)>,
|
||||||
|
(m, m_max0): (usize, usize),
|
||||||
|
metric: DistanceMetric,
|
||||||
|
) {
|
||||||
|
let limit = |layer: usize| if layer == 0 { m_max0 } else { m };
|
||||||
|
|
||||||
|
#[cfg(feature = "parallel")]
|
||||||
|
if overflowed.len() >= PARALLEL_MIN {
|
||||||
|
use rayon::prelude::*;
|
||||||
|
let mut work: Vec<(usize, usize, Vec<usize>)> = overflowed
|
||||||
|
.into_iter()
|
||||||
|
.map(|(layer, node)| (layer, node, std::mem::take(&mut graph[layer][node])))
|
||||||
|
.collect();
|
||||||
|
work.par_iter_mut().for_each(|(layer, node, list)| {
|
||||||
|
prune_connections(vectors, list, *node, limit(*layer), metric);
|
||||||
|
});
|
||||||
|
for (layer, node, list) in work {
|
||||||
|
graph[layer][node] = list;
|
||||||
|
}
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
|
for (layer, node) in overflowed {
|
||||||
|
prune_connections(vectors, &mut graph[layer][node], node, limit(layer), metric);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Fewest independent tasks worth handing to the thread pool.
|
||||||
|
#[cfg(feature = "parallel")]
|
||||||
|
const PARALLEL_MIN: usize = 8;
|
||||||
|
|
||||||
|
/// Add the back-link `neighbor -> new_id` for every selected neighbour, pruning
|
||||||
|
/// each list that overflows.
|
||||||
|
///
|
||||||
|
/// Used by incremental [`HnswIndex::insert`]. (A single insert's handful of
|
||||||
|
/// prunes is too fine-grained to parallelise profitably — measured 1.45x on 16
|
||||||
|
/// cores; bulk builds batch their pruning instead, see `prune_overflowed`.)
|
||||||
|
fn link_back(
|
||||||
|
vectors: &[Vec<f32>],
|
||||||
|
layer: &mut [Vec<usize>],
|
||||||
|
new_id: usize,
|
||||||
|
selected: &[usize],
|
||||||
|
max_conn: usize,
|
||||||
|
metric: DistanceMetric,
|
||||||
|
) {
|
||||||
|
let mut overflowed: Vec<usize> = Vec::new();
|
||||||
|
for &neighbor in selected {
|
||||||
|
layer[neighbor].push(new_id);
|
||||||
|
if layer[neighbor].len() > max_conn {
|
||||||
|
overflowed.push(neighbor);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
for node in overflowed {
|
||||||
|
prune_connections(vectors, &mut layer[node], node, max_conn, metric);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
/// Trim `node`'s neighbour list back to `max_conn` with [`select_neighbors`].
|
/// Trim `node`'s neighbour list back to `max_conn` with [`select_neighbors`].
|
||||||
fn prune_connections(
|
fn prune_connections(
|
||||||
vectors: &[Vec<f32>],
|
vectors: &[Vec<f32>],
|
||||||
@@ -1384,6 +1519,58 @@ mod tests {
|
|||||||
assert!(recall >= 0.95, "incremental recall@10 = {recall}");
|
assert!(recall >= 0.95, "incremental recall@10 = {recall}");
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn deletions_near_the_query_do_not_shrink_or_degrade_results() {
|
||||||
|
let mut vectors = clustered(2040, 16, 20, 11);
|
||||||
|
let queries = vectors.split_off(2000);
|
||||||
|
let mut index = HnswIndex::build_with_metric(&vectors, 8, 40, DistanceMetric::L2);
|
||||||
|
|
||||||
|
let mut short = 0;
|
||||||
|
let mut hits = 0;
|
||||||
|
for q in &queries {
|
||||||
|
// Delete this query's 40 nearest neighbours: more than ef, so every
|
||||||
|
// candidate a plain search collects is a tombstone.
|
||||||
|
let mut exact: Vec<(usize, f32)> = vectors
|
||||||
|
.iter()
|
||||||
|
.enumerate()
|
||||||
|
.filter(|(i, _)| !index.is_deleted(*i))
|
||||||
|
.map(|(i, v)| (i, compute_distance(q, v, DistanceMetric::L2)))
|
||||||
|
.collect();
|
||||||
|
exact.sort_by(|a, b| a.1.total_cmp(&b.1));
|
||||||
|
for &(id, _) in &exact[..40] {
|
||||||
|
index.mark_deleted(id);
|
||||||
|
}
|
||||||
|
let want: Vec<usize> = exact[40..50].iter().map(|e| e.0).collect();
|
||||||
|
|
||||||
|
let got = index.search(q, 10, 32);
|
||||||
|
assert!(got.iter().all(|(id, _)| !index.is_deleted(*id)));
|
||||||
|
short += usize::from(got.len() < 10);
|
||||||
|
hits += got.iter().filter(|(id, _)| want.contains(id)).count();
|
||||||
|
}
|
||||||
|
assert_eq!(short, 0, "searches returned fewer than k live results");
|
||||||
|
let recall = hits as f64 / (10 * queries.len()) as f64;
|
||||||
|
assert!(recall >= 0.9, "recall@10 among live records = {recall}");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn bulk_build_is_deterministic() {
|
||||||
|
// Batched planning runs on a thread pool with the `parallel` feature;
|
||||||
|
// the graph must not depend on scheduling. (It is also the same graph
|
||||||
|
// with and without the feature: both take this exact code path.)
|
||||||
|
let vectors = clustered(2500, 16, 20, 21);
|
||||||
|
let a = HnswIndex::build_with_metric(&vectors, 8, 40, DistanceMetric::Cosine);
|
||||||
|
let b = HnswIndex::build_with_metric(&vectors, 8, 40, DistanceMetric::Cosine);
|
||||||
|
assert_eq!(a.graph_to_bytes(), b.graph_to_bytes());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn batches_stay_a_small_fraction_of_the_graph() {
|
||||||
|
assert_eq!(batch_len(1), 1);
|
||||||
|
assert_eq!(batch_len(15), 1);
|
||||||
|
assert_eq!(batch_len(160), 10);
|
||||||
|
assert_eq!(batch_len(1_000_000), 512);
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn graph_bytes_round_trip_gives_identical_searches() {
|
fn graph_bytes_round_trip_gives_identical_searches() {
|
||||||
let mut vectors = clustered(1260, 16, 12, 9);
|
let mut vectors = clustered(1260, 16, 12, 9);
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
[package]
|
[package]
|
||||||
name = "clawhdf5-bench"
|
name = "clawhdf5-bench"
|
||||||
version = "2.4.0"
|
version = "2.5.0"
|
||||||
edition = "2024"
|
edition = "2024"
|
||||||
description = "Benchmark harnesses for clawhdf5-agent (Track 8)"
|
description = "Benchmark harnesses for clawhdf5-agent (Track 8)"
|
||||||
license = "MIT"
|
license = "MIT"
|
||||||
@@ -13,6 +13,10 @@ path = "src/bin/longmemeval_bench.rs"
|
|||||||
name = "memory_arena"
|
name = "memory_arena"
|
||||||
path = "src/bin/memory_arena.rs"
|
path = "src/bin/memory_arena.rs"
|
||||||
|
|
||||||
|
[[bin]]
|
||||||
|
name = "read_harness"
|
||||||
|
path = "src/bin/read_harness.rs"
|
||||||
|
|
||||||
[[bin]]
|
[[bin]]
|
||||||
name = "search_harness"
|
name = "search_harness"
|
||||||
path = "src/bin/search_harness.rs"
|
path = "src/bin/search_harness.rs"
|
||||||
@@ -53,6 +57,8 @@ harness = false
|
|||||||
[dependencies]
|
[dependencies]
|
||||||
clawhdf5-agent = { path = "../clawhdf5-agent" }
|
clawhdf5-agent = { path = "../clawhdf5-agent" }
|
||||||
clawhdf5-ann = { path = "../clawhdf5-ann" }
|
clawhdf5-ann = { path = "../clawhdf5-ann" }
|
||||||
|
clawhdf5 = { path = "../clawhdf5" }
|
||||||
|
clawhdf5-format = { path = "../clawhdf5-format" }
|
||||||
clawhdf5-io = { path = "../clawhdf5-io" }
|
clawhdf5-io = { path = "../clawhdf5-io" }
|
||||||
mpi = { version = "0.8", optional = true }
|
mpi = { version = "0.8", optional = true }
|
||||||
serde = { workspace = true }
|
serde = { workspace = true }
|
||||||
|
|||||||
@@ -55,44 +55,80 @@ use std::time::{Duration, Instant};
|
|||||||
#[path = "longmemeval_bench/embedder.rs"]
|
#[path = "longmemeval_bench/embedder.rs"]
|
||||||
mod embedder;
|
mod embedder;
|
||||||
|
|
||||||
|
use clawhdf5_agent::bm25::TokenFilter;
|
||||||
|
use clawhdf5_agent::hybrid::Fusion;
|
||||||
use clawhdf5_agent::{AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry};
|
use clawhdf5_agent::{AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry};
|
||||||
use serde::Deserialize;
|
use serde::Deserialize;
|
||||||
use tempfile::TempDir;
|
use tempfile::TempDir;
|
||||||
|
|
||||||
const EMBEDDING_DIM: usize = 384;
|
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}"),
|
||||||
|
};
|
||||||
|
match mode.tokens {
|
||||||
|
TokenFilter::Plain => fusion,
|
||||||
|
TokenFilter::Stemmed => format!("{fusion}_stemmed"),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
/// A retrieval configuration: how much of the score comes from each stage.
|
/// A retrieval configuration: how much of the score comes from each stage.
|
||||||
#[derive(Clone, Copy)]
|
#[derive(Clone, Copy)]
|
||||||
struct Mode {
|
struct Mode {
|
||||||
label: &'static str,
|
label: &'static str,
|
||||||
vector_weight: f32,
|
/// How the two retrieval stages are combined into one ranking.
|
||||||
keyword_weight: f32,
|
fusion: Fusion,
|
||||||
|
/// How keyword tokens are normalised before indexing and querying.
|
||||||
|
tokens: TokenFilter,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Mode {
|
||||||
|
const fn weighted(label: &'static str, vector: f32, keyword: f32) -> Self {
|
||||||
|
Self {
|
||||||
|
label,
|
||||||
|
fusion: Fusion::Weighted { vector, keyword },
|
||||||
|
tokens: TokenFilter::Plain,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
const fn stemmed(mut self, label: &'static str) -> Self {
|
||||||
|
self.label = label;
|
||||||
|
self.tokens = TokenFilter::Stemmed;
|
||||||
|
self
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// The only mode available without real embeddings. Passing zero vectors with
|
/// The only mode available without real embeddings. Passing zero vectors with
|
||||||
/// `vector_weight = 0.0` is what made the vector stage inert.
|
/// `vector_weight = 0.0` is what made the vector stage inert.
|
||||||
const BM25_ONLY: Mode = Mode {
|
const BM25_ONLY: Mode = Mode::weighted("BM25 only (vector stage inert)", 0.0, 1.0);
|
||||||
label: "BM25 only (vector stage inert)",
|
|
||||||
vector_weight: 0.0,
|
|
||||||
keyword_weight: 1.0,
|
|
||||||
};
|
|
||||||
#[cfg(feature = "embeddings")]
|
#[cfg(feature = "embeddings")]
|
||||||
const VECTOR_ONLY: Mode = Mode {
|
const VECTOR_ONLY: Mode = Mode::weighted("Vector only (MiniLM + HNSW)", 1.0, 0.0);
|
||||||
label: "Vector only (MiniLM + HNSW)",
|
|
||||||
vector_weight: 1.0,
|
|
||||||
keyword_weight: 0.0,
|
|
||||||
};
|
|
||||||
/// Tuned by `--sweep` over the full haystack. The former 0.7/0.3 was a
|
/// Tuned by `--sweep` over the full haystack. The former 0.7/0.3 was a
|
||||||
/// documented default that had never been searched, and the sweep found it
|
/// documented default that had never been searched, and the sweep found it
|
||||||
/// strictly dominated: 0.4/0.6 is better on Hit@1, Hit@5, Hit@10 and MRR at
|
/// strictly dominated: 0.4/0.6 is better on Hit@1, Hit@5, Hit@10 and MRR at
|
||||||
/// both granularities.
|
/// both granularities.
|
||||||
#[cfg(feature = "embeddings")]
|
#[cfg(feature = "embeddings")]
|
||||||
const HYBRID: Mode = Mode {
|
const HYBRID: Mode = Mode::weighted("Hybrid (0.4 vector / 0.6 BM25, tuned)", 0.4, 0.6);
|
||||||
label: "Hybrid (0.4 vector / 0.6 BM25, tuned)",
|
|
||||||
vector_weight: 0.4,
|
/// Reciprocal rank fusion, the documented alternative to the weighted sum.
|
||||||
keyword_weight: 0.6,
|
/// It ignores score magnitudes, so there is nothing to tune — which is the
|
||||||
|
/// claim being tested.
|
||||||
|
#[cfg(feature = "embeddings")]
|
||||||
|
const RRF: Mode = Mode {
|
||||||
|
label: "Hybrid (reciprocal rank fusion, k=60)",
|
||||||
|
fusion: Fusion::Rrf { k: 60.0 },
|
||||||
|
tokens: TokenFilter::Plain,
|
||||||
};
|
};
|
||||||
|
|
||||||
|
/// 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");
|
||||||
|
#[cfg(feature = "embeddings")]
|
||||||
|
const HYBRID_STEMMED: Mode = HYBRID.stemmed("Hybrid 0.4/0.6, stemmed tokens");
|
||||||
|
|
||||||
/// Every 0.1 step of vector weight, keyword weight taking the remainder.
|
/// Every 0.1 step of vector weight, keyword weight taking the remainder.
|
||||||
///
|
///
|
||||||
/// Labels are leaked to `&'static str` because `Mode::label` is a `&'static
|
/// Labels are leaked to `&'static str` because `Mode::label` is a `&'static
|
||||||
@@ -104,11 +140,11 @@ fn sweep_modes() -> Vec<Mode> {
|
|||||||
(0..=10)
|
(0..=10)
|
||||||
.map(|i| {
|
.map(|i| {
|
||||||
let v = i as f32 / 10.0;
|
let v = i as f32 / 10.0;
|
||||||
Mode {
|
Mode::weighted(
|
||||||
label: Box::leak(format!("sweep v={v:.1} / k={:.1}", 1.0 - v).into_boxed_str()),
|
Box::leak(format!("sweep v={v:.1} / k={:.1}", 1.0 - v).into_boxed_str()),
|
||||||
vector_weight: v,
|
v,
|
||||||
keyword_weight: 1.0 - v,
|
1.0 - v,
|
||||||
}
|
)
|
||||||
})
|
})
|
||||||
.collect()
|
.collect()
|
||||||
}
|
}
|
||||||
@@ -276,6 +312,7 @@ fn evaluate_question(
|
|||||||
config.compact_threshold = 0.0;
|
config.compact_threshold = 0.0;
|
||||||
|
|
||||||
let mut memory = HDF5Memory::create(config).expect("failed to create HDF5Memory");
|
let mut memory = HDF5Memory::create(config).expect("failed to create HDF5Memory");
|
||||||
|
memory.set_token_filter(mode.tokens);
|
||||||
|
|
||||||
// Build MemoryEntry list from all haystack sessions
|
// Build MemoryEntry list from all haystack sessions
|
||||||
let mut entries: Vec<MemoryEntry> = Vec::new();
|
let mut entries: Vec<MemoryEntry> = Vec::new();
|
||||||
@@ -321,13 +358,7 @@ fn evaluate_question(
|
|||||||
|
|
||||||
let query_emb = embedding_for(embeddings, &q.question);
|
let query_emb = embedding_for(embeddings, &q.question);
|
||||||
let t0 = Instant::now();
|
let t0 = Instant::now();
|
||||||
let results = memory.hybrid_search(
|
let results = memory.hybrid_search_with(&query_emb, &q.question, mode.fusion, top_k);
|
||||||
&query_emb,
|
|
||||||
&q.question,
|
|
||||||
mode.vector_weight,
|
|
||||||
mode.keyword_weight,
|
|
||||||
top_k,
|
|
||||||
);
|
|
||||||
let latency = t0.elapsed();
|
let latency = t0.elapsed();
|
||||||
|
|
||||||
// Session-level recall
|
// Session-level recall
|
||||||
@@ -472,10 +503,7 @@ fn print_report(
|
|||||||
println!(" LongMemEval Benchmark — {}", mode.label);
|
println!(" LongMemEval Benchmark — {}", mode.label);
|
||||||
println!("=================================================================");
|
println!("=================================================================");
|
||||||
println!();
|
println!();
|
||||||
println!(
|
println!("Mode: {}", describe(mode));
|
||||||
"Mode: vector_weight={:.1} / keyword_weight={:.1}",
|
|
||||||
mode.vector_weight, mode.keyword_weight
|
|
||||||
);
|
|
||||||
println!();
|
println!();
|
||||||
println!("Scoring target: RETRIEVAL RECALL (did the gold memory land in top-k).");
|
println!("Scoring target: RETRIEVAL RECALL (did the gold memory land in top-k).");
|
||||||
println!(" No answer is generated or scored. This is NOT the official");
|
println!(" No answer is generated or scored. This is NOT the official");
|
||||||
@@ -602,10 +630,7 @@ fn print_report(
|
|||||||
println!("```json");
|
println!("```json");
|
||||||
println!("{{");
|
println!("{{");
|
||||||
println!(" \"benchmark\": \"longmemeval\",");
|
println!(" \"benchmark\": \"longmemeval\",");
|
||||||
println!(
|
println!(" \"mode\": \"{}\",", describe(mode));
|
||||||
" \"mode\": \"vector_{:.1}_keyword_{:.1}\",",
|
|
||||||
mode.vector_weight, mode.keyword_weight
|
|
||||||
);
|
|
||||||
println!(" \"dataset_variant\": \"{}\",", profile.variant());
|
println!(" \"dataset_variant\": \"{}\",", profile.variant());
|
||||||
println!(" \"scoring_target\": \"retrieval_recall\",");
|
println!(" \"scoring_target\": \"retrieval_recall\",");
|
||||||
println!(" \"k\": 10,");
|
println!(" \"k\": 10,");
|
||||||
@@ -768,18 +793,27 @@ fn main() {
|
|||||||
if sweep {
|
if sweep {
|
||||||
sweep_modes()
|
sweep_modes()
|
||||||
} else {
|
} else {
|
||||||
vec![BM25_ONLY, VECTOR_ONLY, HYBRID]
|
vec![
|
||||||
|
BM25_ONLY,
|
||||||
|
VECTOR_ONLY,
|
||||||
|
HYBRID,
|
||||||
|
RRF,
|
||||||
|
BM25_STEMMED,
|
||||||
|
HYBRID_STEMMED,
|
||||||
|
]
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
#[cfg(not(feature = "embeddings"))]
|
#[cfg(not(feature = "embeddings"))]
|
||||||
{
|
{
|
||||||
vec![BM25_ONLY]
|
vec![BM25_ONLY, BM25_STEMMED]
|
||||||
}
|
}
|
||||||
} else {
|
} else {
|
||||||
if sweep {
|
if sweep {
|
||||||
eprintln!("warning: --sweep needs --embeddings; running BM25 only");
|
eprintln!("warning: --sweep needs --embeddings; running BM25 only");
|
||||||
}
|
}
|
||||||
vec![BM25_ONLY]
|
// Stemming is a property of the keyword stage, so it can be compared
|
||||||
|
// without a model.
|
||||||
|
vec![BM25_ONLY, BM25_STEMMED]
|
||||||
};
|
};
|
||||||
|
|
||||||
for (mode_idx, mode) in modes.iter().enumerate() {
|
for (mode_idx, mode) in modes.iter().enumerate() {
|
||||||
|
|||||||
@@ -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()
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -485,9 +485,13 @@ fn main() {
|
|||||||
|
|
||||||
let mut json = Vec::new();
|
let mut json = Vec::new();
|
||||||
println!("## Search harness");
|
println!("## Search harness");
|
||||||
|
// `--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 {
|
for &n in sizes {
|
||||||
bench_ann(n, &mut json);
|
bench_ann(n, &mut json);
|
||||||
}
|
}
|
||||||
|
}
|
||||||
|
|
||||||
if ann_only {
|
if ann_only {
|
||||||
return;
|
return;
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
[package]
|
[package]
|
||||||
name = "clawhdf5-cli"
|
name = "clawhdf5-cli"
|
||||||
version = "2.4.0"
|
version = "2.5.0"
|
||||||
edition = "2024"
|
edition = "2024"
|
||||||
license = "MIT"
|
license = "MIT"
|
||||||
description = "CLI for clawhdf5 agent memory — create, save, search, recall, stats"
|
description = "CLI for clawhdf5 agent memory — create, save, search, recall, stats"
|
||||||
@@ -14,7 +14,7 @@ name = "clawhdf5"
|
|||||||
path = "src/main.rs"
|
path = "src/main.rs"
|
||||||
|
|
||||||
[dependencies]
|
[dependencies]
|
||||||
clawhdf5-agent = { path = "../clawhdf5-agent", version = "2.4.0" }
|
clawhdf5-agent = { path = "../clawhdf5-agent", version = "2.5.0" }
|
||||||
clap = { version = "4", features = ["derive", "env"] }
|
clap = { version = "4", features = ["derive", "env"] }
|
||||||
serde_json = "1"
|
serde_json = "1"
|
||||||
serde = { workspace = true }
|
serde = { workspace = true }
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
[package]
|
[package]
|
||||||
name = "clawhdf5-derive"
|
name = "clawhdf5-derive"
|
||||||
version = "2.4.0"
|
version = "2.5.0"
|
||||||
edition = "2024"
|
edition = "2024"
|
||||||
description = "Derive macros for rustyhdf5 HDF5 traits"
|
description = "Derive macros for rustyhdf5 HDF5 traits"
|
||||||
license = "MIT"
|
license = "MIT"
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
[package]
|
[package]
|
||||||
name = "clawhdf5-filters"
|
name = "clawhdf5-filters"
|
||||||
version = "2.4.0"
|
version = "2.5.0"
|
||||||
edition = "2024"
|
edition = "2024"
|
||||||
description = "Filter and compression pipeline for clawhdf5"
|
description = "Filter and compression pipeline for clawhdf5"
|
||||||
license = "MIT"
|
license = "MIT"
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
[package]
|
[package]
|
||||||
name = "clawhdf5-format"
|
name = "clawhdf5-format"
|
||||||
version = "2.4.0"
|
version = "2.5.0"
|
||||||
edition = "2024"
|
edition = "2024"
|
||||||
description = "Pure-Rust HDF5 binary format parsing and writing — no C dependencies"
|
description = "Pure-Rust HDF5 binary format parsing and writing — no C dependencies"
|
||||||
license = "MIT"
|
license = "MIT"
|
||||||
@@ -25,7 +25,7 @@ pco = { version = "1.0", optional = true }
|
|||||||
[dev-dependencies]
|
[dev-dependencies]
|
||||||
serde_json = "1"
|
serde_json = "1"
|
||||||
criterion = { workspace = true }
|
criterion = { workspace = true }
|
||||||
clawhdf5-derive = { path = "../clawhdf5-derive", version = "2.4.0" }
|
clawhdf5-derive = { path = "../clawhdf5-derive", version = "2.5.0" }
|
||||||
|
|
||||||
[[bench]]
|
[[bench]]
|
||||||
name = "bench"
|
name = "bench"
|
||||||
|
|||||||
@@ -374,6 +374,11 @@ impl ChunkCache {
|
|||||||
|
|
||||||
// ----- Index operations -----
|
// ----- 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`.
|
/// 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
|
/// 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
|
/// process when the allocation fails; a size taken from the file must surface
|
||||||
/// as an error instead.
|
/// as an error instead.
|
||||||
pub(crate) fn alloc_output(len: usize) -> Result<Vec<u8>, FormatError> {
|
pub(crate) fn alloc_output(len: usize) -> Result<Vec<u8>, FormatError> {
|
||||||
let mut out = Vec::new();
|
if len == 0 {
|
||||||
out.try_reserve_exact(len).map_err(|_| {
|
return Ok(Vec::new());
|
||||||
FormatError::Overflow(format!("cannot allocate {len} bytes for dataset output"))
|
}
|
||||||
})?;
|
let failed =
|
||||||
out.resize(len, 0);
|
|| FormatError::Overflow(format!("cannot allocate {len} bytes for dataset output"));
|
||||||
Ok(out)
|
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> {
|
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.
|
/// 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,
|
/// Every allocated chunk of a chunked dataset, for any supported chunk index,
|
||||||
/// plus the spatial chunk dimensions. Chunks the file never allocated (sparse
|
/// plus the spatial chunk dimensions. Chunks the file never allocated (sparse
|
||||||
/// datasets) are simply absent from the list.
|
/// datasets) are simply absent from the list.
|
||||||
@@ -487,6 +614,18 @@ pub fn list_chunks(
|
|||||||
length_size,
|
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) => {
|
(v, idx) => {
|
||||||
return Err(FormatError::ChunkedReadError(format!(
|
return Err(FormatError::ChunkedReadError(format!(
|
||||||
"unsupported chunked layout version={v}, index_type={idx:?}"
|
"unsupported chunked layout version={v}, index_type={idx:?}"
|
||||||
@@ -629,29 +768,12 @@ pub fn read_chunked_data_cached(
|
|||||||
length_size: u8,
|
length_size: u8,
|
||||||
cache: &ChunkCache,
|
cache: &ChunkCache,
|
||||||
) -> Result<Vec<u8>, FormatError> {
|
) -> Result<Vec<u8>, FormatError> {
|
||||||
let (
|
let (chunk_dimensions, addr_opt) = match layout {
|
||||||
chunk_dimensions,
|
|
||||||
version,
|
|
||||||
chunk_index_type,
|
|
||||||
addr_opt,
|
|
||||||
single_filtered_size,
|
|
||||||
single_filter_mask,
|
|
||||||
) = match layout {
|
|
||||||
DataLayout::Chunked {
|
DataLayout::Chunked {
|
||||||
chunk_dimensions,
|
chunk_dimensions,
|
||||||
btree_address,
|
btree_address,
|
||||||
version,
|
..
|
||||||
chunk_index_type,
|
} => (chunk_dimensions, *btree_address),
|
||||||
single_chunk_filtered_size,
|
|
||||||
single_chunk_filter_mask,
|
|
||||||
} => (
|
|
||||||
chunk_dimensions,
|
|
||||||
*version,
|
|
||||||
*chunk_index_type,
|
|
||||||
*btree_address,
|
|
||||||
*single_chunk_filtered_size,
|
|
||||||
*single_chunk_filter_mask,
|
|
||||||
),
|
|
||||||
_ => {
|
_ => {
|
||||||
return Err(FormatError::ChunkedReadError(
|
return Err(FormatError::ChunkedReadError(
|
||||||
"expected chunked layout".into(),
|
"expected chunked layout".into(),
|
||||||
@@ -688,69 +810,14 @@ pub fn read_chunked_data_cached(
|
|||||||
|
|
||||||
// Populate chunk index on first access
|
// Populate chunk index on first access
|
||||||
if !cache.has_index() {
|
if !cache.has_index() {
|
||||||
let chunks = match (version, chunk_index_type) {
|
let (chunks, _) = list_chunks(
|
||||||
(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,
|
file_data,
|
||||||
&header,
|
layout,
|
||||||
&dataspace.dimensions,
|
dataspace,
|
||||||
spatial_chunk_dims,
|
elem_size,
|
||||||
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,
|
offset_size,
|
||||||
length_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:?}"
|
|
||||||
)));
|
|
||||||
}
|
|
||||||
};
|
|
||||||
cache.populate_index(&chunks, rank);
|
cache.populate_index(&chunks, rank);
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -777,43 +844,20 @@ pub fn read_chunked_data_cached(
|
|||||||
|
|
||||||
let chunk_total_bytes = checked_chunk_byte_len(&chunk_dims, elem_size)?;
|
let chunk_total_bytes = checked_chunk_byte_len(&chunk_dims, elem_size)?;
|
||||||
|
|
||||||
for chunk_info in &chunks {
|
let mut place = |data: &[u8], chunk_info: &ChunkInfo| {
|
||||||
let coord: Vec<u64> = chunk_info.offsets.iter().take(rank).copied().collect();
|
if rank == 0 {
|
||||||
|
let copy_len = data.len().min(output.len());
|
||||||
// Try decompressed cache first
|
output[..copy_len].copy_from_slice(&data[..copy_len]);
|
||||||
let decompressed = if let Some(cached) = cache.get_decompressed_aligned(&coord) {
|
return;
|
||||||
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 chunk_offsets: Vec<usize> = chunk_info
|
let chunk_offsets: Vec<usize> = chunk_info
|
||||||
.offsets
|
.offsets
|
||||||
.iter()
|
.iter()
|
||||||
.take(rank)
|
.take(rank)
|
||||||
.map(|&o| o as usize)
|
.map(|&o| o as usize)
|
||||||
.collect();
|
.collect();
|
||||||
|
|
||||||
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(
|
copy_chunk_to_output(
|
||||||
&decompressed,
|
data,
|
||||||
&mut output,
|
&mut output,
|
||||||
&chunk_offsets,
|
&chunk_offsets,
|
||||||
&chunk_dims,
|
&chunk_dims,
|
||||||
@@ -823,6 +867,63 @@ pub fn read_chunked_data_cached(
|
|||||||
elem_size,
|
elem_size,
|
||||||
rank,
|
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])
|
||||||
|
};
|
||||||
|
|
||||||
|
// 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,
|
cache: &ChunkCache,
|
||||||
sweep: &mut SweepContext,
|
sweep: &mut SweepContext,
|
||||||
) -> Result<Vec<u8>, FormatError> {
|
) -> Result<Vec<u8>, FormatError> {
|
||||||
let (
|
let (chunk_dimensions, addr_opt) = match layout {
|
||||||
chunk_dimensions,
|
|
||||||
version,
|
|
||||||
chunk_index_type,
|
|
||||||
addr_opt,
|
|
||||||
single_filtered_size,
|
|
||||||
single_filter_mask,
|
|
||||||
) = match layout {
|
|
||||||
DataLayout::Chunked {
|
DataLayout::Chunked {
|
||||||
chunk_dimensions,
|
chunk_dimensions,
|
||||||
btree_address,
|
btree_address,
|
||||||
version,
|
..
|
||||||
chunk_index_type,
|
} => (chunk_dimensions, *btree_address),
|
||||||
single_chunk_filtered_size,
|
|
||||||
single_chunk_filter_mask,
|
|
||||||
} => (
|
|
||||||
chunk_dimensions,
|
|
||||||
*version,
|
|
||||||
*chunk_index_type,
|
|
||||||
*btree_address,
|
|
||||||
*single_chunk_filtered_size,
|
|
||||||
*single_chunk_filter_mask,
|
|
||||||
),
|
|
||||||
_ => {
|
_ => {
|
||||||
return Err(FormatError::ChunkedReadError(
|
return Err(FormatError::ChunkedReadError(
|
||||||
"expected chunked layout".into(),
|
"expected chunked layout".into(),
|
||||||
@@ -1044,69 +1128,14 @@ pub fn read_chunked_data_sweep(
|
|||||||
|
|
||||||
// Populate chunk index on first access
|
// Populate chunk index on first access
|
||||||
if !cache.has_index() {
|
if !cache.has_index() {
|
||||||
let chunks = match (version, chunk_index_type) {
|
let (chunks, _) = list_chunks(
|
||||||
(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,
|
file_data,
|
||||||
&header,
|
layout,
|
||||||
&dataspace.dimensions,
|
dataspace,
|
||||||
spatial_chunk_dims,
|
elem_size,
|
||||||
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,
|
offset_size,
|
||||||
length_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:?}"
|
|
||||||
)));
|
|
||||||
}
|
|
||||||
};
|
|
||||||
cache.populate_index(&chunks, rank);
|
cache.populate_index(&chunks, rank);
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -1211,29 +1240,12 @@ pub fn read_chunked_data_indexed(
|
|||||||
length_size: u8,
|
length_size: u8,
|
||||||
cache: &ChunkCache,
|
cache: &ChunkCache,
|
||||||
) -> Result<Vec<u8>, FormatError> {
|
) -> Result<Vec<u8>, FormatError> {
|
||||||
let (
|
let (chunk_dimensions, addr_opt) = match layout {
|
||||||
chunk_dimensions,
|
|
||||||
version,
|
|
||||||
chunk_index_type,
|
|
||||||
addr_opt,
|
|
||||||
single_filtered_size,
|
|
||||||
single_filter_mask,
|
|
||||||
) = match layout {
|
|
||||||
DataLayout::Chunked {
|
DataLayout::Chunked {
|
||||||
chunk_dimensions,
|
chunk_dimensions,
|
||||||
btree_address,
|
btree_address,
|
||||||
version,
|
..
|
||||||
chunk_index_type,
|
} => (chunk_dimensions, *btree_address),
|
||||||
single_chunk_filtered_size,
|
|
||||||
single_chunk_filter_mask,
|
|
||||||
} => (
|
|
||||||
chunk_dimensions,
|
|
||||||
*version,
|
|
||||||
*chunk_index_type,
|
|
||||||
*btree_address,
|
|
||||||
*single_chunk_filtered_size,
|
|
||||||
*single_chunk_filter_mask,
|
|
||||||
),
|
|
||||||
_ => {
|
_ => {
|
||||||
return Err(FormatError::ChunkedReadError(
|
return Err(FormatError::ChunkedReadError(
|
||||||
"expected chunked layout".into(),
|
"expected chunked layout".into(),
|
||||||
@@ -1270,69 +1282,14 @@ pub fn read_chunked_data_indexed(
|
|||||||
|
|
||||||
// Build chunk index on first access
|
// Build chunk index on first access
|
||||||
if !cache.has_chunk_index() {
|
if !cache.has_chunk_index() {
|
||||||
let chunks = match (version, chunk_index_type) {
|
let (chunks, _) = list_chunks(
|
||||||
(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,
|
file_data,
|
||||||
&header,
|
layout,
|
||||||
&dataspace.dimensions,
|
dataspace,
|
||||||
spatial_chunk_dims,
|
elem_size,
|
||||||
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,
|
offset_size,
|
||||||
length_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:?}"
|
|
||||||
)));
|
|
||||||
}
|
|
||||||
};
|
|
||||||
cache.populate_chunk_index(&chunks, rank);
|
cache.populate_chunk_index(&chunks, rank);
|
||||||
// Also populate the legacy index for compatibility
|
// Also populate the legacy index for compatibility
|
||||||
if !cache.has_index() {
|
if !cache.has_index() {
|
||||||
@@ -2284,21 +2241,23 @@ mod tests {
|
|||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[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 values: Vec<f64> = (0..20).map(|i| i as f64).collect();
|
||||||
let (file_data, layout, dataspace) = build_1d_chunked_file(&values, 10);
|
let (file_data, layout, dataspace) = build_1d_chunked_file(&values, 10);
|
||||||
let datatype = make_f64_type();
|
let datatype = make_f64_type();
|
||||||
let cache = ChunkCache::new();
|
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(
|
let raw1 = read_chunked_data_cached(
|
||||||
&file_data, &layout, &dataspace, &datatype, None, 8, 8, &cache,
|
&file_data, &layout, &dataspace, &datatype, None, 8, 8, &cache,
|
||||||
)
|
)
|
||||||
.unwrap();
|
.unwrap();
|
||||||
assert!(cache.has_index());
|
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(
|
let raw2 = read_chunked_data_cached(
|
||||||
&file_data, &layout, &dataspace, &datatype, None, 8, 8, &cache,
|
&file_data, &layout, &dataspace, &datatype, None, 8, 8, &cache,
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -49,6 +49,38 @@ pub struct ChunkOptions {
|
|||||||
pub pcodec: bool,
|
pub pcodec: bool,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Largest chunk the automatic choice produces, in bytes.
|
||||||
|
const AUTO_CHUNK_TARGET_BYTES: u64 = 1 << 20;
|
||||||
|
|
||||||
|
/// Extent assumed for a dimension that is currently empty (an unlimited
|
||||||
|
/// dimension not yet written to) — the same stand-in h5py uses.
|
||||||
|
const AUTO_CHUNK_EMPTY_DIM: u64 = 1024;
|
||||||
|
|
||||||
|
/// Choose chunk dimensions for a dataset nobody specified them for.
|
||||||
|
///
|
||||||
|
/// Asking for compression (or any filter) without chunk dimensions used to
|
||||||
|
/// make the whole dataset one chunk. That defeats the point of chunking: any
|
||||||
|
/// read — even a single row — must decompress everything, and a large dataset
|
||||||
|
/// cannot be decompressed in parallel. Datasets up to the target size stay a
|
||||||
|
/// single chunk, exactly as before; larger ones are split by halving the
|
||||||
|
/// dimensions in turn (so chunks keep roughly the dataset's proportions, the
|
||||||
|
/// approach h5py takes) until a chunk fits the target.
|
||||||
|
pub fn auto_chunk_dims(shape: &[u64], elem_size: usize) -> Vec<u64> {
|
||||||
|
let mut dims: Vec<u64> = shape
|
||||||
|
.iter()
|
||||||
|
.map(|&d| if d == 0 { AUTO_CHUNK_EMPTY_DIM } else { d })
|
||||||
|
.collect();
|
||||||
|
let elem = elem_size.max(1) as u64;
|
||||||
|
let bytes = |dims: &[u64]| dims.iter().fold(elem, |acc, &d| acc.saturating_mul(d));
|
||||||
|
let mut axis = 0;
|
||||||
|
while bytes(&dims) > AUTO_CHUNK_TARGET_BYTES && dims.iter().any(|&d| d > 1) {
|
||||||
|
let i = axis % dims.len();
|
||||||
|
dims[i] = dims[i].div_ceil(2);
|
||||||
|
axis += 1;
|
||||||
|
}
|
||||||
|
dims
|
||||||
|
}
|
||||||
|
|
||||||
impl ChunkOptions {
|
impl ChunkOptions {
|
||||||
/// Whether any chunking option is enabled.
|
/// Whether any chunking option is enabled.
|
||||||
pub fn is_chunked(&self) -> bool {
|
pub fn is_chunked(&self) -> bool {
|
||||||
@@ -135,11 +167,17 @@ impl ChunkOptions {
|
|||||||
|
|
||||||
/// Determine chunk dimensions, using user-specified or auto-computing.
|
/// Determine chunk dimensions, using user-specified or auto-computing.
|
||||||
pub fn resolve_chunk_dims(&self, shape: &[u64]) -> Vec<u64> {
|
pub fn resolve_chunk_dims(&self, shape: &[u64]) -> Vec<u64> {
|
||||||
if let Some(ref dims) = self.chunk_dims {
|
// Without the element size, assume 8 bytes (the widest common scalar);
|
||||||
dims.clone()
|
// the writer uses `resolve_chunk_dims_for`.
|
||||||
} else {
|
self.resolve_chunk_dims_for(shape, 8)
|
||||||
// Auto chunk: use the full dataset shape (single chunk)
|
}
|
||||||
shape.to_vec()
|
|
||||||
|
/// Chunk dimensions for a dataset of `shape` whose elements are `elem_size`
|
||||||
|
/// bytes: the caller's if given, otherwise chosen automatically.
|
||||||
|
pub fn resolve_chunk_dims_for(&self, shape: &[u64], elem_size: usize) -> Vec<u64> {
|
||||||
|
match self.chunk_dims {
|
||||||
|
Some(ref dims) => dims.clone(),
|
||||||
|
None => auto_chunk_dims(shape, elem_size),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -1143,6 +1181,45 @@ mod tests {
|
|||||||
assert_eq!(dims, vec![100, 50]);
|
assert_eq!(dims, vec![100, 50]);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn auto_chunking_splits_only_large_datasets() {
|
||||||
|
let bytes = |dims: &[u64], elem: u64| dims.iter().product::<u64>() * elem;
|
||||||
|
// Up to the target: one chunk, as before.
|
||||||
|
assert_eq!(auto_chunk_dims(&[100, 50], 8), [100, 50]);
|
||||||
|
assert_eq!(auto_chunk_dims(&[131_072], 8), [131_072]); // exactly 1 MiB
|
||||||
|
// Larger: split, keeping proportions, never above the target.
|
||||||
|
let big = auto_chunk_dims(&[4096, 2048], 8);
|
||||||
|
assert!(bytes(&big, 8) <= AUTO_CHUNK_TARGET_BYTES, "{big:?}");
|
||||||
|
assert!(bytes(&big, 8) > AUTO_CHUNK_TARGET_BYTES / 4, "{big:?}");
|
||||||
|
assert_eq!(big[0] / big[1], 2, "proportions kept: {big:?}");
|
||||||
|
// Every dimension stays within the dataset and at least 1.
|
||||||
|
for shape in [
|
||||||
|
vec![10_000_000u64],
|
||||||
|
vec![3, 5_000_000],
|
||||||
|
vec![1, 1, 9_000_000],
|
||||||
|
vec![7; 9],
|
||||||
|
] {
|
||||||
|
let dims = auto_chunk_dims(&shape, 4);
|
||||||
|
assert!(
|
||||||
|
dims.iter().zip(&shape).all(|(c, s)| *c >= 1 && c <= s),
|
||||||
|
"{shape:?} -> {dims:?}"
|
||||||
|
);
|
||||||
|
assert!(
|
||||||
|
bytes(&dims, 4) <= AUTO_CHUNK_TARGET_BYTES,
|
||||||
|
"{shape:?} -> {dims:?}"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
// An empty (unlimited, unwritten) dimension still gets a usable chunk.
|
||||||
|
let growable = auto_chunk_dims(&[0, 128], 8);
|
||||||
|
assert!(growable[0] >= 1 && bytes(&growable, 8) <= AUTO_CHUNK_TARGET_BYTES);
|
||||||
|
// Explicit dimensions always win.
|
||||||
|
let explicit = ChunkOptions {
|
||||||
|
chunk_dims: Some(vec![10, 10]),
|
||||||
|
..Default::default()
|
||||||
|
};
|
||||||
|
assert_eq!(explicit.resolve_chunk_dims_for(&[4096, 2048], 8), [10, 10]);
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn chunk_options_pipeline_deflate() {
|
fn chunk_options_pipeline_deflate() {
|
||||||
// Auto-shuffle is applied before compression by default (matches h5py).
|
// Auto-shuffle is applied before compression by default (matches h5py).
|
||||||
|
|||||||
@@ -307,6 +307,24 @@ pub fn read_raw_data_selection(
|
|||||||
) -> Result<Vec<u8>, FormatError> {
|
) -> Result<Vec<u8>, FormatError> {
|
||||||
use crate::selection::Selection;
|
use crate::selection::Selection;
|
||||||
|
|
||||||
|
crate::partial_read::validate(selection, &dataspace.dimensions)?;
|
||||||
|
|
||||||
|
// Read only what the selection's bounding box touches when that is
|
||||||
|
// possible; everything below is the decode-everything-then-pick path,
|
||||||
|
// kept for the cases `partial_read` declines.
|
||||||
|
if let Some(selected) = crate::partial_read::read_selection(
|
||||||
|
file_data,
|
||||||
|
layout,
|
||||||
|
dataspace,
|
||||||
|
datatype.type_size() as usize,
|
||||||
|
pipeline,
|
||||||
|
offset_size,
|
||||||
|
length_size,
|
||||||
|
selection,
|
||||||
|
)? {
|
||||||
|
return Ok(selected);
|
||||||
|
}
|
||||||
|
|
||||||
match selection {
|
match selection {
|
||||||
Selection::All => {
|
Selection::All => {
|
||||||
return read_raw_data_full(
|
return read_raw_data_full(
|
||||||
@@ -858,6 +876,30 @@ fn get_size(dt: &Datatype) -> usize {
|
|||||||
dt.type_size() as usize
|
dt.type_size() as usize
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Reinterpret little-endian bytes as `count` native values of `T` on a
|
||||||
|
/// little-endian target, in one copy.
|
||||||
|
///
|
||||||
|
/// The buffer is allocated uninitialised and filled by the copy. It used to be
|
||||||
|
/// `vec![0; count]` first, which for a large dataset meant writing every page
|
||||||
|
/// twice (zero it, then overwrite it) — about as expensive as the copy itself.
|
||||||
|
#[cfg(target_endian = "little")]
|
||||||
|
fn native_le_to_vec<T: Copy>(raw: &[u8], count: usize) -> Vec<T> {
|
||||||
|
let bytes = count * core::mem::size_of::<T>();
|
||||||
|
debug_assert!(bytes <= raw.len());
|
||||||
|
let mut result: Vec<T> = Vec::with_capacity(count);
|
||||||
|
// SAFETY: `result` has capacity for `count` values of `T`, i.e. `bytes`
|
||||||
|
// bytes; `raw` holds at least `bytes` bytes (callers derive `count` from
|
||||||
|
// `raw.len() / size_of::<T>()`); the regions cannot overlap because
|
||||||
|
// `result` was just allocated. Every `T` used here (f32/f64/i32/i64) is
|
||||||
|
// valid for any bit pattern, so after the copy all `count` values are
|
||||||
|
// initialised and `set_len` is sound.
|
||||||
|
unsafe {
|
||||||
|
core::ptr::copy_nonoverlapping(raw.as_ptr(), result.as_mut_ptr().cast::<u8>(), bytes);
|
||||||
|
result.set_len(count);
|
||||||
|
}
|
||||||
|
result
|
||||||
|
}
|
||||||
|
|
||||||
/// Convert raw bytes to `f64` values.
|
/// Convert raw bytes to `f64` values.
|
||||||
pub fn read_as_f64(raw: &[u8], datatype: &Datatype) -> Result<Vec<f64>, FormatError> {
|
pub fn read_as_f64(raw: &[u8], datatype: &Datatype) -> Result<Vec<f64>, FormatError> {
|
||||||
// Array datatypes (e.g. an array-typed compound member) are read as a flat
|
// Array datatypes (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];
|
return Ok(native_le_to_vec::<f64>(raw, count));
|
||||||
// SAFETY: On LE platforms, f64 in-memory representation matches LE bytes.
|
|
||||||
// We copy raw bytes directly into the f64 buffer.
|
|
||||||
// SAFETY: The byte slice is properly aligned for this type and the length is divisible by size_of::<T>().
|
|
||||||
unsafe {
|
|
||||||
core::ptr::copy_nonoverlapping(raw.as_ptr(), result.as_mut_ptr() as *mut u8, raw.len());
|
|
||||||
}
|
|
||||||
return Ok(result);
|
|
||||||
}
|
}
|
||||||
|
|
||||||
let order = get_byte_order(datatype);
|
let order = get_byte_order(datatype);
|
||||||
@@ -975,12 +1010,7 @@ pub fn read_as_i64(raw: &[u8], datatype: &Datatype) -> Result<Vec<i64>, FormatEr
|
|||||||
}
|
}
|
||||||
)
|
)
|
||||||
{
|
{
|
||||||
let mut result = vec![0i64; count];
|
return Ok(native_le_to_vec::<i64>(raw, count));
|
||||||
// SAFETY: The byte slice is properly aligned for this type and the length is divisible by size_of::<T>().
|
|
||||||
unsafe {
|
|
||||||
core::ptr::copy_nonoverlapping(raw.as_ptr(), result.as_mut_ptr() as *mut u8, raw.len());
|
|
||||||
}
|
|
||||||
return Ok(result);
|
|
||||||
}
|
}
|
||||||
|
|
||||||
let order = get_byte_order(datatype);
|
let order = get_byte_order(datatype);
|
||||||
@@ -1044,12 +1074,7 @@ pub fn read_as_f32(raw: &[u8], datatype: &Datatype) -> Result<Vec<f32>, FormatEr
|
|||||||
..
|
..
|
||||||
}
|
}
|
||||||
) {
|
) {
|
||||||
let mut result = vec![0.0f32; count];
|
return Ok(native_le_to_vec::<f32>(raw, count));
|
||||||
// SAFETY: The byte slice is properly aligned for this type and the length is divisible by size_of::<T>().
|
|
||||||
unsafe {
|
|
||||||
core::ptr::copy_nonoverlapping(raw.as_ptr(), result.as_mut_ptr() as *mut u8, raw.len());
|
|
||||||
}
|
|
||||||
return Ok(result);
|
|
||||||
}
|
}
|
||||||
|
|
||||||
let order = get_byte_order(datatype);
|
let order = get_byte_order(datatype);
|
||||||
@@ -1126,12 +1151,7 @@ pub fn read_as_i32(raw: &[u8], datatype: &Datatype) -> Result<Vec<i32>, FormatEr
|
|||||||
}
|
}
|
||||||
)
|
)
|
||||||
{
|
{
|
||||||
let mut result = vec![0i32; count];
|
return Ok(native_le_to_vec::<i32>(raw, count));
|
||||||
// SAFETY: The byte slice is properly aligned for this type and the length is divisible by size_of::<T>().
|
|
||||||
unsafe {
|
|
||||||
core::ptr::copy_nonoverlapping(raw.as_ptr(), result.as_mut_ptr() as *mut u8, raw.len());
|
|
||||||
}
|
|
||||||
return Ok(result);
|
|
||||||
}
|
}
|
||||||
|
|
||||||
let order = get_byte_order(datatype);
|
let order = get_byte_order(datatype);
|
||||||
@@ -1407,6 +1427,26 @@ pub fn read_object_references(
|
|||||||
}
|
}
|
||||||
Ok(result)
|
Ok(result)
|
||||||
}
|
}
|
||||||
|
Datatype::Reference {
|
||||||
|
ref_type: crate::datatype::ReferenceType::Object2,
|
||||||
|
size,
|
||||||
|
} => {
|
||||||
|
let elem_size = *size as usize;
|
||||||
|
if elem_size == 0 {
|
||||||
|
return Ok(Vec::new());
|
||||||
|
}
|
||||||
|
if !raw.len().is_multiple_of(elem_size) {
|
||||||
|
return Err(FormatError::DataSizeMismatch {
|
||||||
|
expected: 0,
|
||||||
|
actual: raw.len(),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
raw.chunks_exact(elem_size)
|
||||||
|
.map(|element| {
|
||||||
|
decode_std_object_ref(element).map(|address| ObjectReference { address })
|
||||||
|
})
|
||||||
|
.collect()
|
||||||
|
}
|
||||||
_ => Err(FormatError::TypeMismatch {
|
_ => Err(FormatError::TypeMismatch {
|
||||||
expected: "Reference(Object)",
|
expected: "Reference(Object)",
|
||||||
actual: datatype_name(datatype),
|
actual: datatype_name(datatype),
|
||||||
@@ -1414,6 +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.
|
/// Read region references from raw bytes.
|
||||||
///
|
///
|
||||||
/// Region references encode a dataset selection (hyperslab, point list, etc.)
|
/// Region references encode a dataset selection (hyperslab, point list, etc.)
|
||||||
|
|||||||
@@ -36,8 +36,18 @@ pub enum CharacterSet {
|
|||||||
/// Reference type.
|
/// Reference type.
|
||||||
#[derive(Debug, Clone, PartialEq)]
|
#[derive(Debug, Clone, PartialEq)]
|
||||||
pub enum ReferenceType {
|
pub enum ReferenceType {
|
||||||
|
/// Legacy object reference: the target's object header address.
|
||||||
Object,
|
Object,
|
||||||
|
/// Legacy dataset region reference.
|
||||||
DatasetRegion,
|
DatasetRegion,
|
||||||
|
/// `H5T_STD_REF` object reference (HDF5 1.12+, datatype message version
|
||||||
|
/// 4): a small header followed by an object token. Decoded by
|
||||||
|
/// `data_read::read_object_references`.
|
||||||
|
Object2,
|
||||||
|
/// `H5T_STD_REF` dataset region reference.
|
||||||
|
DatasetRegion2,
|
||||||
|
/// `H5T_STD_REF` attribute reference.
|
||||||
|
Attribute,
|
||||||
}
|
}
|
||||||
|
|
||||||
/// A member of a compound datatype.
|
/// A member of a compound datatype.
|
||||||
@@ -424,9 +434,15 @@ impl Datatype {
|
|||||||
7 => {
|
7 => {
|
||||||
// Reference
|
// Reference
|
||||||
let ref_type_val = bf0 & 0x0F;
|
let ref_type_val = bf0 & 0x0F;
|
||||||
let ref_type = match ref_type_val {
|
// Datatype message version 4 (HDF5 1.12) revised this class:
|
||||||
0 => ReferenceType::Object,
|
// types 2-4 are the new `H5T_STD_REF` references, and the high
|
||||||
1 => ReferenceType::DatasetRegion,
|
// nibble of the first flag byte carries their encoding version.
|
||||||
|
let ref_type = match (ref_type_val, version) {
|
||||||
|
(0, _) => ReferenceType::Object,
|
||||||
|
(1, _) => ReferenceType::DatasetRegion,
|
||||||
|
(2, 4..) => ReferenceType::Object2,
|
||||||
|
(3, 4..) => ReferenceType::DatasetRegion2,
|
||||||
|
(4, 4..) => ReferenceType::Attribute,
|
||||||
_ => return Err(FormatError::InvalidReferenceType(ref_type_val)),
|
_ => return Err(FormatError::InvalidReferenceType(ref_type_val)),
|
||||||
};
|
};
|
||||||
Ok((Datatype::Reference { size, ref_type }, pos))
|
Ok((Datatype::Reference { size, ref_type }, pos))
|
||||||
@@ -1563,6 +1579,28 @@ mod tests {
|
|||||||
assert_eq!(err, FormatError::InvalidCharacterSet(2));
|
assert_eq!(err, FormatError::InvalidCharacterSet(2));
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_reference_v4_std_ref_from_hdf5_2_0() {
|
||||||
|
// Datatype message of an H5T_STD_REF dataset written by HDF5 2.0:
|
||||||
|
// class 7, version 4, type 2 (object), encoding version 1, 18 bytes.
|
||||||
|
let bytes = [0x47, 0x12, 0x00, 0x00, 0x12, 0x00, 0x00, 0x00];
|
||||||
|
let (dt, consumed) = Datatype::parse(&bytes).unwrap();
|
||||||
|
assert_eq!(consumed, 8);
|
||||||
|
assert_eq!(
|
||||||
|
dt,
|
||||||
|
Datatype::Reference {
|
||||||
|
size: 18,
|
||||||
|
ref_type: ReferenceType::Object2
|
||||||
|
}
|
||||||
|
);
|
||||||
|
// The new types are only valid from datatype version 4.
|
||||||
|
let old_version = [0x37, 0x12, 0x00, 0x00, 0x12, 0x00, 0x00, 0x00];
|
||||||
|
assert_eq!(
|
||||||
|
Datatype::parse(&old_version).unwrap_err(),
|
||||||
|
FormatError::InvalidReferenceType(2)
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn test_error_invalid_reference_type() {
|
fn test_error_invalid_reference_type() {
|
||||||
let buf = build_dt_header(7, 1, [5, 0, 0], 8);
|
let buf = build_dt_header(7, 1, [5, 0, 0], 8);
|
||||||
|
|||||||
@@ -117,6 +117,9 @@ pub enum FormatError {
|
|||||||
/// A message is marked shared but was parsed without access to the file,
|
/// A message is marked shared but was parsed without access to the file,
|
||||||
/// so the reference to the real message could not be followed.
|
/// so the reference to the real message could not be followed.
|
||||||
UnresolvedSharedMessage,
|
UnresolvedSharedMessage,
|
||||||
|
/// A selection does not fit the dataset it was applied to (wrong rank, or
|
||||||
|
/// it reaches past a dimension's extent).
|
||||||
|
SelectionOutOfBounds(String),
|
||||||
/// The dataset's raw data is stored in external files (External Data
|
/// The dataset's raw data is stored in external files (External Data
|
||||||
/// Files message), which this reader does not follow.
|
/// Files message), which this reader does not follow.
|
||||||
ExternalDataFilesUnsupported,
|
ExternalDataFilesUnsupported,
|
||||||
@@ -333,6 +336,9 @@ impl fmt::Display for FormatError {
|
|||||||
f,
|
f,
|
||||||
"dataset raw data is stored in external file(s), which is not supported"
|
"dataset raw data is stored in external file(s), which is not supported"
|
||||||
),
|
),
|
||||||
|
FormatError::SelectionOutOfBounds(msg) => {
|
||||||
|
write!(f, "selection out of bounds: {msg}")
|
||||||
|
}
|
||||||
FormatError::UnresolvedSharedMessage => write!(
|
FormatError::UnresolvedSharedMessage => write!(
|
||||||
f,
|
f,
|
||||||
"message is shared but no file data was available to resolve it"
|
"message is shared but no file data was available to resolve it"
|
||||||
|
|||||||
@@ -1221,8 +1221,10 @@ impl FileWriter {
|
|||||||
precompressed: None,
|
precompressed: None,
|
||||||
});
|
});
|
||||||
} else if is_chunked[i] {
|
} else if is_chunked[i] {
|
||||||
let chunk_dims = d.chunk_options.resolve_chunk_dims(&d.ds.dimensions);
|
|
||||||
let elem_size = d.dt.type_size() as usize;
|
let elem_size = d.dt.type_size() as usize;
|
||||||
|
let chunk_dims = d
|
||||||
|
.chunk_options
|
||||||
|
.resolve_chunk_dims_for(&d.ds.dimensions, elem_size);
|
||||||
// Compress once in Pass 1; cache the result so Pass 2 can skip
|
// Compress once in Pass 1; cache the result so Pass 2 can skip
|
||||||
// re-compression and just rebuild the index with real addresses.
|
// re-compression and just rebuild the index with real addresses.
|
||||||
let pre = precompress_chunks(
|
let pre = precompress_chunks(
|
||||||
|
|||||||
@@ -845,9 +845,33 @@ fn shuffle_decompress(data: &[u8], element_size: usize) -> Result<Vec<u8>, Forma
|
|||||||
let num_elements = data.len() / element_size;
|
let num_elements = data.len() / element_size;
|
||||||
let mut result = vec![0u8; data.len()];
|
let mut result = vec![0u8; data.len()];
|
||||||
|
|
||||||
for i in 0..num_elements {
|
// The shuffled stream is `element_size` byte planes of `num_elements`
|
||||||
for j in 0..element_size {
|
// bytes each; un-shuffling interleaves them. This is on the read path of
|
||||||
result[i * element_size + j] = data[j * num_elements + i];
|
// 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());
|
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;
|
pub mod object_header_writer;
|
||||||
#[cfg(feature = "parallel")]
|
#[cfg(feature = "parallel")]
|
||||||
pub mod parallel_read;
|
pub mod parallel_read;
|
||||||
|
pub mod partial_read;
|
||||||
pub mod profiling;
|
pub mod profiling;
|
||||||
pub mod property_list;
|
pub mod property_list;
|
||||||
pub mod selection;
|
pub mod selection;
|
||||||
|
|||||||
@@ -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.
@@ -316,3 +316,97 @@ print('ok')
|
|||||||
// Clean up
|
// Clean up
|
||||||
let _ = std::fs::remove_file(&path);
|
let _ = std::fs::remove_file(&path);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// H5T_STD_REF (HDF5 1.12+ references, datatype message version 4)
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
/// `fixtures/std_ref_hdf5_2_0.h5` (see `gen_std_ref.py`) holds a dataset of
|
||||||
|
/// `H5T_STD_REF` with two object references, written by HDF5 2.0 itself. The
|
||||||
|
/// datatype used to be rejected with `InvalidReferenceType(2)`.
|
||||||
|
#[test]
|
||||||
|
fn std_ref_object_references_from_hdf5_2_0() {
|
||||||
|
use clawhdf5_format::data_layout::DataLayout;
|
||||||
|
use clawhdf5_format::dataspace::Dataspace;
|
||||||
|
use clawhdf5_format::group_v2::resolve_path_any;
|
||||||
|
use clawhdf5_format::message_type::MessageType;
|
||||||
|
use clawhdf5_format::object_header::ObjectHeader;
|
||||||
|
use clawhdf5_format::signature::find_signature;
|
||||||
|
use clawhdf5_format::superblock::Superblock;
|
||||||
|
|
||||||
|
let bytes: &[u8] = include_bytes!("fixtures/std_ref_hdf5_2_0.h5");
|
||||||
|
let sb = Superblock::parse(bytes, find_signature(bytes).unwrap()).unwrap();
|
||||||
|
let (os, ls) = (sb.offset_size, sb.length_size);
|
||||||
|
|
||||||
|
let refs_addr = resolve_path_any(bytes, &sb, "refs").unwrap();
|
||||||
|
let header = ObjectHeader::parse(bytes, refs_addr as usize, os, ls).unwrap();
|
||||||
|
let message = |t: MessageType| {
|
||||||
|
&header
|
||||||
|
.messages
|
||||||
|
.iter()
|
||||||
|
.find(|m| m.msg_type == t)
|
||||||
|
.unwrap()
|
||||||
|
.data
|
||||||
|
};
|
||||||
|
|
||||||
|
let (datatype, _) = Datatype::parse(message(MessageType::Datatype)).unwrap();
|
||||||
|
assert_eq!(
|
||||||
|
datatype,
|
||||||
|
Datatype::Reference {
|
||||||
|
size: 18,
|
||||||
|
ref_type: ReferenceType::Object2
|
||||||
|
}
|
||||||
|
);
|
||||||
|
let dataspace = Dataspace::parse(message(MessageType::Dataspace), ls).unwrap();
|
||||||
|
let layout = DataLayout::parse(message(MessageType::DataLayout), os, ls).unwrap();
|
||||||
|
let raw =
|
||||||
|
clawhdf5_format::data_read::read_raw_data(bytes, &layout, &dataspace, &datatype).unwrap();
|
||||||
|
assert_eq!(raw.len(), 2 * 18);
|
||||||
|
|
||||||
|
// The references point at the objects they were created from.
|
||||||
|
let refs = read_object_references(&raw, &datatype, os).unwrap();
|
||||||
|
let addresses: Vec<u64> = refs.iter().map(|r| r.address).collect();
|
||||||
|
assert_eq!(
|
||||||
|
addresses,
|
||||||
|
[
|
||||||
|
resolve_path_any(bytes, &sb, "target").unwrap(),
|
||||||
|
resolve_path_any(bytes, &sb, "grp").unwrap(),
|
||||||
|
]
|
||||||
|
);
|
||||||
|
// And what they point at is a real object header.
|
||||||
|
for address in addresses {
|
||||||
|
ObjectHeader::parse(bytes, address as usize, os, ls).unwrap();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn std_ref_decoding_rejects_malformed_elements() {
|
||||||
|
let dt = Datatype::Reference {
|
||||||
|
size: 18,
|
||||||
|
ref_type: ReferenceType::Object2,
|
||||||
|
};
|
||||||
|
let mut good = vec![0u8; 18];
|
||||||
|
good[..4].copy_from_slice(&[2, 0, 8, 0xb3]);
|
||||||
|
assert_eq!(
|
||||||
|
read_object_references(&good, &dt, 8).unwrap()[0].address,
|
||||||
|
0xb3
|
||||||
|
);
|
||||||
|
|
||||||
|
// Null reference.
|
||||||
|
assert_eq!(
|
||||||
|
read_object_references(&[0u8; 18], &dt, 8).unwrap()[0].address,
|
||||||
|
u64::MAX
|
||||||
|
);
|
||||||
|
for (what, patch) in [
|
||||||
|
("wrong reference type", (0usize, 3u8)),
|
||||||
|
("external flag", (1, 1)),
|
||||||
|
("token longer than the element", (2, 200)),
|
||||||
|
("zero-length token", (2, 0)),
|
||||||
|
] {
|
||||||
|
let mut bad = good.clone();
|
||||||
|
bad[patch.0] = patch.1;
|
||||||
|
assert!(read_object_references(&bad, &dt, 8).is_err(), "{what}");
|
||||||
|
}
|
||||||
|
// Not a whole number of elements.
|
||||||
|
assert!(read_object_references(&good[..17], &dt, 8).is_err());
|
||||||
|
}
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
[package]
|
[package]
|
||||||
name = "clawhdf5-gpu"
|
name = "clawhdf5-gpu"
|
||||||
version = "2.4.0"
|
version = "2.5.0"
|
||||||
edition = "2024"
|
edition = "2024"
|
||||||
description = "GPU-accelerated vector operations for rustyhdf5 using wgpu compute shaders"
|
description = "GPU-accelerated vector operations for rustyhdf5 using wgpu compute shaders"
|
||||||
license = "MIT"
|
license = "MIT"
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
[package]
|
[package]
|
||||||
name = "clawhdf5-io"
|
name = "clawhdf5-io"
|
||||||
version = "2.4.0"
|
version = "2.5.0"
|
||||||
edition = "2024"
|
edition = "2024"
|
||||||
description = "I/O abstraction layer for rustyhdf5"
|
description = "I/O abstraction layer for rustyhdf5"
|
||||||
license = "MIT"
|
license = "MIT"
|
||||||
@@ -10,7 +10,7 @@ keywords = ["hdf5", "io", "science", "data"]
|
|||||||
categories = ["filesystem", "science"]
|
categories = ["filesystem", "science"]
|
||||||
|
|
||||||
[dependencies]
|
[dependencies]
|
||||||
clawhdf5-format = { path = "../clawhdf5-format", version = "2.4.0" }
|
clawhdf5-format = { path = "../clawhdf5-format", version = "2.5.0" }
|
||||||
memmap2 = { version = "0.9", optional = true }
|
memmap2 = { version = "0.9", optional = true }
|
||||||
libc = { version = "0.2", optional = true }
|
libc = { version = "0.2", optional = true }
|
||||||
tokio = { version = "1", features = ["fs", "io-util"], optional = true }
|
tokio = { version = "1", features = ["fs", "io-util"], optional = true }
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
[package]
|
[package]
|
||||||
name = "clawhdf5-migrate"
|
name = "clawhdf5-migrate"
|
||||||
version = "2.4.0"
|
version = "2.5.0"
|
||||||
edition = "2024"
|
edition = "2024"
|
||||||
description = "CLI to migrate SQLite agent memory databases to HDF5 format"
|
description = "CLI to migrate SQLite agent memory databases to HDF5 format"
|
||||||
license = "MIT"
|
license = "MIT"
|
||||||
@@ -14,9 +14,9 @@ name = "clawhdf5-migrate"
|
|||||||
path = "src/main.rs"
|
path = "src/main.rs"
|
||||||
|
|
||||||
[dependencies]
|
[dependencies]
|
||||||
clawhdf5-agent = { path = "../clawhdf5-agent", version = "2.4.0" }
|
clawhdf5-agent = { path = "../clawhdf5-agent", version = "2.5.0" }
|
||||||
clawhdf5-format = { path = "../clawhdf5-format", version = "2.4.0" }
|
clawhdf5-format = { path = "../clawhdf5-format", version = "2.5.0" }
|
||||||
clawhdf5 = { path = "../clawhdf5", version = "2.4.0" }
|
clawhdf5 = { path = "../clawhdf5", version = "2.5.0" }
|
||||||
rusqlite = { version = "0.31", features = ["bundled"] }
|
rusqlite = { version = "0.31", features = ["bundled"] }
|
||||||
clap = { version = "4", features = ["derive"] }
|
clap = { version = "4", features = ["derive"] }
|
||||||
half = { workspace = true }
|
half = { workspace = true }
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
[package]
|
[package]
|
||||||
name = "clawhdf5-napi"
|
name = "clawhdf5-napi"
|
||||||
version = "2.4.0"
|
version = "2.5.0"
|
||||||
edition = "2024"
|
edition = "2024"
|
||||||
description = "Node.js native addon (napi-rs) exposing clawhdf5-agent to TypeScript/JavaScript"
|
description = "Node.js native addon (napi-rs) exposing clawhdf5-agent to TypeScript/JavaScript"
|
||||||
license = "MIT"
|
license = "MIT"
|
||||||
@@ -10,7 +10,7 @@ repository = "https://git.redclaw.dev/quantumclaw/clawhdf5"
|
|||||||
crate-type = ["cdylib"]
|
crate-type = ["cdylib"]
|
||||||
|
|
||||||
[dependencies]
|
[dependencies]
|
||||||
clawhdf5-agent = { path = "../clawhdf5-agent", version = "2.4.0" }
|
clawhdf5-agent = { path = "../clawhdf5-agent", version = "2.5.0" }
|
||||||
napi = { version = "2", default-features = false, features = ["napi9"] }
|
napi = { version = "2", default-features = false, features = ["napi9"] }
|
||||||
napi-derive = "2"
|
napi-derive = "2"
|
||||||
|
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
[package]
|
[package]
|
||||||
name = "clawhdf5-netcdf4"
|
name = "clawhdf5-netcdf4"
|
||||||
version = "2.4.0"
|
version = "2.5.0"
|
||||||
edition = "2024"
|
edition = "2024"
|
||||||
description = "NetCDF-4 read support built on rustyhdf5 — pure Rust, no C dependencies"
|
description = "NetCDF-4 read support built on rustyhdf5 — pure Rust, no C dependencies"
|
||||||
license = "MIT"
|
license = "MIT"
|
||||||
@@ -10,8 +10,8 @@ keywords = ["netcdf", "netcdf4", "hdf5", "science", "climate"]
|
|||||||
categories = ["parser-implementations", "science"]
|
categories = ["parser-implementations", "science"]
|
||||||
|
|
||||||
[dependencies]
|
[dependencies]
|
||||||
clawhdf5 = { path = "../clawhdf5", version = "2.4.0" }
|
clawhdf5 = { path = "../clawhdf5", version = "2.5.0" }
|
||||||
clawhdf5-format = { path = "../clawhdf5-format", version = "2.4.0" }
|
clawhdf5-format = { path = "../clawhdf5-format", version = "2.5.0" }
|
||||||
|
|
||||||
[dev-dependencies]
|
[dev-dependencies]
|
||||||
tempfile = { workspace = true }
|
tempfile = { workspace = true }
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
[package]
|
[package]
|
||||||
name = "clawhdf5-py"
|
name = "clawhdf5-py"
|
||||||
version = "2.4.0"
|
version = "2.5.0"
|
||||||
edition = "2024"
|
edition = "2024"
|
||||||
description = "Python bindings for rustyhdf5 — a pure-Rust HDF5 library"
|
description = "Python bindings for rustyhdf5 — a pure-Rust HDF5 library"
|
||||||
license = "MIT"
|
license = "MIT"
|
||||||
@@ -14,8 +14,8 @@ name = "clawhdf5"
|
|||||||
crate-type = ["cdylib", "rlib"]
|
crate-type = ["cdylib", "rlib"]
|
||||||
|
|
||||||
[dependencies]
|
[dependencies]
|
||||||
clawhdf5_rs = { path = "../clawhdf5", version = "2.4.0", package = "clawhdf5" }
|
clawhdf5_rs = { path = "../clawhdf5", version = "2.5.0", package = "clawhdf5" }
|
||||||
clawhdf5-format = { path = "../clawhdf5-format", version = "2.4.0" }
|
clawhdf5-format = { path = "../clawhdf5-format", version = "2.5.0" }
|
||||||
pyo3 = "0.29"
|
pyo3 = "0.29"
|
||||||
numpy = "0.29"
|
numpy = "0.29"
|
||||||
|
|
||||||
|
|||||||
@@ -4,7 +4,7 @@ build-backend = "maturin"
|
|||||||
|
|
||||||
[project]
|
[project]
|
||||||
name = "rustyhdf5"
|
name = "rustyhdf5"
|
||||||
version = "2.4.0"
|
version = "2.5.0"
|
||||||
description = "Python bindings for rustyhdf5 — a pure-Rust HDF5 library"
|
description = "Python bindings for rustyhdf5 — a pure-Rust HDF5 library"
|
||||||
requires-python = ">=3.8"
|
requires-python = ">=3.8"
|
||||||
license = { text = "MIT" }
|
license = { text = "MIT" }
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
[package]
|
[package]
|
||||||
name = "clawhdf5"
|
name = "clawhdf5"
|
||||||
version = "2.4.0"
|
version = "2.5.0"
|
||||||
edition = "2024"
|
edition = "2024"
|
||||||
description = "Pure-Rust HDF5 reader/writer — no C dependencies"
|
description = "Pure-Rust HDF5 reader/writer — no C dependencies"
|
||||||
license = "MIT"
|
license = "MIT"
|
||||||
@@ -10,16 +10,16 @@ keywords = ["hdf5", "science", "data", "binary"]
|
|||||||
categories = ["parser-implementations", "science", "encoding"]
|
categories = ["parser-implementations", "science", "encoding"]
|
||||||
|
|
||||||
[dependencies]
|
[dependencies]
|
||||||
clawhdf5-format = { path = "../clawhdf5-format", version = "2.4.0" }
|
clawhdf5-format = { path = "../clawhdf5-format", version = "2.5.0" }
|
||||||
clawhdf5-io = { path = "../clawhdf5-io", version = "2.4.0" }
|
clawhdf5-io = { path = "../clawhdf5-io", version = "2.5.0" }
|
||||||
rayon = { version = "1", optional = true }
|
rayon = { version = "1", optional = true }
|
||||||
|
|
||||||
[dev-dependencies]
|
[dev-dependencies]
|
||||||
tempfile = { workspace = true }
|
tempfile = { workspace = true }
|
||||||
criterion = { workspace = true }
|
criterion = { workspace = true }
|
||||||
clawhdf5-io = { path = "../clawhdf5-io", version = "2.4.0", features = ["mmap"] }
|
clawhdf5-io = { path = "../clawhdf5-io", version = "2.5.0", features = ["mmap"] }
|
||||||
clawhdf5-format = { path = "../clawhdf5-format", version = "2.4.0", features = ["parallel", "fast-checksum"] }
|
clawhdf5-format = { path = "../clawhdf5-format", version = "2.5.0", features = ["parallel", "fast-checksum"] }
|
||||||
clawhdf5-filters = { path = "../clawhdf5-filters", version = "2.4.0" }
|
clawhdf5-filters = { path = "../clawhdf5-filters", version = "2.5.0" }
|
||||||
|
|
||||||
[[bench]]
|
[[bench]]
|
||||||
name = "mmap_bench"
|
name = "mmap_bench"
|
||||||
|
|||||||
@@ -381,8 +381,13 @@ impl<'f> Dataset<'f> {
|
|||||||
|
|
||||||
/// Read all data as `f64` values.
|
/// Read all data as `f64` values.
|
||||||
pub fn read_f64(&self) -> Result<Vec<f64>, Error> {
|
pub fn read_f64(&self) -> Result<Vec<f64>, Error> {
|
||||||
let raw = self.read_raw()?;
|
|
||||||
let dt = self.datatype()?;
|
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)?)
|
Ok(data_read::read_as_f64(&raw, &dt)?)
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -393,29 +398,49 @@ impl<'f> Dataset<'f> {
|
|||||||
///
|
///
|
||||||
/// Read all data as `f32` values.
|
/// Read all data as `f32` values.
|
||||||
pub fn read_f32(&self) -> Result<Vec<f32>, Error> {
|
pub fn read_f32(&self) -> Result<Vec<f32>, Error> {
|
||||||
let raw = self.read_raw()?;
|
|
||||||
let dt = self.datatype()?;
|
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)?)
|
Ok(data_read::read_as_f32(&raw, &dt)?)
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Read all data as `i32` values.
|
/// Read all data as `i32` values.
|
||||||
pub fn read_i32(&self) -> Result<Vec<i32>, Error> {
|
pub fn read_i32(&self) -> Result<Vec<i32>, Error> {
|
||||||
let raw = self.read_raw()?;
|
|
||||||
let dt = self.datatype()?;
|
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)?)
|
Ok(data_read::read_as_i32(&raw, &dt)?)
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Read all data as `i64` values.
|
/// Read all data as `i64` values.
|
||||||
pub fn read_i64(&self) -> Result<Vec<i64>, Error> {
|
pub fn read_i64(&self) -> Result<Vec<i64>, Error> {
|
||||||
let raw = self.read_raw()?;
|
|
||||||
let dt = self.datatype()?;
|
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)?)
|
Ok(data_read::read_as_i64(&raw, &dt)?)
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Read all data as `u64` values.
|
/// Read all data as `u64` values.
|
||||||
pub fn read_u64(&self) -> Result<Vec<u64>, Error> {
|
pub fn read_u64(&self) -> Result<Vec<u64>, Error> {
|
||||||
let raw = self.read_raw()?;
|
|
||||||
let dt = self.datatype()?;
|
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)?)
|
Ok(data_read::read_as_u64(&raw, &dt)?)
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -458,6 +483,7 @@ impl<'f> Dataset<'f> {
|
|||||||
|| (matches!(dl, DataLayout::Chunked { .. })
|
|| (matches!(dl, DataLayout::Chunked { .. })
|
||||||
&& !clawhdf5_format::fill_value::is_default(fill.as_deref()));
|
&& !clawhdf5_format::fill_value::is_default(fill.as_deref()));
|
||||||
if fill_matters {
|
if fill_matters {
|
||||||
|
clawhdf5_format::partial_read::validate(selection, &ds.dimensions)?;
|
||||||
let full = self.read_raw()?;
|
let full = self.read_raw()?;
|
||||||
return Ok(data_read::extract_selection_from_buffer(
|
return Ok(data_read::extract_selection_from_buffer(
|
||||||
&full,
|
&full,
|
||||||
|
|||||||
@@ -913,3 +913,128 @@ with h5py.File("{dst_str}", "r") as f:
|
|||||||
"[(1, 2.5), (3, 4.5)] ('a', 'b') [18446744073709551615, 0, 9223372036854775808] uint64"
|
"[(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]
|
||||||
|
);
|
||||||
|
}
|
||||||
+22
-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
|
## 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
|
**Summary:** HDF5 1.12+ `H5T_STD_REF` references use datatype message version 4
|
||||||
reference types 2–4 (object2 / region2 / attribute), which `Datatype::parse`
|
with reference types 2-4 (object / region / attribute), which `Datatype::parse`
|
||||||
rejects with `InvalidReferenceType`. h5py still writes the legacy v1
|
rejected with `InvalidReferenceType`. h5py still writes the legacy references,
|
||||||
object/region references, which read correctly, so no reproducing file has been
|
so no file had been available to test against.
|
||||||
generated yet; one written with the C API (`H5T_STD_REF`) is needed.
|
|
||||||
|
**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
|
## `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
|
## 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
|
**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
|
with `libver='latest'` indexes its chunks with a version-2 B-tree, and reading it
|
||||||
fails with `ChunkedReadError("unsupported chunked layout version=4,
|
failed with `unsupported chunked layout version=4, index_type=Some(5)`.
|
||||||
index_type=Some(5)")`. Single-chunk, implicit, fixed-array and
|
|
||||||
extensible-array indexes (and the v3 B-tree v1) are supported.
|
|
||||||
|
|
||||||
**Repro:** `f.create_dataset("d", shape=(5, 7), chunks=(2, 3), maxshape=(None, None))`
|
**Fix:** record types 10 (unfiltered) and 11 (filtered) are decoded — address,
|
||||||
with `h5py.File(..., libver='latest')`.
|
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
|
## External links and external raw data are not followed
|
||||||
|
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
{
|
{
|
||||||
"name": "@redclaw/clawhdf5",
|
"name": "@redclaw/clawhdf5",
|
||||||
"version": "2.4.0",
|
"version": "2.5.0",
|
||||||
"description": "Node.js bindings for clawhdf5 — HDF5-backed agent memory with hippocampal consolidation",
|
"description": "Node.js bindings for clawhdf5 — HDF5-backed agent memory with hippocampal consolidation",
|
||||||
"main": "index.js",
|
"main": "index.js",
|
||||||
"types": "index.d.ts",
|
"types": "index.d.ts",
|
||||||
|
|||||||
@@ -63,6 +63,13 @@ run_step "cargo clippy (format feature matrix)" cargo clippy \
|
|||||||
--features parallel,lz4,zstd,pcodec,fast-checksum \
|
--features parallel,lz4,zstd,pcodec,fast-checksum \
|
||||||
-- -D warnings
|
-- -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)
|
# 4. Tests (exclude clawhdf5-py)
|
||||||
run_step "cargo test" cargo test \
|
run_step "cargo test" cargo test \
|
||||||
--workspace \
|
--workspace \
|
||||||
@@ -72,6 +79,10 @@ run_step "cargo test (format feature matrix)" cargo test \
|
|||||||
-p clawhdf5-format \
|
-p clawhdf5-format \
|
||||||
--features parallel,lz4,zstd,pcodec,fast-checksum
|
--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
|
# 5. Python interop suites. The h5py writer tests are #[ignore]d so a plain
|
||||||
# `cargo test` stays hermetic; run them explicitly here.
|
# `cargo test` stays hermetic; run them explicitly here.
|
||||||
if python3 -c "import h5py" >/dev/null 2>&1 || [ "${CLAWHDF5_REQUIRE_INTEROP:-0}" = "1" ]; then
|
if python3 -c "import h5py" >/dev/null 2>&1 || [ "${CLAWHDF5_REQUIRE_INTEROP:-0}" = "1" ]; then
|
||||||
|
|||||||
Reference in New Issue
Block a user