diff --git a/src/mito2/src/manifest/manager.rs b/src/mito2/src/manifest/manager.rs index 8247298d6df..25f9c4c7bef 100644 --- a/src/mito2/src/manifest/manager.rs +++ b/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() } diff --git a/src/mito2/src/memtable.rs b/src/mito2/src/memtable.rs index 5d6a98e377f..e23a8daa8cd 100644 --- a/src/mito2/src/memtable.rs +++ b/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; @@ -62,6 +68,10 @@ impl Memtable for EmptyMemtable { fn id(&self) -> MemtableId { self.id } + + fn write(&self, _kvs: &KeyValues) -> Result<()> { + Ok(()) + } } /// Default memtable builder. diff --git a/src/mito2/src/memtable/key_values.rs b/src/mito2/src/memtable/key_values.rs new file mode 100644 index 00000000000..4ea70c3eae7 --- /dev/null +++ b/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 { + 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 { + 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 { + 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 { + 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, + /// 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]); + } +} diff --git a/src/mito2/src/metadata.rs b/src/mito2/src/metadata.rs index 7f1588d1559..cfd29431b39 100644 --- a/src/mito2/src/metadata.rs +++ b/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}", + ); + } } diff --git a/src/mito2/src/request.rs b/src/mito2/src/request.rs index 10348ac70df..7a303027909 100644 --- a/src/mito2/src/request.rs +++ b/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 { + 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)