use std::{collections::BTreeMap, time::Duration};
use crate::observability::graph_observation::GraphMeasuredPeakRow;
mod aggregate;
mod backend_descriptor;
mod boundary;
mod correlation;
mod descriptor;
mod dynamic_filter;
mod graph;
mod join;
mod json;
mod memory;
mod operator;
mod source_read;
mod task;
use correlation::inherited_query_id;
pub(crate) use query_planning::correlation::CompiledTaskId;
use query_planning::correlation::{QueryId, TraceId};
pub(crate) use self::{
aggregate::AggregateStrategyDecisionMetrics,
correlation::PlanTaskMapping,
dynamic_filter::DynamicFilterRuntimeMetrics,
task::{TaskOperatorMetrics, TaskRuntimeMetrics},
};
pub use self::{
backend_descriptor::{
BACKEND_ATTEMPT_DURATION_BUCKETS_NS, BACKEND_METRIC_DESCRIPTORS, BackendDurationProjectionError, BackendMetric,
BackendMetricDescriptor, BackendMetricKind, BackendMetricLabel, BackendMetricUnit, BackendMetricUpdatePoint,
backend_metric_descriptors, project_backend_duration_histogram,
},
boundary::BoundaryMetrics,
descriptor::{MetricAggregation, MetricDescriptor, MetricDirection, MetricUnit},
dynamic_filter::DynamicFilterSourceSnapshot,
graph::GraphAlgorithmRuntimeMetrics,
join::{JoinOperatorMetrics, JoinStrategyDecisionMetrics},
json::QUERY_METRICS_JSON_SCHEMA,
operator::OperatorMetrics,
source_read::{SourceReadParquetEvidenceMetrics, SourceReadRuntimeMetrics},
task::BudgetedExecutionEvent,
};
use crate::observability::{LakehouseCacheFillEventSnapshot, MetricsCaptureMode, MetricsCompleteness};
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct QueryMetrics {
pub query_id: QueryId,
pub trace_id: Option<TraceId>,
pub boundary: BoundaryMetrics,
pub operators: Vec<OperatorMetrics>,
pub graph_algorithm_runtime_rows: Vec<GraphAlgorithmRuntimeMetrics>,
pub graph_claim_calibration_rows: Vec<GraphMeasuredPeakRow>,
pub join_strategy_decisions: Vec<JoinStrategyDecisionMetrics>,
pub aggregate_strategy_decisions: Vec<AggregateStrategyDecisionMetrics>,
pub dynamic_filter_runtime_rows: Vec<DynamicFilterRuntimeMetrics>,
pub source_read_runtime_rows: Vec<SourceReadRuntimeMetrics>,
pub lakehouse_cache_fill_events: Vec<LakehouseCacheFillEventSnapshot>,
pub input_batches: u64,
pub input_rows: u64,
pub peak_device_memory_bytes: Option<u64>,
pub resident_dataframe_peak_bytes: Option<u64>,
pub allocator_high_water_bytes: Option<u64>,
pub max_task_observed_peak_bytes: Option<u64>,
pub max_task_retained_bytes: Option<u64>,
#[doc(hidden)]
pub task_kind_by_compiled_task_id: BTreeMap<String, String>,
pub task_operator_metrics: BTreeMap<String, TaskOperatorMetrics>,
pub task_runtime_metrics: BTreeMap<String, TaskRuntimeMetrics>,
pub plan_task_mapping_rows: Vec<PlanTaskMapping>,
pub plan_task_mapping_datafusion_shape: Option<String>,
pub budgeted_execution_events: Vec<BudgetedExecutionEvent>,
pub source_chunks_planned: u64,
pub source_chunks_attempted: u64,
pub source_chunks_committed: u64,
pub source_chunks_retried: u64,
pub source_read_prefetches_started: u64,
pub source_read_prefetches_completed: u64,
pub source_read_prefetches_skipped: u64,
pub source_read_prefetches_abandoned: u64,
pub row_groups_planned: u64,
pub row_groups_committed: u64,
pub partial_aggregate_count: u64,
pub partial_aggregate_input_rows: u64,
pub partial_aggregate_rows: u64,
pub partial_aggregate_retained_bytes: u64,
pub final_merge_attempts: u64,
pub final_merge_levels: u64,
pub final_merge_max_fan_in: u64,
pub retry_count: u64,
pub retry_resource_exhausted_count: u64,
pub hash_join_build_count: u64,
pub filtered_join_build_count: u64,
pub one_shot_join_count: u64,
pub hash_join_hash_table_bytes: u64,
pub filtered_join_hash_table_bytes: u64,
pub filtered_join_probe_state_bytes: u64,
pub mixed_join_size_data_bytes: u64,
pub one_shot_join_index_bytes: u64,
pub one_shot_join_materialization_bytes: u64,
pub one_shot_join_output_bytes: u64,
pub parquet_files: u64,
pub parquet_bytes_known: u64,
pub row_groups_selected: u64,
pub read_columns: u64,
pub source_metadata_load_time: Duration,
pub source_read_queue_wait_time: Duration,
pub parquet_read_time: Duration,
pub source_read_time: Duration,
pub operator_time: Duration,
pub chunk_operator_time: Duration,
pub final_merge_time: Duration,
pub final_sort_time: Duration,
pub materialization_time: Duration,
pub compiled_pipeline_attempts: u64,
pub compiled_pipeline_tasks_executed: u64,
pub compiled_pipeline_host_arrow_sinks: u64,
pub dynamic_filter_produced_count: u64,
pub dynamic_filter_skipped_count: u64,
pub dynamic_filter_applied_count: u64,
pub dynamic_filter_row_groups_pruned: u64,
pub dynamic_filter_source_bytes_pruned: u64,
pub dynamic_filter_reader_predicate_rows_before: u64,
pub dynamic_filter_reader_predicate_rows_after: u64,
pub dynamic_filter_reader_predicate_rows_filtered: u64,
pub dynamic_filter_post_decode_rows_before: u64,
pub dynamic_filter_post_decode_rows_after: u64,
pub dynamic_filter_post_decode_rows_filtered: u64,
pub dynamic_filter_source_post_decode_rows_filtered: u64,
pub dynamic_filter_filter_post_decode_rows_filtered: u64,
pub dynamic_filter_join_probe_post_decode_rows_filtered: u64,
pub projection: Option<MetricsProjection>,
pub join_gather_map_bytes: u64,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct MetricsProjection {
pub(crate) revision: Option<u64>,
pub(crate) sealed: bool,
pub completeness: MetricsCompleteness,
}
impl MetricsProjection {
fn merged(self, other: Self) -> Self {
let capture_mode = match (self.completeness.capture_mode, other.completeness.capture_mode) {
(MetricsCaptureMode::Bounded, MetricsCaptureMode::Bounded) => MetricsCaptureMode::Bounded,
_ => MetricsCaptureMode::Disabled,
};
let detail_rows_dropped =
self.completeness.detail_rows_dropped.saturating_add(other.completeness.detail_rows_dropped);
Self {
revision: None,
sealed: self.sealed && other.sealed,
completeness: MetricsCompleteness {
capture_mode,
detail_rows_dropped,
truncated: self.completeness.truncated || other.completeness.truncated,
totals_saturated: self.completeness.totals_saturated || other.completeness.totals_saturated,
},
}
}
}
impl QueryMetrics {
pub fn new(query_id: impl Into<QueryId>) -> Self {
Self { query_id: query_id.into(), ..Default::default() }
}
pub fn query_id(&self) -> &QueryId {
&self.query_id
}
pub fn set_query_id(&mut self, query_id: impl Into<QueryId>) {
self.query_id = query_id.into();
self.stamp_existing_rows_with_query_id();
}
#[must_use]
pub fn with_query_id(mut self, query_id: impl Into<QueryId>) -> Self {
self.set_query_id(query_id);
self
}
#[must_use]
pub fn with_plan_task_mapping_datafusion_shape(mut self, shape: impl Into<String>) -> Self {
let shape = shape.into();
if !shape.is_empty() {
self.plan_task_mapping_datafusion_shape = Some(shape.clone());
for row in &mut self.plan_task_mapping_rows {
if row.datafusion_shape.is_none() {
row.datafusion_shape = Some(shape.clone());
}
}
}
self
}
pub(super) fn inherited_query_id(&self) -> Option<QueryId> {
inherited_query_id(&self.query_id)
}
pub(crate) fn stamp_existing_rows_with_query_id(&mut self) {
let Some(query_id) = self.inherited_query_id() else {
return;
};
for row in &mut self.budgeted_execution_events {
row.query_id = Some(query_id.clone());
}
for row in &mut self.dynamic_filter_runtime_rows {
row.query_id = Some(query_id.clone());
}
for row in &mut self.source_read_runtime_rows {
row.query_id = Some(query_id.clone());
}
for row in &mut self.join_strategy_decisions {
row.query_id = Some(query_id.clone());
}
for row in &mut self.aggregate_strategy_decisions {
row.query_id = Some(query_id.clone());
}
for row in self.task_operator_metrics.values_mut() {
row.query_id = Some(query_id.clone());
}
for row in self.task_runtime_metrics.values_mut() {
row.query_id = Some(query_id.clone());
}
for row in &mut self.plan_task_mapping_rows {
row.query_id = query_id.clone();
}
for row in &mut self.operators {
row.query_id = Some(query_id.clone());
if let Some(join) = &mut row.join {
join.query_id = Some(query_id.clone());
}
}
for row in &mut self.graph_algorithm_runtime_rows {
row.query_id = Some(query_id.clone());
}
for row in &mut self.graph_claim_calibration_rows {
row.query_id = Some(query_id.clone());
}
}
pub fn merge_from(&mut self, mut other: QueryMetrics) {
if !self.query_id.is_assigned() && other.query_id.is_assigned() {
self.query_id = other.query_id.clone();
self.stamp_existing_rows_with_query_id();
}
if self.query_id.is_assigned() && !other.query_id.is_assigned() {
other.set_query_id(self.query_id.clone());
}
if self.trace_id.is_none() {
self.trace_id = other.trace_id.clone();
}
self.boundary.merge_from(other.boundary);
self.operators.extend(other.operators);
self.graph_algorithm_runtime_rows.extend(other.graph_algorithm_runtime_rows);
self.graph_claim_calibration_rows.extend(other.graph_claim_calibration_rows);
self.join_strategy_decisions.extend(other.join_strategy_decisions);
self.aggregate_strategy_decisions.extend(other.aggregate_strategy_decisions);
self.dynamic_filter_runtime_rows.extend(other.dynamic_filter_runtime_rows);
self.source_read_runtime_rows.extend(other.source_read_runtime_rows);
self.lakehouse_cache_fill_events.extend(other.lakehouse_cache_fill_events);
self.input_batches = self.input_batches.saturating_add(other.input_batches);
self.input_rows = self.input_rows.saturating_add(other.input_rows);
self.peak_device_memory_bytes = self.peak_device_memory_bytes.max(other.peak_device_memory_bytes);
self.resident_dataframe_peak_bytes =
self.resident_dataframe_peak_bytes.max(other.resident_dataframe_peak_bytes);
self.allocator_high_water_bytes = self.allocator_high_water_bytes.max(other.allocator_high_water_bytes);
self.max_task_observed_peak_bytes = self.max_task_observed_peak_bytes.max(other.max_task_observed_peak_bytes);
self.max_task_retained_bytes = self.max_task_retained_bytes.max(other.max_task_retained_bytes);
self.task_kind_by_compiled_task_id.extend(other.task_kind_by_compiled_task_id);
self.plan_task_mapping_rows.extend(other.plan_task_mapping_rows);
if self.plan_task_mapping_datafusion_shape.is_none() {
self.plan_task_mapping_datafusion_shape = other.plan_task_mapping_datafusion_shape;
}
for (operator_kind, metrics) in other.task_operator_metrics {
self.task_operator_metrics
.entry(operator_kind.clone())
.or_insert_with(|| TaskOperatorMetrics { operator_kind, ..Default::default() })
.merge_from(metrics);
}
for (compiled_task_id, metrics) in other.task_runtime_metrics {
self.task_runtime_metrics
.entry(compiled_task_id.clone())
.or_insert_with(|| TaskRuntimeMetrics { compiled_task_id, ..Default::default() })
.merge_from(metrics);
}
self.budgeted_execution_events.extend(other.budgeted_execution_events);
self.source_chunks_planned = self.source_chunks_planned.saturating_add(other.source_chunks_planned);
self.source_chunks_attempted = self.source_chunks_attempted.saturating_add(other.source_chunks_attempted);
self.source_chunks_committed = self.source_chunks_committed.saturating_add(other.source_chunks_committed);
self.source_chunks_retried = self.source_chunks_retried.saturating_add(other.source_chunks_retried);
self.source_read_prefetches_started =
self.source_read_prefetches_started.saturating_add(other.source_read_prefetches_started);
self.source_read_prefetches_completed =
self.source_read_prefetches_completed.saturating_add(other.source_read_prefetches_completed);
self.source_read_prefetches_skipped =
self.source_read_prefetches_skipped.saturating_add(other.source_read_prefetches_skipped);
self.source_read_prefetches_abandoned =
self.source_read_prefetches_abandoned.saturating_add(other.source_read_prefetches_abandoned);
self.row_groups_planned = self.row_groups_planned.saturating_add(other.row_groups_planned);
self.row_groups_committed = self.row_groups_committed.saturating_add(other.row_groups_committed);
self.partial_aggregate_count = self.partial_aggregate_count.saturating_add(other.partial_aggregate_count);
self.partial_aggregate_input_rows =
self.partial_aggregate_input_rows.saturating_add(other.partial_aggregate_input_rows);
self.partial_aggregate_rows = self.partial_aggregate_rows.saturating_add(other.partial_aggregate_rows);
self.partial_aggregate_retained_bytes =
self.partial_aggregate_retained_bytes.max(other.partial_aggregate_retained_bytes);
self.final_merge_attempts = self.final_merge_attempts.saturating_add(other.final_merge_attempts);
self.final_merge_levels = self.final_merge_levels.saturating_add(other.final_merge_levels);
self.final_merge_max_fan_in = self.final_merge_max_fan_in.max(other.final_merge_max_fan_in);
self.retry_count = self.retry_count.saturating_add(other.retry_count);
self.retry_resource_exhausted_count =
self.retry_resource_exhausted_count.saturating_add(other.retry_resource_exhausted_count);
self.hash_join_build_count = self.hash_join_build_count.saturating_add(other.hash_join_build_count);
self.filtered_join_build_count = self.filtered_join_build_count.saturating_add(other.filtered_join_build_count);
self.one_shot_join_count = self.one_shot_join_count.saturating_add(other.one_shot_join_count);
self.hash_join_hash_table_bytes =
self.hash_join_hash_table_bytes.saturating_add(other.hash_join_hash_table_bytes);
self.filtered_join_hash_table_bytes =
self.filtered_join_hash_table_bytes.saturating_add(other.filtered_join_hash_table_bytes);
self.filtered_join_probe_state_bytes =
self.filtered_join_probe_state_bytes.saturating_add(other.filtered_join_probe_state_bytes);
self.mixed_join_size_data_bytes =
self.mixed_join_size_data_bytes.saturating_add(other.mixed_join_size_data_bytes);
self.one_shot_join_index_bytes = self.one_shot_join_index_bytes.saturating_add(other.one_shot_join_index_bytes);
self.one_shot_join_materialization_bytes =
self.one_shot_join_materialization_bytes.saturating_add(other.one_shot_join_materialization_bytes);
self.one_shot_join_output_bytes =
self.one_shot_join_output_bytes.saturating_add(other.one_shot_join_output_bytes);
self.parquet_files = self.parquet_files.saturating_add(other.parquet_files);
self.parquet_bytes_known = self.parquet_bytes_known.saturating_add(other.parquet_bytes_known);
self.row_groups_selected = self.row_groups_selected.saturating_add(other.row_groups_selected);
self.read_columns = self.read_columns.saturating_add(other.read_columns);
self.source_metadata_load_time += other.source_metadata_load_time;
self.source_read_queue_wait_time += other.source_read_queue_wait_time;
self.parquet_read_time += other.parquet_read_time;
self.source_read_time += other.source_read_time;
self.operator_time += other.operator_time;
self.chunk_operator_time += other.chunk_operator_time;
self.final_merge_time += other.final_merge_time;
self.final_sort_time += other.final_sort_time;
self.materialization_time += other.materialization_time;
self.compiled_pipeline_attempts =
self.compiled_pipeline_attempts.saturating_add(other.compiled_pipeline_attempts);
self.compiled_pipeline_tasks_executed =
self.compiled_pipeline_tasks_executed.saturating_add(other.compiled_pipeline_tasks_executed);
self.compiled_pipeline_host_arrow_sinks =
self.compiled_pipeline_host_arrow_sinks.saturating_add(other.compiled_pipeline_host_arrow_sinks);
self.dynamic_filter_produced_count =
self.dynamic_filter_produced_count.saturating_add(other.dynamic_filter_produced_count);
self.dynamic_filter_skipped_count =
self.dynamic_filter_skipped_count.saturating_add(other.dynamic_filter_skipped_count);
self.dynamic_filter_applied_count =
self.dynamic_filter_applied_count.saturating_add(other.dynamic_filter_applied_count);
self.dynamic_filter_row_groups_pruned =
self.dynamic_filter_row_groups_pruned.saturating_add(other.dynamic_filter_row_groups_pruned);
self.dynamic_filter_source_bytes_pruned =
self.dynamic_filter_source_bytes_pruned.saturating_add(other.dynamic_filter_source_bytes_pruned);
self.dynamic_filter_reader_predicate_rows_before = self
.dynamic_filter_reader_predicate_rows_before
.saturating_add(other.dynamic_filter_reader_predicate_rows_before);
self.dynamic_filter_reader_predicate_rows_after = self
.dynamic_filter_reader_predicate_rows_after
.saturating_add(other.dynamic_filter_reader_predicate_rows_after);
self.dynamic_filter_reader_predicate_rows_filtered = self
.dynamic_filter_reader_predicate_rows_filtered
.saturating_add(other.dynamic_filter_reader_predicate_rows_filtered);
self.dynamic_filter_post_decode_rows_before =
self.dynamic_filter_post_decode_rows_before.saturating_add(other.dynamic_filter_post_decode_rows_before);
self.dynamic_filter_post_decode_rows_after =
self.dynamic_filter_post_decode_rows_after.saturating_add(other.dynamic_filter_post_decode_rows_after);
self.dynamic_filter_post_decode_rows_filtered = self
.dynamic_filter_post_decode_rows_filtered
.saturating_add(other.dynamic_filter_post_decode_rows_filtered);
self.dynamic_filter_source_post_decode_rows_filtered = self
.dynamic_filter_source_post_decode_rows_filtered
.saturating_add(other.dynamic_filter_source_post_decode_rows_filtered);
self.dynamic_filter_filter_post_decode_rows_filtered = self
.dynamic_filter_filter_post_decode_rows_filtered
.saturating_add(other.dynamic_filter_filter_post_decode_rows_filtered);
self.dynamic_filter_join_probe_post_decode_rows_filtered = self
.dynamic_filter_join_probe_post_decode_rows_filtered
.saturating_add(other.dynamic_filter_join_probe_post_decode_rows_filtered);
self.join_gather_map_bytes = self.join_gather_map_bytes.saturating_add(other.join_gather_map_bytes);
self.projection = match (self.projection, other.projection) {
(Some(left), Some(right)) => Some(left.merged(right)),
(left, right) => left.or(right),
};
self.stamp_existing_rows_with_query_id();
}
}
fn optional_str(value: Option<&str>) -> &str {
value.unwrap_or("none")
}
fn optional_u64(value: Option<u64>) -> String {
value.map_or_else(|| "none".to_string(), |value| value.to_string())
}
fn optional_i64(value: Option<i64>) -> String {
value.map_or_else(|| "none".to_string(), |value| value.to_string())
}