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 = { path = "../clawhdf5", version = "2.0.0" }
|
||||||
clawhdf5-io = { path = "../clawhdf5-io", version = "2.0.0", features = ["mmap"] }
|
clawhdf5-io = { path = "../clawhdf5-io", version = "2.0.0", features = ["mmap"] }
|
||||||
clawhdf5-accel = { path = "../clawhdf5-accel", version = "2.0.0" }
|
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 }
|
clawhdf5-gpu = { path = "../clawhdf5-gpu", version = "2.0.0", optional = true, default-features = false }
|
||||||
serde = { version = "1", features = ["derive"] }
|
serde = { version = "1", features = ["derive"] }
|
||||||
byteorder = "1"
|
byteorder = "1"
|
||||||
@@ -44,9 +45,15 @@ name = "memory_bench"
|
|||||||
harness = false
|
harness = false
|
||||||
|
|
||||||
[features]
|
[features]
|
||||||
default = ["float16"]
|
default = ["float16", "hnsw"]
|
||||||
float16 = ["half"]
|
float16 = ["half"]
|
||||||
parallel = ["rayon"]
|
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 = []
|
agent = []
|
||||||
gpu = ["clawhdf5-gpu/gpu-wgpu"]
|
gpu = ["clawhdf5-gpu/gpu-wgpu"]
|
||||||
fast-math = ["matrixmultiply"]
|
fast-math = ["matrixmultiply"]
|
||||||
|
|||||||
@@ -60,6 +60,22 @@ pub fn hybrid_search(
|
|||||||
};
|
};
|
||||||
let kw_scores = bm25_index.search(query_text, vectors.len());
|
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].
|
// Normalize each set to [0, 1].
|
||||||
let vec_normalized = normalize_scores(&vec_scores);
|
let vec_normalized = normalize_scores(&vec_scores);
|
||||||
let kw_normalized = normalize_scores(&kw_scores);
|
let kw_normalized = normalize_scores(&kw_scores);
|
||||||
|
|||||||
@@ -60,7 +60,16 @@ pub fn cosine_similarity_prenorm(
|
|||||||
use std::path::{Path, PathBuf};
|
use std::path::{Path, PathBuf};
|
||||||
|
|
||||||
use cache::MemoryCache;
|
use cache::MemoryCache;
|
||||||
|
#[cfg(feature = "hnsw")]
|
||||||
|
use clawhdf5_ann::{DistanceMetric, HnswIndex};
|
||||||
use ephemeral::{EphemeralConfig, EphemeralStore};
|
use ephemeral::{EphemeralConfig, EphemeralStore};
|
||||||
|
|
||||||
|
/// HNSW construction parameters used for the agent's vector index. Cosine is the
|
||||||
|
/// agent's similarity metric, so the index is built with cosine distance.
|
||||||
|
#[cfg(feature = "hnsw")]
|
||||||
|
const HNSW_M: usize = 16;
|
||||||
|
#[cfg(feature = "hnsw")]
|
||||||
|
const HNSW_EF_CONSTRUCTION: usize = 64;
|
||||||
// EphemeralEntry and EphemeralStats are part of the crate public API via
|
// EphemeralEntry and EphemeralStats are part of the crate public API via
|
||||||
// the `ephemeral` module; they are not needed directly in lib.rs internals.
|
// the `ephemeral` module; they are not needed directly in lib.rs internals.
|
||||||
#[allow(unused_imports)]
|
#[allow(unused_imports)]
|
||||||
@@ -202,6 +211,22 @@ pub struct HDF5Memory {
|
|||||||
wal: Option<wal::WalFile>,
|
wal: Option<wal::WalFile>,
|
||||||
strategy: Option<Box<dyn MemoryStrategy>>,
|
strategy: Option<Box<dyn MemoryStrategy>>,
|
||||||
pub ephemeral: Option<EphemeralStore>,
|
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 {
|
impl std::fmt::Debug for HDF5Memory {
|
||||||
@@ -235,6 +260,12 @@ impl HDF5Memory {
|
|||||||
wal,
|
wal,
|
||||||
strategy: None,
|
strategy: None,
|
||||||
ephemeral: 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,
|
wal,
|
||||||
strategy: None,
|
strategy: None,
|
||||||
ephemeral: 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(())
|
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.
|
/// Get a reference to the config.
|
||||||
pub fn config(&self) -> &MemoryConfig {
|
pub fn config(&self) -> &MemoryConfig {
|
||||||
&self.config
|
&self.config
|
||||||
@@ -374,6 +515,8 @@ impl HDF5Memory {
|
|||||||
entry.timestamp,
|
entry.timestamp,
|
||||||
entry.session_id,
|
entry.session_id,
|
||||||
);
|
);
|
||||||
|
// In-place embedding change: the index node is stale, force rebuild.
|
||||||
|
self.hnsw_mark_dirty();
|
||||||
let needs_flush = self
|
let needs_flush = self
|
||||||
.wal
|
.wal
|
||||||
.as_ref()
|
.as_ref()
|
||||||
@@ -414,6 +557,7 @@ impl AgentMemory for HDF5Memory {
|
|||||||
entry.session_id,
|
entry.session_id,
|
||||||
entry.tags,
|
entry.tags,
|
||||||
);
|
);
|
||||||
|
self.hnsw_on_insert(idx);
|
||||||
let needs_flush = self
|
let needs_flush = self
|
||||||
.wal
|
.wal
|
||||||
.as_ref()
|
.as_ref()
|
||||||
@@ -440,6 +584,8 @@ impl AgentMemory for HDF5Memory {
|
|||||||
);
|
);
|
||||||
indices.push(idx);
|
indices.push(idx);
|
||||||
}
|
}
|
||||||
|
// Batch inserts rebuild the index once rather than node-by-node.
|
||||||
|
self.hnsw_mark_dirty();
|
||||||
self.flush()?;
|
self.flush()?;
|
||||||
Ok(indices)
|
Ok(indices)
|
||||||
}
|
}
|
||||||
@@ -450,6 +596,7 @@ impl AgentMemory for HDF5Memory {
|
|||||||
"entry {id} not found or already deleted"
|
"entry {id} not found or already deleted"
|
||||||
)));
|
)));
|
||||||
}
|
}
|
||||||
|
self.hnsw_on_delete(id);
|
||||||
self.flush()?;
|
self.flush()?;
|
||||||
|
|
||||||
// Auto-compact if threshold exceeded
|
// Auto-compact if threshold exceeded
|
||||||
@@ -465,6 +612,8 @@ impl AgentMemory for HDF5Memory {
|
|||||||
fn compact(&mut self) -> Result<usize> {
|
fn compact(&mut self) -> Result<usize> {
|
||||||
let (removed, _index_map) = self.cache.compact();
|
let (removed, _index_map) = self.cache.compact();
|
||||||
if removed > 0 {
|
if removed > 0 {
|
||||||
|
// Compaction renumbers cache indices; rebuild the index to match.
|
||||||
|
self.hnsw_mark_dirty();
|
||||||
self.flush()?;
|
self.flush()?;
|
||||||
}
|
}
|
||||||
Ok(removed)
|
Ok(removed)
|
||||||
|
|||||||
@@ -7,6 +7,76 @@ use crate::hybrid;
|
|||||||
use crate::{HDF5Memory, MemoryError, Result, SearchResult};
|
use crate::{HDF5Memory, MemoryError, Result, SearchResult};
|
||||||
|
|
||||||
impl HDF5Memory {
|
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.
|
/// Perform hybrid search combining cosine vector similarity and BM25 keyword search.
|
||||||
pub fn hybrid_search(
|
pub fn hybrid_search(
|
||||||
&mut self,
|
&mut self,
|
||||||
@@ -17,12 +87,9 @@ impl HDF5Memory {
|
|||||||
k: usize,
|
k: usize,
|
||||||
) -> Vec<SearchResult> {
|
) -> Vec<SearchResult> {
|
||||||
let bm25 = bm25::BM25Index::build(&self.cache.chunks, &self.cache.tombstones);
|
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_embedding,
|
||||||
query_text,
|
query_text,
|
||||||
&self.cache.embeddings,
|
|
||||||
&self.cache.chunks,
|
|
||||||
&self.cache.tombstones,
|
|
||||||
&bm25,
|
&bm25,
|
||||||
vector_weight,
|
vector_weight,
|
||||||
keyword_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.
|
/// HNSW (Hierarchical Navigable Small World) approximate nearest neighbor index.
|
||||||
///
|
///
|
||||||
/// Supports building an index from vectors, searching for nearest neighbors,
|
/// Supports building an index from vectors, incremental insertion and soft
|
||||||
/// and serializing/deserializing to HDF5 format.
|
/// deletion, searching for nearest neighbors, and serializing/deserializing to
|
||||||
|
/// HDF5 format.
|
||||||
#[derive(Debug, Clone)]
|
#[derive(Debug, Clone)]
|
||||||
pub struct HnswIndex {
|
pub struct HnswIndex {
|
||||||
/// All vectors in the index.
|
/// All vectors in the index.
|
||||||
vectors: Vec<Vec<f32>>,
|
vectors: Vec<Vec<f32>>,
|
||||||
/// Adjacency lists per layer. `graph[layer][node]` = list of neighbor IDs.
|
/// Adjacency lists per layer. `graph[layer][node]` = list of neighbor IDs.
|
||||||
graph: Vec<Vec<Vec<usize>>>,
|
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 node ID.
|
||||||
entry_point: usize,
|
entry_point: usize,
|
||||||
/// Maximum number of connections per node (per layer).
|
/// Maximum number of connections per node (per layer).
|
||||||
@@ -291,6 +305,7 @@ impl HnswIndex {
|
|||||||
Self {
|
Self {
|
||||||
vectors: vectors.to_vec(),
|
vectors: vectors.to_vec(),
|
||||||
graph,
|
graph,
|
||||||
|
deleted: vec![false; n],
|
||||||
entry_point,
|
entry_point,
|
||||||
m,
|
m,
|
||||||
m_max0,
|
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.
|
/// Search the index for the `k` nearest neighbors to the query vector.
|
||||||
///
|
///
|
||||||
/// # Parameters
|
/// # Parameters
|
||||||
@@ -328,11 +500,13 @@ 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
|
// 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);
|
let candidates = search_layer(&self.vectors, &self.graph[0], query, ep, ef, self.metric);
|
||||||
|
|
||||||
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()
|
||||||
@@ -388,11 +562,16 @@ impl HnswIndex {
|
|||||||
.set_attr("layer", AttrValue::I64(layer_idx as i64));
|
.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
|
// Store config as attributes on a small dataset
|
||||||
let node_levels_i32: Vec<i32> = self.node_levels.iter().map(|&l| l as i32).collect();
|
let node_levels_i32: Vec<i32> = self.node_levels.iter().map(|&l| l as i32).collect();
|
||||||
group
|
group
|
||||||
.create_dataset("config")
|
.create_dataset("config")
|
||||||
.with_i32_data(&node_levels_i32)
|
.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("m", AttrValue::I64(self.m as i64))
|
||||||
.set_attr(
|
.set_attr(
|
||||||
"ef_construction",
|
"ef_construction",
|
||||||
@@ -426,6 +605,15 @@ impl HnswIndex {
|
|||||||
let config_dt = read_dataset_datatype(data, &sb, "ann/config")?;
|
let config_dt = read_dataset_datatype(data, &sb, "ann/config")?;
|
||||||
let node_levels_i32 = read_as_i32(&config_raw, &config_dt)?;
|
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 m = get_attr_i64(&config_attrs, "m")? as usize;
|
||||||
let ef_construction = get_attr_i64(&config_attrs, "ef_construction")? 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;
|
let entry_point = get_attr_i64(&config_attrs, "entry_point")? as usize;
|
||||||
@@ -488,9 +676,22 @@ impl HnswIndex {
|
|||||||
graph.push(layer_graph);
|
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 {
|
Ok(Self {
|
||||||
vectors,
|
vectors,
|
||||||
graph,
|
graph,
|
||||||
|
deleted,
|
||||||
entry_point,
|
entry_point,
|
||||||
m,
|
m,
|
||||||
m_max0: m * 2,
|
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> {
|
fn get_attr_string(attrs: &[(String, AttrValue)], name: &str) -> Result<String, FormatError> {
|
||||||
for (n, v) in attrs {
|
for (n, v) in attrs {
|
||||||
if n == name {
|
if n == name {
|
||||||
@@ -1089,4 +1300,139 @@ mod tests {
|
|||||||
let d = compute_distance(&a, &b, DistanceMetric::Cosine);
|
let d = compute_distance(&a, &b, DistanceMetric::Cosine);
|
||||||
assert!((d - 1.0).abs() < 1e-6); // zero vector -> distance 1
|
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]
|
[dependencies]
|
||||||
clawhdf5_rs = { path = "../clawhdf5", version = "2.0.0", package = "clawhdf5" }
|
clawhdf5_rs = { path = "../clawhdf5", version = "2.0.0", package = "clawhdf5" }
|
||||||
clawhdf5-format = { path = "../clawhdf5-format", version = "2.0.0" }
|
clawhdf5-format = { path = "../clawhdf5-format", version = "2.0.0" }
|
||||||
pyo3 = "0.23"
|
pyo3 = "0.28"
|
||||||
numpy = "0.23"
|
numpy = "0.28"
|
||||||
|
|
||||||
[features]
|
[features]
|
||||||
extension-module = ["pyo3/extension-module"]
|
extension-module = ["pyo3/extension-module"]
|
||||||
|
|||||||
@@ -44,7 +44,7 @@ impl PyAttrs {
|
|||||||
|
|
||||||
#[pymethods]
|
#[pymethods]
|
||||||
impl PyAttrs {
|
impl PyAttrs {
|
||||||
fn __getitem__(&self, py: Python<'_>, key: &str) -> PyResult<PyObject> {
|
fn __getitem__(&self, py: Python<'_>, key: &str) -> PyResult<Py<PyAny>> {
|
||||||
match &self.inner {
|
match &self.inner {
|
||||||
AttrsInner::Read(map) => match map.get(key) {
|
AttrsInner::Read(map) => match map.get(key) {
|
||||||
Some(val) => Ok(attr_value_to_py(py, val)),
|
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 keys = self.keys(py)?;
|
||||||
let iter = keys.call_method0(py, "__iter__")?;
|
let iter = keys.call_method0(py, "__iter__")?;
|
||||||
Ok(iter)
|
Ok(iter)
|
||||||
@@ -112,7 +112,7 @@ impl PyAttrs {
|
|||||||
}
|
}
|
||||||
|
|
||||||
/// Return attribute names as a list.
|
/// 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 {
|
let names: Vec<String> = match &self.inner {
|
||||||
AttrsInner::Read(map) => map.keys().cloned().collect(),
|
AttrsInner::Read(map) => map.keys().cloned().collect(),
|
||||||
AttrsInner::Write(store) => store
|
AttrsInner::Write(store) => store
|
||||||
@@ -127,8 +127,8 @@ impl PyAttrs {
|
|||||||
}
|
}
|
||||||
|
|
||||||
/// Return attribute values as a list.
|
/// Return attribute values as a list.
|
||||||
fn values(&self, py: Python<'_>) -> PyResult<PyObject> {
|
fn values(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
|
||||||
let vals: Vec<PyObject> = match &self.inner {
|
let vals: Vec<Py<PyAny>> = match &self.inner {
|
||||||
AttrsInner::Read(map) => map.values().map(|v| attr_value_to_py(py, v)).collect(),
|
AttrsInner::Read(map) => map.values().map(|v| attr_value_to_py(py, v)).collect(),
|
||||||
AttrsInner::Write(store) => store
|
AttrsInner::Write(store) => store
|
||||||
.lock()
|
.lock()
|
||||||
@@ -145,8 +145,8 @@ impl PyAttrs {
|
|||||||
}
|
}
|
||||||
|
|
||||||
/// Return attribute (key, value) pairs as a list of tuples.
|
/// Return attribute (key, value) pairs as a list of tuples.
|
||||||
fn items(&self, py: Python<'_>) -> PyResult<PyObject> {
|
fn items(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
|
||||||
let pairs: Vec<(String, PyObject)> = match &self.inner {
|
let pairs: Vec<(String, Py<PyAny>)> = match &self.inner {
|
||||||
AttrsInner::Read(map) => map
|
AttrsInner::Read(map) => map
|
||||||
.iter()
|
.iter()
|
||||||
.map(|(k, v)| (k.clone(), attr_value_to_py(py, v)))
|
.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 {
|
impl PyDataset {
|
||||||
/// The shape of the dataset as a tuple.
|
/// The shape of the dataset as a tuple.
|
||||||
#[getter]
|
#[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))?;
|
let tuple = pyo3::types::PyTuple::new(py, self.cached_shape.iter().map(|&d| d as usize))?;
|
||||||
Ok(tuple.into_any().unbind())
|
Ok(tuple.into_any().unbind())
|
||||||
}
|
}
|
||||||
@@ -88,7 +88,7 @@ impl PyDataset {
|
|||||||
///
|
///
|
||||||
/// The full dataset is always read from the underlying file; the index
|
/// The full dataset is always read from the underlying file; the index
|
||||||
/// is then applied on the resulting numpy array.
|
/// 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 arr = self.read_as_numpy(py)?;
|
||||||
let indexed = arr.get_item(key)?;
|
let indexed = arr.get_item(key)?;
|
||||||
Ok(indexed.unbind())
|
Ok(indexed.unbind())
|
||||||
@@ -112,7 +112,7 @@ impl PyDataset {
|
|||||||
/// Read the full dataset and return it as a numpy array (or list for strings).
|
/// 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
|
/// 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
|
/// during the potentially expensive operation. The numpy array
|
||||||
/// construction still happens with the GIL held.
|
/// construction still happens with the GIL held.
|
||||||
fn read_as_numpy<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
|
fn read_as_numpy<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
|
||||||
@@ -123,7 +123,7 @@ impl PyDataset {
|
|||||||
match &self.cached_dtype {
|
match &self.cached_dtype {
|
||||||
DType::F64 => {
|
DType::F64 => {
|
||||||
let data = py
|
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)?;
|
.map_err(to_py_err)?;
|
||||||
let nd = ArrayD::from_shape_vec(IxDyn(&shape), data)
|
let nd = ArrayD::from_shape_vec(IxDyn(&shape), data)
|
||||||
.map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))?;
|
.map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))?;
|
||||||
@@ -132,7 +132,7 @@ impl PyDataset {
|
|||||||
}
|
}
|
||||||
DType::F32 => {
|
DType::F32 => {
|
||||||
let data = py
|
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)?;
|
.map_err(to_py_err)?;
|
||||||
let nd = ArrayD::from_shape_vec(IxDyn(&shape), data)
|
let nd = ArrayD::from_shape_vec(IxDyn(&shape), data)
|
||||||
.map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))?;
|
.map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))?;
|
||||||
@@ -141,7 +141,7 @@ impl PyDataset {
|
|||||||
}
|
}
|
||||||
DType::I32 => {
|
DType::I32 => {
|
||||||
let data = py
|
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)?;
|
.map_err(to_py_err)?;
|
||||||
let nd = ArrayD::from_shape_vec(IxDyn(&shape), data)
|
let nd = ArrayD::from_shape_vec(IxDyn(&shape), data)
|
||||||
.map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))?;
|
.map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))?;
|
||||||
@@ -150,7 +150,7 @@ impl PyDataset {
|
|||||||
}
|
}
|
||||||
DType::I64 => {
|
DType::I64 => {
|
||||||
let data = py
|
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)?;
|
.map_err(to_py_err)?;
|
||||||
let nd = ArrayD::from_shape_vec(IxDyn(&shape), data)
|
let nd = ArrayD::from_shape_vec(IxDyn(&shape), data)
|
||||||
.map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))?;
|
.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
|
// Try zero-copy first (contiguous layout), fall back to
|
||||||
// read_u64 + cast for chunked/compact datasets.
|
// read_u64 + cast for chunked/compact datasets.
|
||||||
let data: Vec<u8> = py
|
let data: Vec<u8> = py
|
||||||
.allow_threads(|| {
|
.detach(|| {
|
||||||
let ds = file.dataset(path)?;
|
let ds = file.dataset(path)?;
|
||||||
match ds.read_u8_zerocopy() {
|
match ds.read_u8_zerocopy() {
|
||||||
Ok(slice) => Ok(slice.to_vec()),
|
Ok(slice) => Ok(slice.to_vec()),
|
||||||
@@ -179,7 +179,7 @@ impl PyDataset {
|
|||||||
}
|
}
|
||||||
DType::U64 => {
|
DType::U64 => {
|
||||||
let data = py
|
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)?;
|
.map_err(to_py_err)?;
|
||||||
let nd = ArrayD::from_shape_vec(IxDyn(&shape), data)
|
let nd = ArrayD::from_shape_vec(IxDyn(&shape), data)
|
||||||
.map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))?;
|
.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
|
// String reads need the GIL for PyList construction, but we
|
||||||
// release it during the Rust I/O portion.
|
// release it during the Rust I/O portion.
|
||||||
let data = py
|
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)?;
|
.map_err(to_py_err)?;
|
||||||
let list = PyList::new(py, &data)?;
|
let list = PyList::new(py, &data)?;
|
||||||
Ok(list.into_any())
|
Ok(list.into_any())
|
||||||
|
|||||||
@@ -102,7 +102,7 @@ impl PyFile {
|
|||||||
}
|
}
|
||||||
|
|
||||||
/// Get a child object (dataset or group) by path.
|
/// 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()?;
|
let file = self.read_file()?;
|
||||||
// Try dataset first
|
// Try dataset first
|
||||||
match file.dataset(key) {
|
match file.dataset(key) {
|
||||||
@@ -130,7 +130,7 @@ impl PyFile {
|
|||||||
}
|
}
|
||||||
|
|
||||||
/// List the names of all children in the root group.
|
/// 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 file = self.read_file()?;
|
||||||
let root = file.root();
|
let root = file.root();
|
||||||
let mut names = root.datasets().map_err(to_py_err)?;
|
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.
|
/// 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 state = self.write_state_mut()?;
|
||||||
let group_state = Arc::new(Mutex::new(WriteGroupState {
|
let group_state = Arc::new(Mutex::new(WriteGroupState {
|
||||||
name: name.to_string(),
|
name: name.to_string(),
|
||||||
|
|||||||
@@ -57,7 +57,7 @@ impl PyGroup {
|
|||||||
#[pymethods]
|
#[pymethods]
|
||||||
impl PyGroup {
|
impl PyGroup {
|
||||||
/// Get a child object (dataset or subgroup) by name or path.
|
/// 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 {
|
match &self.inner {
|
||||||
GroupInner::Read { file, path } => {
|
GroupInner::Read { file, path } => {
|
||||||
let full_path = if path.is_empty() {
|
let full_path = if path.is_empty() {
|
||||||
@@ -96,7 +96,7 @@ impl PyGroup {
|
|||||||
}
|
}
|
||||||
|
|
||||||
/// List the names of all children (datasets and subgroups).
|
/// 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 {
|
match &self.inner {
|
||||||
GroupInner::Read { file, path } => {
|
GroupInner::Read { file, path } => {
|
||||||
let group = if path.is_empty() {
|
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.
|
/// 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 {
|
match val {
|
||||||
clawhdf5_rs::AttrValue::F64(v) => v.into_pyobject(py).unwrap().into_any().unbind(),
|
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(),
|
clawhdf5_rs::AttrValue::I64(v) => v.into_pyobject(py).unwrap().into_any().unbind(),
|
||||||
|
|||||||
Reference in New Issue
Block a user