nodedb_vector/collection/
stats.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Live statistics aggregation for `VectorCollection`.
4
5use crate::index_config::IndexType;
6
7use super::lifecycle::VectorCollection;
8
9impl VectorCollection {
10    /// Collect live statistics from all segments.
11    pub fn stats(&self) -> nodedb_types::VectorIndexStats {
12        let growing_vectors = self.growing.len();
13        let sealed_vectors: usize = self.sealed.iter().map(|s| s.index.len()).sum();
14        let building_vectors: usize = self.building.iter().map(|s| s.flat.len()).sum();
15
16        let ivf_vectors = self.ivf.as_ref().map_or(0, |ivf| ivf.len());
17        let tombstone_count: usize = self.ivf.as_ref().map_or(0, |ivf| ivf.tombstone_count())
18            + self
19                .sealed
20                .iter()
21                .map(|s| s.index.tombstone_count())
22                .sum::<usize>()
23            + self.growing.tombstone_count()
24            + self
25                .building
26                .iter()
27                .map(|s| s.flat.tombstone_count())
28                .sum::<usize>();
29
30        let total = growing_vectors + sealed_vectors + building_vectors + ivf_vectors;
31        let tombstone_ratio = if total > 0 {
32            tombstone_count as f64 / total as f64
33        } else {
34            0.0
35        };
36
37        let quantization = if let Some(ref dispatch) = self.codec_dispatch {
38            match dispatch.quantization() {
39                "rabitq" => nodedb_types::VectorIndexQuantization::RaBitQ,
40                "bbq" => nodedb_types::VectorIndexQuantization::Bbq,
41                _ => nodedb_types::VectorIndexQuantization::None,
42            }
43        } else if self.ivf.is_some() || self.sealed.iter().any(|s| s.pq.is_some()) {
44            nodedb_types::VectorIndexQuantization::Pq
45        } else if self.sealed.iter().any(|s| s.sq8.is_some()) {
46            nodedb_types::VectorIndexQuantization::Sq8
47        } else {
48            nodedb_types::VectorIndexQuantization::None
49        };
50
51        let index_type = match self.index_config.index_type {
52            IndexType::HnswPq => nodedb_types::VectorIndexType::HnswPq,
53            IndexType::IvfPq => nodedb_types::VectorIndexType::IvfPq,
54            IndexType::Hnsw => nodedb_types::VectorIndexType::Hnsw,
55        };
56
57        let hnsw_mem: usize = self
58            .sealed
59            .iter()
60            .map(|s| s.index.memory_usage_bytes())
61            .sum();
62        let sq8_mem: usize = self
63            .sealed
64            .iter()
65            .filter_map(|s| s.sq8.as_ref().map(|(_, data)| data.len()))
66            .sum();
67        let growing_mem = growing_vectors * self.dim * std::mem::size_of::<f32>();
68        let building_mem = building_vectors * self.dim * std::mem::size_of::<f32>();
69        let ivf_mem = self.ivf.as_ref().map_or(0, |ivf| ivf.memory_bytes());
70        let memory_bytes = hnsw_mem + sq8_mem + growing_mem + building_mem + ivf_mem;
71
72        let disk_bytes: usize = self
73            .sealed
74            .iter()
75            .filter_map(|s| s.mmap_vectors.as_ref().map(|m| m.file_size()))
76            .sum();
77
78        let metric_name = format!("{:?}", self.params.metric).to_lowercase();
79
80        nodedb_types::VectorIndexStats {
81            sealed_count: self.sealed.len(),
82            building_count: self.building.len(),
83            growing_vectors,
84            sealed_vectors,
85            live_count: self.live_count(),
86            tombstone_count,
87            tombstone_ratio,
88            quantization,
89            memory_bytes,
90            disk_bytes,
91            build_in_progress: !self.building.is_empty(),
92            index_type,
93            hnsw_m: self.params.m,
94            hnsw_m0: self.params.m0,
95            hnsw_ef_construction: self.params.ef_construction,
96            metric: metric_name,
97            dimensions: self.dim,
98            seal_threshold: self.seal_threshold,
99            mmap_segment_count: self.mmap_segment_count,
100            // `arena_bytes` is populated by the Data Plane handler which
101            // has access to `nodedb_mem::CollectionArenaHandle`. The field
102            // is always `None` here; callers overwrite it after calling
103            // `stats()` when a dedicated arena handle is available.
104            arena_bytes: None,
105            ivf: self.is_ivf().then(|| nodedb_types::VectorIvfStats {
106                training_threshold: self.ivf_training_threshold(),
107                trained: self.ivf.is_some(),
108                trained_on: self.ivf.as_ref().map_or(0, |ivf| ivf.trained_on()),
109                trained_at_ms: self.ivf.as_ref().map_or(0, |ivf| ivf.trained_at_ms()),
110                indexed_vectors: ivf_vectors,
111                cells: self.ivf.as_ref().map_or(0, |ivf| ivf.n_cells()),
112                nprobe: self.index_config.ivf_nprobe,
113            }),
114            // The owning core fills this from its build queue.
115            builds_queued: 0,
116            builds_completed: self.builds_completed,
117            builds_failed: self.builds_failed,
118        }
119    }
120}
nodedb_vector/collection/stats.rs

nodedb_vector/collection/
stats.rs
