Skip to content

Add name filter to metrics #23719

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
gabotechs wants to merge 1 commit into
base: main
Choose a base branch
from
+81 −0
Draft

Add name filter to metrics #23719

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
31 changes: 31 additions & 0 deletions datafusion/physical-expr-common/src/metrics/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -418,6 +418,21 @@ impl MetricsSet {
.collect::<Vec<_>>();
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::<Vec<_>>();
Self { metrics }
}
}

impl Display for MetricsSet {
Expand Down Expand Up @@ -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"
);
}
}
50 changes: 50 additions & 0 deletions datafusion/physical-plan/src/display.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<Vec<MetricCategory>>,
/// Optional filter by metric names. Only metric names in this list
/// will be rendered.
metric_names: Option<Vec<String>>,
// (TreeRender) Maximum total width of the rendered tree
tree_maximum_render_width: usize,
/// Optional summary totals (currently only used by `pgjson`) — the total
Expand Down Expand Up @@ -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,
}
Expand All @@ -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,
}
Expand All @@ -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,
}
Expand Down Expand Up @@ -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<String>) -> 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;
Expand Down Expand Up @@ -279,6 +294,7 @@ impl<'a> DisplayableExecutionPlan<'a> {
show_schema: bool,
metric_types: Vec<MetricType>,
metric_categories: Option<Vec<MetricCategory>>,
metric_names: Option<Vec<String>>,
}
impl fmt::Display for Wrapper<'_> {
fn fmt(&self, f: &mut Formatter) -> fmt::Result {
Expand All @@ -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)
}
Expand All @@ -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(),
}
}

Expand All @@ -324,6 +342,7 @@ impl<'a> DisplayableExecutionPlan<'a> {
show_statistics: bool,
metric_types: Vec<MetricType>,
metric_categories: Option<Vec<MetricCategory>>,
metric_names: Option<Vec<String>>,
}
impl fmt::Display for Wrapper<'_> {
fn fmt(&self, f: &mut Formatter) -> fmt::Result {
Expand All @@ -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(),
};
Expand All @@ -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(),
}
}

Expand Down Expand Up @@ -403,6 +424,7 @@ impl<'a> DisplayableExecutionPlan<'a> {
show_schema: bool,
metric_types: Vec<MetricType>,
metric_categories: Option<Vec<MetricCategory>>,
metric_names: Option<Vec<String>>,
summary: Option<AnalyzeSummary>,
}
impl fmt::Display for Wrapper<'_> {
Expand All @@ -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,
Expand Down Expand Up @@ -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,
}
}
Expand All @@ -460,6 +484,7 @@ impl<'a> DisplayableExecutionPlan<'a> {
show_schema: bool,
metric_types: Vec<MetricType>,
metric_categories: Option<Vec<MetricCategory>>,
metric_names: Option<Vec<String>>,
}

impl fmt::Display for Wrapper<'_> {
Expand All @@ -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(())
Expand All @@ -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(),
}
}

Expand Down Expand Up @@ -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<'_, '_> {
Expand All @@ -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=[]")?;
Expand All @@ -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=[]")?;
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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()
Expand All @@ -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()
Expand Down Expand Up @@ -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<u32, serde_json::Value>,
parent_ids: Vec<u32>,
next_id: u32,
Expand Down Expand Up @@ -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();
Expand Down
Loading