nodedb_lite/storage/pagedb_storage/segment_ext/
vector.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! `VectorSegmentExt` over pagedb segments.
4
5use async_trait::async_trait;
6use pagedb::PagedbError;
7use pagedb::vfs::Vfs;
8
9use crate::error::LiteError;
10use crate::storage::pagedb_storage::types::PagedbStorage;
11
12#[async_trait]
13impl<V: Vfs + Clone + Send + Sync + 'static> crate::storage::vector_segment_ext::VectorSegmentExt
14    for PagedbStorage<V>
15where
16    <V as Vfs>::LockHandle: Sync,
17    <V as Vfs>::File: Sync,
18{
19    async fn write_vector_segment(
20        &self,
21        collection_name: &str,
22        dim: usize,
23        vectors: &[Vec<f32>],
24        surrogate_ids: &[u64],
25    ) -> Result<(), LiteError> {
26        use crate::engine::vector::pagedb_backing::build_ndvs_bytes;
27        use pagedb::{RealmId, SegmentKind};
28
29        let ndvs = build_ndvs_bytes(dim, vectors, surrogate_ids)?;
30
31        // Chunk the NDVS bytes into page-sized pieces.
32        // Page body capacity = 4096 - ENVELOPE_OVERHEAD (40).
33        const PAGE_BODY_CAP: usize = 4096 - 40;
34        let chunks: Vec<&[u8]> = ndvs.chunks(PAGE_BODY_CAP).collect();
35
36        let realm = RealmId::new([0u8; 16]);
37        let segment_name = format!("vec/hnsw/{collection_name}");
38
39        let mut writer = self
40            .db
41            .create_segment(realm, SegmentKind::Unspecified)
42            .await
43            .map_err(LiteError::from)?;
44        writer
45            .append_extent(&chunks)
46            .await
47            .map_err(LiteError::from)?;
48        let meta = writer.seal().await.map_err(LiteError::from)?;
49
50        let mut txn = self.db.begin_write().await.map_err(LiteError::from)?;
51
52        // Replace if already linked (atomic swap).
53        let already_exists = txn.link_segment(&segment_name, &meta).await;
54        match already_exists {
55            Ok(()) => {}
56            Err(PagedbError::AlreadyLinked) => {
57                // Use replace_segment to atomically swap old → new.
58                txn.replace_segment(&segment_name, &meta)
59                    .await
60                    .map_err(LiteError::from)?;
61            }
62            Err(e) => return Err(LiteError::from(e)),
63        }
64
65        txn.commit().await.map(|_| ()).map_err(LiteError::from)
66    }
67
68    async fn open_vector_segment(
69        &self,
70        collection_name: &str,
71    ) -> Result<Option<crate::engine::vector::pagedb_backing::PagedbBacking>, LiteError> {
72        use crate::engine::vector::pagedb_backing::PagedbBacking;
73
74        let segment_name = format!("vec/hnsw/{collection_name}");
75        let txn = self.db.begin_read().await.map_err(LiteError::from)?;
76
77        let reader = match txn.open_segment(&segment_name).await {
78            Ok(r) => r,
79            Err(PagedbError::NotFound) => return Ok(None),
80            Err(e) => return Err(LiteError::from(e)),
81        };
82
83        let backing = PagedbBacking::open(reader).await?;
84        Ok(Some(backing))
85    }
86
87    async fn delete_vector_segment(&self, collection_name: &str) -> Result<(), LiteError> {
88        let segment_name = format!("vec/hnsw/{collection_name}");
89        let mut txn = self.db.begin_write().await.map_err(LiteError::from)?;
90        match txn.unlink_segment(&segment_name).await {
91            Ok(()) => {}
92            Err(PagedbError::NotLinked) => return Ok(()), // already gone
93            Err(e) => return Err(LiteError::from(e)),
94        }
95        txn.commit().await.map(|_| ()).map_err(LiteError::from)
96    }
97}
nodedb_lite/storage/pagedb_storage/segment_ext/vector.rs

nodedb_lite/storage/pagedb_storage/segment_ext/
vector.rs
