diff --git a/benchmarks/src/stats.rs b/benchmarks/src/stats.rs index 764800bad..7571b8666 100644 --- a/benchmarks/src/stats.rs +++ b/benchmarks/src/stats.rs @@ -1,7 +1,7 @@ use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion}; use datafusion::physical_plan::ExecutionPlan; -use datafusion::physical_plan::metrics::MetricsSet; -use datafusion_distributed::{NetworkBoundaryExt, Stage}; +use datafusion::physical_plan::metrics::{MetricValue, MetricsSet}; +use datafusion_distributed::{NetworkBoundaryExt, QErrorMetric, STATS_Q_ERROR_METRIC, Stage}; use sketches_ddsketch::{Config, DDSketch}; use std::sync::Arc; @@ -11,20 +11,17 @@ pub struct StatsEstimationQError { pub p95: f64, } -/// Computes P50 and P95 q-error between the sampled byte estimate and the actual output bytes for -/// every dynamically sampled stage boundary in `plan`. +/// Collects P50 and P95 of the q-error recorded at every dynamically sampled stage boundary in +/// `plan`. pub fn stats_estimation_q_error(plan: &Arc) -> Option { let mut boundary_q_errors = DDSketch::new(Config::defaults()); let _ = plan.apply(|node| { if let Some(boundary) = node.as_network_boundary() && let Stage::Local(input_stage) = boundary.input_stage() - && let Some(sampled_bytes) = metric_total(&input_stage.metrics_set, "sampled_bytes") - && let Some(actual_bytes) = node - .metrics() - .and_then(|metrics| metric_total(&metrics, "output_bytes")) + && let Some(q_error) = q_error_metric_value(&input_stage.metrics_set) { - boundary_q_errors.add(q_error(sampled_bytes, actual_bytes)); + boundary_q_errors.add(q_error); } Ok(TreeNodeRecursion::Continue) }); @@ -39,19 +36,14 @@ fn q_error_percentiles(q_errors: &DDSketch) -> Option { }) } -fn metric_total(metrics: &MetricsSet, name: &str) -> Option { - metrics - .sum(|metric| metric.value().name() == name) - .map(|value| value.as_usize()) -} - -/// Q-error is the standard cardinality-estimation metric because it treats equal-factor over- and -/// underestimates symmetrically. See https://www.vldb.org/pvldb/vol2/vldb09-657.pdf and -/// https://vldb.org/pvldb/vol9/p204-leis.pdf. -fn q_error(estimated: usize, actual: usize) -> f64 { - let estimated = estimated.max(1) as f64; - let actual = actual.max(1) as f64; - (estimated / actual).max(actual / estimated) +fn q_error_metric_value(metrics: &MetricsSet) -> Option { + metrics.iter().find_map(|metric| match metric.value() { + MetricValue::Custom { name, value } if name == STATS_Q_ERROR_METRIC => value + .as_any() + .downcast_ref::() + .map(QErrorMetric::value), + _ => None, + }) } pub fn median(mut values: Vec) -> Option { @@ -90,4 +82,12 @@ mod tests { assert!((49.0..=51.0).contains(&percentiles.p50)); assert!((94.0..=96.0).contains(&percentiles.p95)); } + + #[test] + fn reads_q_error_metric_value() { + let mut metrics = MetricsSet::new(); + metrics.push(QErrorMetric::new_metric(STATS_Q_ERROR_METRIC, 2.75)); + + assert_eq!(q_error_metric_value(&metrics), Some(2.75)); + } } diff --git a/src/common/mod.rs b/src/common/mod.rs index 07463ca0b..8929512d9 100644 --- a/src/common/mod.rs +++ b/src/common/mod.rs @@ -1,5 +1,6 @@ mod children_helpers; mod once_lock; +mod record_batch; mod recursion; mod task_context_helpers; mod time; @@ -8,6 +9,7 @@ mod vec; pub(crate) use children_helpers::require_one_child; pub(crate) use once_lock::OnceLockResult; +pub(crate) use record_batch::logical_record_batch_size; pub(crate) use recursion::TreeNodeExt; pub(crate) use task_context_helpers::task_ctx_with_extension; pub(crate) use time::now_ns; diff --git a/src/common/record_batch.rs b/src/common/record_batch.rs new file mode 100644 index 000000000..9d919e21c --- /dev/null +++ b/src/common/record_batch.rs @@ -0,0 +1,10 @@ +use datafusion::arrow::array::RecordBatch; + +/// Returns the logical size of a batch's slices, excluding unused backing-buffer capacity. +pub(crate) fn logical_record_batch_size(batch: &RecordBatch) -> usize { + batch + .columns() + .iter() + .map(|column| column.to_data().get_slice_memory_size().unwrap_or(0)) + .sum() +} diff --git a/src/coordinator/mod.rs b/src/coordinator/mod.rs index c1a8a8dd2..374352eba 100644 --- a/src/coordinator/mod.rs +++ b/src/coordinator/mod.rs @@ -6,4 +6,6 @@ mod prepare_static_plan; mod query_coordinator; pub use distributed::DistributedExec; +pub use prepare_dynamic_plan::{ESTIMATED_OUTPUT_BYTES_METRIC, ESTIMATED_PCT_SAMPLED_METRIC}; + pub(crate) use metrics_store::MetricsStore; diff --git a/src/coordinator/prepare_dynamic_plan.rs b/src/coordinator/prepare_dynamic_plan.rs index 0c0cb23e9..00bc4d942 100644 --- a/src/coordinator/prepare_dynamic_plan.rs +++ b/src/coordinator/prepare_dynamic_plan.rs @@ -24,6 +24,9 @@ use std::any::TypeId; use std::sync::Arc; use tokio_stream::wrappers::UnboundedReceiverStream; +pub const ESTIMATED_PCT_SAMPLED_METRIC: &str = "estimated_pct_sampled"; +pub const ESTIMATED_OUTPUT_BYTES_METRIC: &str = "estimated_output_bytes"; + pub(super) async fn prepare_dynamic_plan( query_coordinator: &QueryCoordinator, base_plan: &Arc, @@ -286,7 +289,7 @@ async fn gather_runtime_statistics( } else if let Some(estimated_driver_path_leaf_rows) = estimated_driver_path_leaf_rows(plan) { // The stage is still producing. Estimate how far along it is from the fraction of the // driver-path leaf rows consumed so far. - (rows_pulled_from_leafs as f32 / estimated_driver_path_leaf_rows as f32).min(1.0) + rows_pulled_from_leafs as f32 / estimated_driver_path_leaf_rows as f32 } else { // We can't measure progress (no leaf-row estimate, or nothing pulled from the leaves // yet even though we're not at EOS): fall back rather than dividing by ~0. @@ -294,10 +297,11 @@ async fn gather_runtime_statistics( }; new_metrics.push(MaxGaugeMetric::new_metric( - "estimated_pct_sampled", + ESTIMATED_PCT_SAMPLED_METRIC, (estimated_pct_sampled * 100.) as usize, )); + let estimated_pct_sampled = estimated_pct_sampled.min(1.0); let total_num_rows = (rows_ready as f32 / estimated_pct_sampled) as usize; if total_num_rows == 0 { @@ -308,7 +312,7 @@ async fn gather_runtime_statistics( let total_byte_size: usize = per_col_byte_size.iter().sum(); new_metrics.push(BytesCounterMetric::new_metric( - "estimated_output_bytes", + ESTIMATED_OUTPUT_BYTES_METRIC, total_byte_size, )); diff --git a/src/execution_plans/sampler.rs b/src/execution_plans/sampler.rs index 8eb5c3642..b4c3c5acc 100644 --- a/src/execution_plans/sampler.rs +++ b/src/execution_plans/sampler.rs @@ -1,4 +1,4 @@ -use crate::common::{TreeNodeExt, require_one_child, vec_cast}; +use crate::common::{TreeNodeExt, logical_record_batch_size, require_one_child, vec_cast}; use crate::{ BytesCounterMetric, BytesMetricExt, GaugeMetricExt, LatencyMetricExt, LoadInfo, MaxGaugeMetric, MaxLatencyMetric, P50LatencyMetric, @@ -6,6 +6,7 @@ use crate::{ use datafusion::arrow::array::Array; use datafusion::arrow::array::ArrayRef; use datafusion::arrow::record_batch::RecordBatch; +use datafusion::common::format::MetricCategory; use datafusion::common::runtime::SpawnedTask; use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion}; use datafusion::common::{DataFusionError, Result, exec_err}; @@ -13,13 +14,13 @@ use datafusion::common::{HashSet, ScalarValue}; use datafusion::execution::memory_pool::{MemoryConsumer, MemoryReservation}; use datafusion::execution::{SendableRecordBatchStream, TaskContext}; use datafusion::physical_expr_common::metrics::{Gauge, MetricValue, MetricsSet}; -use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, MetricBuilder, Time}; +use datafusion::physical_plan::metrics::{Count, ExecutionPlanMetricsSet, MetricBuilder, Time}; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use datafusion::physical_plan::{ DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, PlanProperties, }; use futures::stream::FusedStream; -use futures::{Stream, StreamExt, TryFutureExt}; +use futures::{Stream, StreamExt, TryFutureExt, TryStreamExt}; use std::collections::VecDeque; use std::fmt::{Debug, Formatter}; use std::pin::Pin; @@ -68,6 +69,8 @@ pub(crate) struct SamplerExecMetrics { bytes_ready: BytesCounterMetric, /// Elapsed compute while sampling. elapsed_compute: Time, + /// Number of output bytes. + output_bytes: Count, } impl SamplerExecMetrics { @@ -88,6 +91,13 @@ impl SamplerExecMetrics { bdr().build(MetricValue::ElapsedCompute(time.clone())); time }, + output_bytes: { + let count = Count::new(); + bdr() + .with_category(MetricCategory::Bytes) + .build(MetricValue::OutputBytes(count.clone())); + count + }, } } } @@ -273,11 +283,16 @@ impl PartitionSampler { Ok(peek.chain(input_stream).boxed()) }); + let output_bytes = self.metrics.output_bytes.clone(); + let stream = async move { task.await .map_err(|err| DataFusionError::Internal(err.to_string()))? } - .try_flatten_stream(); + .try_flatten_stream() + .inspect_ok(move |b| { + output_bytes.add(logical_record_batch_size(b)); + }); self.stream .lock() diff --git a/src/lib.rs b/src/lib.rs index b98856d13..f5885bc29 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -17,7 +17,9 @@ mod worker_resolver; #[cfg(feature = "grpc")] pub use arrow_ipc::CompressionType; -pub use coordinator::DistributedExec; +pub use coordinator::{ + DistributedExec, ESTIMATED_OUTPUT_BYTES_METRIC, ESTIMATED_PCT_SAMPLED_METRIC, +}; pub use distributed_ext::{DistributedExt, DistributedGetterExt}; pub use distributed_planner::{ DistributedConfig, NetworkBoundary, NetworkBoundaryExt, SessionStateBuilderExt, @@ -31,7 +33,7 @@ pub use metrics::{ AvgLatencyMetric, BytesCounterMetric, BytesMetricExt, DISTRIBUTED_DATAFUSION_TASK_ID_LABEL, DistributedMetricsFormat, FirstLatencyMetric, GaugeMetricExt, LatencyMetricExt, MaxGaugeMetric, MaxLatencyMetric, MinLatencyMetric, P50LatencyMetric, P75LatencyMetric, P95LatencyMetric, - P99LatencyMetric, rewrite_distributed_plan_with_metrics, + P99LatencyMetric, QErrorMetric, STATS_Q_ERROR_METRIC, rewrite_distributed_plan_with_metrics, }; #[cfg(any(feature = "integration", test))] diff --git a/src/metrics/mod.rs b/src/metrics/mod.rs index fd1139426..fb94081ea 100644 --- a/src/metrics/mod.rs +++ b/src/metrics/mod.rs @@ -1,6 +1,7 @@ mod bytes_metric; mod latency_metric; mod max_gauge_metric; +mod q_error_metric; mod task_metrics_collector; mod task_metrics_rewriter; @@ -10,8 +11,11 @@ pub use latency_metric::{ P50LatencyMetric, P75LatencyMetric, P95LatencyMetric, P99LatencyMetric, }; pub use max_gauge_metric::{GaugeMetricExt, MaxGaugeMetric}; +pub use q_error_metric::QErrorMetric; pub(crate) use task_metrics_collector::collect_plan_metrics; -pub use task_metrics_rewriter::{DistributedMetricsFormat, rewrite_distributed_plan_with_metrics}; +pub use task_metrics_rewriter::{ + DistributedMetricsFormat, STATS_Q_ERROR_METRIC, rewrite_distributed_plan_with_metrics, +}; /// Label used to annotate metrics in execution plan nodes with the task in which they were executed. /// Note that the same task id may be used in multiple stages. pub const DISTRIBUTED_DATAFUSION_TASK_ID_LABEL: &str = "task_id"; diff --git a/src/metrics/q_error_metric.rs b/src/metrics/q_error_metric.rs new file mode 100644 index 000000000..d4cf1cacf --- /dev/null +++ b/src/metrics/q_error_metric.rs @@ -0,0 +1,109 @@ +use datafusion::physical_plan::Metric; +use datafusion::physical_plan::metrics::{CustomMetricValue, MetricValue}; +use std::any::Any; +use std::borrow::Cow; +use std::fmt::{Debug, Display, Formatter}; +use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering::Relaxed}; + +/// A floating-point q-error metric. +/// +/// When multiple values are aggregated, the largest q-error is retained so a bad estimate is not +/// hidden. A q-error of `1.0` represents a perfect estimate and is therefore the neutral value. +#[derive(Clone)] +pub struct QErrorMetric { + value: Arc, +} + +impl Default for QErrorMetric { + fn default() -> Self { + Self::from_value(1.0) + } +} + +impl QErrorMetric { + pub fn new_metric(name: impl Into>, value: f64) -> Arc { + Arc::new(Metric::new( + MetricValue::Custom { + name: name.into(), + value: Arc::new(Self::from_value(value)), + }, + None, + )) + } + + pub fn from_value(value: f64) -> Self { + Self { + value: Arc::new(AtomicU64::new(value.to_bits())), + } + } + + pub fn value(&self) -> f64 { + f64::from_bits(self.value.load(Relaxed)) + } +} + +impl Debug for QErrorMetric { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.debug_tuple("QErrorMetric").field(&self.value()).finish() + } +} + +impl Display for QErrorMetric { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + let value = format!("{:.2}", self.value()); + write!(f, "{}x", value.trim_end_matches('0').trim_end_matches('.')) + } +} + +impl CustomMetricValue for QErrorMetric { + fn new_empty(&self) -> Arc { + Arc::new(Self::default()) + } + + fn aggregate(&self, other: Arc) { + let Some(other) = other.as_any().downcast_ref::() else { + return; + }; + let other = other.value(); + let _ = self.value.fetch_update(Relaxed, Relaxed, |value| { + Some(f64::from_bits(value).max(other).to_bits()) + }); + } + + fn as_any(&self) -> &dyn Any { + self + } + + fn is_eq(&self, other: &Arc) -> bool { + other + .as_any() + .downcast_ref::() + .is_some_and(|other| other.value().to_bits() == self.value().to_bits()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn default_is_a_perfect_estimate() { + assert_eq!(QErrorMetric::default().value(), 1.0); + } + + #[test] + fn display_formats_as_a_factor() { + assert_eq!(QErrorMetric::from_value(1.456).to_string(), "1.46x"); + assert_eq!(QErrorMetric::from_value(1.5).to_string(), "1.5x"); + assert_eq!(QErrorMetric::from_value(1.0).to_string(), "1x"); + } + + #[test] + fn aggregate_retains_the_largest_q_error() { + let metric = QErrorMetric::from_value(2.0); + metric.aggregate(Arc::new(QErrorMetric::from_value(1.5))); + metric.aggregate(Arc::new(QErrorMetric::from_value(4.25))); + assert_eq!(metric.value(), 4.25); + } +} diff --git a/src/metrics/task_metrics_rewriter.rs b/src/metrics/task_metrics_rewriter.rs index 8913ff875..22d2a0383 100644 --- a/src/metrics/task_metrics_rewriter.rs +++ b/src/metrics/task_metrics_rewriter.rs @@ -1,11 +1,11 @@ use crate::common::TreeNodeExt; -use crate::coordinator::{DistributedExec, MetricsStore}; +use crate::coordinator::{DistributedExec, ESTIMATED_OUTPUT_BYTES_METRIC, MetricsStore}; use crate::distributed_planner::NetworkBoundaryExt; -use crate::execution_plans::MetricsWrapperExec; +use crate::execution_plans::{MetricsWrapperExec, SamplerExec}; use crate::metrics::DISTRIBUTED_DATAFUSION_TASK_ID_LABEL; use crate::metrics::collect_plan_metrics; use crate::stage::{LocalStage, Stage}; -use crate::{DistributedTaskContext, TaskKey}; +use crate::{DistributedTaskContext, QErrorMetric, TaskKey}; use datafusion::common::HashMap; use datafusion::common::plan_err; use datafusion::common::tree_node::Transformed; @@ -28,6 +28,8 @@ pub enum DistributedMetricsFormat { PerTask, } +pub const STATS_Q_ERROR_METRIC: &str = "stats_q_error"; + impl DistributedMetricsFormat { pub(crate) fn to_rewrite_ctx(self, task_id: u64) -> RewriteCtx { match self { @@ -78,9 +80,9 @@ pub async fn rewrite_distributed_plan_with_metrics( let network_boundary = network_boundary.with_input_stage(Stage::Local(LocalStage { query_id: stage.query_id, num: stage.num, - plan: plan_with_metrics, tasks: stage.tasks, - metrics_set: stage.metrics_set.clone(), + metrics_set: stage_metrics(stage.metrics_set.clone(), &plan_with_metrics), + plan: plan_with_metrics, }))?; let network_boundary = MetricsWrapperExec::new(network_boundary, plan.metrics().unwrap_or_default()); @@ -92,6 +94,51 @@ pub async fn rewrite_distributed_plan_with_metrics( plan.with_new_children(vec![transformed.data]) } +fn stage_metrics( + mut stage_metrics: MetricsSet, + plan_with_metrics: &Arc, +) -> MetricsSet { + let mut sampler_output_bytes = None; + let _ = plan_with_metrics.apply(|plan| { + if plan.is::() { + if let Some(output_bytes) = plan + .metrics() + .and_then(|metrics| metric_total(&metrics, "output_bytes")) + { + sampler_output_bytes = Some(output_bytes) + } + Ok(TreeNodeRecursion::Stop) + } else { + Ok(TreeNodeRecursion::Continue) + } + }); + + if let Some(estimated_byte_size) = metric_total(&stage_metrics, ESTIMATED_OUTPUT_BYTES_METRIC) + && let Some(sampler_output_bytes) = sampler_output_bytes + { + stage_metrics.push(QErrorMetric::new_metric( + STATS_Q_ERROR_METRIC, + q_error(estimated_byte_size, sampler_output_bytes), + )); + } + stage_metrics +} + +fn metric_total(metrics: &MetricsSet, name: &str) -> Option { + metrics + .sum(|metric| metric.value().name() == name) + .map(|value| value.as_usize()) +} + +/// Q-error is the standard cardinality-estimation metric because it treats equal-factor over- and +/// underestimates symmetrically. See https://www.vldb.org/pvldb/vol2/vldb09-657.pdf and +/// https://vldb.org/pvldb/vol9/p204-leis.pdf. +fn q_error(estimated: usize, actual: usize) -> f64 { + let estimated = estimated.max(1) as f64; + let actual = actual.max(1) as f64; + (estimated / actual).max(actual / estimated) +} + /// Extra information for rewriting local plans. #[derive(Default)] pub struct RewriteCtx { @@ -661,4 +708,12 @@ mod tests { assert_eq!(labels[1].name(), DISTRIBUTED_DATAFUSION_TASK_ID_LABEL); assert_eq!(labels[1].value(), "42"); } + + #[test] + fn q_error_is_symmetric_and_handles_zero() { + assert_eq!(super::q_error(200, 100), 2.0); + assert_eq!(super::q_error(100, 200), 2.0); + assert_eq!(super::q_error(0, 10), 10.0); + assert_eq!(super::q_error(0, 0), 1.0); + } } diff --git a/src/protocol/grpc/generated/worker.rs b/src/protocol/grpc/generated/worker.rs index d1f0fb4a9..56a143ff9 100644 --- a/src/protocol/grpc/generated/worker.rs +++ b/src/protocol/grpc/generated/worker.rs @@ -80,8 +80,8 @@ pub struct LoadInfo { /// The amount of rows that were pulled from leaf nodes while this partition was sampling data. #[prost(uint64, tag = "8")] pub rows_pulled_from_leaf: u64, - /// Whether the sampled partition stream reached end-of-stream by the time this LoadInfo was - /// captured. + /// Whether the sampled partition stream reached end-of-stream (i.e. the partition finished + /// producing all of its output) by the time this LoadInfo was captured. #[prost(bool, tag = "9")] pub reached_eos: bool, } @@ -259,13 +259,13 @@ pub struct Metric { pub partition: ::core::option::Option, #[prost( oneof = "metric::Value", - tags = "10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34" + tags = "10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35" )] pub value: ::core::option::Option, } /// Nested message and enum types in `Metric`. pub mod metric { - #[derive(Clone, PartialEq, Eq, Hash, ::prost::Oneof)] + #[derive(Clone, PartialEq, ::prost::Oneof)] pub enum Value { #[prost(message, tag = "10")] OutputRows(super::OutputRows), @@ -317,6 +317,8 @@ pub mod metric { CustomP99Latency(super::PercentileLatency), #[prost(message, tag = "34")] CustomMaxGauge(super::MaxGauge), + #[prost(message, tag = "35")] + CustomQError(super::QError), } } /// A MetricsSet is a protobuf mirror of datafusion::physical_plan::metrics::MetricsSet. It represents @@ -466,6 +468,13 @@ pub struct MaxGauge { #[prost(uint64, tag = "2")] pub value: u64, } +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct QError { + #[prost(string, tag = "1")] + pub name: ::prost::alloc::string::String, + #[prost(double, tag = "2")] + pub value: f64, +} /// Generated client implementations. pub mod worker_service_client { #![allow( diff --git a/src/protocol/grpc/metrics_proto.rs b/src/protocol/grpc/metrics_proto.rs index 6db8cda92..95a5f0477 100644 --- a/src/protocol/grpc/metrics_proto.rs +++ b/src/protocol/grpc/metrics_proto.rs @@ -12,6 +12,7 @@ use std::sync::Arc; use crate::{ AvgLatencyMetric, BytesCounterMetric, FirstLatencyMetric, MaxGaugeMetric, MaxLatencyMetric, MinLatencyMetric, P50LatencyMetric, P75LatencyMetric, P95LatencyMetric, P99LatencyMetric, + QErrorMetric, }; /// df_metrics_set_to_proto converts a [MetricsSet] to a [pb::MetricsSet]. @@ -247,6 +248,15 @@ pub fn df_metric_to_proto(metric: Arc) -> Result() { + Ok(pb::Metric { + value: Some(pb::metric::Value::CustomQError(pb::QError { + name: name.to_string(), + value: q_error.value(), + })), + partition, + labels, + }) } else { internal_err!("{}", CUSTOM_METRICS_NOT_SUPPORTED) } @@ -573,6 +583,17 @@ pub fn metric_proto_to_df(metric: pb::Metric) -> Result, DataFusionE labels, ))) } + Some(pb::metric::Value::CustomQError(q_error)) => { + let value = QErrorMetric::from_value(q_error.value); + Ok(Arc::new(Metric::new_with_labels( + MetricValue::Custom { + name: Cow::Owned(q_error.name), + value: Arc::new(value), + }, + partition, + labels, + ))) + } None => internal_err!("proto metric is missing the metric field"), } } @@ -1304,4 +1325,42 @@ mod tests { _ => panic!("expected Custom metrics"), } } + + #[test] + fn test_q_error_metric_roundtrip() { + let mut metrics_set = MetricsSet::new(); + metrics_set.push(Arc::new(Metric::new_with_labels( + MetricValue::Custom { + name: Cow::Borrowed("stats_q_error"), + value: Arc::new(QErrorMetric::from_value(std::f64::consts::PI)), + }, + Some(2), + vec![Label::new("stage", "4")], + ))); + + let proto = df_metrics_set_to_proto(&metrics_set).unwrap(); + let Some(pb::metric::Value::CustomQError(q_error)) = &proto.metrics[0].value else { + panic!("expected CustomQError metric"); + }; + assert_eq!(q_error.name, "stats_q_error"); + assert_eq!(q_error.value.to_bits(), std::f64::consts::PI.to_bits()); + + let roundtrip = metrics_set_proto_to_df(&proto).unwrap(); + let metric = roundtrip.iter().next().unwrap(); + assert_eq!(metric.partition(), Some(2)); + assert_eq!(metric.labels(), &[Label::new("stage", "4")]); + let MetricValue::Custom { name, value } = metric.value() else { + panic!("expected Custom metric"); + }; + assert_eq!(name, "stats_q_error"); + assert_eq!( + value + .as_any() + .downcast_ref::() + .unwrap() + .value() + .to_bits(), + std::f64::consts::PI.to_bits() + ); + } } diff --git a/src/protocol/grpc/worker.proto b/src/protocol/grpc/worker.proto index 24f5c9871..5a3d748dd 100644 --- a/src/protocol/grpc/worker.proto +++ b/src/protocol/grpc/worker.proto @@ -253,6 +253,7 @@ message Metric { PercentileLatency custom_p95_latency = 32; PercentileLatency custom_p99_latency = 33; MaxGauge custom_max_gauge = 34; + QError custom_q_error = 35; } } @@ -363,4 +364,9 @@ message PercentileLatency { message MaxGauge { string name = 1; uint64 value = 2; -} \ No newline at end of file +} + +message QError { + string name = 1; + double value = 2; +} diff --git a/src/stage.rs b/src/stage.rs index 7903a48a9..3e9d253e6 100644 --- a/src/stage.rs +++ b/src/stage.rs @@ -1047,6 +1047,7 @@ pub(crate) fn find_all_stages(plan: &Arc) -> Vec<&Stage> { #[cfg(test)] mod tests { use super::*; + use crate::QErrorMetric; use crate::test_utils::mock_exec::MockExec; use datafusion::arrow::datatypes::{DataType, Field, Schema}; use datafusion::physical_expr::expressions::Column; @@ -1109,6 +1110,12 @@ mod tests { assert_eq!(format_metrics_by_task(&set), "output_rows=100"); } + #[test] + fn format_metrics_by_task_displays_q_error_as_a_factor() { + let set = metrics_set([QErrorMetric::new_metric("stats_q_error", 1.456)]); + assert_eq!(format_metrics_by_task(&set), "stats_q_error=1.46x"); + } + fn single_column_schema() -> Arc { Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])) } diff --git a/src/worker/worker_connection_pool.rs b/src/worker/worker_connection_pool.rs index 9aa87e03d..058019c6f 100644 --- a/src/worker/worker_connection_pool.rs +++ b/src/worker/worker_connection_pool.rs @@ -1,3 +1,4 @@ +use crate::common::logical_record_batch_size; use crate::distributed_planner::ProducerHead; use crate::passthrough_headers::get_passthrough_headers; use crate::stage::RemoteStage; @@ -227,15 +228,6 @@ impl WorkerConnectionPool { } } -/// Returns the logical size of a batch's slices, excluding unused backing-buffer capacity. -fn logical_record_batch_size(batch: &RecordBatch) -> usize { - batch - .columns() - .iter() - .map(|column| column.to_data().get_slice_memory_size().unwrap_or(0)) - .sum() -} - impl Debug for WorkerConnectionPool { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { f.debug_struct("WorkerConnections") diff --git a/tests/metrics_collection.rs b/tests/metrics_collection.rs index b7e6a4e7f..b4e7eb538 100644 --- a/tests/metrics_collection.rs +++ b/tests/metrics_collection.rs @@ -6,6 +6,7 @@ mod tests { use datafusion::common::{Result, assert_contains}; use datafusion::execution::SessionState; use datafusion::physical_plan::display::DisplayableExecutionPlan; + use datafusion::physical_plan::metrics::MetricValue; use datafusion::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; use datafusion::physical_plan::{ExecutionPlan, execute_stream}; use datafusion::prelude::SessionContext; @@ -17,7 +18,8 @@ mod tests { }; use datafusion_distributed::{ DefaultSessionBuilder, DistributedExt, DistributedLeafExec, DistributedMetricsFormat, - NetworkCoalesceExec, NetworkShuffleExec, WorkerQueryContext, display_plan_ascii, + NetworkBoundaryExt, NetworkCoalesceExec, NetworkShuffleExec, QErrorMetric, + STATS_Q_ERROR_METRIC, Stage, WorkerQueryContext, display_plan_ascii, rewrite_distributed_plan_with_metrics, }; use futures::TryStreamExt; @@ -373,6 +375,49 @@ mod tests { Ok(()) } + #[tokio::test] + async fn test_dynamic_stats_q_error_is_collected_at_stage_level() + -> Result<(), Box> { + let (mut d_ctx, _guard, _) = start_localhost_context(3, DefaultSessionBuilder).await; + d_ctx.set_distributed_dynamic_task_count(true)?; + + let query = + r#"SELECT count(*), "RainToday" FROM weather GROUP BY "RainToday" ORDER BY count(*)"#; + + let s_ctx = SessionContext::default(); + let (_, d_physical) = execute(&s_ctx, &d_ctx, query).await?; + let d_physical = + rewrite_with_metrics(d_physical, DistributedMetricsFormat::Aggregated).await; + + let mut q_errors = vec![]; + d_physical.apply(|node| { + if let Some(boundary) = node.as_network_boundary() + && let Stage::Local(stage) = boundary.input_stage() + { + q_errors.extend(stage.metrics_set.iter().filter_map(|metric| { + let MetricValue::Custom { name, value } = metric.value() else { + return None; + }; + (name == STATS_Q_ERROR_METRIC) + .then(|| value.as_any().downcast_ref::()) + .flatten() + .map(QErrorMetric::value) + })); + } + Ok(TreeNodeRecursion::Continue) + })?; + + assert_eq!(q_errors.len(), 1); + assert!(q_errors[0].is_finite()); + assert!(q_errors[0] >= 1.0); + assert_contains!( + display_plan_ascii(d_physical.as_ref(), true), + "stats_q_error=" + ); + + Ok(()) + } + /// Looks for an [ExecutionPlan] that matches the provided type parameter `T1` in /// the left node and `T2` in the right node and compares its metrics. /// There might be more than one, so `index` determines which one is compared.