1use std::cell::RefCell;
7use std::collections::HashSet;
8use std::fs::{File, OpenOptions, TryLockError};
9use std::path::{Path, PathBuf};
10use std::sync::atomic::{AtomicU8, AtomicUsize, Ordering};
11use std::sync::Arc;
12
13use crate::core::sync::RwLock;
14
15use super::{
16 kernels::dot, vertex::SearchScratch, BuildParams, Graph, Metric, PreparedRows, Residency,
17 VGraph, VSearchParams,
18};
19use crate::core::durable_file::{atomic_write, sync_directory, sync_mapped_file};
20use crate::core::error::WarmError;
21
22mod delta;
23mod delta_rows;
24mod diagnostics;
25mod format;
26mod merge;
27use delta::Delta;
28use format::{BaseRows, BaseRowsWriter};
29
30const MAX_DIM: usize = 1 << 20;
31const MAX_PAYLOAD: usize = 64 << 20;
32const MAX_OPS: usize = 65536;
33const THREE_BYTE_IDS_BELOW: usize = (1 << 24) - 1;
35
36#[derive(Debug, thiserror::Error)]
38pub enum UpdateError {
39 #[error("unknown vector id {0}")]
40 UnknownId(u32),
41 #[error("vector id {0} already exists")]
42 AlreadyExists(u32),
43 #[error("id {0} occurs twice in one batch")]
44 DuplicateIdInBatch(u32),
45 #[error("expected dimension {expected}, got {actual}")]
46 DimensionMismatch { expected: usize, actual: usize },
47 #[error("invalid vector: {0}")]
48 InvalidVector(&'static str),
49 #[error("id u32::MAX is reserved")]
50 IdOutOfRange,
51 #[error("vector index allocation or batch limit exceeded")]
52 Capacity,
53 #[error("corrupt vector log at byte {offset}: {reason}")]
54 CorruptLog { offset: u64, reason: String },
55 #[error("corrupt vector base: {0}")]
56 CorruptBase(String),
57 #[error("unsupported vector format: {0}")]
58 UnsupportedFormat(&'static str),
59 #[error("another writer is active")]
60 WriterBusy,
61 #[error("index is already open by a writer")]
62 Locked,
63 #[error("write outcome uncertain; reopen the index before further writes")]
64 RecoveryRequired,
65 #[error(transparent)]
66 Io(#[from] std::io::Error),
67 #[error(transparent)]
68 Persistence(#[from] anyhow::Error),
69 #[error(transparent)]
70 Search(#[from] WarmError),
71}
72
73#[derive(Clone, Copy, Debug)]
75pub struct ConsolidationPolicy {
76 pub delta_ratio: f64,
77 pub tombstone_ratio: f64,
78 pub min_delta: usize,
79}
80impl Default for ConsolidationPolicy {
81 fn default() -> Self {
82 Self {
83 delta_ratio: 0.1,
84 tombstone_ratio: 0.2,
85 min_delta: 1024,
86 }
87 }
88}
89
90#[derive(Clone, Copy, Debug)]
92pub enum ConsolidationMode {
93 Rebuild,
95 Merge { rebuild_after: f64 },
102}
103
104#[derive(Clone, Copy, Debug, PartialEq, Eq)]
106pub enum Consolidation {
107 Unchanged,
109 Merged,
110 Rebuilt,
111}
112
113#[derive(Clone)]
114pub(super) struct Encoding {
115 pub(super) bits: u8,
116 pub(super) residual_bits: u8,
117 pub(super) align: bool,
118 pub(super) reorder: bool,
119 pub(super) id_bytes: u8,
120 pub(super) screen: bool,
121 pub(super) block_bits: u8,
122}
123impl Default for Encoding {
124 fn default() -> Self {
125 Self {
126 bits: 8,
127 residual_bits: 8,
128 align: true,
129 reorder: false,
130 id_bytes: 4,
131 screen: false,
132 block_bits: 0,
133 }
134 }
135}
136#[derive(Clone)]
137struct Manifest {
138 identity: [u8; 16],
139 dim: usize,
140 sequence: u64,
141 build: BuildParams,
142 encoding: Encoding,
143}
144impl Manifest {
145 fn vertex_path(&self, dir: &Path) -> PathBuf {
146 dir.join(format!("base-{}.vertex", self.sequence))
147 }
148 fn rows_path(&self, dir: &Path) -> PathBuf {
149 dir.join(format!("rows-{}.bin", self.sequence))
150 }
151}
152
153#[derive(Clone)]
154enum Operation {
155 Insert(u32, Vec<f32>),
156 Delete(u32),
157 Replace(u32, Vec<f32>),
158}
159impl Operation {
160 fn id(&self) -> u32 {
161 match self {
162 Self::Insert(id, _) | Self::Delete(id) | Self::Replace(id, _) => *id,
163 }
164 }
165 fn tag(&self) -> u8 {
166 match self {
167 Self::Insert(..) => 1,
168 Self::Delete(..) => 2,
169 Self::Replace(..) => 3,
170 }
171 }
172 fn vector(&self) -> Option<&[f32]> {
173 match self {
174 Self::Insert(_, v) | Self::Replace(_, v) => Some(v),
175 Self::Delete(_) => None,
176 }
177 }
178}
179struct Base {
180 graph: Option<VGraph>,
181 rows: BaseRows,
182 dense_ids: bool,
183}
184impl Base {
185 fn open(dir: &Path, m: &Manifest) -> Result<Self, UpdateError> {
186 let rows = BaseRows::open(&m.rows_path(dir), m)?;
187 let graph = if rows.ids.is_empty() {
188 None
189 } else {
190 let graph = VGraph::open(&m.vertex_path(dir), Residency::Mapped)?;
191 if graph.len() != rows.ids.len()
192 || graph.dim() != m.dim
193 || graph.metric() != m.build.metric.into()
194 {
195 return Err(UpdateError::CorruptBase(
196 "vertex/rows shape mismatch".into(),
197 ));
198 }
199 Some(graph)
200 };
201 let dense_ids = rows.ids.iter().enumerate().all(|(i, &id)| i == id as usize);
202 Ok(Self {
203 graph,
204 rows,
205 dense_ids,
206 })
207 }
208}
209#[derive(Clone)]
210struct State {
211 base: Arc<Base>,
212 tombstones: Vec<u64>,
213 deleted: usize,
214 delta: Delta,
215 pending: Vec<Operation>,
216 touched: HashSet<u32>,
217 pending_bytes: usize,
218 manifest: Manifest,
219 sequence: u64,
220 log_end: u64,
221}
222impl State {
223 fn new(base: Base, manifest: Manifest) -> Self {
224 Self {
225 tombstones: vec![0; base.rows.ids.len().div_ceil(64)],
226 base: Arc::new(base),
227 deleted: 0,
228 delta: Delta::default(),
229 pending: Vec::new(),
230 touched: HashSet::new(),
231 pending_bytes: 0,
232 sequence: manifest.sequence,
233 manifest,
234 log_end: format::LOG_HEADER,
235 }
236 }
237 fn dead(&self, slot: usize) -> bool {
238 self.tombstones[slot / 64] & (1 << (slot % 64)) != 0
239 }
240 fn base_slot(&self, id: u32) -> Option<usize> {
241 self.base
242 .rows
243 .ids
244 .binary_search(&id)
245 .ok()
246 .filter(|&s| !self.dead(s))
247 }
248 fn contains(&self, id: u32) -> bool {
249 self.delta.live.contains_key(&id) || self.base_slot(id).is_some()
250 }
251 fn len(&self) -> usize {
252 self.base.rows.ids.len() - self.deleted + self.delta.live.len()
253 }
254 fn validate_op(&self, op: &Operation) -> Result<(), UpdateError> {
255 if op.id() == u32::MAX {
256 return Err(UpdateError::IdOutOfRange);
257 }
258 if let Some(v) = op.vector() {
259 validate_vector(v, self.manifest.dim, self.manifest.build.metric)?;
260 }
261 match op {
262 Operation::Insert(id, _) if self.contains(*id) => Err(UpdateError::AlreadyExists(*id)),
263 Operation::Delete(id) | Operation::Replace(id, _) if !self.contains(*id) => {
264 Err(UpdateError::UnknownId(*id))
265 }
266 _ => Ok(()),
267 }
268 }
269 fn apply(&mut self, op: &Operation) {
270 let id = op.id();
271 if !matches!(op, Operation::Insert(..)) {
272 self.delta.remove(id);
273 if let Some(slot) = self.base_slot(id) {
274 self.tombstones[slot / 64] |= 1 << (slot % 64);
275 self.deleted += 1;
276 }
277 }
278 if let Some(row) = op.vector() {
279 let mut row = row.to_vec();
280 super::prepare_rows(&mut row, self.manifest.dim, self.manifest.build.metric);
281 self.delta.insert(id, &row, &self.manifest.build);
282 }
283 }
284 fn replay(&mut self, ops: Vec<Operation>) -> Result<(), UpdateError> {
285 for op in &ops {
286 self.validate_op(op)?;
287 }
288 for op in &ops {
289 self.apply(op);
290 }
291 Ok(())
292 }
293}
294
295thread_local! { static SCRATCH: RefCell<Option<SearchScratch>> = const { RefCell::new(None) }; }
296
297pub struct UpdatableVGraph {
301 dir: PathBuf,
302 _lock: IndexLock,
303 writer: AtomicU8,
304 state: RwLock<State>,
305 diagnostics: diagnostics::Diagnostics,
306 churn: AtomicUsize,
308}
309struct Writer<'a>(&'a AtomicU8);
310impl Writer<'_> {
311 fn uncertain(&self) {
312 self.0.store(2, Ordering::Release);
313 }
314}
315impl Drop for Writer<'_> {
316 fn drop(&mut self) {
317 let _ = self
318 .0
319 .compare_exchange(1, 0, Ordering::Release, Ordering::Relaxed);
320 }
321}
322struct IndexLock(File);
323impl Drop for IndexLock {
324 fn drop(&mut self) {
325 if let Err(error) = self.0.unlock() {
327 tracing::warn!(%error, "failed to release vector index lock");
328 }
329 }
330}
331fn lock(dir: &Path) -> Result<IndexLock, UpdateError> {
332 let file = OpenOptions::new()
333 .create(true)
334 .truncate(false)
335 .read(true)
336 .write(true)
337 .open(dir.join(".lock"))?;
338 file.try_lock().map_err(|e| match e {
339 TryLockError::WouldBlock => UpdateError::Locked,
340 TryLockError::Error(e) => UpdateError::Io(e),
341 })?;
342 Ok(IndexLock(file))
343}
344fn validate_vector(row: &[f32], dim: usize, metric: Metric) -> Result<(), UpdateError> {
345 if row.len() != dim {
346 return Err(UpdateError::DimensionMismatch {
347 expected: dim,
348 actual: row.len(),
349 });
350 }
351 if row.iter().any(|v| !v.is_finite()) {
352 return Err(UpdateError::InvalidVector("non-finite coordinate"));
353 }
354 if !dot(row, row).is_finite() {
355 return Err(UpdateError::InvalidVector("squared norm overflows f32"));
356 }
357 if metric == Metric::Angular && row.iter().all(|&v| v == 0.0) {
358 return Err(UpdateError::InvalidVector("zero angular vector"));
359 }
360 Ok(())
361}
362fn distance(q: &[f32], row: &[f32], metric: Metric) -> f32 {
363 match metric {
364 Metric::Angular => -dot(q, row),
365 Metric::Euclidean => dot(row, row) * 0.5 - dot(q, row),
366 }
367}
368
369fn neighbour_id_bytes(configured: u8, rows: usize) -> u8 {
373 match configured {
374 3 | 4 if rows < THREE_BYTE_IDS_BELOW => 3,
375 3 | 4 => 4,
376 other => other,
377 }
378}
379
380fn widened_width(requested: usize, allocated: usize, live: usize) -> usize {
381 ((requested.min(allocated) as u128 * allocated as u128).div_ceil(live as u128))
382 .min(allocated as u128) as usize
383}
384
385impl UpdatableVGraph {
386 pub fn create(
388 path: &Path,
389 data: &[f32],
390 dim: usize,
391 params: &BuildParams,
392 ) -> Result<Self, UpdateError> {
393 Self::create_with_encoding(path, data, dim, params, Encoding::default())
394 }
395 fn create_with_encoding(
396 path: &Path,
397 data: &[f32],
398 dim: usize,
399 params: &BuildParams,
400 encoding: Encoding,
401 ) -> Result<Self, UpdateError> {
402 if dim == 0
403 || dim > MAX_DIM
404 || !data.len().is_multiple_of(dim)
405 || data.len() / dim >= u32::MAX as usize
406 {
407 return Err(UpdateError::Capacity);
408 }
409 let manifest = Manifest {
410 dim,
411 identity: rand::random(),
412 sequence: 0,
413 build: params.clone(),
414 encoding,
415 };
416 Manifest::decode(&manifest.encode())?;
417 for row in data.chunks_exact(dim) {
418 validate_vector(row, dim, params.metric)?;
419 }
420 std::fs::create_dir(path)?;
421 let file = lock(path)?;
422 let mut vectors = data.to_vec();
423 super::prepare_rows(&mut vectors, dim, params.metric);
424 let ids: Vec<u32> = (0..(data.len() / dim) as u32).collect();
425 write_base(path, &manifest, &ids, vectors)?;
426 format::rotate(&path.join("updates.log"), &manifest)?;
427 atomic_write(&path.join("CURRENT"), &manifest.encode())?;
428 sync_directory(
429 path.parent()
430 .filter(|p| !p.as_os_str().is_empty())
431 .unwrap_or_else(|| Path::new(".")),
432 )?;
433 let base = Base::open(path, &manifest)?;
434 Ok(Self {
435 dir: path.to_path_buf(),
436 _lock: file,
437 writer: AtomicU8::new(0),
438 state: RwLock::new(State::new(base, manifest)),
439 diagnostics: diagnostics::Diagnostics::from_env(),
440 churn: AtomicUsize::new(0),
441 })
442 }
443 pub fn create_from_graph(
445 path: &Path,
446 graph: &Graph,
447 vertex: &VGraph,
448 ) -> Result<Self, UpdateError> {
449 if graph.n != vertex.len()
450 || graph.dim != vertex.dim()
451 || vertex.metric() != graph.metric.into()
452 {
453 return Err(UpdateError::CorruptBase(
454 "graph and vertex encoding differ".into(),
455 ));
456 }
457 let manifest = Manifest {
458 identity: rand::random(),
459 dim: graph.dim,
460 sequence: 0,
461 build: graph.build.clone(),
462 encoding: vertex.update_encoding(),
463 };
464 Manifest::decode(&manifest.encode())?;
465 for row in graph.vectors.chunks_exact(graph.dim) {
466 validate_vector(row, graph.dim, graph.metric)?;
467 }
468 std::fs::create_dir(path)?;
469 let file = lock(path)?;
470 vertex.save(&manifest.vertex_path(path))?;
471 let ids: Vec<_> = (0..graph.n as u32).collect();
472 BaseRows::write(&manifest.rows_path(path), &manifest, &ids, &graph.vectors)?;
473 format::rotate(&path.join("updates.log"), &manifest)?;
474 atomic_write(&path.join("CURRENT"), &manifest.encode())?;
475 sync_directory(
476 path.parent()
477 .filter(|p| !p.as_os_str().is_empty())
478 .unwrap_or_else(|| Path::new(".")),
479 )?;
480 let base = Base::open(path, &manifest)?;
481 Ok(Self {
482 dir: path.to_path_buf(),
483 _lock: file,
484 writer: AtomicU8::new(0),
485 state: RwLock::new(State::new(base, manifest)),
486 diagnostics: diagnostics::Diagnostics::from_env(),
487 churn: AtomicUsize::new(0),
488 })
489 }
490 pub fn open(path: &Path) -> Result<Self, UpdateError> {
492 let file = lock(path)?;
493 let manifest = Manifest::read(&path.join("CURRENT"))?;
494 let base = Base::open(path, &manifest)?;
495 let mut state = State::new(base, manifest.clone());
496 let log = path.join("updates.log");
497 let (sequence, end) = format::replay(&log, &manifest, |ops| state.replay(ops))?;
498 let log_file = OpenOptions::new().write(true).open(&log)?;
499 if log_file.metadata()?.len() != end {
500 log_file.set_len(end)?;
501 log_file.sync_all()?;
502 sync_mapped_file(&log_file)?;
503 }
504 state.sequence = sequence;
505 state.log_end = end;
506 cleanup_generations(path, state.manifest.sequence)?;
507 Ok(Self {
508 dir: path.to_path_buf(),
509 _lock: file,
510 writer: AtomicU8::new(0),
511 state: RwLock::new(state),
512 diagnostics: diagnostics::Diagnostics::from_env(),
513 churn: AtomicUsize::new(0),
514 })
515 }
516 fn writer(&self) -> Result<Writer<'_>, UpdateError> {
517 self.writer
518 .compare_exchange(0, 1, Ordering::Acquire, Ordering::Relaxed)
519 .map_err(|v| {
520 if v == 2 {
521 UpdateError::RecoveryRequired
522 } else {
523 UpdateError::WriterBusy
524 }
525 })?;
526 Ok(Writer(&self.writer))
527 }
528 pub fn len(&self) -> usize {
529 self.state.read().len()
530 }
531 pub fn is_empty(&self) -> bool {
532 self.len() == 0
533 }
534 pub fn dim(&self) -> usize {
535 self.state.read().manifest.dim
536 }
537 pub fn insert(&self, id: u32, vector: &[f32]) -> Result<(), UpdateError> {
539 self.stage(id, Some(vector), 1)
540 }
541 pub fn delete(&self, id: u32) -> Result<(), UpdateError> {
543 self.stage(id, None, 2)
544 }
545 pub fn replace(&self, id: u32, vector: &[f32]) -> Result<(), UpdateError> {
547 self.stage(id, Some(vector), 3)
548 }
549 fn stage(&self, id: u32, vector: Option<&[f32]>, tag: u8) -> Result<(), UpdateError> {
550 let _writer = self.writer()?;
551 let mut s = self.state.write();
552 if id == u32::MAX {
553 return Err(UpdateError::IdOutOfRange);
554 }
555 if s.touched.contains(&id) {
556 return Err(UpdateError::DuplicateIdInBatch(id));
557 }
558 if let Some(v) = vector {
559 validate_vector(v, s.manifest.dim, s.manifest.build.metric)?;
560 }
561 let size = 5 + vector.map_or(0, |v| v.len() * 4);
562 if s.pending.len() >= MAX_OPS
563 || s.pending_bytes + size > MAX_PAYLOAD
564 || s.delta.len() >= (u32::MAX - 1) as usize
565 {
566 return Err(UpdateError::Capacity);
567 }
568 let op = match tag {
569 1 => Operation::Insert(id, vector.expect("insert vector").to_vec()),
570 2 => Operation::Delete(id),
571 _ => Operation::Replace(id, vector.expect("replace vector").to_vec()),
572 };
573 s.validate_op(&op)?;
574 s.apply(&op);
575 s.touched.insert(id);
576 s.pending_bytes += size;
577 s.pending.push(op);
578 Ok(())
579 }
580 pub fn delete_batch(&self, ids: &[u32]) -> Result<(), UpdateError> {
583 self.stage_all(ids.iter().map(|&id| Operation::Delete(id)).collect())
584 }
585 pub fn replace_batch(&self, ids: &[u32], vectors: &[f32]) -> Result<(), UpdateError> {
588 let dim = self.dim();
589 if vectors.len() != ids.len() * dim {
590 return Err(UpdateError::DimensionMismatch {
591 expected: ids.len() * dim,
592 actual: vectors.len(),
593 });
594 }
595 self.stage_all(
596 ids.iter()
597 .zip(vectors.chunks_exact(dim.max(1)))
598 .map(|(&id, row)| Operation::Replace(id, row.to_vec()))
599 .collect(),
600 )
601 }
602 fn stage_all(&self, ops: Vec<Operation>) -> Result<(), UpdateError> {
605 let _writer = self.writer()?;
606 let mut s = self.state.write();
607 let size: usize = ops
608 .iter()
609 .map(|op| 5 + op.vector().map_or(0, |v| v.len() * 4))
610 .sum();
611 if s.pending.len() + ops.len() > MAX_OPS
612 || s.pending_bytes + size > MAX_PAYLOAD
613 || s.delta.len() + ops.len() >= (u32::MAX - 1) as usize
614 {
615 return Err(UpdateError::Capacity);
616 }
617 let mut batch = HashSet::with_capacity(ops.len());
618 for op in &ops {
619 let id = op.id();
620 if s.touched.contains(&id) || !batch.insert(id) {
621 return Err(UpdateError::DuplicateIdInBatch(id));
622 }
623 s.validate_op(op)?;
624 }
625 for op in ops {
626 s.apply(&op);
627 s.touched.insert(op.id());
628 s.pending.push(op);
629 }
630 s.pending_bytes += size;
631 Ok(())
632 }
633 pub fn batch_capacity(&self) -> usize {
635 MAX_OPS.min(MAX_PAYLOAD / (5 + 4 * self.dim()))
636 }
637 pub fn insert_batch(&self, ids: &[u32], vectors: &[f32]) -> Result<(), UpdateError> {
640 let _writer = self.writer()?;
641 let mut s = self.state.write();
642 let dim = s.manifest.dim;
643 if vectors.len() != ids.len() * dim {
644 return Err(UpdateError::DimensionMismatch {
645 expected: ids.len() * dim,
646 actual: vectors.len(),
647 });
648 }
649 let size = (5 + 4 * dim) * ids.len();
650 if s.pending.len() + ids.len() > MAX_OPS
651 || s.pending_bytes + size > MAX_PAYLOAD
652 || s.delta.len() + ids.len() >= (u32::MAX - 1) as usize
653 {
654 return Err(UpdateError::Capacity);
655 }
656 let mut batch = HashSet::with_capacity(ids.len());
657 for (&id, row) in ids.iter().zip(vectors.chunks_exact(dim)) {
658 if id == u32::MAX {
659 return Err(UpdateError::IdOutOfRange);
660 }
661 if s.touched.contains(&id) || !batch.insert(id) {
662 return Err(UpdateError::DuplicateIdInBatch(id));
663 }
664 if s.contains(id) {
665 return Err(UpdateError::AlreadyExists(id));
666 }
667 validate_vector(row, dim, s.manifest.build.metric)?;
668 }
669 let mut rows = vectors.to_vec();
670 super::prepare_rows(&mut rows, dim, s.manifest.build.metric);
671 let build = s.manifest.build.clone();
672 s.delta.insert_batch(ids, &rows, &build);
673 for (&id, row) in ids.iter().zip(vectors.chunks_exact(dim)) {
674 s.pending.push(Operation::Insert(id, row.to_vec()));
675 s.touched.insert(id);
676 }
677 s.pending_bytes += size;
678 Ok(())
679 }
680 pub fn commit(&self) -> Result<(), UpdateError> {
682 let writer = self.writer()?;
683 self.commit_with(&writer)
684 }
685 fn commit_with(&self, writer: &Writer<'_>) -> Result<(), UpdateError> {
686 self.commit_with_checkpoints(writer, |_| Ok(()))
687 }
688 fn commit_with_checkpoints(
689 &self,
690 writer: &Writer<'_>,
691 mut checkpoint: impl FnMut(u8) -> Result<(), UpdateError>,
692 ) -> Result<(), UpdateError> {
693 let (bytes, end, sequence) = {
694 let s = self.state.read();
695 if s.pending.is_empty() {
696 return Ok(());
697 }
698 let sequence = s.sequence.checked_add(1).ok_or(UpdateError::Capacity)?;
699 (
700 format::frame(&s.pending, sequence, s.manifest.dim)?,
701 s.log_end,
702 sequence,
703 )
704 };
705 if let Err(e) = format::append(&self.dir.join("updates.log"), end, &bytes, &mut checkpoint)
706 {
707 writer.uncertain();
708 return Err(e);
709 }
710 let mut s = self.state.write();
711 s.sequence = sequence;
712 s.log_end += bytes.len() as u64;
713 s.pending.clear();
714 s.touched.clear();
715 s.pending_bytes = 0;
716 drop(s);
717 if let Err(error) = checkpoint(3) {
718 writer.uncertain();
719 return Err(error);
720 }
721 Ok(())
722 }
723 pub fn search(
725 &self,
726 query: &[f32],
727 k: usize,
728 params: VSearchParams,
729 ) -> Result<Vec<u32>, UpdateError> {
730 let mut out = Vec::new();
731 self.search_into(query, k, params, &mut out)?;
732 Ok(out)
733 }
734 pub fn search_into(
736 &self,
737 query: &[f32],
738 k: usize,
739 params: VSearchParams,
740 out: &mut Vec<u32>,
741 ) -> Result<(), UpdateError> {
742 let s = self.state.read();
743 validate_vector(query, s.manifest.dim, s.manifest.build.metric)?;
744 let k = k.min(s.len());
745 out.clear();
746 if k == 0 {
747 return Ok(());
748 }
749 self.diagnostics.query();
750 if s.delta.len() == 0 && s.deleted == 0 {
751 if let Some(graph) = &s.base.graph {
752 SCRATCH.with(|cache| {
753 graph.search_reusing(&mut cache.borrow_mut(), query, k, params, out)
754 })?;
755 if !s.base.dense_ids {
756 for slot in out {
757 *slot = s.base.rows.ids[*slot as usize];
758 }
759 }
760 }
761 return Ok(());
762 }
763 if let Some(graph) = s
764 .base
765 .graph
766 .as_ref()
767 .filter(|_| s.deleted < s.base.rows.ids.len())
768 {
769 let live = graph.len() - s.deleted;
770 let mut take = widened_width(params.ef.max(params.rerank).max(k), graph.len(), live);
771 loop {
772 let mut widened = params;
773 widened.ef = take;
774 widened.rerank = 0;
775 widened.patience = 0;
776 SCRATCH.with(|cache| {
777 graph.search_reusing(&mut cache.borrow_mut(), query, take, widened, out)
778 })?;
779 if take == graph.len()
780 || out.iter().filter(|&&slot| !s.dead(slot as usize)).count() >= k.min(live)
781 {
782 break;
783 }
784 take = take.saturating_mul(2).min(graph.len());
785 }
786 }
787 let metric = s.manifest.build.metric;
788 let mut candidates =
789 Vec::with_capacity(out.len() + s.delta.live.len().min(params.ef.max(k)));
790 for &slot in out.iter() {
791 let slot = slot as usize;
792 if !s.dead(slot) {
793 candidates.push((
794 distance(query, s.base.rows.row(slot)?, metric),
795 s.base.rows.ids[slot],
796 ));
797 }
798 }
799 if !s.delta.live.is_empty() {
800 let mut ef = widened_width(
801 params.ef.max(params.rerank).max(k),
802 s.delta.len(),
803 s.delta.live.len(),
804 );
805 loop {
806 let found = if self.diagnostics.enabled {
807 let (found, visits) = s.delta.search_profiled(query, ef, metric);
808 self.diagnostics.delta_search(visits, found.len());
809 found
810 } else {
811 s.delta.search(query, ef, metric)
812 };
813 if ef == s.delta.len() || found.len() >= k.min(s.delta.live.len()) {
814 candidates.extend(found);
815 break;
816 }
817 ef = ef.saturating_mul(2).min(s.delta.len());
818 }
819 }
820 candidates.sort_unstable_by(|a, b| a.0.total_cmp(&b.0).then(a.1.cmp(&b.1)));
821 candidates.truncate(k);
822 out.clear();
823 out.extend(candidates.into_iter().map(|(_, id)| id));
824 Ok(())
825 }
826 pub fn streaming_diagnostics(&self) -> Option<std::collections::BTreeMap<String, f64>> {
830 self.diagnostics.snapshot(&self.state.read())
831 }
832 pub fn needs_consolidation(&self, policy: ConsolidationPolicy) -> Result<bool, UpdateError> {
833 if !policy.delta_ratio.is_finite()
834 || !policy.tombstone_ratio.is_finite()
835 || !(0.0..=1.0).contains(&policy.delta_ratio)
836 || !(0.0..=1.0).contains(&policy.tombstone_ratio)
837 {
838 return Err(UpdateError::InvalidVector(
839 "consolidation ratios must be in 0..=1",
840 ));
841 }
842 let s = self.state.read();
843 let base = s.base.rows.ids.len().max(1) as f64;
844 Ok(
845 (s.delta.len() >= policy.min_delta
846 && s.delta.len() as f64 / base >= policy.delta_ratio)
847 || (s.deleted > 0 && s.deleted as f64 / base >= policy.tombstone_ratio),
848 )
849 }
850 pub fn consolidate(&self) -> Result<(), UpdateError> {
852 self.consolidate_inner(|_| Ok(()))
853 }
854 pub fn consolidate_with(&self, mode: ConsolidationMode) -> Result<Consolidation, UpdateError> {
856 self.consolidate_mode(mode, |_| Ok(()))
857 }
858 fn consolidate_inner(
859 &self,
860 checkpoint: impl FnMut(u8) -> Result<(), UpdateError>,
861 ) -> Result<(), UpdateError> {
862 self.consolidate_mode(ConsolidationMode::Rebuild, checkpoint)
863 .map(drop)
864 }
865 fn spill_delta_rows(&self, _writer: &Writer<'_>) -> Result<(), UpdateError> {
866 let source = {
867 let state = self.state.read();
868 if state.manifest.encoding.screen || state.manifest.encoding.block_bits != 0 {
869 None
870 } else {
871 state.delta.vectors.spill_source()
872 }
873 };
874 if let Some(source) = source {
875 let mapped = source.mapped()?;
876 let mut state = self.state.write();
877 assert!(
878 state.delta.vectors.same_backing(&source),
879 "delta changed while its writer was held"
880 );
881 state.delta.vectors = mapped;
882 }
883 Ok(())
884 }
885
886 fn consolidate_mode(
887 &self,
888 mode: ConsolidationMode,
889 mut checkpoint: impl FnMut(u8) -> Result<(), UpdateError>,
890 ) -> Result<Consolidation, UpdateError> {
891 let writer = self.writer()?;
892 self.commit_with(&writer)?;
893 checkpoint(0)?;
894 self.spill_delta_rows(&writer)?;
895 let snapshot = self.state.read();
898 if snapshot.delta.len() == 0 && snapshot.deleted == 0 {
899 return Ok(Consolidation::Unchanged);
900 }
901 let mut manifest = snapshot.manifest.clone();
902 manifest.sequence = snapshot.sequence;
903 let mut ids: Vec<_> = snapshot
904 .base
905 .rows
906 .ids
907 .iter()
908 .enumerate()
909 .filter(|(slot, _)| !snapshot.dead(*slot))
910 .map(|(_, id)| *id)
911 .chain(snapshot.delta.live.keys().copied())
912 .collect();
913 ids.sort_unstable();
914 manifest.encoding.id_bytes = neighbour_id_bytes(manifest.encoding.id_bytes, ids.len());
915 let base = Arc::clone(&snapshot.base);
916 let live_base = base.rows.ids.len() - snapshot.deleted;
917 let changed = snapshot.delta.live.len() + snapshot.deleted;
918 let merge = match mode {
919 ConsolidationMode::Merge { rebuild_after } => {
920 let churn = self.churn.load(Ordering::Relaxed) + changed;
921 base.graph.is_some()
922 && live_base > 0
923 && snapshot.delta.live.len() <= live_base
924 && !manifest.encoding.reorder
925 && !manifest.encoding.screen
926 && manifest.encoding.block_bits == 0
927 && (churn as f64) < rebuild_after * ids.len() as f64
928 }
929 ConsolidationMode::Rebuild => false,
930 };
931 let plan = merge.then(|| {
932 let (new_to_old, old_to_new) =
933 merge::slot_maps(&base.rows.ids, |slot| !snapshot.dead(slot), &ids);
934 (new_to_old, old_to_new)
935 });
936 drop(snapshot);
937 if !manifest.encoding.screen && manifest.encoding.block_bits == 0 {
938 self.write_mapped_base(&manifest, &ids, &base, plan, &mut checkpoint)?;
939 } else {
940 let snapshot = self.state.read();
941 let mut rows = Vec::new();
942 rows.try_reserve_exact(
943 ids.len()
944 .checked_mul(manifest.dim)
945 .ok_or(UpdateError::Capacity)?,
946 )
947 .map_err(|_| UpdateError::Capacity)?;
948 for &id in &ids {
949 if let Some(&slot) = snapshot.delta.live.get(&id) {
950 rows.extend_from_slice(snapshot.delta.row(slot, manifest.dim));
951 } else {
952 rows.extend_from_slice(
953 snapshot
954 .base
955 .rows
956 .row(snapshot.base_slot(id).expect("live base id"))?,
957 );
958 }
959 }
960 drop(snapshot);
961 write_base_with_checkpoints(&self.dir, &manifest, &ids, rows, &mut checkpoint)?;
962 }
963 drop(base);
964 checkpoint(1)?;
965 let base = Base::open(&self.dir, &manifest)?;
966 let publication = (|| {
967 atomic_write(&self.dir.join("CURRENT"), &manifest.encode())?;
968 checkpoint(2)?;
969 checkpoint(3)?;
970 format::rotate(&self.dir.join("updates.log"), &manifest)?;
971 checkpoint(4)?;
972 Ok::<_, UpdateError>(())
973 })();
974 if let Err(e) = publication {
975 writer.uncertain();
976 return Err(e);
977 }
978 *self.state.write() = State::new(base, manifest.clone());
979 if merge {
980 self.churn.fetch_add(changed, Ordering::Relaxed);
981 } else {
982 self.churn.store(0, Ordering::Relaxed);
983 }
984 if let Err(error) = checkpoint(5) {
985 writer.uncertain();
986 return Err(error);
987 }
988 cleanup_generations(&self.dir, manifest.sequence)?;
989 Ok(if merge {
990 Consolidation::Merged
991 } else {
992 Consolidation::Rebuilt
993 })
994 }
995 fn write_mapped_base(
996 &self,
997 manifest: &Manifest,
998 ids: &[u32],
999 base: &Base,
1000 plan: Option<(Vec<u32>, Vec<u32>)>,
1001 mut checkpoint: impl FnMut(u8) -> Result<(), UpdateError>,
1002 ) -> Result<(), UpdateError> {
1003 checkpoint(6)?;
1004 let delta_rows = self.state.read().delta.vectors.clone();
1005 let t = super::prof::start();
1006 let mut output = BaseRowsWriter::new(&manifest.rows_path(&self.dir), manifest, ids)?;
1007 let mut prepared = PreparedRows::new(ids.len(), manifest.dim)?;
1008 let chunk_rows = ((4 << 20) / 4 / manifest.dim).max(1);
1009 let mut rows = Vec::new();
1010 rows.try_reserve_exact(
1011 chunk_rows
1012 .min(ids.len())
1013 .checked_mul(manifest.dim)
1014 .ok_or(UpdateError::Capacity)?,
1015 )
1016 .map_err(|_| UpdateError::Capacity)?;
1017 if let Some(graph) = &base.graph {
1018 graph.release_pages()?;
1019 }
1020 for chunk in ids.chunks(chunk_rows) {
1021 rows.clear();
1022 let snapshot = self.state.read();
1023 for &id in chunk {
1024 if let Some(&slot) = snapshot.delta.live.get(&id) {
1025 rows.extend_from_slice(snapshot.delta.row(slot, manifest.dim));
1026 } else {
1027 rows.extend_from_slice(
1028 base.rows
1029 .row(snapshot.base_slot(id).expect("live base id"))?,
1030 );
1031 }
1032 }
1033 drop(snapshot);
1034 output.append(&rows)?;
1035 prepared.append(&rows)?;
1036 base.rows.release_pages()?;
1037 delta_rows.release_pages()?;
1038 }
1039 output.finish()?;
1040 super::prof::report("rows", t);
1041 checkpoint(7)?;
1042 if let Some(graph) = &base.graph {
1043 graph.release_pages()?;
1044 }
1045 if !ids.is_empty() {
1046 let graph = match (plan, &base.graph) {
1047 (Some((new_to_old, old_to_new)), Some(old)) => {
1048 let plan = merge::MergePlan {
1049 old,
1050 new_to_old,
1051 old_to_new,
1052 };
1053 let (dim, build) = (manifest.dim, &manifest.build);
1054 prepared.build_with(build, |rows, norms| {
1055 merge::merge_structure(rows, norms, dim, build, &plan)
1056 })?
1057 }
1058 _ => prepared.build(&manifest.build)?,
1059 };
1060 checkpoint(8)?;
1061 let t = super::prof::start();
1062 let e = &manifest.encoding;
1063 VGraph::save_prepared(
1064 &graph,
1065 &manifest.vertex_path(&self.dir),
1066 e.bits,
1067 e.residual_bits,
1068 e.align,
1069 e.reorder,
1070 e.id_bytes,
1071 )?;
1072 super::prof::report("encode", t);
1073 checkpoint(9)?;
1074 }
1075 sync_directory(&self.dir)?;
1076 checkpoint(10)?;
1077 Ok(())
1078 }
1079}
1080fn cleanup_generations(dir: &Path, active: u64) -> Result<(), UpdateError> {
1081 for entry in std::fs::read_dir(dir)? {
1082 let entry = entry?;
1083 let name = entry.file_name();
1084 let Some(name) = name.to_str() else { continue };
1085 let sequence = name
1086 .strip_prefix("base-")
1087 .and_then(|s| s.strip_suffix(".vertex"))
1088 .or_else(|| {
1089 name.strip_prefix("rows-")
1090 .and_then(|s| s.strip_suffix(".bin"))
1091 })
1092 .and_then(|s| s.parse::<u64>().ok());
1093 if sequence.is_some_and(|s| s != active) {
1094 std::fs::remove_file(entry.path())?;
1095 }
1096 }
1097 sync_directory(dir)?;
1098 Ok(())
1099}
1100
1101fn write_base(
1104 dir: &Path,
1105 manifest: &Manifest,
1106 ids: &[u32],
1107 rows: Vec<f32>,
1108) -> Result<(), UpdateError> {
1109 write_base_with_checkpoints(dir, manifest, ids, rows, |_| Ok(()))
1110}
1111
1112fn write_base_with_checkpoints(
1113 dir: &Path,
1114 manifest: &Manifest,
1115 ids: &[u32],
1116 rows: Vec<f32>,
1117 mut checkpoint: impl FnMut(u8) -> Result<(), UpdateError>,
1118) -> Result<(), UpdateError> {
1119 checkpoint(6)?;
1120 BaseRows::write(&manifest.rows_path(dir), manifest, ids, &rows)?;
1121 checkpoint(7)?;
1122 if !ids.is_empty() {
1123 let graph = Graph::build_owned(rows, manifest.dim, &manifest.build)?;
1124 let e = &manifest.encoding;
1125 checkpoint(8)?;
1126 if !e.screen && e.block_bits == 0 {
1127 VGraph::save_from_graph(
1128 &graph,
1129 &manifest.vertex_path(dir),
1130 e.bits,
1131 e.residual_bits,
1132 e.align,
1133 e.reorder,
1134 e.id_bytes,
1135 )?;
1136 } else {
1137 let mut vertex = VGraph::from_graph_with(
1138 &graph,
1139 e.bits,
1140 e.residual_bits,
1141 e.align,
1142 e.reorder,
1143 e.id_bytes,
1144 e.screen,
1145 )?;
1146 if e.block_bits != 0 {
1147 vertex = vertex.with_blocks(&graph, e.block_bits)?;
1148 }
1149 vertex.save(&manifest.vertex_path(dir))?;
1150 }
1151 checkpoint(9)?;
1152 }
1153 sync_directory(dir)?;
1154 checkpoint(10)?;
1155 Ok(())
1156}
1157
1158#[cfg(test)]
1159mod tests;
1160
1161#[cfg(test)]
1162mod batch_tests {
1163 use super::*;
1164
1165 #[test]
1168 fn rejected_batches_stage_nothing() {
1169 let dir = tempfile::tempdir().unwrap();
1170 let path = dir.path().join("index");
1171 let params = BuildParams {
1172 metric: Metric::Euclidean,
1173 degree: 8,
1174 build_ef: 16,
1175 ..Default::default()
1176 };
1177 let index = UpdatableVGraph::create(&path, &[], 2, ¶ms).unwrap();
1178 let rows: Vec<f32> = (0..8).flat_map(|i| [i as f32, 0.0]).collect();
1179 index
1180 .insert_batch(&(0..8).collect::<Vec<_>>(), &rows)
1181 .unwrap();
1182 index.commit().unwrap();
1183 let search = VSearchParams {
1184 ef: 64,
1185 rerank: 64,
1186 ..Default::default()
1187 };
1188 assert!(index.delete_batch(&[1, 1]).is_err());
1189 assert!(index.delete_batch(&[2, 99]).is_err());
1190 assert!(index
1191 .replace_batch(&[3, 4], &[3.0, 0.0, f32::NAN, 0.0])
1192 .is_err());
1193 assert_eq!(index.len(), 8);
1194 assert_eq!(index.search(&[2.0, 0.0], 1, search).unwrap(), vec![2]);
1195 assert_eq!(index.search(&[1.0, 0.0], 1, search).unwrap(), vec![1]);
1196 index.delete_batch(&[1, 2]).unwrap();
1197 index.replace_batch(&[3], &[100.0, 0.0]).unwrap();
1198 index.commit().unwrap();
1199 assert_eq!(index.len(), 6);
1200 assert_eq!(index.search(&[100.0, 0.0], 1, search).unwrap(), vec![3]);
1201 }
1202}