From 95ab15c4dfa3b4b15086b4692812589b1b50bfc8 Mon Sep 17 00:00:00 2001
From: Renjie Liu <liurenjie2008@gmail.com>
Date: Sun, 9 Apr 2023 17:30:45 +0800
Subject: [PATCH] feat: Estimate size for data chunk (#9050)

---
 src/common/src/array/column.rs       |  7 +++++++
 src/common/src/array/data_chunk.rs   | 11 +++++++++++
 src/common/src/array/stream_chunk.rs |  8 ++++++++
 src/common/src/array/vis.rs          | 10 ++++++++++
 4 files changed, 36 insertions(+)

diff --git a/src/common/src/array/column.rs b/src/common/src/array/column.rs
index d7640879a957c..0ca785cd88b2f 100644
--- a/src/common/src/array/column.rs
+++ b/src/common/src/array/column.rs
@@ -19,6 +19,7 @@ use risingwave_pb::data::PbColumn;
 
 use super::{Array, ArrayError, ArrayResult, I64Array};
 use crate::array::{ArrayImpl, ArrayRef};
+use crate::collection::estimate_size::EstimateSize;
 
 /// A [`Column`] consists of its logical data type
 /// and its corresponding physical array implementation,
@@ -119,6 +120,12 @@ impl From<ArrayImpl> for Column {
     }
 }
 
+impl EstimateSize for Column {
+    fn estimated_heap_size(&self) -> usize {
+        self.array.estimated_heap_size()
+    }
+}
+
 #[cfg(test)]
 mod tests {
 
diff --git a/src/common/src/array/data_chunk.rs b/src/common/src/array/data_chunk.rs
index d0f2ea5b80b8d..943d2a349a3e8 100644
--- a/src/common/src/array/data_chunk.rs
+++ b/src/common/src/array/data_chunk.rs
@@ -25,6 +25,7 @@ use crate::array::column::Column;
 use crate::array::data_chunk_iter::RowRef;
 use crate::array::ArrayBuilderImpl;
 use crate::buffer::{Bitmap, BitmapBuilder};
+use crate::collection::estimate_size::EstimateSize;
 use crate::hash::HashCode;
 use crate::row::Row;
 use crate::types::struct_type::StructType;
@@ -527,6 +528,16 @@ impl From<StructArray> for DataChunk {
     }
 }
 
+impl EstimateSize for DataChunk {
+    fn estimated_heap_size(&self) -> usize {
+        self.columns
+            .iter()
+            .map(Column::estimated_heap_size)
+            .sum::<usize>()
+            + self.vis2.estimated_heap_size()
+    }
+}
+
 /// Test utilities for [`DataChunk`].
 pub trait DataChunkTestExt {
     /// Parse a chunk from string.
diff --git a/src/common/src/array/stream_chunk.rs b/src/common/src/array/stream_chunk.rs
index c9d0cca6c91e0..12271552c7ef7 100644
--- a/src/common/src/array/stream_chunk.rs
+++ b/src/common/src/array/stream_chunk.rs
@@ -13,6 +13,7 @@
 // limitations under the License.
 
 use std::fmt;
+use std::mem::size_of;
 
 use itertools::Itertools;
 use risingwave_pb::data::{PbOp, PbStreamChunk};
@@ -21,6 +22,7 @@ use super::{ArrayResult, DataChunkTestExt};
 use crate::array::column::Column;
 use crate::array::{DataChunk, Vis};
 use crate::buffer::Bitmap;
+use crate::collection::estimate_size::EstimateSize;
 use crate::row::{OwnedRow, Row};
 use crate::types::to_text::ToText;
 use crate::types::DataType;
@@ -291,6 +293,12 @@ impl fmt::Debug for StreamChunk {
     }
 }
 
+impl EstimateSize for StreamChunk {
+    fn estimated_heap_size(&self) -> usize {
+        self.data.estimated_heap_size() + self.ops.capacity() * size_of::<Op>()
+    }
+}
+
 /// Test utilities for [`StreamChunk`].
 pub trait StreamChunkTestExt: Sized {
     fn from_pretty(s: &str) -> Self;
diff --git a/src/common/src/array/vis.rs b/src/common/src/array/vis.rs
index 9d65e92c7da06..6040f5344977e 100644
--- a/src/common/src/array/vis.rs
+++ b/src/common/src/array/vis.rs
@@ -16,6 +16,7 @@ use auto_enums::auto_enum;
 use itertools::repeat_n;
 
 use crate::buffer::{Bitmap, BitmapBuilder};
+use crate::collection::estimate_size::EstimateSize;
 
 /// `Vis` is a visibility bitmap of rows.
 #[derive(Clone, PartialEq, Debug)]
@@ -97,6 +98,15 @@ impl Vis {
     }
 }
 
+impl EstimateSize for Vis {
+    fn estimated_heap_size(&self) -> usize {
+        match self {
+            Vis::Bitmap(bitmap) => bitmap.estimated_heap_size(),
+            Vis::Compact(_) => 0,
+        }
+    }
+}
+
 impl<'a, 'b> std::ops::BitAnd<&'b Vis> for &'a Vis {
     type Output = Vis;