nodedb_vector/vamana/
shard.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! `Shard` trait — abstracts local-mmap vs future remote-RPC partition fetch.
4//!
5//! Vamana's beam-search rerank step needs full-precision FP32 vectors for
6//! exact distance computation.  On a single node these live in an mmap region;
7//! in a distributed deployment they would come from a remote RPC or RDMA read.
8//! This trait provides that seam without coupling search to any I/O mechanism.
9
10/// Abstraction over a single partition's vector store.
11///
12/// `LocalShard` is the single-node implementation backed by an in-memory
13/// `Vec`.  Future impls (RDMA, QUIC RPC) implement the same interface.
14pub trait Shard: Send + Sync {
15    /// Vector dimensionality.
16    fn dim(&self) -> usize;
17
18    /// Fetch the full-precision FP32 vector at the given node index.
19    ///
20    /// Returns `None` if `idx` is out of range.  Local = mmap read (or
21    /// `Vec` copy in tests); remote = future RPC/RDMA fetch.
22    fn fetch_fp32(&self, idx: usize) -> Option<Vec<f32>>;
23}
24
25/// Single-node shard backed by an in-memory vector slice.
26///
27/// In production this would wrap a mmap pointer to the SSD-resident
28/// full-precision block described in `storage::VamanaStorageLayout`.
29/// For the current tier we hold owned `Vec<Vec<f32>>` so the type compiles
30/// and tests pass without io_uring plumbing.
31pub struct LocalShard {
32    /// Vector dimensionality.
33    pub dim: usize,
34    /// Dense FP32 vector storage indexed by node order.
35    pub vectors: Vec<Vec<f32>>,
36}
37
38impl LocalShard {
39    /// Create a `LocalShard` from a slice of FP32 vectors.
40    ///
41    /// All vectors must have the same length (`dim`).  The constructor
42    /// is infallible; callers are expected to validate dimensions upstream.
43    pub fn new(dim: usize, vectors: Vec<Vec<f32>>) -> Self {
44        Self { dim, vectors }
45    }
46}
47
48impl Shard for LocalShard {
49    fn dim(&self) -> usize {
50        self.dim
51    }
52
53    fn fetch_fp32(&self, idx: usize) -> Option<Vec<f32>> {
54        self.vectors.get(idx).cloned()
55    }
56}
57
58#[cfg(test)]
59mod tests {
60    use super::*;
61
62    #[test]
63    fn local_shard_returns_correct_vector() {
64        let vecs = vec![
65            vec![1.0_f32, 2.0, 3.0],
66            vec![4.0_f32, 5.0, 6.0],
67            vec![7.0_f32, 8.0, 9.0],
68        ];
69        let shard = LocalShard::new(3, vecs.clone());
70
71        assert_eq!(shard.dim(), 3);
72        assert_eq!(shard.fetch_fp32(0), Some(vecs[0].clone()));
73        assert_eq!(shard.fetch_fp32(1), Some(vecs[1].clone()));
74        assert_eq!(shard.fetch_fp32(2), Some(vecs[2].clone()));
75    }
76
77    #[test]
78    fn local_shard_out_of_range_returns_none() {
79        let shard = LocalShard::new(2, vec![vec![1.0, 2.0]]);
80        assert!(shard.fetch_fp32(1).is_none());
81        assert!(shard.fetch_fp32(100).is_none());
82    }
83
84    #[test]
85    fn local_shard_empty_returns_none() {
86        let shard = LocalShard::new(4, vec![]);
87        assert!(shard.fetch_fp32(0).is_none());
88    }
89}
nodedb_vector/vamana/shard.rs

nodedb_vector/vamana/
shard.rs
