docs(physical-plan): add doctest examples for StatisticsRegistry, ExtendedStatistics, and compute_extended

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::<StatisticsExec>().is_some()
+//!     }
 //!
-//! // Register custom provider (higher priority)
-//! registry.register(Arc::new(MyHistogramProvider));
+//!     fn compute_statistics(
+//!         &self,
+//!         plan: &dyn ExecutionPlan,
+//!         _child_stats: &[ExtendedStatistics],
+//!     ) -> Result<StatisticsResult> {
+//!         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::<HistogramStats>() {
-///     // Use histogram for better estimation
-/// }
+/// let correlation = stats.get_extension::<ColumnCorrelation>().unwrap();
+/// assert_eq!(correlation.coefficient, 0.92);
 /// ```
 #[derive(Debug, Clone, Default)]
 pub struct ExtendedStatistics {

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::<StatisticsExec>().is_some()
+    ///     }
+    ///
+    ///     fn compute_statistics(
+    ///         &self,
+    ///         plan: &dyn ExecutionPlan,
+    ///         _child_stats: &[ExtendedStatistics],
+    ///     ) -> Result<StatisticsResult> {
+    ///         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::<ColumnCorrelation>().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,