From 88e8f37766f552dcafbd683ec1e529b73aee9ea5 Mon Sep 17 00:00:00 2001
From: evenyag <realevenyag@gmail.com>
Date: Mon, 25 Sep 2023 17:06:25 +0800
Subject: [PATCH] feat: consider vec size

---
 src/mito2/src/cache/cache_size.rs | 18 ++++++++++++++----
 1 file changed, 14 insertions(+), 4 deletions(-)

diff --git a/src/mito2/src/cache/cache_size.rs b/src/mito2/src/cache/cache_size.rs
index 0dbd1ef5524c..1620f8203ba1 100644
--- a/src/mito2/src/cache/cache_size.rs
+++ b/src/mito2/src/cache/cache_size.rs
@@ -107,16 +107,26 @@ fn row_group_meta_heap_size(meta: &RowGroupMetaData) -> usize {
 
 /// Returns estimated size of [ParquetColumnIndex] allocated from heap.
 fn parquet_column_index_heap_size(column_index: &ParquetColumnIndex) -> usize {
-    column_index.iter().map(|index| index.len()).sum::<usize>() * mem::size_of::<Index>()
+    column_index
+        .iter()
+        .map(|row_group| row_group.len() * mem::size_of::<Index>() + mem::size_of_val(row_group))
+        .sum()
 }
 
 /// Returns estimated size of [ParquetOffsetIndex] allocated from heap.
 fn parquet_offset_index_heap_size(offset_index: &ParquetOffsetIndex) -> usize {
     offset_index
         .iter()
-        .map(|index| index.iter().map(|index| index.len()).sum::<usize>())
-        .sum::<usize>()
-        * mem::size_of::<PageLocation>()
+        .map(|row_group| {
+            row_group
+                .iter()
+                .map(|column| {
+                    column.len() * mem::size_of::<PageLocation>() + mem::size_of_val(column)
+                })
+                .sum::<usize>()
+                + mem::size_of_val(row_group)
+        })
+        .sum()
 }
 
 #[cfg(test)]