From a365d9fdd517db895fee32ff86ac679dc767e44a Mon Sep 17 00:00:00 2001 From: Richard Chien Date: Thu, 21 Mar 2024 16:43:56 +0800 Subject: [PATCH 1/2] rename `VecWithKvSize` to `EstimatedVec` and impl `EstimateSize` for it Signed-off-by: Richard Chien --- .../estimate_size/src/collections/mod.rs | 2 +- .../estimate_size/src/collections/vec.rs | 18 ++++++++++-------- src/stream/src/executor/lookup/cache.rs | 8 ++++---- src/stream/src/executor/lookup/impl_.rs | 4 ++-- 4 files changed, 17 insertions(+), 15 deletions(-) diff --git a/src/common/estimate_size/src/collections/mod.rs b/src/common/estimate_size/src/collections/mod.rs index 0a3acc1db310..3e8864f4549e 100644 --- a/src/common/estimate_size/src/collections/mod.rs +++ b/src/common/estimate_size/src/collections/mod.rs @@ -24,7 +24,7 @@ pub use hashmap::EstimatedHashMap; pub mod btreemap; pub use btreemap::EstimatedBTreeMap; pub mod vec; -pub use vec::VecWithKvSize; +pub use vec::EstimatedVec; mod private { use super::*; diff --git a/src/common/estimate_size/src/collections/vec.rs b/src/common/estimate_size/src/collections/vec.rs index 6d6338bee996..6064cbfbfe3a 100644 --- a/src/common/estimate_size/src/collections/vec.rs +++ b/src/common/estimate_size/src/collections/vec.rs @@ -15,12 +15,12 @@ use crate::EstimateSize; #[derive(Clone)] -pub struct VecWithKvSize { +pub struct EstimatedVec { inner: Vec, kv_heap_size: usize, } -impl Default for VecWithKvSize { +impl Default for EstimatedVec { fn default() -> Self { Self { inner: vec![], @@ -29,13 +29,15 @@ impl Default for VecWithKvSize { } } -impl VecWithKvSize { - pub fn new() -> Self { - Default::default() +impl EstimateSize for EstimatedVec { + fn estimated_heap_size(&self) -> usize { + self.kv_heap_size } +} - pub fn get_kv_size(&self) -> usize { - self.kv_heap_size +impl EstimatedVec { + pub fn new() -> Self { + Default::default() } pub fn push(&mut self, value: T) { @@ -54,7 +56,7 @@ impl VecWithKvSize { } } -impl IntoIterator for VecWithKvSize { +impl IntoIterator for EstimatedVec { type IntoIter = std::vec::IntoIter; type Item = T; diff --git a/src/stream/src/executor/lookup/cache.rs b/src/stream/src/executor/lookup/cache.rs index aa3507c87e73..7e6544473366 100644 --- a/src/stream/src/executor/lookup/cache.rs +++ b/src/stream/src/executor/lookup/cache.rs @@ -16,7 +16,7 @@ use std::collections::HashSet; use risingwave_common::array::{Op, StreamChunk}; use risingwave_common::row::{OwnedRow, Row, RowExt}; -use risingwave_common_estimate_size::collections::VecWithKvSize; +use risingwave_common_estimate_size::collections::EstimatedVec; use risingwave_common_estimate_size::{EstimateSize, KvSize}; use crate::cache::{new_unbounded, ManagedLruCache}; @@ -35,7 +35,7 @@ impl LookupCache { } /// Update a key after lookup cache misses. - pub fn batch_update(&mut self, key: OwnedRow, value: VecWithKvSize) { + pub fn batch_update(&mut self, key: OwnedRow, value: EstimatedVec) { self.data.push(key, LookupEntryState::new(value)); } @@ -115,8 +115,8 @@ impl LookupEntryState { } } - fn new(value: VecWithKvSize) -> Self { - let kv_heap_size = value.get_kv_size(); + fn new(value: EstimatedVec) -> Self { + let kv_heap_size = value.estimated_heap_size(); Self { inner: HashSet::from_iter(value), kv_heap_size: KvSize::with_size(kv_heap_size), diff --git a/src/stream/src/executor/lookup/impl_.rs b/src/stream/src/executor/lookup/impl_.rs index 14347e9cf1e1..20873a8c7338 100644 --- a/src/stream/src/executor/lookup/impl_.rs +++ b/src/stream/src/executor/lookup/impl_.rs @@ -21,7 +21,7 @@ use risingwave_common::row::{OwnedRow, Row, RowExt}; use risingwave_common::util::epoch::EpochPair; use risingwave_common::util::iter_util::ZipEqDebug; use risingwave_common::util::sort_util::ColumnOrder; -use risingwave_common_estimate_size::collections::VecWithKvSize; +use risingwave_common_estimate_size::collections::EstimatedVec; use risingwave_hummock_sdk::HummockReadEpoch; use risingwave_storage::store::PrefetchOptions; use risingwave_storage::table::batch_table::storage_table::StorageTable; @@ -370,7 +370,7 @@ impl LookupExecutor { tracing::debug!(target: "events::stream::lookup::lookup_row", "{:?}", lookup_row); - let mut all_rows = VecWithKvSize::new(); + let mut all_rows = EstimatedVec::new(); // Drop the stream. { let all_data_iter = match self.arrangement.use_current_epoch { From 302f3c959fc940f44a554229d3e56230b85a900e Mon Sep 17 00:00:00 2001 From: Richard Chien Date: Thu, 21 Mar 2024 18:10:54 +0800 Subject: [PATCH 2/2] rename `kv_heap_size` to `heap_size` Signed-off-by: Richard Chien --- src/common/estimate_size/src/collections/vec.rs | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/src/common/estimate_size/src/collections/vec.rs b/src/common/estimate_size/src/collections/vec.rs index 6064cbfbfe3a..ac9f766ee1ec 100644 --- a/src/common/estimate_size/src/collections/vec.rs +++ b/src/common/estimate_size/src/collections/vec.rs @@ -17,21 +17,21 @@ use crate::EstimateSize; #[derive(Clone)] pub struct EstimatedVec { inner: Vec, - kv_heap_size: usize, + heap_size: usize, } impl Default for EstimatedVec { fn default() -> Self { Self { inner: vec![], - kv_heap_size: 0, + heap_size: 0, } } } impl EstimateSize for EstimatedVec { fn estimated_heap_size(&self) -> usize { - self.kv_heap_size + self.heap_size } } @@ -41,9 +41,7 @@ impl EstimatedVec { } pub fn push(&mut self, value: T) { - self.kv_heap_size = self - .kv_heap_size - .saturating_add(value.estimated_heap_size()); + self.heap_size = self.heap_size.saturating_add(value.estimated_heap_size()); self.inner.push(value); }