diff --git a/src/metric-engine/src/engine/put.rs b/src/metric-engine/src/engine/put.rs index 60f7cc8cd0..d9462c60df 100644 --- a/src/metric-engine/src/engine/put.rs +++ b/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] diff --git a/src/mito2/src/sst/parquet/flat_format.rs b/src/mito2/src/sst/parquet/flat_format.rs index 50af9152ee..b6cae690da 100644 --- a/src/mito2/src/sst/parquet/flat_format.rs +++ b/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, - metric_value_split_columns: Vec, } 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 { + pub(crate) fn convert_batch(&self, batch: RecordBatch) -> Result { 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); }; diff --git a/src/mito2/src/sst/parquet/format.rs b/src/mito2/src/sst/parquet/format.rs index a6338c1a1c..972cda83de 100644 --- a/src/mito2/src/sst/parquet/format.rs +++ b/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); } diff --git a/src/mito2/src/sst/parquet/writer.rs b/src/mito2/src/sst/parquet/writer.rs index cde4d6428d..01712b4b4f 100644 --- a/src/mito2/src/sst/parquet/writer.rs +++ b/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 { + fn convert_batch(&self, batch: RecordBatch) -> Result { 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 { 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>, opts: &WriteOptions, ) -> Result> { 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)