Skip to content

Commit

Permalink
feat: Estimate size for data chunk (risingwavelabs#9050)
Browse files Browse the repository at this point in the history
liurenjie1024 authored Apr 9, 2023

Verified

This commit was created on GitHub.com and signed with GitHub’s verified signature. The key has expired.
1 parent ced72de commit 95ab15c
Showing 4 changed files with 36 additions and 0 deletions.
7 changes: 7 additions & 0 deletions src/common/src/array/column.rs
Original file line number Diff line number Diff line change
@@ -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 {

11 changes: 11 additions & 0 deletions src/common/src/array/data_chunk.rs
Original file line number Diff line number Diff line change
@@ -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.
8 changes: 8 additions & 0 deletions src/common/src/array/stream_chunk.rs
Original file line number Diff line number Diff line change
@@ -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;
10 changes: 10 additions & 0 deletions src/common/src/array/vis.rs
Original file line number Diff line number Diff line change
@@ -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;

0 comments on commit 95ab15c

Please sign in to comment.