kosiew commented on code in PR #25699:
URL: https://github.com/apache/datafusion/pull/25699#discussion_r4192584360
##########
datafusion/physical-plan/src/statistics.rs:
##########
@@ -307,6 +307,73 @@ 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));
+ ///
Review Comment:
Optional: could we also assert `correlation.coefficient == 0.92` here? That
would make this example show how to read the correlation value as well as the
column indexes.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
