diff --git a/datafusion/core/Cargo.toml b/datafusion/core/Cargo.toml index 222c0ec688b78..7275c4b674d20 100644 --- a/datafusion/core/Cargo.toml +++ b/datafusion/core/Cargo.toml @@ -340,3 +340,7 @@ required-features = ["parquet"] [[bench]] harness = false name = "reset_plan_states" + +[[bench]] +harness = false +name = "partition_metrics" diff --git a/datafusion/core/benches/partition_metrics.rs b/datafusion/core/benches/partition_metrics.rs new file mode 100644 index 0000000000000..ceef897cd196f --- /dev/null +++ b/datafusion/core/benches/partition_metrics.rs @@ -0,0 +1,190 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Execute a SQL query over a partitioned table, then report operator metrics. +//! Run with `cargo bench -p datafusion --bench partition_metrics`. + +use std::hint::black_box; +use std::sync::Arc; +use std::time::Duration; + +use arrow::array::Int32Array; +use arrow::datatypes::{DataType, Field, Schema}; +use arrow::record_batch::RecordBatch; +use criterion::{BatchSize, BenchmarkId, Criterion, criterion_group, criterion_main}; +use datafusion::datasource::MemTable; +use datafusion::prelude::{SessionConfig, SessionContext}; +use datafusion_common::Result; +use datafusion_physical_plan::ExecutionPlan; +use datafusion_physical_plan::metrics::MetricsSet; +use futures::TryStreamExt; + +async fn query(partitions: usize) -> Result<(SessionContext, Arc)> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from_iter_values(0..32))], + )?; + let table = MemTable::try_new(schema, vec![vec![batch]; partitions])?; + let ctx = SessionContext::new_with_config( + SessionConfig::new() + .with_target_partitions(partitions) + .with_batch_size(16), + ); + ctx.register_table("t", Arc::new(table))?; + let plan = ctx + .sql("SELECT a + 1 AS b FROM t WHERE a < 16") + .await? + .create_physical_plan() + .await?; + assert_eq!(plan.properties().partitioning.partition_count(), partitions); + Ok((ctx, plan)) +} + +fn nodes(plan: &Arc) -> Vec> { + let mut result = vec![Arc::clone(plan)]; + for child in plan.children() { + result.extend(nodes(child)); + } + result +} + +fn select(metrics: &MetricsSet, partition: usize, indexed: bool) -> usize { + if indexed { + metrics + .for_partition(partition) + .iter() + .map(|m| { + black_box(m.value().as_usize()); + 1 + }) + .sum() + } else { + metrics + .iter() + .filter(|m| m.partition() == Some(partition)) + .map(|m| { + black_box(m.value().as_usize()); + 1 + }) + .sum() + } +} + +async fn execute(ctx: &SessionContext, plan: &Arc, partition: usize) { + let batches: Vec<_> = plan + .execute(partition, ctx.task_ctx()) + .unwrap() + .try_collect() + .await + .unwrap(); + assert_eq!(batches.iter().map(|b| b.num_rows()).sum::(), 16); +} + +fn benchmarks(c: &mut Criterion) { + let runtime = tokio::runtime::Runtime::new().unwrap(); + let mut retrieval = c.benchmark_group("partition_metrics/sql_retrieval"); + retrieval + .sample_size(20) + .warm_up_time(Duration::from_secs(1)) + .measurement_time(Duration::from_secs(3)); + let mut expected_count = None; + for partitions in [1, 64, 1024, 8192] { + let (ctx, plan) = runtime.block_on(query(partitions)).unwrap(); + runtime.block_on(async { + for p in 0..partitions { + execute(&ctx, &plan, p).await; + } + }); + let nodes = nodes(&plan); + let count: usize = nodes + .iter() + .filter_map(|n| n.metrics()) + .map(|m| m.for_partition(0).iter().count()) + .sum(); + assert!(count > 0); + assert_eq!(*expected_count.get_or_insert(count), count); + for (label, indexed) in [("full_then_filter", false), ("partition", true)] { + retrieval.bench_with_input( + BenchmarkId::new(label, partitions), + &indexed, + |b, &indexed| { + b.iter(|| { + for node in &nodes { + if let Some(metrics) = node.metrics() { + black_box(select(&metrics, 0, indexed)); + } + } + }); + }, + ); + } + } + retrieval.finish(); + + let mut execution = c.benchmark_group("partition_metrics/sql_execution"); + execution + .sample_size(10) + .warm_up_time(Duration::from_secs(1)) + .measurement_time(Duration::from_secs(3)); + for partitions in [64, 1024] { + for (label, report) in [ + ("execute_only", None), + ("full_then_filter", Some(false)), + ("partition", Some(true)), + ] { + execution.bench_with_input( + BenchmarkId::new(label, partitions), + &report, + |b, &report| { + b.iter_batched( + || runtime.block_on(query(partitions)).unwrap(), + |(ctx, plan)| { + let nodes = nodes(&plan); + runtime.block_on(async { + let mut previous_snapshots = Vec::new(); + for partition in 0..partitions { + // Keep the previous snapshots alive while executing and + // registering the next partition's metrics. + execute(&ctx, &plan, partition).await; + if let Some(indexed) = report { + let snapshots: Vec<_> = nodes + .iter() + .filter_map(|n| n.metrics()) + .collect(); + for metrics in &snapshots { + black_box(select( + metrics, partition, indexed, + )); + } + black_box(&previous_snapshots); + previous_snapshots = snapshots; + } + } + }); + }, + BatchSize::PerIteration, + ); + }, + ); + } + } + execution.finish(); +} + +criterion_group!(benches, benchmarks); +criterion_main!(benches); diff --git a/datafusion/core/src/dataframe/parquet.rs b/datafusion/core/src/dataframe/parquet.rs index 1685dff23dff1..a9f46e31ec4d9 100644 --- a/datafusion/core/src/dataframe/parquet.rs +++ b/datafusion/core/src/dataframe/parquet.rs @@ -490,6 +490,21 @@ mod tests { let metrics = plan .metrics() .expect("DataSinkExec should return metrics from ParquetSink"); + let selected = plan.metrics().unwrap().for_partition(0); + let expected: Vec<_> = metrics + .iter() + .filter(|metric| metric.partition() == Some(0)) + .collect(); + assert_eq!(selected.iter().count(), expected.len()); + for (actual, expected) in selected.iter().zip(expected) { + assert!(Arc::ptr_eq(actual, expected)); + } + // Sink-wide row counts are intentionally absent from a partition snapshot. + assert!( + selected + .iter() + .all(|metric| metric.value().name() != "rows_written") + ); let aggregated = metrics.aggregate_by_name(); // rows_written should be 100 diff --git a/datafusion/datasource-parquet/src/source.rs b/datafusion/datasource-parquet/src/source.rs index 3df0e0e7e14f9..6ee9918640bb0 100644 --- a/datafusion/datasource-parquet/src/source.rs +++ b/datafusion/datasource-parquet/src/source.rs @@ -1358,6 +1358,42 @@ mod tests { use arrow::datatypes::Schema; use datafusion_physical_expr::expressions::lit; + #[test] + fn partition_metrics_exclude_derived_plan_metrics() { + use datafusion_datasource::file_scan_config::FileScanConfigBuilder; + use datafusion_datasource::source::DataSourceExec; + use datafusion_execution::object_store::ObjectStoreUrl; + use datafusion_physical_plan::ExecutionPlan; + use datafusion_physical_plan::metrics::MetricBuilder; + + let source = Arc::new(ParquetSource::new(Arc::new(Schema::empty()))); + let metrics = source.metrics().clone(); + let config = + FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), source) + .build(); + let plan: Arc = DataSourceExec::from_data_source(config); + assert_eq!(plan.metrics().unwrap().for_partition(0).iter().count(), 0); + MetricBuilder::new(&metrics).output_rows(0).add(10); + MetricBuilder::new(&metrics).output_rows(1).add(20); + MetricBuilder::new(&metrics).global_counter("global").add(1); + let selected = plan.metrics().unwrap().for_partition(0); + assert_eq!(selected.output_rows(), Some(10)); + assert!(selected.iter().all(|m| m.partition() == Some(0))); + let full = plan.metrics().unwrap(); + assert_eq!(full.output_rows(), Some(30)); + assert!( + full.iter().any( + |m| m.value().name() == "output_rows_skew" && m.partition().is_none() + ) + ); + MetricBuilder::new(&metrics).output_rows(0).add(5); + assert_eq!(selected.output_rows(), Some(10)); + assert_eq!( + plan.metrics().unwrap().for_partition(0).output_rows(), + Some(15) + ); + } + #[test] fn test_reverse_scan_default_value() { use arrow::datatypes::Schema; diff --git a/datafusion/ffi/src/execution_plan.rs b/datafusion/ffi/src/execution_plan.rs index 056d29d16aecf..ffba61a072658 100644 --- a/datafusion/ffi/src/execution_plan.rs +++ b/datafusion/ffi/src/execution_plan.rs @@ -866,6 +866,9 @@ pub mod tests { let observed = metric_foreign.metrics().expect("metrics should be present"); assert_eq!(observed.output_rows(), Some(42)); + assert_eq!(observed.for_partition(0).output_rows(), Some(11)); + assert_eq!(observed.for_partition(1).output_rows(), Some(31)); + assert_eq!(observed.for_partition(2).iter().count(), 0); Ok(()) } diff --git a/datafusion/ffi/tests/ffi_execution_plan.rs b/datafusion/ffi/tests/ffi_execution_plan.rs index 28828695d01af..6cad71fa13b07 100644 --- a/datafusion/ffi/tests/ffi_execution_plan.rs +++ b/datafusion/ffi/tests/ffi_execution_plan.rs @@ -83,6 +83,17 @@ mod tests { // across that boundary - not just the in-process From conversions // covered by physical_expr::metrics's roundtrip tests. let metrics = plan.metrics().expect("plan should report metrics"); + let selected = metrics.for_partition(0); + let expected: Vec<_> = metrics + .iter() + .filter(|m| m.partition() == Some(0)) + .collect(); + assert!(!expected.is_empty()); + assert_eq!(selected.iter().count(), expected.len()); + for (actual, expected) in selected.iter().zip(expected) { + assert!(Arc::ptr_eq(actual, expected)); + } + assert_eq!(metrics.for_partition(usize::MAX).iter().count(), 0); // Assert the transported Bytes category, the generic variant/name/ // value, and (below) the byte-formatted display output - MetricValue diff --git a/datafusion/physical-expr-common/src/metrics/mod.rs b/datafusion/physical-expr-common/src/metrics/mod.rs index c00fcff70514b..f47cb2c5434f0 100644 --- a/datafusion/physical-expr-common/src/metrics/mod.rs +++ b/datafusion/physical-expr-common/src/metrics/mod.rs @@ -22,12 +22,14 @@ mod builder; mod custom; mod elapsed_compute; mod expression; +mod snapshot; mod value; use datafusion_common::HashMap; pub use datafusion_common::format::{MetricCategory, MetricType}; use datafusion_common::human_readable_size; use parking_lot::Mutex; +use snapshot::{Registry, Snapshot}; use std::{ borrow::Cow, fmt::{self, Debug, Display}, @@ -228,9 +230,18 @@ impl Metric { } /// A snapshot of the metrics for a particular execution plan. +/// +/// Snapshots returned by [`ExecutionPlanMetricsSet::clone_inner`] record a fixed +/// registration boundary and defer copying metric handles until iteration. Their +/// membership is fixed while metric values remain live. Cloning these snapshots +/// is O(1); cloning an owned set copies its handles. +/// +/// A registry-backed snapshot retains its source registry, including later +/// registrations. Full iteration caches its original members in O(n) time; +/// [`Self::for_partition`] can instead use the registry's incremental index. #[derive(Default, Debug, Clone)] pub struct MetricsSet { - metrics: Vec>, + metrics: Snapshot, } impl MetricsSet { @@ -239,11 +250,35 @@ impl MetricsSet { Default::default() } - /// Add the specified metric + /// Add the specified metric. + /// + /// Mutating a deferred snapshot first copies its members into an independent + /// vector. Subsequent additions append to that vector. pub fn push(&mut self, metric: Arc) { self.metrics.push(metric) } + /// Return a snapshot containing only metrics with `partition == Some(partition)`. + /// + /// For registry-backed snapshots, this incrementally indexes registrations + /// not processed by earlier partition reads, then clones only matching handles. + /// Each registration is indexed at most once across snapshots of the registry. + /// Owned sets instead filter their handles in O(n) time. + /// + /// Unpartitioned metrics are excluded; an unknown partition yields an empty set. + /// Registration order, duplicates and shared metric values are preserved. The + /// result owns its handles and does not retain the source registry. It never + /// acquires metrics registered after this snapshot was created. + /// + /// Index updates and selection hold the registry lock and can delay concurrent + /// registration. This does not avoid work already performed by the provider, + /// such as full metrics transport over FFI. + pub fn for_partition(&self, partition: usize) -> Self { + Self { + metrics: self.metrics.for_partition(partition), + } + } + /// Returns an iterator across all metrics pub fn iter(&self) -> impl Iterator> { self.metrics.iter() @@ -360,10 +395,7 @@ impl MetricsSet { }); } - let new_metrics = map - .into_iter() - .map(|(_k, v)| Arc::new(v)) - .collect::>(); + let new_metrics = map.into_iter().map(|(_k, v)| Arc::new(v)).collect(); Self { metrics: new_metrics, @@ -371,14 +403,15 @@ impl MetricsSet { } /// Sort the order of metrics so the "most useful" show up first - pub fn sorted_for_display(mut self) -> Self { - self.metrics.sort_unstable_by_key(|metric| { + pub fn sorted_for_display(self) -> Self { + let mut metrics: Vec<_> = self.into_iter().collect(); + metrics.sort_unstable_by_key(|metric| { ( metric.value().display_sort_key(), metric.value().name().to_owned(), ) }); - self + metrics.into_iter().collect() } /// Remove all timestamp metrics (for more compact display) @@ -388,7 +421,7 @@ impl MetricsSet { let metrics = metrics .into_iter() .filter(|m| !m.value.is_timestamp()) - .collect::>(); + .collect(); Self { metrics } } @@ -397,14 +430,14 @@ impl MetricsSet { /// [`MetricType`] appears in `allowed`. pub fn filter_by_metric_types(self, allowed: &[MetricType]) -> Self { if allowed.is_empty() { - return Self { metrics: vec![] }; + return Self::new(); } let metrics = self .metrics .into_iter() .filter(|metric| allowed.contains(&metric.metric_type())) - .collect::>(); + .collect(); Self { metrics } } @@ -418,7 +451,7 @@ impl MetricsSet { /// removed. pub fn filter_by_categories(self, allowed: &[MetricCategory]) -> Self { if allowed.is_empty() { - return Self { metrics: vec![] }; + return Self::new(); } let metrics = self @@ -430,7 +463,7 @@ impl MetricsSet { .unwrap_or(MetricCategory::Uncategorized); allowed.contains(&cat) }) - .collect::>(); + .collect(); Self { metrics } } @@ -438,14 +471,14 @@ impl MetricsSet { /// 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![] }; + return Self::new(); } let metrics = self .metrics .into_iter() .filter(|metric| names.iter().any(|name| name == metric.value().name())) - .collect::>(); + .collect(); Self { metrics } } } @@ -509,33 +542,37 @@ impl FromIterator> for MetricsSet { /// underlying metrics set #[derive(Default, Debug, Clone)] pub struct ExecutionPlanMetricsSet { - inner: Arc>, + inner: Arc>, } impl ExecutionPlanMetricsSet { /// Create a new empty shared metrics set pub fn new() -> Self { - Self { - inner: Arc::new(Mutex::new(MetricsSet::new())), - } + Self::default() } /// Add the specified metric to the underlying metric set pub fn register(&self, metric: Arc) { - self.inner.lock().push(metric) + self.inner.lock().metrics.push(metric) } - /// Return a clone of the inner [`MetricsSet`] + /// Return a snapshot with the current registration boundary in O(1). + /// + /// Metric handles are copied on iteration or partition selection. Later + /// registrations leave this snapshot's membership unchanged, while values + /// remain shared. Retaining a snapshot also retains the source registry. pub fn clone_inner(&self) -> MetricsSet { - let guard = self.inner.lock(); - (*guard).clone() + let end = self.inner.lock().metrics.len(); + MetricsSet { + metrics: Snapshot::new(Arc::clone(&self.inner), end), + } } } impl From for ExecutionPlanMetricsSet { fn from(metrics: MetricsSet) -> Self { Self { - inner: Arc::new(Mutex::new(metrics)), + inner: Arc::new(Mutex::new(Registry::new(metrics.into_iter().collect()))), } } } @@ -737,6 +774,134 @@ mod tests { assert_eq!(borrowed.to_string(), shared.to_string()); } + #[test] + fn selected_snapshot_mutations_are_independent() { + let registry = ExecutionPlanMetricsSet::new(); + MetricBuilder::new(®istry).output_rows(0).add(3); + MetricBuilder::new(®istry).output_rows(1).add(7); + MetricBuilder::new(®istry) + .global_counter("global") + .add(11); + let full = registry.clone_inner(); + let selected = full.for_partition(0); + assert_eq!(selected.for_partition(0).output_rows(), Some(3)); + assert_eq!(selected.for_partition(1).iter().count(), 0); + let mut modified = selected.clone(); + let count = Count::new(); + count.add(13); + modified.push(Arc::new(Metric::new( + MetricValue::OutputRows(count), + Some(1), + ))); + assert_eq!(modified.output_rows(), Some(16)); + assert_eq!(modified.for_partition(1).output_rows(), Some(13)); + assert_eq!(selected.output_rows(), Some(3)); + assert_eq!(full.output_rows(), Some(10)); + assert_eq!(registry.clone_inner().iter().count(), 3); + assert_eq!(modified.clone().into_iter().count(), 2); + assert_eq!( + modified.sorted_for_display().for_partition(0).output_rows(), + Some(3) + ); + } + + #[test] + fn partition_snapshots_preserve_registration_and_shared_values() { + let metrics = ExecutionPlanMetricsSet::new(); + assert_eq!(metrics.clone_inner().for_partition(0).iter().count(), 0); + let first = MetricBuilder::new(&metrics).output_rows(0); + first.add(11); + MetricBuilder::new(&metrics).global_counter("global").add(7); + MetricBuilder::new(&metrics).output_rows(usize::MAX).add(99); + // Same name and partition must not overwrite the earlier metric. + MetricBuilder::new(&metrics).output_rows(0).add(13); + let snapshot = metrics.clone_inner().for_partition(0); + assert_eq!(snapshot.output_rows(), Some(24)); + assert_eq!(snapshot.iter().count(), 2); + assert_eq!(snapshot.aggregate_by_name().output_rows(), Some(24)); + assert_eq!( + metrics + .clone_inner() + .for_partition(usize::MAX) + .output_rows(), + Some(99) + ); + assert_eq!(metrics.clone_inner().for_partition(1).iter().count(), 0); + + let shared = metrics.clone(); + first.add(1); + MetricBuilder::new(&shared).output_rows(0).add(17); + assert_eq!(snapshot.output_rows(), Some(25)); + assert_eq!( + shared.clone_inner().for_partition(0).output_rows(), + Some(42) + ); + assert_eq!( + metrics.clone_inner().for_partition(0).output_rows(), + Some(42) + ); + + let full = metrics.clone_inner(); + let imported = ExecutionPlanMetricsSet::from(full.clone()); + for (original, copied) in full.iter().zip(imported.clone_inner().iter()) { + assert!(Arc::ptr_eq(original, copied)); + } + for partition in [0, 1, usize::MAX] { + let expected: Vec<_> = full + .iter() + .filter(|metric| metric.partition() == Some(partition)) + .collect(); + let selected = imported.clone_inner().for_partition(partition); + assert_eq!(expected.len(), selected.iter().count()); + for (original, copied) in expected.into_iter().zip(selected.iter()) { + assert!(Arc::ptr_eq(original, copied)); + } + } + // From shares metric values, but creates an independent registration set. + MetricBuilder::new(&imported).output_rows(0).add(3); + assert_eq!( + imported.clone_inner().for_partition(0).output_rows(), + Some(45) + ); + assert_eq!( + metrics.clone_inner().for_partition(0).output_rows(), + Some(42) + ); + assert_eq!(full.iter().filter(|m| m.partition().is_none()).count(), 1); + } + + #[test] + fn partition_snapshots_during_registration() { + let metrics = ExecutionPlanMetricsSet::new(); + let barrier = std::sync::Barrier::new(5); + std::thread::scope(|scope| { + for partition in 0..4 { + let metrics = &metrics; + let barrier = &barrier; + scope.spawn(move || { + barrier.wait(); + for _ in 0..1000 { + MetricBuilder::new(metrics).output_rows(partition).add(1); + } + }); + } + barrier.wait(); + for _ in 0..1000 { + for partition in 0..4 { + let selected = metrics.clone_inner().for_partition(partition); + assert!(selected.iter().all(|m| m.partition() == Some(partition))); + assert!(selected.iter().count() <= 1000); + } + } + }); + for partition in 0..4 { + let selected = metrics.clone_inner().for_partition(partition); + assert_eq!(selected.iter().count(), 1000); + assert_eq!(selected.output_rows(), Some(1000)); + } + assert_eq!(metrics.clone_inner().output_rows(), Some(4000)); + } + #[test] fn test_output_rows() { let metrics = ExecutionPlanMetricsSet::new(); diff --git a/datafusion/physical-expr-common/src/metrics/snapshot.rs b/datafusion/physical-expr-common/src/metrics/snapshot.rs new file mode 100644 index 0000000000000..f57979394e68f --- /dev/null +++ b/datafusion/physical-expr-common/src/metrics/snapshot.rs @@ -0,0 +1,301 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Fixed-membership snapshots of an append-only registry. +//! +//! Registration only appends to a vector. Partition readers share an index that +//! catches up with registration on demand; snapshots remember their original end +//! position even when a later reader has advanced the index past that position. + +use super::Metric; +use parking_lot::Mutex; +use std::collections::HashMap; +use std::fmt; +use std::sync::{Arc, OnceLock}; + +#[derive(Debug, Default)] +pub(super) struct Registry { + pub(super) metrics: Vec>, + // No index allocation or maintenance on the registration path. + index: Option>, +} + +#[derive(Debug, Default)] +struct PartitionIndex { + // Number of registry entries already examined, including global metrics. + indexed: usize, + // Partition ID -> positions in Registry::metrics, in registration order. + positions: HashMap>, +} + +impl Registry { + pub(super) fn new(metrics: Vec>) -> Self { + Self { + metrics, + index: None, + } + } + + fn select(&mut self, partition: usize, end: usize) -> Vec> { + let index = self.index.get_or_insert_with(Default::default); + for position in index.indexed..end { + if let Some(partition) = self.metrics[position].partition() { + index.positions.entry(partition).or_default().push(position); + } + } + index.indexed = index.indexed.max(end); + let Some(positions) = index.positions.get(&partition) else { + return Vec::new(); + }; + // Another snapshot may already have indexed registrations after our end. + let len = positions.partition_point(|&position| position < end); + positions[..len] + .iter() + .map(|&i| Arc::clone(&self.metrics[i])) + .collect() + } +} + +#[derive(Clone)] +pub(super) enum Snapshot { + Owned(Vec>), + Deferred(Arc), +} + +pub(super) struct Deferred { + registry: Arc>, + end: usize, + flat: OnceLock>>, +} + +impl Default for Snapshot { + fn default() -> Self { + Self::Owned(Vec::new()) + } +} + +impl fmt::Debug for Snapshot { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_list().entries(self.iter()).finish() + } +} + +impl Deferred { + fn materialize(&self) -> Vec> { + self.registry.lock().metrics[..self.end].to_vec() + } +} + +impl Snapshot { + pub(super) fn new(registry: Arc>, end: usize) -> Self { + Self::Deferred(Arc::new(Deferred { + registry, + end, + flat: OnceLock::new(), + })) + } + + pub(super) fn for_partition(&self, partition: usize) -> Self { + Self::Owned(match self { + Self::Owned(metrics) => metrics + .iter() + .filter(|m| m.partition() == Some(partition)) + .cloned() + .collect(), + Self::Deferred(snapshot) => { + snapshot.registry.lock().select(partition, snapshot.end) + } + }) + } + + pub(super) fn push(&mut self, metric: Arc) { + if let Self::Deferred(_) = self { + *self = Self::Owned(std::mem::take(self).into_iter().collect()); + } + let Self::Owned(metrics) = self else { + unreachable!() + }; + metrics.push(metric); + } + + pub(super) fn iter(&self) -> std::slice::Iter<'_, Arc> { + match self { + Self::Owned(metrics) => metrics.iter(), + Self::Deferred(snapshot) => { + snapshot.flat.get_or_init(|| snapshot.materialize()).iter() + } + } + } +} + +impl IntoIterator for Snapshot { + type Item = Arc; + type IntoIter = std::vec::IntoIter; + + fn into_iter(self) -> Self::IntoIter { + let metrics = match self { + Self::Owned(metrics) => metrics, + Self::Deferred(snapshot) => match Arc::try_unwrap(snapshot) { + Ok(mut snapshot) => snapshot + .flat + .take() + .unwrap_or_else(|| snapshot.materialize()), + Err(snapshot) => { + snapshot.flat.get_or_init(|| snapshot.materialize()).clone() + } + }, + }; + metrics.into_iter() + } +} + +impl<'a> IntoIterator for &'a Snapshot { + type Item = &'a Arc; + type IntoIter = std::slice::Iter<'a, Arc>; + + fn into_iter(self) -> Self::IntoIter { + self.iter() + } +} + +impl Extend> for Snapshot { + fn extend>>(&mut self, iter: I) { + for metric in iter { + self.push(metric); + } + } +} + +impl FromIterator> for Snapshot { + fn from_iter>>(iter: I) -> Self { + Self::Owned(iter.into_iter().collect()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::metrics::{Count, ExecutionPlanMetricsSet, MetricValue}; + + fn metric(partition: Option) -> Arc { + Arc::new(Metric::new( + MetricValue::OutputRows(Count::new()), + partition, + )) + } + + #[test] + fn out_of_order_snapshots_preserve_membership_and_registration_order() { + let registry = ExecutionPlanMetricsSet::new(); + let mut expected = Vec::new(); + let mut snapshots = Vec::new(); + for i in 0..1024 { + let partition = match i % 5 { + 0 => None, + 1 => Some(usize::MAX), + 2 => Some(0), + _ => Some(i * 17), + }; + let metric = metric(partition); + expected.push(Arc::clone(&metric)); + registry.register(metric); + if i % 31 == 0 { + snapshots.push((registry.clone_inner(), expected.len())); + } + } + // Reading full snapshots must not build the partition index. + assert_eq!(registry.clone_inner().iter().count(), 1024); + assert!(registry.inner.lock().index.is_none()); + // Advance the shared index before reading older, unmaterialized snapshots. + registry.clone_inner().for_partition(0); + for (snapshot, len) in snapshots.into_iter().rev() { + for partition in [0, 17, 51, 999999, usize::MAX] { + let selected = snapshot.for_partition(partition); + let expected: Vec<_> = expected[..len] + .iter() + .filter(|m| m.partition() == Some(partition)) + .collect(); + assert_eq!(selected.iter().count(), expected.len()); + for (actual, expected) in selected.iter().zip(expected) { + assert!(Arc::ptr_eq(actual, expected)); + } + } + assert_eq!(snapshot.iter().count(), len); + for (actual, expected) in snapshot.iter().zip(&expected[..len]) { + assert!(Arc::ptr_eq(actual, expected)); + } + } + } + + #[test] + fn registration_leaves_index_and_retained_snapshots_unchanged() { + let registry = ExecutionPlanMetricsSet::new(); + registry.register(metric(Some(0))); + let empty_partition = registry.clone_inner(); + let old = registry.clone_inner(); + assert_eq!(old.for_partition(0).iter().count(), 1); + registry.register(metric(Some(0))); + registry.register(metric(Some(1))); + { + let registry = registry.inner.lock(); + let index = registry.index.as_ref().unwrap(); + assert_eq!(index.indexed, 1); + assert_eq!(index.positions[&0], vec![0]); + assert!(!index.positions.contains_key(&1)); + } + assert_eq!(registry.clone_inner().for_partition(0).iter().count(), 2); + assert_eq!(registry.clone_inner().for_partition(1).iter().count(), 1); + assert_eq!(empty_partition.for_partition(1).iter().count(), 0); + assert_eq!(old.for_partition(0).iter().count(), 1); + assert_eq!(old.iter().count(), 1); + let registry = registry.inner.lock(); + let index = registry.index.as_ref().unwrap(); + assert_eq!(index.indexed, 3); + assert_eq!(index.positions[&0], vec![0, 1]); + } + + #[test] + fn selected_result_does_not_retain_registry() { + let registry = ExecutionPlanMetricsSet::new(); + registry.register(metric(Some(0))); + registry.register(metric(Some(1))); + let weak = Arc::downgrade(®istry.inner); + let snapshot = registry.clone_inner(); + let selected = snapshot.for_partition(0); + drop(registry); + assert!(weak.upgrade().is_some()); + drop(snapshot); + assert!(weak.upgrade().is_none()); + assert_eq!(selected.iter().count(), 1); + } + + #[test] + fn mutating_full_snapshot_detaches_from_registry() { + let registry = ExecutionPlanMetricsSet::new(); + registry.register(metric(Some(0))); + let original = registry.clone_inner(); + let mut modified = original.clone(); + modified.push(metric(Some(1))); + registry.register(metric(Some(2))); + assert_eq!(original.iter().count(), 1); + assert_eq!(modified.iter().count(), 2); + assert_eq!(modified.for_partition(1).iter().count(), 1); + assert_eq!(modified.for_partition(2).iter().count(), 0); + assert_eq!(registry.clone_inner().for_partition(1).iter().count(), 0); + assert_eq!(registry.clone_inner().for_partition(2).iter().count(), 1); + } +} diff --git a/datafusion/physical-plan/src/execution_plan.rs b/datafusion/physical-plan/src/execution_plan.rs index 978059b9faf6b..2f5680e659e0d 100644 --- a/datafusion/physical-plan/src/execution_plan.rs +++ b/datafusion/physical-plan/src/execution_plan.rs @@ -706,10 +706,9 @@ pub trait ExecutionPlan: Any + Debug + DisplayAs + Send + Sync { /// [`MetricsSet`]s may change as execution progresses, the /// specific metrics will not. /// - /// Once `self.execute()` has returned (technically the future is - /// resolved) for all available partitions, the set of metrics - /// should be complete. If this function is called prior to - /// `execute()` new metrics may appear in subsequent calls. + /// New metrics may be registered during execution, including while streams + /// returned by [`Self::execute`] are being polled. Call again to obtain + /// metrics registered since the previous snapshot. fn metrics(&self) -> Option { None } diff --git a/datafusion/physical-plan/tests/metrics/plan.rs b/datafusion/physical-plan/tests/metrics/plan.rs new file mode 100644 index 0000000000000..d09c8d981f181 --- /dev/null +++ b/datafusion/physical-plan/tests/metrics/plan.rs @@ -0,0 +1,118 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Shared fixture: real streaming, filter and projection operators. + +use std::sync::Arc; + +use arrow::array::Int32Array; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use arrow::record_batch::RecordBatch; +use datafusion_common::{Result, ScalarValue}; +use datafusion_execution::{SendableRecordBatchStream, TaskContext}; +use datafusion_expr::Operator; +use datafusion_physical_expr::PhysicalExpr; +use datafusion_physical_expr::expressions::{BinaryExpr, Column, Literal}; +use datafusion_physical_plan::ExecutionPlan; +use datafusion_physical_plan::filter::FilterExec; +use datafusion_physical_plan::projection::ProjectionExec; +use datafusion_physical_plan::stream::RecordBatchStreamAdapter; +use datafusion_physical_plan::streaming::{PartitionStream, StreamingTableExec}; +use futures::StreamExt; +use tokio::sync::Notify; + +#[derive(Debug)] +struct InputPartition { + schema: SchemaRef, + batch: RecordBatch, + gate: Option>, +} + +impl PartitionStream for InputPartition { + fn schema(&self) -> &SchemaRef { + &self.schema + } + + fn execute(&self, _: Arc) -> SendableRecordBatchStream { + let first = self.batch.clone(); + let second = self.batch.clone(); + let gate = self.gate.clone(); + let stream = futures::stream::once(async move { Ok(first) }).chain( + futures::stream::once(async move { + if let Some(gate) = gate { + gate.notified().await; + } + Ok(second) + }), + ); + Box::pin(RecordBatchStreamAdapter::new( + Arc::clone(&self.schema), + stream, + )) + } +} + +pub type SharedPlan = (Vec>, Arc); + +/// Partition 0 emits one batch, then waits while other partitions finish. +/// Returns nodes in root-to-leaf order for per-node reporting. +pub fn shared_plan(partitions: usize) -> Result { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from_iter_values(0..32))], + )?; + let gate = Arc::new(Notify::new()); + let inputs = (0..partitions) + .map(|partition| { + Arc::new(InputPartition { + schema: Arc::clone(&schema), + batch: batch.clone(), + gate: (partition == 0).then(|| Arc::clone(&gate)), + }) as Arc + }) + .collect(); + // The limit enables StreamingTableExec's own baseline metrics without + // truncating this finite input. + let source: Arc = Arc::new(StreamingTableExec::try_new( + schema, + inputs, + None, + [], + false, + Some(usize::MAX), + )?); + let column: Arc = Arc::new(Column::new("a", 0)); + let predicate = Arc::new(BinaryExpr::new( + Arc::clone(&column), + Operator::Lt, + Arc::new(Literal::new(ScalarValue::Int32(Some(16)))), + )); + let filter: Arc = Arc::new( + FilterExec::try_new(predicate, Arc::clone(&source))?.with_batch_size(16)?, + ); + let expression: Arc = Arc::new(BinaryExpr::new( + column, + Operator::Plus, + Arc::new(Literal::new(ScalarValue::Int32(Some(1)))), + )); + let projection: Arc = Arc::new(ProjectionExec::try_new( + vec![(expression, "b".to_string())], + Arc::clone(&filter), + )?); + Ok((vec![projection, filter, source], gate)) +} diff --git a/datafusion/physical-plan/tests/partition_metrics.rs b/datafusion/physical-plan/tests/partition_metrics.rs new file mode 100644 index 0000000000000..3d4e97bce316f --- /dev/null +++ b/datafusion/physical-plan/tests/partition_metrics.rs @@ -0,0 +1,101 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#[path = "metrics/plan.rs"] +mod plan; + +use std::sync::Arc; + +use arrow::array::Int32Array; +use datafusion_common::Result; +use datafusion_execution::TaskContext; +use futures::{FutureExt, StreamExt, TryStreamExt}; + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn shared_tree_partition_metrics() -> Result<()> { + tokio::time::timeout(std::time::Duration::from_secs(30), run_shared_tree()) + .await + .expect("shared tree did not complete") +} + +async fn run_shared_tree() -> Result<()> { + let (nodes, gate) = plan::shared_plan(16)?; + let context = Arc::new(TaskContext::default()); + for node in &nodes { + assert_eq!(node.metrics().unwrap().for_partition(0).iter().count(), 0); + } + let mut long_stream = nodes[0].execute(0, Arc::clone(&context))?; + let first = long_stream.next().await.unwrap()?; + assert_eq!(first.num_rows(), 16); + assert!(long_stream.next().now_or_never().is_none()); + let early = nodes[0].metrics().unwrap().for_partition(0); + assert_eq!(early.output_rows(), Some(16)); + + // Concurrent execution and registration on the same Arc. + let mut tasks = tokio::task::JoinSet::new(); + for partition in 1..16 { + let root = Arc::clone(&nodes[0]); + let context = Arc::clone(&context); + tasks.spawn(async move { + let batches: Vec<_> = root.execute(partition, context)?.try_collect().await?; + for batch in &batches { + let values = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(values.values().as_ref(), &(1..17).collect::>()); + } + assert_eq!(batches.iter().map(|b| b.num_rows()).sum::(), 32); + Ok::<_, datafusion_common::DataFusionError>(()) + }); + } + while let Some(result) = tasks.join_next().await { + result.unwrap()?; + } + // Other completed partitions do not affect the still-running partition. + assert_eq!(early.output_rows(), Some(16)); + assert_eq!( + nodes[0].metrics().unwrap().for_partition(0).output_rows(), + Some(16) + ); + gate.notify_one(); + let remaining: Vec<_> = long_stream.try_collect().await?; + assert_eq!(remaining.iter().map(|b| b.num_rows()).sum::(), 16); + assert_eq!(early.output_rows(), Some(32)); + + for (node_index, node) in nodes.iter().enumerate() { + let rows = if node_index == 2 { 64 } else { 32 }; + let full = node.metrics().unwrap(); + assert_eq!(full.output_rows(), Some(rows * 16)); + for partition in 0..16 { + let selected = node.metrics().unwrap().for_partition(partition); + assert_eq!(selected.output_rows(), Some(rows)); + assert_eq!(selected.aggregate_by_name().output_rows(), Some(rows)); + let expected: Vec<_> = full + .iter() + .filter(|m| m.partition() == Some(partition)) + .collect(); + assert_eq!(selected.iter().count(), expected.len()); + for (actual, expected) in selected.iter().zip(expected) { + assert!(Arc::ptr_eq(actual, expected)); + } + } + assert_eq!(node.metrics().unwrap().for_partition(16).iter().count(), 0); + } + Ok(()) +}