Skip to content

Commit bd0852b

Browse files
gabotechsgabotechs
committed
Add name filter to metrics
1 parent e8a65f2 commit bd0852b

2 files changed

Lines changed: 81 additions & 0 deletions

File tree

‎datafusion/physical-expr-common/src/metrics/mod.rs‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -418,6 +418,21 @@ impl MetricsSet {
418418
.collect::<Vec<_>>();
419419
Self { metrics }
420420
}
421+
422+
/// Returns a new `MetricsSet` filtered by metric name.
423+
/// Only metrics with the names appearing the list will be kept.
424+
pub fn filter_by_names(self, names: &[String]) -> Self {
425+
if names.is_empty() {
426+
return Self { metrics: vec![] };
427+
}
428+
429+
let metrics = self
430+
.metrics
431+
.into_iter()
432+
.filter(|metric| names.iter().any(|name| name == metric.value().name()))
433+
.collect::<Vec<_>>();
434+
Self { metrics }
435+
}
421436
}
422437

423438
impl Display for MetricsSet {
@@ -966,4 +981,20 @@ mod tests {
966981
metric_names(&metrics)
967982
);
968983
}
984+
985+
#[test]
986+
fn test_filter_by_names() {
987+
let metrics = ExecutionPlanMetricsSet::new();
988+
MetricBuilder::new(&metrics).output_rows(0);
989+
MetricBuilder::new(&metrics).counter("custom_counter", 0);
990+
991+
let names = vec!["output_rows".to_string()];
992+
let filtered = metrics.clone_inner().filter_by_names(&names);
993+
994+
assert_eq!(filtered.iter().count(), 1);
995+
assert_eq!(
996+
filtered.iter().next().unwrap().value().name(),
997+
"output_rows"
998+
);
999+
}
9691000
}

‎datafusion/physical-plan/src/display.rs‎

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -129,6 +129,9 @@ pub struct DisplayableExecutionPlan<'a> {
129129
/// Optional filter by semantic category (rows / bytes / timing).
130130
/// `None` means show all categories; `Some(vec![])` means plan-only.
131131
metric_categories: Option<Vec<MetricCategory>>,
132+
/// Optional filter by metric names. Only metric names in this list
133+
/// will be rendered.
134+
metric_names: Option<Vec<String>>,
132135
// (TreeRender) Maximum total width of the rendered tree
133136
tree_maximum_render_width: usize,
134137
/// Optional summary totals (currently only used by `pgjson`) — the total
@@ -159,6 +162,7 @@ impl<'a> DisplayableExecutionPlan<'a> {
159162
show_schema: false,
160163
metric_types: Self::default_metric_types(),
161164
metric_categories: None,
165+
metric_names: None,
162166
tree_maximum_render_width: 240,
163167
summary: None,
164168
}
@@ -175,6 +179,7 @@ impl<'a> DisplayableExecutionPlan<'a> {
175179
show_schema: false,
176180
metric_types: Self::default_metric_types(),
177181
metric_categories: None,
182+
metric_names: None,
178183
tree_maximum_render_width: 240,
179184
summary: None,
180185
}
@@ -191,6 +196,7 @@ impl<'a> DisplayableExecutionPlan<'a> {
191196
show_schema: false,
192197
metric_types: Self::default_metric_types(),
193198
metric_categories: None,
199+
metric_names: None,
194200
tree_maximum_render_width: 240,
195201
summary: None,
196202
}
@@ -234,6 +240,15 @@ impl<'a> DisplayableExecutionPlan<'a> {
234240
self
235241
}
236242

243+
/// Specify which metric names to include.
244+
///
245+
/// - An empty vector means plan-only — suppress all metrics.
246+
/// - `vec!["metric_1"]` means show only the metric named `metric_1`.
247+
pub fn set_metric_names(mut self, metric_names: Vec<String>) -> Self {
248+
self.metric_names = Some(metric_names);
249+
self
250+
}
251+
237252
/// Set the maximum render width for the tree format
238253
pub fn set_tree_maximum_render_width(mut self, width: usize) -> Self {
239254
self.tree_maximum_render_width = width;
@@ -279,6 +294,7 @@ impl<'a> DisplayableExecutionPlan<'a> {
279294
show_schema: bool,
280295
metric_types: Vec<MetricType>,
281296
metric_categories: Option<Vec<MetricCategory>>,
297+
metric_names: Option<Vec<String>>,
282298
}
283299
impl fmt::Display for Wrapper<'_> {
284300
fn fmt(&self, f: &mut Formatter) -> fmt::Result {
@@ -291,6 +307,7 @@ impl<'a> DisplayableExecutionPlan<'a> {
291307
show_schema: self.show_schema,
292308
metric_types: &self.metric_types,
293309
metric_categories: self.metric_categories.as_deref(),
310+
metric_names: self.metric_names.as_deref(),
294311
};
295312
accept(self.plan, &mut visitor)
296313
}
@@ -303,6 +320,7 @@ impl<'a> DisplayableExecutionPlan<'a> {
303320
show_schema: self.show_schema,
304321
metric_types: self.metric_types.clone(),
305322
metric_categories: self.metric_categories.clone(),
323+
metric_names: self.metric_names.clone(),
306324
}
307325
}
308326

@@ -324,6 +342,7 @@ impl<'a> DisplayableExecutionPlan<'a> {
324342
show_statistics: bool,
325343
metric_types: Vec<MetricType>,
326344
metric_categories: Option<Vec<MetricCategory>>,
345+
metric_names: Option<Vec<String>>,
327346
}
328347
impl fmt::Display for Wrapper<'_> {
329348
fn fmt(&self, f: &mut Formatter) -> fmt::Result {
@@ -336,6 +355,7 @@ impl<'a> DisplayableExecutionPlan<'a> {
336355
show_statistics: self.show_statistics,
337356
metric_types: &self.metric_types,
338357
metric_categories: self.metric_categories.as_deref(),
358+
metric_names: self.metric_names.as_deref(),
339359
graphviz_builder: GraphvizBuilder::default(),
340360
parents: Vec::new(),
341361
};
@@ -355,6 +375,7 @@ impl<'a> DisplayableExecutionPlan<'a> {
355375
show_statistics: self.show_statistics,
356376
metric_types: self.metric_types.clone(),
357377
metric_categories: self.metric_categories.clone(),
378+
metric_names: self.metric_names.clone(),
358379
}
359380
}
360381

@@ -403,6 +424,7 @@ impl<'a> DisplayableExecutionPlan<'a> {
403424
show_schema: bool,
404425
metric_types: Vec<MetricType>,
405426
metric_categories: Option<Vec<MetricCategory>>,
427+
metric_names: Option<Vec<String>>,
406428
summary: Option<AnalyzeSummary>,
407429
}
408430
impl fmt::Display for Wrapper<'_> {
@@ -413,6 +435,7 @@ impl<'a> DisplayableExecutionPlan<'a> {
413435
show_schema: self.show_schema,
414436
metric_types: &self.metric_types,
415437
metric_categories: self.metric_categories.as_deref(),
438+
metric_names: self.metric_names.as_deref(),
416439
objects: HashMap::new(),
417440
parent_ids: Vec::new(),
418441
next_id: 0,
@@ -446,6 +469,7 @@ impl<'a> DisplayableExecutionPlan<'a> {
446469
show_schema: self.show_schema,
447470
metric_types: self.metric_types.clone(),
448471
metric_categories: self.metric_categories.clone(),
472+
metric_names: self.metric_names.clone(),
449473
summary: self.summary,
450474
}
451475
}
@@ -460,6 +484,7 @@ impl<'a> DisplayableExecutionPlan<'a> {
460484
show_schema: bool,
461485
metric_types: Vec<MetricType>,
462486
metric_categories: Option<Vec<MetricCategory>>,
487+
metric_names: Option<Vec<String>>,
463488
}
464489

465490
impl fmt::Display for Wrapper<'_> {
@@ -473,6 +498,7 @@ impl<'a> DisplayableExecutionPlan<'a> {
473498
show_schema: self.show_schema,
474499
metric_types: &self.metric_types,
475500
metric_categories: self.metric_categories.as_deref(),
501+
metric_names: self.metric_names.as_deref(),
476502
};
477503
visitor.pre_visit(self.plan)?;
478504
Ok(())
@@ -486,6 +512,7 @@ impl<'a> DisplayableExecutionPlan<'a> {
486512
show_schema: self.show_schema,
487513
metric_types: self.metric_types.clone(),
488514
metric_categories: self.metric_categories.clone(),
515+
metric_names: self.metric_names.clone(),
489516
}
490517
}
491518

@@ -544,6 +571,8 @@ struct IndentVisitor<'a, 'b> {
544571
metric_types: &'a [MetricType],
545572
/// Optional filter by semantic category (rows / bytes / timing).
546573
metric_categories: Option<&'a [MetricCategory]>,
574+
/// Optional filter by metric name.
575+
metric_names: Option<&'a [String]>,
547576
}
548577

549578
impl ExecutionPlanVisitor for IndentVisitor<'_, '_> {
@@ -563,6 +592,9 @@ impl ExecutionPlanVisitor for IndentVisitor<'_, '_> {
563592
if let Some(cats) = self.metric_categories {
564593
metrics = metrics.filter_by_categories(cats);
565594
}
595+
if let Some(names) = self.metric_names {
596+
metrics = metrics.filter_by_names(names);
597+
}
566598
write!(self.f, ", metrics=[{metrics}]")?;
567599
} else {
568600
write!(self.f, ", metrics=[]")?;
@@ -574,6 +606,9 @@ impl ExecutionPlanVisitor for IndentVisitor<'_, '_> {
574606
if let Some(cats) = self.metric_categories {
575607
metrics = metrics.filter_by_categories(cats);
576608
}
609+
if let Some(names) = self.metric_names {
610+
metrics = metrics.filter_by_names(names);
611+
}
577612
write!(self.f, ", metrics=[{metrics}]")?;
578613
} else {
579614
write!(self.f, ", metrics=[]")?;
@@ -616,6 +651,8 @@ struct GraphvizVisitor<'a, 'b> {
616651
metric_types: &'a [MetricType],
617652
/// Optional filter by semantic category
618653
metric_categories: Option<&'a [MetricCategory]>,
654+
/// Optional filter by metric name.
655+
metric_names: Option<&'a [String]>,
619656

620657
graphviz_builder: GraphvizBuilder,
621658
/// Used to record parent node ids when visiting a plan.
@@ -660,6 +697,9 @@ impl ExecutionPlanVisitor for GraphvizVisitor<'_, '_> {
660697
if let Some(cats) = self.metric_categories {
661698
metrics = metrics.filter_by_categories(cats);
662699
}
700+
if let Some(names) = self.metric_names {
701+
metrics = metrics.filter_by_names(names);
702+
}
663703
format!("metrics=[{metrics}]")
664704
} else {
665705
"metrics=[]".to_string()
@@ -671,6 +711,9 @@ impl ExecutionPlanVisitor for GraphvizVisitor<'_, '_> {
671711
if let Some(cats) = self.metric_categories {
672712
metrics = metrics.filter_by_categories(cats);
673713
}
714+
if let Some(names) = self.metric_names {
715+
metrics = metrics.filter_by_names(names);
716+
}
674717
format!("metrics=[{metrics}]")
675718
} else {
676719
"metrics=[]".to_string()
@@ -729,6 +772,7 @@ struct PgJsonExecutionPlanVisitor<'a> {
729772
show_schema: bool,
730773
metric_types: &'a [MetricType],
731774
metric_categories: Option<&'a [MetricCategory]>,
775+
metric_names: Option<&'a [String]>,
732776
objects: HashMap<u32, serde_json::Value>,
733777
parent_ids: Vec<u32>,
734778
next_id: u32,
@@ -813,6 +857,12 @@ impl PgJsonExecutionPlanVisitor<'_> {
813857
metrics
814858
};
815859

860+
let metrics = if let Some(names) = self.metric_names {
861+
metrics.filter_by_names(names)
862+
} else {
863+
metrics
864+
};
865+
816866
// Build the Extras bucket, while extracting PG-canonical keys to the
817867
// top level.
818868
let mut extras = serde_json::Map::new();

0 commit comments

Comments
 (0)