fix: correct output of prewarm

Signed-off-by: usamoi <[email protected]>
tensorchord · Jan 22, 2025 · 1103caf · 1103caf
1 parent 345e3bf
commit 1103caf
Showing 1 changed file with 37 additions and 20 deletions.
diff --git a/src/algorithm/prewarm.rs b/src/algorithm/prewarm.rs
@@ -21,26 +21,29 @@ pub fn prewarm<O: Operator>(relation: impl RelationRead + Clone, height: i32) ->
     }
     type State = Vec<u32>;
     let mut state: State = {
-        let mut results = Vec::new();
-        let counter = 1_usize;
+        let mut nodes = Vec::new();
         {
             vectors::vector_access_1::<O, _>(relation.clone(), root_mean, ());
-            results.push(root_first);
+            nodes.push(root_first);
         }
-        writeln!(message, "number of tuples: {}", results.len()).unwrap();
-        writeln!(message, "number of pages: {}", counter).unwrap();
-        results
+        writeln!(message, "------------------------").unwrap();
+        writeln!(message, "number of nodes: {}", nodes.len()).unwrap();
+        writeln!(message, "number of tuples: {}", 1).unwrap();
+        writeln!(message, "number of pages: {}", 1).unwrap();
+        nodes
     };
     let mut step = |state: State| {
-        let mut counter = 0_usize;
-        let mut results = Vec::new();
+        let mut counter_pages = 0_usize;
+        let mut counter_tuples = 0_usize;
+        let mut nodes = Vec::new();
         for list in state {
             let mut current = list;
             while current != u32::MAX {
-                counter += 1;
+                counter_pages += 1;
                 pgrx::check_for_interrupts!();
                 let h1_guard = relation.read(current);
                 for i in 1..=h1_guard.len() {
+                    counter_tuples += 1;
                     let h1_tuple = h1_guard
                         .get(i)
                         .expect("data corruption")
@@ -51,7 +54,7 @@ pub fn prewarm<O: Operator>(relation: impl RelationRead + Clone, height: i32) ->
                                 vectors::vector_access_1::<O, _>(relation.clone(), mean, ());
                             }
                             for first in h1_tuple.first().iter().copied() {
-                                results.push(first);
+                                nodes.push(first);
                             }
                         }
                         H1TupleReader::_1(_) => (),
@@ -60,16 +63,19 @@ pub fn prewarm<O: Operator>(relation: impl RelationRead + Clone, height: i32) ->
                 current = h1_guard.get_opaque().next;
             }
         }
-        writeln!(message, "number of tuples: {}", results.len()).unwrap();
-        writeln!(message, "number of pages: {}", counter).unwrap();
-        results
+        writeln!(message, "------------------------").unwrap();
+        writeln!(message, "number of nodes: {}", nodes.len()).unwrap();
+        writeln!(message, "number of tuples: {}", counter_tuples).unwrap();
+        writeln!(message, "number of pages: {}", counter_pages).unwrap();
+        nodes
     };
     for _ in (std::cmp::max(1, prewarm_max_height)..height_of_root).rev() {
         state = step(state);
     }
     if prewarm_max_height == 0 {
-        let mut counter = 0_usize;
-        let mut results = Vec::new();
+        let mut counter_pages = 0_usize;
+        let mut counter_tuples = 0_usize;
+        let mut counter_nodes = 0_usize;
         for list in state {
             let jump_guard = relation.read(list);
             let jump_tuple = jump_guard
@@ -79,21 +85,32 @@ pub fn prewarm<O: Operator>(relation: impl RelationRead + Clone, height: i32) ->
             let first = jump_tuple.first();
             let mut current = first;
             while current != u32::MAX {
-                counter += 1;
+                counter_pages += 1;
                 pgrx::check_for_interrupts!();
                 let h0_guard = relation.read(current);
                 for i in 1..=h0_guard.len() {
-                    let _h0_tuple = h0_guard
+                    counter_tuples += 1;
+                    let h0_tuple = h0_guard
                         .get(i)
                         .expect("data corruption")
                         .pipe(read_tuple::<H0Tuple>);
-                    results.push(());
+                    match h0_tuple {
+                        H0TupleReader::_0(_h0_tuple) => {
+                            counter_nodes += 1;
+                        }
+                        H0TupleReader::_1(_h0_tuple) => {
+                            counter_nodes += 32;
+                        }
+                        H0TupleReader::_2(_) => (),
+                    }
                 }
                 current = h0_guard.get_opaque().next;
             }
         }
-        writeln!(message, "number of tuples: {}", results.len()).unwrap();
-        writeln!(message, "number of pages: {}", counter).unwrap();
+        writeln!(message, "------------------------").unwrap();
+        writeln!(message, "number of nodes: {}", counter_nodes).unwrap();
+        writeln!(message, "number of tuples: {}", counter_tuples).unwrap();
+        writeln!(message, "number of pages: {}", counter_pages).unwrap();
     }
     message
 }