test: add unit test

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!()
+        }
+    }
 }

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),