perf(hashkey): refine the memory reserving for SerializedKey (#9556)

Co-authored-by: stonepage <[email protected]>

src/common/src/array/data_chunk.rs

@@ -431,6 +431,20 @@ impl DataChunk {
         DataChunk::new(columns, indexes.len())
     }
 
+    fn partition_sizes_for_columns(&self, col_indices: &[usize]) -> (usize, Vec<&Column>) {
+        let mut col_variable: Vec<&Column> = vec![];
+        let mut row_len_fixed: usize = 0;
+        for i in col_indices {
+            let col = &self.columns[*i];
+            if let Some(field_len) = try_get_exact_serialize_datum_size(&col.array()) {
+                row_len_fixed += field_len;
+            } else {
+                col_variable.push(col);
+            }
+        }
+        (row_len_fixed, col_variable)
+    }
+
     /// Partition fixed size datums and variable length ones.
     /// ---
     /// In some cases, we have fixed size for the entire column,
@@ -472,6 +486,35 @@ impl DataChunk {
         )
     }
 
+    pub fn compute_key_sizes_by_columns(&self, column_indices: &[usize]) -> Vec<usize> {
+        let (row_len_fixed, col_variable) = self.partition_sizes_for_columns(column_indices);
+        let mut sizes: Vec<usize> = Vec::with_capacity(self.capacity());
+        let update_sizes = |sizes: &mut Vec<usize>, col_variable, i| unsafe {
+            sizes.push(row_len_fixed + Self::compute_size_of_variable_cols_in_row(col_variable, i))
+        };
+        match &self.vis2 {
+            Vis::Bitmap(vis) => {
+                let rows_num = vis.len();
+                for i in 0..rows_num {
+                    // SAFETY(value_at_unchecked): the idx is always in bound.
+                    unsafe {
+                        if vis.is_set_unchecked(i) {
+                            update_sizes(&mut sizes, &col_variable, i);
+                        } else {
+                            sizes.push(0)
+                        }
+                    }
+                }
+            }
+            Vis::Compact(rows_num) => {
+                for i in 0..*rows_num {
+                    update_sizes(&mut sizes, &col_variable, i);
+                }
+            }
+        }
+        sizes
+    }
+
     /// Serialize each row into value encoding bytes.
     ///
     /// The returned vector's size is `self.capacity()` and for the invisible row will give a empty

src/common/src/hash/key.rs

@@ -911,12 +911,15 @@ impl<B: NullBitmap> HashKey for SerializedKey<B> {
         data_chunk: &DataChunk,
         hash_codes: Vec<XxHash64HashCode>,
     ) -> Vec<Self> {
-        let estimated_key_size = data_chunk.estimate_value_encoding_size(column_idxes);
+        let estimated_key_sizes = data_chunk.compute_key_sizes_by_columns(column_idxes);
 
         // Construct serializers for each row.
         let mut serializers: Vec<Self::S> = hash_codes
             .into_iter()
-            .map(|hashcode| Self::S::from_hash_code(hashcode, estimated_key_size))
+            .zip_eq_fast(estimated_key_sizes)
+            .map(|(hashcode, estimated_key_size)| {
+                Self::S::from_hash_code(hashcode, estimated_key_size)
+            })
             .collect();
 
         for column_idx in column_idxes {