feat: show index size in index info table (#222)

* feat: show index size in index info table

Signed-off-by: Keming <kemingyang@tensorchord.ai>

* fix clippy

Signed-off-by: Keming <kemingyang@tensorchord.ai>

* compute all the file for sealed segments

Signed-off-by: Keming <kemingyang@tensorchord.ai>

* fix dir_size

Signed-off-by: Keming <kemingyang@tensorchord.ai>

* display the sum size in the table

Signed-off-by: Keming <kemingyang@tensorchord.ai>

---------

Signed-off-by: Keming <kemingyang@tensorchord.ai>

Author: kemingy

crates/service/src/index/mod.rs

@@ -62,13 +62,24 @@ pub struct IndexOptions {
     pub indexing: IndexingOptions,
 }
 
+#[derive(Debug, Serialize, Deserialize)]
+#[serde(deny_unknown_fields)]
+pub struct SegmentSizeInfo {
+    pub id: Uuid,
+    #[serde(rename = "type")]
+    pub typ: String,
+    pub length: usize,
+    pub size: u64,
+}
+
 #[derive(Debug, Serialize, Deserialize)]
 pub struct IndexStat {
     pub indexing: bool,
     pub sealed: Vec<u32>,
     pub growing: Vec<u32>,
     pub write: u32,
     pub options: IndexOptions,
+    pub sizes: Vec<SegmentSizeInfo>,
 }
 
 pub struct Index<S: G> {
@@ -239,6 +250,12 @@ impl<S: G> Index<S> {
             growing: view.growing.values().map(|x| x.len()).collect(),
             write: view.write.as_ref().map(|(_, x)| x.len()).unwrap_or(0),
             options: self.options().clone(),
+            sizes: view
+                .sealed
+                .values()
+                .map(|x| x.size())
+                .chain(view.growing.values().map(|x| x.size()))
+                .collect(),
         }
     }
 }

crates/service/src/index/segments/growing.rs

@@ -3,6 +3,7 @@
 use super::SegmentTracker;
 use crate::index::IndexOptions;
 use crate::index::IndexTracker;
+use crate::index::SegmentSizeInfo;
 use crate::prelude::*;
 use crate::utils::dir_ops::sync_dir;
 use crate::utils::file_wal::FileWal;
@@ -140,6 +141,14 @@ impl<S: G> GrowingSegment<S> {
     pub fn len(&self) -> u32 {
         self.len.load(Ordering::Acquire) as u32
     }
+    pub fn size(&self) -> SegmentSizeInfo {
+        SegmentSizeInfo {
+            id: self.uuid,
+            typ: "growing".to_string(),
+            length: self.len() as usize,
+            size: (self.len() as u64) * (std::mem::size_of::<Log<S>>() as u64),
+        }
+    }
     pub fn vector(&self, i: u32) -> &[S::Scalar] {
         let i = i as usize;
         if i >= self.len.load(Ordering::Acquire) {

crates/service/src/index/segments/sealed.rs

@@ -1,9 +1,9 @@
 use super::growing::GrowingSegment;
 use super::SegmentTracker;
 use crate::index::indexing::{DynamicIndexIter, DynamicIndexing};
-use crate::index::{IndexOptions, IndexTracker};
+use crate::index::{IndexOptions, IndexTracker, SegmentSizeInfo};
 use crate::prelude::*;
-use crate::utils::dir_ops::sync_dir;
+use crate::utils::dir_ops::{dir_size, sync_dir};
 use serde::{Deserialize, Serialize};
 use std::path::PathBuf;
 use std::sync::Arc;
@@ -58,6 +58,22 @@ impl<S: G> SealedSegment<S> {
     pub fn len(&self) -> u32 {
         self.indexing.len()
     }
+    pub fn size(&self) -> SegmentSizeInfo {
+        let mut info = SegmentSizeInfo {
+            id: self.uuid,
+            typ: "sealed".to_string(),
+            length: self.len() as usize,
+            size: 0,
+        };
+        let path = self._tracker.path.join("indexing");
+        match dir_size(&path) {
+            Ok(size) => info.size = size as u64,
+            Err(e) => {
+                panic!("Failed to get size of {:?}: {}", path, e);
+            }
+        }
+        info
+    }
     pub fn vector(&self, i: u32) -> &[S::Scalar] {
         self.indexing.vector(i)
     }

crates/service/src/utils/dir_ops.rs

@@ -1,7 +1,29 @@
-use std::fs::File;
+use std::fs::{read_dir, File};
+use std::io;
 use std::path::Path;
 
 pub fn sync_dir(path: impl AsRef<Path>) {
     let file = File::open(path).expect("Failed to sync dir.");
     file.sync_all().expect("Failed to sync dir.");
 }
+
+pub fn dir_size(dir: &Path) -> io::Result<usize> {
+    let mut size = 0;
+    if dir.is_dir() {
+        for entry in read_dir(dir)? {
+            let entry = entry?;
+            let path = entry.path();
+            let name = path.file_name().unwrap().to_string_lossy();
+            if name.starts_with('.') {
+                // ignore hidden files
+                continue;
+            }
+            if path.is_dir() {
+                size += dir_size(&path)?;
+            } else {
+                size += entry.metadata()?.len() as usize;
+            }
+        }
+    }
+    Ok(size)
+}

docs/indexing.md

@@ -119,6 +119,7 @@ We also provide a view `pg_vector_index_info` to monitor the progress of indexin
 | idx_sealed   | int8[] | The number of tuples in each sealed segment.  |
 | idx_growing  | int8[] | The number of tuples in each growing segment. |
 | idx_write    | int8   | The number of tuples in write buffer.         |
+| idx_size     | int8   | The byte size for all the segments.           |
 | idx_config   | text   | The configuration of the index.               |
 
 ## Examples

src/index/views.rs

@@ -9,6 +9,7 @@ CREATE TYPE VectorIndexStat AS (
     idx_sealed BIGINT[],
     idx_growing BIGINT[],
     idx_write BIGINT,
+    idx_size BIGINT,
     idx_options TEXT
 );",
     name = "create_composites",
@@ -40,6 +41,11 @@ fn vector_stat(oid: pgrx::pg_sys::Oid) -> pgrx::composite_type!("VectorIndexStat
     })
     .unwrap();
     res.set_by_name("idx_write", stat.write as i64).unwrap();
+    res.set_by_name(
+        "idx_size",
+        stat.sizes.iter().map(|x| x.size as i64).sum::<i64>(),
+    )
+    .unwrap();
     res.set_by_name("idx_options", serde_json::to_string(&stat.options))
         .unwrap();
     res