fix(mito2): split metric values in primary-key SSTs

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

src/metric-engine/src/engine/put.rs

@@ -762,6 +762,7 @@ impl MetricEngineInner {
 mod tests {
     use std::collections::HashSet;
 
+    use api::v1::region::{StrictWindow, compact_request};
     use api::v1::value::ValueData;
     use api::v1::{ColumnDataType, ColumnSchema as PbColumnSchema};
     use common_error::ext::ErrorExt;
@@ -785,7 +786,8 @@ mod tests {
     use store_api::path_utils::table_dir;
     use store_api::region_engine::RegionEngine;
     use store_api::region_request::{
-        EnterStagingRequest, RegionRequest, StagingPartitionDirective,
+        AlterKind, EnterStagingRequest, RegionAlterRequest, RegionCompactRequest, RegionRequest,
+        SetRegionOption, StagingPartitionDirective,
     };
     use store_api::storage::ScanRequest;
     use store_api::storage::consts::PRIMARY_KEY_COLUMN_NAME;
@@ -1201,33 +1203,39 @@ mod tests {
     }
 
     #[tokio::test]
-    async fn test_metric_value_split_roundtrip_after_flush() {
+    async fn test_metric_value_split_across_sst_formats_and_compaction() {
         let env = TestEnv::new().await;
         env.init_metric_region().await;
 
         let schema = test_util::row_schema_with_tags(&["job"]);
-        let rows = [
+        let build_rows = |samples: &[(i64, f64, &str)]| {
+            samples
+                .iter()
+                .map(|(timestamp, value, job)| Row {
+                    values: vec![
+                        ValueData::TimestampMillisecondValue(*timestamp).into(),
+                        ValueData::F64Value(*value).into(),
+                        ValueData::StringValue(job.to_string()).into(),
+                    ],
+                })
+                .collect()
+        };
+        let rows = build_rows(&[
             (0, 1.0, "integer"),
             (1, 2.0, "integer"),
             (0, 1.5, "float"),
             (1, 2.0, "float"),
-        ]
-        .into_iter()
-        .map(|(timestamp, value, job)| Row {
-            values: vec![
-                ValueData::TimestampMillisecondValue(timestamp).into(),
-                ValueData::F64Value(value).into(),
-                ValueData::StringValue(job.to_string()).into(),
-            ],
-        })
-        .collect();
+        ]);
 
         let logical_region_id = env.default_logical_region_id();
         env.metric()
             .handle_request(
                 logical_region_id,
                 RegionRequest::Put(RegionPutRequest {
-                    rows: Rows { schema, rows },
+                    rows: Rows {
+                        schema: schema.clone(),
+                        rows,
+                    },
                     hint: None,
                     partition_expr_version: None,
                 }),
@@ -1258,28 +1266,6 @@ mod tests {
             "metric physical reads should hide split companion column"
         );
 
-        let raw_physical_batches = RecordBatches::try_collect(
-            env.mito()
-                .scan_to_stream(
-                    to_data_region_id(physical_region_id),
-                    ScanRequest::default(),
-                )
-                .await
-                .unwrap(),
-        )
-        .await
-        .unwrap();
-        let physical_rows = collect_value_rows(&raw_physical_batches, Some(&int_column_name));
-        assert_eq!(
-            physical_rows,
-            vec![
-                ("float".to_string(), 0, Some(1.5), None),
-                ("float".to_string(), 1, None, Some(2)),
-                ("integer".to_string(), 0, None, Some(1)),
-                ("integer".to_string(), 1, None, Some(2)),
-            ]
-        );
-
         let logical_batches = RecordBatches::try_collect(
             env.metric()
                 .scan_to_stream(logical_region_id, ScanRequest::default())
@@ -1321,6 +1307,98 @@ mod tests {
                 ("integer".to_string(), 1, Some(2.0), None),
             ]
         );
+
+        env.metric()
+            .handle_request(
+                physical_region_id,
+                RegionRequest::Alter(RegionAlterRequest {
+                    kind: AlterKind::SetRegionOptions {
+                        options: vec![SetRegionOption::Format("primary_key".to_string())],
+                    },
+                }),
+            )
+            .await
+            .unwrap();
+
+        let rows = build_rows(&[
+            (2, 3.0, "integer"),
+            (3, 4.0, "integer"),
+            (2, 3.5, "float"),
+            (3, 4.0, "float"),
+        ]);
+        env.metric()
+            .handle_request(
+                logical_region_id,
+                RegionRequest::Put(RegionPutRequest {
+                    rows: Rows { schema, rows },
+                    hint: None,
+                    partition_expr_version: None,
+                }),
+            )
+            .await
+            .unwrap();
+        env.metric()
+            .handle_request(physical_region_id, RegionRequest::Flush(Default::default()))
+            .await
+            .unwrap();
+
+        let expected_physical_rows = vec![
+            ("float".to_string(), 0, Some(1.5), None),
+            ("float".to_string(), 1, None, Some(2)),
+            ("float".to_string(), 2, Some(3.5), None),
+            ("float".to_string(), 3, None, Some(4)),
+            ("integer".to_string(), 0, None, Some(1)),
+            ("integer".to_string(), 1, None, Some(2)),
+            ("integer".to_string(), 2, None, Some(3)),
+            ("integer".to_string(), 3, None, Some(4)),
+        ];
+        let data_region_id = to_data_region_id(physical_region_id);
+        let raw_physical_batches = RecordBatches::try_collect(
+            env.mito()
+                .scan_to_stream(data_region_id, ScanRequest::default())
+                .await
+                .unwrap(),
+        )
+        .await
+        .unwrap();
+        assert_eq!(
+            collect_value_rows(&raw_physical_batches, Some(&int_column_name)),
+            expected_physical_rows
+        );
+        assert_eq!(
+            env.mito().region_statistic(data_region_id).unwrap().sst_num,
+            2
+        );
+
+        env.mito()
+            .handle_request(
+                data_region_id,
+                RegionRequest::Compact(RegionCompactRequest {
+                    options: compact_request::Options::StrictWindow(StrictWindow {
+                        window_seconds: 0,
+                    }),
+                    parallelism: None,
+                }),
+            )
+            .await
+            .unwrap();
+        assert_eq!(
+            env.mito().region_statistic(data_region_id).unwrap().sst_num,
+            1
+        );
+
+        let raw_physical_batches = RecordBatches::try_collect(
+            env.mito()
+                .scan_to_stream(data_region_id, ScanRequest::default())
+                .await
+                .unwrap(),
+        )
+        .await
+        .unwrap();
+        assert_eq!(
+            collect_value_rows(&raw_physical_batches, Some(&int_column_name)),
+            expected_physical_rows
+        );
     }
 
     #[tokio::test]

src/mito2/src/sst/parquet/flat_format.rs

@@ -57,9 +57,6 @@ use crate::sst::parquet::format::{
     FIXED_POS_COLUMN_NUM, FormatProjection, INTERNAL_COLUMN_NUM, PrimaryKeyArray,
     PrimaryKeyReadFormat, StatValues, column_null_counts, column_values,
 };
-use crate::sst::parquet::metric_value_split::{
-    MetricValueSplitColumn, metric_value_split_columns, split_metric_value_columns,
-};
 use crate::sst::parquet::read_columns::ParquetReadColumns;
 use crate::sst::{
     FlatSchemaOptions, flat_sst_arrow_schema_column_num, tag_maybe_to_dictionary_field,
@@ -71,18 +68,15 @@ pub(crate) struct FlatWriteFormat {
     /// SST file schema.
     arrow_schema: SchemaRef,
     override_sequence: Option<SequenceNumber>,
-    metric_value_split_columns: Vec<MetricValueSplitColumn>,
 }
 
 impl FlatWriteFormat {
     /// Creates a new helper.
     pub(crate) fn new(metadata: RegionMetadataRef, options: &FlatSchemaOptions) -> FlatWriteFormat {
         let arrow_schema = to_flat_sst_arrow_schema(&metadata, options);
-        let metric_value_split_columns = metric_value_split_columns(&metadata, &arrow_schema);
         FlatWriteFormat {
             arrow_schema,
             override_sequence: None,
-            metric_value_split_columns,
         }
     }
 
@@ -102,11 +96,9 @@ impl FlatWriteFormat {
     }
 
     /// Convert `batch` to a arrow record batch to store in parquet.
-    pub(crate) fn convert_batch(&self, batch: &RecordBatch) -> Result<RecordBatch> {
+    pub(crate) fn convert_batch(&self, batch: RecordBatch) -> Result<RecordBatch> {
         debug_assert_eq!(batch.num_columns(), self.arrow_schema.fields().len());
 
-        let batch = split_metric_value_columns(batch, &self.metric_value_split_columns)?;
-
         let Some(override_sequence) = self.override_sequence else {
             return Ok(batch);
         };

src/mito2/src/sst/parquet/format.rs

@@ -1288,7 +1288,7 @@ mod tests {
         let expect_record =
             RecordBatch::try_new(build_test_flat_sst_schema_with_field_ids(), columns).unwrap();
 
-        let actual = format.convert_batch(&batch).unwrap();
+        let actual = format.convert_batch(batch).unwrap();
         assert_eq!(expect_record, actual);
     }
 
@@ -1319,7 +1319,7 @@ mod tests {
         )
         .unwrap();
 
-        let actual = format.convert_batch(&batch).unwrap();
+        let actual = format.convert_batch(batch).unwrap();
         assert_eq!(expected_record, actual);
     }
 

src/mito2/src/sst/parquet/writer.rs

@@ -57,7 +57,9 @@ use crate::sst::file::RegionFileId;
 use crate::sst::index::{IndexOutput, Indexer, IndexerBuilder};
 use crate::sst::parquet::flat_format::{FlatWriteFormat, time_index_column_index};
 use crate::sst::parquet::format::PrimaryKeyWriteFormat;
-use crate::sst::parquet::metric_value_split::metric_value_split_columns;
+use crate::sst::parquet::metric_value_split::{
+    MetricValueSplitColumn, metric_value_split_columns, split_metric_value_columns,
+};
 use crate::sst::parquet::{PARQUET_METADATA_KEY, SstInfo, WriteOptions};
 use crate::sst::{
     DEFAULT_WRITE_BUFFER_SIZE, DEFAULT_WRITE_CONCURRENCY, FlatSchemaOptions, SeriesEstimator,
@@ -75,11 +77,11 @@ enum FlatBatchConverter {
 }
 
 impl FlatBatchConverter {
-    fn convert_batch(&self, batch: &RecordBatch) -> Result<RecordBatch> {
+    fn convert_batch(&self, batch: RecordBatch) -> Result<RecordBatch> {
         match self {
             FlatBatchConverter::Flat(f) => f.convert_batch(batch),
             FlatBatchConverter::PrimaryKey { format, num_fields } => {
-                format.convert_flat_batch(batch, *num_fields)
+                format.convert_flat_batch(&batch, *num_fields)
             }
         }
     }
@@ -338,9 +340,10 @@ where
     ) -> Result<SstInfoArray> {
         let mut results = smallvec![];
         let mut stats = SourceStats::default();
+        let mut split_columns = None;
 
         while let Some(record_batch) = self
-            .write_next_flat_batch(&mut source, converter, opts)
+            .write_next_flat_batch(&mut source, converter, &mut split_columns, opts)
             .await
             .transpose()
         {
@@ -420,6 +423,7 @@ where
         &mut self,
         source: &mut FlatSource,
         converter: &FlatBatchConverter,
+        split_columns: &mut Option<Vec<MetricValueSplitColumn>>,
         opts: &WriteOptions,
     ) -> Result<Option<RecordBatch>> {
         let start = Instant::now();
@@ -428,7 +432,11 @@ where
         };
         self.metrics.iter_source += start.elapsed();
 
-        let arrow_batch = converter.convert_batch(&record_batch)?;
+        let split_columns = split_columns.get_or_insert_with(|| {
+            metric_value_split_columns(&self.metadata, record_batch.schema_ref())
+        });
+        let split_batch = split_metric_value_columns(&record_batch, split_columns)?;
+        let arrow_batch = converter.convert_batch(split_batch)?;
 
         let start = Instant::now();
         self.maybe_init_writer(arrow_batch.schema_ref(), opts)