fix(metric-engine): report query load under physical region id (#8355)

* fix(metric-engine): report query load under physical region id

Propagate the physical region ID through the scanner, record batch
stream, and query engine so that query-load metrics (CPU time, scanned
bytes) are attributed to the correct physical region rather than always
to the logical region.

- `src/store-api/src/region_engine.rs` — add `query_load_region_id` to
  `ScannerProperties` and `set_query_load_region_id` to `RegionScanner`
  trait
- `src/mito2/src/read/seq_scan.rs`,
  `src/mito2/src/read/series_scan.rs`,
  `src/mito2/src/read/unordered_scan.rs` — implement the new trait
  method on each scanner
- `src/common/recordbatch/src/adapter.rs` — carry region id through
  `RecordBatchStreamAdapter` into `RecordBatchMetrics`
- `src/table/src/table/scan.rs` — expose region id on `RegionScanExec`
- `src/query/src/datafusion.rs` — extract region id from the physical
  plan and set it on the output stream
- `src/query/src/dist_plan/merge_scan.rs` — use metrics-contained region
  id in query-load reporting, falling back to the logical region id
- `src/metric-engine/src/engine/read.rs` — set region id on the metric
  engine scanner

Signed-off-by: Lei, HUANG <ratuthomm@gmail.com>

* fix(query): satisfy clippy for query load region id

Signed-off-by: Lei, HUANG <ratuthomm@gmail.com>

* fix(query): ignore missing query load region ids

Signed-off-by: Lei, HUANG <ratuthomm@gmail.com>

---------

Signed-off-by: Lei, HUANG <ratuthomm@gmail.com>
Co-authored-by: Lei, HUANG <ratuthomm@gmail.com>

src/common/recordbatch/src/adapter.rs

@@ -217,6 +217,7 @@ pub struct RecordBatchStreamAdapter {
     metrics: Option<BaselineMetrics>,
     /// Aggregated plan-level metrics. Resolved after an [ExecutionPlan] is finished.
     metrics_2: Metrics,
+    query_load_region_id: Option<u64>,
     /// 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<u64>) {
+        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<PlanMetrics>,
+    /// 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<u64>,
     /// Per-region watermark for incremental-read checkpoint advancement.
     ///
     /// The watermark is the latest sequence (`seq`) this query round safely read

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)
     }

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<u64> {
         self.stream_ctx.input.snapshot_sequence
     }

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<u64> {
         self.stream_ctx.input.snapshot_sequence
     }

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<u64> {
         self.stream_ctx.input.snapshot_sequence
     }

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<dyn ExecutionPlan>) -> Option<u64> {
+    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::<RegionScanExec>()
+            && 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<QueryEngineState>,
     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<RegionId>,
+    ) -> Arc<dyn ExecutionPlan> {
+        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<dyn PhysicalExpr>,
+            Arc::new(Column::new("ts", 0)) as Arc<dyn PhysicalExpr>,
+        )];
+        let plan: Arc<dyn ExecutionPlan> = 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 {

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);
+    }
 }

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<RegionId>,
 }
 
 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<RegionId> {
+        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<SequenceNumber> {
         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<SequenceNumber> {
         self.snapshot_sequence
     }

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<u64> {
+        self.scanner
+            .lock()
+            .unwrap()
+            .properties()
+            .query_load_region_id()
+            .map(|region_id| region_id.as_u64())
+    }
 }
 
 impl ExecutionPlan for RegionScanExec {