diff --git a/datafusion-examples/README.md b/datafusion-examples/README.md index 86cfffe1a80e8..6c7aad31a9233 100644 --- a/datafusion-examples/README.md +++ b/datafusion-examples/README.md @@ -214,6 +214,16 @@ cargo run --example dataframe -- dataframe | frontend | [`sql_ops/frontend.rs`](examples/sql_ops/frontend.rs) | Build LogicalPlans from SQL | | query | [`sql_ops/query.rs`](examples/sql_ops/query.rs) | Query data using SQL | +## Statistics Examples + +### Group: `statistics` + +#### Category: Single Process + +| Subcommand | File Path | Description | +| ------------ | ------------------------------------------------------------------- | ----------------------------------------------------------------------- | +| join_reorder | [`statistics/join_reorder.rs`](examples/statistics/join_reorder.rs) | Supply and refine column statistics via a provider to flip a join order | + ## UDF Examples ### Group: `udf` diff --git a/datafusion-examples/examples/statistics/join_reorder.rs b/datafusion-examples/examples/statistics/join_reorder.rs new file mode 100644 index 0000000000000..0eccc93630af2 --- /dev/null +++ b/datafusion-examples/examples/statistics/join_reorder.rs @@ -0,0 +1,179 @@ +// 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. + +//! Plug refined cardinality estimation into the optimizer via the `StatisticsRegistry`. +//! +//! DataFusion's built-in statistics are intentionally simple defaults; the registry +//! is the seam for plugging in more refined estimation (formulas or heuristics +//! suited to your system or data) without changing the core. +//! +//! `(SELECT user_id FROM events WHERE amount < 50 GROUP BY user_id) JOIN dims`: a +//! provider supplies the column stats a catalog knows (`amount` range, `user_id` +//! distinct count); the built-in `FilterStatisticsProvider` then refines the +//! post-filter distinct count with the survival formula +//! `NDV * (1 - (1 - selectivity)^(rows / NDV))` (Yao/Cardenas) to ~32, below `dims` +//! (48), flipping the join build side. Core's simpler `min(NDV, rows)` cap would +//! give 50 (> 48) and keep the other order; the refinement is the point. The +//! ground-truth query prints the true surviving distinct count (below 48), +//! confirming the flip. + +use std::sync::Arc; + +use datafusion::arrow::array::Int32Array; +use datafusion::arrow::datatypes::{DataType, Field, Schema}; +use datafusion::arrow::record_batch::RecordBatch; +use datafusion::arrow::util::pretty::pretty_format_batches; +use datafusion::catalog::MemTable; +use datafusion::common::Result; +use datafusion::common::ScalarValue; +use datafusion::common::stats::Precision; +use datafusion::execution::SessionStateBuilder; +use datafusion::physical_plan::ExecutionPlan; +use datafusion::physical_plan::operator_statistics::{ + ClosureStatisticsProvider, ExtendedStatistics, StatisticsRegistry, StatisticsResult, +}; +use datafusion::physical_plan::statistics::StatisticsArgs; +use datafusion::prelude::*; +use rand::{Rng, SeedableRng, rngs::StdRng}; + +/// Matches the base `events` scan: a leaf carrying both `user_id` and `amount`. +fn catalog_matches(plan: &dyn ExecutionPlan) -> bool { + let schema = plan.schema(); + plan.children().is_empty() + && schema.index_of("user_id").is_ok() + && schema.index_of("amount").is_ok() +} + +/// Injects the catalog-known `amount` range (for filter selectivity) and +/// `user_id` distinct count (for post-filter NDV) the in-memory table lacks. +fn catalog_stats( + plan: &dyn ExecutionPlan, + _child_stats: &[ExtendedStatistics], +) -> Result { + let schema = plan.schema(); + let user_id = schema.index_of("user_id")?; + let amount = schema.index_of("amount")?; + let mut stats = (*plan.statistics_from_inputs(&[], &StatisticsArgs::new())?).clone(); + stats.column_statistics[amount].min_value = + Precision::Inexact(ScalarValue::Int32(Some(0))); + stats.column_statistics[amount].max_value = + Precision::Inexact(ScalarValue::Int32(Some(999))); + stats.column_statistics[user_id].distinct_count = Precision::Inexact(100); + Ok(StatisticsResult::Computed(ExtendedStatistics::new(stats))) +} + +fn int_col(values: &[i32]) -> Arc { + Arc::new(Int32Array::from_iter_values(values.iter().copied())) +} + +fn mem_table(fields: &[(&str, Arc)]) -> Result> { + let schema = Arc::new(Schema::new( + fields + .iter() + .map(|(name, _)| Field::new(*name, DataType::Int32, false)) + .collect::>(), + )); + let cols = fields.iter().map(|(_, col)| Arc::clone(col) as _).collect(); + let batch = RecordBatch::try_new(Arc::clone(&schema), cols)?; + Ok(Arc::new(MemTable::try_new(schema, vec![vec![batch]])?)) +} + +fn build_ctx(with_registry: bool) -> Result { + let config = SessionConfig::new() + .with_target_partitions(4) + .set_bool("datafusion.explain.physical_plan_only", true) + .set_bool("datafusion.explain.show_statistics", true) + // Force Partitioned hash joins so statistics alone drive the build side. + .set_usize( + "datafusion.optimizer.hash_join_single_partition_threshold", + 1, + ) + .set_usize( + "datafusion.optimizer.hash_join_single_partition_threshold_rows", + 1, + ); + + let mut builder = SessionStateBuilder::new() + .with_config(config) + .with_default_features(); + if with_registry { + let mut registry = StatisticsRegistry::default_with_builtin_providers(); + registry.register(Arc::new(ClosureStatisticsProvider::with_matches( + catalog_matches, + catalog_stats, + ))); + builder = builder.with_statistics_registry(registry); + } + let ctx = SessionContext::new_with_state(builder.build()); + + let n = 1000i32; + let user_ids: Vec = (0..n).map(|v| v % 100).collect(); + // `amount` independent of `user_id` so the filter keeps a representative sample. + let mut rng = StdRng::seed_from_u64(2024); + let amounts: Vec = (0..n).map(|_| rng.random_range(0..1000)).collect(); + ctx.register_table( + "events", + mem_table(&[ + ("user_id", int_col(&user_ids)), + ("amount", int_col(&amounts)), + ])?, + )?; + ctx.register_table( + "dims", + mem_table(&[ + ("user_id", int_col(&(0..48).collect::>())), + ("label", int_col(&(0..48).collect::>())), + ])?, + )?; + Ok(ctx) +} + +const QUERY: &str = "SELECT e.user_id, d.label \ + FROM (SELECT user_id FROM events WHERE amount < 50 GROUP BY user_id) e \ + JOIN dims d ON e.user_id = d.user_id"; + +async fn explain(ctx: &SessionContext) -> Result { + let batches = ctx + .sql(&format!("EXPLAIN {QUERY}")) + .await? + .collect() + .await?; + Ok(pretty_format_batches(&batches)?.to_string()) +} + +pub async fn join_reorder() -> Result<()> { + let truth_query = "SELECT count(DISTINCT user_id) AS true_distinct_users \ + FROM events WHERE amount < 50"; + println!("-- Ground truth --\n{truth_query}\n"); + let truth = build_ctx(false)?.sql(truth_query).await?.collect().await?; + println!("{}\n", pretty_format_batches(&truth)?); + + println!("-- Query --\n{QUERY}\n"); + println!( + "A hash join builds its in-memory hash table from one input and probes with\n\ + the other, so the smaller input should be the build side. Default estimation\n\ + sizes the grouped `events` at 1000 rows and builds from `dims`; the\n\ + registry's refined ~32 estimate is below `dims` (48 rows), so it flips the\n\ + build side to `events`. The ground-truth count above (also below 48)\n\ + confirms `events` really is the smaller, cheaper side.\n" + ); + println!("-- Without the registry (default estimation) --"); + println!("{}\n", explain(&build_ctx(false)?).await?); + println!("-- With the registry (built-in refinement) --"); + println!("{}", explain(&build_ctx(true)?).await?); + Ok(()) +} diff --git a/datafusion-examples/examples/statistics/main.rs b/datafusion-examples/examples/statistics/main.rs new file mode 100644 index 0000000000000..313e18645d50f --- /dev/null +++ b/datafusion-examples/examples/statistics/main.rs @@ -0,0 +1,84 @@ +// 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. + +//! # Pluggable operator statistics (`StatisticsRegistry`) +//! +//! These examples show how a `StatisticsProvider` registered on the session can +//! change a physical plan by supplying statistics the built-in estimators +//! cannot express. +//! +//! ## Usage +//! ```bash +//! cargo run --example statistics -- [all|join_reorder] +//! ``` +//! +//! Each subcommand runs a corresponding example: +//! - `all`: run all examples included in this module +//! +//! - `join_reorder` +//! (file: join_reorder.rs, desc: Supply and refine column statistics via a provider to flip a join order) + +mod join_reorder; + +use datafusion::error::{DataFusionError, Result}; +use strum::{IntoEnumIterator, VariantNames}; +use strum_macros::{Display, EnumIter, EnumString, VariantNames}; + +#[derive(EnumIter, EnumString, Display, VariantNames)] +#[strum(serialize_all = "snake_case")] +enum ExampleKind { + All, + JoinReorder, +} + +impl ExampleKind { + const EXAMPLE_NAME: &str = "statistics"; + + fn runnable() -> impl Iterator { + ExampleKind::iter().filter(|v| !matches!(v, ExampleKind::All)) + } + + async fn run(&self) -> Result<()> { + match self { + ExampleKind::All => { + for example in ExampleKind::runnable() { + println!("Running example: {example}"); + Box::pin(example.run()).await?; + } + Ok(()) + } + ExampleKind::JoinReorder => join_reorder::join_reorder().await, + } + } +} + +#[tokio::main] +async fn main() -> Result<()> { + let usage = format!( + "Usage: cargo run --example {} -- [{}]", + ExampleKind::EXAMPLE_NAME, + ExampleKind::VARIANTS.join("|") + ); + + let example: ExampleKind = std::env::args() + .nth(1) + .unwrap_or_else(|| ExampleKind::All.to_string()) + .parse() + .map_err(|_| DataFusionError::Execution(format!("Unknown example. {usage}")))?; + + example.run().await +} diff --git a/datafusion/common/src/config.rs b/datafusion/common/src/config.rs index 81f573fc2a23e..197ee9b29a92e 100644 --- a/datafusion/common/src/config.rs +++ b/datafusion/common/src/config.rs @@ -1568,11 +1568,10 @@ config_namespace! { /// query is used. pub join_reordering: bool, default = true - /// When set to true, the physical plan optimizer uses the pluggable - /// `StatisticsRegistry` for statistics propagation across operators. - /// This enables more accurate cardinality estimates compared to each - /// operator's built-in `partition_statistics`. - pub use_statistics_registry: bool, default = false + /// (Deprecated) Ignored: the physical plan optimizer always consults the + /// session's pluggable `StatisticsRegistry` (register providers on the + /// `SessionState`; with none it is a no-op). + pub use_statistics_registry: bool, warn = "`use_statistics_registry` is deprecated and ignored; the StatisticsRegistry is always consulted, register providers on the SessionState", default = false /// When set to true, the physical plan optimizer will prefer HashJoin over SortMergeJoin. /// HashJoin can work more efficiently than SortMergeJoin but consumes more memory diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index 4e914556b4cc0..c6cfcdc420ff2 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -2738,6 +2738,10 @@ impl DefaultPhysicalPlanner { let explain_format = &e.explain_format; // Statement-level override wins over session config for show_statistics. let show_statistics = e.show_statistics.unwrap_or(config.show_statistics); + let statistics_registry = session_state + .statistics_registry() + .cloned() + .unwrap_or_default(); if !e.logical_optimization_succeeded { return Ok(Arc::new(ExplainExec::new( @@ -2801,6 +2805,15 @@ impl DefaultPhysicalPlanner { } if !config.logical_plan_only && e.logical_optimization_succeeded { + let render_indent = + |plan: &dyn ExecutionPlan, show_statistics: bool, show_schema: bool| { + displayable(plan) + .set_show_statistics(show_statistics) + .set_statistics_registry(statistics_registry.clone()) + .set_show_schema(show_schema) + .indent(e.verbose) + .to_string() + }; match self .create_initial_plan(e.plan.as_ref(), session_state) .await @@ -2809,11 +2822,11 @@ impl DefaultPhysicalPlanner { // Include statistics / schema if enabled stringified_plans.push(StringifiedPlan::new( InitialPhysicalPlan, - displayable(input.as_ref()) - .set_show_statistics(show_statistics) - .set_show_schema(config.show_schema) - .indent(e.verbose) - .to_string(), + render_indent( + input.as_ref(), + show_statistics, + config.show_schema, + ), )); // Show statistics + schema in verbose output even if not @@ -2822,15 +2835,14 @@ impl DefaultPhysicalPlanner { if !show_statistics { stringified_plans.push(StringifiedPlan::new( InitialPhysicalPlanWithStats, - displayable(input.as_ref()) - .set_show_statistics(true) - .indent(e.verbose) - .to_string(), + render_indent(input.as_ref(), true, false), )); } if !config.show_schema { stringified_plans.push(StringifiedPlan::new( InitialPhysicalPlanWithSchema, + // Schema only: statistics are off, so this + // renders without the registry. displayable(input.as_ref()) .set_show_schema(true) .indent(e.verbose) @@ -2847,11 +2859,7 @@ impl DefaultPhysicalPlanner { let plan_type = OptimizedPhysicalPlan { optimizer_name }; stringified_plans.push(StringifiedPlan::new( plan_type, - displayable(plan) - .set_show_statistics(show_statistics) - .set_show_schema(config.show_schema) - .indent(e.verbose) - .to_string(), + render_indent(plan, show_statistics, config.show_schema), )); }, ); @@ -2860,11 +2868,11 @@ impl DefaultPhysicalPlanner { // This plan will includes statistics if show_statistics is on stringified_plans.push(StringifiedPlan::new( FinalPhysicalPlan, - displayable(input.as_ref()) - .set_show_statistics(show_statistics) - .set_show_schema(config.show_schema) - .indent(e.verbose) - .to_string(), + render_indent( + input.as_ref(), + show_statistics, + config.show_schema, + ), )); // Show statistics + schema in verbose output even if not @@ -2873,10 +2881,7 @@ impl DefaultPhysicalPlanner { if !show_statistics { stringified_plans.push(StringifiedPlan::new( FinalPhysicalPlanWithStats, - displayable(input.as_ref()) - .set_show_statistics(true) - .indent(e.verbose) - .to_string(), + render_indent(input.as_ref(), true, false), )); } if !config.show_schema { @@ -2940,11 +2945,16 @@ impl DefaultPhysicalPlanner { ExplainAnalyzeCategories::All => None, ExplainAnalyzeCategories::Only(cats) => Some(cats), }; + let statistics_registry = session_state + .statistics_registry() + .cloned() + .unwrap_or_default(); Ok(Arc::new( AnalyzeExec::builder(a.verbose, show_statistics, input, schema) .with_metric_types(metric_types) .with_metric_categories(metric_categories) .with_format(a.format.clone()) + .with_statistics_registry(statistics_registry) .build(), )) } diff --git a/datafusion/physical-optimizer/src/join_selection.rs b/datafusion/physical-optimizer/src/join_selection.rs index 42736f8205089..947fcf855df19 100644 --- a/datafusion/physical-optimizer/src/join_selection.rs +++ b/datafusion/physical-optimizer/src/join_selection.rs @@ -57,17 +57,17 @@ impl JoinSelection { } } -/// Get statistics for a plan node, using the registry if available. +/// Get statistics for a plan node, consulting the registry's providers if one +/// is available. fn get_stats( plan: &dyn ExecutionPlan, registry: Option<&StatisticsRegistry>, ) -> Result> { - if let Some(reg) = registry { - reg.compute(plan) - .map(|s| Arc::::clone(s.base_arc())) - } else { - StatisticsContext::new().compute(plan, &StatisticsArgs::new()) - } + let ctx = match registry { + Some(reg) => StatisticsContext::new_with_registry(reg.clone()), + None => StatisticsContext::new(), + }; + ctx.compute(plan, &StatisticsArgs::new()) } // TODO: We need some performance test for Right Semi/Right Join swap to Left Semi/Left Join in case that the right side is smaller but not much smaller. @@ -153,16 +153,7 @@ impl PhysicalOptimizerRule for JoinSelection { context: &dyn PhysicalOptimizerContext, ) -> Result> { let config = context.config_options(); - let mut default_registry = None; - let registry: Option<&StatisticsRegistry> = - if config.optimizer.use_statistics_registry { - Some(context.statistics_registry().unwrap_or_else(|| { - default_registry - .insert(StatisticsRegistry::default_with_builtin_providers()) - })) - } else { - None - }; + let registry = context.statistics_registry(); let subrules: Vec> = vec![ Box::new(hash_join_convert_symmetric_subrule), Box::new(hash_join_swap_subrule), diff --git a/datafusion/physical-plan/src/analyze.rs b/datafusion/physical-plan/src/analyze.rs index 31e0a27410ff9..a4cfb76614698 100644 --- a/datafusion/physical-plan/src/analyze.rs +++ b/datafusion/physical-plan/src/analyze.rs @@ -27,6 +27,7 @@ use super::{ use crate::display::DisplayableExecutionPlan; use crate::execution_plan::EvaluationType; use crate::metrics::{MetricCategory, MetricType}; +use crate::operator_statistics::StatisticsRegistry; use crate::{DisplayFormatType, ExecutionPlan, Partitioning}; use arrow::{array::StringBuilder, datatypes::SchemaRef, record_batch::RecordBatch}; @@ -54,6 +55,9 @@ pub struct AnalyzeExec { metric_categories: Option>, /// Output format for the rendered plan + metrics. format: ExplainFormat, + /// Registry consulted when rendering displayed statistics, so `EXPLAIN + /// ANALYZE` reflects the same provider-refined stats as plain `EXPLAIN`. + statistics_registry: StatisticsRegistry, /// The input plan (the plan being analyzed) pub(crate) input: Arc, /// The output schema for RecordBatches of this exec node @@ -72,6 +76,7 @@ pub struct AnalyzeExecBuilder { metric_types: Vec, metric_categories: Option>, format: ExplainFormat, + statistics_registry: StatisticsRegistry, } impl AnalyzeExecBuilder { @@ -89,6 +94,7 @@ impl AnalyzeExecBuilder { metric_types: vec![MetricType::Summary, MetricType::Dev], metric_categories: None, format: ExplainFormat::Indent, + statistics_registry: StatisticsRegistry::new(), } } @@ -110,6 +116,11 @@ impl AnalyzeExecBuilder { self } + pub fn with_statistics_registry(mut self, registry: StatisticsRegistry) -> Self { + self.statistics_registry = registry; + self + } + pub fn build(self) -> AnalyzeExec { let cache = AnalyzeExec::compute_properties(&self.input, Arc::clone(&self.schema)); @@ -119,6 +130,7 @@ impl AnalyzeExecBuilder { metric_types: self.metric_types, metric_categories: self.metric_categories, format: self.format, + statistics_registry: self.statistics_registry, input: self.input, schema: self.schema, cache: Arc::new(cache), @@ -233,6 +245,7 @@ impl ExecutionPlan for AnalyzeExec { .with_metric_types(self.metric_types.clone()) .with_metric_categories(self.metric_categories.clone()) .with_format(self.format.clone()) + .with_statistics_registry(self.statistics_registry.clone()) .build(), )) } @@ -272,6 +285,7 @@ impl ExecutionPlan for AnalyzeExec { let metric_types = self.metric_types.clone(); let metric_categories = self.metric_categories.clone(); let format = self.format.clone(); + let statistics_registry = self.statistics_registry.clone(); // future that gathers the results from all the tasks in the // JoinSet that computes the overall row count and final @@ -295,6 +309,7 @@ impl ExecutionPlan for AnalyzeExec { &metric_types, metric_categories.as_deref(), &format, + &statistics_registry, ) }; @@ -413,6 +428,7 @@ fn create_output_batch( metric_types: &[MetricType], metric_categories: Option<&[MetricCategory]>, format: &ExplainFormat, + statistics_registry: &StatisticsRegistry, ) -> Result { let mut type_builder = StringBuilder::with_capacity(1, 1024); let mut plan_builder = StringBuilder::with_capacity(1, 1024); @@ -425,6 +441,7 @@ fn create_output_batch( .set_metric_types(metric_types.to_vec()) .set_metric_categories(metric_categories.map(|c| c.to_vec())) .set_show_statistics(show_statistics) + .set_statistics_registry(statistics_registry.clone()) .indent(verbose) .to_string(); plan_builder.append_value(annotated_plan); @@ -437,6 +454,7 @@ fn create_output_batch( .set_metric_types(metric_types.to_vec()) .set_metric_categories(metric_categories.map(|c| c.to_vec())) .set_show_statistics(show_statistics) + .set_statistics_registry(statistics_registry.clone()) .indent(verbose) .to_string(); plan_builder.append_value(annotated_plan); diff --git a/datafusion/physical-plan/src/display.rs b/datafusion/physical-plan/src/display.rs index 34493a5f51742..59ba5fc8f6d3f 100644 --- a/datafusion/physical-plan/src/display.rs +++ b/datafusion/physical-plan/src/display.rs @@ -32,6 +32,7 @@ use datafusion_physical_expr::LexOrdering; use crate::metrics::{MetricCategory, MetricType, MetricValue}; use crate::render_tree::RenderTree; +use crate::operator_statistics::StatisticsRegistry; use crate::statistics::{StatisticsArgs, StatisticsContext}; use super::{ExecutionPlan, ExecutionPlanVisitor, accept}; @@ -122,6 +123,7 @@ pub struct DisplayableExecutionPlan<'a> { show_metrics: ShowMetrics, /// If statistics should be displayed show_statistics: bool, + registry: StatisticsRegistry, /// If schema should be displayed. See [`Self::set_show_schema`] show_schema: bool, /// Which metric categories should be included when rendering @@ -155,43 +157,31 @@ impl<'a> DisplayableExecutionPlan<'a> { /// Create a wrapper around an [`ExecutionPlan`] which can be /// pretty printed in a variety of ways pub fn new(inner: &'a dyn ExecutionPlan) -> Self { - Self { - inner, - show_metrics: ShowMetrics::None, - show_statistics: false, - show_schema: false, - metric_types: Self::default_metric_types(), - metric_categories: None, - metric_names: None, - tree_maximum_render_width: 240, - summary: None, - } + Self::with_show_metrics(inner, ShowMetrics::None) } /// Create a wrapper around an [`ExecutionPlan`] which can be /// pretty printed in a variety of ways that also shows aggregated /// metrics pub fn with_metrics(inner: &'a dyn ExecutionPlan) -> Self { - Self { - inner, - show_metrics: ShowMetrics::Aggregated, - show_statistics: false, - show_schema: false, - metric_types: Self::default_metric_types(), - metric_categories: None, - metric_names: None, - tree_maximum_render_width: 240, - summary: None, - } + Self::with_show_metrics(inner, ShowMetrics::Aggregated) } /// Create a wrapper around an [`ExecutionPlan`] which can be /// pretty printed in a variety of ways that also shows all low /// level metrics pub fn with_full_metrics(inner: &'a dyn ExecutionPlan) -> Self { + Self::with_show_metrics(inner, ShowMetrics::Full) + } + + fn with_show_metrics( + inner: &'a dyn ExecutionPlan, + show_metrics: ShowMetrics, + ) -> Self { Self { inner, - show_metrics: ShowMetrics::Full, + registry: StatisticsRegistry::new(), + show_metrics, show_statistics: false, show_schema: false, metric_types: Self::default_metric_types(), @@ -217,6 +207,12 @@ impl<'a> DisplayableExecutionPlan<'a> { self } + /// Set the [`StatisticsRegistry`] consulted when computing displayed statistics. + pub fn set_statistics_registry(mut self, registry: StatisticsRegistry) -> Self { + self.registry = registry; + self + } + /// Specify which metric types should be rendered alongside the plan pub fn set_metric_types(mut self, metric_types: Vec) -> Self { self.metric_types = metric_types; @@ -294,6 +290,7 @@ impl<'a> DisplayableExecutionPlan<'a> { plan: &'a dyn ExecutionPlan, show_metrics: ShowMetrics, show_statistics: bool, + registry: StatisticsRegistry, show_schema: bool, metric_types: Vec, metric_categories: Option>, @@ -307,6 +304,7 @@ impl<'a> DisplayableExecutionPlan<'a> { indent: 0, show_metrics: self.show_metrics, show_statistics: self.show_statistics, + registry: self.registry.clone(), show_schema: self.show_schema, metric_types: &self.metric_types, metric_categories: self.metric_categories.as_deref(), @@ -320,6 +318,7 @@ impl<'a> DisplayableExecutionPlan<'a> { plan: self.inner, show_metrics: self.show_metrics, show_statistics: self.show_statistics, + registry: self.registry.clone(), show_schema: self.show_schema, metric_types: self.metric_types.clone(), metric_categories: self.metric_categories.clone(), @@ -343,6 +342,7 @@ impl<'a> DisplayableExecutionPlan<'a> { plan: &'a dyn ExecutionPlan, show_metrics: ShowMetrics, show_statistics: bool, + registry: StatisticsRegistry, metric_types: Vec, metric_categories: Option>, metric_names: Option>, @@ -356,6 +356,7 @@ impl<'a> DisplayableExecutionPlan<'a> { t, show_metrics: self.show_metrics, show_statistics: self.show_statistics, + registry: self.registry.clone(), metric_types: &self.metric_types, metric_categories: self.metric_categories.as_deref(), metric_names: self.metric_names.as_deref(), @@ -376,6 +377,7 @@ impl<'a> DisplayableExecutionPlan<'a> { plan: self.inner, show_metrics: self.show_metrics, show_statistics: self.show_statistics, + registry: self.registry.clone(), metric_types: self.metric_types.clone(), metric_categories: self.metric_categories.clone(), metric_names: self.metric_names.clone(), @@ -484,6 +486,7 @@ impl<'a> DisplayableExecutionPlan<'a> { plan: &'a dyn ExecutionPlan, show_metrics: ShowMetrics, show_statistics: bool, + registry: StatisticsRegistry, show_schema: bool, metric_types: Vec, metric_categories: Option>, @@ -498,6 +501,7 @@ impl<'a> DisplayableExecutionPlan<'a> { indent: 0, show_metrics: self.show_metrics, show_statistics: self.show_statistics, + registry: self.registry.clone(), show_schema: self.show_schema, metric_types: &self.metric_types, metric_categories: self.metric_categories.as_deref(), @@ -512,6 +516,7 @@ impl<'a> DisplayableExecutionPlan<'a> { plan: self.inner, show_metrics: self.show_metrics, show_statistics: self.show_statistics, + registry: self.registry.clone(), show_schema: self.show_schema, metric_types: self.metric_types.clone(), metric_categories: self.metric_categories.clone(), @@ -568,6 +573,7 @@ struct IndentVisitor<'a, 'b> { show_metrics: ShowMetrics, /// If statistics should be displayed show_statistics: bool, + registry: StatisticsRegistry, /// If schema should be displayed show_schema: bool, /// Which metric types should be rendered @@ -619,7 +625,7 @@ impl ExecutionPlanVisitor for IndentVisitor<'_, '_> { } } if self.show_statistics { - let stats = StatisticsContext::new() + let stats = StatisticsContext::new_with_registry(self.registry.clone()) .compute(plan, &StatisticsArgs::new()) .map_err(|_e| fmt::Error)?; write!(self.f, ", statistics=[{stats}]")?; @@ -650,6 +656,7 @@ struct GraphvizVisitor<'a, 'b> { show_metrics: ShowMetrics, /// If statistics should be displayed show_statistics: bool, + registry: StatisticsRegistry, /// Which metric types should be rendered metric_types: &'a [MetricType], /// Optional filter by semantic category @@ -725,7 +732,7 @@ impl ExecutionPlanVisitor for GraphvizVisitor<'_, '_> { }; let statistics = if self.show_statistics { - let stats = StatisticsContext::new() + let stats = StatisticsContext::new_with_registry(self.registry.clone()) .compute(plan, &StatisticsArgs::new()) .map_err(|_e| fmt::Error)?; format!("statistics=[{stats}]") diff --git a/datafusion/physical-plan/src/operator_statistics/mod.rs b/datafusion/physical-plan/src/operator_statistics/mod.rs index 142768fcf49d2..2c9086280704e 100644 --- a/datafusion/physical-plan/src/operator_statistics/mod.rs +++ b/datafusion/physical-plan/src/operator_statistics/mod.rs @@ -23,9 +23,9 @@ //! //! # Overview //! -//! The default implementation delegates to each operator's built-in -//! `partition_statistics`. Users can register custom [`StatisticsProvider`] -//! implementations to: +//! With no providers registered, statistics come from each operator's built-in +//! [`ExecutionPlan::statistics_from_inputs`]. Users can register custom +//! [`StatisticsProvider`] implementations to: //! //! 1. Provide statistics for custom [`ExecutionPlan`] implementations //! 2. Override default estimation with advanced approaches (e.g., histograms) @@ -48,42 +48,32 @@ //! 5. [`JoinStatisticsProvider`] - NDV-based join output estimation (hash, sort-merge, cross) //! 6. [`LimitStatisticsProvider`] - caps output at the fetch limit (local and global) //! 7. [`UnionStatisticsProvider`] - sums input row counts -//! 8. [`DefaultStatisticsProvider`] - fallback to `partition_statistics(None)` //! -//! # Relationship to [#20184](https://github.com/apache/datafusion/issues/20184) +//! # Statistics walk //! -//! This module performs its own bottom-up tree walk in [`StatisticsRegistry::compute`], -//! separate from the walk optimizer rules do via `transform_up`. This means existing -//! rules that call `partition_statistics` directly bypass the registry. +//! Providers plug into the single [`StatisticsContext`] walk: at each node the +//! chain is consulted before falling back to the operator's built-in +//! [`ExecutionPlan::statistics_from_inputs`], with children already resolved by +//! that same walk. The first provider that returns a computed result sets the +//! node's statistics. The walk carries [`ExtendedStatistics`]; see +//! [`StatisticsContext`] for how a provider's extensions propagate. //! -//! [#20184](https://github.com/apache/datafusion/issues/20184) adds a `child_stats` -//! parameter to `partition_statistics`. Once it lands, the registry can feed enriched -//! **base** [`Statistics`] into operators' built-in `partition_statistics` calls, -//! removing redundancy for the base-stats path (row counts, column stats). However, -//! the separate registry walk is still required for [`ExtendedStatistics`] extension -//! propagation: `partition_statistics` returns `Arc`, so extensions -//! (histograms, sketches, etc.) are stripped at that boundary and can only flow -//! through the registry walk. -//! -//! If [`Statistics`] itself were extended to carry a type-erased extension map -//! (similar to [`ExtendedStatistics`]), the registry walk could be dropped entirely: -//! extensions would flow naturally through `partition_statistics(child_stats)` and -//! the registry would become a pure chain-of-responsibility on top of the existing -//! traversal with no separate walk needed. +//! [`StatisticsContext`]: crate::statistics::StatisticsContext //! //! # Example //! //! ```ignore //! use datafusion_physical_plan::operator_statistics::*; //! -//! // Create registry with default provider +//! // Create an empty registry //! let mut registry = StatisticsRegistry::new(); //! //! // Register custom provider (higher priority) //! registry.register(Arc::new(MyHistogramProvider)); //! //! // Compute statistics through the chain -//! let stats = registry.compute(plan.as_ref())?; +//! let stats = StatisticsContext::new_with_registry(registry) +//! .compute_extended(plan.as_ref(), &StatisticsArgs::new())?; //! ``` use std::fmt::{self, Debug}; @@ -94,7 +84,14 @@ use datafusion_common::stats::Precision; use datafusion_common::{Result, Statistics}; use crate::ExecutionPlan; -use crate::statistics::{StatisticsArgs, StatisticsContext}; +use crate::aggregates::{AggregateExec, AggregateMode}; +use crate::execution_plan::CardinalityEffect; +use crate::filter::FilterExec; +use crate::joins::{CrossJoinExec, HashJoinExec, JoinOnRef, SortMergeJoinExec}; +use crate::limit::{GlobalLimitExec, LocalLimitExec}; +use crate::projection::ProjectionExec; +use crate::statistics::{ChildStats, StatisticsArgs, StatisticsContext}; +use crate::union::UnionExec; // ============================================================================ // ExtendedStatistics: Statistics with type-safe extensions @@ -177,6 +174,19 @@ impl ExtendedStatistics { pub fn merge_extensions(&mut self, other: &ExtendedStatistics) { self.extensions.merge(&other.extensions); } + + /// Create from base statistics plus an existing extension map. + pub(crate) fn new_with_extensions( + base: Arc, + extensions: Extensions, + ) -> Self { + Self { base, extensions } + } + + /// Returns the extension map. + pub(crate) fn extensions(&self) -> &Extensions { + &self.extensions + } } impl From for ExtendedStatistics { @@ -223,23 +233,25 @@ pub enum StatisticsResult { /// struct MyStatisticsProvider; /// /// impl StatisticsProvider for MyStatisticsProvider { +/// fn matches(&self, plan: &dyn ExecutionPlan) -> bool { +/// plan.downcast_ref::().is_some() +/// } +/// /// fn compute_statistics( /// &self, /// plan: &dyn ExecutionPlan, /// child_stats: &[ExtendedStatistics], /// ) -> Result { -/// if let Some(my_exec) = plan.downcast_ref::() { -/// // Custom logic for MyCustomExec -/// Ok(StatisticsResult::Computed(/* ... */)) -/// } else { -/// // Let next provider handle it -/// Ok(StatisticsResult::Delegate) -/// } +/// // matches() guaranteed this node is a MyCustomExec +/// Ok(StatisticsResult::Computed(/* ... */)) /// } /// } /// ``` pub trait StatisticsProvider: Debug + Send + Sync { - /// Compute statistics for an [`ExecutionPlan`] node. + /// Compute *overall* (all-partitions) statistics for a node. + /// + /// The partition-agnostic entry point; per-partition requests go through + /// [`Self::compute_statistics_with_args`] (whose default delegates per partition). /// /// # Arguments /// * `plan` - The execution plan node to compute statistics for @@ -253,21 +265,89 @@ pub trait StatisticsProvider: Debug + Send + Sync { &self, plan: &dyn ExecutionPlan, child_stats: &[ExtendedStatistics], - ) -> Result; + ) -> Result { + let _ = (plan, child_stats); + Ok(StatisticsResult::Delegate) + } + + /// Compute statistics for `args.partition()` (or overall when `None`). + /// + /// The default is overall-only: forwards to [`Self::compute_statistics`] for + /// `None`, and returns `Delegate` for a specific partition so the operator's + /// own per-partition statistics are used. Override to provide per-partition stats. + fn compute_statistics_with_args( + &self, + plan: &dyn ExecutionPlan, + child_stats: &[ExtendedStatistics], + args: &StatisticsArgs, + ) -> Result { + if args.partition().is_some() { + Ok(StatisticsResult::Delegate) + } else { + self.compute_statistics(plan, child_stats) + } + } + + /// Which child statistics this provider needs for `plan`, one entry per + /// child in `plan.children()` order. + /// + /// Resolved independently of the operator's own + /// [`ExecutionPlan::child_stats_requests`], so a provider can refine an + /// operator that itself declares [`ChildStats::Skip`]. Results are memoized + /// per `(node, partition)`, so requesting a child already resolved for the + /// operator costs nothing. + /// + /// The default requests each child's overall (`None`) statistics, so + /// refining from child statistics works without modifying the operator. + /// Override to [`ChildStats::Skip`] a child the provider does not need. A + /// provider that computes per-partition statistics (overriding + /// [`Self::compute_statistics_with_args`]) must also override this to request + /// the matching partition, which the default ignores. + fn child_stats_requests( + &self, + plan: &dyn ExecutionPlan, + _partition: Option, + ) -> Vec { + plan.children() + .iter() + .map(|_| ChildStats::At(None)) + .collect() + } + + /// Whether this provider handles `plan`. Checked before its + /// [`Self::child_stats_requests`] are resolved, so a non-matching provider + /// never triggers child statistics computation. + fn matches(&self, _plan: &dyn ExecutionPlan) -> bool { + true + } } -/// Default statistics provider that delegates to each operator's built-in -/// `partition_statistics` implementation. +/// Deprecated statistics provider that delegates to each operator's built-in +/// [`ExecutionPlan::statistics_from_inputs`]. +/// +/// Redundant now that the statistics walk falls back to `statistics_from_inputs` +/// natively when the provider chain delegates or is empty, so a terminal +/// "default" provider is unnecessary. Kept for backward compatibility; do not +/// add it to a provider chain. +#[deprecated( + since = "55.0.0", + note = "redundant: the statistics walk falls back to `statistics_from_inputs` when the provider chain delegates or is empty; register no terminal provider" +)] #[derive(Debug, Default)] pub struct DefaultStatisticsProvider; +#[expect(deprecated)] impl StatisticsProvider for DefaultStatisticsProvider { fn compute_statistics( &self, plan: &dyn ExecutionPlan, - _child_stats: &[ExtendedStatistics], + child_stats: &[ExtendedStatistics], ) -> Result { - let base = StatisticsContext::new().compute(plan, &StatisticsArgs::new())?; + let child_base: Vec> = child_stats + .iter() + .map(|c| Arc::clone(c.base_arc())) + .collect(); + let base = plan.statistics_from_inputs(&child_base, &StatisticsArgs::new())?; Ok(StatisticsResult::Computed(ExtendedStatistics::new_arc( base, ))) @@ -299,7 +379,7 @@ impl StatisticsRegistry { /// Create a new empty registry. /// /// With no providers, `compute()` falls back to each plan node's - /// built-in `partition_statistics()`. Register providers to enhance + /// built-in `statistics_from_inputs()`. Register providers to enhance /// statistics (e.g., inject NDV, use histograms). pub fn new() -> Self { Self { @@ -322,7 +402,6 @@ impl StatisticsRegistry { /// 5. [`JoinStatisticsProvider`] /// 6. [`LimitStatisticsProvider`] /// 7. [`UnionStatisticsProvider`] - /// 8. [`DefaultStatisticsProvider`] pub fn default_with_builtin_providers() -> Self { Self::with_providers(vec![ Arc::new(FilterStatisticsProvider), @@ -332,7 +411,6 @@ impl StatisticsRegistry { Arc::new(JoinStatisticsProvider), Arc::new(LimitStatisticsProvider), Arc::new(UnionStatisticsProvider), - Arc::new(DefaultStatisticsProvider), ]) } @@ -346,52 +424,40 @@ impl StatisticsRegistry { &self.providers } - /// Compute extended statistics for a plan through the provider chain. + /// Compute extended statistics for `plan` through the provider chain. /// - /// Performs a bottom-up tree walk: child statistics are computed recursively - /// and passed to providers, mirroring how `partition_statistics` composes - /// operators. Once [#20184](https://github.com/apache/datafusion/issues/20184) - /// lands, the registry can feed enriched base stats directly into - /// `partition_statistics(child_stats)`, removing the need for a separate walk. + /// Thin wrapper over the single [`StatisticsContext`] walk (the same path the + /// optimizer and EXPLAIN use), so there is one traversal implementation. + /// Provider extensions are preserved; see [`StatisticsContext`] for how they + /// propagate up the tree. /// - /// If no providers are registered, falls back to the plan's built-in - /// `partition_statistics(None)` with no overhead. + /// Each provider's `child_stats` come from its own + /// [`StatisticsProvider::child_stats_requests`], so a provider can read child + /// statistics for an existing operator without that operator declaring + /// anything. + #[deprecated( + since = "55.0.0", + note = "use `StatisticsContext::new_with_registry(registry).compute_extended(plan, &StatisticsArgs::new())`" + )] pub fn compute(&self, plan: &dyn ExecutionPlan) -> Result { - // Fast path: no providers registered, skip the walk entirely - if self.providers.is_empty() { - let base = StatisticsContext::new().compute(plan, &StatisticsArgs::new())?; - return Ok(ExtendedStatistics::new_arc(base)); - } - - let children = plan.children(); - - // For leaf nodes, try providers with empty child stats. - // For non-leaf nodes, recursively compute enhanced child stats first. - let child_stats: Vec = if children.is_empty() { - Vec::new() - } else { - children - .iter() - .map(|child| self.compute(child.as_ref())) - .collect::>>()? - }; - - for provider in &self.providers { - match provider.compute_statistics(plan, &child_stats)? { - StatisticsResult::Computed(stats) => return Ok(stats), - StatisticsResult::Delegate => continue, - } - } - // Fallback: use plan's built-in stats - let base = StatisticsContext::new().compute(plan, &StatisticsArgs::new())?; - Ok(ExtendedStatistics::new_arc(base)) + Ok(StatisticsContext::new_with_registry(self.clone()) + .compute_extended(plan, &StatisticsArgs::new())? + .as_ref() + .clone()) } /// Compute statistics and return only the base Statistics (no extensions). /// /// Convenience method for callers that don't need extensions. + #[deprecated( + since = "55.0.0", + note = "use `StatisticsContext::new_with_registry(registry).compute(plan, &StatisticsArgs::new())`" + )] pub fn compute_base(&self, plan: &dyn ExecutionPlan) -> Result { - Ok(self.compute(plan)?.base().clone()) + Ok(StatisticsContext::new_with_registry(self.clone()) + .compute(plan, &StatisticsArgs::new())? + .as_ref() + .clone()) } } @@ -471,7 +537,7 @@ pub fn ndv_after_selectivity( /// Rescale `total_byte_size` proportionally after overriding `num_rows`. /// /// When a provider replaces `num_rows` but keeps the rest of the stats from -/// `partition_statistics`, the original `total_byte_size` becomes inconsistent. +/// `statistics_from_inputs`, the original `total_byte_size` becomes inconsistent. /// This function adjusts it by the ratio `new_rows / old_rows`, preserving the /// average bytes-per-row from the original estimate. fn rescale_byte_size(stats: &mut Statistics, new_num_rows: Precision) { @@ -496,27 +562,29 @@ fn rescale_byte_size(stats: &mut Statistics, new_num_rows: Precision) { }; } -/// Fetches base statistics from the operator's built-in `partition_statistics`, -/// overrides `num_rows` with the registry-computed estimate, and rescales -/// `total_byte_size` proportionally. -/// -/// Used by providers that compute a better row count but cannot yet propagate -/// column-level stats (NDV, min/max) through the operator — pending #20184. +/// Overrides the operator's built-in output statistics (computed from the +/// pre-resolved `child_stats`) with `num_rows`, rescaling `total_byte_size` +/// proportionally. Used by providers that refine only the row count. fn computed_with_row_count( plan: &dyn ExecutionPlan, num_rows: Precision, + child_stats: &[ExtendedStatistics], ) -> Result { + let child_base: Vec> = child_stats + .iter() + .map(|c| Arc::clone(c.base_arc())) + .collect(); let mut base = Arc::unwrap_or_clone( - StatisticsContext::new().compute(plan, &StatisticsArgs::new())?, + plan.statistics_from_inputs(&child_base, &StatisticsArgs::new())?, ); rescale_byte_size(&mut base, num_rows); Ok(StatisticsResult::Computed(ExtendedStatistics::new(base))) } -/// Statistics provider for [`FilterExec`](crate::filter::FilterExec) that uses +/// Statistics provider for [`FilterExec`] that uses /// pre-computed enhanced child statistics from the registry walk. /// -/// Unlike the default provider (which calls `partition_statistics` and gets raw +/// Unlike the built-in fallback (which calls `statistics_from_inputs` and gets raw /// child stats), this provider receives enhanced child stats that may include /// NDV overrides injected at the scan level. It applies the same selectivity /// estimation logic as `FilterExec::statistics_helper`, then additionally @@ -526,13 +594,15 @@ fn computed_with_row_count( pub struct FilterStatisticsProvider; impl StatisticsProvider for FilterStatisticsProvider { + fn matches(&self, plan: &dyn ExecutionPlan) -> bool { + plan.downcast_ref::().is_some() + } + fn compute_statistics( &self, plan: &dyn ExecutionPlan, child_stats: &[ExtendedStatistics], ) -> Result { - use crate::filter::FilterExec; - let Some(filter) = plan.downcast_ref::() else { return Ok(StatisticsResult::Delegate); }; @@ -572,7 +642,7 @@ impl StatisticsProvider for FilterStatisticsProvider { } } -/// Statistics provider for [`ProjectionExec`](crate::projection::ProjectionExec) +/// Statistics provider for [`ProjectionExec`] /// that uses pre-computed enhanced child statistics from the registry walk. /// /// Maps enhanced child column statistics to output columns based on the @@ -582,13 +652,15 @@ impl StatisticsProvider for FilterStatisticsProvider { pub struct ProjectionStatisticsProvider; impl StatisticsProvider for ProjectionStatisticsProvider { + fn matches(&self, plan: &dyn ExecutionPlan) -> bool { + plan.downcast_ref::().is_some() + } + fn compute_statistics( &self, plan: &dyn ExecutionPlan, child_stats: &[ExtendedStatistics], ) -> Result { - use crate::projection::ProjectionExec; - let Some(proj) = plan.downcast_ref::() else { return Ok(StatisticsResult::Delegate); }; @@ -608,23 +680,26 @@ impl StatisticsProvider for ProjectionStatisticsProvider { } /// Statistics provider for single-input operators with -/// [`CardinalityEffect::Equal`](crate::execution_plan::CardinalityEffect::Equal). +/// [`CardinalityEffect::Equal`]. /// /// These operators (Sort, Repartition, CoalescePartitions, etc.) don't /// transform statistics, so we pass through the enhanced child stats directly. -/// This avoids the fallback calling `partition_statistics(None)` which would +/// This avoids the fallback calling `statistics_from_inputs` (overall) which would /// trigger a redundant internal recursion with raw (non-enhanced) stats. #[derive(Debug, Default)] pub struct PassthroughStatisticsProvider; impl StatisticsProvider for PassthroughStatisticsProvider { + fn matches(&self, plan: &dyn ExecutionPlan) -> bool { + plan.children().len() == 1 + && matches!(plan.cardinality_effect(), CardinalityEffect::Equal) + } + fn compute_statistics( &self, plan: &dyn ExecutionPlan, child_stats: &[ExtendedStatistics], ) -> Result { - use crate::execution_plan::CardinalityEffect; - if child_stats.len() != 1 || !matches!(plan.cardinality_effect(), CardinalityEffect::Equal) { @@ -644,7 +719,7 @@ impl StatisticsProvider for PassthroughStatisticsProvider { } } -/// Statistics provider for [`AggregateExec`](crate::aggregates::AggregateExec) +/// Statistics provider for [`AggregateExec`] /// that estimates output cardinality from the NDV of GROUP BY columns. /// /// For each GROUP BY column, looks up `distinct_count` from the enhanced @@ -654,7 +729,7 @@ impl StatisticsProvider for PassthroughStatisticsProvider { /// produce overestimates. /// /// For GROUPING SETS / CUBE / ROLLUP, delegates to the built-in -/// `partition_statistics`, which handles per-set NDV estimation correctly. +/// `statistics_from_inputs`, which handles per-set NDV estimation correctly. /// /// Delegates when: /// - The plan is not an `AggregateExec` @@ -666,16 +741,17 @@ impl StatisticsProvider for PassthroughStatisticsProvider { pub struct AggregateStatisticsProvider; impl StatisticsProvider for AggregateStatisticsProvider { + fn matches(&self, plan: &dyn ExecutionPlan) -> bool { + plan.downcast_ref::().is_some() + } + fn compute_statistics( &self, plan: &dyn ExecutionPlan, child_stats: &[ExtendedStatistics], ) -> Result { - use crate::aggregates::AggregateExec; use datafusion_physical_expr::expressions::Column; - use crate::aggregates::AggregateMode; - let Some(agg) = plan.downcast_ref::() else { return Ok(StatisticsResult::Delegate); }; @@ -732,7 +808,7 @@ impl StatisticsProvider for AggregateStatisticsProvider { let num_rows = Precision::Inexact(estimate); - computed_with_row_count(plan, num_rows) + computed_with_row_count(plan, num_rows, child_stats) } } @@ -758,12 +834,17 @@ impl StatisticsProvider for AggregateStatisticsProvider { pub struct JoinStatisticsProvider; impl StatisticsProvider for JoinStatisticsProvider { + fn matches(&self, plan: &dyn ExecutionPlan) -> bool { + plan.downcast_ref::().is_some() + || plan.downcast_ref::().is_some() + || plan.downcast_ref::().is_some() + } + fn compute_statistics( &self, plan: &dyn ExecutionPlan, child_stats: &[ExtendedStatistics], ) -> Result { - use crate::joins::{CrossJoinExec, HashJoinExec, SortMergeJoinExec}; use datafusion_common::JoinType; use datafusion_physical_expr::expressions::Column; @@ -780,8 +861,6 @@ impl StatisticsProvider for JoinStatisticsProvider { return Ok(StatisticsResult::Delegate); }; - use crate::joins::JoinOnRef; - /// Estimate equi-join output using NDV of join key columns: /// left_rows * right_rows / product(max(left_ndv_i, right_ndv_i)) /// Falls back to Cartesian product if any key lacks NDV on both sides. @@ -867,12 +946,12 @@ impl StatisticsProvider for JoinStatisticsProvider { Precision::Inexact(estimated) }; - computed_with_row_count(plan, num_rows) + computed_with_row_count(plan, num_rows, child_stats) } } -/// Statistics provider for [`LocalLimitExec`](crate::limit::LocalLimitExec) and -/// [`GlobalLimitExec`](crate::limit::GlobalLimitExec). +/// Statistics provider for [`LocalLimitExec`] and +/// [`GlobalLimitExec`]. /// /// Caps output row count at the limit value, accounting for any leading skip offset /// in `GlobalLimitExec`. @@ -880,13 +959,16 @@ impl StatisticsProvider for JoinStatisticsProvider { pub struct LimitStatisticsProvider; impl StatisticsProvider for LimitStatisticsProvider { + fn matches(&self, plan: &dyn ExecutionPlan) -> bool { + plan.downcast_ref::().is_some() + || plan.downcast_ref::().is_some() + } + fn compute_statistics( &self, plan: &dyn ExecutionPlan, child_stats: &[ExtendedStatistics], ) -> Result { - use crate::limit::{GlobalLimitExec, LocalLimitExec}; - if child_stats.is_empty() { return Ok(StatisticsResult::Delegate); } @@ -917,24 +999,26 @@ impl StatisticsProvider for LimitStatisticsProvider { }, }; - computed_with_row_count(plan, num_rows) + computed_with_row_count(plan, num_rows, child_stats) } } -/// Statistics provider for [`UnionExec`](crate::union::UnionExec). +/// Statistics provider for [`UnionExec`]. /// /// Sums row counts across all inputs. #[derive(Debug, Default)] pub struct UnionStatisticsProvider; impl StatisticsProvider for UnionStatisticsProvider { + fn matches(&self, plan: &dyn ExecutionPlan) -> bool { + plan.downcast_ref::().is_some() + } + fn compute_statistics( &self, plan: &dyn ExecutionPlan, child_stats: &[ExtendedStatistics], ) -> Result { - use crate::union::UnionExec; - if plan.downcast_ref::().is_none() { return Ok(StatisticsResult::Delegate); } @@ -956,7 +1040,7 @@ impl StatisticsProvider for UnionStatisticsProvider { }, )?; - computed_with_row_count(plan, total) + computed_with_row_count(plan, total, child_stats) } } @@ -964,6 +1048,8 @@ type ProviderFn = dyn Fn(&dyn ExecutionPlan, &[ExtendedStatistics]) -> Result bool + Send + Sync; + /// A [`StatisticsProvider`] backed by a user-supplied closure. /// /// Useful for injecting custom statistics in tests or for cardinality feedback @@ -978,18 +1064,18 @@ type ProviderFn = dyn Fn(&dyn ExecutionPlan, &[ExtendedStatistics]) -> Result().is_some() { +/// let provider = ClosureStatisticsProvider::with_matches( +/// |plan| plan.downcast_ref::().is_some(), +/// |plan, child_stats| { /// Ok(StatisticsResult::Computed(ExtendedStatistics::from(Statistics { /// num_rows: Precision::Inexact(42), /// ..Statistics::new_unknown(plan.schema().as_ref()) /// }))) -/// } else { -/// Ok(StatisticsResult::Delegate) -/// } -/// }); +/// }, +/// ); /// ``` pub struct ClosureStatisticsProvider { + matches_fn: Box, f: Box, } @@ -1001,7 +1087,24 @@ impl ClosureStatisticsProvider { + Sync + 'static, ) -> Self { - Self { f: Box::new(f) } + Self { + matches_fn: Box::new(|_| true), + f: Box::new(f), + } + } + + /// Like [`Self::new`] but applies only where `matches` returns true. + pub fn with_matches( + matches: impl Fn(&dyn ExecutionPlan) -> bool + Send + Sync + 'static, + f: impl Fn(&dyn ExecutionPlan, &[ExtendedStatistics]) -> Result + + Send + + Sync + + 'static, + ) -> Self { + Self { + matches_fn: Box::new(matches), + f: Box::new(f), + } } } @@ -1012,6 +1115,10 @@ impl Debug for ClosureStatisticsProvider { } impl StatisticsProvider for ClosureStatisticsProvider { + fn matches(&self, plan: &dyn ExecutionPlan) -> bool { + (self.matches_fn)(plan) + } + fn compute_statistics( &self, plan: &dyn ExecutionPlan, @@ -1025,12 +1132,16 @@ impl StatisticsProvider for ClosureStatisticsProvider { mod tests { use super::*; use crate::filter::FilterExec; + use crate::joins::JoinOn; use crate::projection::ProjectionExec; - use crate::statistics::StatisticsArgs; - use crate::{DisplayAs, DisplayFormatType, PlanProperties}; + use crate::statistics::{ChildStats, StatisticsArgs, StatisticsContext}; + use crate::{ + DisplayAs, DisplayFormatType, PlanProperties, SendableRecordBatchStream, + }; use arrow::datatypes::{DataType, Field, Schema}; use datafusion_common::stats::Precision; - use datafusion_common::{ColumnStatistics, ScalarValue}; + use datafusion_common::{ColumnStatistics, ScalarValue, internal_err}; + use datafusion_execution::TaskContext; use datafusion_expr::Operator; use datafusion_physical_expr::PhysicalExpr; use datafusion_physical_expr::expressions::{BinaryExpr, Column, Literal, col, lit}; @@ -1039,6 +1150,17 @@ mod tests { use crate::execution_plan::{Boundedness, EmissionType}; + /// Compute statistics via [`StatisticsContext`] (replaces the deprecated + /// `compute`/`compute_base`). + fn compute( + registry: &StatisticsRegistry, + plan: &dyn ExecutionPlan, + ) -> Result { + Ok((*StatisticsContext::new_with_registry(registry.clone()) + .compute_extended(plan, &StatisticsArgs::new())?) + .clone()) + } + fn make_schema() -> Arc { Arc::new(Schema::new(vec![ Field::new("a", DataType::Int32, false), @@ -1120,8 +1242,8 @@ mod tests { fn execute( &self, _partition: usize, - _context: Arc, - ) -> Result { + _context: Arc, + ) -> Result { unimplemented!() } @@ -1143,14 +1265,33 @@ mod tests { #[test] fn test_default_provider() -> Result<()> { - let engine = StatisticsRegistry::new(); + let registry = StatisticsRegistry::new(); let source = make_source(1000); - let stats = engine.compute(source.as_ref())?; + let stats = compute(®istry, source.as_ref())?; assert!(matches!(stats.base.num_rows, Precision::Exact(1000))); Ok(()) } + #[test] + fn test_compute_preserves_provider_extensions() -> Result<()> { + #[derive(Debug, Clone, PartialEq)] + struct Sketch(u32); + + let source = make_source(1000); + let provider = ClosureStatisticsProvider::new(|plan, _child_stats| { + let mut ext = + ExtendedStatistics::new(Statistics::new_unknown(&plan.schema())); + ext.set_extension(Sketch(42)); + Ok(StatisticsResult::Computed(ext)) + }); + let registry = StatisticsRegistry::with_providers(vec![Arc::new(provider)]); + + let stats = compute(®istry, source.as_ref())?; + assert_eq!(stats.get_extension::(), Some(&Sketch(42))); + Ok(()) + } + #[test] fn test_custom_chain_configuration() -> Result<()> { let source = make_source(1000); @@ -1159,10 +1300,10 @@ mod tests { let custom_only = StatisticsRegistry::with_providers(vec![Arc::new(CustomStatisticsProvider)]); // CustomStatisticsProvider only handles CustomExec, delegates for others - // With no default provider, filter returns fallback statistics + // With no provider handling FilterExec, it returns the built-in fallback statistics let filter: Arc = Arc::new(FilterExec::try_new(lit(true), Arc::clone(&source))?); - let stats = custom_only.compute(filter.as_ref())?; + let stats = compute(&custom_only, filter.as_ref())?; // Falls back to plan.statistics() since no provider handles it assert!(stats.base.num_rows.get_value().is_some()); @@ -1173,7 +1314,7 @@ mod tests { }) as Arc]); // OverrideFilterProvider handles filters, built-in fallback handles the rest - let stats = with_override.compute(filter.as_ref())?; + let stats = compute(&with_override, filter.as_ref())?; assert!(matches!(stats.base.num_rows, Precision::Inexact(250))); // Verify chain inspection @@ -1222,8 +1363,8 @@ mod tests { fn execute( &self, _partition: usize, - _context: Arc, - ) -> Result { + _context: Arc, + ) -> Result { unimplemented!() } } @@ -1247,17 +1388,301 @@ mod tests { #[test] fn test_custom_provider_for_custom_exec() -> Result<()> { - let mut engine = StatisticsRegistry::new(); - engine.register(Arc::new(CustomStatisticsProvider)); + let mut registry = StatisticsRegistry::new(); + registry.register(Arc::new(CustomStatisticsProvider)); let source = make_source(1000); let custom: Arc = Arc::new(CustomExec { input: source }); - let stats = engine.compute(custom.as_ref())?; + let stats = compute(®istry, custom.as_ref())?; assert!(matches!(stats.base.num_rows, Precision::Exact(1000))); Ok(()) } + /// A provider that overrides `child_stats_requests` to need no children. + #[derive(Debug)] + struct NoChildStatsProvider; + + impl StatisticsProvider for NoChildStatsProvider { + fn child_stats_requests( + &self, + plan: &dyn ExecutionPlan, + _partition: Option, + ) -> Vec { + plan.children().iter().map(|_| ChildStats::Skip).collect() + } + + fn compute_statistics( + &self, + plan: &dyn ExecutionPlan, + child_stats: &[ExtendedStatistics], + ) -> Result { + if plan.downcast_ref::().is_some() { + Ok(StatisticsResult::Computed(child_stats[0].clone())) + } else { + Ok(StatisticsResult::Delegate) + } + } + } + + #[test] + fn test_provider_opts_out_of_child_stats() -> Result<()> { + let mut registry = StatisticsRegistry::new(); + registry.register(Arc::new(NoChildStatsProvider)); + + let source = make_source(1000); + let custom: Arc = Arc::new(CustomExec { input: source }); + + // The provider requested no children, so it sees the unknown placeholder + // rather than the source's row count. + let stats = compute(®istry, custom.as_ref())?; + assert!(stats.base.num_rows.get_value().is_none()); + Ok(()) + } + + /// Errors when its statistics are computed, so resolving it is detectable. + #[derive(Debug)] + struct ErroringExec { + input: Arc, + } + + impl DisplayAs for ErroringExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "ErroringExec") + } + } + + impl ExecutionPlan for ErroringExec { + fn name(&self) -> &str { + "ErroringExec" + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + Ok(Arc::new(ErroringExec { + input: Arc::clone(&children[0]), + })) + } + + fn properties(&self) -> &Arc { + self.input.properties() + } + + fn statistics_from_inputs( + &self, + _input_stats: &[Arc], + _args: &StatisticsArgs, + ) -> Result> { + internal_err!("child statistics should not be computed") + } + + fn execute( + &self, + _partition: usize, + _context: Arc, + ) -> Result { + unimplemented!() + } + } + + /// Computes a fixed row count for `ProjectionExec`, requesting no children. + #[derive(Debug)] + struct SkipComputeProvider; + + impl StatisticsProvider for SkipComputeProvider { + fn matches(&self, plan: &dyn ExecutionPlan) -> bool { + plan.downcast_ref::().is_some() + } + + fn child_stats_requests( + &self, + plan: &dyn ExecutionPlan, + _partition: Option, + ) -> Vec { + plan.children().iter().map(|_| ChildStats::Skip).collect() + } + + fn compute_statistics( + &self, + plan: &dyn ExecutionPlan, + _child_stats: &[ExtendedStatistics], + ) -> Result { + if plan.downcast_ref::().is_none() { + return Ok(StatisticsResult::Delegate); + } + let mut stats = Statistics::new_unknown(plan.schema().as_ref()); + stats.num_rows = Precision::Exact(7); + Ok(StatisticsResult::Computed(ExtendedStatistics::from(stats))) + } + } + + /// Never matches, but if consulted it resolves the child and returns its row + /// count, leaking an observable (non-error) value. + #[derive(Debug)] + struct NonMatchingLeakProvider; + + impl StatisticsProvider for NonMatchingLeakProvider { + fn matches(&self, _plan: &dyn ExecutionPlan) -> bool { + false + } + + fn compute_statistics( + &self, + plan: &dyn ExecutionPlan, + child_stats: &[ExtendedStatistics], + ) -> Result { + let mut stats = Statistics::new_unknown(plan.schema().as_ref()); + if let Some(child) = child_stats.first() { + stats.num_rows = child.base.num_rows; + } + Ok(StatisticsResult::Computed(ExtendedStatistics::from(stats))) + } + } + + /// Default `matches()` (always true) and `child_stats_requests` (`At(None)`), + /// but always delegates. + #[derive(Debug)] + struct DelegatingProvider; + + impl StatisticsProvider for DelegatingProvider { + fn compute_statistics( + &self, + _plan: &dyn ExecutionPlan, + _child_stats: &[ExtendedStatistics], + ) -> Result { + Ok(StatisticsResult::Delegate) + } + } + + /// A `ProjectionExec` (which requests `At`) over a child that errors when its + /// statistics are computed. + fn projection_over_erroring_child() -> Result> { + let schema = make_schema(); + let child: Arc = Arc::new(ErroringExec { + input: make_source(1000), + }); + Ok(Arc::new(ProjectionExec::try_new( + vec![(col("a", &schema)?, "a".to_string())], + child, + )?)) + } + + /// A `ProjectionExec` (which requests `At`) over a source with `num_rows` rows. + fn projection_over_source(num_rows: usize) -> Result> { + let schema = make_schema(); + Ok(Arc::new(ProjectionExec::try_new( + vec![(col("a", &schema)?, "a".to_string())], + make_source(num_rows), + )?)) + } + + #[test] + fn test_provider_override_skips_operator_child_walk() -> Result<()> { + // A matching provider requests Skip, so the erroring child is never + // resolved and its error never surfaces. + let parent = projection_over_erroring_child()?; + let mut registry = StatisticsRegistry::new(); + registry.register(Arc::new(SkipComputeProvider)); + + let stats = compute(®istry, parent.as_ref())?; + assert!(matches!(stats.base.num_rows, Precision::Exact(7))); + Ok(()) + } + + #[test] + fn test_matches_gate_skips_non_matching_provider() -> Result<()> { + // A non-matching provider precedes the handling one. The gate must skip it + // before it is consulted, otherwise it leaks the child row count (999) + // instead of the later provider's 7. + let parent = projection_over_source(999)?; + let registry = StatisticsRegistry::with_providers(vec![ + Arc::new(NonMatchingLeakProvider), + Arc::new(SkipComputeProvider), + ]); + + let stats = compute(®istry, parent.as_ref())?; + assert!(matches!(stats.base.num_rows, Precision::Exact(7))); + Ok(()) + } + + #[test] + fn test_delegating_provider_child_error_does_not_block_later_provider() -> Result<()> + { + // A delegating provider that keeps the default matcher resolves the + // erroring child speculatively; that failure must not stop the later + // child-skipping provider from succeeding. + let parent = projection_over_erroring_child()?; + let registry = StatisticsRegistry::with_providers(vec![ + Arc::new(DelegatingProvider), + Arc::new(SkipComputeProvider), + ]); + + let stats = compute(®istry, parent.as_ref())?; + assert!(matches!(stats.base.num_rows, Precision::Exact(7))); + Ok(()) + } + + /// Returns the wrong number of child stat requests. + #[derive(Debug)] + struct MalformedRequestProvider; + + impl StatisticsProvider for MalformedRequestProvider { + fn child_stats_requests( + &self, + _plan: &dyn ExecutionPlan, + _partition: Option, + ) -> Vec { + Vec::new() + } + } + + /// Requests a non-existent partition of its child. + #[derive(Debug)] + struct InvalidPartitionProvider; + + impl StatisticsProvider for InvalidPartitionProvider { + fn child_stats_requests( + &self, + plan: &dyn ExecutionPlan, + _partition: Option, + ) -> Vec { + plan.children() + .iter() + .map(|_| ChildStats::At(Some(usize::MAX))) + .collect() + } + } + + #[test] + fn test_malformed_child_stats_requests_fails_fast() -> Result<()> { + // A wrong-length request set is a provider bug: it must error, not be + // skipped in favor of the later provider. + let parent = projection_over_source(10)?; + let registry = StatisticsRegistry::with_providers(vec![ + Arc::new(MalformedRequestProvider), + Arc::new(SkipComputeProvider), + ]); + assert!(compute(®istry, parent.as_ref()).is_err()); + Ok(()) + } + + #[test] + fn test_invalid_partition_request_fails_fast() -> Result<()> { + let parent = projection_over_source(10)?; + let registry = StatisticsRegistry::with_providers(vec![ + Arc::new(InvalidPartitionProvider), + Arc::new(SkipComputeProvider), + ]); + assert!(compute(®istry, parent.as_ref()).is_err()); + Ok(()) + } + #[derive(Debug)] struct OverrideFilterProvider { fixed_selectivity: f64, @@ -1293,8 +1718,8 @@ mod tests { #[test] fn test_override_builtin_operator() -> Result<()> { - let mut engine = StatisticsRegistry::new(); - engine.register(Arc::new(OverrideFilterProvider { + let mut registry = StatisticsRegistry::new(); + registry.register(Arc::new(OverrideFilterProvider { fixed_selectivity: 0.1, })); @@ -1302,20 +1727,20 @@ mod tests { let filter: Arc = Arc::new(FilterExec::try_new(lit(true), source)?); - let stats = engine.compute(filter.as_ref())?; + let stats = compute(®istry, filter.as_ref())?; assert!(matches!(stats.base.num_rows, Precision::Inexact(100))); Ok(()) } #[test] fn test_filter_statistics_propagation() -> Result<()> { - let engine = StatisticsRegistry::new(); + let registry = StatisticsRegistry::new(); let source = make_source(1000); let predicate = lit(true); let filter: Arc = Arc::new(FilterExec::try_new(predicate, source)?); - let stats = engine.compute(filter.as_ref())?; + let stats = compute(®istry, filter.as_ref())?; assert!(stats.base.num_rows.get_value().unwrap_or(&0) <= &1000); Ok(()) } @@ -1363,11 +1788,9 @@ mod tests { let filter: Arc = Arc::new(FilterExec::try_new(predicate, source)?); - let registry = StatisticsRegistry::with_providers(vec![ - Arc::new(FilterStatisticsProvider), - Arc::new(DefaultStatisticsProvider), - ]); - let stats = registry.compute(filter.as_ref())?; + let registry = + StatisticsRegistry::with_providers(vec![Arc::new(FilterStatisticsProvider)]); + let stats = compute(®istry, filter.as_ref())?; let output_ndv_a = stats.base.column_statistics[0] .distinct_count @@ -1397,7 +1820,7 @@ mod tests { #[test] fn test_projection_statistics_propagation() -> Result<()> { - let engine = StatisticsRegistry::new(); + let registry = StatisticsRegistry::new(); let source = make_source(1000); let schema = make_schema(); let proj: Arc = Arc::new(ProjectionExec::try_new( @@ -1405,7 +1828,7 @@ mod tests { source, )?); - let stats = engine.compute(proj.as_ref())?; + let stats = compute(®istry, proj.as_ref())?; assert!(matches!(stats.base.num_rows, Precision::Exact(1000))); Ok(()) } @@ -1414,12 +1837,12 @@ mod tests { fn test_passthrough_statistics_propagation() -> Result<()> { use crate::coalesce_partitions::CoalescePartitionsExec; - let engine = StatisticsRegistry::new(); + let registry = StatisticsRegistry::new(); let source = make_source(1000); let coalesce: Arc = Arc::new(CoalescePartitionsExec::new(source)); - let stats = engine.compute(coalesce.as_ref())?; + let stats = compute(®istry, coalesce.as_ref())?; // PassthroughStatisticsProvider should propagate child row count unchanged assert_eq!(stats.base.num_rows, Precision::Exact(1000)); Ok(()) @@ -1427,11 +1850,11 @@ mod tests { #[test] fn test_chain_priority() -> Result<()> { - let mut engine = StatisticsRegistry::new(); - engine.register(Arc::new(OverrideFilterProvider { + let mut registry = StatisticsRegistry::new(); + registry.register(Arc::new(OverrideFilterProvider { fixed_selectivity: 0.5, })); - engine.register(Arc::new(CustomStatisticsProvider)); + registry.register(Arc::new(CustomStatisticsProvider)); let source = make_source(1000); @@ -1439,13 +1862,13 @@ mod tests { let custom: Arc = Arc::new(CustomExec { input: Arc::clone(&source), }); - let stats = engine.compute(custom.as_ref())?; + let stats = compute(®istry, custom.as_ref())?; assert!(matches!(stats.base.num_rows, Precision::Exact(1000))); // FilterExec: CustomStatisticsProvider delegates, OverrideFilterProvider handles let filter: Arc = Arc::new(FilterExec::try_new(lit(true), source)?); - let stats = engine.compute(filter.as_ref())?; + let stats = compute(®istry, filter.as_ref())?; assert!(matches!(stats.base.num_rows, Precision::Inexact(500))); Ok(()) @@ -1553,11 +1976,10 @@ mod tests { )]); let agg = make_aggregate(source, group_by)?; - let registry = StatisticsRegistry::with_providers(vec![ - Arc::new(AggregateStatisticsProvider), - Arc::new(DefaultStatisticsProvider), - ]); - let stats = registry.compute(agg.as_ref())?; + let registry = StatisticsRegistry::with_providers(vec![Arc::new( + AggregateStatisticsProvider, + )]); + let stats = compute(®istry, agg.as_ref())?; assert_eq!(stats.base.num_rows, Precision::Inexact(10)); Ok(()) } @@ -1571,11 +1993,10 @@ mod tests { ]); let agg = make_aggregate(source, group_by)?; - let registry = StatisticsRegistry::with_providers(vec![ - Arc::new(AggregateStatisticsProvider), - Arc::new(DefaultStatisticsProvider), - ]); - let stats = registry.compute(agg.as_ref())?; + let registry = StatisticsRegistry::with_providers(vec![Arc::new( + AggregateStatisticsProvider, + )]); + let stats = compute(®istry, agg.as_ref())?; // 10 * 5 = 50 assert_eq!(stats.base.num_rows, Precision::Inexact(50)); Ok(()) @@ -1591,11 +2012,10 @@ mod tests { ]); let agg = make_aggregate(source, group_by)?; - let registry = StatisticsRegistry::with_providers(vec![ - Arc::new(AggregateStatisticsProvider), - Arc::new(DefaultStatisticsProvider), - ]); - let stats = registry.compute(agg.as_ref())?; + let registry = StatisticsRegistry::with_providers(vec![Arc::new( + AggregateStatisticsProvider, + )]); + let stats = compute(®istry, agg.as_ref())?; assert_eq!(stats.base.num_rows, Precision::Inexact(500)); Ok(()) } @@ -1610,12 +2030,11 @@ mod tests { )]); let agg = make_aggregate(source, group_by)?; - let registry = StatisticsRegistry::with_providers(vec![ - Arc::new(AggregateStatisticsProvider), - Arc::new(DefaultStatisticsProvider), - ]); - let stats = registry.compute(agg.as_ref())?; - // Delegates to DefaultStatisticsProvider, which calls partition_statistics + let registry = StatisticsRegistry::with_providers(vec![Arc::new( + AggregateStatisticsProvider, + )]); + let stats = compute(®istry, agg.as_ref())?; + // Delegates, falling back to the operator's built-in statistics_from_inputs assert!( stats.base.num_rows.get_value().is_some() || matches!(stats.base.num_rows, Precision::Absent) @@ -1635,11 +2054,10 @@ mod tests { let group_by = PhysicalGroupBy::new_single(vec![(expr, "sum".to_string())]); let agg = make_aggregate(source, group_by)?; - let registry = StatisticsRegistry::with_providers(vec![ - Arc::new(AggregateStatisticsProvider), - Arc::new(DefaultStatisticsProvider), - ]); - let stats = registry.compute(agg.as_ref())?; + let registry = StatisticsRegistry::with_providers(vec![Arc::new( + AggregateStatisticsProvider, + )]); + let stats = compute(®istry, agg.as_ref())?; // Should delegate (expression is not a Column) assert!( stats.base.num_rows.get_value().is_some() @@ -1676,13 +2094,12 @@ mod tests { ); let agg = make_aggregate(source, group_by)?; - let registry = StatisticsRegistry::with_providers(vec![ - Arc::new(AggregateStatisticsProvider), - Arc::new(DefaultStatisticsProvider), - ]); - let stats = registry.compute(agg.as_ref())?; - // Multiple grouping sets: provider delegates to DefaultStatisticsProvider, - // which calls the built-in partition_statistics for correct per-set + let registry = StatisticsRegistry::with_providers(vec![Arc::new( + AggregateStatisticsProvider, + )]); + let stats = compute(®istry, agg.as_ref())?; + // Multiple grouping sets: provider delegates, so the operator's built-in + // statistics_from_inputs computes the correct per-set // NDV estimation. The exact value depends on the built-in implementation. assert!( stats.base.num_rows.get_value().is_some() @@ -1709,12 +2126,11 @@ mod tests { source.schema(), )?); - let registry = StatisticsRegistry::with_providers(vec![ - Arc::new(AggregateStatisticsProvider), - Arc::new(DefaultStatisticsProvider), - ]); - let stats = registry.compute(agg.as_ref())?; - // Should fall through to DefaultStatisticsProvider (partition_statistics). + let registry = StatisticsRegistry::with_providers(vec![Arc::new( + AggregateStatisticsProvider, + )]); + let stats = compute(®istry, agg.as_ref())?; + // Should fall through to the operator's built-in statistics_from_inputs. // The exact value depends on the built-in implementation. assert!( stats.base.num_rows.get_value().is_some() @@ -1757,7 +2173,7 @@ mod tests { right: Arc, ) -> Result> { let _schema = make_schema(); - let on: crate::joins::JoinOn = vec![( + let on: JoinOn = vec![( Arc::new(Column::new("a", 0)) as Arc, Arc::new(Column::new("a", 0)) as Arc, )]; @@ -1782,11 +2198,9 @@ mod tests { let right = make_source_with_ndv_2col(500, Some(50)); let join = make_hash_join(left, right)?; - let registry = StatisticsRegistry::with_providers(vec![ - Arc::new(JoinStatisticsProvider), - Arc::new(DefaultStatisticsProvider), - ]); - let stats = registry.compute(join.as_ref())?; + let registry = + StatisticsRegistry::with_providers(vec![Arc::new(JoinStatisticsProvider)]); + let stats = compute(®istry, join.as_ref())?; assert_eq!(stats.base.num_rows, Precision::Inexact(5000)); Ok(()) } @@ -1820,7 +2234,7 @@ mod tests { let right = make_source_ndv_b(500, 25); // Join on column "b" (index 1) - let on: crate::joins::JoinOn = vec![( + let on: JoinOn = vec![( Arc::new(Column::new("b", 1)) as Arc, Arc::new(Column::new("b", 1)) as Arc, )]; @@ -1836,11 +2250,9 @@ mod tests { false, )?); - let registry = StatisticsRegistry::with_providers(vec![ - Arc::new(JoinStatisticsProvider), - Arc::new(DefaultStatisticsProvider), - ]); - let stats = registry.compute(join.as_ref())?; + let registry = + StatisticsRegistry::with_providers(vec![Arc::new(JoinStatisticsProvider)]); + let stats = compute(®istry, join.as_ref())?; assert_eq!(stats.base.num_rows, Precision::Inexact(10_000)); Ok(()) } @@ -1876,7 +2288,7 @@ mod tests { let left = make_source_2ndv(1000, 100, 20); let right = make_source_2ndv(500, 50, 10); - let on: crate::joins::JoinOn = vec![ + let on: JoinOn = vec![ ( Arc::new(Column::new("a", 0)) as Arc, Arc::new(Column::new("a", 0)) as Arc, @@ -1898,11 +2310,9 @@ mod tests { false, )?); - let registry = StatisticsRegistry::with_providers(vec![ - Arc::new(JoinStatisticsProvider), - Arc::new(DefaultStatisticsProvider), - ]); - let stats = registry.compute(join.as_ref())?; + let registry = + StatisticsRegistry::with_providers(vec![Arc::new(JoinStatisticsProvider)]); + let stats = compute(®istry, join.as_ref())?; assert_eq!(stats.base.num_rows, Precision::Inexact(250)); Ok(()) } @@ -1914,11 +2324,9 @@ mod tests { let right = make_source_with_ndv_2col(200, None); let join = make_hash_join(left, right)?; - let registry = StatisticsRegistry::with_providers(vec![ - Arc::new(JoinStatisticsProvider), - Arc::new(DefaultStatisticsProvider), - ]); - let stats = registry.compute(join.as_ref())?; + let registry = + StatisticsRegistry::with_providers(vec![Arc::new(JoinStatisticsProvider)]); + let stats = compute(®istry, join.as_ref())?; assert_eq!(stats.base.num_rows, Precision::Inexact(20_000)); Ok(()) } @@ -1939,12 +2347,10 @@ mod tests { None, )?); - let registry = StatisticsRegistry::with_providers(vec![ - Arc::new(JoinStatisticsProvider), - Arc::new(DefaultStatisticsProvider), - ]); - let stats = registry.compute(join.as_ref())?; - // Provider delegates; result comes from built-in partition_statistics. + let registry = + StatisticsRegistry::with_providers(vec![Arc::new(JoinStatisticsProvider)]); + let stats = compute(®istry, join.as_ref())?; + // Provider delegates; result comes from built-in statistics_from_inputs. assert!( stats.base.num_rows.get_value().is_some() || matches!(stats.base.num_rows, Precision::Absent) @@ -1957,7 +2363,7 @@ mod tests { right: Arc, join_type: JoinType, ) -> Result> { - let on: crate::joins::JoinOn = vec![( + let on: JoinOn = vec![( Arc::new(Column::new("a", 0)) as Arc, Arc::new(Column::new("a", 0)) as Arc, )]; @@ -1984,11 +2390,9 @@ mod tests { let left = make_source_with_ndv_2col(left_rows, left_ndv); let right = make_source_with_ndv_2col(right_rows, right_ndv); let join = make_hash_join_typed(left, right, join_type)?; - let registry = StatisticsRegistry::with_providers(vec![ - Arc::new(JoinStatisticsProvider), - Arc::new(DefaultStatisticsProvider), - ]); - Ok(registry.compute(join.as_ref())?.base.num_rows) + let registry = + StatisticsRegistry::with_providers(vec![Arc::new(JoinStatisticsProvider)]); + Ok(compute(®istry, join.as_ref())?.base.num_rows) } #[test] @@ -2071,11 +2475,9 @@ mod tests { let right = make_source(200); let join: Arc = Arc::new(CrossJoinExec::new(left, right)); - let registry = StatisticsRegistry::with_providers(vec![ - Arc::new(JoinStatisticsProvider), - Arc::new(DefaultStatisticsProvider), - ]); - let stats = registry.compute(join.as_ref())?; + let registry = + StatisticsRegistry::with_providers(vec![Arc::new(JoinStatisticsProvider)]); + let stats = compute(®istry, join.as_ref())?; // Both inputs have Exact row counts -> result is also Exact assert_eq!(stats.base.num_rows, Precision::Exact(20_000)); Ok(()) @@ -2093,11 +2495,9 @@ mod tests { let source = make_source(1000); let limit: Arc = Arc::new(LocalLimitExec::new(source, 100)); - let registry = StatisticsRegistry::with_providers(vec![ - Arc::new(LimitStatisticsProvider), - Arc::new(DefaultStatisticsProvider), - ]); - let stats = registry.compute(limit.as_ref())?; + let registry = + StatisticsRegistry::with_providers(vec![Arc::new(LimitStatisticsProvider)]); + let stats = compute(®istry, limit.as_ref())?; assert_eq!(stats.base.num_rows, Precision::Exact(100)); Ok(()) } @@ -2108,11 +2508,9 @@ mod tests { let source = make_source(50); let limit: Arc = Arc::new(LocalLimitExec::new(source, 200)); - let registry = StatisticsRegistry::with_providers(vec![ - Arc::new(LimitStatisticsProvider), - Arc::new(DefaultStatisticsProvider), - ]); - let stats = registry.compute(limit.as_ref())?; + let registry = + StatisticsRegistry::with_providers(vec![Arc::new(LimitStatisticsProvider)]); + let stats = compute(®istry, limit.as_ref())?; assert_eq!(stats.base.num_rows, Precision::Exact(50)); Ok(()) } @@ -2124,11 +2522,9 @@ mod tests { let limit: Arc = Arc::new(GlobalLimitExec::new(source, 200, Some(100))); - let registry = StatisticsRegistry::with_providers(vec![ - Arc::new(LimitStatisticsProvider), - Arc::new(DefaultStatisticsProvider), - ]); - let stats = registry.compute(limit.as_ref())?; + let registry = + StatisticsRegistry::with_providers(vec![Arc::new(LimitStatisticsProvider)]); + let stats = compute(®istry, limit.as_ref())?; assert_eq!(stats.base.num_rows, Precision::Exact(100)); Ok(()) } @@ -2140,11 +2536,9 @@ mod tests { let limit: Arc = Arc::new(GlobalLimitExec::new(source, 200, Some(50))); - let registry = StatisticsRegistry::with_providers(vec![ - Arc::new(LimitStatisticsProvider), - Arc::new(DefaultStatisticsProvider), - ]); - let stats = registry.compute(limit.as_ref())?; + let registry = + StatisticsRegistry::with_providers(vec![Arc::new(LimitStatisticsProvider)]); + let stats = compute(®istry, limit.as_ref())?; assert_eq!(stats.base.num_rows, Precision::Exact(0)); Ok(()) } @@ -2156,11 +2550,9 @@ mod tests { let source = make_source_with_precision(Precision::Inexact(1000)); let limit: Arc = Arc::new(LocalLimitExec::new(source, 100)); - let registry = StatisticsRegistry::with_providers(vec![ - Arc::new(LimitStatisticsProvider), - Arc::new(DefaultStatisticsProvider), - ]); - let stats = registry.compute(limit.as_ref())?; + let registry = + StatisticsRegistry::with_providers(vec![Arc::new(LimitStatisticsProvider)]); + let stats = compute(®istry, limit.as_ref())?; assert_eq!(stats.base.num_rows, Precision::Inexact(100)); Ok(()) } @@ -2179,11 +2571,9 @@ mod tests { fn test_union_provider_sums_rows() -> Result<()> { let union = UnionExec::try_new(vec![make_source(300), make_source(700)])?; - let registry = StatisticsRegistry::with_providers(vec![ - Arc::new(UnionStatisticsProvider), - Arc::new(DefaultStatisticsProvider), - ]); - let stats = registry.compute(union.as_ref())?; + let registry = + StatisticsRegistry::with_providers(vec![Arc::new(UnionStatisticsProvider)]); + let stats = compute(®istry, union.as_ref())?; assert_eq!(stats.base.num_rows, Precision::Exact(1000)); Ok(()) } @@ -2196,11 +2586,9 @@ mod tests { make_source(300), ])?; - let registry = StatisticsRegistry::with_providers(vec![ - Arc::new(UnionStatisticsProvider), - Arc::new(DefaultStatisticsProvider), - ]); - let stats = registry.compute(union.as_ref())?; + let registry = + StatisticsRegistry::with_providers(vec![Arc::new(UnionStatisticsProvider)]); + let stats = compute(®istry, union.as_ref())?; assert_eq!(stats.base.num_rows, Precision::Exact(600)); Ok(()) } @@ -2213,11 +2601,9 @@ mod tests { make_source_with_precision(Precision::Absent), ])?; - let registry = StatisticsRegistry::with_providers(vec![ - Arc::new(UnionStatisticsProvider), - Arc::new(DefaultStatisticsProvider), - ]); - let stats = registry.compute(union.as_ref())?; + let registry = + StatisticsRegistry::with_providers(vec![Arc::new(UnionStatisticsProvider)]); + let stats = compute(®istry, union.as_ref())?; assert_eq!(stats.base.num_rows, Precision::Absent); Ok(()) } @@ -2243,15 +2629,12 @@ mod tests { } }); - let registry = StatisticsRegistry::with_providers(vec![ - Arc::new(provider), - Arc::new(DefaultStatisticsProvider), - ]); + let registry = StatisticsRegistry::with_providers(vec![Arc::new(provider)]); let source = make_source(1000); let filter: Arc = Arc::new(FilterExec::try_new(lit(true), source)?); - let stats = registry.compute(filter.as_ref())?; + let stats = compute(®istry, filter.as_ref())?; assert_eq!(stats.base.num_rows, Precision::Inexact(42)); Ok(()) } @@ -2292,8 +2675,8 @@ mod tests { let filter_b: Arc = Arc::new(FilterExec::try_new(lit(true), make_source(200))?); - let stats_a = registry.compute(filter_a.as_ref())?; - let stats_b = registry.compute(filter_b.as_ref())?; + let stats_a = compute(®istry, filter_a.as_ref())?; + let stats_b = compute(®istry, filter_b.as_ref())?; assert_eq!(stats_a.base.num_rows, Precision::Inexact(100)); assert_eq!(stats_b.base.num_rows, Precision::Inexact(50)); diff --git a/datafusion/physical-plan/src/statistics.rs b/datafusion/physical-plan/src/statistics.rs index 9246d7d9f5a9c..efd8af0d726b0 100644 --- a/datafusion/physical-plan/src/statistics.rs +++ b/datafusion/physical-plan/src/statistics.rs @@ -21,14 +21,29 @@ //! [`ExecutionPlan::statistics_from_inputs`]. use crate::ExecutionPlan; +use crate::displayable; +use crate::operator_statistics::{ + ExtendedStatistics, StatisticsRegistry, StatisticsResult, +}; +use datafusion_common::extensions::Extensions; use datafusion_common::{ Result, Statistics, assert_eq_or_internal_err, assert_or_internal_err, }; +use log::debug; use std::cell::RefCell; use std::collections::HashMap; use std::rc::Rc; use std::sync::Arc; +type CacheKey = (usize, Option); + +fn cache_key(plan: &dyn ExecutionPlan, partition: Option) -> CacheKey { + ( + plan as *const dyn ExecutionPlan as *const () as usize, + partition, + ) +} + /// Per-call memoization cache for statistics computation. /// /// Keyed by `(plan node pointer address, partition)`. Shared across @@ -37,34 +52,15 @@ use std::sync::Arc; /// The pointer-based key is safe within a single synchronous walk: /// all `Arc` nodes are held by the plan tree for /// the duration of the walk, so addresses cannot be reused. +/// +/// Core statistics and provider extensions are cached separately: the +/// `statistics` map is the hot path (populated on every walk); the `extensions` +/// map is populated only when a provider returns non-empty extensions, so a walk +/// with no providers never touches it. #[derive(Debug, Default)] -struct StatsCache(HashMap<(usize, Option), Arc>); - -impl StatsCache { - fn get( - &self, - plan: &dyn ExecutionPlan, - partition: Option, - ) -> Option<&Arc> { - let key = ( - plan as *const dyn ExecutionPlan as *const () as usize, - partition, - ); - self.0.get(&key) - } - - fn insert( - &mut self, - plan: &dyn ExecutionPlan, - partition: Option, - stats: Arc, - ) { - let key = ( - plan as *const dyn ExecutionPlan as *const () as usize, - partition, - ); - self.0.insert(key, stats); - } +struct StatsCache { + statistics: HashMap>, + extensions: HashMap, } /// Arguments passed to [`ExecutionPlan::statistics_from_inputs`] carrying @@ -118,8 +114,22 @@ pub enum ChildStats { /// Owns the bottom-up traversal and per-walk memoization cache for statistics /// computation. Call [`StatisticsContext::compute`] to walk a plan tree. +/// +/// An optional [`StatisticsRegistry`] plugs providers into the walk: at each node +/// they are consulted before the operator's built-in +/// [`ExecutionPlan::statistics_from_inputs`]. An empty registry is the built-in +/// computation. +/// +/// The walk carries [`ExtendedStatistics`]. A node has extensions only if a +/// provider `Computed` them for it; a node that falls back to the built-in +/// [`ExecutionPlan::statistics_from_inputs`] has none. So extensions propagate +/// upward only through an unbroken chain of provider-handled nodes: a single +/// built-in node yields no extensions and hides those of everything beneath it. +/// [`Self::compute_extended`] observes extensions; [`Self::compute`] returns core +/// [`Statistics`] only. pub struct StatisticsContext { cache: Rc>, + registry: StatisticsRegistry, } impl Default for StatisticsContext { @@ -129,10 +139,16 @@ impl Default for StatisticsContext { } impl StatisticsContext { - /// Creates a context with an empty cache. + /// Creates a context with an empty cache and no statistics providers. pub fn new() -> Self { + Self::new_with_registry(StatisticsRegistry::new()) + } + + /// Creates a context whose walk consults `registry`'s provider chain. + pub fn new_with_registry(registry: StatisticsRegistry) -> Self { Self { cache: Rc::new(RefCell::new(StatsCache::default())), + registry, } } @@ -143,15 +159,49 @@ impl StatisticsContext { /// (which rewrite the plan) when reusing one context across them, so stale /// pointer keys cannot collide. pub fn reset_cache(&self) { - self.cache.borrow_mut().0.clear(); + let mut cache = self.cache.borrow_mut(); + cache.statistics.clear(); + cache.extensions.clear(); + } + + /// Computes the core [`Statistics`] for `plan`, discarding any + /// provider-supplied extensions (see [`Self::compute_extended`]). + /// + /// With no providers registered this is the plain built-in walk: only the + /// `statistics` cache is touched, so it carries no extension overhead. + pub fn compute( + &self, + plan: &dyn ExecutionPlan, + args: &StatisticsArgs, + ) -> Result> { + self.compute_base(plan, args) + } + + /// 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). + pub fn compute_extended( + &self, + plan: &dyn ExecutionPlan, + args: &StatisticsArgs, + ) -> Result> { + let statistics = self.compute_base(plan, args)?; + let extensions = self + .cached_extensions(plan, args.partition()) + .unwrap_or_default(); + Ok(Arc::new(ExtendedStatistics::new_with_extensions( + statistics, extensions, + ))) } - /// Computes statistics for `plan`, resolving children first and passing - /// the results to [`ExecutionPlan::statistics_from_inputs`]. + /// Bottom-up walk producing the node's core statistics, resolving children + /// first and consulting the provider chain before the operator's built-in + /// [`ExecutionPlan::statistics_from_inputs`]. Any extensions a provider + /// attaches are recorded in the extension cache for [`Self::compute_extended`]. /// /// When `args.partition()` is `Some(idx)`, `idx` is validated against the /// plan's partition count. - pub fn compute( + fn compute_base( &self, plan: &dyn ExecutionPlan, args: &StatisticsArgs, @@ -168,12 +218,35 @@ impl StatisticsContext { ); } - if let Some(cached) = self.cache.borrow().get(plan, partition) { - return Ok(Arc::clone(cached)); + if let Some(cached) = self.cached_statistics(plan, partition) { + return Ok(cached); } let children = plan.children(); - let requests = plan.child_stats_requests(partition); + // Try providers before resolving the operator's own children, so a + // provider that overrides this node is not blocked by the fallback walk. + let statistics = match self.try_provider_stats(plan, &children, args)? { + Some(statistics) => statistics, + None => { + let requests = plan.child_stats_requests(partition); + self.validate_child_requests(plan, &children, &requests)?; + let child_statistics = + self.resolve_children(plan, &children, &requests)?; + plan.statistics_from_inputs(&child_statistics, args)? + } + }; + self.store_statistics(plan, partition, Arc::clone(&statistics)); + Ok(statistics) + } + + /// Validates child stat `requests` against `plan`'s children: the count must + /// match, and each `At(Some(idx))` must be a valid partition of that child. + fn validate_child_requests( + &self, + plan: &dyn ExecutionPlan, + children: &[&Arc], + requests: &[ChildStats], + ) -> Result<()> { assert_eq_or_internal_err!( requests.len(), children.len(), @@ -182,24 +255,181 @@ impl StatisticsContext { requests.len(), children.len() ); - let child_stats = children + for (child, directive) in children.iter().zip(requests) { + if let ChildStats::At(Some(idx)) = directive { + let count = child.properties().partitioning.partition_count(); + assert_or_internal_err!( + *idx < count, + "{} requested invalid partition {idx} for child {} with {count} partitions", + plan.name(), + child.name() + ); + } + } + Ok(()) + } + + /// Resolves each child's core statistics per `requests`: computes the child + /// at the requested partition (memoized), or supplies a + /// [`Statistics::new_unknown`] placeholder for [`ChildStats::Skip`]. Callers + /// must validate `requests` via [`Self::validate_child_requests`] first. + fn resolve_children( + &self, + plan: &dyn ExecutionPlan, + children: &[&Arc], + requests: &[ChildStats], + ) -> Result>> { + children .iter() .zip(requests) - .map(|(child, directive)| match directive { - ChildStats::At(p) => { - self.compute(child.as_ref(), &StatisticsArgs::new().with_partition(p)) - } + .enumerate() + .map(|(i, (child, directive))| match directive { + ChildStats::At(p) => self + .compute_base(child.as_ref(), &StatisticsArgs::new().with_partition(*p)) + .map_err(|e| { + e.context(format!( + "computing statistics for child {i} ({}) of {} at partition {p:?}", + child.name(), + plan.name() + )) + }), ChildStats::Skip => { Ok(Arc::new(Statistics::new_unknown(child.schema().as_ref()))) } }) - .collect::>>()?; + .collect() + } + + /// Runs the provider chain, returning the first `Computed` result's core + /// statistics (and recording its extensions), or `None` if the chain is empty + /// or all delegate. A partition-blind provider applies only to overall stats + /// (its default `compute_statistics_with_args` delegates per partition). + /// + /// Each provider's child statistics come from its own + /// [`child_stats_requests`](crate::operator_statistics::StatisticsProvider::child_stats_requests) + /// and are memoized, so a walk with no providers pays nothing. + fn try_provider_stats( + &self, + plan: &dyn ExecutionPlan, + children: &[&Arc], + args: &StatisticsArgs, + ) -> Result>> { + let providers = self.registry.providers(); + if providers.is_empty() { + return Ok(None); + } + let partition = args.partition(); + for provider in providers { + if !provider.matches(plan) { + continue; + } + let requests = provider.child_stats_requests(plan, partition); + self.validate_child_requests(plan, children, &requests)?; + // A provider's child walk is speculative: on failure, skip the provider + // so a later one or the operator fallback can handle the node. Not + // error-swallowing, whoever genuinely needs the child resolves it again + // and the error resurfaces there; a matched provider's own `compute` + // error below stays fatal. + let child_statistics = match self.resolve_children(plan, children, &requests) + { + Ok(child_statistics) => child_statistics, + Err(e) => { + debug!( + "Statistics provider {provider:?} skipped for {}: child statistics resolution failed: {e}", + displayable(plan).one_line().to_string().trim_end() + ); + continue; + } + }; + let child_extended = + self.child_extended_stats(children, &requests, &child_statistics); + if let StatisticsResult::Computed(computed) = + provider.compute_statistics_with_args(plan, &child_extended, args)? + { + if !computed.extensions().is_empty() { + self.store_extensions(plan, partition, computed.extensions().clone()); + } + return Ok(Some(Arc::clone(computed.base_arc()))); + } + } + Ok(None) + } + + /// Pairs each child's core statistics with any extensions cached for it, + /// producing the [`ExtendedStatistics`] the provider chain consumes. Called + /// only when providers exist, so an empty registry pays no extension cost. + fn child_extended_stats( + &self, + children: &[&Arc], + requests: &[ChildStats], + child_statistics: &[Arc], + ) -> Vec { + children + .iter() + .zip(requests) + .zip(child_statistics) + .map(|((child, directive), statistics)| { + let extensions = match directive { + ChildStats::At(p) => self.cached_extensions(child.as_ref(), *p), + ChildStats::Skip => None, + }; + match extensions { + Some(extensions) => ExtendedStatistics::new_with_extensions( + Arc::clone(statistics), + extensions, + ), + None => ExtendedStatistics::new_arc(Arc::clone(statistics)), + } + }) + .collect() + } - let result = plan.statistics_from_inputs(&child_stats, args)?; + fn cached_statistics( + &self, + plan: &dyn ExecutionPlan, + partition: Option, + ) -> Option> { + self.cache + .borrow() + .statistics + .get(&cache_key(plan, partition)) + .cloned() + } + + fn store_statistics( + &self, + plan: &dyn ExecutionPlan, + partition: Option, + statistics: Arc, + ) { + self.cache + .borrow_mut() + .statistics + .insert(cache_key(plan, partition), statistics); + } + + fn cached_extensions( + &self, + plan: &dyn ExecutionPlan, + partition: Option, + ) -> Option { + self.cache + .borrow() + .extensions + .get(&cache_key(plan, partition)) + .cloned() + } + + fn store_extensions( + &self, + plan: &dyn ExecutionPlan, + partition: Option, + extensions: Extensions, + ) { self.cache .borrow_mut() - .insert(plan, partition, Arc::clone(&result)); - Ok(result) + .extensions + .insert(cache_key(plan, partition), extensions); } } @@ -207,10 +437,97 @@ impl StatisticsContext { mod tests { use super::*; use crate::coalesce_partitions::CoalescePartitionsExec; + use crate::operator_statistics::StatisticsProvider; use crate::test::exec::StatisticsExec; + use crate::union::UnionExec; use arrow::datatypes::{DataType, Field, Schema}; use datafusion_common::{ColumnStatistics, stats::Precision}; + /// Overall-only provider: sets a fixed row count for any node. + #[derive(Debug)] + struct FixedRowCountProvider(usize); + impl StatisticsProvider for FixedRowCountProvider { + fn compute_statistics( + &self, + plan: &dyn ExecutionPlan, + _child_stats: &[ExtendedStatistics], + ) -> Result { + let mut stats = Statistics::new_unknown(&plan.schema()); + stats.num_rows = Precision::Exact(self.0); + Ok(StatisticsResult::Computed(ExtendedStatistics::new(stats))) + } + } + + /// Partition-aware provider: encodes the requested partition in the row count. + #[derive(Debug)] + struct PartitionRowCountProvider; + impl StatisticsProvider for PartitionRowCountProvider { + fn compute_statistics_with_args( + &self, + plan: &dyn ExecutionPlan, + _child_stats: &[ExtendedStatistics], + args: &StatisticsArgs, + ) -> Result { + let marker = 700 + args.partition().map_or(0, |p| p + 1); + let mut stats = Statistics::new_unknown(&plan.schema()); + stats.num_rows = Precision::Exact(marker); + Ok(StatisticsResult::Computed(ExtendedStatistics::new(stats))) + } + } + + #[derive(Debug, Clone, PartialEq)] + struct Tag(u32); + + /// Leaf provider: sets a row count and attaches a `Tag` extension. + #[derive(Debug)] + struct TagLeafProvider { + rows: usize, + tag: u32, + } + impl StatisticsProvider for TagLeafProvider { + fn compute_statistics( + &self, + plan: &dyn ExecutionPlan, + child_stats: &[ExtendedStatistics], + ) -> Result { + if !child_stats.is_empty() { + return Ok(StatisticsResult::Delegate); + } + let mut stats = Statistics::new_unknown(&plan.schema()); + stats.num_rows = Precision::Exact(self.rows); + let mut extended = ExtendedStatistics::new(stats); + extended.set_extension(Tag(self.tag)); + Ok(StatisticsResult::Computed(extended)) + } + } + + /// Non-leaf provider: re-emits a `Tag` doubled from the first child's `Tag`, + /// proving the child's extension reached this provider. + #[derive(Debug)] + struct TagDoublingProvider; + impl StatisticsProvider for TagDoublingProvider { + fn compute_statistics( + &self, + plan: &dyn ExecutionPlan, + child_stats: &[ExtendedStatistics], + ) -> Result { + let Some(Tag(v)) = child_stats.first().and_then(|c| c.get_extension::()) + else { + return Ok(StatisticsResult::Delegate); + }; + let mut extended = + ExtendedStatistics::new(Statistics::new_unknown(&plan.schema())); + extended.set_extension(Tag(v * 2)); + Ok(StatisticsResult::Computed(extended)) + } + } + + fn ctx_with(provider: Arc) -> StatisticsContext { + StatisticsContext::new_with_registry(StatisticsRegistry::with_providers(vec![ + provider, + ])) + } + fn make_stats_leaf(num_rows: usize) -> Arc { let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); let col_stats = vec![ColumnStatistics { @@ -256,7 +573,7 @@ mod tests { let args = StatisticsArgs::new(); let s1 = ctx.compute(leaf.as_ref(), &args).unwrap(); - assert!(!ctx.cache.borrow().0.is_empty()); + assert!(!ctx.cache.borrow().statistics.is_empty()); let s2 = ctx.compute(leaf.as_ref(), &args).unwrap(); assert!(Arc::ptr_eq(&s1, &s2)); @@ -267,8 +584,82 @@ mod tests { let leaf = make_stats_leaf(10); let ctx = StatisticsContext::new(); let _ = ctx.compute(leaf.as_ref(), &StatisticsArgs::new()).unwrap(); - assert!(!ctx.cache.borrow().0.is_empty()); + assert!(!ctx.cache.borrow().statistics.is_empty()); ctx.reset_cache(); - assert!(ctx.cache.borrow().0.is_empty()); + assert!(ctx.cache.borrow().statistics.is_empty()); + } + + #[test] + fn partition_aware_provider_applies_per_partition() { + let leaf = make_stats_leaf(10); + let ctx = ctx_with(Arc::new(PartitionRowCountProvider)); + + let per_part = ctx + .compute( + leaf.as_ref(), + &StatisticsArgs::new().with_partition(Some(0)), + ) + .unwrap(); + assert_eq!(per_part.num_rows, Precision::Exact(701)); + } + + #[test] + fn extensions_reach_parent_provider() { + let leaf = make_stats_leaf(100); + let parent: Arc = Arc::new(CoalescePartitionsExec::new(leaf)); + let ctx = StatisticsContext::new_with_registry( + StatisticsRegistry::with_providers(vec![ + Arc::new(TagLeafProvider { rows: 100, tag: 7 }), + Arc::new(TagDoublingProvider), + ]), + ); + let extended = ctx + .compute_extended(parent.as_ref(), &StatisticsArgs::new()) + .unwrap(); + assert_eq!(extended.get_extension::(), Some(&Tag(14))); + } + + #[test] + fn builtin_fallback_drops_extensions() { + let leaf = make_stats_leaf(100); + let parent: Arc = + Arc::new(CoalescePartitionsExec::new(Arc::clone(&leaf))); + let ctx = ctx_with(Arc::new(TagLeafProvider { rows: 100, tag: 7 })); + + let leaf_extended = ctx + .compute_extended(leaf.as_ref(), &StatisticsArgs::new()) + .unwrap(); + assert_eq!(leaf_extended.get_extension::(), Some(&Tag(7))); + + let parent_extended = ctx + .compute_extended(parent.as_ref(), &StatisticsArgs::new()) + .unwrap(); + assert_eq!(parent_extended.get_extension::(), None); + assert_eq!(parent_extended.base().num_rows, Precision::Exact(100)); + } + + #[test] + fn per_partition_union_with_registry_no_out_of_bounds() { + // Two 2-partition inputs -> 4 output partitions. Union owns output + // partition 3 via its second input (owning_input(3) = (1, 1)); the first + // input is Skipped, so the walk supplies a placeholder for it (never + // resolving it at partition 3, which is out of that input's 0..2 range). + // An overall-only provider delegates for a specific partition, so p3 keeps + // the operator's honest per-partition answer while the overall request + // picks up the provider's row count. + let union = + UnionExec::try_new(vec![make_stats_leaf(10), make_stats_leaf(20)]).unwrap(); + let ctx = ctx_with(Arc::new(FixedRowCountProvider(999))); + + let p3 = ctx + .compute( + union.as_ref(), + &StatisticsArgs::new().with_partition(Some(3)), + ) + .unwrap(); + assert_eq!(p3.num_rows, Precision::Absent); + + let overall = ctx.compute(union.as_ref(), &StatisticsArgs::new()).unwrap(); + assert_eq!(overall.num_rows, Precision::Exact(999)); } } diff --git a/datafusion/sqllogictest/src/test_context.rs b/datafusion/sqllogictest/src/test_context.rs index 99c3179ef1056..700cd1de4b6ae 100644 --- a/datafusion/sqllogictest/src/test_context.rs +++ b/datafusion/sqllogictest/src/test_context.rs @@ -60,6 +60,7 @@ use async_trait::async_trait; use datafusion::common::cast::as_float64_array; use datafusion::execution::SessionStateBuilder; use datafusion::execution::runtime_env::RuntimeEnv; +use datafusion::physical_plan::operator_statistics::StatisticsRegistry; use log::info; use sqlparser::ast; use tempfile::TempDir; @@ -130,6 +131,15 @@ impl TestContext { state_builder.with_type_planner(Arc::new(SqlLogicTestTypePlanner)); } + if matches!( + relative_path.file_name().and_then(|name| name.to_str()), + Some("statistics_registry.slt") + ) { + state_builder = state_builder.with_statistics_registry( + StatisticsRegistry::default_with_builtin_providers(), + ); + } + let state = state_builder.build(); let mut test_ctx = TestContext::new(SessionContext::new_with_state(state)); diff --git a/datafusion/sqllogictest/test_files/information_schema.slt b/datafusion/sqllogictest/test_files/information_schema.slt index 77acaa4747f9d..00cc98aa4dbac 100644 --- a/datafusion/sqllogictest/test_files/information_schema.slt +++ b/datafusion/sqllogictest/test_files/information_schema.slt @@ -499,7 +499,7 @@ datafusion.optimizer.repartition_windows true Should DataFusion repartition data datafusion.optimizer.skip_failed_rules false When set to true, the logical plan optimizer will produce warning messages if any optimization rules produce errors and then proceed to the next rule. When set to false, any rules that produce errors will cause the query to fail datafusion.optimizer.subset_repartition_threshold 4 Partition count threshold for subset satisfaction optimization. When the current partition count is >= this threshold, DataFusion will skip repartitioning if the required partitioning expression is a subset of the current partition expression such as Hash(a) satisfies Hash(a, b). When the current partition count is < this threshold, DataFusion will repartition to increase parallelism even when subset satisfaction applies. Set to 0 to always repartition (disable subset satisfaction optimization). Set to a high value to always use subset satisfaction. Example (subset_repartition_threshold = 4): ```text Hash([a]) satisfies Hash([a, b]) because (Hash([a, b]) is subset of Hash([a]) If current partitions (3) < threshold (4), repartition: AggregateExec: mode=FinalPartitioned, gby=[a, b], aggr=[SUM(x)] RepartitionExec: partitioning=Hash([a, b], 8), input_partitions=3 AggregateExec: mode=Partial, gby=[a, b], aggr=[SUM(x)] DataSourceExec: file_groups={...}, output_partitioning=Hash([a], 3) If current partitions (8) >= threshold (4), use subset satisfaction: AggregateExec: mode=SinglePartitioned, gby=[a, b], aggr=[SUM(x)] DataSourceExec: file_groups={...}, output_partitioning=Hash([a], 8) ``` datafusion.optimizer.top_down_join_key_reordering true When set to true, the physical plan optimizer will run a top down process to reorder the join keys -datafusion.optimizer.use_statistics_registry false When set to true, the physical plan optimizer uses the pluggable `StatisticsRegistry` for statistics propagation across operators. This enables more accurate cardinality estimates compared to each operator's built-in `partition_statistics`. +datafusion.optimizer.use_statistics_registry false (Deprecated) Ignored: the physical plan optimizer always consults the session's pluggable `StatisticsRegistry` (register providers on the `SessionState`; with none it is a no-op). datafusion.runtime.file_statistics_cache_limit 20M Maximum memory to use for file statistics cache. Supports suffixes K (kilobytes), M (megabytes), and G (gigabytes) or '0' for 0. Example: '2G' for 2 gigabytes. datafusion.runtime.list_files_cache_limit 1M Maximum memory to use for list files cache. Supports suffixes K (kilobytes), M (megabytes), and G (gigabytes) or '0' for 0. Example: '2G' for 2 gigabytes. datafusion.runtime.list_files_cache_ttl NULL TTL (time-to-live) of the entries in the list file cache. Supports units m (minutes), and s (seconds). Example: '2m' for 2 minutes. diff --git a/datafusion/sqllogictest/test_files/statistics_registry.slt b/datafusion/sqllogictest/test_files/statistics_registry.slt index 89258bec299c1..9e8dbc8293c90 100644 --- a/datafusion/sqllogictest/test_files/statistics_registry.slt +++ b/datafusion/sqllogictest/test_files/statistics_registry.slt @@ -22,8 +22,9 @@ # orders (10 rows, same distribution) # dim_small (50 rows) # -# Built-in: 10*10 / NDV(3) = 33 < 50 -> keeps inner join on build side (wrong; actual = 66) -# Registry: 10*10 = 100 > 50 -> swaps dim_small to build side (correct) +# Built-in operator stats give 10*10 / NDV(3) = 33 < 50, which would keep the +# inner join on the build side (wrong; actual output is 66). The registry's +# conservative estimate 10*10 = 100 > 50 swaps dim_small to the build side. # # Parquet files written by COPY TO carry min/max stats (NDV=3 via range) but no # distinct_count, so the registry falls back to the cartesian product upper bound. @@ -87,11 +88,8 @@ CREATE EXTERNAL TABLE dim_small STORED AS PARQUET LOCATION 'test_files/scratch/statistics_registry/dim_small.parquet'; -# -- Without registry -------------------------------------------------------- -# Built-in estimate 33 < 50: inner join stays on left (suboptimal; actual output is 66 rows) - -statement ok -set datafusion.optimizer.use_statistics_registry = false; +# -- Registry-driven join ordering ------------------------------------------- +# Conservative estimate 100 > 50: dim_small correctly swapped to the build side. query TT EXPLAIN SELECT o.order_id, c.region_id, d.label @@ -100,19 +98,21 @@ JOIN orders o ON c.customer_id = o.customer_id JOIN dim_small d ON o.small_id = d.small_id; ---- physical_plan -01)HashJoinExec: mode=Partitioned, join_type=Inner, on=[(small_id@2, small_id@0)], projection=[order_id@1, region_id@0, label@4] -02)--RepartitionExec: partitioning=Hash([small_id@2], 4), input_partitions=1, maintains_sort_order=true -03)----HashJoinExec: mode=Partitioned, join_type=Inner, on=[(customer_id@0, customer_id@1)], projection=[region_id@1, order_id@2, small_id@4] -04)------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/statistics_registry/customers.parquet]]}, projection=[customer_id, region_id], output_ordering=[region_id@1 ASC NULLS LAST], file_type=parquet -05)------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/statistics_registry/orders.parquet]]}, projection=[order_id, customer_id, small_id], output_ordering=[order_id@0 ASC NULLS LAST], file_type=parquet, predicate=DynamicFilter [ empty ], dynamic_rg_pruning=eligible -06)--RepartitionExec: partitioning=Hash([small_id@0], 4), input_partitions=1, maintains_sort_order=true -07)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/statistics_registry/dim_small.parquet]]}, projection=[small_id, label], output_ordering=[small_id@0 ASC NULLS LAST], file_type=parquet, predicate=DynamicFilter [ empty ], dynamic_rg_pruning=eligible +01)HashJoinExec: mode=Partitioned, join_type=Inner, on=[(small_id@0, small_id@2)], projection=[order_id@3, region_id@2, label@1] +02)--RepartitionExec: partitioning=Hash([small_id@0], 4), input_partitions=1, maintains_sort_order=true +03)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/statistics_registry/dim_small.parquet]]}, projection=[small_id, label], output_ordering=[small_id@0 ASC NULLS LAST], file_type=parquet +04)--RepartitionExec: partitioning=Hash([small_id@2], 4), input_partitions=1, maintains_sort_order=true +05)----HashJoinExec: mode=Partitioned, join_type=Inner, on=[(customer_id@0, customer_id@1)], projection=[region_id@1, order_id@2, small_id@4] +06)------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/statistics_registry/customers.parquet]]}, projection=[customer_id, region_id], output_ordering=[region_id@1 ASC NULLS LAST], file_type=parquet +07)------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/statistics_registry/orders.parquet]]}, projection=[order_id, customer_id, small_id], output_ordering=[order_id@0 ASC NULLS LAST], file_type=parquet, predicate=DynamicFilter [ empty ] AND DynamicFilter [ empty ], dynamic_rg_pruning=eligible -# -- With registry ----------------------------------------------------------- -# Conservative estimate 100 > 50: dim_small correctly swapped to build side +# -- Registry reflected in EXPLAIN statistics -------------------------------- +# With show_statistics on, the displayed statistics column reflects the +# registry's estimates (the top join's conservative Rows), not the built-in +# empty-registry path. statement ok -set datafusion.optimizer.use_statistics_registry = true; +set datafusion.explain.show_statistics = true; query TT EXPLAIN SELECT o.order_id, c.region_id, d.label @@ -121,29 +121,31 @@ JOIN orders o ON c.customer_id = o.customer_id JOIN dim_small d ON o.small_id = d.small_id; ---- physical_plan -01)HashJoinExec: mode=Partitioned, join_type=Inner, on=[(small_id@0, small_id@2)], projection=[order_id@3, region_id@2, label@1] -02)--RepartitionExec: partitioning=Hash([small_id@0], 4), input_partitions=1, maintains_sort_order=true -03)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/statistics_registry/dim_small.parquet]]}, projection=[small_id, label], output_ordering=[small_id@0 ASC NULLS LAST], file_type=parquet -04)--RepartitionExec: partitioning=Hash([small_id@2], 4), input_partitions=1, maintains_sort_order=true -05)----HashJoinExec: mode=Partitioned, join_type=Inner, on=[(customer_id@0, customer_id@1)], projection=[region_id@1, order_id@2, small_id@4] -06)------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/statistics_registry/customers.parquet]]}, projection=[customer_id, region_id], output_ordering=[region_id@1 ASC NULLS LAST], file_type=parquet -07)------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/statistics_registry/orders.parquet]]}, projection=[order_id, customer_id, small_id], output_ordering=[order_id@0 ASC NULLS LAST], file_type=parquet, predicate=DynamicFilter [ empty ] AND DynamicFilter [ empty ], dynamic_rg_pruning=eligible - -# -- Verify results are identical regardless of join order -------------------- +01)HashJoinExec: mode=Partitioned, join_type=Inner, on=[(small_id@0, small_id@2)], projection=[order_id@3, region_id@2, label@1], statistics=[Rows=Inexact(5000), Bytes=Absent, [(Col[0]: Min=Inexact(Int32(1)) Max=Inexact(Int32(10)) Null=Inexact(0) ScanBytes=Inexact(40)),(Col[1]: Min=Exact(Int32(1)) Max=Exact(Int32(10)) Null=Exact(0) ScanBytes=Exact(40)),(Col[2]: Min=Exact(Int32(1)) Max=Exact(Int32(50)) Null=Exact(0) ScanBytes=Exact(200))]] +02)--RepartitionExec: partitioning=Hash([small_id@0], 4), input_partitions=1, maintains_sort_order=true, statistics=[Rows=Exact(50), Bytes=Exact(400), [(Col[0]: Min=Exact(Int32(1)) Max=Exact(Int32(50)) Null=Exact(0) ScanBytes=Exact(200)),(Col[1]: Min=Exact(Int32(1)) Max=Exact(Int32(50)) Null=Exact(0) ScanBytes=Exact(200))]] +03)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/statistics_registry/dim_small.parquet]]}, projection=[small_id, label], output_ordering=[small_id@0 ASC NULLS LAST], file_type=parquet, statistics=[Rows=Exact(50), Bytes=Exact(400), [(Col[0]: Min=Exact(Int32(1)) Max=Exact(Int32(50)) Null=Exact(0) ScanBytes=Exact(200)),(Col[1]: Min=Exact(Int32(1)) Max=Exact(Int32(50)) Null=Exact(0) ScanBytes=Exact(200))]] +04)--RepartitionExec: partitioning=Hash([small_id@2], 4), input_partitions=1, maintains_sort_order=true, statistics=[Rows=Inexact(100), Bytes=Absent, [(Col[0]: Min=Exact(Int32(1)) Max=Exact(Int32(10)) Null=Exact(0) ScanBytes=Exact(40)),(Col[1]: Min=Inexact(Int32(1)) Max=Inexact(Int32(10)) Null=Inexact(0) ScanBytes=Inexact(40)),(Col[2]: Min=Inexact(Int32(1)) Max=Inexact(Int32(10)) Null=Inexact(0) ScanBytes=Inexact(40))]] +05)----HashJoinExec: mode=Partitioned, join_type=Inner, on=[(customer_id@0, customer_id@1)], projection=[region_id@1, order_id@2, small_id@4], statistics=[Rows=Inexact(100), Bytes=Absent, [(Col[0]: Min=Exact(Int32(1)) Max=Exact(Int32(10)) Null=Exact(0) ScanBytes=Exact(40)),(Col[1]: Min=Inexact(Int32(1)) Max=Inexact(Int32(10)) Null=Inexact(0) ScanBytes=Inexact(40)),(Col[2]: Min=Inexact(Int32(1)) Max=Inexact(Int32(10)) Null=Inexact(0) ScanBytes=Inexact(40))]] +06)------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/statistics_registry/customers.parquet]]}, projection=[customer_id, region_id], output_ordering=[region_id@1 ASC NULLS LAST], file_type=parquet, statistics=[Rows=Exact(10), Bytes=Exact(80), [(Col[0]: Min=Exact(Int32(1)) Max=Exact(Int32(3)) Null=Exact(0) ScanBytes=Exact(40)),(Col[1]: Min=Exact(Int32(1)) Max=Exact(Int32(10)) Null=Exact(0) ScanBytes=Exact(40))]] +07)------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/statistics_registry/orders.parquet]]}, projection=[order_id, customer_id, small_id], output_ordering=[order_id@0 ASC NULLS LAST], file_type=parquet, predicate=DynamicFilter [ empty ] AND DynamicFilter [ empty ], dynamic_rg_pruning=eligible, statistics=[Rows=Inexact(10), Bytes=Inexact(120), [(Col[0]: Min=Inexact(Int32(1)) Max=Inexact(Int32(10)) Null=Inexact(0) ScanBytes=Inexact(40)),(Col[1]: Min=Inexact(Int32(1)) Max=Inexact(Int32(3)) Null=Inexact(0) ScanBytes=Inexact(40)),(Col[2]: Min=Inexact(Int32(1)) Max=Inexact(Int32(10)) Null=Inexact(0) ScanBytes=Inexact(40))]] + +# -- Registry reflected in EXPLAIN ANALYZE statistics ------------------------ +# EXPLAIN ANALYZE renders the same registry-refined statistics as plain EXPLAIN +# (the top join's conservative Rows=Inexact(5000)) alongside runtime metrics. +# Timings are volatile, so match only the deterministic statistics fragment. -statement ok -set datafusion.optimizer.use_statistics_registry = false; - -query I -SELECT count(*) +query TT +EXPLAIN ANALYZE SELECT o.order_id, c.region_id, d.label FROM customers c JOIN orders o ON c.customer_id = o.customer_id JOIN dim_small d ON o.small_id = d.small_id; ---- -66 +Plan with Metricsstatistics=[Rows=Inexact(5000) statement ok -set datafusion.optimizer.use_statistics_registry = true; +set datafusion.explain.show_statistics = false; + +# -- Correctness ------------------------------------------------------------- query I SELECT count(*) @@ -158,9 +160,6 @@ JOIN dim_small d ON o.small_id = d.small_id; statement ok set datafusion.explain.physical_plan_only = false; -statement ok -set datafusion.optimizer.use_statistics_registry = false; - statement ok set datafusion.optimizer.hash_join_single_partition_threshold = 1048576; diff --git a/docs/source/library-user-guide/upgrading/55.0.0.md b/docs/source/library-user-guide/upgrading/55.0.0.md index 6097c8dc717df..c4f8fb285753e 100644 --- a/docs/source/library-user-guide/upgrading/55.0.0.md +++ b/docs/source/library-user-guide/upgrading/55.0.0.md @@ -454,6 +454,52 @@ let stats = plan.partition_statistics(None)?; let stats = StatisticsContext::new().compute(plan.as_ref(), &StatisticsArgs::new())?; ``` +### `datafusion.optimizer.use_statistics_registry` is deprecated and ignored + +The `datafusion.optimizer.use_statistics_registry` config flag is deprecated and +now ignored (setting it emits a deprecation warning). The pluggable +`StatisticsRegistry` is always consulted during the single statistics walk: +providers registered on the session take effect directly. With no providers +registered the registry is a no-op, so default behavior is unchanged from the +previous default (`use_statistics_registry = false`). + +**Who is affected:** + +- Anyone who set `datafusion.optimizer.use_statistics_registry`. The setting is + still accepted (with a deprecation warning) but has no effect, and will be + removed in a future release. + +**Migration guide:** + +- Drop any `use_statistics_registry` setting. +- To use statistics providers, register them on the session with + `SessionStateBuilder::with_statistics_registry(...)`. To keep the built-in + NDV-aware providers the flag previously enabled, register + `StatisticsRegistry::default_with_builtin_providers()`. +- To opt out, register nothing (the default). + +### `DefaultStatisticsProvider` is deprecated + +`datafusion_physical_plan::operator_statistics::DefaultStatisticsProvider` is +deprecated. It is redundant: the statistics walk (`StatisticsContext`) falls back +to each operator's `statistics_from_inputs` when the provider chain delegates or +is empty, so a terminal "default" provider is no longer needed. It is no longer +part of `StatisticsRegistry::default_with_builtin_providers()`. + +**Migration guide:** + +- Remove `DefaultStatisticsProvider` from any custom provider chain; register no + terminal provider instead (the walk falls back on its own). + +### `StatisticsRegistry::compute` and `compute_base` are deprecated + +Use the walk instead: + +```rust +StatisticsContext::new_with_registry(registry) + .compute_extended(plan, &StatisticsArgs::new())?; // or .compute(...) for core Statistics +``` + ### `DdlStatement::CreateExternalTable` and `CreateFunction` are now boxed The two largest variants of `datafusion_expr::DdlStatement` are now diff --git a/docs/source/user-guide/configs.md b/docs/source/user-guide/configs.md index e01af3476b94c..5b967e28896fc 100644 --- a/docs/source/user-guide/configs.md +++ b/docs/source/user-guide/configs.md @@ -169,7 +169,7 @@ The following configuration settings are available: | datafusion.optimizer.max_passes | 3 | Number of times that the optimizer will attempt to optimize the plan | | datafusion.optimizer.top_down_join_key_reordering | true | When set to true, the physical plan optimizer will run a top down process to reorder the join keys | | datafusion.optimizer.join_reordering | true | When set to true, the physical plan optimizer may swap join inputs based on statistics. When set to false, statistics-driven join input reordering is disabled and the original join order in the query is used. | -| datafusion.optimizer.use_statistics_registry | false | When set to true, the physical plan optimizer uses the pluggable `StatisticsRegistry` for statistics propagation across operators. This enables more accurate cardinality estimates compared to each operator's built-in `partition_statistics`. | +| datafusion.optimizer.use_statistics_registry | false | (Deprecated) Ignored: the physical plan optimizer always consults the session's pluggable `StatisticsRegistry` (register providers on the `SessionState`; with none it is a no-op). | | datafusion.optimizer.prefer_hash_join | true | When set to true, the physical plan optimizer will prefer HashJoin over SortMergeJoin. HashJoin can work more efficiently than SortMergeJoin but consumes more memory | | datafusion.optimizer.enable_piecewise_merge_join | false | When set to true, piecewise merge join is enabled. PiecewiseMergeJoin is currently experimental. Physical planner will opt for PiecewiseMergeJoin when there is only one range filter. | | datafusion.optimizer.hash_join_single_partition_threshold | 1048576 | The maximum estimated size in bytes for one input side of a HashJoin will be collected into a single partition |