feat: Add write method to memtable trait (#2123)

* feat: validate semantic type

* feat: define KeyValues

* test: test semantic type check

* feat: impl KeyValues

* test: test KeyValues

* feat: Add write to Memtable

* style: fix clippy

* docs: more comment

src/mito2/src/manifest/manager.rs

@@ -429,7 +429,8 @@ mod test {
                 ),
                 semantic_type: SemanticType::Field,
                 column_id: 251,
-            });
+            })
+            .primary_key(vec![36]);
         builder.build().unwrap()
     }
 

src/mito2/src/memtable.rs

@@ -14,12 +14,15 @@
 
 //! Memtables are write buffers for regions.
 
+pub mod key_values;
 pub(crate) mod version;
 
 use std::fmt;
 use std::sync::atomic::{AtomicU32, Ordering};
 use std::sync::Arc;
 
+use crate::error::Result;
+use crate::memtable::key_values::KeyValues;
 use crate::metadata::RegionMetadataRef;
 
 /// Id for memtables.
@@ -31,6 +34,9 @@ pub type MemtableId = u32;
 pub trait Memtable: Send + Sync + fmt::Debug {
     /// Returns the id of this memtable.
     fn id(&self) -> MemtableId;
+
+    /// Write key values into the memtable.
+    fn write(&self, kvs: &KeyValues) -> Result<()>;
 }
 
 pub type MemtableRef = Arc<dyn Memtable>;
@@ -62,6 +68,10 @@ impl Memtable for EmptyMemtable {
     fn id(&self) -> MemtableId {
         self.id
     }
+
+    fn write(&self, _kvs: &KeyValues) -> Result<()> {
+        Ok(())
+    }
 }
 
 /// Default memtable builder.

src/mito2/src/memtable/key_values.rs

@@ -0,0 +1,379 @@
+// Copyright 2023 Greptime Team
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::collections::HashMap;
+
+use greptime_proto::v1::mito::{Mutation, OpType};
+use greptime_proto::v1::{Row, Rows, Value};
+use store_api::storage::SequenceNumber;
+
+use crate::metadata::{RegionMetadata, SemanticType};
+
+/// Key value view of a mutation.
+#[derive(Debug)]
+pub struct KeyValues {
+    /// Mutation to read.
+    ///
+    /// This mutation must be a valid mutation and rows in the mutation
+    /// must not be `None`.
+    mutation: Mutation,
+    /// Key value read helper.
+    helper: ReadRowHelper,
+}
+
+impl KeyValues {
+    /// Creates [KeyValues] from specific `mutation`.
+    ///
+    /// Returns `None` if `rows` of the `mutation` is `None`.
+    pub fn new(metadata: &RegionMetadata, mutation: Mutation) -> Option<KeyValues> {
+        let rows = mutation.rows.as_ref()?;
+        let helper = ReadRowHelper::new(metadata, rows);
+
+        Some(KeyValues { mutation, helper })
+    }
+
+    /// Returns a key value iterator.
+    pub fn iter(&self) -> impl Iterator<Item = KeyValue> {
+        let rows = self.mutation.rows.as_ref().unwrap();
+        rows.rows.iter().enumerate().map(|(idx, row)| {
+            KeyValue {
+                row,
+                helper: &self.helper,
+                sequence: self.mutation.sequence + idx as u64, // Calculate sequence for each row.
+                // Safety: This is a valid mutation.
+                op_type: OpType::from_i32(self.mutation.op_type).unwrap(),
+            }
+        })
+    }
+
+    /// Returns number of rows.
+    pub fn num_rows(&self) -> usize {
+        // Safety: rows is not None.
+        self.mutation.rows.as_ref().unwrap().rows.len()
+    }
+}
+
+/// Key value view of a row.
+///
+/// A key value view divides primary key columns and field (value) columns.
+/// Primary key columns have the same order as region's primary key. Field
+/// columns are ordered by their position in the region schema (The same order
+/// as users defined while creating the region).
+#[derive(Debug)]
+pub struct KeyValue<'a> {
+    row: &'a Row,
+    helper: &'a ReadRowHelper,
+    sequence: SequenceNumber,
+    op_type: OpType,
+}
+
+impl<'a> KeyValue<'a> {
+    /// Get primary key columns.
+    pub fn primary_keys(&self) -> impl Iterator<Item = &Value> {
+        self.helper.indices[..self.helper.num_primary_key_column]
+            .iter()
+            .map(|idx| &self.row.values[*idx])
+    }
+
+    /// Get field columns.
+    pub fn fields(&self) -> impl Iterator<Item = &Value> {
+        self.helper.indices[self.helper.num_primary_key_column + 1..]
+            .iter()
+            .map(|idx| &self.row.values[*idx])
+    }
+
+    /// Get timestamp.
+    pub fn timestamp(&self) -> Value {
+        // Timestamp is primitive, we clone it.
+        let index = self.helper.indices[self.helper.num_primary_key_column];
+        self.row.values[index].clone()
+    }
+
+    /// Get number of primary key columns.
+    pub fn num_primary_keys(&self) -> usize {
+        self.helper.num_primary_key_column
+    }
+
+    /// Get number of field columns.
+    pub fn num_fields(&self) -> usize {
+        self.row.values.len() - self.helper.num_primary_key_column - 1
+    }
+
+    /// Get sequence.
+    pub fn sequence(&self) -> SequenceNumber {
+        self.sequence
+    }
+
+    /// Get op type.
+    pub fn op_type(&self) -> OpType {
+        self.op_type
+    }
+}
+
+/// Helper to read rows in key, value order.
+#[derive(Debug)]
+struct ReadRowHelper {
+    /// Key and value column indices.
+    ///
+    /// `indices[..num_primary_key_column]` are primary key columns, `indices[num_primary_key_column]`
+    /// is the timestamp column and remainings are field columns.
+    indices: Vec<usize>,
+    /// Number of primary key columns.
+    num_primary_key_column: usize,
+}
+
+impl ReadRowHelper {
+    /// Creates a [ReadRowHelper] for specific `rows`.
+    ///
+    /// # Panics
+    /// The `rows` must fill their missing columns first and have same columns with `metadata`.
+    /// Otherwise this method will panic.
+    fn new(metadata: &RegionMetadata, rows: &Rows) -> ReadRowHelper {
+        assert_eq!(
+            metadata.column_metadatas.len(),
+            rows.schema.len(),
+            "Length mismatch, column_metas: {:?}, rows_schema: {:?}",
+            metadata.column_metadatas,
+            rows.schema
+        );
+
+        // Build a name to index mapping for rows.
+        let name_to_index: HashMap<_, _> = rows
+            .schema
+            .iter()
+            .enumerate()
+            .map(|(index, col)| (&col.column_name, index))
+            .collect();
+        let mut indices = Vec::with_capacity(metadata.column_metadatas.len());
+
+        // Get primary key indices.
+        for pk_id in &metadata.primary_key {
+            // Safety: Id comes from primary key.
+            let column = metadata.column_by_id(*pk_id).unwrap();
+            let index = name_to_index.get(&column.column_schema.name).unwrap();
+            indices.push(*index);
+        }
+        // Get timestamp index.
+        let ts_index = name_to_index
+            .get(&metadata.time_index_column().column_schema.name)
+            .unwrap();
+        indices.push(*ts_index);
+        // Iterate columns and find field columns.
+        for column in metadata.column_metadatas.iter() {
+            if column.semantic_type == SemanticType::Field {
+                // Get index in request for each field column.
+                let index = name_to_index.get(&column.column_schema.name).unwrap();
+                indices.push(*index);
+            }
+        }
+
+        ReadRowHelper {
+            indices,
+            num_primary_key_column: metadata.primary_key.len(),
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use datatypes::prelude::ConcreteDataType;
+    use datatypes::schema::ColumnSchema;
+    use greptime_proto::v1;
+    use greptime_proto::v1::{value, ColumnDataType, Value};
+    use store_api::storage::RegionId;
+
+    use super::*;
+    use crate::metadata::{ColumnMetadata, RegionMetadataBuilder};
+
+    const TS_NAME: &str = "ts";
+    const START_SEQ: SequenceNumber = 100;
+
+    /// Creates a region: `ts, k0, k1, ..., v0, v1, ...`
+    fn new_region_metadata(num_tag: usize, num_field: usize) -> RegionMetadata {
+        let mut builder = RegionMetadataBuilder::new(RegionId::new(1, 1), 1);
+        let mut column_id = 0;
+        builder.push_column_metadata(ColumnMetadata {
+            column_schema: ColumnSchema::new(
+                TS_NAME,
+                ConcreteDataType::timestamp_millisecond_datatype(),
+                false,
+            ),
+            semantic_type: SemanticType::Timestamp,
+            column_id,
+        });
+        // For simplicity, we use the same data type for tag/field columns.
+        let mut primary_key = Vec::with_capacity(num_tag);
+        for i in 0..num_tag {
+            column_id += 1;
+            builder.push_column_metadata(ColumnMetadata {
+                column_schema: ColumnSchema::new(
+                    format!("k{i}"),
+                    ConcreteDataType::int64_datatype(),
+                    true,
+                ),
+                semantic_type: SemanticType::Tag,
+                column_id,
+            });
+            primary_key.push(i as u32 + 1);
+        }
+        for i in 0..num_field {
+            column_id += 1;
+            builder.push_column_metadata(ColumnMetadata {
+                column_schema: ColumnSchema::new(
+                    format!("v{i}"),
+                    ConcreteDataType::int64_datatype(),
+                    true,
+                ),
+                semantic_type: SemanticType::Field,
+                column_id,
+            });
+        }
+        builder.primary_key(primary_key);
+        builder.build().unwrap()
+    }
+
+    /// Creates rows `[ 0, 1, ..., n ] x num_rows`
+    fn new_rows(column_names: &[&str], num_rows: usize) -> Rows {
+        let mut rows = Vec::with_capacity(num_rows);
+        for _ in 0..num_rows {
+            // For simplicity, we use i64 for timestamp millisecond type. This violates the column schema
+            // but it's acceptable for tests.
+            let values: Vec<_> = (0..column_names.len())
+                .map(|idx| i64_value(idx as i64))
+                .collect();
+            rows.push(Row { values });
+        }
+
+        let schema = column_names
+            .iter()
+            .map(|column_name| {
+                let datatype = if *column_name == TS_NAME {
+                    ColumnDataType::TimestampMillisecond as i32
+                } else {
+                    ColumnDataType::Int64 as i32
+                };
+                let semantic_type = if column_name.starts_with('k') {
+                    SemanticType::Tag as i32
+                } else if column_name.starts_with('v') {
+                    SemanticType::Field as i32
+                } else {
+                    SemanticType::Timestamp as i32
+                };
+                v1::ColumnSchema {
+                    column_name: column_name.to_string(),
+                    datatype,
+                    semantic_type,
+                }
+            })
+            .collect();
+
+        Rows { rows, schema }
+    }
+
+    fn new_mutation(column_names: &[&str], num_rows: usize) -> Mutation {
+        let rows = new_rows(column_names, num_rows);
+        Mutation {
+            op_type: OpType::Put as i32,
+            sequence: START_SEQ,
+            rows: Some(rows),
+        }
+    }
+
+    fn i64_value(data: i64) -> Value {
+        Value {
+            value: Some(value::Value::I64Value(data)),
+        }
+    }
+
+    fn check_key_values(kvs: &KeyValues, num_rows: usize, keys: &[i64], ts: i64, values: &[i64]) {
+        assert_eq!(num_rows, kvs.num_rows());
+        let mut expect_seq = START_SEQ;
+        let expect_ts = i64_value(ts);
+        for kv in kvs.iter() {
+            assert_eq!(expect_seq, kv.sequence());
+            expect_seq += 1;
+            assert_eq!(OpType::Put, kv.op_type);
+            assert_eq!(keys.len(), kv.num_primary_keys());
+            assert_eq!(values.len(), kv.num_fields());
+
+            assert_eq!(expect_ts, kv.timestamp());
+            let expect_keys: Vec<_> = keys.iter().map(|k| i64_value(*k)).collect();
+            let actual_keys: Vec<_> = kv.primary_keys().cloned().collect();
+            assert_eq!(expect_keys, actual_keys);
+            let expect_values: Vec<_> = values.iter().map(|v| i64_value(*v)).collect();
+            let actual_values: Vec<_> = kv.fields().cloned().collect();
+            assert_eq!(expect_values, actual_values);
+        }
+    }
+
+    #[test]
+    fn test_empty_key_values() {
+        let meta = new_region_metadata(1, 1);
+        let mutation = Mutation {
+            op_type: OpType::Put as i32,
+            sequence: 100,
+            rows: None,
+        };
+        let kvs = KeyValues::new(&meta, mutation);
+        assert!(kvs.is_none());
+    }
+
+    #[test]
+    fn test_ts_only() {
+        let meta = new_region_metadata(0, 0);
+        let mutation = new_mutation(&["ts"], 2);
+        let kvs = KeyValues::new(&meta, mutation).unwrap();
+        check_key_values(&kvs, 2, &[], 0, &[]);
+    }
+
+    #[test]
+    fn test_no_field() {
+        let meta = new_region_metadata(2, 0);
+        // The value of each row:
+        // k1=0, ts=1, k0=2,
+        let mutation = new_mutation(&["k1", "ts", "k0"], 3);
+        let kvs = KeyValues::new(&meta, mutation).unwrap();
+        // KeyValues
+        // keys: [k0=2, k1=0]
+        // ts: 1,
+        check_key_values(&kvs, 3, &[2, 0], 1, &[]);
+    }
+
+    #[test]
+    fn test_no_tag() {
+        let meta = new_region_metadata(0, 2);
+        // The value of each row:
+        // v1=0, v0=1, ts=2,
+        let mutation = new_mutation(&["v1", "v0", "ts"], 3);
+        let kvs = KeyValues::new(&meta, mutation).unwrap();
+        // KeyValues (note that v0 is in front of v1 in region schema)
+        // ts: 2,
+        // fields: [v0=1, v1=0]
+        check_key_values(&kvs, 3, &[], 2, &[1, 0]);
+    }
+
+    #[test]
+    fn test_tag_field() {
+        let meta = new_region_metadata(2, 2);
+        // The value of each row:
+        // k0=0, v0=1, ts=2, k1=3, v1=4,
+        let mutation = new_mutation(&["k0", "v0", "ts", "k1", "v1"], 3);
+        let kvs = KeyValues::new(&meta, mutation).unwrap();
+        // KeyValues
+        // keys: [k0=0, k1=3]
+        // ts: 2,
+        // fields: [v0=1, v1=4]
+        check_key_values(&kvs, 3, &[0, 3], 2, &[1, 4]);
+    }
+}

src/mito2/src/metadata.rs

@@ -194,11 +194,16 @@ impl RegionMetadata {
                     }
                 );
 
+                // Safety: Column with specific id must exist.
+                let column = self.column_by_id(*column_id).unwrap();
                 // Checks duplicate.
                 ensure!(
                     !pk_ids.contains(&column_id),
                     InvalidMetaSnafu {
-                        reason: format!("duplicate column {} in primary key", id_names[column_id]),
+                        reason: format!(
+                            "duplicate column {} in primary key",
+                            column.column_schema.name
+                        ),
                     }
                 );
 
@@ -208,7 +213,18 @@ impl RegionMetadata {
                     InvalidMetaSnafu {
                         reason: format!(
                             "column {} is already a time index column",
-                            id_names[column_id]
+                            column.column_schema.name,
+                        ),
+                    }
+                );
+
+                // Checks semantic type.
+                ensure!(
+                    column.semantic_type == SemanticType::Tag,
+                    InvalidMetaSnafu {
+                        reason: format!(
+                            "semantic type of column {} should be Tag, not {:?}",
+                            column.column_schema.name, column.semantic_type
                         ),
                     }
                 );
@@ -217,6 +233,23 @@ impl RegionMetadata {
             }
         }
 
+        // Checks tag semantic type.
+        let num_tag = self
+            .column_metadatas
+            .iter()
+            .filter(|col| col.semantic_type == SemanticType::Tag)
+            .count();
+        ensure!(
+            num_tag == self.primary_key.len(),
+            InvalidMetaSnafu {
+                reason: format!(
+                    "number of primary key columns {} not equal to tag columns {}",
+                    self.primary_key.len(),
+                    num_tag
+                ),
+            }
+        );
+
         Ok(())
     }
 }
@@ -597,4 +630,63 @@ mod test {
             "unexpected err: {err}",
         );
     }
+
+    #[test]
+    fn test_primary_key_semantic_type() {
+        let mut builder = create_builder();
+        builder
+            .push_column_metadata(ColumnMetadata {
+                column_schema: ColumnSchema::new(
+                    "ts",
+                    ConcreteDataType::timestamp_millisecond_datatype(),
+                    false,
+                ),
+                semantic_type: SemanticType::Timestamp,
+                column_id: 1,
+            })
+            .push_column_metadata(ColumnMetadata {
+                column_schema: ColumnSchema::new("a", ConcreteDataType::float64_datatype(), true),
+                semantic_type: SemanticType::Field,
+                column_id: 2,
+            })
+            .primary_key(vec![2]);
+        let err = builder.build().unwrap_err();
+        assert!(
+            err.to_string()
+                .contains("semantic type of column a should be Tag, not Field"),
+            "unexpected err: {err}",
+        );
+    }
+
+    #[test]
+    fn test_primary_key_tag_num() {
+        let mut builder = create_builder();
+        builder
+            .push_column_metadata(ColumnMetadata {
+                column_schema: ColumnSchema::new(
+                    "ts",
+                    ConcreteDataType::timestamp_millisecond_datatype(),
+                    false,
+                ),
+                semantic_type: SemanticType::Timestamp,
+                column_id: 1,
+            })
+            .push_column_metadata(ColumnMetadata {
+                column_schema: ColumnSchema::new("a", ConcreteDataType::string_datatype(), true),
+                semantic_type: SemanticType::Tag,
+                column_id: 2,
+            })
+            .push_column_metadata(ColumnMetadata {
+                column_schema: ColumnSchema::new("b", ConcreteDataType::string_datatype(), true),
+                semantic_type: SemanticType::Tag,
+                column_id: 3,
+            })
+            .primary_key(vec![2]);
+        let err = builder.build().unwrap_err();
+        assert!(
+            err.to_string()
+                .contains("number of primary key columns 1 not equal to tag columns 2"),
+            "unexpected err: {err}",
+        );
+    }
 }

src/mito2/src/request.rs

@@ -128,6 +128,11 @@ impl WriteRequest {
         unimplemented!()
     }
 
+    /// Get column index by name.
+    pub(crate) fn column_index_by_name(&self, name: &str) -> Option<usize> {
+        self.name_to_index.get(name).copied()
+    }
+
     /// Checks schema of rows.
     ///
     /// If column with default value is missing, it returns a special [FillDefault](crate::error::Error::FillDefault)