nodedb_vector/vamana/
node_fetcher.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! `NodeFetcher` trait — abstracts in-memory vs io_uring-backed FP32 vector
4//! retrieval for Vamana beam-search and rerank.
5//!
6//! The trait is intentionally *not* `Send + Sync` so that `IoUringNodeFetcher`
7//! (Data Plane, `!Send`) can implement it without wrapping in `Arc<Mutex<…>>`.
8//! The `InMemoryFetcher` impl is `Send + Sync` (plain memory slice), but the
9//! trait bound is left absent to keep the interface Data-Plane-compatible.
10//!
11//! ## Pre-fetch contract
12//!
13//! `prefetch_batch` submits I/O for a set of node indices *before* compute
14//! touches them.  On the in-memory path this is a no-op.  On the io_uring
15//! path it issues `IORING_OP_READ` SQEs in parallel; `fetch_fp32` then
16//! collects completions.  This satisfies the TPC Page Fault Hazard rule:
17//! pages are pre-warmed asynchronously so the reactor thread never blocks
18//! on a major fault.
19
20/// Abstraction over FP32 vector storage for a Vamana index partition.
21///
22/// Two impls exist:
23/// - [`InMemoryFetcher`] — plain `Vec<Vec<f32>>`, available on all targets.
24/// - `IoUringNodeFetcher` — Data Plane only, lives in `nodedb::data`.
25pub trait NodeFetcher {
26    /// Vector dimensionality.
27    fn dim(&self) -> usize;
28
29    /// Submit pre-fetch hints for `indices` so their data is ready before
30    /// `fetch_fp32` is called.
31    ///
32    /// On the in-memory path this is a no-op.  On the io_uring path this
33    /// issues `IORING_OP_READ` SQEs without blocking — completion is polled
34    /// lazily inside `fetch_fp32`.
35    fn prefetch_batch(&mut self, indices: &[u32]);
36
37    /// Retrieve the full-precision FP32 vector for node `idx`.
38    ///
39    /// On the in-memory path this copies from a slice.  On the io_uring path
40    /// this collects the pre-fetched completion or falls back to a synchronous
41    /// read if the pre-fetch was not issued.
42    ///
43    /// Returns `None` if `idx` is out of range.
44    fn fetch_fp32(&mut self, idx: u32) -> Option<Vec<f32>>;
45}
46
47/// In-memory `NodeFetcher` backed by a contiguous `Vec<Vec<f32>>`.
48///
49/// Used on all targets (including WASM) and in unit tests.  In production
50/// an mmap-backed variant may wrap this type with the same interface.
51pub struct InMemoryFetcher {
52    dim: usize,
53    vectors: Vec<Vec<f32>>,
54}
55
56impl InMemoryFetcher {
57    /// Create from an owned vector slice.  All entries must have length `dim`.
58    pub fn new(dim: usize, vectors: Vec<Vec<f32>>) -> Self {
59        Self { dim, vectors }
60    }
61
62    /// Create from a borrowed slice, cloning each vector.
63    pub fn from_slice(dim: usize, vectors: &[Vec<f32>]) -> Self {
64        Self {
65            dim,
66            vectors: vectors.to_vec(),
67        }
68    }
69
70    /// Number of nodes stored.
71    pub fn len(&self) -> usize {
72        self.vectors.len()
73    }
74
75    /// Returns `true` if no vectors are stored.
76    pub fn is_empty(&self) -> bool {
77        self.vectors.is_empty()
78    }
79}
80
81impl NodeFetcher for InMemoryFetcher {
82    fn dim(&self) -> usize {
83        self.dim
84    }
85
86    /// No-op: all data is already in RAM.
87    fn prefetch_batch(&mut self, _indices: &[u32]) {}
88
89    fn fetch_fp32(&mut self, idx: u32) -> Option<Vec<f32>> {
90        self.vectors.get(idx as usize).cloned()
91    }
92}
93
94#[cfg(test)]
95mod tests {
96    use super::*;
97
98    #[test]
99    fn in_memory_fetcher_returns_correct_vector() {
100        let vecs = vec![
101            vec![1.0_f32, 2.0, 3.0],
102            vec![4.0_f32, 5.0, 6.0],
103            vec![7.0_f32, 8.0, 9.0],
104        ];
105        let mut f = InMemoryFetcher::new(3, vecs.clone());
106        assert_eq!(f.dim(), 3);
107        assert_eq!(f.fetch_fp32(0), Some(vecs[0].clone()));
108        assert_eq!(f.fetch_fp32(2), Some(vecs[2].clone()));
109    }
110
111    #[test]
112    fn in_memory_fetcher_oob_returns_none() {
113        let mut f = InMemoryFetcher::new(2, vec![vec![1.0, 2.0]]);
114        assert!(f.fetch_fp32(1).is_none());
115        assert!(f.fetch_fp32(100).is_none());
116    }
117
118    #[test]
119    fn in_memory_fetcher_empty() {
120        let mut f = InMemoryFetcher::new(4, vec![]);
121        assert!(f.is_empty());
122        assert!(f.fetch_fp32(0).is_none());
123    }
124
125    #[test]
126    fn prefetch_batch_is_noop() {
127        let mut f = InMemoryFetcher::new(2, vec![vec![1.0, 2.0], vec![3.0, 4.0]]);
128        // Should not panic or error.
129        f.prefetch_batch(&[0, 1]);
130        assert_eq!(f.fetch_fp32(0), Some(vec![1.0, 2.0]));
131    }
132}
nodedb_vector/vamana/node_fetcher.rs

nodedb_vector/vamana/
node_fetcher.rs
