From 1d243ebe19957732c3ddd12e1c03d17ae3754188 Mon Sep 17 00:00:00 2001 From: Alessandro Solimando Date: Thu, 24 Sep 2026 10:37:55 +0200 Subject: [PATCH 1/3] docs(physical-plan): add doctest examples for StatisticsRegistry, ExtendedStatistics, and compute_extended Add real, compiled `# Example` doctests to `StatisticsContext::compute_extended`, the `StatisticsRegistry` module doc, and the `ExtendedStatistics` struct doc. The `StatisticsRegistry`/`ExtendedStatistics` examples were previously wrapped in ` ```ignore ` blocks referencing undefined names, so they were never compiled or checked. `compute_extended` had no example at all. The shared running example is a cross-column correlation extension (`ColumnCorrelation`), a real, node-level statistic (unlike a per-column histogram, which has no home in the current single node-level `Extensions` field). Closes #25698. --- .../src/operator_statistics/mod.rs | 75 ++++++++++++++----- datafusion/physical-plan/src/statistics.rs | 67 +++++++++++++++++ 2 files changed, 123 insertions(+), 19 deletions(-) diff --git a/datafusion/physical-plan/src/operator_statistics/mod.rs b/datafusion/physical-plan/src/operator_statistics/mod.rs index 07a0fa49d64b3..682df8e4805d3 100644 --- a/datafusion/physical-plan/src/operator_statistics/mod.rs +++ b/datafusion/physical-plan/src/operator_statistics/mod.rs @@ -62,19 +62,51 @@ //! //! # Example //! -//! ```ignore -//! use datafusion_physical_plan::operator_statistics::*; +//! ``` +//! use std::sync::Arc; +//! use arrow::datatypes::Schema; +//! use datafusion_common::stats::Precision; +//! use datafusion_common::{Result, Statistics}; +//! use datafusion_physical_plan::ExecutionPlan; +//! use datafusion_physical_plan::operator_statistics::{ +//! ExtendedStatistics, StatisticsProvider, StatisticsRegistry, StatisticsResult, +//! }; +//! use datafusion_physical_plan::statistics::{StatisticsArgs, StatisticsContext}; +//! use datafusion_physical_plan::test::exec::StatisticsExec; //! -//! // Create an empty registry -//! let mut registry = StatisticsRegistry::new(); +//! #[derive(Debug)] +//! struct MyProvider; +//! +//! impl StatisticsProvider for MyProvider { +//! fn matches(&self, plan: &dyn ExecutionPlan) -> bool { +//! plan.downcast_ref::().is_some() +//! } //! -//! // Register custom provider (higher priority) -//! registry.register(Arc::new(MyHistogramProvider)); +//! fn compute_statistics( +//! &self, +//! plan: &dyn ExecutionPlan, +//! _child_stats: &[ExtendedStatistics], +//! ) -> Result { +//! let base = plan.statistics_from_inputs(&[], &StatisticsArgs::new())?; +//! Ok(StatisticsResult::Computed(ExtendedStatistics::new_arc(base))) +//! } +//! } +//! +//! let plan = StatisticsExec::new(Statistics::new_unknown(&Schema::empty()), Schema::empty()); +//! +//! let mut registry = StatisticsRegistry::new(); +//! registry.register(Arc::new(MyProvider)); //! -//! // Compute statistics through the chain //! let stats = StatisticsContext::new_with_registry(registry) -//! .compute_extended(plan.as_ref(), &StatisticsArgs::new())?; +//! .compute_extended(&plan, &StatisticsArgs::new())?; +//! assert_eq!(stats.base().num_rows, Precision::Absent); +//! # Ok::<(), datafusion_common::DataFusionError>(()) //! ``` +//! +//! See [`StatisticsContext::compute_extended`] for an example that also +//! attaches a custom extension to the computed statistics. +//! +//! [`StatisticsContext::compute_extended`]: crate::statistics::StatisticsContext::compute_extended use std::fmt::{self, Debug}; use std::sync::Arc; @@ -104,21 +136,26 @@ use crate::union::UnionExec; /// /// # Example /// -/// ```ignore -/// // Define a custom statistics extension +/// ``` +/// use arrow::datatypes::Schema; +/// use datafusion_common::Statistics; +/// use datafusion_physical_plan::operator_statistics::ExtendedStatistics; +/// +/// // Pearson correlation between two columns, identified by schema index. /// #[derive(Debug, Clone)] -/// struct HistogramStats { -/// buckets: Vec<(i64, i64, usize)>, // (min, max, count) +/// struct ColumnCorrelation { +/// columns: (usize, usize), +/// coefficient: f64, /// } /// -/// // Set extension in a planner -/// let mut stats = ExtendedStatistics::from(base_stats); -/// stats.set_extension(HistogramStats { buckets: vec![] }); +/// let mut stats = ExtendedStatistics::from(Statistics::new_unknown(&Schema::empty())); +/// stats.set_extension(ColumnCorrelation { +/// columns: (0, 1), +/// coefficient: 0.92, +/// }); /// -/// // Retrieve in a consumer -/// if let Some(hist) = stats.get_extension::() { -/// // Use histogram for better estimation -/// } +/// let correlation = stats.get_extension::().unwrap(); +/// assert_eq!(correlation.coefficient, 0.92); /// ``` #[derive(Debug, Clone, Default)] pub struct ExtendedStatistics { diff --git a/datafusion/physical-plan/src/statistics.rs b/datafusion/physical-plan/src/statistics.rs index c1a8ad29bc5fa..e81b18d703104 100644 --- a/datafusion/physical-plan/src/statistics.rs +++ b/datafusion/physical-plan/src/statistics.rs @@ -181,6 +181,73 @@ impl StatisticsContext { /// Computes the [`ExtendedStatistics`] for `plan`: the core statistics plus /// any extensions a provider attached to this node (see the type-level docs /// for how extensions propagate up the tree). + /// + /// # Example + /// + /// ``` + /// use std::sync::Arc; + /// use arrow::datatypes::{DataType, Field, Schema}; + /// use datafusion_common::{ColumnStatistics, Result, Statistics}; + /// use datafusion_common::stats::Precision; + /// use datafusion_physical_plan::ExecutionPlan; + /// use datafusion_physical_plan::operator_statistics::{ + /// ExtendedStatistics, StatisticsProvider, StatisticsRegistry, StatisticsResult, + /// }; + /// use datafusion_physical_plan::statistics::{StatisticsArgs, StatisticsContext}; + /// use datafusion_physical_plan::test::exec::StatisticsExec; + /// + /// // Pearson correlation between two columns, identified by schema index. + /// #[derive(Debug, Clone)] + /// struct ColumnCorrelation { + /// columns: (usize, usize), + /// coefficient: f64, + /// } + /// + /// #[derive(Debug)] + /// struct CorrelationProvider; + /// + /// impl StatisticsProvider for CorrelationProvider { + /// fn matches(&self, plan: &dyn ExecutionPlan) -> bool { + /// plan.downcast_ref::().is_some() + /// } + /// + /// fn compute_statistics( + /// &self, + /// plan: &dyn ExecutionPlan, + /// _child_stats: &[ExtendedStatistics], + /// ) -> Result { + /// let base = plan.statistics_from_inputs(&[], &StatisticsArgs::new())?; + /// let mut extended = ExtendedStatistics::new_arc(base); + /// extended.set_extension(ColumnCorrelation { + /// columns: (0, 1), + /// coefficient: 0.92, + /// }); + /// Ok(StatisticsResult::Computed(extended)) + /// } + /// } + /// + /// let schema = Schema::new(vec![ + /// Field::new("city", DataType::Utf8, false), + /// Field::new("zip_code", DataType::Utf8, false), + /// ]); + /// let stats = Statistics { + /// num_rows: Precision::Exact(1000), + /// total_byte_size: Precision::Absent, + /// column_statistics: vec![ColumnStatistics::new_unknown(); 2], + /// }; + /// let plan = StatisticsExec::new(stats, schema); + /// + /// let mut registry = StatisticsRegistry::new(); + /// registry.register(Arc::new(CorrelationProvider)); + /// let context = StatisticsContext::new_with_registry(registry); + /// + /// let extended = context.compute_extended(&plan, &StatisticsArgs::new())?; + /// assert_eq!(extended.base().num_rows, Precision::Exact(1000)); + /// + /// let correlation = extended.get_extension::().unwrap(); + /// assert_eq!(correlation.columns, (0, 1)); + /// # Ok::<(), datafusion_common::DataFusionError>(()) + /// ``` pub fn compute_extended( &self, plan: &dyn ExecutionPlan, From 1c3458a27b4b34f319cb02358b4ccdc0abe52efe Mon Sep 17 00:00:00 2001 From: Alessandro Solimando Date: Thu, 24 Sep 2026 22:39:25 +0200 Subject: [PATCH 2/3] docs(physical-plan): hide use statements in the new doctests Prefix the use statements in the compute_extended, StatisticsRegistry, and ExtendedStatistics doctests with # so they compile but do not show in the rendered documentation, per alamb's review on PR #25689. --- .../src/operator_statistics/mod.rs | 26 +++++++++---------- datafusion/physical-plan/src/statistics.rs | 20 +++++++------- 2 files changed, 23 insertions(+), 23 deletions(-) diff --git a/datafusion/physical-plan/src/operator_statistics/mod.rs b/datafusion/physical-plan/src/operator_statistics/mod.rs index 682df8e4805d3..d53d3e27f570b 100644 --- a/datafusion/physical-plan/src/operator_statistics/mod.rs +++ b/datafusion/physical-plan/src/operator_statistics/mod.rs @@ -63,16 +63,16 @@ //! # Example //! //! ``` -//! use std::sync::Arc; -//! use arrow::datatypes::Schema; -//! use datafusion_common::stats::Precision; -//! use datafusion_common::{Result, Statistics}; -//! use datafusion_physical_plan::ExecutionPlan; -//! use datafusion_physical_plan::operator_statistics::{ -//! ExtendedStatistics, StatisticsProvider, StatisticsRegistry, StatisticsResult, -//! }; -//! use datafusion_physical_plan::statistics::{StatisticsArgs, StatisticsContext}; -//! use datafusion_physical_plan::test::exec::StatisticsExec; +//! # use std::sync::Arc; +//! # use arrow::datatypes::Schema; +//! # use datafusion_common::stats::Precision; +//! # use datafusion_common::{Result, Statistics}; +//! # use datafusion_physical_plan::ExecutionPlan; +//! # use datafusion_physical_plan::operator_statistics::{ +//! # ExtendedStatistics, StatisticsProvider, StatisticsRegistry, StatisticsResult, +//! # }; +//! # use datafusion_physical_plan::statistics::{StatisticsArgs, StatisticsContext}; +//! # use datafusion_physical_plan::test::exec::StatisticsExec; //! //! #[derive(Debug)] //! struct MyProvider; @@ -137,9 +137,9 @@ use crate::union::UnionExec; /// # Example /// /// ``` -/// use arrow::datatypes::Schema; -/// use datafusion_common::Statistics; -/// use datafusion_physical_plan::operator_statistics::ExtendedStatistics; +/// # use arrow::datatypes::Schema; +/// # use datafusion_common::Statistics; +/// # use datafusion_physical_plan::operator_statistics::ExtendedStatistics; /// /// // Pearson correlation between two columns, identified by schema index. /// #[derive(Debug, Clone)] diff --git a/datafusion/physical-plan/src/statistics.rs b/datafusion/physical-plan/src/statistics.rs index e81b18d703104..fcc22682bbf3c 100644 --- a/datafusion/physical-plan/src/statistics.rs +++ b/datafusion/physical-plan/src/statistics.rs @@ -185,16 +185,16 @@ impl StatisticsContext { /// # Example /// /// ``` - /// use std::sync::Arc; - /// use arrow::datatypes::{DataType, Field, Schema}; - /// use datafusion_common::{ColumnStatistics, Result, Statistics}; - /// use datafusion_common::stats::Precision; - /// use datafusion_physical_plan::ExecutionPlan; - /// use datafusion_physical_plan::operator_statistics::{ - /// ExtendedStatistics, StatisticsProvider, StatisticsRegistry, StatisticsResult, - /// }; - /// use datafusion_physical_plan::statistics::{StatisticsArgs, StatisticsContext}; - /// use datafusion_physical_plan::test::exec::StatisticsExec; + /// # use std::sync::Arc; + /// # use arrow::datatypes::{DataType, Field, Schema}; + /// # use datafusion_common::{ColumnStatistics, Result, Statistics}; + /// # use datafusion_common::stats::Precision; + /// # use datafusion_physical_plan::ExecutionPlan; + /// # use datafusion_physical_plan::operator_statistics::{ + /// # ExtendedStatistics, StatisticsProvider, StatisticsRegistry, StatisticsResult, + /// # }; + /// # use datafusion_physical_plan::statistics::{StatisticsArgs, StatisticsContext}; + /// # use datafusion_physical_plan::test::exec::StatisticsExec; /// /// // Pearson correlation between two columns, identified by schema index. /// #[derive(Debug, Clone)] From 2bc67a13e3c933b6f7963bcd1782dbe86e6b1ead Mon Sep 17 00:00:00 2001 From: Alessandro Solimando Date: Tue, 6 Oct 2026 09:57:35 +0200 Subject: [PATCH 3/3] docs(physical-plan): assert the correlation coefficient in compute_extended doctest Also demonstrates reading the coefficient, not just the column indexes. Addresses kosiew's review suggestion on PR #25699. --- datafusion/physical-plan/src/statistics.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/datafusion/physical-plan/src/statistics.rs b/datafusion/physical-plan/src/statistics.rs index 7b243577f0cde..7124c9107dd0a 100644 --- a/datafusion/physical-plan/src/statistics.rs +++ b/datafusion/physical-plan/src/statistics.rs @@ -372,6 +372,7 @@ impl StatisticsContext { /// /// let correlation = extended.get_extension::().unwrap(); /// assert_eq!(correlation.columns, (0, 1)); + /// assert_eq!(correlation.coefficient, 0.92); /// # Ok::<(), datafusion_common::DataFusionError>(()) /// ``` pub fn compute_extended(