diff --git a/src/storage/src/memtable.rs b/src/storage/src/memtable.rs index 2cb20e06c1..0280194a3a 100644 --- a/src/storage/src/memtable.rs +++ b/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, diff --git a/src/storage/src/memtable/btree.rs b/src/storage/src/memtable/btree.rs index de33f1098d..ca4db39ba5 100644 --- a/src/storage/src/memtable/btree.rs +++ b/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, sequence: SequenceNumber, @@ -309,7 +311,7 @@ impl InnerKey { } } -#[derive(Clone)] +#[derive(Clone, Debug)] struct RowValue { values: Vec, } diff --git a/src/storage/src/memtable/schema.rs b/src/storage/src/memtable/schema.rs index bcda7ab3c1..e23c210de0 100644 --- a/src/storage/src/memtable/schema.rs +++ b/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, } diff --git a/src/storage/src/memtable/tests.rs b/src/storage/src/memtable/tests.rs new file mode 100644 index 0000000000..a6adfa19f6 --- /dev/null +++ b/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], +) -> 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], +) -> 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], +) { + 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], +) { + 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 + ); +} diff --git a/src/storage/src/metadata.rs b/src/storage/src/metadata.rs index d222d7b14a..8bab1c513c 100644 --- a/src/storage/src/metadata.rs +++ b/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; -#[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, } -#[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; -#[derive(Clone)] +#[derive(Clone, Debug)] pub struct ColumnFamiliesMetadata { /// Map column family id to column family metadata. id_to_cfs: HashMap, @@ -133,7 +133,7 @@ impl ColumnFamiliesMetadata { } } -#[derive(Clone)] +#[derive(Clone, Debug)] pub struct ColumnFamilyMetadata { /// Column family name. pub name: String, diff --git a/src/store-api/src/storage/descriptors.rs b/src/store-api/src/storage/descriptors.rs index 97d5b08b60..b9d0ecb2b8 100644 --- a/src/store-api/src/storage/descriptors.rs +++ b/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, /// 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,