diff --git a/src/common/recordbatch/src/adapter.rs b/src/common/recordbatch/src/adapter.rs index 15c07773146..e6a620ab6c5 100644 --- a/src/common/recordbatch/src/adapter.rs +++ b/src/common/recordbatch/src/adapter.rs @@ -217,6 +217,7 @@ pub struct RecordBatchStreamAdapter { metrics: Option, /// Aggregated plan-level metrics. Resolved after an [ExecutionPlan] is finished. metrics_2: Metrics, + query_load_region_id: Option, /// Display plan and metrics in verbose mode. explain_verbose: bool, span: Span, @@ -239,6 +240,7 @@ impl RecordBatchStreamAdapter { stream, metrics: None, metrics_2: Metrics::Unavailable, + query_load_region_id: None, explain_verbose: false, span: Span::current(), }) @@ -253,6 +255,7 @@ impl RecordBatchStreamAdapter { stream, metrics: None, metrics_2: Metrics::Unavailable, + query_load_region_id: None, explain_verbose: false, span: subspan, }) @@ -262,6 +265,10 @@ impl RecordBatchStreamAdapter { self.metrics_2 = Metrics::Unresolved(plan) } + pub fn set_query_load_region_id(&mut self, region_id: Option) { + self.query_load_region_id = region_id; + } + /// Set the verbose mode for displaying plan and metrics. pub fn set_explain_verbose(&mut self, verbose: bool) { self.explain_verbose = verbose; @@ -312,6 +319,8 @@ impl Stream for RecordBatchStreamAdapter { { let mut metric_collector = MetricCollector::new(self.explain_verbose); accept(df_plan.as_ref(), &mut metric_collector).unwrap(); + metric_collector.record_batch_metrics.query_load_region_id = + self.query_load_region_id; self.metrics_2 = Metrics::PartialResolved( df_plan.clone(), metric_collector.record_batch_metrics, @@ -328,6 +337,8 @@ impl Stream for RecordBatchStreamAdapter { { let mut metric_collector = MetricCollector::new(self.explain_verbose); accept(df_plan.as_ref(), &mut metric_collector).unwrap(); + metric_collector.record_batch_metrics.query_load_region_id = + self.query_load_region_id; self.metrics_2 = Metrics::Resolved(metric_collector.record_batch_metrics); } Poll::Ready(None) @@ -448,6 +459,9 @@ pub struct RecordBatchMetrics { // Detailed per-plan metrics /// An ordered list of plan metrics, from top to bottom in post-order. pub plan_metrics: Vec, + /// Region id that should receive query-load metrics for this scan. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub query_load_region_id: Option, /// Per-region watermark for incremental-read checkpoint advancement. /// /// The watermark is the latest sequence (`seq`) this query round safely read diff --git a/src/metric-engine/src/engine/read.rs b/src/metric-engine/src/engine/read.rs index 178c2e5b92e..0537d17c73b 100644 --- a/src/metric-engine/src/engine/read.rs +++ b/src/metric-engine/src/engine/read.rs @@ -85,6 +85,7 @@ impl MetricEngineInner { .await .context(MitoReadOperationSnafu)?; scanner.set_logical_region(true); + scanner.set_query_load_region_id(data_region_id); Ok(scanner) } diff --git a/src/mito2/src/read/seq_scan.rs b/src/mito2/src/read/seq_scan.rs index 6bc796c85c5..88349fd63be 100644 --- a/src/mito2/src/read/seq_scan.rs +++ b/src/mito2/src/read/seq_scan.rs @@ -573,6 +573,10 @@ impl RegionScanner for SeqScan { self.properties.set_logical_region(logical_region); } + fn set_query_load_region_id(&mut self, region_id: store_api::storage::RegionId) { + self.properties.set_query_load_region_id(region_id); + } + fn snapshot_sequence(&self) -> Option { self.stream_ctx.input.snapshot_sequence } diff --git a/src/mito2/src/read/series_scan.rs b/src/mito2/src/read/series_scan.rs index 8f554b6aedd..856b57331a8 100644 --- a/src/mito2/src/read/series_scan.rs +++ b/src/mito2/src/read/series_scan.rs @@ -381,6 +381,10 @@ impl RegionScanner for SeriesScan { self.properties.set_logical_region(logical_region); } + fn set_query_load_region_id(&mut self, region_id: store_api::storage::RegionId) { + self.properties.set_query_load_region_id(region_id); + } + fn snapshot_sequence(&self) -> Option { self.stream_ctx.input.snapshot_sequence } diff --git a/src/mito2/src/read/unordered_scan.rs b/src/mito2/src/read/unordered_scan.rs index a48a7048bdc..0c54084bff6 100644 --- a/src/mito2/src/read/unordered_scan.rs +++ b/src/mito2/src/read/unordered_scan.rs @@ -355,6 +355,10 @@ impl RegionScanner for UnorderedScan { self.properties.set_logical_region(logical_region); } + fn set_query_load_region_id(&mut self, region_id: store_api::storage::RegionId) { + self.properties.set_query_load_region_id(region_id); + } + fn snapshot_sequence(&self) -> Option { self.stream_ctx.input.snapshot_sequence } diff --git a/src/query/src/datafusion.rs b/src/query/src/datafusion.rs index 7be6656c6fc..22c64aab4a6 100644 --- a/src/query/src/datafusion.rs +++ b/src/query/src/datafusion.rs @@ -49,6 +49,7 @@ use snafu::{OptionExt, ResultExt, ensure}; use sqlparser::ast::AnalyzeFormat; use table::TableRef; use table::requests::{DeleteRequest, InsertRequest}; +use table::table::scan::{REGION_SCAN_EXEC_NAME, RegionScanExec}; use tracing::Span; use crate::analyze::DistAnalyzeExec; @@ -78,6 +79,27 @@ pub const QUERY_PARALLELISM_HINT: &str = "query_parallelism"; /// Whether to fallback to the original plan when failed to push down. pub const QUERY_FALLBACK_HINT: &str = "query_fallback"; +fn query_load_region_id(plan: &Arc) -> Option { + let mut region_id = None; + let mut stack = vec![plan.clone()]; + + while let Some(plan) = stack.pop() { + if plan.name() == REGION_SCAN_EXEC_NAME + && let Some(scan) = plan.as_any().downcast_ref::() + && let Some(scan_region_id) = scan.query_load_region_id() + { + match region_id { + Some(region_id) if region_id != scan_region_id => return None, + Some(_) => {} + None => region_id = Some(scan_region_id), + } + } + stack.extend(plan.children().into_iter().cloned()); + } + + region_id +} + pub struct DatafusionQueryEngine { state: Arc, plugins: Plugins, @@ -587,6 +609,7 @@ impl QueryExecutor for DatafusionQueryEngine { .map_err(BoxedError::new) .context(QueryExecutionSnafu)?; stream.set_metrics2(plan.clone()); + stream.set_query_load_region_id(query_load_region_id(plan)); stream.set_explain_verbose(explain_verbose); let stream = OnDone::new(Box::pin(stream), move || { let exec_cost = exec_timer.stop_and_record(); @@ -620,6 +643,7 @@ impl QueryExecutor for DatafusionQueryEngine { .map_err(BoxedError::new) .context(QueryExecutionSnafu)?; stream.set_metrics2(plan.clone()); + stream.set_query_load_region_id(query_load_region_id(plan)); stream.set_explain_verbose(explain_verbose); let stream = OnDone::new(Box::pin(stream), move || { let exec_cost = exec_timer.stop_and_record(); @@ -753,6 +777,10 @@ mod tests { fn set_logical_region(&mut self, logical_region: bool) { self.properties.set_logical_region(logical_region); } + + fn set_query_load_region_id(&mut self, region_id: store_api::storage::RegionId) { + self.properties.set_query_load_region_id(region_id); + } } impl DisplayAs for RecordingScanner { @@ -761,6 +789,72 @@ mod tests { } } + fn build_query_load_region_scan( + query_load_region_id: Option, + ) -> Arc { + let schema = Arc::new(datatypes::schema::Schema::new(vec![ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + )])); + + let mut metadata_builder = RegionMetadataBuilder::new(RegionId::new(1024, 1)); + metadata_builder + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + semantic_type: SemanticType::Timestamp, + column_id: 1, + }) + .primary_key(vec![]); + let metadata = Arc::new(metadata_builder.build().unwrap()); + let mut scanner = RecordingScanner::new( + schema, + metadata, + Arc::new(AtomicUsize::new(0)), + Arc::new(AtomicUsize::new(0)), + ); + if let Some(region_id) = query_load_region_id { + scanner.set_query_load_region_id(region_id); + } + + Arc::new(RegionScanExec::new(Box::new(scanner), ScanRequest::default(), None).unwrap()) + } + + #[test] + fn query_load_region_id_ignores_scans_without_region_id() { + let query_load_region_id = RegionId::new(1024, 42); + let scan_without_region_id = build_query_load_region_scan(None); + let scan_with_region_id = build_query_load_region_scan(Some(query_load_region_id)); + let on: JoinOn = vec![( + Arc::new(Column::new("ts", 0)) as Arc, + Arc::new(Column::new("ts", 0)) as Arc, + )]; + let plan: Arc = Arc::new( + HashJoinExec::try_new( + scan_without_region_id, + scan_with_region_id, + on, + None, + &JoinType::Inner, + None, + PartitionMode::CollectLeft, + NullEquality::NullEqualsNull, + false, + ) + .unwrap(), + ); + + assert_eq!( + super::query_load_region_id(&plan), + Some(query_load_region_id.as_u64()) + ); + } + async fn create_test_engine() -> QueryEngineRef { let catalog_manager = catalog::memory::new_memory_catalog_manager().unwrap(); let req = RegisterTableRequest { diff --git a/src/query/src/dist_plan/merge_scan.rs b/src/query/src/dist_plan/merge_scan.rs index 9615bc2a132..262e391b556 100644 --- a/src/query/src/dist_plan/merge_scan.rs +++ b/src/query/src/dist_plan/merge_scan.rs @@ -943,6 +943,13 @@ fn report_region_query_load(region_id: RegionId, load: &ReadItem) { .inc_by(load.table_scan); } +fn query_load_region_id(default_region_id: RegionId, metrics: &RecordBatchMetrics) -> RegionId { + metrics + .query_load_region_id + .map(RegionId::from_u64) + .unwrap_or(default_region_id) +} + impl Drop for MergeScanExec { fn drop(&mut self) { if !self.enable_per_region_metrics { @@ -952,7 +959,7 @@ impl Drop for MergeScanExec { let metrics = self.sub_stage_metrics.lock().unwrap(); for (region_id, metrics) in metrics.iter() { let load = region_scan_load(metrics); - report_region_query_load(*region_id, &load); + report_region_query_load(query_load_region_id(*region_id, metrics), &load); } } } @@ -1371,4 +1378,89 @@ mod tests { let _ = REGION_QUERY_CPU_TIME.remove_label_values(&labels); let _ = REGION_QUERY_SCANNED_BYTES.remove_label_values(&labels); } + + #[test] + fn merge_scan_reports_query_load_with_metrics_region_id() { + use store_api::metrics::{REGION_QUERY_CPU_TIME, REGION_QUERY_SCANNED_BYTES}; + + let logical_region_id = RegionId::new(1024, 10002); + let physical_region_id = RegionId::new(1024, 1); + let logical_region_id_label = logical_region_id.to_string(); + let physical_region_id_label = physical_region_id.to_string(); + let logical_labels = [&logical_region_id_label]; + let physical_labels = [&physical_region_id_label]; + let _ = REGION_QUERY_CPU_TIME.remove_label_values(&logical_labels); + let _ = REGION_QUERY_SCANNED_BYTES.remove_label_values(&logical_labels); + let _ = REGION_QUERY_CPU_TIME.remove_label_values(&physical_labels); + let _ = REGION_QUERY_SCANNED_BYTES.remove_label_values(&physical_labels); + + let plan = LogicalPlanBuilder::empty(true) + .project(vec![lit(1i32).alias("col1")]) + .unwrap() + .build() + .unwrap(); + let schema = plan.schema().as_arrow().clone(); + let exec = MergeScanExec::new( + &SessionStateBuilder::new().build(), + TableName::new("catalog", "schema", "table"), + vec![logical_region_id], + plan, + &schema, + Arc::new(TestRegionQueryHandler), + QueryContext::arc(), + 1, + AliasMapping::new(), + None, + true, + ) + .unwrap(); + + let metrics = RecordBatchMetrics { + elapsed_compute: 42, + query_load_region_id: Some(physical_region_id.as_u64()), + plan_metrics: vec![PlanMetrics { + plan: "RegionScanExec: region=1".to_string(), + plan_name: REGION_SCAN_EXEC_NAME.to_string(), + level: 0, + metrics: vec![("output_bytes".to_string(), 24)], + }], + ..Default::default() + }; + exec.sub_stage_metrics + .lock() + .unwrap() + .insert(logical_region_id, metrics); + + drop(exec); + + assert_eq!( + REGION_QUERY_CPU_TIME + .with_label_values(&logical_labels) + .get(), + 0 + ); + assert_eq!( + REGION_QUERY_SCANNED_BYTES + .with_label_values(&logical_labels) + .get(), + 0 + ); + assert_eq!( + REGION_QUERY_CPU_TIME + .with_label_values(&physical_labels) + .get(), + 42 + ); + assert_eq!( + REGION_QUERY_SCANNED_BYTES + .with_label_values(&physical_labels) + .get(), + 24 + ); + + let _ = REGION_QUERY_CPU_TIME.remove_label_values(&logical_labels); + let _ = REGION_QUERY_SCANNED_BYTES.remove_label_values(&logical_labels); + let _ = REGION_QUERY_CPU_TIME.remove_label_values(&physical_labels); + let _ = REGION_QUERY_SCANNED_BYTES.remove_label_values(&physical_labels); + } } diff --git a/src/store-api/src/region_engine.rs b/src/store-api/src/region_engine.rs index e5ab05e5e7b..ce3d1bf1a8f 100644 --- a/src/store-api/src/region_engine.rs +++ b/src/store-api/src/region_engine.rs @@ -313,6 +313,9 @@ pub struct ScannerProperties { /// Whether the scanner is scanning a logical region. logical_region: bool, + + /// Region id that should receive query-load metrics for this scanner. + query_load_region_id: Option, } impl ScannerProperties { @@ -337,6 +340,7 @@ impl ScannerProperties { distinguish_partition_range: false, target_partitions: 0, logical_region: false, + query_load_region_id: None, } } @@ -384,6 +388,16 @@ impl ScannerProperties { pub fn set_logical_region(&mut self, logical_region: bool) { self.logical_region = logical_region; } + + /// Returns the region id that should receive query-load metrics. + pub fn query_load_region_id(&self) -> Option { + self.query_load_region_id + } + + /// Sets the region id that should receive query-load metrics. + pub fn set_query_load_region_id(&mut self, region_id: RegionId) { + self.query_load_region_id = Some(region_id); + } } /// Request to override the scanner properties. @@ -471,6 +485,9 @@ pub trait RegionScanner: Debug + DisplayAs + Send { /// Sets whether the scanner is reading a logical region. fn set_logical_region(&mut self, logical_region: bool); + /// Sets the region id that should receive query-load metrics. + fn set_query_load_region_id(&mut self, region_id: RegionId); + fn snapshot_sequence(&self) -> Option { None } @@ -1044,6 +1061,10 @@ impl RegionScanner for SinglePartitionScanner { self.properties.set_logical_region(logical_region); } + fn set_query_load_region_id(&mut self, region_id: RegionId) { + self.properties.set_query_load_region_id(region_id); + } + fn snapshot_sequence(&self) -> Option { self.snapshot_sequence } diff --git a/src/table/src/table/scan.rs b/src/table/src/table/scan.rs index 20338d0cfd8..a5f06a40d3b 100644 --- a/src/table/src/table/scan.rs +++ b/src/table/src/table/scan.rs @@ -220,6 +220,7 @@ impl RegionScanExec { )); let append_mode = scanner_props.append_mode(); let total_rows = scanner_props.total_rows(); + Ok(Self { scanner: Arc::new(Mutex::new(scanner)), arrow_schema, @@ -359,6 +360,15 @@ impl RegionScanExec { .unwrap() .add_dyn_filter_to_predicate(filter_exprs) } + + pub fn query_load_region_id(&self) -> Option { + self.scanner + .lock() + .unwrap() + .properties() + .query_load_region_id() + .map(|region_id| region_id.as_u64()) + } } impl ExecutionPlan for RegionScanExec {