From 727bdb8b8641b6ba5993245bca62fc5f900212aa Mon Sep 17 00:00:00 2001
From: evenyag <realevenyag@gmail.com>
Date: Thu, 9 Jun 2022 17:30:17 +0800
Subject: [PATCH] feat: Add sequences and value_types to Batch

---
 src/storage/src/memtable.rs        |  4 +++-
 src/storage/src/memtable/btree.rs  | 10 ++++++++--
 src/store-api/src/storage/types.rs |  6 ++++++
 3 files changed, 17 insertions(+), 3 deletions(-)

diff --git a/src/storage/src/memtable.rs b/src/storage/src/memtable.rs
index b0aff45bc66..2cb20e06c13 100644
--- a/src/storage/src/memtable.rs
+++ b/src/storage/src/memtable.rs
@@ -5,7 +5,7 @@ mod schema;
 use std::mem;
 use std::sync::Arc;
 
-use datatypes::vectors::VectorRef;
+use datatypes::vectors::{UInt64Vector, UInt8Vector, VectorRef};
 use store_api::storage::{SequenceNumber, ValueType};
 
 use crate::error::Result;
@@ -58,6 +58,8 @@ pub enum RowOrdering {
 
 pub struct Batch {
     pub keys: Vec<VectorRef>,
+    pub sequences: UInt64Vector,
+    pub value_types: UInt8Vector,
     pub values: Vec<VectorRef>,
 }
 
diff --git a/src/storage/src/memtable/btree.rs b/src/storage/src/memtable/btree.rs
index 56982c4cf86..de33f1098d5 100644
--- a/src/storage/src/memtable/btree.rs
+++ b/src/storage/src/memtable/btree.rs
@@ -3,9 +3,9 @@ use std::collections::{btree_map, BTreeMap};
 use std::ops::Bound;
 use std::sync::{Arc, RwLock};
 
+use datatypes::prelude::*;
 use datatypes::value::Value;
-use datatypes::vectors::VectorBuilder;
-use datatypes::vectors::VectorRef;
+use datatypes::vectors::{UInt64VectorBuilder, UInt8VectorBuilder, VectorBuilder};
 use store_api::storage::{SequenceNumber, ValueType};
 
 use crate::error::Result;
@@ -101,9 +101,13 @@ impl BTreeIterator {
         let iter = MapIterWrapper::new(iter);
 
         let mut keys = Vec::with_capacity(self.ctx.batch_size);
+        let mut sequences = UInt64VectorBuilder::with_capacity(self.ctx.batch_size);
+        let mut value_types = UInt8VectorBuilder::with_capacity(self.ctx.batch_size);
         let mut values = Vec::with_capacity(self.ctx.batch_size);
         for (inner_key, row_value) in iter.take(self.ctx.batch_size) {
             keys.push(inner_key);
+            sequences.push(Some(inner_key.sequence));
+            value_types.push(Some(inner_key.value_type.as_u8()));
             values.push(row_value);
         }
 
@@ -114,6 +118,8 @@ impl BTreeIterator {
 
         Some(Batch {
             keys: Self::keys_to_vectors(&keys),
+            sequences: sequences.finish(),
+            value_types: value_types.finish(),
             values: Self::values_to_vectors(&values),
         })
     }
diff --git a/src/store-api/src/storage/types.rs b/src/store-api/src/storage/types.rs
index e2679b1e6d9..361ba7caca6 100644
--- a/src/store-api/src/storage/types.rs
+++ b/src/store-api/src/storage/types.rs
@@ -10,3 +10,9 @@ pub enum ValueType {
     /// Put operation.
     Put,
 }
+
+impl ValueType {
+    pub fn as_u8(&self) -> u8 {
+        *self as u8
+    }
+}