feat: integrate HNSW into agent search, fix Python 3.14 build #1
@@ -14,6 +14,7 @@ clawhdf5-format = { path = "../clawhdf5-format", version = "2.0.0", features = [
|
||||
clawhdf5 = { path = "../clawhdf5", version = "2.0.0" }
|
||||
clawhdf5-io = { path = "../clawhdf5-io", version = "2.0.0", features = ["mmap"] }
|
||||
clawhdf5-accel = { path = "../clawhdf5-accel", version = "2.0.0" }
|
||||
clawhdf5-ann = { path = "../clawhdf5-ann", version = "2.0.0", optional = true }
|
||||
clawhdf5-gpu = { path = "../clawhdf5-gpu", version = "2.0.0", optional = true, default-features = false }
|
||||
serde = { version = "1", features = ["derive"] }
|
||||
byteorder = "1"
|
||||
@@ -44,9 +45,15 @@ name = "memory_bench"
|
||||
harness = false
|
||||
|
||||
[features]
|
||||
default = ["float16"]
|
||||
default = ["float16", "hnsw"]
|
||||
float16 = ["half"]
|
||||
parallel = ["rayon"]
|
||||
# HNSW approximate-nearest-neighbour acceleration for the vector stage of
|
||||
# hybrid_search. On by default; the index is rebuilt from the cache on demand
|
||||
# and stays self-consistent with the persisted memory store. Disable with
|
||||
# `--no-default-features` (plus re-enabling other defaults) to force the exact
|
||||
# linear cosine scan.
|
||||
hnsw = ["clawhdf5-ann"]
|
||||
agent = []
|
||||
gpu = ["clawhdf5-gpu/gpu-wgpu"]
|
||||
fast-math = ["matrixmultiply"]
|
||||
|
||||
@@ -60,6 +60,22 @@ pub fn hybrid_search(
|
||||
};
|
||||
let kw_scores = bm25_index.search(query_text, vectors.len());
|
||||
|
||||
merge_vector_keyword(vec_scores, kw_scores, vector_weight, keyword_weight, k)
|
||||
}
|
||||
|
||||
/// Merge pre-computed vector-similarity and keyword scores into a single ranking.
|
||||
///
|
||||
/// Both score sets are independently min-max normalized to [0, 1] and combined
|
||||
/// with the given weights. This is the shared core of [`hybrid_search`]; it is
|
||||
/// also used by the optional HNSW path, which supplies vector scores from an
|
||||
/// approximate-nearest-neighbour index instead of a full linear scan.
|
||||
pub fn merge_vector_keyword(
|
||||
vec_scores: Vec<(usize, f32)>,
|
||||
kw_scores: Vec<(usize, f32)>,
|
||||
vector_weight: f32,
|
||||
keyword_weight: f32,
|
||||
k: usize,
|
||||
) -> Vec<(usize, f32)> {
|
||||
// Normalize each set to [0, 1].
|
||||
let vec_normalized = normalize_scores(&vec_scores);
|
||||
let kw_normalized = normalize_scores(&kw_scores);
|
||||
|
||||
@@ -60,7 +60,16 @@ pub fn cosine_similarity_prenorm(
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
use cache::MemoryCache;
|
||||
#[cfg(feature = "hnsw")]
|
||||
use clawhdf5_ann::{DistanceMetric, HnswIndex};
|
||||
use ephemeral::{EphemeralConfig, EphemeralStore};
|
||||
|
||||
/// HNSW construction parameters used for the agent's vector index. Cosine is the
|
||||
/// agent's similarity metric, so the index is built with cosine distance.
|
||||
#[cfg(feature = "hnsw")]
|
||||
const HNSW_M: usize = 16;
|
||||
#[cfg(feature = "hnsw")]
|
||||
const HNSW_EF_CONSTRUCTION: usize = 64;
|
||||
// EphemeralEntry and EphemeralStats are part of the crate public API via
|
||||
// the `ephemeral` module; they are not needed directly in lib.rs internals.
|
||||
#[allow(unused_imports)]
|
||||
@@ -202,6 +211,22 @@ pub struct HDF5Memory {
|
||||
wal: Option<wal::WalFile>,
|
||||
strategy: Option<Box<dyn MemoryStrategy>>,
|
||||
pub ephemeral: Option<EphemeralStore>,
|
||||
/// Optional HNSW index accelerating the vector stage of `hybrid_search`.
|
||||
/// `None` when the store isn't indexable (no/zero-dim/mixed-dim embeddings);
|
||||
/// rebuilt from the cache whenever it drifts out of sync (see
|
||||
/// [`HDF5Memory::ensure_hnsw_fresh`]).
|
||||
#[cfg(feature = "hnsw")]
|
||||
hnsw: Option<HnswIndex>,
|
||||
/// Set when an in-place update/compaction may have invalidated `hnsw`,
|
||||
/// forcing a rebuild before the next search.
|
||||
#[cfg(feature = "hnsw")]
|
||||
hnsw_dirty: bool,
|
||||
/// Cache length the current `hnsw` value reflects. A mismatch with the live
|
||||
/// cache length triggers a rebuild — this both picks up unhooked cache
|
||||
/// pushes and avoids re-attempting to build an unindexable store every
|
||||
/// search.
|
||||
#[cfg(feature = "hnsw")]
|
||||
hnsw_synced_len: usize,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for HDF5Memory {
|
||||
@@ -235,6 +260,12 @@ impl HDF5Memory {
|
||||
wal,
|
||||
strategy: None,
|
||||
ephemeral: None,
|
||||
#[cfg(feature = "hnsw")]
|
||||
hnsw: None,
|
||||
#[cfg(feature = "hnsw")]
|
||||
hnsw_dirty: false,
|
||||
#[cfg(feature = "hnsw")]
|
||||
hnsw_synced_len: 0,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -262,6 +293,14 @@ impl HDF5Memory {
|
||||
wal,
|
||||
strategy: None,
|
||||
ephemeral: None,
|
||||
// Existing data is loaded from disk + WAL replay; mark the index
|
||||
// dirty so it is (re)built from the cache on the first search.
|
||||
#[cfg(feature = "hnsw")]
|
||||
hnsw: None,
|
||||
#[cfg(feature = "hnsw")]
|
||||
hnsw_dirty: true,
|
||||
#[cfg(feature = "hnsw")]
|
||||
hnsw_synced_len: 0,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -284,6 +323,108 @@ impl HDF5Memory {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// ---- HNSW index maintenance --------------------------------------------
|
||||
//
|
||||
// The index mirrors the cache: HNSW node id == cache index, kept aligned by
|
||||
// appending to both in lock-step and mirroring deletes. The incremental
|
||||
// hooks below are an optimization for the hot path; correctness is
|
||||
// guaranteed by `ensure_hnsw_fresh`, which rebuilds from the cache whenever
|
||||
// the index length drifts from the cache length (covering any mutation path
|
||||
// that doesn't call a hook, e.g. consolidation pushes).
|
||||
|
||||
/// Build an HNSW index over the entire cache, re-applying tombstones as
|
||||
/// soft-deletions so node ids stay aligned with cache indices.
|
||||
///
|
||||
/// Returns `None` for stores that aren't usefully indexable: no embeddings,
|
||||
/// a zero embedding dimension, or embeddings of mixed dimension (in which
|
||||
/// case `hybrid_search` keeps using the linear scan).
|
||||
#[cfg(feature = "hnsw")]
|
||||
fn build_hnsw_from_cache(&self) -> Option<HnswIndex> {
|
||||
let dim = self.cache.embedding_dim;
|
||||
if dim == 0 || self.cache.embeddings.is_empty() {
|
||||
return None;
|
||||
}
|
||||
if self.cache.embeddings.iter().any(|e| e.len() != dim) {
|
||||
return None;
|
||||
}
|
||||
let mut index = HnswIndex::build_with_metric(
|
||||
&self.cache.embeddings,
|
||||
HNSW_M,
|
||||
HNSW_EF_CONSTRUCTION,
|
||||
DistanceMetric::Cosine,
|
||||
);
|
||||
for (i, &t) in self.cache.tombstones.iter().enumerate() {
|
||||
if t != 0 {
|
||||
index.mark_deleted(i);
|
||||
}
|
||||
}
|
||||
Some(index)
|
||||
}
|
||||
|
||||
/// Ensure the HNSW index reflects the current cache. Rebuilds when marked
|
||||
/// dirty or when the cache length no longer matches what the index reflects.
|
||||
#[cfg(feature = "hnsw")]
|
||||
fn ensure_hnsw_fresh(&mut self) {
|
||||
let n = self.cache.embeddings.len();
|
||||
if self.hnsw_dirty || self.hnsw_synced_len != n {
|
||||
self.hnsw = self.build_hnsw_from_cache();
|
||||
self.hnsw_synced_len = n;
|
||||
self.hnsw_dirty = false;
|
||||
}
|
||||
}
|
||||
|
||||
/// Incrementally index the embedding just pushed at `idx`. Falls back to a
|
||||
/// rebuild (via the dirty flag) for the first vector, dimension mismatches,
|
||||
/// or id drift.
|
||||
#[cfg(feature = "hnsw")]
|
||||
fn hnsw_on_insert(&mut self, idx: usize) {
|
||||
if self.hnsw_dirty {
|
||||
return; // a rebuild is already pending; it will pick this up
|
||||
}
|
||||
let emb_len = self.cache.embeddings[idx].len();
|
||||
match self.hnsw.as_mut() {
|
||||
Some(index) if emb_len == index.dimension() => {
|
||||
let id = index.insert(self.cache.embeddings[idx].clone());
|
||||
if id == idx {
|
||||
self.hnsw_synced_len = self.cache.embeddings.len();
|
||||
} else {
|
||||
self.hnsw_dirty = true;
|
||||
}
|
||||
}
|
||||
// Dimension mismatch, first-ever vector, or no index yet: defer to a
|
||||
// rebuild, which decides indexability uniformly.
|
||||
_ => self.hnsw_dirty = true,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(not(feature = "hnsw"))]
|
||||
#[inline]
|
||||
fn hnsw_on_insert(&mut self, _idx: usize) {}
|
||||
|
||||
/// Mirror a cache deletion into the index.
|
||||
#[cfg(feature = "hnsw")]
|
||||
fn hnsw_on_delete(&mut self, id: usize) {
|
||||
if let Some(index) = self.hnsw.as_mut() {
|
||||
index.mark_deleted(id);
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(not(feature = "hnsw"))]
|
||||
#[inline]
|
||||
fn hnsw_on_delete(&mut self, _id: usize) {}
|
||||
|
||||
/// Mark the index for rebuild after a mutation that may have changed
|
||||
/// existing embeddings or renumbered ids (in-place update, compaction).
|
||||
#[cfg(feature = "hnsw")]
|
||||
#[inline]
|
||||
fn hnsw_mark_dirty(&mut self) {
|
||||
self.hnsw_dirty = true;
|
||||
}
|
||||
|
||||
#[cfg(not(feature = "hnsw"))]
|
||||
#[inline]
|
||||
fn hnsw_mark_dirty(&mut self) {}
|
||||
|
||||
/// Get a reference to the config.
|
||||
pub fn config(&self) -> &MemoryConfig {
|
||||
&self.config
|
||||
@@ -374,6 +515,8 @@ impl HDF5Memory {
|
||||
entry.timestamp,
|
||||
entry.session_id,
|
||||
);
|
||||
// In-place embedding change: the index node is stale, force rebuild.
|
||||
self.hnsw_mark_dirty();
|
||||
let needs_flush = self
|
||||
.wal
|
||||
.as_ref()
|
||||
@@ -414,6 +557,7 @@ impl AgentMemory for HDF5Memory {
|
||||
entry.session_id,
|
||||
entry.tags,
|
||||
);
|
||||
self.hnsw_on_insert(idx);
|
||||
let needs_flush = self
|
||||
.wal
|
||||
.as_ref()
|
||||
@@ -440,6 +584,8 @@ impl AgentMemory for HDF5Memory {
|
||||
);
|
||||
indices.push(idx);
|
||||
}
|
||||
// Batch inserts rebuild the index once rather than node-by-node.
|
||||
self.hnsw_mark_dirty();
|
||||
self.flush()?;
|
||||
Ok(indices)
|
||||
}
|
||||
@@ -450,6 +596,7 @@ impl AgentMemory for HDF5Memory {
|
||||
"entry {id} not found or already deleted"
|
||||
)));
|
||||
}
|
||||
self.hnsw_on_delete(id);
|
||||
self.flush()?;
|
||||
|
||||
// Auto-compact if threshold exceeded
|
||||
@@ -465,6 +612,8 @@ impl AgentMemory for HDF5Memory {
|
||||
fn compact(&mut self) -> Result<usize> {
|
||||
let (removed, _index_map) = self.cache.compact();
|
||||
if removed > 0 {
|
||||
// Compaction renumbers cache indices; rebuild the index to match.
|
||||
self.hnsw_mark_dirty();
|
||||
self.flush()?;
|
||||
}
|
||||
Ok(removed)
|
||||
|
||||
@@ -7,6 +7,76 @@ use crate::hybrid;
|
||||
use crate::{HDF5Memory, MemoryError, Result, SearchResult};
|
||||
|
||||
impl HDF5Memory {
|
||||
/// Vector + keyword scoring stage of [`HDF5Memory::hybrid_search`].
|
||||
///
|
||||
/// Without the `hnsw` feature this is a full linear cosine scan (the exact
|
||||
/// previous behaviour, also used as the correctness oracle in tests). With
|
||||
/// `hnsw` enabled and an index available, the vector candidates come from an
|
||||
/// approximate-nearest-neighbour search over an over-fetched pool, then merge
|
||||
/// with BM25 via the shared [`hybrid::merge_vector_keyword`].
|
||||
#[cfg(feature = "hnsw")]
|
||||
fn vector_keyword_search(
|
||||
&mut self,
|
||||
query_embedding: &[f32],
|
||||
query_text: &str,
|
||||
bm25: &bm25::BM25Index,
|
||||
vector_weight: f32,
|
||||
keyword_weight: f32,
|
||||
k: usize,
|
||||
) -> Vec<(usize, f32)> {
|
||||
self.ensure_hnsw_fresh();
|
||||
match self.hnsw.as_ref() {
|
||||
Some(index)
|
||||
if !index.is_empty() && index.dimension() == query_embedding.len() =>
|
||||
{
|
||||
// Over-fetch so the merge sees a useful vector pool; cosine
|
||||
// distance from the index converts back to similarity (1 - d).
|
||||
let pool = (k * 8).max(64);
|
||||
let vec_scores: Vec<(usize, f32)> = index
|
||||
.search(query_embedding, pool, pool)
|
||||
.into_iter()
|
||||
.map(|(id, dist)| (id, 1.0 - dist))
|
||||
.collect();
|
||||
let kw_scores = bm25.search(query_text, self.cache.len());
|
||||
hybrid::merge_vector_keyword(vec_scores, kw_scores, vector_weight, keyword_weight, k)
|
||||
}
|
||||
_ => hybrid::hybrid_search(
|
||||
query_embedding,
|
||||
query_text,
|
||||
&self.cache.embeddings,
|
||||
&self.cache.chunks,
|
||||
&self.cache.tombstones,
|
||||
bm25,
|
||||
vector_weight,
|
||||
keyword_weight,
|
||||
k,
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(not(feature = "hnsw"))]
|
||||
fn vector_keyword_search(
|
||||
&mut self,
|
||||
query_embedding: &[f32],
|
||||
query_text: &str,
|
||||
bm25: &bm25::BM25Index,
|
||||
vector_weight: f32,
|
||||
keyword_weight: f32,
|
||||
k: usize,
|
||||
) -> Vec<(usize, f32)> {
|
||||
hybrid::hybrid_search(
|
||||
query_embedding,
|
||||
query_text,
|
||||
&self.cache.embeddings,
|
||||
&self.cache.chunks,
|
||||
&self.cache.tombstones,
|
||||
bm25,
|
||||
vector_weight,
|
||||
keyword_weight,
|
||||
k,
|
||||
)
|
||||
}
|
||||
|
||||
/// Perform hybrid search combining cosine vector similarity and BM25 keyword search.
|
||||
pub fn hybrid_search(
|
||||
&mut self,
|
||||
@@ -17,12 +87,9 @@ impl HDF5Memory {
|
||||
k: usize,
|
||||
) -> Vec<SearchResult> {
|
||||
let bm25 = bm25::BM25Index::build(&self.cache.chunks, &self.cache.tombstones);
|
||||
let scored = hybrid::hybrid_search(
|
||||
let scored = self.vector_keyword_search(
|
||||
query_embedding,
|
||||
query_text,
|
||||
&self.cache.embeddings,
|
||||
&self.cache.chunks,
|
||||
&self.cache.tombstones,
|
||||
&bm25,
|
||||
vector_weight,
|
||||
keyword_weight,
|
||||
|
||||
@@ -0,0 +1,163 @@
|
||||
//! Integration tests for the optional HNSW-accelerated vector search path.
|
||||
//!
|
||||
//! These only run when the crate is built with `--features hnsw`. They drive the
|
||||
//! real `HDF5Memory` API (save / save_batch / delete / hybrid_search) and check
|
||||
//! the approximate results against a brute-force cosine oracle, plus confirm that
|
||||
//! deletions are honoured end-to-end.
|
||||
#![cfg(feature = "hnsw")]
|
||||
|
||||
use clawhdf5_agent::{AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry};
|
||||
use tempfile::TempDir;
|
||||
|
||||
/// Deterministic splitmix64 so tests are reproducible without an RNG crate.
|
||||
fn splitmix64(state: &mut u64) -> u64 {
|
||||
*state = state.wrapping_add(0x9e37_79b9_7f4a_7c15);
|
||||
let mut z = *state;
|
||||
z = (z ^ (z >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9);
|
||||
z = (z ^ (z >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb);
|
||||
z ^ (z >> 31)
|
||||
}
|
||||
|
||||
fn make_vector(seed: &mut u64, dim: usize) -> Vec<f32> {
|
||||
(0..dim)
|
||||
.map(|_| (splitmix64(seed) >> 40) as f32 / 16_777_216.0 - 0.5)
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn cosine(a: &[f32], b: &[f32]) -> f32 {
|
||||
let dot: f32 = a.iter().zip(b).map(|(x, y)| x * y).sum();
|
||||
let na: f32 = a.iter().map(|x| x * x).sum::<f32>().sqrt();
|
||||
let nb: f32 = b.iter().map(|x| x * x).sum::<f32>().sqrt();
|
||||
if na == 0.0 || nb == 0.0 {
|
||||
0.0
|
||||
} else {
|
||||
dot / (na * nb)
|
||||
}
|
||||
}
|
||||
|
||||
fn entry(chunk: &str, embedding: Vec<f32>, tags: &str) -> MemoryEntry {
|
||||
MemoryEntry {
|
||||
chunk: chunk.to_string(),
|
||||
embedding,
|
||||
source_channel: "test".to_string(),
|
||||
timestamp: 0.0,
|
||||
session_id: "s".to_string(),
|
||||
tags: tags.to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
fn new_memory(dir: &TempDir, dim: usize) -> HDF5Memory {
|
||||
let config = MemoryConfig::new(dir.path().join("mem.h5"), "agent", dim);
|
||||
HDF5Memory::create(config).unwrap()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn hnsw_matches_bruteforce_oracle() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let dim = 16;
|
||||
let n = 250;
|
||||
let mut mem = new_memory(&dir, dim);
|
||||
|
||||
let mut seed = 0xC0FF_EE12_3456_789A;
|
||||
let vectors: Vec<Vec<f32>> = (0..n).map(|_| make_vector(&mut seed, dim)).collect();
|
||||
for (i, v) in vectors.iter().enumerate() {
|
||||
mem.save(entry(&format!("chunk {i}"), v.clone(), &format!("k{i}")))
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
// Vector-only query: keyword weight 0 isolates the HNSW vector stage.
|
||||
let query = make_vector(&mut seed, dim);
|
||||
let k = 10;
|
||||
let results = mem.hybrid_search(&query, "", 1.0, 0.0, k);
|
||||
assert_eq!(results.len(), k, "should return k results");
|
||||
|
||||
// Brute-force cosine top-k oracle.
|
||||
let mut oracle: Vec<(usize, f32)> = vectors
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, v)| (i, cosine(&query, v)))
|
||||
.collect();
|
||||
oracle.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
|
||||
let oracle_ids: std::collections::HashSet<usize> =
|
||||
oracle.iter().take(k).map(|(i, _)| *i).collect();
|
||||
let hnsw_ids: std::collections::HashSet<usize> =
|
||||
results.iter().map(|r| r.index).collect();
|
||||
|
||||
let overlap = oracle_ids.intersection(&hnsw_ids).count();
|
||||
assert!(
|
||||
overlap >= 7,
|
||||
"HNSW recall too low vs brute force: {overlap}/{k} (hnsw={hnsw_ids:?})"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn deleted_entry_excluded_from_search() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let dim = 8;
|
||||
let mut mem = new_memory(&dir, dim);
|
||||
|
||||
let mut seed = 42;
|
||||
let vectors: Vec<Vec<f32>> = (0..60).map(|_| make_vector(&mut seed, dim)).collect();
|
||||
for (i, v) in vectors.iter().enumerate() {
|
||||
mem.save(entry(&format!("c{i}"), v.clone(), &format!("t{i}")))
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
// Query exactly equal to vector 5 — it must be the top hit.
|
||||
let query = vectors[5].clone();
|
||||
let top = mem.hybrid_search(&query, "", 1.0, 0.0, 1);
|
||||
assert_eq!(top[0].index, 5, "exact match should rank first");
|
||||
|
||||
mem.delete(5).unwrap();
|
||||
|
||||
let after = mem.hybrid_search(&query, "", 1.0, 0.0, 5);
|
||||
assert!(
|
||||
after.iter().all(|r| r.index != 5),
|
||||
"deleted entry must not appear in results"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn incremental_inserts_after_search_are_found() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let dim = 8;
|
||||
let mut mem = new_memory(&dir, dim);
|
||||
|
||||
let mut seed = 7;
|
||||
// First batch, then a search to force the index to build.
|
||||
for i in 0..40 {
|
||||
let v = make_vector(&mut seed, dim);
|
||||
mem.save(entry(&format!("a{i}"), v, &format!("a{i}"))).unwrap();
|
||||
}
|
||||
let _ = mem.hybrid_search(&make_vector(&mut seed, dim), "", 1.0, 0.0, 5);
|
||||
|
||||
// Now insert a distinctive vector incrementally and confirm we can find it.
|
||||
let needle = vec![10.0f32; dim];
|
||||
let idx = mem
|
||||
.save(entry("needle", needle.clone(), "needle"))
|
||||
.unwrap();
|
||||
let hits = mem.hybrid_search(&needle, "", 1.0, 0.0, 1);
|
||||
assert_eq!(hits[0].index, idx, "incrementally inserted vector must be found");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn save_batch_then_search_is_consistent() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let dim = 8;
|
||||
let mut mem = new_memory(&dir, dim);
|
||||
|
||||
let mut seed = 99;
|
||||
let vectors: Vec<Vec<f32>> = (0..50).map(|_| make_vector(&mut seed, dim)).collect();
|
||||
let entries: Vec<MemoryEntry> = vectors
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, v)| entry(&format!("b{i}"), v.clone(), &format!("b{i}")))
|
||||
.collect();
|
||||
mem.save_batch(entries).unwrap();
|
||||
|
||||
// Exact-match queries should resolve to themselves after a batch insert.
|
||||
for probe in [0usize, 17, 49] {
|
||||
let hits = mem.hybrid_search(&vectors[probe], "", 1.0, 0.0, 1);
|
||||
assert_eq!(hits[0].index, probe, "batch-inserted vector {probe} not found");
|
||||
}
|
||||
}
|
||||
@@ -153,16 +153,30 @@ impl Ord for FarCandidate {
|
||||
}
|
||||
}
|
||||
|
||||
/// On-disk format version for the serialized HNSW index.
|
||||
///
|
||||
/// - Version 1: original layout (`vectors`, `graph_layer_*`, `config`), no
|
||||
/// deletion support and no explicit version tag.
|
||||
/// - Version 2: adds a `format_version` attribute and a `deleted` bitset dataset
|
||||
/// so live insert/delete state survives a save/load round-trip.
|
||||
///
|
||||
/// Files written before this constant existed are treated as version 1 on load.
|
||||
pub const HNSW_FORMAT_VERSION: i64 = 2;
|
||||
|
||||
/// HNSW (Hierarchical Navigable Small World) approximate nearest neighbor index.
|
||||
///
|
||||
/// Supports building an index from vectors, searching for nearest neighbors,
|
||||
/// and serializing/deserializing to HDF5 format.
|
||||
/// Supports building an index from vectors, incremental insertion and soft
|
||||
/// deletion, searching for nearest neighbors, and serializing/deserializing to
|
||||
/// HDF5 format.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct HnswIndex {
|
||||
/// All vectors in the index.
|
||||
vectors: Vec<Vec<f32>>,
|
||||
/// Adjacency lists per layer. `graph[layer][node]` = list of neighbor IDs.
|
||||
graph: Vec<Vec<Vec<usize>>>,
|
||||
/// Soft-deletion flags, one per node. Deleted nodes remain in the graph for
|
||||
/// connectivity but are never returned from [`HnswIndex::search`].
|
||||
deleted: Vec<bool>,
|
||||
/// Entry point node ID.
|
||||
entry_point: usize,
|
||||
/// Maximum number of connections per node (per layer).
|
||||
@@ -291,6 +305,7 @@ impl HnswIndex {
|
||||
Self {
|
||||
vectors: vectors.to_vec(),
|
||||
graph,
|
||||
deleted: vec![false; n],
|
||||
entry_point,
|
||||
m,
|
||||
m_max0,
|
||||
@@ -300,6 +315,163 @@ impl HnswIndex {
|
||||
}
|
||||
}
|
||||
|
||||
/// Create an empty index with the given parameters. Used as the starting
|
||||
/// point for incremental [`HnswIndex::insert`] and as the result of
|
||||
/// [`HnswIndex::compact`] when every vector has been deleted.
|
||||
pub fn new(m: usize, ef_construction: usize, metric: DistanceMetric) -> Self {
|
||||
assert!(m >= 2, "m must be at least 2");
|
||||
Self {
|
||||
vectors: Vec::new(),
|
||||
graph: Vec::new(),
|
||||
deleted: Vec::new(),
|
||||
entry_point: 0,
|
||||
m,
|
||||
m_max0: m * 2,
|
||||
ef_construction,
|
||||
node_levels: Vec::new(),
|
||||
metric,
|
||||
}
|
||||
}
|
||||
|
||||
/// Insert a single vector into the index incrementally and return its id.
|
||||
///
|
||||
/// The id is the vector's position in insertion order and is stable for the
|
||||
/// life of the index (until [`HnswIndex::compact`] renumbers survivors).
|
||||
/// Inserting into an empty index seeds the entry point.
|
||||
///
|
||||
/// # Panics
|
||||
/// Panics if `vector`'s dimension does not match the existing vectors.
|
||||
pub fn insert(&mut self, vector: Vec<f32>) -> usize {
|
||||
let id = self.vectors.len();
|
||||
|
||||
// Seed an empty index.
|
||||
if id == 0 {
|
||||
let node_level = assign_level(0, self.m);
|
||||
self.vectors.push(vector);
|
||||
self.deleted.push(false);
|
||||
self.node_levels.push(node_level);
|
||||
self.graph = (0..=node_level).map(|_| vec![Vec::new(); 1]).collect();
|
||||
self.entry_point = 0;
|
||||
return 0;
|
||||
}
|
||||
|
||||
assert_eq!(
|
||||
vector.len(),
|
||||
self.vectors[0].len(),
|
||||
"insert dimension mismatch"
|
||||
);
|
||||
|
||||
let node_level = assign_level(id, self.m);
|
||||
self.vectors.push(vector);
|
||||
self.deleted.push(false);
|
||||
self.node_levels.push(node_level);
|
||||
|
||||
// Grow every existing layer with an empty adjacency slot for `id`, and
|
||||
// add any brand-new top layers this node introduces.
|
||||
for layer in self.graph.iter_mut() {
|
||||
layer.push(Vec::new());
|
||||
}
|
||||
while self.graph.len() <= node_level {
|
||||
self.graph.push(vec![Vec::new(); id + 1]);
|
||||
}
|
||||
|
||||
let ep_level = self.node_levels[self.entry_point];
|
||||
let mut ep = self.entry_point;
|
||||
|
||||
// Phase 1: greedy descent from the top down to node_level + 1.
|
||||
for layer in (node_level + 1..=ep_level).rev() {
|
||||
ep = greedy_closest(&self.vectors, &self.graph[layer], &self.vectors[id], ep, self.metric);
|
||||
}
|
||||
|
||||
// Phase 2: search and connect from min(node_level, ep_level) down to 0.
|
||||
let bottom = node_level.min(ep_level);
|
||||
for layer in (0..=bottom).rev() {
|
||||
let max_conn = if layer == 0 { self.m_max0 } else { self.m };
|
||||
let neighbors = search_layer(
|
||||
&self.vectors,
|
||||
&self.graph[layer],
|
||||
&self.vectors[id],
|
||||
ep,
|
||||
self.ef_construction,
|
||||
self.metric,
|
||||
);
|
||||
let selected: Vec<usize> = neighbors.iter().take(max_conn).map(|c| c.id).collect();
|
||||
self.graph[layer][id] = selected.clone();
|
||||
for &neighbor in &selected {
|
||||
self.graph[layer][neighbor].push(id);
|
||||
if self.graph[layer][neighbor].len() > max_conn {
|
||||
prune_connections(
|
||||
&self.vectors,
|
||||
&mut self.graph[layer][neighbor],
|
||||
neighbor,
|
||||
max_conn,
|
||||
self.metric,
|
||||
);
|
||||
}
|
||||
}
|
||||
if !selected.is_empty() {
|
||||
ep = selected[0];
|
||||
}
|
||||
}
|
||||
|
||||
// Promote the entry point if this node sits on a higher layer.
|
||||
if node_level > ep_level {
|
||||
self.entry_point = id;
|
||||
}
|
||||
|
||||
id
|
||||
}
|
||||
|
||||
/// Soft-delete the vector with the given id. The node stays in the graph so
|
||||
/// traversal/connectivity is preserved, but it will never be returned from
|
||||
/// [`HnswIndex::search`]. Idempotent; out-of-range ids are ignored.
|
||||
///
|
||||
/// Returns `true` if the id existed and was not already deleted.
|
||||
pub fn mark_deleted(&mut self, id: usize) -> bool {
|
||||
if id >= self.deleted.len() || self.deleted[id] {
|
||||
return false;
|
||||
}
|
||||
self.deleted[id] = true;
|
||||
true
|
||||
}
|
||||
|
||||
/// Returns whether the vector with the given id is soft-deleted.
|
||||
pub fn is_deleted(&self, id: usize) -> bool {
|
||||
self.deleted.get(id).copied().unwrap_or(false)
|
||||
}
|
||||
|
||||
/// Number of soft-deleted vectors still occupying the index.
|
||||
pub fn deleted_count(&self) -> usize {
|
||||
self.deleted.iter().filter(|&&d| d).count()
|
||||
}
|
||||
|
||||
/// Number of live (non-deleted) vectors.
|
||||
pub fn active_len(&self) -> usize {
|
||||
self.vectors.len() - self.deleted_count()
|
||||
}
|
||||
|
||||
/// Rebuild the index from scratch, dropping all soft-deleted vectors and
|
||||
/// renumbering the survivors into a compact `0..active_len` id space.
|
||||
///
|
||||
/// Returns a mapping from old id to new id (`None` for dropped vectors) so
|
||||
/// callers can rewrite any external id references they keep.
|
||||
pub fn compact(&mut self) -> Vec<Option<usize>> {
|
||||
let mut mapping = vec![None; self.vectors.len()];
|
||||
let mut surviving: Vec<Vec<f32>> = Vec::with_capacity(self.active_len());
|
||||
for (old, v) in self.vectors.iter().enumerate() {
|
||||
if !self.deleted[old] {
|
||||
mapping[old] = Some(surviving.len());
|
||||
surviving.push(v.clone());
|
||||
}
|
||||
}
|
||||
*self = if surviving.is_empty() {
|
||||
Self::new(self.m, self.ef_construction, self.metric)
|
||||
} else {
|
||||
Self::build_with_metric(&surviving, self.m, self.ef_construction, self.metric)
|
||||
};
|
||||
mapping
|
||||
}
|
||||
|
||||
/// Search the index for the `k` nearest neighbors to the query vector.
|
||||
///
|
||||
/// # Parameters
|
||||
@@ -328,11 +500,13 @@ impl HnswIndex {
|
||||
ep = greedy_closest(&self.vectors, &self.graph[layer], query, ep, self.metric);
|
||||
}
|
||||
|
||||
// Search layer 0 with ef candidates
|
||||
// Search layer 0 with ef candidates. Deleted nodes are still traversed
|
||||
// (they remain valid graph waypoints) but are filtered from the result.
|
||||
let candidates = search_layer(&self.vectors, &self.graph[0], query, ep, ef, self.metric);
|
||||
|
||||
candidates
|
||||
.into_iter()
|
||||
.filter(|c| !self.deleted[c.id])
|
||||
.take(k)
|
||||
.map(|c| (c.id, c.distance))
|
||||
.collect()
|
||||
@@ -388,11 +562,16 @@ impl HnswIndex {
|
||||
.set_attr("layer", AttrValue::I64(layer_idx as i64));
|
||||
}
|
||||
|
||||
// Soft-deletion bitset (format version 2+): 0 = live, 1 = deleted.
|
||||
let deleted_i32: Vec<i32> = self.deleted.iter().map(|&d| d as i32).collect();
|
||||
group.create_dataset("deleted").with_i32_data(&deleted_i32);
|
||||
|
||||
// Store config as attributes on a small dataset
|
||||
let node_levels_i32: Vec<i32> = self.node_levels.iter().map(|&l| l as i32).collect();
|
||||
group
|
||||
.create_dataset("config")
|
||||
.with_i32_data(&node_levels_i32)
|
||||
.set_attr("format_version", AttrValue::I64(HNSW_FORMAT_VERSION))
|
||||
.set_attr("m", AttrValue::I64(self.m as i64))
|
||||
.set_attr(
|
||||
"ef_construction",
|
||||
@@ -426,6 +605,15 @@ impl HnswIndex {
|
||||
let config_dt = read_dataset_datatype(data, &sb, "ann/config")?;
|
||||
let node_levels_i32 = read_as_i32(&config_raw, &config_dt)?;
|
||||
|
||||
// Files written before format version 2 have no version attribute; treat
|
||||
// them as version 1. Reject anything newer than we understand.
|
||||
let format_version = get_attr_i64_opt(&config_attrs, "format_version").unwrap_or(1);
|
||||
if format_version > HNSW_FORMAT_VERSION {
|
||||
return Err(FormatError::SerializationError(format!(
|
||||
"unsupported HNSW format version {format_version} (this build understands up to {HNSW_FORMAT_VERSION})"
|
||||
)));
|
||||
}
|
||||
|
||||
let m = get_attr_i64(&config_attrs, "m")? as usize;
|
||||
let ef_construction = get_attr_i64(&config_attrs, "ef_construction")? as usize;
|
||||
let entry_point = get_attr_i64(&config_attrs, "entry_point")? as usize;
|
||||
@@ -488,9 +676,22 @@ impl HnswIndex {
|
||||
graph.push(layer_graph);
|
||||
}
|
||||
|
||||
// Deleted bitset (version 2+). Older files default every node to live.
|
||||
let deleted = if format_version >= 2 {
|
||||
let deleted_raw = read_dataset_raw(data, &sb, "ann/deleted")?;
|
||||
let deleted_dt = read_dataset_datatype(data, &sb, "ann/deleted")?;
|
||||
let deleted_i32 = read_as_i32(&deleted_raw, &deleted_dt)?;
|
||||
let mut deleted: Vec<bool> = deleted_i32.iter().map(|&d| d != 0).collect();
|
||||
deleted.resize(n, false);
|
||||
deleted
|
||||
} else {
|
||||
vec![false; n]
|
||||
};
|
||||
|
||||
Ok(Self {
|
||||
vectors,
|
||||
graph,
|
||||
deleted,
|
||||
entry_point,
|
||||
m,
|
||||
m_max0: m * 2,
|
||||
@@ -809,6 +1010,16 @@ fn get_attr_i64(attrs: &[(String, AttrValue)], name: &str) -> Result<i64, Format
|
||||
)))
|
||||
}
|
||||
|
||||
/// Like [`get_attr_i64`] but returns `None` when the attribute is absent or not
|
||||
/// an integer, instead of erroring. Used for optional/back-compat attributes.
|
||||
fn get_attr_i64_opt(attrs: &[(String, AttrValue)], name: &str) -> Option<i64> {
|
||||
attrs.iter().find(|(n, _)| n == name).and_then(|(_, v)| match v {
|
||||
AttrValue::I64(val) => Some(*val),
|
||||
AttrValue::U64(val) => Some(*val as i64),
|
||||
_ => None,
|
||||
})
|
||||
}
|
||||
|
||||
fn get_attr_string(attrs: &[(String, AttrValue)], name: &str) -> Result<String, FormatError> {
|
||||
for (n, v) in attrs {
|
||||
if n == name {
|
||||
@@ -1089,4 +1300,139 @@ mod tests {
|
||||
let d = compute_distance(&a, &b, DistanceMetric::Cosine);
|
||||
assert!((d - 1.0).abs() < 1e-6); // zero vector -> distance 1
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn insert_into_empty_index() {
|
||||
let mut index = HnswIndex::new(4, 16, DistanceMetric::L2);
|
||||
assert!(index.is_empty());
|
||||
let id = index.insert(vec![1.0, 0.0, 0.0]);
|
||||
assert_eq!(id, 0);
|
||||
assert_eq!(index.len(), 1);
|
||||
let results = index.search(&[1.0, 0.0, 0.0], 1, 16);
|
||||
assert_eq!(results, vec![(0, 0.0)]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn incremental_insert_matches_batch_recall() {
|
||||
// Build one index incrementally and one in batch from the same vectors,
|
||||
// then confirm the incremental index has acceptable recall vs brute force.
|
||||
let vectors = make_random_vectors(120, 8, 2024);
|
||||
let mut incremental = HnswIndex::new(16, 64, DistanceMetric::L2);
|
||||
for v in &vectors {
|
||||
incremental.insert(v.clone());
|
||||
}
|
||||
assert_eq!(incremental.len(), vectors.len());
|
||||
|
||||
let query = &vectors[60];
|
||||
let k = 10;
|
||||
let results = incremental.search(query, k, 64);
|
||||
|
||||
let mut brute: Vec<(usize, f32)> = vectors
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, v)| (i, compute_distance(query, v, DistanceMetric::L2)))
|
||||
.collect();
|
||||
brute.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap());
|
||||
brute.truncate(k);
|
||||
|
||||
let hnsw_ids: HashSet<usize> = results.iter().map(|r| r.0).collect();
|
||||
let brute_ids: HashSet<usize> = brute.iter().map(|r| r.0).collect();
|
||||
let overlap = hnsw_ids.intersection(&brute_ids).count();
|
||||
assert!(
|
||||
overlap >= k * 8 / 10,
|
||||
"incremental HNSW recall too low: {overlap}/{k}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mark_deleted_excludes_from_search() {
|
||||
let vectors = vec![
|
||||
vec![1.0, 0.0],
|
||||
vec![0.9, 0.1],
|
||||
vec![0.0, 1.0],
|
||||
vec![0.1, 0.9],
|
||||
];
|
||||
let mut index = HnswIndex::build(&vectors, 4, 16);
|
||||
|
||||
// Exact match on vector 0 before deletion.
|
||||
let before = index.search(&[1.0, 0.0], 1, 16);
|
||||
assert_eq!(before[0].0, 0);
|
||||
|
||||
assert!(index.mark_deleted(0));
|
||||
assert!(index.is_deleted(0));
|
||||
assert!(!index.mark_deleted(0)); // idempotent
|
||||
assert_eq!(index.deleted_count(), 1);
|
||||
assert_eq!(index.active_len(), 3);
|
||||
|
||||
// Vector 0 must no longer be returned; nearest is now vector 1.
|
||||
let after = index.search(&[1.0, 0.0], 2, 16);
|
||||
assert!(after.iter().all(|(id, _)| *id != 0));
|
||||
assert_eq!(after[0].0, 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn compact_drops_deleted_and_renumbers() {
|
||||
let vectors = make_random_vectors(20, 4, 4242);
|
||||
let mut index = HnswIndex::build(&vectors, 8, 32);
|
||||
index.mark_deleted(3);
|
||||
index.mark_deleted(7);
|
||||
index.mark_deleted(11);
|
||||
|
||||
let mapping = index.compact();
|
||||
assert_eq!(mapping.len(), 20);
|
||||
assert_eq!(index.len(), 17);
|
||||
assert_eq!(index.deleted_count(), 0);
|
||||
// Deleted ids map to None; survivors map to a dense 0..17 range.
|
||||
assert!(mapping[3].is_none() && mapping[7].is_none() && mapping[11].is_none());
|
||||
let mut new_ids: Vec<usize> = mapping.iter().filter_map(|m| *m).collect();
|
||||
new_ids.sort_unstable();
|
||||
assert_eq!(new_ids, (0..17).collect::<Vec<_>>());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn compact_all_deleted_yields_empty_index() {
|
||||
let vectors = make_random_vectors(5, 3, 1);
|
||||
let mut index = HnswIndex::build(&vectors, 4, 16);
|
||||
for i in 0..5 {
|
||||
index.mark_deleted(i);
|
||||
}
|
||||
index.compact();
|
||||
assert!(index.is_empty());
|
||||
assert!(index.search(&[0.0, 0.0, 0.0], 3, 16).is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn versioned_roundtrip_preserves_deletions() {
|
||||
let vectors = make_random_vectors(30, 4, 8080);
|
||||
let mut index = HnswIndex::build(&vectors, 6, 24);
|
||||
index.mark_deleted(5);
|
||||
index.mark_deleted(12);
|
||||
|
||||
let bytes = index.to_hdf5_bytes().unwrap();
|
||||
let loaded = HnswIndex::load_from_hdf5(&bytes).unwrap();
|
||||
|
||||
assert_eq!(loaded.len(), index.len());
|
||||
assert!(loaded.is_deleted(5));
|
||||
assert!(loaded.is_deleted(12));
|
||||
assert_eq!(loaded.deleted_count(), 2);
|
||||
// A deleted vector is not returned even after a round-trip.
|
||||
let results = loaded.search(&vectors[5], 5, 24);
|
||||
assert!(results.iter().all(|(id, _)| *id != 5));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn insert_then_save_load_search() {
|
||||
let mut index = HnswIndex::new(8, 32, DistanceMetric::Cosine);
|
||||
let vectors = make_random_vectors(25, 5, 31337);
|
||||
for v in &vectors {
|
||||
index.insert(v.clone());
|
||||
}
|
||||
let bytes = index.to_hdf5_bytes().unwrap();
|
||||
let loaded = HnswIndex::load_from_hdf5(&bytes).unwrap();
|
||||
assert_eq!(loaded.len(), 25);
|
||||
assert_eq!(loaded.metric(), DistanceMetric::Cosine);
|
||||
let results = loaded.search(&vectors[0], 3, 32);
|
||||
assert_eq!(results.len(), 3);
|
||||
assert_eq!(results[0].0, 0);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -16,8 +16,8 @@ crate-type = ["cdylib", "rlib"]
|
||||
[dependencies]
|
||||
clawhdf5_rs = { path = "../clawhdf5", version = "2.0.0", package = "clawhdf5" }
|
||||
clawhdf5-format = { path = "../clawhdf5-format", version = "2.0.0" }
|
||||
pyo3 = "0.23"
|
||||
numpy = "0.23"
|
||||
pyo3 = "0.28"
|
||||
numpy = "0.28"
|
||||
|
||||
[features]
|
||||
extension-module = ["pyo3/extension-module"]
|
||||
|
||||
@@ -44,7 +44,7 @@ impl PyAttrs {
|
||||
|
||||
#[pymethods]
|
||||
impl PyAttrs {
|
||||
fn __getitem__(&self, py: Python<'_>, key: &str) -> PyResult<PyObject> {
|
||||
fn __getitem__(&self, py: Python<'_>, key: &str) -> PyResult<Py<PyAny>> {
|
||||
match &self.inner {
|
||||
AttrsInner::Read(map) => match map.get(key) {
|
||||
Some(val) => Ok(attr_value_to_py(py, val)),
|
||||
@@ -100,7 +100,7 @@ impl PyAttrs {
|
||||
}
|
||||
}
|
||||
|
||||
fn __iter__(&self, py: Python<'_>) -> PyResult<PyObject> {
|
||||
fn __iter__(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
|
||||
let keys = self.keys(py)?;
|
||||
let iter = keys.call_method0(py, "__iter__")?;
|
||||
Ok(iter)
|
||||
@@ -112,7 +112,7 @@ impl PyAttrs {
|
||||
}
|
||||
|
||||
/// Return attribute names as a list.
|
||||
fn keys(&self, py: Python<'_>) -> PyResult<PyObject> {
|
||||
fn keys(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
|
||||
let names: Vec<String> = match &self.inner {
|
||||
AttrsInner::Read(map) => map.keys().cloned().collect(),
|
||||
AttrsInner::Write(store) => store
|
||||
@@ -127,8 +127,8 @@ impl PyAttrs {
|
||||
}
|
||||
|
||||
/// Return attribute values as a list.
|
||||
fn values(&self, py: Python<'_>) -> PyResult<PyObject> {
|
||||
let vals: Vec<PyObject> = match &self.inner {
|
||||
fn values(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
|
||||
let vals: Vec<Py<PyAny>> = match &self.inner {
|
||||
AttrsInner::Read(map) => map.values().map(|v| attr_value_to_py(py, v)).collect(),
|
||||
AttrsInner::Write(store) => store
|
||||
.lock()
|
||||
@@ -145,8 +145,8 @@ impl PyAttrs {
|
||||
}
|
||||
|
||||
/// Return attribute (key, value) pairs as a list of tuples.
|
||||
fn items(&self, py: Python<'_>) -> PyResult<PyObject> {
|
||||
let pairs: Vec<(String, PyObject)> = match &self.inner {
|
||||
fn items(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
|
||||
let pairs: Vec<(String, Py<PyAny>)> = match &self.inner {
|
||||
AttrsInner::Read(map) => map
|
||||
.iter()
|
||||
.map(|(k, v)| (k.clone(), attr_value_to_py(py, v)))
|
||||
|
||||
@@ -65,7 +65,7 @@ fn dtype_to_numpy_str(dt: &DType) -> &'static str {
|
||||
impl PyDataset {
|
||||
/// The shape of the dataset as a tuple.
|
||||
#[getter]
|
||||
fn shape(&self, py: Python<'_>) -> PyResult<PyObject> {
|
||||
fn shape(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
|
||||
let tuple = pyo3::types::PyTuple::new(py, self.cached_shape.iter().map(|&d| d as usize))?;
|
||||
Ok(tuple.into_any().unbind())
|
||||
}
|
||||
@@ -88,7 +88,7 @@ impl PyDataset {
|
||||
///
|
||||
/// The full dataset is always read from the underlying file; the index
|
||||
/// is then applied on the resulting numpy array.
|
||||
fn __getitem__<'py>(&self, py: Python<'py>, key: &Bound<'py, PyAny>) -> PyResult<PyObject> {
|
||||
fn __getitem__<'py>(&self, py: Python<'py>, key: &Bound<'py, PyAny>) -> PyResult<Py<PyAny>> {
|
||||
let arr = self.read_as_numpy(py)?;
|
||||
let indexed = arr.get_item(key)?;
|
||||
Ok(indexed.unbind())
|
||||
@@ -112,7 +112,7 @@ impl PyDataset {
|
||||
/// Read the full dataset and return it as a numpy array (or list for strings).
|
||||
///
|
||||
/// For numeric types, the Rust I/O (file reading + decompression) is
|
||||
/// performed inside `py.allow_threads()` so that the GIL is released
|
||||
/// performed inside `py.detach()` so that the GIL is released
|
||||
/// during the potentially expensive operation. The numpy array
|
||||
/// construction still happens with the GIL held.
|
||||
fn read_as_numpy<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
|
||||
@@ -123,7 +123,7 @@ impl PyDataset {
|
||||
match &self.cached_dtype {
|
||||
DType::F64 => {
|
||||
let data = py
|
||||
.allow_threads(|| file.dataset(path).and_then(|ds| ds.read_f64()))
|
||||
.detach(|| file.dataset(path).and_then(|ds| ds.read_f64()))
|
||||
.map_err(to_py_err)?;
|
||||
let nd = ArrayD::from_shape_vec(IxDyn(&shape), data)
|
||||
.map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))?;
|
||||
@@ -132,7 +132,7 @@ impl PyDataset {
|
||||
}
|
||||
DType::F32 => {
|
||||
let data = py
|
||||
.allow_threads(|| file.dataset(path).and_then(|ds| ds.read_f32()))
|
||||
.detach(|| file.dataset(path).and_then(|ds| ds.read_f32()))
|
||||
.map_err(to_py_err)?;
|
||||
let nd = ArrayD::from_shape_vec(IxDyn(&shape), data)
|
||||
.map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))?;
|
||||
@@ -141,7 +141,7 @@ impl PyDataset {
|
||||
}
|
||||
DType::I32 => {
|
||||
let data = py
|
||||
.allow_threads(|| file.dataset(path).and_then(|ds| ds.read_i32()))
|
||||
.detach(|| file.dataset(path).and_then(|ds| ds.read_i32()))
|
||||
.map_err(to_py_err)?;
|
||||
let nd = ArrayD::from_shape_vec(IxDyn(&shape), data)
|
||||
.map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))?;
|
||||
@@ -150,7 +150,7 @@ impl PyDataset {
|
||||
}
|
||||
DType::I64 => {
|
||||
let data = py
|
||||
.allow_threads(|| file.dataset(path).and_then(|ds| ds.read_i64()))
|
||||
.detach(|| file.dataset(path).and_then(|ds| ds.read_i64()))
|
||||
.map_err(to_py_err)?;
|
||||
let nd = ArrayD::from_shape_vec(IxDyn(&shape), data)
|
||||
.map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))?;
|
||||
@@ -161,7 +161,7 @@ impl PyDataset {
|
||||
// Try zero-copy first (contiguous layout), fall back to
|
||||
// read_u64 + cast for chunked/compact datasets.
|
||||
let data: Vec<u8> = py
|
||||
.allow_threads(|| {
|
||||
.detach(|| {
|
||||
let ds = file.dataset(path)?;
|
||||
match ds.read_u8_zerocopy() {
|
||||
Ok(slice) => Ok(slice.to_vec()),
|
||||
@@ -179,7 +179,7 @@ impl PyDataset {
|
||||
}
|
||||
DType::U64 => {
|
||||
let data = py
|
||||
.allow_threads(|| file.dataset(path).and_then(|ds| ds.read_u64()))
|
||||
.detach(|| file.dataset(path).and_then(|ds| ds.read_u64()))
|
||||
.map_err(to_py_err)?;
|
||||
let nd = ArrayD::from_shape_vec(IxDyn(&shape), data)
|
||||
.map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))?;
|
||||
@@ -190,7 +190,7 @@ impl PyDataset {
|
||||
// String reads need the GIL for PyList construction, but we
|
||||
// release it during the Rust I/O portion.
|
||||
let data = py
|
||||
.allow_threads(|| file.dataset(path).and_then(|ds| ds.read_string()))
|
||||
.detach(|| file.dataset(path).and_then(|ds| ds.read_string()))
|
||||
.map_err(to_py_err)?;
|
||||
let list = PyList::new(py, &data)?;
|
||||
Ok(list.into_any())
|
||||
|
||||
@@ -102,7 +102,7 @@ impl PyFile {
|
||||
}
|
||||
|
||||
/// Get a child object (dataset or group) by path.
|
||||
fn __getitem__(&self, py: Python<'_>, key: &str) -> PyResult<PyObject> {
|
||||
fn __getitem__(&self, py: Python<'_>, key: &str) -> PyResult<Py<PyAny>> {
|
||||
let file = self.read_file()?;
|
||||
// Try dataset first
|
||||
match file.dataset(key) {
|
||||
@@ -130,7 +130,7 @@ impl PyFile {
|
||||
}
|
||||
|
||||
/// List the names of all children in the root group.
|
||||
fn keys(&self, py: Python<'_>) -> PyResult<PyObject> {
|
||||
fn keys(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
|
||||
let file = self.read_file()?;
|
||||
let root = file.root();
|
||||
let mut names = root.datasets().map_err(to_py_err)?;
|
||||
@@ -175,7 +175,7 @@ impl PyFile {
|
||||
}
|
||||
|
||||
/// Create a group (write mode only). Returns a `Group` handle.
|
||||
fn create_group(&mut self, py: Python<'_>, name: &str) -> PyResult<PyObject> {
|
||||
fn create_group(&mut self, py: Python<'_>, name: &str) -> PyResult<Py<PyAny>> {
|
||||
let state = self.write_state_mut()?;
|
||||
let group_state = Arc::new(Mutex::new(WriteGroupState {
|
||||
name: name.to_string(),
|
||||
|
||||
@@ -57,7 +57,7 @@ impl PyGroup {
|
||||
#[pymethods]
|
||||
impl PyGroup {
|
||||
/// Get a child object (dataset or subgroup) by name or path.
|
||||
fn __getitem__(&self, py: Python<'_>, key: &str) -> PyResult<PyObject> {
|
||||
fn __getitem__(&self, py: Python<'_>, key: &str) -> PyResult<Py<PyAny>> {
|
||||
match &self.inner {
|
||||
GroupInner::Read { file, path } => {
|
||||
let full_path = if path.is_empty() {
|
||||
@@ -96,7 +96,7 @@ impl PyGroup {
|
||||
}
|
||||
|
||||
/// List the names of all children (datasets and subgroups).
|
||||
fn keys(&self, py: Python<'_>) -> PyResult<PyObject> {
|
||||
fn keys(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
|
||||
match &self.inner {
|
||||
GroupInner::Read { file, path } => {
|
||||
let group = if path.is_empty() {
|
||||
|
||||
@@ -114,7 +114,7 @@ pub(crate) fn py_to_attr_value(val: &Bound<'_, PyAny>) -> PyResult<OwnedAttrValu
|
||||
}
|
||||
|
||||
/// Convert an `AttrValue` (from the Rust lib) to a Python object.
|
||||
pub(crate) fn attr_value_to_py(py: Python<'_>, val: &clawhdf5_rs::AttrValue) -> PyObject {
|
||||
pub(crate) fn attr_value_to_py(py: Python<'_>, val: &clawhdf5_rs::AttrValue) -> Py<PyAny> {
|
||||
match val {
|
||||
clawhdf5_rs::AttrValue::F64(v) => v.into_pyobject(py).unwrap().into_any().unbind(),
|
||||
clawhdf5_rs::AttrValue::I64(v) => v.into_pyobject(py).unwrap().into_any().unbind(),
|
||||
|
||||
Reference in New Issue
Block a user