From c79581ecf8b0aaad27fb03c65acf2fc91a527b22 Mon Sep 17 00:00:00 2001 From: Ruihang Xia Date: Mon, 13 Jul 2026 20:16:39 +0800 Subject: [PATCH] refactor(mito2): own split metric value reads Signed-off-by: Ruihang Xia --- Cargo.lock | 1 - src/metric-engine/Cargo.toml | 1 - src/metric-engine/src/data_region.rs | 8 +- src/metric-engine/src/engine/bulk_insert.rs | 5 +- src/metric-engine/src/engine/put.rs | 34 +- src/metric-engine/src/engine/read.rs | 398 ++++++------------ src/metric-engine/src/lib.rs | 1 - src/mito2/AGENTS.md | 1 + src/mito2/src/engine.rs | 17 +- src/mito2/src/lib.rs | 1 + src/mito2/src/metric_value.rs | 104 +++++ src/mito2/src/read.rs | 1 + src/mito2/src/read/scan_region.rs | 12 + .../src => mito2/src/read}/value_split.rs | 208 ++++++++- .../src/sst/parquet/metric_value_split.rs | 50 +-- 15 files changed, 495 insertions(+), 347 deletions(-) create mode 100644 src/mito2/src/metric_value.rs rename src/{metric-engine/src => mito2/src/read}/value_split.rs (58%) diff --git a/Cargo.lock b/Cargo.lock index da503213c1e..e00e30e28a6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -8332,7 +8332,6 @@ dependencies = [ "common-wal", "criterion 0.7.0", "datafusion", - "datafusion-physical-expr", "datatypes", "futures-util", "fxhash", diff --git a/src/metric-engine/Cargo.toml b/src/metric-engine/Cargo.toml index 27f6b0a7c14..0005a4dc543 100644 --- a/src/metric-engine/Cargo.toml +++ b/src/metric-engine/Cargo.toml @@ -26,7 +26,6 @@ common-telemetry.workspace = true common-time.workspace = true chrono.workspace = true datafusion.workspace = true -datafusion-physical-expr.workspace = true datatypes.workspace = true futures-util.workspace = true humantime-serde.workspace = true diff --git a/src/metric-engine/src/data_region.rs b/src/metric-engine/src/data_region.rs index 48976619c93..d8e76db9b65 100644 --- a/src/metric-engine/src/data_region.rs +++ b/src/metric-engine/src/data_region.rs @@ -101,7 +101,7 @@ impl DataRegion { // retrieve underlying version let region_metadata = self .mito - .get_metadata(region_id) + .get_physical_metadata(region_id) .await .context(MitoReadOperationSnafu)?; @@ -214,7 +214,7 @@ impl DataRegion { let data_region_id = utils::to_data_region_id(physical_region_id); let metadata = self .mito - .get_metadata(data_region_id) + .get_physical_metadata(data_region_id) .await .context(MitoReadOperationSnafu)?; Ok(metadata.column_metadatas.clone()) @@ -268,7 +268,7 @@ mod test { let current_version = env .mito() - .get_metadata(utils::to_data_region_id(env.default_physical_region_id())) + .get_physical_metadata(utils::to_data_region_id(env.default_physical_region_id())) .await .unwrap() .schema_version; @@ -306,7 +306,7 @@ mod test { let new_metadata = env .mito() - .get_metadata(utils::to_data_region_id(env.default_physical_region_id())) + .get_physical_metadata(utils::to_data_region_id(env.default_physical_region_id())) .await .unwrap(); let column_names = new_metadata diff --git a/src/metric-engine/src/engine/bulk_insert.rs b/src/metric-engine/src/engine/bulk_insert.rs index 8ae8d567fb6..4b355106ea1 100644 --- a/src/metric-engine/src/engine/bulk_insert.rs +++ b/src/metric-engine/src/engine/bulk_insert.rs @@ -25,7 +25,6 @@ use snafu::{OptionExt, ResultExt, ensure}; use store_api::codec::PrimaryKeyEncoding; use store_api::metadata::RegionMetadataRef; use store_api::metric_engine_consts::is_metric_engine_value_int_column; -use store_api::region_engine::RegionEngine; use store_api::region_request::{AffectedRows, RegionBulkInsertsRequest, RegionRequest}; use store_api::storage::RegionId; use store_api::storage::consts::PRIMARY_KEY_COLUMN_NAME; @@ -163,7 +162,7 @@ impl MetricEngineInner { ) -> Result { let metadata = self .mito - .get_metadata(data_region_id) + .get_physical_metadata(data_region_id) .await .context(error::MitoReadOperationSnafu)?; let batch_schema = batch.schema(); @@ -222,7 +221,7 @@ impl MetricEngineInner { async fn physical_schema_version(&self, region_id: RegionId) -> Result { Ok(self .mito - .get_metadata(region_id) + .get_physical_metadata(region_id) .await .context(error::MitoReadOperationSnafu)? .schema_version) diff --git a/src/metric-engine/src/engine/put.rs b/src/metric-engine/src/engine/put.rs index d9462c60dff..ea62279159d 100644 --- a/src/metric-engine/src/engine/put.rs +++ b/src/metric-engine/src/engine/put.rs @@ -1251,6 +1251,7 @@ mod tests { .unwrap(); let physical_region_id = env.default_physical_region_id(); + let data_region_id = to_data_region_id(physical_region_id); let visible_physical_stream = env .metric() .scan_to_stream(physical_region_id, ScanRequest::default()) @@ -1258,6 +1259,24 @@ mod tests { .unwrap(); let int_column_name = store_api::metric_engine_consts::metric_engine_value_int_column_name(greptime_value()); + let visible_mito_metadata = env.mito().get_metadata(data_region_id).await.unwrap(); + let physical_mito_metadata = env + .mito() + .get_physical_metadata(data_region_id) + .await + .unwrap(); + assert!( + visible_mito_metadata + .column_by_name(&int_column_name) + .is_none(), + "mito query metadata should hide split companion columns" + ); + assert!( + physical_mito_metadata + .column_by_name(&int_column_name) + .is_some(), + "mito physical metadata should retain split companion columns" + ); assert!( visible_physical_stream .schema() @@ -1266,6 +1285,20 @@ mod tests { "metric physical reads should hide split companion column" ); + let logical_scanner = env + .metric() + .handle_query(logical_region_id, ScanRequest::default()) + .await + .unwrap(); + assert_eq!(logical_scanner.metadata().region_id, logical_region_id); + assert!( + logical_scanner + .metadata() + .column_by_name(&int_column_name) + .is_none(), + "logical scanner metadata should not expose split companion columns" + ); + let logical_batches = RecordBatches::try_collect( env.metric() .scan_to_stream(logical_region_id, ScanRequest::default()) @@ -1352,7 +1385,6 @@ mod tests { ("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()) diff --git a/src/metric-engine/src/engine/read.rs b/src/metric-engine/src/engine/read.rs index f5dd2a32cf1..94cb43e7b8d 100644 --- a/src/metric-engine/src/engine/read.rs +++ b/src/metric-engine/src/engine/read.rs @@ -12,27 +12,24 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::collections::HashSet; +use std::fmt; use std::sync::Arc; use api::v1::SemanticType; +use common_error::ext::BoxedError; +use common_recordbatch::SendableRecordBatchStream; use common_telemetry::{debug, error, tracing}; -use datafusion::arrow::datatypes::DataType as ArrowDataType; -use datafusion::common::tree_node::{Transformed, TreeNode, TreeNodeRecursion, TreeNodeRewriter}; -use datafusion::common::{Column, Result as DataFusionResult}; -use datafusion::functions::expr_fn::coalesce; -use datafusion::logical_expr::expr_fn::cast; -use datafusion::logical_expr::utils::expr_to_columns; -use datafusion::logical_expr::{self, Expr}; -use datatypes::prelude::ConcreteDataType; +use datafusion::logical_expr; +use datafusion::physical_expr::PhysicalExpr; +use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet; +use datafusion::physical_plan::{DisplayAs, DisplayFormatType}; use snafu::{OptionExt, ResultExt}; use store_api::metadata::{RegionMetadataBuilder, RegionMetadataRef}; -use store_api::metric_engine_consts::{ - DATA_SCHEMA_TABLE_ID_COLUMN_NAME, is_metric_engine_value_int_column, - metric_engine_value_int_column_name, +use store_api::metric_engine_consts::DATA_SCHEMA_TABLE_ID_COLUMN_NAME; +use store_api::region_engine::{ + PrepareRequest, QueryScanContext, RegionEngine, RegionScanner, RegionScannerRef, + ScannerProperties, }; -use store_api::region_engine::{RegionEngine, RegionScannerRef}; -use store_api::region_request::AlterKind; use store_api::storage::{RegionId, ScanRequest, SequenceNumber}; use crate::engine::MetricEngineInner; @@ -42,7 +39,6 @@ use crate::error::{ }; use crate::metrics::MITO_OPERATION_ELAPSED; use crate::utils; -use crate::value_split::{ValueColumnProjection, ValueSplitProjectionMapper, ValueSplitScanner}; impl MetricEngineInner { #[tracing::instrument(skip_all)] @@ -73,32 +69,10 @@ impl MetricEngineInner { .with_label_values(&["read_physical"]) .start_timer(); - let data_region_id = utils::to_data_region_id(region_id); - let physical_metadata = self - .mito - .get_metadata(data_region_id) + self.mito + .handle_query(region_id, request) .await - .context(MitoReadOperationSnafu)?; - let visible_metadata = visible_physical_region_metadata(&physical_metadata)?; - let (request, mapper) = self.transform_request_with_mapper( - data_region_id, - request, - &visible_metadata, - &physical_metadata, - None, - )?; - - let scanner = self - .mito - .handle_query(data_region_id, request) - .await - .context(MitoReadOperationSnafu)?; - - Ok(Box::new(ValueSplitScanner::new( - scanner, - visible_metadata, - mapper, - ))) + .context(MitoReadOperationSnafu) } async fn read_logical_region( @@ -115,8 +89,8 @@ impl MetricEngineInner { let logical_metadata = self .logical_region_metadata(physical_region_id, logical_region_id) .await?; - let (request, mapper) = self - .transform_logical_request_with_mapper( + let request = self + .transform_request_with_metadata( physical_region_id, logical_region_id, request, @@ -131,10 +105,9 @@ impl MetricEngineInner { scanner.set_logical_region(true); scanner.set_query_load_region_id(data_region_id); - Ok(Box::new(ValueSplitScanner::new( + Ok(Box::new(LogicalRegionScanner::new( scanner, logical_metadata, - mapper, ))) } @@ -156,13 +129,10 @@ impl MetricEngineInner { self.state.read().unwrap().exist_physical_region(region_id); if is_reading_physical_region { - let data_region_id = utils::to_data_region_id(region_id); - let physical_metadata = self - .mito - .get_metadata(data_region_id) + self.mito + .get_metadata(region_id) .await - .context(MitoReadOperationSnafu)?; - visible_physical_region_metadata(&physical_metadata) + .context(MitoReadOperationSnafu) } else { let physical_region_id = self.get_physical_region_id(region_id).await?; self.logical_region_metadata(physical_region_id, region_id) @@ -198,152 +168,60 @@ impl MetricEngineInner { let logical_metadata = self .logical_region_metadata(physical_region_id, logical_region_id) .await?; - self.transform_logical_request_with_mapper( + self.transform_request_with_metadata( physical_region_id, logical_region_id, request, &logical_metadata, ) .await - .map(|(request, _)| request) } - async fn transform_logical_request_with_mapper( + async fn transform_request_with_metadata( &self, physical_region_id: RegionId, logical_region_id: RegionId, - request: ScanRequest, + mut request: ScanRequest, logical_metadata: &RegionMetadataRef, - ) -> Result<(ScanRequest, ValueSplitProjectionMapper)> { + ) -> Result { + let logical_projection = match request.projection_input.as_ref() { + Some(projection_input) => projection_input.projection.clone(), + None => (0..logical_metadata.column_metadatas.len()).collect(), + }; let data_region_id = utils::to_data_region_id(physical_region_id); let physical_metadata = self .mito .get_metadata(data_region_id) .await .context(MitoReadOperationSnafu)?; - self.transform_request_with_mapper( - logical_region_id, - request, - logical_metadata, - &physical_metadata, - Some(logical_region_id), - ) - } - fn transform_request_with_mapper( - &self, - region_id: RegionId, - mut request: ScanRequest, - visible_metadata: &RegionMetadataRef, - physical_metadata: &RegionMetadataRef, - logical_region_id: Option, - ) -> Result<(ScanRequest, ValueSplitProjectionMapper)> { - let split_value_columns = split_value_columns(visible_metadata, physical_metadata); - let mut residual_column_names = HashSet::new(); - let residual_filters = request - .filters - .iter() - .filter(|filter| { - let mut columns = HashSet::new(); - let is_residual = expr_to_columns(filter, &mut columns).is_ok() - && columns - .iter() - .any(|column| split_value_columns.contains(&column.name)); - if is_residual { - residual_column_names.extend(columns.into_iter().map(|column| column.name)); - } - is_residual - }) - .cloned() - .collect::>(); - let mut visible_projection = match request.projection_input.as_ref() { - Some(projection_input) => projection_input.projection.clone(), - None => (0..visible_metadata.column_metadatas.len()).collect(), - }; - let visible_columns = visible_projection.len(); - let mut projected = visible_projection.iter().copied().collect::>(); - for (index, column) in visible_metadata.column_metadatas.iter().enumerate() { - if residual_column_names.contains(&column.column_schema.name) && projected.insert(index) - { - visible_projection.push(index); - } + let mut physical_projection = Vec::with_capacity(logical_projection.len()); + for logical_index in logical_projection { + let logical_column = logical_metadata + .column_metadatas + .get(logical_index) + .with_context(|| InvalidRequestSnafu { + region_id: logical_region_id, + reason: format!("projection index {logical_index} is out of bound"), + })?; + let name = &logical_column.column_schema.name; + let physical_index = + physical_metadata + .column_index_by_name(name) + .with_context(|| InvalidRequestSnafu { + region_id: logical_region_id, + reason: format!("column {name} is missing from physical metadata"), + })?; + physical_projection.push(physical_index); } - let (physical_projection, mapper) = self.transform_projection_with_mapper( - region_id, - &visible_projection, - visible_metadata, - physical_metadata, - visible_columns, - residual_filters, - )?; // Top-level projections are indices; nested paths are column names. request.projection_input.get_or_insert_default().projection = physical_projection; - request.filters = request - .filters - .into_iter() - .map(|filter| rewrite_metric_value_filter(region_id, filter, &split_value_columns)) - .collect::>>()?; - if let Some(logical_region_id) = logical_region_id { - request.filters.push( - logical_expr::col(DATA_SCHEMA_TABLE_ID_COLUMN_NAME) - .eq(logical_expr::lit(logical_region_id.table_id())), - ); - } - - Ok((request, mapper)) - } - - fn transform_projection_with_mapper( - &self, - logical_region_id: RegionId, - origin_projection: &[usize], - logical_metadata: &RegionMetadataRef, - physical_metadata: &RegionMetadataRef, - visible_columns: usize, - residual_filters: Vec, - ) -> Result<(Vec, ValueSplitProjectionMapper)> { - let mut physical_projection = Vec::with_capacity(origin_projection.len()); - let mut output_columns = Vec::with_capacity(origin_projection.len()); - - for logical_idx in origin_projection { - let logical_column = logical_metadata - .column_metadatas - .get(*logical_idx) - .with_context(|| InvalidRequestSnafu { - region_id: logical_region_id, - reason: format!("projection index {} is out of bound", logical_idx), - })?; - let name = &logical_column.column_schema.name; - // Safety: logical columns is a strict subset of physical columns - let float_index = physical_metadata.column_index_by_name(name).unwrap(); - let input_float_index = physical_projection.len(); - physical_projection.push(float_index); - - if logical_column.semantic_type == SemanticType::Field - && logical_column.column_schema.data_type == ConcreteDataType::float64_datatype() - && let Some(int_index) = physical_metadata - .column_index_by_name(&metric_engine_value_int_column_name(name)) - { - let input_int_index = physical_projection.len(); - physical_projection.push(int_index); - output_columns.push(ValueColumnProjection::Split { - float_index: input_float_index, - int_index: input_int_index, - output_schema: logical_column.column_schema.clone(), - }); - } else { - output_columns.push(ValueColumnProjection::Direct { - input_index: input_float_index, - output_schema: logical_column.column_schema.clone(), - }); - } - } - - Ok(( - physical_projection, - ValueSplitProjectionMapper::new(output_columns, visible_columns, residual_filters), - )) + request.filters.push( + logical_expr::col(DATA_SCHEMA_TABLE_ID_COLUMN_NAME) + .eq(logical_expr::lit(logical_region_id.table_id())), + ); + Ok(request) } pub async fn logical_region_metadata( @@ -379,116 +257,82 @@ impl MetricEngineInner { } } -fn visible_physical_region_metadata( - physical_metadata: &RegionMetadataRef, -) -> Result { - let visible_columns = physical_metadata - .column_metadatas - .iter() - .filter(|column| !is_metric_engine_value_int_column(&column.column_schema.name)) - .cloned() - .collect::>(); - if visible_columns.len() == physical_metadata.column_metadatas.len() { - return Ok(physical_metadata.clone()); +struct LogicalRegionScanner { + inner: RegionScannerRef, + metadata: RegionMetadataRef, +} + +impl LogicalRegionScanner { + fn new(inner: RegionScannerRef, metadata: RegionMetadataRef) -> Self { + Self { inner, metadata } + } +} + +impl fmt::Debug for LogicalRegionScanner { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("LogicalRegionScanner") + .field("inner", &self.inner) + .finish() + } +} + +impl DisplayAs for LogicalRegionScanner { + fn fmt_as(&self, t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { + self.inner.fmt_as(t, f) + } +} + +impl RegionScanner for LogicalRegionScanner { + fn name(&self) -> &str { + self.inner.name() } - let primary_key = physical_metadata.primary_key.clone(); - let mut builder = RegionMetadataBuilder::from_existing((**physical_metadata).clone()); - builder - .alter(AlterKind::SyncColumns { - column_metadatas: visible_columns, - }) - .context(InvalidMetadataSnafu)?; - builder.primary_key(primary_key); - builder.build().map(Arc::new).context(InvalidMetadataSnafu) -} - -fn split_value_columns( - logical_metadata: &RegionMetadataRef, - physical_metadata: &RegionMetadataRef, -) -> HashSet { - logical_metadata - .column_metadatas - .iter() - .filter(|column| { - column.semantic_type == SemanticType::Field - && column.column_schema.data_type == ConcreteDataType::float64_datatype() - }) - .filter_map(|column| { - let value_name = &column.column_schema.name; - physical_metadata - .column_by_name(&metric_engine_value_int_column_name(value_name)) - .filter(|int_column| { - int_column.semantic_type == SemanticType::Field - && int_column.column_schema.data_type == ConcreteDataType::int64_datatype() - }) - .map(|_| value_name.clone()) - }) - .collect() -} - -fn rewrite_metric_value_filter( - logical_region_id: RegionId, - filter: Expr, - split_value_columns: &HashSet, -) -> Result { - if split_value_columns.is_empty() { - return Ok(filter); + fn properties(&self) -> &ScannerProperties { + self.inner.properties() } - let filter_display = filter.to_string(); - let mut rewriter = MetricValueFilterRewriter { - split_value_columns, - }; - filter - .rewrite(&mut rewriter) - .map(|rewritten| rewritten.data) - .map_err(|err| { - InvalidRequestSnafu { - region_id: logical_region_id, - reason: format!("failed to rewrite metric value filter {filter_display}: {err}"), - } - .build() - }) -} - -struct MetricValueFilterRewriter<'a> { - split_value_columns: &'a HashSet, -} - -impl TreeNodeRewriter for MetricValueFilterRewriter<'_> { - type Node = Expr; - - fn f_down(&mut self, expr: Expr) -> DataFusionResult> { - let recursion = if matches!( - expr, - Expr::Exists(_) | Expr::InSubquery(_) | Expr::ScalarSubquery(_) - ) { - TreeNodeRecursion::Jump - } else { - TreeNodeRecursion::Continue - }; - - Ok(Transformed::new(expr, false, recursion)) + fn schema(&self) -> datatypes::schema::SchemaRef { + self.inner.schema() } - fn f_up(&mut self, expr: Expr) -> DataFusionResult> { - let Expr::Column(column) = expr else { - return Ok(Transformed::no(expr)); - }; + fn metadata(&self) -> RegionMetadataRef { + self.metadata.clone() + } - if !self.split_value_columns.contains(&column.name) { - return Ok(Transformed::no(Expr::Column(column))); - } + fn prepare(&mut self, request: PrepareRequest) -> std::result::Result<(), BoxedError> { + self.inner.prepare(request) + } - let int_column = Column { - relation: column.relation.clone(), - name: metric_engine_value_int_column_name(&column.name), - spans: column.spans.clone(), - }; - let float_expr = Expr::Column(column); - let int_expr = cast(Expr::Column(int_column), ArrowDataType::Float64); - Ok(Transformed::yes(coalesce(vec![int_expr, float_expr]))) + fn scan_partition( + &self, + ctx: &QueryScanContext, + metrics_set: &ExecutionPlanMetricsSet, + partition: usize, + ) -> std::result::Result { + self.inner.scan_partition(ctx, metrics_set, partition) + } + + fn has_predicate_without_region(&self) -> bool { + self.inner.has_predicate_without_region() + } + + fn add_dyn_filter_to_predicate( + &mut self, + filter_exprs: Vec>, + ) -> Vec { + self.inner.add_dyn_filter_to_predicate(filter_exprs) + } + + fn set_logical_region(&mut self, logical_region: bool) { + self.inner.set_logical_region(logical_region); + } + + fn set_query_load_region_id(&mut self, region_id: RegionId) { + self.inner.set_query_load_region_id(region_id); + } + + fn snapshot_sequence(&self) -> Option { + self.inner.snapshot_sequence() } } @@ -570,7 +414,7 @@ mod test { assert_eq!( scan_req.projection_indices().unwrap(), - &[12, 11, 10, 9, 0, 1, 2, 5] + &[11, 10, 9, 8, 0, 1, 4] ); assert_eq!(scan_req.filters.len(), 1); assert_eq!( @@ -589,7 +433,7 @@ mod test { .unwrap(); assert_eq!( scan_req.projection_indices().unwrap(), - &[12, 11, 10, 9, 0, 1, 2, 5] + &[11, 10, 9, 8, 0, 1, 4] ); } } diff --git a/src/metric-engine/src/lib.rs b/src/metric-engine/src/lib.rs index 5da7a1d554e..d209eb7588f 100644 --- a/src/metric-engine/src/lib.rs +++ b/src/metric-engine/src/lib.rs @@ -64,4 +64,3 @@ pub mod row_modifier; #[cfg(test)] mod test_util; mod utils; -mod value_split; diff --git a/src/mito2/AGENTS.md b/src/mito2/AGENTS.md index c25ede60ed0..c5ca1e497ce 100644 --- a/src/mito2/AGENTS.md +++ b/src/mito2/AGENTS.md @@ -22,6 +22,7 @@ snapshot isolation). It implements the `RegionEngine` trait from `store-api`. | `request` | `src/mito2/src/request.rs` | `WriteRequest`/`RegionRequest` types and result channels | | `wal` | `src/mito2/src/wal.rs` | Write-ahead log wrapper over `log-store` | | `memtable` | `src/mito2/src/memtable/` | In-memory write buffers (time-series / bulk / partition) | +| `metric_value` | `src/mito2/src/metric_value.rs` | Metric float/int companion discovery and logical schema visibility | | `flush` | `src/mito2/src/flush.rs` | `FlushScheduler`, `WriteBufferManager`, memtable → SST | | `compaction` | `src/mito2/src/compaction/` | TWCS picker, strict-window manual picker, compactor, memory control | | `access_layer` | `src/mito2/src/access_layer.rs` | SST read/write over the object store | diff --git a/src/mito2/src/engine.rs b/src/mito2/src/engine.rs index 94911814b56..0845ab910c6 100644 --- a/src/mito2/src/engine.rs +++ b/src/mito2/src/engine.rs @@ -142,6 +142,7 @@ use crate::extension::BoxedExtensionRangeProviderFactory; use crate::gc::GcLimiterRef; use crate::manifest::action::RegionEdit; use crate::memtable::MemtableStats; +use crate::metric_value::visible_region_metadata; use crate::metrics::{ HANDLE_REQUEST_ELAPSED, SCAN_MEMORY_EXHAUSTED_TOTAL, SCAN_MEMORY_USAGE_BYTES, SCAN_REQUESTS_REJECTED_TOTAL, @@ -401,6 +402,14 @@ impl MitoEngine { .map(|r| r.primary_key_encoding()) } + /// Returns persisted region metadata, including engine-internal columns. + pub async fn get_physical_metadata( + &self, + region_id: RegionId, + ) -> std::result::Result { + self.inner.get_metadata(region_id).map_err(BoxedError::new) + } + /// Handle substrait query and return a stream of record batches /// /// Notice that the output stream's ordering is not guranateed. If order @@ -1258,7 +1267,7 @@ impl RegionEngine for MitoEngine { ) -> Result { self.scan_region(region_id, request) .map_err(BoxedError::new)? - .region_scanner() + .query_region_scanner() .await .map_err(BoxedError::new) } @@ -1281,7 +1290,11 @@ impl RegionEngine for MitoEngine { &self, region_id: RegionId, ) -> std::result::Result { - self.inner.get_metadata(region_id).map_err(BoxedError::new) + let metadata = self + .inner + .get_metadata(region_id) + .map_err(BoxedError::new)?; + visible_region_metadata(&metadata).map_err(BoxedError::new) } /// Stop the engine. diff --git a/src/mito2/src/lib.rs b/src/mito2/src/lib.rs index a452e106c7a..07f91d96e0b 100644 --- a/src/mito2/src/lib.rs +++ b/src/mito2/src/lib.rs @@ -36,6 +36,7 @@ pub mod flush; pub mod gc; pub mod manifest; pub mod memtable; +mod metric_value; mod metrics; pub mod read; pub mod region; diff --git a/src/mito2/src/metric_value.rs b/src/mito2/src/metric_value.rs new file mode 100644 index 00000000000..fc812322b26 --- /dev/null +++ b/src/mito2/src/metric_value.rs @@ -0,0 +1,104 @@ +// 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::HashSet; +use std::sync::Arc; + +use api::v1::SemanticType; +use datatypes::prelude::ConcreteDataType; +use snafu::ResultExt; +use store_api::metadata::{RegionMetadata, RegionMetadataBuilder, RegionMetadataRef}; +use store_api::metric_engine_consts::{ + DATA_SCHEMA_TABLE_ID_COLUMN_NAME, DATA_SCHEMA_TSID_COLUMN_NAME, + metric_engine_value_int_column_name, +}; +use store_api::region_request::AlterKind; +use store_api::storage::consts::ReservedColumnId; + +use crate::error::{InvalidMetadataSnafu, Result}; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct MetricValueColumn { + pub(crate) value_index: usize, + pub(crate) int_index: usize, +} + +pub(crate) fn metric_value_columns(metadata: &RegionMetadata) -> Vec { + if !is_metric_engine_data_region(metadata) { + return Vec::new(); + } + + metadata + .field_columns() + .filter(|column| column.column_schema.data_type == ConcreteDataType::float64_datatype()) + .filter_map(|value_column| { + let int_name = metric_engine_value_int_column_name(&value_column.column_schema.name); + let int_column = metadata.column_by_name(&int_name)?; + if int_column.semantic_type != SemanticType::Field + || int_column.column_schema.data_type != ConcreteDataType::int64_datatype() + { + return None; + } + + Some(MetricValueColumn { + value_index: metadata.column_index_by_name(&value_column.column_schema.name)?, + int_index: metadata.column_index_by_name(&int_name)?, + }) + }) + .collect() +} + +pub(crate) fn visible_region_metadata(metadata: &RegionMetadataRef) -> Result { + let split_columns = metric_value_columns(metadata); + if split_columns.is_empty() { + return Ok(metadata.clone()); + } + + let int_indices = split_columns + .into_iter() + .map(|column| column.int_index) + .collect::>(); + let visible_columns = metadata + .column_metadatas + .iter() + .enumerate() + .filter(|(index, _)| !int_indices.contains(index)) + .map(|(_, column)| column.clone()) + .collect::>(); + + let primary_key = metadata.primary_key.clone(); + let mut builder = RegionMetadataBuilder::from_existing((**metadata).clone()); + builder + .alter(AlterKind::SyncColumns { + column_metadatas: visible_columns, + }) + .context(InvalidMetadataSnafu)?; + builder.primary_key(primary_key); + builder.build().map(Arc::new).context(InvalidMetadataSnafu) +} + +pub(crate) fn is_metric_engine_data_region(metadata: &RegionMetadata) -> bool { + let has_internal_tag = |name, column_id| { + metadata.column_by_name(name).is_some_and(|column| { + column.semantic_type == SemanticType::Tag + && column.column_id == column_id + && metadata.primary_key.contains(&column_id) + }) + }; + + has_internal_tag( + DATA_SCHEMA_TABLE_ID_COLUMN_NAME, + ReservedColumnId::table_id(), + ) && has_internal_tag(DATA_SCHEMA_TSID_COLUMN_NAME, ReservedColumnId::tsid()) +} diff --git a/src/mito2/src/read.rs b/src/mito2/src/read.rs index 192bc92b167..6d86dd0d225 100644 --- a/src/mito2/src/read.rs +++ b/src/mito2/src/read.rs @@ -36,6 +36,7 @@ pub(crate) mod seq_scan; pub mod series_scan; pub mod stream; pub(crate) mod unordered_scan; +pub(crate) mod value_split; use std::collections::HashMap; use std::sync::Arc; diff --git a/src/mito2/src/read/scan_region.rs b/src/mito2/src/read/scan_region.rs index 9e1f3b3acf0..d7b6da750f2 100644 --- a/src/mito2/src/read/scan_region.rs +++ b/src/mito2/src/read/scan_region.rs @@ -72,6 +72,7 @@ use crate::read::seq_scan::SeqScan; use crate::read::series_scan::SeriesScan; use crate::read::stream::ScanBatchStream; use crate::read::unordered_scan::UnorderedScan; +use crate::read::value_split::{ValueSplitScanner, prepare_value_split_scan}; use crate::read::{BoxedRecordBatchStream, RecordBatch}; use crate::region::options::MergeMode; use crate::region::version::VersionRef; @@ -373,6 +374,17 @@ impl ScanRegion { } } + /// Returns a public query scanner that hides metric value companion columns. + pub(crate) async fn query_region_scanner(mut self) -> Result { + let mapper = + prepare_value_split_scan(self.region_id(), &mut self.request, &self.version.metadata)?; + let scanner = self.region_scanner().await?; + match mapper { + Some(mapper) => Ok(Box::new(ValueSplitScanner::new(scanner, mapper))), + None => Ok(scanner), + } + } + /// Scan sequentially. #[tracing::instrument(skip_all, fields(region_id = %self.region_id()))] pub(crate) async fn seq_scan(self) -> Result { diff --git a/src/metric-engine/src/value_split.rs b/src/mito2/src/read/value_split.rs similarity index 58% rename from src/metric-engine/src/value_split.rs rename to src/mito2/src/read/value_split.rs index 2928ae26681..2c3d9ca4dfe 100644 --- a/src/metric-engine/src/value_split.rs +++ b/src/mito2/src/read/value_split.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::collections::HashSet; use std::fmt; use std::pin::Pin; use std::sync::Arc; @@ -26,23 +27,200 @@ use common_recordbatch::filter::batch_filter; use common_recordbatch::{ DfRecordBatch, OrderOption, RecordBatch, RecordBatchStream, SendableRecordBatchStream, }; -use datafusion::common::ToDFSchema; +use datafusion::arrow::datatypes::DataType as ArrowDataType; +use datafusion::common::tree_node::{Transformed, TreeNode, TreeNodeRecursion, TreeNodeRewriter}; +use datafusion::common::{Column, Result as DataFusionResult, ToDFSchema}; use datafusion::execution::context::ExecutionProps; +use datafusion::functions::expr_fn::coalesce; use datafusion::logical_expr::Expr; -use datafusion::logical_expr::utils::conjunction; -use datafusion::physical_expr::create_physical_expr; +use datafusion::logical_expr::expr_fn::cast; +use datafusion::logical_expr::utils::{conjunction, expr_to_columns}; +use datafusion::physical_expr::{PhysicalExpr, create_physical_expr}; use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet; use datafusion::physical_plan::{DisplayAs, DisplayFormatType}; -use datafusion_physical_expr::PhysicalExpr; use datatypes::arrow::array::{Array, ArrayRef, Float64Array, Float64Builder, Int64Array}; use datatypes::schema::{ColumnSchema, Schema, SchemaRef}; -use futures_util::Stream; +use futures::Stream; use snafu::{OptionExt, ResultExt}; use store_api::metadata::RegionMetadataRef; +use store_api::metric_engine_consts::metric_engine_value_int_column_name; use store_api::region_engine::{ PrepareRequest, QueryScanContext, RegionScanner, RegionScannerRef, ScannerProperties, }; -use store_api::storage::{RegionId, SequenceNumber}; +use store_api::storage::{RegionId, ScanRequest, SequenceNumber}; + +use crate::error::{InvalidRequestSnafu, Result}; +use crate::metric_value::{metric_value_columns, visible_region_metadata}; + +pub(crate) fn prepare_value_split_scan( + region_id: RegionId, + request: &mut ScanRequest, + physical_metadata: &RegionMetadataRef, +) -> Result> { + let split_columns = metric_value_columns(physical_metadata); + if split_columns.is_empty() { + return Ok(None); + } + + let visible_metadata = visible_region_metadata(physical_metadata)?; + let split_value_columns = split_columns + .into_iter() + .map(|column| { + physical_metadata.column_metadatas[column.value_index] + .column_schema + .name + .clone() + }) + .collect::>(); + + let mut residual_column_names = HashSet::new(); + let residual_filters = request + .filters + .iter() + .filter(|filter| { + let mut columns = HashSet::new(); + let is_residual = expr_to_columns(filter, &mut columns).is_ok() + && columns + .iter() + .any(|column| split_value_columns.contains(&column.name)); + if is_residual { + residual_column_names.extend(columns.into_iter().map(|column| column.name)); + } + is_residual + }) + .cloned() + .collect::>(); + + let mut visible_projection = match request.projection_input.as_ref() { + Some(projection_input) => projection_input.projection.clone(), + None => (0..visible_metadata.column_metadatas.len()).collect(), + }; + let visible_columns = visible_projection.len(); + let mut projected = visible_projection.iter().copied().collect::>(); + for (index, column) in visible_metadata.column_metadatas.iter().enumerate() { + if residual_column_names.contains(&column.column_schema.name) && projected.insert(index) { + visible_projection.push(index); + } + } + + let mut physical_projection = Vec::with_capacity(visible_projection.len()); + let mut output_columns = Vec::with_capacity(visible_projection.len()); + for visible_index in visible_projection { + let visible_column = visible_metadata + .column_metadatas + .get(visible_index) + .with_context(|| InvalidRequestSnafu { + region_id, + reason: format!("projection index {visible_index} is out of bound"), + })?; + let value_name = &visible_column.column_schema.name; + let physical_index = physical_metadata + .column_index_by_name(value_name) + .with_context(|| InvalidRequestSnafu { + region_id, + reason: format!("column {value_name} is missing from physical metadata"), + })?; + let input_value_index = physical_projection.len(); + physical_projection.push(physical_index); + + if split_value_columns.contains(value_name) { + let int_name = metric_engine_value_int_column_name(value_name); + let int_index = physical_metadata + .column_index_by_name(&int_name) + .with_context(|| InvalidRequestSnafu { + region_id, + reason: format!("column {int_name} is missing from physical metadata"), + })?; + let input_int_index = physical_projection.len(); + physical_projection.push(int_index); + output_columns.push(ValueColumnProjection::Split { + float_index: input_value_index, + int_index: input_int_index, + output_schema: visible_column.column_schema.clone(), + }); + } else { + output_columns.push(ValueColumnProjection::Direct { + input_index: input_value_index, + output_schema: visible_column.column_schema.clone(), + }); + } + } + + request.projection_input.get_or_insert_default().projection = physical_projection; + request.filters = request + .filters + .drain(..) + .map(|filter| rewrite_metric_value_filter(region_id, filter, &split_value_columns)) + .collect::>>()?; + + Ok(Some(ValueSplitProjectionMapper::new( + visible_metadata, + output_columns, + visible_columns, + residual_filters, + ))) +} + +fn rewrite_metric_value_filter( + region_id: RegionId, + filter: Expr, + split_value_columns: &HashSet, +) -> Result { + let filter_display = filter.to_string(); + let mut rewriter = MetricValueFilterRewriter { + split_value_columns, + }; + filter + .rewrite(&mut rewriter) + .map(|rewritten| rewritten.data) + .map_err(|err| { + InvalidRequestSnafu { + region_id, + reason: format!("failed to rewrite metric value filter {filter_display}: {err}"), + } + .build() + }) +} + +struct MetricValueFilterRewriter<'a> { + split_value_columns: &'a HashSet, +} + +impl TreeNodeRewriter for MetricValueFilterRewriter<'_> { + type Node = Expr; + + fn f_down(&mut self, expr: Expr) -> DataFusionResult> { + let recursion = if matches!( + expr, + Expr::Exists(_) | Expr::InSubquery(_) | Expr::ScalarSubquery(_) + ) { + TreeNodeRecursion::Jump + } else { + TreeNodeRecursion::Continue + }; + + Ok(Transformed::new(expr, false, recursion)) + } + + fn f_up(&mut self, expr: Expr) -> DataFusionResult> { + let Expr::Column(column) = expr else { + return Ok(Transformed::no(expr)); + }; + + if !self.split_value_columns.contains(&column.name) { + return Ok(Transformed::no(Expr::Column(column))); + } + + let int_column = Column { + relation: column.relation.clone(), + name: metric_engine_value_int_column_name(&column.name), + spans: column.spans.clone(), + }; + let float_expr = Expr::Column(column); + let int_expr = cast(Expr::Column(int_column), ArrowDataType::Float64); + Ok(Transformed::yes(coalesce(vec![int_expr, float_expr]))) + } +} #[derive(Clone)] pub(crate) enum ValueColumnProjection { @@ -59,6 +237,7 @@ pub(crate) enum ValueColumnProjection { #[derive(Clone)] pub(crate) struct ValueSplitProjectionMapper { + metadata: RegionMetadataRef, output_schema: SchemaRef, working_schema: SchemaRef, columns: Vec, @@ -69,6 +248,7 @@ pub(crate) struct ValueSplitProjectionMapper { impl ValueSplitProjectionMapper { pub(crate) fn new( + metadata: RegionMetadataRef, columns: Vec, visible_columns: usize, residual_filters: Vec, @@ -90,6 +270,7 @@ impl ValueSplitProjectionMapper { .collect::>(); Self { + metadata, output_schema: Arc::new(Schema::new(output_columns)), working_schema: Arc::new(Schema::new(working_columns)), columns, @@ -196,21 +377,12 @@ fn coalesce_value_columns(float_col: &ArrayRef, int_col: &ArrayRef) -> RecordBat pub(crate) struct ValueSplitScanner { inner: RegionScannerRef, - logical_metadata: RegionMetadataRef, mapper: ValueSplitProjectionMapper, } impl ValueSplitScanner { - pub(crate) fn new( - inner: RegionScannerRef, - logical_metadata: RegionMetadataRef, - mapper: ValueSplitProjectionMapper, - ) -> Self { - Self { - inner, - logical_metadata, - mapper, - } + pub(crate) fn new(inner: RegionScannerRef, mapper: ValueSplitProjectionMapper) -> Self { + Self { inner, mapper } } } @@ -243,7 +415,7 @@ impl RegionScanner for ValueSplitScanner { } fn metadata(&self) -> RegionMetadataRef { - self.logical_metadata.clone() + self.mapper.metadata.clone() } fn prepare(&mut self, request: PrepareRequest) -> Result<(), BoxedError> { diff --git a/src/mito2/src/sst/parquet/metric_value_split.rs b/src/mito2/src/sst/parquet/metric_value_split.rs index 8dd6f5b43fd..c2150a33c44 100644 --- a/src/mito2/src/sst/parquet/metric_value_split.rs +++ b/src/mito2/src/sst/parquet/metric_value_split.rs @@ -14,23 +14,17 @@ use std::sync::Arc; -use api::v1::SemanticType; use datatypes::arrow::array::{ Array, ArrayRef, Float64Array, Float64Builder, Int64Array, Int64Builder, }; use datatypes::arrow::buffer::{BooleanBuffer, Buffer, NullBuffer, ScalarBuffer}; use datatypes::arrow::datatypes::SchemaRef; use datatypes::arrow::record_batch::RecordBatch; -use datatypes::prelude::ConcreteDataType; use snafu::{OptionExt, ResultExt}; use store_api::metadata::RegionMetadata; -use store_api::metric_engine_consts::{ - DATA_SCHEMA_TABLE_ID_COLUMN_NAME, DATA_SCHEMA_TSID_COLUMN_NAME, - metric_engine_value_int_column_name, -}; -use store_api::storage::consts::ReservedColumnId; use crate::error::{InvalidRecordBatchSnafu, NewRecordBatchSnafu, Result}; +use crate::metric_value::metric_value_columns; #[derive(Debug, Clone)] pub(crate) struct MetricValueSplitColumn { @@ -42,26 +36,15 @@ pub(crate) fn metric_value_split_columns( metadata: &RegionMetadata, arrow_schema: &SchemaRef, ) -> Vec { - if !is_metric_engine_data_region(metadata) { - return vec![]; - } - - metadata - .field_columns() - .filter(|column| column.column_schema.data_type == ConcreteDataType::float64_datatype()) - .filter_map(|float_column| { - let int_name = metric_engine_value_int_column_name(&float_column.column_schema.name); - let int_column = metadata.column_by_name(&int_name)?; - if int_column.semantic_type != SemanticType::Field - || int_column.column_schema.data_type != ConcreteDataType::int64_datatype() - { - return None; - } - + metric_value_columns(metadata) + .into_iter() + .filter_map(|column| { + let float_column = &metadata.column_metadatas[column.value_index]; + let int_column = &metadata.column_metadatas[column.int_index]; let float_index = arrow_schema .index_of(&float_column.column_schema.name) .ok()?; - let int_index = arrow_schema.index_of(&int_name).ok()?; + let int_index = arrow_schema.index_of(&int_column.column_schema.name).ok()?; Some(MetricValueSplitColumn { float_index, int_index, @@ -70,21 +53,6 @@ pub(crate) fn metric_value_split_columns( .collect() } -fn is_metric_engine_data_region(metadata: &RegionMetadata) -> bool { - let has_internal_tag = |name, column_id| { - metadata.column_by_name(name).is_some_and(|column| { - column.semantic_type == SemanticType::Tag - && column.column_id == column_id - && metadata.primary_key.contains(&column_id) - }) - }; - - has_internal_tag( - DATA_SCHEMA_TABLE_ID_COLUMN_NAME, - ReservedColumnId::table_id(), - ) && has_internal_tag(DATA_SCHEMA_TSID_COLUMN_NAME, ReservedColumnId::tsid()) -} - pub(crate) fn split_metric_value_columns( batch: &RecordBatch, split_columns: &[MetricValueSplitColumn], @@ -247,12 +215,16 @@ fn integer_value(value: f64) -> Option { #[cfg(test)] mod tests { + use api::v1::SemanticType; use datatypes::arrow::array::{Float64Array, Int64Array}; + use datatypes::prelude::ConcreteDataType; use datatypes::schema::ColumnSchema; use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder}; + use store_api::metric_engine_consts::metric_engine_value_int_column_name; use store_api::storage::RegionId; use super::*; + use crate::metric_value::is_metric_engine_data_region; fn column_metadata( column_id: u32,