Add stats-q-error to stage level metrics

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<dyn ExecutionPlan>) -> Option<StatsEstimationQError> {
     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<StatsEstimationQError> {
     })
 }
 
-fn metric_total(metrics: &MetricsSet, name: &str) -> Option<usize> {
-    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<f64> {
+    metrics.iter().find_map(|metric| match metric.value() {
+        MetricValue::Custom { name, value } if name == STATS_Q_ERROR_METRIC => value
+            .as_any()
+            .downcast_ref::<QErrorMetric>()
+            .map(QErrorMetric::value),
+        _ => None,
+    })
 }
 
 pub fn median(mut values: Vec<f64>) -> Option<f64> {
@@ -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));
+    }
 }

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;

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()
+}

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;

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<dyn ExecutionPlan>,
@@ -294,7 +297,7 @@ async fn gather_runtime_statistics(
     };
 
     new_metrics.push(MaxGaugeMetric::new_metric(
-        "estimated_pct_sampled",
+        ESTIMATED_PCT_SAMPLED_METRIC,
         (estimated_pct_sampled * 100.) as usize,
     ));
 
@@ -308,7 +311,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,
     ));
 

src/execution_plans/sampler.rs

@@ -1,25 +1,26 @@
-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,
 };
 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};
 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()

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))]

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";

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<AtomicU64>,
+}
+
+impl Default for QErrorMetric {
+    fn default() -> Self {
+        Self::from_value(1.0)
+    }
+}
+
+impl QErrorMetric {
+    pub fn new_metric(name: impl Into<Cow<'static, str>>, value: f64) -> Arc<Metric> {
+        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<dyn CustomMetricValue> {
+        Arc::new(Self::default())
+    }
+
+    fn aggregate(&self, other: Arc<dyn CustomMetricValue + 'static>) {
+        let Some(other) = other.as_any().downcast_ref::<Self>() 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<dyn CustomMetricValue>) -> bool {
+        other
+            .as_any()
+            .downcast_ref::<Self>()
+            .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);
+    }
+}