test: Add simple write/iter test for memtable

src/storage/src/memtable.rs

@@ -1,6 +1,8 @@
 mod btree;
 mod inserter;
 mod schema;
+#[cfg(test)]
+mod tests;
 
 use std::mem;
 use std::sync::Arc;
@@ -47,7 +49,7 @@ impl Default for IterContext {
 }
 
 /// The ordering of the iterator output.
-#[derive(Debug)]
+#[derive(Debug, PartialEq)]
 pub enum RowOrdering {
     /// The output rows are unordered.
     Unordered,

src/storage/src/memtable/btree.rs

@@ -257,6 +257,8 @@ impl<'a> IterRow<'a> {
                 .collect(),
         };
 
+        self.index += 1;
+
         (inner_key, row_value)
     }
 }
@@ -277,7 +279,7 @@ impl<'a> Iterator for IterRow<'a> {
     }
 }
 
-#[derive(Clone, PartialEq, Eq)]
+#[derive(Clone, Debug, PartialEq, Eq)]
 struct InnerKey {
     row_key: Vec<Value>,
     sequence: SequenceNumber,
@@ -309,7 +311,7 @@ impl InnerKey {
     }
 }
 
-#[derive(Clone)]
+#[derive(Clone, Debug)]
 struct RowValue {
     values: Vec<Value>,
 }

src/storage/src/memtable/schema.rs

@@ -1,6 +1,6 @@
 use crate::metadata::{ColumnMetadata, ColumnsRowKeyMetadataRef};
 
-#[derive(Clone)]
+#[derive(Clone, Debug, PartialEq)]
 pub struct MemtableSchema {
     columns_row_key: ColumnsRowKeyMetadataRef,
 }

src/storage/src/memtable/tests.rs

@@ -0,0 +1,187 @@
+use datatypes::prelude::*;
+use datatypes::type_id::LogicalTypeId;
+use datatypes::vectors::{Int64VectorBuilder, UInt64VectorBuilder};
+
+use super::*;
+use crate::metadata::RegionMetadata;
+use crate::test_util::descriptor_util::RegionDescBuilder;
+
+// Schema for testing memtable:
+// - key: Int64(timestamp), UInt64(version),
+// - value: UInt64
+fn schema_for_test() -> MemtableSchema {
+    // Just build a region desc and use its columns_row_key metadata.
+    let desc = RegionDescBuilder::new("test")
+        .push_value_column(("v1", LogicalTypeId::UInt64, true))
+        .build();
+    let metadata: RegionMetadata = desc.try_into().unwrap();
+
+    MemtableSchema::new(metadata.columns_row_key)
+}
+
+fn kvs_for_test_with_index(
+    sequence: SequenceNumber,
+    value_type: ValueType,
+    start_index_in_batch: usize,
+    keys: &[(i64, u64)],
+    values: &[Option<u64>],
+) -> KeyValues {
+    assert_eq!(keys.len(), values.len());
+
+    let mut key_builders = (
+        Int64VectorBuilder::with_capacity(keys.len()),
+        UInt64VectorBuilder::with_capacity(keys.len()),
+    );
+    for key in keys {
+        key_builders.0.push(Some(key.0));
+        key_builders.1.push(Some(key.1));
+    }
+    let keys = vec![
+        Arc::new(key_builders.0.finish()) as _,
+        Arc::new(key_builders.1.finish()) as _,
+    ];
+
+    let mut value_builder = UInt64VectorBuilder::with_capacity(values.len());
+    for value in values {
+        value_builder.push(*value);
+    }
+    let values = vec![Arc::new(value_builder.finish()) as _];
+
+    KeyValues {
+        sequence,
+        value_type,
+        start_index_in_batch,
+        keys,
+        values,
+    }
+}
+
+fn kvs_for_test(
+    sequence: SequenceNumber,
+    value_type: ValueType,
+    keys: &[(i64, u64)],
+    values: &[Option<u64>],
+) -> KeyValues {
+    kvs_for_test_with_index(sequence, value_type, 0, keys, values)
+}
+
+fn write_kvs(
+    memtable: &dyn Memtable,
+    sequence: SequenceNumber,
+    value_type: ValueType,
+    keys: &[(i64, u64)],
+    values: &[Option<u64>],
+) {
+    let kvs = kvs_for_test(sequence, value_type, keys, values);
+
+    memtable.write(&kvs).unwrap();
+}
+
+fn check_iter_content(
+    iter: &mut dyn BatchIterator,
+    keys: &[(i64, u64)],
+    sequences: &[u64],
+    value_types: &[ValueType],
+    values: &[Option<u64>],
+) {
+    let mut index = 0;
+    while let Some(batch) = iter.next().unwrap() {
+        assert_eq!(2, batch.keys.len());
+        assert_eq!(1, batch.values.len());
+        let row_num = batch.keys[0].len();
+        assert_eq!(row_num, batch.keys[1].len());
+        assert_eq!(row_num, batch.sequences.len());
+        assert_eq!(row_num, batch.value_types.len());
+        assert_eq!(row_num, batch.values[0].len());
+
+        for i in 0..row_num {
+            let (k0, k1) = (batch.keys[0].get(i), batch.keys[1].get(i));
+            let sequence = batch.sequences.get_data(i).unwrap();
+            let value_type = batch.value_types.get_data(i).unwrap();
+            let v = batch.values[0].get(i);
+
+            assert_eq!(Value::from(keys[index].0), k0);
+            assert_eq!(Value::from(keys[index].1), k1);
+            assert_eq!(sequences[index], sequence);
+            assert_eq!(value_types[index].as_u8(), value_type);
+            assert_eq!(Value::from(values[index]), v);
+
+            index += 1;
+        }
+    }
+
+    assert_eq!(keys.len(), index);
+}
+
+// TODO(yingwen): Check size of the returned batch.
+
+#[test]
+fn test_write_iter_memtable() {
+    let builder = DefaultMemtableBuilder {};
+    let schema = schema_for_test();
+    let mem = builder.build(schema.clone());
+
+    write_kvs(
+        &*mem,
+        10, // sequence
+        ValueType::Put,
+        &[
+            (1000, 1),
+            (1000, 2),
+            (1001, 1),
+            (2002, 1),
+            (2003, 1),
+            (2003, 5),
+        ], // keys
+        &[Some(1), Some(2), Some(3), Some(7), Some(8), Some(9)], // values
+    );
+    write_kvs(
+        &*mem,
+        11, // sequence
+        ValueType::Put,
+        &[(1002, 1), (1003, 1), (1004, 1)], // keys
+        &[None, Some(5), None],             // values
+    );
+
+    let mut iter = mem.iter(IterContext { batch_size: 4 }).unwrap();
+    assert_eq!(schema, *iter.schema());
+    assert_eq!(RowOrdering::Key, iter.ordering());
+
+    check_iter_content(
+        &mut *iter,
+        &[
+            (1000, 1),
+            (1000, 2),
+            (1001, 1),
+            (1002, 1),
+            (1003, 1),
+            (1004, 1),
+            (2002, 1),
+            (2003, 1),
+            (2003, 5),
+        ], // keys
+        &[10, 10, 10, 11, 11, 11, 10, 10, 10], // sequences
+        &[
+            ValueType::Put,
+            ValueType::Put,
+            ValueType::Put,
+            ValueType::Put,
+            ValueType::Put,
+            ValueType::Put,
+            ValueType::Put,
+            ValueType::Put,
+            ValueType::Put,
+        ], // value types
+        &[
+            Some(1),
+            Some(2),
+            Some(3),
+            None,
+            Some(5),
+            None,
+            Some(7),
+            Some(8),
+            Some(9),
+        ], // values
+    );
+}

src/storage/src/metadata.rs

@@ -48,7 +48,7 @@ pub type VersionNumber = u32;
 // TODO(yingwen): Make some fields of metadata private.
 
 /// In memory metadata of region.
-#[derive(Clone)]
+#[derive(Clone, Debug)]
 pub struct RegionMetadata {
     /// Schema of the region.
     ///
@@ -66,13 +66,13 @@ pub struct RegionMetadata {
 
 pub type RegionMetadataRef = Arc<RegionMetadata>;
 
-#[derive(Clone)]
+#[derive(Clone, Debug, PartialEq)]
 pub struct ColumnMetadata {
     pub cf_id: ColumnFamilyId,
     pub desc: ColumnDescriptor,
 }
 
-#[derive(Clone)]
+#[derive(Clone, Debug, PartialEq)]
 pub struct ColumnsMetadata {
     /// All columns, in `(key columns, timestamp, [version,] value columns)` order.
     ///
@@ -82,7 +82,7 @@ pub struct ColumnsMetadata {
     pub name_to_col_index: HashMap<String, usize>,
 }
 
-#[derive(Default, Clone)]
+#[derive(Clone, Debug, Default, PartialEq)]
 pub struct RowKeyMetadata {
     /// Exclusive end index of row key columns.
     row_key_end: usize,
@@ -93,7 +93,7 @@ pub struct RowKeyMetadata {
     pub enable_version_column: bool,
 }
 
-#[derive(Clone)]
+#[derive(Clone, Debug, PartialEq)]
 pub struct ColumnsRowKeyMetadata {
     columns: ColumnsMetadata,
     row_key: RowKeyMetadata,
@@ -121,7 +121,7 @@ impl ColumnsRowKeyMetadata {
 
 pub type ColumnsRowKeyMetadataRef = Arc<ColumnsRowKeyMetadata>;
 
-#[derive(Clone)]
+#[derive(Clone, Debug)]
 pub struct ColumnFamiliesMetadata {
     /// Map column family id to column family metadata.
     id_to_cfs: HashMap<ColumnFamilyId, ColumnFamilyMetadata>,
@@ -133,7 +133,7 @@ impl ColumnFamiliesMetadata {
     }
 }
 
-#[derive(Clone)]
+#[derive(Clone, Debug)]
 pub struct ColumnFamilyMetadata {
     /// Column family name.
     pub name: String,

src/store-api/src/storage/descriptors.rs

@@ -9,7 +9,7 @@ pub type ColumnFamilyId = u32;
 
 // TODO(yingwen): Validate default value has same type with column, and name is a valid column name.
 /// A [ColumnDescriptor] contains information to create a column.
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, PartialEq)]
 pub struct ColumnDescriptor {
     pub id: ColumnId,
     pub name: String,
@@ -29,7 +29,7 @@ impl From<&ColumnDescriptor> for ColumnSchema {
 }
 
 /// A [RowKeyDescriptor] contains information about row key.
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, PartialEq)]
 pub struct RowKeyDescriptor {
     pub columns: Vec<ColumnDescriptor>,
     /// Timestamp key column.
@@ -41,7 +41,7 @@ pub struct RowKeyDescriptor {
 }
 
 /// A [ColumnFamilyDescriptor] contains information to create a column family.
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, PartialEq)]
 pub struct ColumnFamilyDescriptor {
     pub cf_id: ColumnFamilyId,
     pub name: String,
@@ -50,7 +50,7 @@ pub struct ColumnFamilyDescriptor {
 }
 
 /// A [RegionDescriptor] contains information to create a region.
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, PartialEq)]
 pub struct RegionDescriptor {
     /// Region name.
     pub name: String,