diff --git a/datafusion/physical-expr-common/src/metrics/mod.rs b/datafusion/physical-expr-common/src/metrics/mod.rs index d6048a0fcd338..f510714dae2a4 100644 --- a/datafusion/physical-expr-common/src/metrics/mod.rs +++ b/datafusion/physical-expr-common/src/metrics/mod.rs @@ -418,6 +418,21 @@ impl MetricsSet { .collect::>(); Self { metrics } } + + /// Returns a new `MetricsSet` filtered by metric name. + /// Only metrics with the names appearing the list will be kept. + pub fn filter_by_names(self, names: &[String]) -> Self { + if names.is_empty() { + return Self { metrics: vec![] }; + } + + let metrics = self + .metrics + .into_iter() + .filter(|metric| names.iter().any(|name| name == metric.value().name())) + .collect::>(); + Self { metrics } + } } impl Display for MetricsSet { @@ -966,4 +981,20 @@ mod tests { metric_names(&metrics) ); } + + #[test] + fn test_filter_by_names() { + let metrics = ExecutionPlanMetricsSet::new(); + MetricBuilder::new(&metrics).output_rows(0); + MetricBuilder::new(&metrics).counter("custom_counter", 0); + + let names = vec!["output_rows".to_string()]; + let filtered = metrics.clone_inner().filter_by_names(&names); + + assert_eq!(filtered.iter().count(), 1); + assert_eq!( + filtered.iter().next().unwrap().value().name(), + "output_rows" + ); + } } diff --git a/datafusion/physical-plan/src/display.rs b/datafusion/physical-plan/src/display.rs index 6a4d09057bec9..d14be3575a361 100644 --- a/datafusion/physical-plan/src/display.rs +++ b/datafusion/physical-plan/src/display.rs @@ -129,6 +129,9 @@ pub struct DisplayableExecutionPlan<'a> { /// Optional filter by semantic category (rows / bytes / timing). /// `None` means show all categories; `Some(vec![])` means plan-only. metric_categories: Option>, + /// Optional filter by metric names. Only metric names in this list + /// will be rendered. + metric_names: Option>, // (TreeRender) Maximum total width of the rendered tree tree_maximum_render_width: usize, /// Optional summary totals (currently only used by `pgjson`) — the total @@ -159,6 +162,7 @@ impl<'a> DisplayableExecutionPlan<'a> { show_schema: false, metric_types: Self::default_metric_types(), metric_categories: None, + metric_names: None, tree_maximum_render_width: 240, summary: None, } @@ -175,6 +179,7 @@ impl<'a> DisplayableExecutionPlan<'a> { show_schema: false, metric_types: Self::default_metric_types(), metric_categories: None, + metric_names: None, tree_maximum_render_width: 240, summary: None, } @@ -191,6 +196,7 @@ impl<'a> DisplayableExecutionPlan<'a> { show_schema: false, metric_types: Self::default_metric_types(), metric_categories: None, + metric_names: None, tree_maximum_render_width: 240, summary: None, } @@ -234,6 +240,15 @@ impl<'a> DisplayableExecutionPlan<'a> { self } + /// Specify which metric names to include. + /// + /// - An empty vector means plan-only — suppress all metrics. + /// - `vec!["metric_1"]` means show only the metric named `metric_1`. + pub fn set_metric_names(mut self, metric_names: Vec) -> Self { + self.metric_names = Some(metric_names); + self + } + /// Set the maximum render width for the tree format pub fn set_tree_maximum_render_width(mut self, width: usize) -> Self { self.tree_maximum_render_width = width; @@ -279,6 +294,7 @@ impl<'a> DisplayableExecutionPlan<'a> { show_schema: bool, metric_types: Vec, metric_categories: Option>, + metric_names: Option>, } impl fmt::Display for Wrapper<'_> { fn fmt(&self, f: &mut Formatter) -> fmt::Result { @@ -291,6 +307,7 @@ impl<'a> DisplayableExecutionPlan<'a> { show_schema: self.show_schema, metric_types: &self.metric_types, metric_categories: self.metric_categories.as_deref(), + metric_names: self.metric_names.as_deref(), }; accept(self.plan, &mut visitor) } @@ -303,6 +320,7 @@ impl<'a> DisplayableExecutionPlan<'a> { show_schema: self.show_schema, metric_types: self.metric_types.clone(), metric_categories: self.metric_categories.clone(), + metric_names: self.metric_names.clone(), } } @@ -324,6 +342,7 @@ impl<'a> DisplayableExecutionPlan<'a> { show_statistics: bool, metric_types: Vec, metric_categories: Option>, + metric_names: Option>, } impl fmt::Display for Wrapper<'_> { fn fmt(&self, f: &mut Formatter) -> fmt::Result { @@ -336,6 +355,7 @@ impl<'a> DisplayableExecutionPlan<'a> { show_statistics: self.show_statistics, metric_types: &self.metric_types, metric_categories: self.metric_categories.as_deref(), + metric_names: self.metric_names.as_deref(), graphviz_builder: GraphvizBuilder::default(), parents: Vec::new(), }; @@ -355,6 +375,7 @@ impl<'a> DisplayableExecutionPlan<'a> { show_statistics: self.show_statistics, metric_types: self.metric_types.clone(), metric_categories: self.metric_categories.clone(), + metric_names: self.metric_names.clone(), } } @@ -403,6 +424,7 @@ impl<'a> DisplayableExecutionPlan<'a> { show_schema: bool, metric_types: Vec, metric_categories: Option>, + metric_names: Option>, summary: Option, } impl fmt::Display for Wrapper<'_> { @@ -413,6 +435,7 @@ impl<'a> DisplayableExecutionPlan<'a> { show_schema: self.show_schema, metric_types: &self.metric_types, metric_categories: self.metric_categories.as_deref(), + metric_names: self.metric_names.as_deref(), objects: HashMap::new(), parent_ids: Vec::new(), next_id: 0, @@ -446,6 +469,7 @@ impl<'a> DisplayableExecutionPlan<'a> { show_schema: self.show_schema, metric_types: self.metric_types.clone(), metric_categories: self.metric_categories.clone(), + metric_names: self.metric_names.clone(), summary: self.summary, } } @@ -460,6 +484,7 @@ impl<'a> DisplayableExecutionPlan<'a> { show_schema: bool, metric_types: Vec, metric_categories: Option>, + metric_names: Option>, } impl fmt::Display for Wrapper<'_> { @@ -473,6 +498,7 @@ impl<'a> DisplayableExecutionPlan<'a> { show_schema: self.show_schema, metric_types: &self.metric_types, metric_categories: self.metric_categories.as_deref(), + metric_names: self.metric_names.as_deref(), }; visitor.pre_visit(self.plan)?; Ok(()) @@ -486,6 +512,7 @@ impl<'a> DisplayableExecutionPlan<'a> { show_schema: self.show_schema, metric_types: self.metric_types.clone(), metric_categories: self.metric_categories.clone(), + metric_names: self.metric_names.clone(), } } @@ -544,6 +571,8 @@ struct IndentVisitor<'a, 'b> { metric_types: &'a [MetricType], /// Optional filter by semantic category (rows / bytes / timing). metric_categories: Option<&'a [MetricCategory]>, + /// Optional filter by metric name. + metric_names: Option<&'a [String]>, } impl ExecutionPlanVisitor for IndentVisitor<'_, '_> { @@ -563,6 +592,9 @@ impl ExecutionPlanVisitor for IndentVisitor<'_, '_> { if let Some(cats) = self.metric_categories { metrics = metrics.filter_by_categories(cats); } + if let Some(names) = self.metric_names { + metrics = metrics.filter_by_names(names); + } write!(self.f, ", metrics=[{metrics}]")?; } else { write!(self.f, ", metrics=[]")?; @@ -574,6 +606,9 @@ impl ExecutionPlanVisitor for IndentVisitor<'_, '_> { if let Some(cats) = self.metric_categories { metrics = metrics.filter_by_categories(cats); } + if let Some(names) = self.metric_names { + metrics = metrics.filter_by_names(names); + } write!(self.f, ", metrics=[{metrics}]")?; } else { write!(self.f, ", metrics=[]")?; @@ -616,6 +651,8 @@ struct GraphvizVisitor<'a, 'b> { metric_types: &'a [MetricType], /// Optional filter by semantic category metric_categories: Option<&'a [MetricCategory]>, + /// Optional filter by metric name. + metric_names: Option<&'a [String]>, graphviz_builder: GraphvizBuilder, /// Used to record parent node ids when visiting a plan. @@ -660,6 +697,9 @@ impl ExecutionPlanVisitor for GraphvizVisitor<'_, '_> { if let Some(cats) = self.metric_categories { metrics = metrics.filter_by_categories(cats); } + if let Some(names) = self.metric_names { + metrics = metrics.filter_by_names(names); + } format!("metrics=[{metrics}]") } else { "metrics=[]".to_string() @@ -671,6 +711,9 @@ impl ExecutionPlanVisitor for GraphvizVisitor<'_, '_> { if let Some(cats) = self.metric_categories { metrics = metrics.filter_by_categories(cats); } + if let Some(names) = self.metric_names { + metrics = metrics.filter_by_names(names); + } format!("metrics=[{metrics}]") } else { "metrics=[]".to_string() @@ -729,6 +772,7 @@ struct PgJsonExecutionPlanVisitor<'a> { show_schema: bool, metric_types: &'a [MetricType], metric_categories: Option<&'a [MetricCategory]>, + metric_names: Option<&'a [String]>, objects: HashMap, parent_ids: Vec, next_id: u32, @@ -813,6 +857,12 @@ impl PgJsonExecutionPlanVisitor<'_> { metrics }; + let metrics = if let Some(names) = self.metric_names { + metrics.filter_by_names(names) + } else { + metrics + }; + // Build the Extras bucket, while extracting PG-canonical keys to the // top level. let mut extras = serde_json::Map::new();