feat: Estimate size for data chunk (#9050)

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 {
 

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.

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;

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;