feat(examples): add StatisticsRegistry examples (override + custom stats)

asolimando · asolimando · commit 1f2d39d7a112 · 2026-07-16T21:04:01.000+02:00
Two runnable examples under datafusion-examples/examples/statistics, each showing
a registry-driven plan flip via before/after EXPLAIN:

- override: a provider corrects an intermediate join's row-count estimate with
  plain Statistics, flipping the outer join's build side.
- custom_stats: a custom NDV sketch is attached at scans, merged across a join
  through the ExtendedStatistics extension map, and consumed to flip the plan -
  something a scalar distinct_count cannot express.
diff --git a/datafusion-examples/README.md b/datafusion-examples/README.md
@@ -213,6 +213,17 @@ 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                                                                                 |
+| ------------ | ----------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------- |
+| custom_stats | [`statistics/statistics_registry_custom_stats.rs`](examples/statistics/statistics_registry_custom_stats.rs) | Propagate a custom NDV sketch through the extension map and consume it at joins             |
+| override     | [`statistics/statistics_registry_override.rs`](examples/statistics/statistics_registry_override.rs)         | Correct a built-in operator's row-count estimate with plain Statistics to flip a join order |
+
 ## UDF Examples
 
 ### Group: `udf`
diff --git a/datafusion-examples/examples/statistics/main.rs b/datafusion-examples/examples/statistics/main.rs
@@ -0,0 +1,94 @@
+// 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|override|custom_stats]
+//! ```
+//!
+//! Each subcommand runs a corresponding example:
+//! - `all`: run all examples included in this module
+//!
+//! - `custom_stats`
+//!   (file: statistics_registry_custom_stats.rs, desc: Propagate a custom NDV sketch through the extension map and consume it at joins)
+//!
+//! - `override`
+//!   (file: statistics_registry_override.rs, desc: Correct a built-in operator's row-count estimate with plain Statistics to flip a join order)
+
+mod statistics_registry_custom_stats;
+mod statistics_registry_override;
+
+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,
+    Override,
+    CustomStats,
+}
+
+impl ExampleKind {
+    const EXAMPLE_NAME: &str = "statistics";
+
+    fn runnable() -> impl Iterator<Item = ExampleKind> {
+        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::Override => {
+                statistics_registry_override::statistics_registry_override().await
+            }
+            ExampleKind::CustomStats => {
+                statistics_registry_custom_stats::statistics_registry_custom_stats().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-examples/examples/statistics/statistics_registry_custom_stats.rs b/datafusion-examples/examples/statistics/statistics_registry_custom_stats.rs
@@ -0,0 +1,251 @@
+// 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.
+
+//! Propagate a custom NDV sketch through the extension map and consume it at joins.
+//!
+//! Core [`Statistics`] carries a *scalar* `distinct_count` at leaves but cannot
+//! correctly combine NDV across a join (that needs the underlying sketches to be
+//! unioned/intersected), so it cannot propagate accurate NDV up a join chain. The
+//! [`ExtendedStatistics`] extension map can carry the sketch itself, which a
+//! provider merges as it walks.
+//!
+//! `(a ⋈ b ON k) ⋈ c ON k`. A producer provider attaches each scan's key sketch;
+//! a join provider intersects the two child sketches for the join key, emits the
+//! merged sketch, and estimates cardinality from its NDV. With no sketches the
+//! all-distinct fallback underestimates `a ⋈ b` (50 rows), so it stays on the
+//! build side; the sketch yields the true NDV and estimates 1000 rows, flipping
+//! `c` onto the build side. The merged `a ⋈ b` sketch is what the `⋈ c` join then
+//! consumes.
+//!
+//! Extensions are not forwarded by the walk itself, so `PassthroughStatisticsProvider`
+//! is registered to carry the scan sketches through the hash repartitions into the
+//! join; without it (or a custom passthrough) they would stop at the first operator
+//! that does not produce them.
+//!
+//! For clarity the "sketch" is an exact set of key values; a real deployment would
+//! use a bounded-space HLL/theta sketch with the same merge semantics.
+
+use std::collections::BTreeSet;
+use std::sync::Arc;
+
+use datafusion::arrow::array::Int64Array;
+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::stats::Precision;
+use datafusion::common::{Result, Statistics};
+use datafusion::execution::SessionStateBuilder;
+use datafusion::physical_plan::ExecutionPlan;
+use datafusion::physical_plan::joins::HashJoinExec;
+use datafusion::physical_plan::operator_statistics::{
+    ExtendedStatistics, PassthroughStatisticsProvider, StatisticsProvider,
+    StatisticsRegistry, StatisticsResult,
+};
+use datafusion::prelude::*;
+
+/// A custom statistic: the set of distinct join-key values. Stand-in for an
+/// HLL/theta sketch (same `ndv`/`intersect` semantics, bounded space).
+#[derive(Clone, Debug)]
+struct KeySketch {
+    values: BTreeSet<i64>,
+}
+
+impl KeySketch {
+    fn of(values: impl IntoIterator<Item = i64>) -> Self {
+        Self {
+            values: values.into_iter().collect(),
+        }
+    }
+    fn ndv(&self) -> usize {
+        self.values.len()
+    }
+    /// Merge for an equi-join key: only matched values survive.
+    fn intersect(&self, other: &Self) -> Self {
+        Self {
+            values: self.values.intersection(&other.values).copied().collect(),
+        }
+    }
+}
+
+/// Attaches each base table's key sketch (and row count) to its scan, as an
+/// external catalog would. Matched by a table's marker column.
+#[derive(Debug)]
+struct CatalogSketches {
+    tables: Vec<(&'static str, usize, KeySketch)>,
+}
+
+impl StatisticsProvider for CatalogSketches {
+    fn compute_statistics(
+        &self,
+        plan: &dyn ExecutionPlan,
+        child_stats: &[ExtendedStatistics],
+    ) -> Result<StatisticsResult> {
+        // Only leaf scans.
+        if !child_stats.is_empty() {
+            return Ok(StatisticsResult::Delegate);
+        }
+        let schema = plan.schema();
+        let Some((_, rows, sketch)) = self
+            .tables
+            .iter()
+            .find(|(marker, _, _)| schema.field_with_name(marker).is_ok())
+        else {
+            return Ok(StatisticsResult::Delegate);
+        };
+        let mut base = Statistics::new_unknown(&schema);
+        base.num_rows = Precision::Exact(*rows);
+        let mut ext = ExtendedStatistics::new(base);
+        ext.set_extension(sketch.clone());
+        Ok(StatisticsResult::Computed(ext))
+    }
+}
+
+/// Estimates an equi-join's cardinality from the child key sketches and emits the
+/// merged (intersected) sketch so the next join can consume it.
+#[derive(Debug)]
+struct SketchJoinCardinality;
+
+impl StatisticsProvider for SketchJoinCardinality {
+    fn compute_statistics(
+        &self,
+        plan: &dyn ExecutionPlan,
+        child_stats: &[ExtendedStatistics],
+    ) -> Result<StatisticsResult> {
+        if plan.downcast_ref::<HashJoinExec>().is_none() || child_stats.len() != 2 {
+            return Ok(StatisticsResult::Delegate);
+        }
+        let (Some(left), Some(right)) = (
+            child_stats[0].get_extension::<KeySketch>(),
+            child_stats[1].get_extension::<KeySketch>(),
+        ) else {
+            return Ok(StatisticsResult::Delegate);
+        };
+        let (Some(&l_rows), Some(&r_rows)) = (
+            child_stats[0].base().num_rows.get_value(),
+            child_stats[1].base().num_rows.get_value(),
+        ) else {
+            return Ok(StatisticsResult::Delegate);
+        };
+        let merged = left.intersect(right);
+        let ndv = left.ndv().max(right.ndv()).max(1);
+        let mut base = Statistics::new_unknown(&plan.schema());
+        base.num_rows = Precision::Inexact(l_rows.saturating_mul(r_rows) / ndv);
+        let mut ext = ExtendedStatistics::new(base);
+        ext.set_extension(merged);
+        Ok(StatisticsResult::Computed(ext))
+    }
+}
+
+fn mem_table(cols: &[(&str, Vec<i64>)]) -> Result<Arc<MemTable>> {
+    let schema = Arc::new(Schema::new(
+        cols.iter()
+            .map(|(name, _)| Field::new(*name, DataType::Int64, false))
+            .collect::<Vec<_>>(),
+    ));
+    let arrays = cols
+        .iter()
+        .map(|(_, v)| Arc::new(Int64Array::from(v.clone())) as _)
+        .collect();
+    let batch = RecordBatch::try_new(Arc::clone(&schema), arrays)?;
+    Ok(Arc::new(MemTable::try_new(schema, vec![vec![batch]])?))
+}
+
+async fn build_ctx(with_registry: bool) -> Result<SessionContext> {
+    let config = SessionConfig::new()
+        .with_target_partitions(4)
+        .set_bool("datafusion.explain.physical_plan_only", true)
+        .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 registry = StatisticsRegistry::with_providers(vec![
+            Arc::new(SketchJoinCardinality),
+            // Carries the scan sketches through the hash RepartitionExecs into the
+            // join (extensions are not forwarded by the walk itself).
+            Arc::new(PassthroughStatisticsProvider),
+            Arc::new(CatalogSketches {
+                tables: vec![
+                    ("va", 50, KeySketch::of(0..50)),
+                    ("vb", 1000, KeySketch::of(0..50)),
+                    ("vc", 500, KeySketch::of(0..500)),
+                ],
+            }),
+        ]);
+        builder = builder.with_statistics_registry(registry);
+    }
+
+    let ctx = SessionContext::new_with_state(builder.build());
+    ctx.register_table(
+        "a",
+        mem_table(&[("k", (0..50).collect()), ("va", (0..50).collect())])?,
+    )?;
+    ctx.register_table(
+        "b",
+        mem_table(&[
+            ("k", (0..1000).map(|i| i % 50).collect()),
+            ("vb", (0..1000).collect()),
+        ])?,
+    )?;
+    ctx.register_table(
+        "c",
+        mem_table(&[("k", (0..500).collect()), ("vc", (0..500).collect())])?,
+    )?;
+    Ok(ctx)
+}
+
+const QUERY: &str = "SELECT a.va, b.vb, c.vc \
+     FROM a \
+     JOIN b ON a.k = b.k \
+     JOIN c ON b.k = c.k";
+
+async fn explain(ctx: &SessionContext) -> Result<String> {
+    let batches = ctx
+        .sql(&format!("EXPLAIN {QUERY}"))
+        .await?
+        .collect()
+        .await?;
+    Ok(pretty_format_batches(&batches)?.to_string())
+}
+
+pub async fn statistics_registry_custom_stats() -> Result<()> {
+    println!(
+        "Watch the build (first) input of the OUTER join (on `c`). Without sketches the\n\
+         all-distinct fallback underestimates `a ⋈ b` (50 rows) and keeps it on the build\n\
+         side; the key sketch gives the true NDV (a ⋈ b = 1000 rows), moving `c` to build.\n"
+    );
+
+    println!("-- Without the sketch providers (core has no distinct_count) --");
+    println!("{}\n", explain(&build_ctx(false).await?).await?);
+
+    println!(
+        "-- With CatalogSketches + SketchJoinCardinality (sketch merged a ⋈ b -> ⋈ c) --"
+    );
+    println!("{}", explain(&build_ctx(true).await?).await?);
+
+    Ok(())
+}
diff --git a/datafusion-examples/examples/statistics/statistics_registry_override.rs b/datafusion-examples/examples/statistics/statistics_registry_override.rs
