From 48801c60d365b31c1119a4edf8efaca03f78d170 Mon Sep 17 00:00:00 2001
From: evenyag <realevenyag@gmail.com>
Date: Thu, 22 Feb 2024 21:32:01 +0800
Subject: [PATCH] chore: add methods to get pk weights

---
 src/mito2/src/memtable/merge_tree/dict.rs | 29 ++++++++++++++++-------
 1 file changed, 20 insertions(+), 9 deletions(-)

diff --git a/src/mito2/src/memtable/merge_tree/dict.rs b/src/mito2/src/memtable/merge_tree/dict.rs
index b17dd7798910..5c1c3c3a57f6 100644
--- a/src/mito2/src/memtable/merge_tree/dict.rs
+++ b/src/mito2/src/memtable/merge_tree/dict.rs
@@ -181,15 +181,30 @@ impl DictBuilderReader {
         self.key_by_pk_index(pk_index)
     }
 
+    /// Gets the key by the pk index.
+    pub fn key_by_pk_index(&self, pk_index: PkIndex) -> &[u8] {
+        let block_idx = pk_index / MAX_KEYS_PER_BLOCK;
+        self.blocks[block_idx as usize].key_by_pk_index(pk_index)
+    }
+
+    /// Returns pk weights to sort a data part and replaces pk indices.
+    pub(crate) fn pk_weights_to_sort_data(&self) -> Vec<u16> {
+        compute_pk_weights(&self.sorted_pk_indices)
+    }
+
     /// Returns pk indices sorted by keys.
     pub(crate) fn sorted_pk_index(&self) -> &[PkIndex] {
         &self.sorted_pk_indices
     }
+}
 
-    fn key_by_pk_index(&self, pk_index: PkIndex) -> &[u8] {
-        let block_idx = pk_index / MAX_KEYS_PER_BLOCK;
-        self.blocks[block_idx as usize].key_by_pk_index(pk_index)
+/// Returns pk weights to sort a data part and replaces pk indices.
+fn compute_pk_weights(sorted_pk_indices: &[PkIndex]) -> Vec<u16> {
+    let mut pk_weights = vec![0; sorted_pk_indices.len()];
+    for (weight, pk_index) in sorted_pk_indices.iter().enumerate() {
+        pk_weights[*pk_index as usize] = weight as u16;
     }
+    pk_weights
 }
 
 /// A key dictionary.
@@ -223,13 +238,9 @@ impl KeyDict {
         self.pk_to_index.get(key).copied()
     }
 
-    /// Sets the pk weights to sort a data part and replaces pk indices.
+    /// Returns pk weights to sort a data part and replaces pk indices.
     pub(crate) fn pk_weights_to_sort_data(&self) -> Vec<u16> {
-        let mut pk_weights = vec![0; self.key_positions.len()];
-        for (weight, pk_index) in self.key_positions.iter().enumerate() {
-            pk_weights[*pk_index as usize] = weight as u16;
-        }
-        pk_weights
+        compute_pk_weights(&self.key_positions)
     }
 
     /// Returns the shared memory size.