From c78e14be94bbabf42840d6129190dc737a034084 Mon Sep 17 00:00:00 2001 From: "Lei, HUANG" Date: Tue, 14 May 2024 02:22:19 +0000 Subject: [PATCH] test: add unit test --- src/mito2/src/memtable/partition_tree.rs | 173 ++++++++++++++++++++++- src/mito2/src/test_util/memtable_util.rs | 2 +- 2 files changed, 173 insertions(+), 2 deletions(-) diff --git a/src/mito2/src/memtable/partition_tree.rs b/src/mito2/src/memtable/partition_tree.rs index fccbf8387c4a..d2435e0790c0 100644 --- a/src/mito2/src/memtable/partition_tree.rs +++ b/src/mito2/src/memtable/partition_tree.rs @@ -311,14 +311,24 @@ impl MemtableBuilder for PartitionTreeMemtableBuilder { #[cfg(test)] mod tests { + use api::v1::value::ValueData; + use api::v1::{Row, Rows, SemanticType}; use common_time::Timestamp; use datafusion_common::{Column, ScalarValue}; use datafusion_expr::{BinaryExpr, Expr, Operator}; + use datatypes::data_type::ConcreteDataType; use datatypes::scalars::ScalarVector; + use datatypes::schema::ColumnSchema; + use datatypes::value::Value; use datatypes::vectors::Int64Vector; + use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder}; + use store_api::storage::RegionId; use super::*; - use crate::test_util::memtable_util::{self, collect_iter_timestamps}; + use crate::row_converter::{McmpRowCodec, RowCodec, SortField}; + use crate::test_util::memtable_util::{ + self, collect_iter_timestamps, region_metadata_to_row_schema, + }; #[test] fn test_memtable_sorted_input() { @@ -562,4 +572,165 @@ mod tests { assert!(config.dedup); assert_eq!(PartitionTreeConfig::default(), config); } + + fn metadata_for_metric_engine() -> RegionMetadataRef { + let mut builder = RegionMetadataBuilder::new(RegionId::new(123, 456)); + builder + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "__table_id", + ConcreteDataType::uint32_datatype(), + false, + ), + semantic_type: SemanticType::Tag, + column_id: 2147483652, + }) + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "__ts_id", + ConcreteDataType::uint64_datatype(), + false, + ), + semantic_type: SemanticType::Tag, + column_id: 2147483651, + }) + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "test_label", + ConcreteDataType::string_datatype(), + false, + ), + semantic_type: SemanticType::Tag, + column_id: 2, + }) + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "greptime_timestamp", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + semantic_type: SemanticType::Timestamp, + column_id: 0, + }) + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "greptime_value", + ConcreteDataType::float64_datatype(), + true, + ), + semantic_type: SemanticType::Field, + column_id: 1, + }) + .primary_key(vec![2147483652, 2147483651, 2]); + let region_metadata = builder.build().unwrap(); + Arc::new(region_metadata) + } + + fn build_key_values( + metadata: RegionMetadataRef, + labels: &[&str], + table_id: &[u32], + ts_id: &[u64], + ts: &[i64], + values: &[f64], + sequence: u64, + ) -> KeyValues { + let column_schema = region_metadata_to_row_schema(&metadata); + + let rows = ts + .iter() + .zip(table_id.iter()) + .zip(ts_id.iter()) + .zip(labels.iter()) + .zip(values.iter()) + .map(|((((ts, table_id), ts_id), label), val)| Row { + values: vec![ + api::v1::Value { + value_data: Some(ValueData::U32Value(*table_id)), + }, + api::v1::Value { + value_data: Some(ValueData::U64Value(*ts_id)), + }, + api::v1::Value { + value_data: Some(ValueData::StringValue(label.to_string())), + }, + api::v1::Value { + value_data: Some(ValueData::TimestampMillisecondValue(*ts)), + }, + api::v1::Value { + value_data: Some(ValueData::F64Value(*val)), + }, + ], + }) + .collect(); + let mutation = api::v1::Mutation { + op_type: 1, + sequence, + rows: Some(Rows { + schema: column_schema, + rows, + }), + }; + KeyValues::new(metadata.as_ref(), mutation).unwrap() + } + #[test] + fn test_write_freeze() { + let metadata = metadata_for_metric_engine(); + let memtable = PartitionTreeMemtableBuilder::new( + PartitionTreeConfig { + index_max_keys_per_shard: 40, + ..Default::default() + }, + None, + ) + .build(1, &metadata); + + let codec = McmpRowCodec::new( + metadata + .primary_key_columns() + .map(|c| SortField::new(c.column_schema.data_type.clone())) + .collect(), + ); + + memtable + .write(&build_key_values( + metadata.clone(), + &["daily", "10min", "daily", "10min"], + &[1025, 1025, 1025, 1025], + &[ + 16442255374049317291, + 5686004715529701024, + 16442255374049317291, + 5686004715529701024, + ], + &[1712070000000, 1712717731000, 1712761200000, 1712761200000], + &[0.0, 0.0, 0.0, 0.0], + 1, + )) + .unwrap(); + + memtable.freeze().unwrap(); + let new_memtable = memtable.fork(2, &metadata); + + new_memtable + .write(&build_key_values( + metadata.clone(), + &["10min"], + &[1025], + &[5686004715529701024], + &[1714643131000], + &[0.1], + 2, + )) + .unwrap(); + + let mut reader = new_memtable.iter(None, None).unwrap(); + let batch = reader.next().unwrap().unwrap(); + let pk = codec.decode(batch.primary_key()).unwrap(); + if let Value::String(s) = &pk[2] { + assert_eq!("10min", s.as_utf8()); + } else { + unreachable!() + } + } } diff --git a/src/mito2/src/test_util/memtable_util.rs b/src/mito2/src/test_util/memtable_util.rs index 1628ecf56e7c..b2764ba88957 100644 --- a/src/mito2/src/test_util/memtable_util.rs +++ b/src/mito2/src/test_util/memtable_util.rs @@ -115,7 +115,7 @@ pub fn metadata_with_primary_key( enable_table_id: bool, ) -> RegionMetadataRef { let mut builder = RegionMetadataBuilder::new(RegionId::new(123, 456)); - let maybe_table_id = if enable_table_id { "table_id" } else { "k1" }; + let maybe_table_id = if enable_table_id { "__table_id" } else { "k1" }; builder .push_column_metadata(ColumnMetadata { column_schema: ColumnSchema::new("k0", ConcreteDataType::string_datatype(), false),