1use alloc::format;
38use alloc::vec::Vec;
39
40use plugmem_arena::{
41 Arena, ArenaCfg, BlobHeap, BlobHeapBuilder, BlobHeapCfg, BlobId, ChunkPool, ChunkPoolCfg,
42 ListHandle, ShardMode,
43};
44
45use crate::error::Error;
46use crate::id::{FactId, NONE_U32};
47use crate::index::IdListIndex;
48use crate::index::bm25::Bm25Index;
49use crate::index::hnsw::{HnswGraph, HnswScratch};
50use crate::index::vecpool::VecPool;
51use crate::journal::Op;
52use crate::memory::persist::Sections;
53use crate::memory::shards::ShardLayout;
54use crate::model::{
55 EdgeHistorySlot, EdgeSlot, EntityByName, EntityRecord, FactAux, FactRecord, TemporalSlot,
56};
57use crate::snapshot::SnapshotSink;
58use crate::storage::{Scratch, Storage};
59use crate::tokenizer::Tokenizer;
60
61use super::Memory;
62
63pub(crate) const TOKENIZER_INDEX_VERSION: u32 = 2;
69
70const AUTO_HNSW_INSERT_BUDGET: usize = 4096;
71const NO_HNSW_INSERT_LIMIT: u32 = u32::MAX;
72
73#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
75#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
76pub enum MaintenanceMode {
77 #[default]
80 Auto,
81 Compact,
83 ReindexText,
85 OptimizeVectors,
87 Full,
89}
90
91#[derive(Clone, Copy, Debug, PartialEq, Eq)]
93#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
94pub struct MaintenanceOptions {
95 pub mode: MaintenanceMode,
97 pub max_hnsw_inserts: Option<usize>,
99}
100
101impl Default for MaintenanceOptions {
102 fn default() -> Self {
103 Self::auto()
104 }
105}
106
107impl MaintenanceOptions {
108 pub const fn auto() -> Self {
110 Self {
111 mode: MaintenanceMode::Auto,
112 max_hnsw_inserts: Some(AUTO_HNSW_INSERT_BUDGET),
113 }
114 }
115
116 pub const fn full() -> Self {
118 Self {
119 mode: MaintenanceMode::Full,
120 max_hnsw_inserts: None,
121 }
122 }
123
124 pub(crate) fn from_journal(mode: u8, max_hnsw_inserts: u32) -> Result<Self, Error> {
125 let mode = match mode {
126 0 => MaintenanceMode::Auto,
127 1 => MaintenanceMode::Compact,
128 2 => MaintenanceMode::ReindexText,
129 3 => MaintenanceMode::OptimizeVectors,
130 4 => MaintenanceMode::Full,
131 _ => return Err(Error::Corrupt("journal maintenance mode is invalid")),
132 };
133 let max_hnsw_inserts = if max_hnsw_inserts == NO_HNSW_INSERT_LIMIT {
134 None
135 } else {
136 Some(max_hnsw_inserts as usize)
137 };
138 Ok(Self {
139 mode,
140 max_hnsw_inserts,
141 })
142 }
143
144 pub(crate) fn journal_mode(self) -> u8 {
145 match self.mode {
146 MaintenanceMode::Auto => 0,
147 MaintenanceMode::Compact => 1,
148 MaintenanceMode::ReindexText => 2,
149 MaintenanceMode::OptimizeVectors => 3,
150 MaintenanceMode::Full => 4,
151 }
152 }
153
154 pub(crate) fn journal_max_hnsw_inserts(self) -> u32 {
155 self.max_hnsw_inserts
156 .map(|n| u32::try_from(n).unwrap_or(u32::MAX - 1))
157 .unwrap_or(NO_HNSW_INSERT_LIMIT)
158 }
159}
160
161const LIVE_DENSE_SLACK: usize = 8;
164
165impl<'a> Memory<'a> {
166 pub(super) fn fact_ids_ascending(&self) -> Vec<u32> {
180 let mut ids: Vec<u32> = self.facts.iter().map(|rec| rec.id.0).collect();
181 ids.sort_unstable();
182 ids
183 }
184
185 fn dense_id_span(&self) -> usize {
189 let cap = self
190 .facts
191 .len()
192 .saturating_add(1)
193 .saturating_mul(LIVE_DENSE_SLACK);
194 (self.next_fact as usize).min(cap)
195 }
196}
197
198fn scratch_err<E: core::fmt::Debug>(e: E) -> Error {
200 Error::Storage(format!("{e:?}"))
201}
202
203trait PoolSink {
209 fn push_text(&mut self, bytes: &[u8]) -> Result<BlobId, Error>;
211 fn push_vector(&mut self, src: &VecPool<'_>, slot: u32) -> Result<u32, Error>;
213}
214
215struct OwnedPools {
217 texts: BlobHeap<'static>,
218 vecs: VecPool<'static>,
219}
220
221impl PoolSink for OwnedPools {
222 fn push_text(&mut self, bytes: &[u8]) -> Result<BlobId, Error> {
223 Ok(self.texts.push(bytes)?)
224 }
225
226 fn push_vector(&mut self, src: &VecPool<'_>, slot: u32) -> Result<u32, Error> {
227 Ok(self.vecs.copy_slot(src, slot))
228 }
229}
230
231struct StreamPools<'s, T: Scratch, V: Scratch> {
235 text_scratch: &'s mut T,
236 text_index: BlobHeapBuilder,
237 vec_scratch: &'s mut V,
238 vec_count: u32,
239}
240
241impl<T: Scratch, V: Scratch> PoolSink for StreamPools<'_, T, V> {
242 fn push_text(&mut self, bytes: &[u8]) -> Result<BlobId, Error> {
243 self.text_scratch.write(bytes).map_err(scratch_err)?;
244 Ok(self.text_index.push_len(bytes.len())?)
245 }
246
247 fn push_vector(&mut self, src: &VecPool<'_>, slot: u32) -> Result<u32, Error> {
248 self.vec_scratch
249 .write(src.slot_bytes(slot as usize))
250 .map_err(scratch_err)?;
251 let new = self.vec_count;
252 self.vec_count += 1;
253 Ok(new)
254 }
255}
256
257struct RebuildMeta {
261 facts: Arena<'static, FactRecord>,
262 fact_aux: Arena<'static, FactAux>,
263 entities: Arena<'static, EntityRecord>,
264 by_name: Arena<'static, EntityByName>,
272 temporal: Arena<'static, TemporalSlot>,
273 tag_lists: ChunkPool<'static>,
274 metas: BlobHeap<'static>,
277 bm25: Bm25Index<'static>,
278 tags_idx: IdListIndex<'static>,
279 entity_facts: IdListIndex<'static>,
280}
281
282#[derive(Clone, Debug, Default, PartialEq, Eq)]
284#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
285pub struct MaintainReport {
286 pub purged: usize,
289 pub bytes_before: usize,
291 pub bytes_after: usize,
293 pub no_op: bool,
296 pub tombstones_before: usize,
298 pub facts_before: usize,
300 pub facts_after: usize,
302 pub vectors_before: usize,
304 pub vectors_after: usize,
306 pub hnsw_indexed_before: u32,
308 pub hnsw_indexed_after: u32,
310 pub shards_before: ShardLayout,
312 pub shards_after: ShardLayout,
315 pub structural_compacted: bool,
317 pub bm25_compacted: bool,
319 pub bm25_reindexed: bool,
321 pub hnsw_rebuilt: bool,
323 pub hnsw_remapped: bool,
325 pub hnsw_inserted: u32,
327 pub edges_compacted: bool,
331 pub edges_before: usize,
333 pub edge_versions_before: usize,
335}
336
337struct RepackedEdges {
339 out: Arena<'static, EdgeSlot>,
340 inn: Arena<'static, EdgeSlot>,
341 hist_out: Arena<'static, EdgeHistorySlot>,
342 hist_in: Arena<'static, EdgeHistorySlot>,
343}
344
345impl RepackedEdges {
346 fn pool_bytes(&self) -> usize {
347 self.out.pool_bytes()
348 + self.inn.pool_bytes()
349 + self.hist_out.pool_bytes()
350 + self.hist_in.pool_bytes()
351 }
352}
353
354struct Rebuilt {
357 facts: Arena<'static, FactRecord>,
358 entities: Arena<'static, EntityRecord>,
359 by_name: Arena<'static, EntityByName>,
360 fact_aux: Arena<'static, FactAux>,
361 texts: BlobHeap<'static>,
362 metas: BlobHeap<'static>,
363 tag_lists: ChunkPool<'static>,
364 bm25: Bm25Index<'static>,
365 tags_idx: IdListIndex<'static>,
366 entity_facts: IdListIndex<'static>,
367 temporal: Arena<'static, TemporalSlot>,
368 vecs: VecPool<'static>,
369 hnsw: HnswGraph<'static>,
370 edges: Option<RepackedEdges>,
371 bm25_tokenizer_version: u32,
372 layout: ShardLayout,
375 report: MaintainReport,
376}
377
378#[derive(Clone, Copy)]
379enum Bm25Policy {
380 Compact,
381 Reindex,
382}
383
384#[derive(Clone, Copy)]
385struct WorkPlan {
386 compact: bool,
387 bm25_reindex: bool,
388 optimize_vectors: bool,
389 hnsw_full_rebuild: bool,
390 repack_edges: bool,
391 max_hnsw_inserts: Option<usize>,
392 layout: ShardLayout,
405}
406
407#[derive(Clone, Copy, Default)]
408struct GraphWork {
409 rebuilt: bool,
410 remapped: bool,
411 inserted: u32,
412}
413
414impl WorkPlan {
415 fn needs_work(self) -> bool {
416 self.compact || self.bm25_reindex || self.optimize_vectors || self.repack_edges
417 }
418}
419
420fn hnsw_target(start: u32, total: u32, max_hnsw_inserts: Option<usize>) -> u32 {
421 let Some(max) = max_hnsw_inserts else {
422 return total;
423 };
424 let max = u32::try_from(max).unwrap_or(u32::MAX);
425 start.saturating_add(max).min(total)
426}
427
428impl Memory<'_> {
429 pub fn maintain<S: Storage>(
442 &mut self,
443 store: &mut S,
444 now: u64,
445 ) -> Result<MaintainReport, Error> {
446 self.maintain_with_options(store, now, MaintenanceOptions::auto())
447 }
448
449 pub fn maintain_with_options<S: Storage>(
451 &mut self,
452 store: &mut S,
453 now: u64,
454 options: MaintenanceOptions,
455 ) -> Result<MaintainReport, Error> {
456 let plan = self.work_plan(options);
457 let bytes_before = self.satellite_bytes(plan.repack_edges);
458 let mut report = self.report_skeleton(bytes_before);
459 if !plan.needs_work() {
460 report.no_op = true;
461 return Ok(report);
462 }
463
464 if !plan.compact {
465 let mut bm25 = None;
466 if plan.bm25_reindex {
467 bm25 = Some(self.reindex_bm25_from_text()?);
468 report.bm25_reindexed = true;
469 }
470 let mut hnsw = None;
471 if plan.optimize_vectors {
472 let (graph, work) =
473 self.optimize_graph(plan.hnsw_full_rebuild, plan.max_hnsw_inserts)?;
474 hnsw = Some(graph);
475 report.hnsw_rebuilt = work.rebuilt;
476 report.hnsw_remapped = work.remapped;
477 report.hnsw_inserted = work.inserted;
478 }
479 let mut entry = Vec::new();
481 Op::Maintain {
482 now,
483 mode: options.journal_mode(),
484 max_hnsw_inserts: options.journal_max_hnsw_inserts(),
485 }
486 .encode(&mut entry);
487 store
488 .append_journal(&entry)
489 .map_err(|e| Error::Storage(format!("{e:?}")))?;
490 if let Some(bm25) = bm25 {
491 self.bm25 = bm25;
492 self.bm25_tokenizer_version = TOKENIZER_INDEX_VERSION;
493 }
494 if let Some(hnsw) = hnsw {
495 self.hnsw = hnsw;
496 }
497 report.bytes_after = self.satellite_bytes(plan.repack_edges);
498 report.hnsw_indexed_after = self.hnsw.indexed();
499 return Ok(report);
500 }
501
502 let (rebuilt, _) = self.rebuild(plan)?;
503 report = rebuilt.report.clone();
504 let mut entry = Vec::new();
507 Op::Maintain {
508 now,
509 mode: options.journal_mode(),
510 max_hnsw_inserts: options.journal_max_hnsw_inserts(),
511 }
512 .encode(&mut entry);
513 store
514 .append_journal(&entry)
515 .map_err(|e| Error::Storage(format!("{e:?}")))?;
516 self.install(rebuilt);
517 Ok(report)
518 }
519
520 pub(super) fn replay_maintain_with_options(
521 &mut self,
522 options: MaintenanceOptions,
523 ) -> Result<(), Error> {
524 let plan = self.work_plan(options);
525 if !plan.needs_work() {
526 return Ok(());
527 }
528 if !plan.compact {
529 if plan.bm25_reindex {
530 self.bm25 = self.reindex_bm25_from_text()?;
531 self.bm25_tokenizer_version = TOKENIZER_INDEX_VERSION;
532 }
533 if plan.optimize_vectors {
534 let (hnsw, _) =
535 self.optimize_graph(plan.hnsw_full_rebuild, plan.max_hnsw_inserts)?;
536 self.hnsw = hnsw;
537 }
538 return Ok(());
539 }
540 let (rebuilt, _) = self.rebuild(plan)?;
541 self.install(rebuilt);
542 Ok(())
543 }
544
545 fn satellite_bytes(&self, with_edges: bool) -> usize {
550 let edges = if with_edges { self.edge_bytes() } else { 0 };
551 edges
552 + self.facts.pool_bytes()
553 + self.fact_aux.pool_bytes()
554 + self.entities.pool_bytes()
555 + self.hnsw.pool_bytes()
556 + self.texts.pool_bytes()
557 + self.metas.pool_bytes()
558 + self.tag_lists.pool_bytes()
559 + self.bm25.pool_bytes()
560 + self.tags_idx.pool_bytes()
561 + self.entity_facts.pool_bytes()
562 + self.temporal.pool_bytes()
563 + self.vecs.pool_bytes()
564 }
565
566 fn edge_bytes(&self) -> usize {
567 self.edges_out.pool_bytes()
568 + self.edges_in.pool_bytes()
569 + self.edges_hist_out.pool_bytes()
570 + self.edges_hist_in.pool_bytes()
571 }
572
573 pub fn maintenance_needed(&self, options: MaintenanceOptions) -> bool {
575 self.work_plan(options).needs_work()
576 }
577
578 pub fn maintenance_preview(
582 &self,
583 options: MaintenanceOptions,
584 bytes_before: usize,
585 ) -> MaintainReport {
586 let mut report = self.report_skeleton(bytes_before);
587 if !self.maintenance_needed(options) {
588 report.no_op = true;
589 }
590 report
591 }
592
593 fn report_skeleton(&self, bytes_before: usize) -> MaintainReport {
594 MaintainReport {
595 purged: 0,
596 bytes_before,
597 bytes_after: bytes_before,
598 no_op: false,
599 tombstones_before: self.tombstones,
600 facts_before: self.facts.len(),
601 facts_after: self.facts.len(),
602 vectors_before: self.vecs.len(),
603 vectors_after: self.vecs.len(),
604 hnsw_indexed_before: self.hnsw.indexed(),
605 hnsw_indexed_after: self.hnsw.indexed(),
606 shards_before: ShardLayout::of_config(&self.cfg),
607 shards_after: ShardLayout::of_config(&self.cfg),
608 structural_compacted: false,
609 bm25_compacted: false,
610 bm25_reindexed: false,
611 hnsw_rebuilt: false,
612 hnsw_remapped: false,
613 hnsw_inserted: 0,
614 edges_compacted: false,
615 edges_before: self.edges_out.len(),
616 edge_versions_before: self.edges_hist_out.len(),
617 }
618 }
619
620 fn work_plan(&self, options: MaintenanceOptions) -> WorkPlan {
621 let tokenizer_stale = self.bm25_tokenizer_version != TOKENIZER_INDEX_VERSION;
622 let has_tombstones = self.tombstones != 0;
623 let compaction_due = has_tombstones || self.bm25.needs_resummarize();
632 let graph_tail = self.cfg.dim != 0
633 && self.vecs.len() >= self.cfg.flat_to_hnsw
634 && self.hnsw.indexed() < self.vecs.len() as u32;
635 let layout = self.target_layout();
641 let stored_layout = ShardLayout::of_config(&self.cfg);
642 let relayout_due = stored_layout.compacted_groups_earn_rebuild(&layout);
643 let edges_need_relayout = stored_layout.edges_earn_rebuild(&layout);
648 match options.mode {
649 MaintenanceMode::Auto => WorkPlan {
650 compact: compaction_due || relayout_due,
651 bm25_reindex: tokenizer_stale,
652 optimize_vectors: graph_tail,
653 hnsw_full_rebuild: false,
654 repack_edges: edges_need_relayout,
655 max_hnsw_inserts: options.max_hnsw_inserts,
656 layout,
657 },
658 MaintenanceMode::Compact => WorkPlan {
659 compact: compaction_due || relayout_due,
660 bm25_reindex: tokenizer_stale,
661 optimize_vectors: has_tombstones && self.hnsw.indexed() != 0,
662 hnsw_full_rebuild: false,
663 repack_edges: edges_need_relayout,
664 max_hnsw_inserts: options.max_hnsw_inserts,
665 layout,
666 },
667 MaintenanceMode::ReindexText => WorkPlan {
668 compact: compaction_due || relayout_due,
669 bm25_reindex: true,
670 optimize_vectors: has_tombstones && self.hnsw.indexed() != 0,
671 hnsw_full_rebuild: false,
672 repack_edges: edges_need_relayout,
673 max_hnsw_inserts: options.max_hnsw_inserts,
674 layout,
675 },
676 MaintenanceMode::OptimizeVectors => WorkPlan {
679 compact: false,
680 bm25_reindex: false,
681 optimize_vectors: self.cfg.dim != 0
682 && self.vecs.len() >= self.cfg.flat_to_hnsw
683 && (graph_tail || self.hnsw.indexed() == 0),
684 hnsw_full_rebuild: self.hnsw.indexed() == 0,
685 repack_edges: false,
686 max_hnsw_inserts: options.max_hnsw_inserts,
687 layout,
688 },
689 MaintenanceMode::Full => WorkPlan {
690 compact: true,
691 bm25_reindex: true,
692 optimize_vectors: self.cfg.dim != 0 && self.vecs.len() >= self.cfg.flat_to_hnsw,
693 hnsw_full_rebuild: true,
694 repack_edges: !self.edges_hist_out.is_empty() || edges_need_relayout,
695 max_hnsw_inserts: None,
696 layout,
697 },
698 }
699 }
700
701 fn repack_edges(&self, layout: &ShardLayout) -> Result<RepackedEdges, Error> {
716 let ord =
717 ArenaCfg::new(layout.edges, ShardMode::Ordered).with_max_bytes(self.cfg.max_bytes);
718 let mut out = Arena::new(ord)?;
719 let mut inn = Arena::new(ord)?;
720 let mut hist_out = Arena::new(ord)?;
721 let mut hist_in = Arena::new(ord)?;
722 for edge in self.edges_out.iter() {
725 out.insert(&edge)?;
726 }
727 for edge in self.edges_in.iter() {
728 inn.insert(&edge)?;
729 }
730 for version in self.edges_hist_out.iter() {
731 hist_out.insert(&version)?;
732 }
733 for version in self.edges_hist_in.iter() {
734 hist_in.insert(&version)?;
735 }
736 Ok(RepackedEdges {
737 out,
738 inn,
739 hist_out,
740 hist_in,
741 })
742 }
743
744 fn reindex_bm25_from_text(&self) -> Result<Bm25Index<'static>, Error> {
751 let cfg = &self.cfg;
752 let mut bm25 = Bm25Index::new(cfg.shards_postings, cfg.max_bytes)?;
753 let mut tokenizer = Tokenizer::new();
754 let mut tf: Vec<(u32, u8)> = Vec::new();
755 for fid in self.fact_ids_ascending() {
756 let id = FactId(fid);
757 let Some(rec) = self.facts.get(&fid.to_be_bytes()) else {
758 continue;
759 };
760 if rec.is_tombstone() {
761 continue;
762 }
763 let text = core::str::from_utf8(self.texts.get(rec.text))
764 .map_err(|_| Error::Corrupt("maintain: fact text is not UTF-8"))?;
765 tf.clear();
766 let terms = &self.terms;
767 let tf_ref = &mut tf;
768 tokenizer.tokenize(text, &mut |token| {
769 if let Some(term) = terms.lookup(token) {
770 match tf_ref.iter_mut().find(|(t, _)| *t == term.0) {
771 Some((_, c)) => *c = c.saturating_add(1),
772 None => tf_ref.push((term.0, 1)),
773 }
774 }
775 });
776 bm25.index_doc(id, &tf)?;
777 }
778 Ok(bm25)
779 }
780
781 fn optimize_graph(
782 &self,
783 full_rebuild: bool,
784 max_hnsw_inserts: Option<usize>,
785 ) -> Result<(HnswGraph<'static>, GraphWork), Error> {
786 let total = self.vecs.len() as u32;
787 let mut graph = if full_rebuild || self.hnsw.indexed() == 0 {
788 HnswGraph::new(self.cfg.hnsw_m, self.cfg.hnsw_m0, self.cfg.max_bytes)?
789 } else {
790 self.hnsw.to_owned(self.cfg.max_bytes)?
791 };
792 let start = graph.indexed();
793 let target = hnsw_target(start, total, max_hnsw_inserts);
794 let mut scratch = HnswScratch::default();
795 graph.insert_bulk(
796 &self.vecs,
797 target,
798 self.cfg.hnsw_ef_construction,
799 &mut scratch,
800 )?;
801 Ok((
802 graph,
803 GraphWork {
804 rebuilt: full_rebuild || self.hnsw.indexed() == 0,
805 remapped: !full_rebuild && self.hnsw.indexed() != 0,
806 inserted: target.saturating_sub(start),
807 },
808 ))
809 }
810
811 fn rebuild(&self, plan: WorkPlan) -> Result<(Rebuilt, usize), Error> {
815 let cfg = &self.cfg;
816 let blob = BlobHeapCfg::new()
817 .with_max_bytes(cfg.max_bytes)
818 .with_max_blob(cfg.max_blob);
819 let mut pools = OwnedPools {
820 texts: BlobHeap::new(blob),
821 vecs: VecPool::new(cfg.dim, cfg.max_bytes),
822 };
823 let bm25_policy = if plan.bm25_reindex {
824 Bm25Policy::Reindex
825 } else {
826 Bm25Policy::Compact
827 };
828 let (m, vec_map, purged) = self.rebuild_parts(&mut pools, bm25_policy, &plan.layout)?;
829 let (hnsw, graph_work) = self.rebuild_graph(
830 &vec_map,
831 &pools.vecs,
832 plan.hnsw_full_rebuild,
833 plan.max_hnsw_inserts,
834 )?;
835 let edges = plan
836 .repack_edges
837 .then(|| self.repack_edges(&plan.layout))
838 .transpose()?;
839 let mut report = self.report_skeleton(self.satellite_bytes(plan.repack_edges));
840 report.purged = purged;
841 report.bytes_after = edges.as_ref().map_or(0, RepackedEdges::pool_bytes)
842 + m.facts.pool_bytes()
843 + m.fact_aux.pool_bytes()
844 + m.entities.pool_bytes()
845 + hnsw.pool_bytes()
846 + pools.texts.pool_bytes()
847 + m.metas.pool_bytes()
848 + m.tag_lists.pool_bytes()
849 + m.bm25.pool_bytes()
850 + m.tags_idx.pool_bytes()
851 + m.entity_facts.pool_bytes()
852 + m.temporal.pool_bytes()
853 + pools.vecs.pool_bytes();
854 report.facts_after = m.facts.len();
855 report.vectors_after = pools.vecs.len();
856 report.hnsw_indexed_after = hnsw.indexed();
857 report.structural_compacted = true;
858 report.bm25_compacted = matches!(bm25_policy, Bm25Policy::Compact);
859 report.bm25_reindexed = matches!(bm25_policy, Bm25Policy::Reindex);
860 report.hnsw_rebuilt = graph_work.rebuilt;
861 report.hnsw_remapped = graph_work.remapped;
862 report.hnsw_inserted = graph_work.inserted;
863 report.edges_compacted = edges.is_some();
864 let layout = ShardLayout::of_config(&self.cfg).realized(&plan.layout, edges.is_some());
865 report.shards_before = ShardLayout::of_config(&self.cfg);
866 report.shards_after = layout;
867 Ok((
868 Rebuilt {
869 facts: m.facts,
870 entities: m.entities,
871 by_name: m.by_name,
872 fact_aux: m.fact_aux,
873 texts: pools.texts,
874 metas: m.metas,
875 tag_lists: m.tag_lists,
876 bm25: m.bm25,
877 tags_idx: m.tags_idx,
878 entity_facts: m.entity_facts,
879 temporal: m.temporal,
880 vecs: pools.vecs,
881 hnsw,
882 edges,
883 bm25_tokenizer_version: if plan.bm25_reindex {
884 TOKENIZER_INDEX_VERSION
885 } else {
886 self.bm25_tokenizer_version
887 },
888 layout,
889 report,
890 },
891 purged,
892 ))
893 }
894
895 fn rebuild_parts<P: PoolSink>(
902 &self,
903 pools: &mut P,
904 bm25_policy: Bm25Policy,
905 layout: &ShardLayout,
906 ) -> Result<(RebuildMeta, alloc::vec::Vec<u32>, usize), Error> {
907 let cfg = &self.cfg;
908 let uni =
909 |shards: usize| ArenaCfg::new(shards, ShardMode::Uniform).with_max_bytes(cfg.max_bytes);
910 let ord =
911 |shards: usize| ArenaCfg::new(shards, ShardMode::Ordered).with_max_bytes(cfg.max_bytes);
912
913 let mut entities = Arena::new(uni(layout.entities))?;
914 let mut by_name = Arena::new(ord(layout.entities))?;
918 for entry in self.by_name.iter() {
919 by_name.insert(&entry)?;
920 }
921 let mut facts = Arena::new(uni(layout.facts))?;
922 let mut fact_aux = Arena::new(uni(layout.facts))?;
923 let mut tag_lists = ChunkPool::new(ChunkPoolCfg::new().with_max_bytes(cfg.max_bytes));
924 let dense = self.dense_id_span();
928 let mut live = alloc::vec![false; dense];
929 for rec in self.facts.iter() {
930 let at = rec.id.0 as usize;
931 if at < dense && !rec.is_tombstone() {
932 live[at] = true;
933 }
934 }
935 let is_live = |id: FactId| match live.get(id.0 as usize) {
936 Some(&flag) => flag,
937 None => self
938 .facts
939 .get(&id.0.to_be_bytes())
940 .is_some_and(|rec| !rec.is_tombstone()),
941 };
942 let mut bm25 = match bm25_policy {
943 Bm25Policy::Compact => {
944 self.bm25
945 .compact_live(layout.postings, cfg.max_bytes, is_live)?
946 }
947 Bm25Policy::Reindex => Bm25Index::new(layout.postings, cfg.max_bytes)?,
948 };
949 let mut tags_idx = IdListIndex::new(layout.postings, cfg.max_bytes)?;
950 let mut entity_facts = IdListIndex::new(layout.entities, cfg.max_bytes)?;
951 let mut temporal = Arena::new(ord(layout.temporal))?;
952 let mut metas = BlobHeap::new(
953 BlobHeapCfg::new()
954 .with_max_bytes(cfg.max_bytes)
955 .with_max_blob(cfg.max_blob),
956 );
957
958 for eid in 0..self.next_entity {
961 let rec = self
962 .entities
963 .get(&eid.to_be_bytes())
964 .ok_or(Error::Corrupt("maintain: entity id gap"))?;
965 let name_id = pools.push_text(self.texts.get(rec.name))?;
966 entities.insert(&EntityRecord {
967 name: name_id,
968 ..rec
969 })?;
970 }
971
972 let mut tokenizer = Tokenizer::new();
980 let mut tf: Vec<(u32, u8)> = Vec::new();
981
982 let mut vec_map = alloc::vec![NONE_U32; self.vecs.len()];
985
986 let mut purged = 0usize;
987 for fid in self.fact_ids_ascending() {
988 let id = FactId(fid);
989 let Some(rec) = self.facts.get(&fid.to_be_bytes()) else {
992 continue;
993 };
994
995 if rec.is_tombstone() {
996 purged += 1;
999 continue;
1000 }
1001
1002 let text_bytes = self.texts.get(rec.text);
1004 let text_id = pools.push_text(text_bytes)?;
1005 if matches!(bm25_policy, Bm25Policy::Reindex) {
1006 let text = core::str::from_utf8(text_bytes)
1007 .map_err(|_| Error::Corrupt("maintain: fact text is not UTF-8"))?;
1008
1009 tf.clear();
1010 let terms = &self.terms;
1011 let tf_ref = &mut tf;
1012 tokenizer.tokenize(text, &mut |token| {
1013 if let Some(term) = terms.lookup(token) {
1014 match tf_ref.iter_mut().find(|(t, _)| *t == term.0) {
1015 Some((_, c)) => *c = c.saturating_add(1),
1016 None => tf_ref.push((term.0, 1)),
1017 }
1018 }
1019 });
1020 bm25.index_doc(id, &tf)?;
1021 }
1022
1023 let aux = self
1028 .fact_aux
1029 .get(&fid.to_be_bytes())
1030 .ok_or(Error::Corrupt("maintain: fact aux gap"))?;
1031 let mut tags = ListHandle::EMPTY;
1032 for chunk in self.tag_lists.iter(&aux.tags) {
1033 for raw in chunk.chunks_exact(4) {
1034 let term = u32::from_be_bytes(raw.try_into().unwrap());
1035 tag_lists.push(&mut tags, &term.to_be_bytes())?;
1036 tags_idx.push(term, id, 0)?;
1037 }
1038 }
1039 let meta = if aux.meta.0 == NONE_U32 {
1042 BlobId(NONE_U32)
1043 } else {
1044 metas.push(self.metas.get(aux.meta))?
1045 };
1046 fact_aux.insert(&FactAux { id, tags, meta })?;
1047
1048 if let Some(entity) = rec.entity.some() {
1050 entity_facts.push(entity.0, id, 0)?;
1051 }
1052 temporal.insert(&TemporalSlot {
1053 recorded_at: rec.recorded_at,
1054 fact: id,
1055 })?;
1056
1057 let vector = if rec.has_vector() {
1059 let new_slot = pools.push_vector(&self.vecs, rec.vector)?;
1060 vec_map[rec.vector as usize] = new_slot;
1061 new_slot
1062 } else {
1063 NONE_U32
1064 };
1065 facts.insert(&FactRecord {
1066 text: text_id,
1067 vector,
1068 ..rec
1069 })?;
1070 }
1071
1072 Ok((
1073 RebuildMeta {
1074 facts,
1075 by_name,
1076 fact_aux,
1077 entities,
1078 temporal,
1079 tag_lists,
1080 metas,
1081 bm25,
1082 tags_idx,
1083 entity_facts,
1084 },
1085 vec_map,
1086 purged,
1087 ))
1088 }
1089
1090 pub fn snapshot_disk_first<T: Scratch, V: Scratch, Sk: SnapshotSink>(
1104 &self,
1105 created_at: u64,
1106 text_scratch: &mut T,
1107 vec_scratch: &mut V,
1108 sink: Sk,
1109 ) -> Result<usize, Error> {
1110 let options = MaintenanceOptions::auto();
1111 if !self.maintenance_needed(options) {
1112 self.write_snapshot_with(&self.sections(), created_at, sink)?;
1113 return Ok(0);
1114 }
1115 Ok(self
1116 .snapshot_disk_first_with_options(created_at, text_scratch, vec_scratch, sink, options)?
1117 .purged)
1118 }
1119
1120 pub fn snapshot_disk_first_with_options<T: Scratch, V: Scratch, Sk: SnapshotSink>(
1124 &self,
1125 created_at: u64,
1126 text_scratch: &mut T,
1127 vec_scratch: &mut V,
1128 sink: Sk,
1129 options: MaintenanceOptions,
1130 ) -> Result<MaintainReport, Error> {
1131 let plan = self.work_plan(options);
1132 let mut report = self.report_skeleton(self.satellite_bytes(plan.repack_edges));
1133 if !plan.needs_work() {
1134 report.no_op = true;
1135 return Ok(report);
1136 }
1137 let cfg = &self.cfg;
1138 let blob = BlobHeapCfg::new()
1139 .with_max_bytes(cfg.max_bytes)
1140 .with_max_blob(cfg.max_blob);
1141 let mut pools = StreamPools {
1142 text_scratch,
1143 text_index: BlobHeapBuilder::new(blob),
1144 vec_scratch,
1145 vec_count: 0,
1146 };
1147 let bm25_policy = if plan.bm25_reindex {
1148 Bm25Policy::Reindex
1149 } else {
1150 Bm25Policy::Compact
1151 };
1152 let (m, vec_map, purged) = self.rebuild_parts(&mut pools, bm25_policy, &plan.layout)?;
1153
1154 let StreamPools {
1157 text_scratch,
1158 text_index,
1159 vec_scratch,
1160 ..
1161 } = pools;
1162 let mut text_index_bytes = Vec::new();
1163 text_index.dump_index(&mut text_index_bytes);
1164 let text_pool = text_scratch.freeze().map_err(scratch_err)?;
1165 let vec_pool = vec_scratch.freeze().map_err(scratch_err)?;
1166 let texts = BlobHeap::load_borrowed(blob, &text_index_bytes, text_pool)?;
1167 let vecs = VecPool::from_parts_borrowed(cfg.dim, cfg.max_bytes, vec_pool)?;
1168 let (hnsw, graph_work) = self.rebuild_graph(
1169 &vec_map,
1170 &vecs,
1171 plan.hnsw_full_rebuild,
1172 plan.max_hnsw_inserts,
1173 )?;
1174
1175 let edges = plan
1179 .repack_edges
1180 .then(|| self.repack_edges(&plan.layout))
1181 .transpose()?;
1182 let sections = Sections {
1183 facts: &m.facts,
1184 fact_aux: &m.fact_aux,
1185 entities: &m.entities,
1186 by_name: &m.by_name,
1187 temporal: &m.temporal,
1188 texts: &texts,
1189 metas: &m.metas,
1190 tag_lists: &m.tag_lists,
1191 bm25: &m.bm25,
1192 tags_idx: &m.tags_idx,
1193 entity_facts: &m.entity_facts,
1194 vecs: &vecs,
1195 hnsw: &hnsw,
1196 edges_out: edges.as_ref().map_or(&self.edges_out, |e| &e.out),
1197 edges_in: edges.as_ref().map_or(&self.edges_in, |e| &e.inn),
1198 edges_hist_out: edges.as_ref().map_or(&self.edges_hist_out, |e| &e.hist_out),
1199 edges_hist_in: edges.as_ref().map_or(&self.edges_hist_in, |e| &e.hist_in),
1200 layout: ShardLayout::of_config(&self.cfg).realized(&plan.layout, edges.is_some()),
1201 };
1202 self.write_snapshot_with(§ions, created_at, sink)?;
1203 report.purged = purged;
1204 report.edges_compacted = edges.is_some();
1205 report.bytes_after = edges.as_ref().map_or(0, RepackedEdges::pool_bytes)
1206 + m.facts.pool_bytes()
1207 + m.fact_aux.pool_bytes()
1208 + m.entities.pool_bytes()
1209 + hnsw.pool_bytes()
1210 + texts.pool_bytes()
1211 + m.metas.pool_bytes()
1212 + m.tag_lists.pool_bytes()
1213 + m.bm25.pool_bytes()
1214 + m.tags_idx.pool_bytes()
1215 + m.entity_facts.pool_bytes()
1216 + m.temporal.pool_bytes()
1217 + vecs.pool_bytes();
1218 report.facts_after = m.facts.len();
1219 report.vectors_after = vecs.len();
1220 report.hnsw_indexed_after = hnsw.indexed();
1221 report.structural_compacted = true;
1222 report.bm25_compacted = matches!(bm25_policy, Bm25Policy::Compact);
1223 report.bm25_reindexed = matches!(bm25_policy, Bm25Policy::Reindex);
1224 report.hnsw_rebuilt = graph_work.rebuilt;
1225 report.hnsw_remapped = graph_work.remapped;
1226 report.hnsw_inserted = graph_work.inserted;
1227 Ok(report)
1228 }
1229
1230 fn rebuild_graph(
1241 &self,
1242 vec_map: &[u32],
1243 pool: &VecPool<'_>,
1244 full_rebuild: bool,
1245 max_hnsw_inserts: Option<usize>,
1246 ) -> Result<(HnswGraph<'static>, GraphWork), Error> {
1247 let cfg = &self.cfg;
1248 let mut graph: HnswGraph<'static> = HnswGraph::new(cfg.hnsw_m, cfg.hnsw_m0, cfg.max_bytes)?;
1249 let total = pool.len() as u32;
1250 if cfg.dim == 0 || (total as usize) < cfg.flat_to_hnsw {
1251 return Ok((graph, GraphWork::default()));
1252 }
1253 let old_indexed = self.hnsw.indexed() as usize;
1254 let dead = vec_map[..old_indexed]
1255 .iter()
1256 .filter(|&&m| m == NONE_U32)
1257 .count();
1258 let mut scratch = HnswScratch::default();
1259 let mut work = GraphWork::default();
1260 if old_indexed > 0 && !full_rebuild {
1261 graph = self.hnsw.remapped(vec_map, pool, cfg.max_bytes)?;
1262 work.remapped = true;
1263 } else if old_indexed > 0 || total > 0 {
1264 work.rebuilt = true;
1265 }
1266 if full_rebuild && dead * 10 > old_indexed {
1267 work.rebuilt = true;
1268 }
1269 let start = graph.indexed();
1270 let target = hnsw_target(start, total, max_hnsw_inserts);
1271 graph.insert_bulk(pool, target, cfg.hnsw_ef_construction, &mut scratch)?;
1272 work.inserted = target.saturating_sub(start);
1273 Ok((graph, work))
1274 }
1275
1276 fn install(&mut self, r: Rebuilt) {
1279 self.facts = r.facts;
1280 self.entities = r.entities;
1281 self.by_name = r.by_name;
1282 self.fact_aux = r.fact_aux;
1283 self.texts = r.texts;
1284 self.metas = r.metas;
1285 self.tag_lists = r.tag_lists;
1286 self.bm25 = r.bm25;
1287 self.tags_idx = r.tags_idx;
1288 self.entity_facts = r.entity_facts;
1289 self.temporal = r.temporal;
1290 self.vecs = r.vecs;
1291 self.hnsw = r.hnsw;
1292 if let Some(edges) = r.edges {
1293 self.edges_out = edges.out;
1294 self.edges_in = edges.inn;
1295 self.edges_hist_out = edges.hist_out;
1296 self.edges_hist_in = edges.hist_in;
1297 }
1298 self.tombstones = 0;
1299 self.bm25_tokenizer_version = r.bm25_tokenizer_version;
1300 r.layout.apply(&mut self.cfg);
1303 }
1304}