diff --git a/src/common/recordbatch/src/adapter.rs b/src/common/recordbatch/src/adapter.rs index a69eedf7ebf..15c07773146 100644 --- a/src/common/recordbatch/src/adapter.rs +++ b/src/common/recordbatch/src/adapter.rs @@ -366,6 +366,7 @@ impl ExecutionPlanVisitor for MetricCollector { let Some(metric) = plan.metrics() else { self.record_batch_metrics.plan_metrics.push(PlanMetrics { plan: plan.name().to_string(), + plan_name: plan.name().to_string(), level: self.current_level, metrics: vec![], }); @@ -380,6 +381,7 @@ impl ExecutionPlanVisitor for MetricCollector { .timestamps_removed(); let mut plan_metric = PlanMetrics { plan: one_line(plan, self.verbose).to_string(), + plan_name: plan.name().to_string(), level: self.current_level, metrics: Vec::with_capacity(metric.iter().size_hint().0), }; @@ -519,6 +521,9 @@ impl Display for RecordBatchMetrics { pub struct PlanMetrics { /// The plan name pub plan: String, + /// The stable execution plan name. + #[serde(default)] + pub plan_name: String, /// The level of the plan, starts from 0 pub level: usize, /// An ordered key-value list of metrics. @@ -916,6 +921,22 @@ mod test { assert_eq!(metrics.memory_usage, 34); } + #[test] + fn test_plan_metrics_deserializes_without_plan_name() { + let metrics: RecordBatchMetrics = serde_json::from_value(json!({ + "elapsed_compute": 12, + "memory_usage": 34, + "plan_metrics": [{ + "plan": "SeqScan: region=1", + "level": 0, + "metrics": [] + }] + })) + .unwrap(); + + assert_eq!(metrics.plan_metrics[0].plan_name, ""); + } + #[test] fn test_recordbatch_metrics_region_watermarks_serde_roundtrip() { let metrics = RecordBatchMetrics { diff --git a/src/query/src/dist_plan/merge_scan.rs b/src/query/src/dist_plan/merge_scan.rs index 0b03d4f53ba..5d0db51e024 100644 --- a/src/query/src/dist_plan/merge_scan.rs +++ b/src/query/src/dist_plan/merge_scan.rs @@ -47,6 +47,7 @@ use meter_core::data::ReadItem; use meter_macros::read_meter; use session::context::QueryContextRef; use store_api::storage::RegionId; +use table::table::scan::REGION_SCAN_EXEC_NAME; use table::table_name::TableName; use tokio::time::Instant; use tracing::{Instrument, Span}; @@ -496,13 +497,14 @@ impl MergeScanExec { // process metrics after all data is drained. if let Some(metrics) = stream.metrics() { + let output_bytes = scan_output_bytes(&metrics); let (c, s) = parse_catalog_and_schema_from_db_string(&dbname); let value = read_meter!( c, s, ReadItem { cpu_time: metrics.elapsed_compute as u64, - table_scan: metrics.memory_usage as u64 + table_scan: output_bytes as u64, }, current_channel as u8 ); @@ -916,6 +918,17 @@ impl DisplayAs for MergeScanExec { } } +/// Extract total `output_bytes` from [`RegionScanExec`] plan nodes. +fn scan_output_bytes(metrics: &RecordBatchMetrics) -> usize { + metrics + .plan_metrics + .iter() + .filter(|pm| pm.plan_name == REGION_SCAN_EXEC_NAME) + .flat_map(|pm| &pm.metrics) + .filter_map(|(name, value)| (name == "output_bytes").then_some(*value)) + .sum() +} + #[derive(Debug, Clone)] struct MergeScanMetric { /// Nanosecond elapsed till the scan operator is ready to emit data @@ -969,6 +982,7 @@ mod tests { use async_trait::async_trait; use common_query::request::INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY; + use common_recordbatch::adapter::{PlanMetrics, RecordBatchMetrics}; use datafusion::config::ConfigOptions; use datafusion::execution::SessionStateBuilder; use datafusion::physical_plan::filter_pushdown::ChildFilterPushdownResult; @@ -1210,4 +1224,57 @@ mod tests { assert_eq!(exec.captured_remote_dyn_filters().len(), 1); assert!(matches!(propagation.filters.as_slice(), [PushedDown::Yes])); } + + #[test] + fn scan_output_bytes_uses_plan_name() { + let metrics = RecordBatchMetrics { + plan_metrics: vec![PlanMetrics { + plan: "SeqScan: region=1".to_string(), + plan_name: REGION_SCAN_EXEC_NAME.to_string(), + level: 0, + metrics: vec![("output_bytes".to_string(), 42)], + }], + ..Default::default() + }; + + assert_eq!(scan_output_bytes(&metrics), 42); + } + + #[test] + fn scan_output_bytes_defaults_to_zero_without_region_scan() { + let metrics = RecordBatchMetrics { + plan_metrics: vec![PlanMetrics { + plan: "ProjectionExec".to_string(), + plan_name: "ProjectionExec".to_string(), + level: 0, + metrics: vec![("output_bytes".to_string(), 42)], + }], + ..Default::default() + }; + + assert_eq!(scan_output_bytes(&metrics), 0); + } + + #[test] + fn scan_output_bytes_sums_multiple_region_scans() { + let metrics = RecordBatchMetrics { + 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(), 42)], + }, + PlanMetrics { + plan: "RegionScanExec: region=2".to_string(), + plan_name: REGION_SCAN_EXEC_NAME.to_string(), + level: 0, + metrics: vec![("output_bytes".to_string(), 18)], + }, + ], + ..Default::default() + }; + + assert_eq!(scan_output_bytes(&metrics), 60); + } } diff --git a/src/table/src/table/metrics.rs b/src/table/src/table/metrics.rs index 23356a408d8..b76be231f1f 100644 --- a/src/table/src/table/metrics.rs +++ b/src/table/src/table/metrics.rs @@ -15,7 +15,7 @@ use std::time::Duration; use datafusion::physical_plan::metrics::{ - Count, ExecutionPlanMetricsSet, Gauge, MetricBuilder, ScopedTimerGuard, Time, Timestamp, + Count, ExecutionPlanMetricsSet, MetricBuilder, ScopedTimerGuard, Time, Timestamp, }; /// This metrics struct is used to record and hold metrics like memory usage @@ -25,10 +25,10 @@ use datafusion::physical_plan::metrics::{ pub struct StreamMetrics { /// Timestamp when the stream finished end_time: Timestamp, - /// Used memory in bytes - mem_used: Gauge, /// Number of rows in output output_rows: Count, + /// Number of bytes in output + output_bytes: Count, /// Elapsed time used to `poll` the stream poll_elapsed: Time, /// Elapsed time used to `.await`ing the stream @@ -36,28 +36,28 @@ pub struct StreamMetrics { } impl StreamMetrics { - /// Create a new MemoryUsageMetrics structure, and set `start_time` to now + /// Create a new [`StreamMetrics`] structure, and set `start_time` to now. pub fn new(metrics: &ExecutionPlanMetricsSet, partition: usize) -> Self { let start_time = MetricBuilder::new(metrics).start_timestamp(partition); start_time.record(); Self { end_time: MetricBuilder::new(metrics).end_timestamp(partition), - mem_used: MetricBuilder::new(metrics).mem_used(partition), output_rows: MetricBuilder::new(metrics).output_rows(partition), + output_bytes: MetricBuilder::new(metrics).output_bytes(partition), poll_elapsed: MetricBuilder::new(metrics).subset_time("elapsed_poll", partition), await_elapsed: MetricBuilder::new(metrics).subset_time("elapsed_await", partition), } } - pub fn record_mem_usage(&self, mem_used: usize) { - self.mem_used.add(mem_used); - } - pub fn record_output(&self, num_rows: usize) { self.output_rows.add(num_rows); } + pub fn record_output_bytes(&self, num_bytes: usize) { + self.output_bytes.add(num_bytes); + } + /// Record the end time of the query pub fn try_done(&self) { if self.end_time.value().is_none() { diff --git a/src/table/src/table/scan.rs b/src/table/src/table/scan.rs index d301202e172..20338d0cfd8 100644 --- a/src/table/src/table/scan.rs +++ b/src/table/src/table/scan.rs @@ -54,6 +54,9 @@ use store_api::storage::{ScanRequest, TimeSeriesDistribution}; use crate::table::metrics::StreamMetrics; +/// The name of [`RegionScanExec`] plan node. +pub const REGION_SCAN_EXEC_NAME: &str = "RegionScanExec"; + /// A plan to read multiple partitions from a region of a table. #[derive(Clone)] pub struct RegionScanExec { @@ -453,7 +456,7 @@ impl ExecutionPlan for RegionScanExec { } fn name(&self) -> &str { - "RegionScanExec" + REGION_SCAN_EXEC_NAME } fn handle_child_pushdown_result( @@ -519,8 +522,8 @@ impl Stream for StreamWithMetricWrapper { Ok(record_batch) => { // we don't record elapsed time here // since it's calling storage api involving I/O ops - this.metric - .record_mem_usage(record_batch.buffer_memory_size()); + let batch_bytes = record_batch.buffer_memory_size(); + this.metric.record_output_bytes(batch_bytes); this.metric.record_output(record_batch.num_rows()); Poll::Ready(Some(Ok(record_batch.into_df_record_batch()))) }