feat: define structs inside the tree

src/mito2/src/memtable/merge_tree.rs

@@ -39,7 +39,7 @@ use crate::memtable::{
     MemtableRef, MemtableStats,
 };
 
-/// Id of a shard, unique inside a tree.
+/// Id of a shard, only unique inside a partition.
 type ShardId = u32;
 /// Index of a primary key in a shard.
 type PkIndex = u16;

src/mito2/src/memtable/merge_tree/data.rs

@@ -13,3 +13,9 @@
 // limitations under the License.
 
 //! Data part of a shard.
+
+/// Buffer to store columns not in the primary key.
+pub struct DataBuffer {}
+
+/// Data parts under a shard.
+pub struct DataParts {}

src/mito2/src/memtable/merge_tree/dict.rs

@@ -13,3 +13,16 @@
 // limitations under the License.
 
 //! Key dictionary of a shard.
+
+use std::sync::Arc;
+
+/// Builder to build a key dictionary.
+pub struct KeyDictBuilder {}
+
+/// Buffer to store unsorted primary keys.
+pub struct KeyBuffer {}
+
+/// A key dictionary.
+pub struct KeyDict {}
+
+pub type KeyDictRef = Arc<KeyDict>;

src/mito2/src/memtable/merge_tree/partition.rs

@@ -13,3 +13,30 @@
 // limitations under the License.
 
 //! Partition of a merge tree.
+//!
+//! We only support partitioning the tree by pre-defined internal columns.
+
+use std::sync::{Arc, RwLock};
+
+use crate::memtable::merge_tree::shard::Shard;
+use crate::memtable::merge_tree::shard_builder::ShardBuilder;
+use crate::memtable::merge_tree::ShardId;
+
+/// Key of a partition.
+pub type PartitionKey = u32;
+
+/// A tree partition.
+pub struct Partition {
+    inner: RwLock<Inner>,
+}
+
+pub type PartitionRef = Arc<Partition>;
+
+/// Inner struct of the partition.
+struct Inner {
+    /// Shard whose dictionary is active.
+    shard_builder: ShardBuilder,
+    next_shard_id: ShardId,
+    /// Shards with frozon dictionary.
+    shards: Vec<Shard>,
+}

src/mito2/src/memtable/merge_tree/shard.rs

@@ -13,3 +13,16 @@
 // limitations under the License.
 
 //! Shard in a partition.
+
+use crate::memtable::merge_tree::data::DataParts;
+use crate::memtable::merge_tree::dict::KeyDictRef;
+use crate::memtable::merge_tree::ShardId;
+
+/// Shard stores data related to the same key dictionary.
+pub struct Shard {
+    shard_id: ShardId,
+    /// Key dictionary of the shard. `None` if the schema of the tree doesn't have a primary key.
+    key_dict: Option<KeyDictRef>,
+    /// Data in the shard.
+    data_parts: DataParts,
+}

src/mito2/src/memtable/merge_tree/shard_builder.rs

@@ -13,3 +13,15 @@
 // limitations under the License.
 
 //! Builder of a shard.
+
+use crate::memtable::merge_tree::data::DataBuffer;
+use crate::memtable::merge_tree::dict::KeyDictBuilder;
+
+/// Builder to write keys and data to a shard that the key dictionary
+/// is still active.
+pub struct ShardBuilder {
+    /// Builder for the key dictionary.
+    dict_builder: KeyDictBuilder,
+    /// Buffer to store data.
+    data_buffer: DataBuffer,
+}

src/mito2/src/memtable/merge_tree/tree.rs

@@ -14,14 +14,16 @@
 
 //! Implementation of the merge tree.
 
-use std::sync::Arc;
+use std::collections::BTreeMap;
+use std::sync::{Arc, RwLock};
 
 use store_api::metadata::RegionMetadataRef;
 use store_api::storage::ColumnId;
 use table::predicate::Predicate;
 
 use crate::error::Result;
 use crate::memtable::merge_tree::metrics::WriteMetrics;
+use crate::memtable::merge_tree::partition::{PartitionKey, PartitionRef};
 use crate::memtable::merge_tree::MergeTreeConfig;
 use crate::memtable::{BoxedBatchIterator, KeyValues};
 use crate::row_converter::{McmpRowCodec, SortField};
@@ -34,6 +36,8 @@ pub struct MergeTree {
     pub(crate) metadata: RegionMetadataRef,
     /// Primary key codec.
     row_codec: Arc<McmpRowCodec>,
+    /// Partitions in the tree.
+    partitions: RwLock<BTreeMap<PartitionKey, PartitionRef>>,
 }
 
 impl MergeTree {
@@ -50,6 +54,7 @@ impl MergeTree {
             config: config.clone(),
             metadata,
             row_codec: Arc::new(row_codec),
+            partitions: Default::default(),
         }
     }