From 2c18ce9848b298ca9f2ebf0b3f0542da6f6c5399 Mon Sep 17 00:00:00 2001
From: Richard Chien <stdrc@outlook.com>
Date: Thu, 21 Mar 2024 16:43:56 +0800
Subject: [PATCH] rename `VecWithKvSize` to `EstimatedVec` and impl
 `EstimateSize` for it

Signed-off-by: Richard Chien <stdrc@outlook.com>
---
 .../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 0a3acc1db310e..3e8864f4549e1 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 6d6338bee996d..6064cbfbfe3a6 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<T: EstimateSize> {
+pub struct EstimatedVec<T: EstimateSize> {
     inner: Vec<T>,
     kv_heap_size: usize,
 }
 
-impl<T: EstimateSize> Default for VecWithKvSize<T> {
+impl<T: EstimateSize> Default for EstimatedVec<T> {
     fn default() -> Self {
         Self {
             inner: vec![],
@@ -29,13 +29,15 @@ impl<T: EstimateSize> Default for VecWithKvSize<T> {
     }
 }
 
-impl<T: EstimateSize> VecWithKvSize<T> {
-    pub fn new() -> Self {
-        Default::default()
+impl<T: EstimateSize> EstimateSize for EstimatedVec<T> {
+    fn estimated_heap_size(&self) -> usize {
+        self.kv_heap_size
     }
+}
 
-    pub fn get_kv_size(&self) -> usize {
-        self.kv_heap_size
+impl<T: EstimateSize> EstimatedVec<T> {
+    pub fn new() -> Self {
+        Default::default()
     }
 
     pub fn push(&mut self, value: T) {
@@ -54,7 +56,7 @@ impl<T: EstimateSize> VecWithKvSize<T> {
     }
 }
 
-impl<T: EstimateSize> IntoIterator for VecWithKvSize<T> {
+impl<T: EstimateSize> IntoIterator for EstimatedVec<T> {
     type IntoIter = std::vec::IntoIter<Self::Item>;
     type Item = T;
 
diff --git a/src/stream/src/executor/lookup/cache.rs b/src/stream/src/executor/lookup/cache.rs
index aa3507c87e738..7e6544473366b 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<OwnedRow>) {
+    pub fn batch_update(&mut self, key: OwnedRow, value: EstimatedVec<OwnedRow>) {
         self.data.push(key, LookupEntryState::new(value));
     }
 
@@ -115,8 +115,8 @@ impl LookupEntryState {
         }
     }
 
-    fn new(value: VecWithKvSize<OwnedRow>) -> Self {
-        let kv_heap_size = value.get_kv_size();
+    fn new(value: EstimatedVec<OwnedRow>) -> 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 14347e9cf1e13..20873a8c7338f 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<S: StateStore> LookupExecutor<S> {
 
         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 {