diff --git a/datafusion/physical-plan/src/operator_statistics/mod.rs b/datafusion/physical-plan/src/operator_statistics/mod.rs index e09aa6b86a1a7..49aff1011144d 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 b4194d61e6407..7124c9107dd0a 100644 --- a/datafusion/physical-plan/src/statistics.rs +++ b/datafusion/physical-plan/src/statistics.rs @@ -307,6 +307,74 @@ impl StatisticsContext { /// for how extensions propagate up the tree). As with [`Self::compute`], /// the root is not memoized and must be a standalone plan node; see /// [`Self::compute_extended_arc`]. + /// + /// # 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)); + /// assert_eq!(correlation.coefficient, 0.92); + /// # Ok::<(), datafusion_common::DataFusionError>(()) + /// ``` pub fn compute_extended( &self, plan: &dyn ExecutionPlan,