feat: integrate HNSW into agent search, fix Python 3.14 build #1

Merged
osobh merged 1 commits from feat/hnsw-agent-integration-py314 into main 2026-06-03 00:46:09 +00:00
12 changed files with 781 additions and 33 deletions
+8 -1
View File
@@ -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"]
+16
View File
@@ -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);
+149
View File
@@ -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)
+71 -4
View File
@@ -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");
}
}
+349 -3
View File
@@ -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);
}
} }
+2 -2
View File
@@ -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"]
+7 -7
View File
@@ -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)))
+10 -10
View File
@@ -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())
+3 -3
View File
@@ -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(),
+2 -2
View File
@@ -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() {
+1 -1
View File
@@ -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(),