fix(cubestore): Query Cache - correct size calculation (complex TableValue) (#6400)

Author: ovr

rust/cubestore/cubestore/src/queryplanner/providers/query_cache.rs

@@ -1,5 +1,5 @@
 use crate::queryplanner::project_schema;
-use crate::sql::cache::SqlResultCache;
+use crate::sql::cache::{sql_result_cache_sizeof, SqlResultCache};
 use arrow::array::{Array, Int64Builder, StringBuilder};
 use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
 use arrow::record_batch::RecordBatch;
@@ -11,7 +11,6 @@ use datafusion::logical_plan::Expr;
 use datafusion::physical_plan::memory::MemoryExec;
 use datafusion::physical_plan::Partitioning;
 use datafusion::physical_plan::{ExecutionPlan, SendableRecordBatchStream};
-use deepsize::DeepSizeOf;
 use std::any::Any;
 use std::fmt;
 use std::fmt::Formatter;
@@ -145,7 +144,9 @@ impl ExecutionPlan for InfoSchemaQueryCacheTableExec {
         for (k, v) in self.cache.iter() {
             builder.add_row(
                 k.get_query(),
-                v.deep_size_of().try_into().unwrap_or(i64::MAX),
+                sql_result_cache_sizeof(&k, &v)
+                    .try_into()
+                    .unwrap_or(i64::MAX),
             );
         }
 

rust/cubestore/cubestore/src/sql/cache.rs

@@ -71,7 +71,7 @@ pub struct SqlResultCache {
     result_cache: Cache<SqlResultCacheKey, Arc<DataFrame>>,
 }
 
-fn dataframe_sizeof(key: &SqlResultCacheKey, df: &Arc<DataFrame>) -> u32 {
+pub fn sql_result_cache_sizeof(key: &SqlResultCacheKey, df: &Arc<DataFrame>) -> u32 {
     (key.deep_size_of() + df.deep_size_of())
         .try_into()
         .unwrap_or(u32::MAX)
@@ -89,7 +89,7 @@ impl SqlResultCache {
             queue_cache: Mutex::new(lru::LruCache::new(1000)),
             result_cache: cache_builder
                 .max_capacity(capacity_bytes)
-                .weigher(dataframe_sizeof)
+                .weigher(sql_result_cache_sizeof)
                 .build(),
         }
     }

rust/cubestore/cubestore/src/table/mod.rs

@@ -33,10 +33,6 @@ pub enum TableValue {
 }
 
 impl DeepSizeOf for TableValue {
-    fn deep_size_of(&self) -> usize {
-        32
-    }
-
     fn deep_size_of_children(&self, context: &mut Context) -> usize {
         match self {
             TableValue::Null => 0,
@@ -229,6 +225,7 @@ pub fn cmp_same_types(l: &TableValue, r: &TableValue) -> Ordering {
 mod tests {
     use crate::table::{TableValue, TimestampValue};
     use crate::util::decimal::Decimal;
+    use deepsize::DeepSizeOf;
     use serde::{Deserialize, Serialize};
 
     #[test]
@@ -257,4 +254,16 @@ mod tests {
             assert_eq!(v, &v2);
         }
     }
+
+    #[test]
+    fn table_value_deep_size_of() {
+        for (v, expected_size) in [
+            (TableValue::Null, 32_usize),
+            (TableValue::Int(1), 32_usize),
+            (TableValue::Decimal(Decimal::new(1)), 32_usize),
+            (TableValue::String("foo".into()), 35_usize),
+        ] {
+            assert_eq!(v.deep_size_of(), expected_size, "size for {:?}", v);
+        }
+    }
 }