From 7700a167f22dabb642636425da10ba2b4edb96b8 Mon Sep 17 00:00:00 2001
From: evenyag <realevenyag@gmail.com>
Date: Tue, 14 Jun 2022 16:53:12 +0800
Subject: [PATCH] chores: Address CR comment

---
 src/datatypes/src/vectors/mutable.rs |   6 --
 src/storage/src/memtable.rs          |   2 +-
 src/storage/src/memtable/btree.rs    | 116 ++++++++++++++-------------
 3 files changed, 63 insertions(+), 61 deletions(-)

diff --git a/src/datatypes/src/vectors/mutable.rs b/src/datatypes/src/vectors/mutable.rs
index 39a7bcc370..ff9ed18a38 100644
--- a/src/datatypes/src/vectors/mutable.rs
+++ b/src/datatypes/src/vectors/mutable.rs
@@ -15,11 +15,5 @@ pub trait MutableVector: Send + Sync {
 
     fn as_mut_any(&mut self) -> &mut dyn Any;
 
-    // /// Push a value into the mutable vector.
-    // ///
-    // /// # Panics
-    // /// Panics if the data type of the value differs from the mutable vector's data type.
-    // fn push_value(&mut self, value: &Value);
-
     fn to_vector(&mut self) -> VectorRef;
 }
diff --git a/src/storage/src/memtable.rs b/src/storage/src/memtable.rs
index ea0e10c4df..d828398adf 100644
--- a/src/storage/src/memtable.rs
+++ b/src/storage/src/memtable.rs
@@ -26,7 +26,7 @@ pub trait Memtable: Send + Sync {
     /// Panics if the schema of key/value differs from memtable's schema.
     fn write(&self, kvs: &KeyValues) -> Result<()>;
 
-    /// Iterators the memtable.
+    /// Iterates the memtable.
     // TODO(yingwen): Consider passing a projector (does column projection).
     fn iter(&self, ctx: IterContext) -> Result<BatchIteratorPtr>;
 
diff --git a/src/storage/src/memtable/btree.rs b/src/storage/src/memtable/btree.rs
index bd13dcf724..4d3464a4d8 100644
--- a/src/storage/src/memtable/btree.rs
+++ b/src/storage/src/memtable/btree.rs
@@ -119,64 +119,12 @@ impl BTreeIterator {
         self.last_key = keys.last().map(|k| (*k).clone());
 
         Some(Batch {
-            keys: Self::keys_to_vectors(&keys),
+            keys: rows_to_vectors(keys.as_slice()),
             sequences: sequences.finish(),
             value_types: value_types.finish(),
-            values: Self::values_to_vectors(&values),
+            values: rows_to_vectors(values.as_slice()),
         })
     }
-
-    // Assumes column num of all row key is equal.
-    fn keys_to_vectors(keys: &[&InnerKey]) -> Vec<VectorRef> {
-        if keys.is_empty() {
-            return Vec::new();
-        }
-
-        let column_num = keys[0].row_key.len();
-        let row_num = keys.len();
-        let mut builders = Vec::with_capacity(column_num);
-        for v in &keys[0].row_key {
-            builders.push(VectorBuilder::with_capacity(v.data_type(), row_num));
-        }
-
-        let mut vectors = Vec::with_capacity(column_num);
-        for (col_idx, builder) in builders.iter_mut().enumerate() {
-            for row_key in keys {
-                let value = &row_key.row_key[col_idx];
-                builder.push(value);
-            }
-
-            vectors.push(builder.finish());
-        }
-
-        vectors
-    }
-
-    // Assumes column num of all row value is equal.
-    fn values_to_vectors(values: &[&RowValue]) -> Vec<VectorRef> {
-        if values.is_empty() {
-            return Vec::new();
-        }
-
-        let column_num = values[0].values.len();
-        let row_num = values.len();
-        let mut builders = Vec::with_capacity(column_num);
-        for v in &values[0].values {
-            builders.push(VectorBuilder::with_capacity(v.data_type(), row_num));
-        }
-
-        let mut vectors = Vec::with_capacity(column_num);
-        for (col_idx, builder) in builders.iter_mut().enumerate() {
-            for row_value in values {
-                let value = &row_value.values[col_idx];
-                builder.push(value);
-            }
-
-            vectors.push(builder.finish());
-        }
-
-        vectors
-    }
 }
 
 /// `MapIterWrapper` removes same user key with elder sequence.
@@ -317,3 +265,63 @@ impl InnerKey {
 struct RowValue {
     values: Vec<Value>,
 }
+
+trait RowsProvider {
+    fn row_num(&self) -> usize;
+
+    fn column_num(&self) -> usize {
+        self.row_by_index(0).len()
+    }
+
+    fn is_empty(&self) -> bool {
+        self.row_num() == 0
+    }
+
+    fn row_by_index(&self, idx: usize) -> &Vec<Value>;
+}
+
+impl<'a> RowsProvider for &'a [&InnerKey] {
+    fn row_num(&self) -> usize {
+        self.len()
+    }
+
+    fn row_by_index(&self, idx: usize) -> &Vec<Value> {
+        &self[idx].row_key
+    }
+}
+
+impl<'a> RowsProvider for &'a [&RowValue] {
+    fn row_num(&self) -> usize {
+        self.len()
+    }
+
+    fn row_by_index(&self, idx: usize) -> &Vec<Value> {
+        &self[idx].values
+    }
+}
+
+fn rows_to_vectors<T: RowsProvider>(provider: T) -> Vec<VectorRef> {
+    if provider.is_empty() {
+        return Vec::new();
+    }
+
+    let column_num = provider.column_num();
+    let row_num = provider.row_num();
+    let mut builders = Vec::with_capacity(column_num);
+    for v in provider.row_by_index(0) {
+        builders.push(VectorBuilder::with_capacity(v.data_type(), row_num));
+    }
+
+    let mut vectors = Vec::with_capacity(column_num);
+    for (col_idx, builder) in builders.iter_mut().enumerate() {
+        for row_idx in 0..row_num {
+            let row = provider.row_by_index(row_idx);
+            let value = &row[col_idx];
+            builder.push(value);
+        }
+
+        vectors.push(builder.finish());
+    }
+
+    vectors
+}