feat[substrait]: translate scan statistics via RelCommon hints

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: etiennepelissier

datafusion/substrait/src/logical_plan/consumer/rel/read_rel.rs

@@ -16,6 +16,7 @@
 // under the License.
 
 use crate::logical_plan::consumer::SubstraitConsumer;
+use crate::logical_plan::consumer::SubstraitHints;
 use crate::logical_plan::consumer::from_substrait_literal;
 use crate::logical_plan::consumer::from_substrait_named_struct;
 use crate::logical_plan::consumer::utils::ensure_schema_compatibility;
@@ -40,12 +41,51 @@ pub async fn from_read_rel(
     consumer: &impl SubstraitConsumer,
     read: &ReadRel,
 ) -> datafusion::common::Result<LogicalPlan> {
+    // proto3 defaults all numeric fields to 0, so treat 0.0 (or negative) as
+    // "not provided" to avoid injecting bogus stats when a producer sets
+    // hint.stats for other fields but leaves a field unset.
+    // Note: zero-row / zero-size values will also be silently dropped — this
+    // is an accepted limitation of the proto3 default.
+    // We also reject non-finite values (Inf, NaN) and values at or above
+    // usize::MAX, since the consumer will later cast these to usize.  Rust's
+    // float-to-integer cast saturates rather than panicking (since 1.45), but
+    // silently capping to usize::MAX would produce nonsensical statistics.
+    let hints = {
+        let stats = read
+            .common
+            .as_ref()
+            .and_then(|c| c.hint.as_ref())
+            .and_then(|h| h.stats.as_ref());
+        // usize::MAX as f64 may round up slightly above the true maximum, so
+        // using it as a strict upper bound excludes values that would overflow.
+        let max_hint: f64 = usize::MAX as f64;
+        let row_count = stats.and_then(|s| {
+            if s.row_count > 0.0 && s.row_count.is_finite() && s.row_count < max_hint {
+                Some(s.row_count)
+            } else {
+                None
+            }
+        });
+        let record_size = stats.and_then(|s| {
+            if s.record_size > 0.0 && s.record_size.is_finite() && s.record_size < max_hint {
+                Some(s.record_size)
+            } else {
+                None
+            }
+        });
+        SubstraitHints {
+            row_count,
+            record_size,
+        }
+    };
+
     async fn read_with_schema(
         consumer: &impl SubstraitConsumer,
         table_ref: TableReference,
         schema: DFSchema,
         projection: &Option<MaskExpression>,
         filter: &Option<Box<Expression>>,
+        hints: SubstraitHints,
     ) -> datafusion::common::Result<LogicalPlan> {
         let schema = schema.replace_qualifier(table_ref.clone());
 
@@ -57,14 +97,17 @@ pub async fn from_read_rel(
         };
 
         let plan = {
-            let provider = match consumer.resolve_table_ref(&table_ref).await? {
-                Some(ref provider) => Arc::clone(provider),
+            let provider = match consumer
+                .resolve_table_ref(&table_ref, hints)
+                .await?
+            {
+                Some(provider) => provider,
                 _ => return plan_err!("No table named '{table_ref}'"),
             };
 
             LogicalPlanBuilder::scan_with_filters(
                 table_ref,
-                provider_as_source(Arc::clone(&provider)),
+                provider_as_source(provider),
                 None,
                 filters,
             )?
@@ -110,6 +153,7 @@ pub async fn from_read_rel(
                 substrait_schema,
                 &read.projection,
                 &read.filter,
+                hints,
             )
             .await
         }
@@ -213,6 +257,7 @@ pub async fn from_read_rel(
                 substrait_schema,
                 &read.projection,
                 &read.filter,
+                hints,
             )
             .await
         }

datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs

@@ -24,13 +24,16 @@ use super::{
 };
 use crate::extensions::Extensions;
 use async_trait::async_trait;
-use datafusion::arrow::datatypes::DataType;
-use datafusion::catalog::TableProvider;
+use datafusion::arrow::datatypes::{DataType, SchemaRef};
+use datafusion::catalog::{Session, TableProvider};
+use datafusion::common::stats::Precision;
 use datafusion::common::{
-    DFSchema, ScalarValue, TableReference, not_impl_err, substrait_err,
+    DFSchema, ScalarValue, Statistics, TableReference, not_impl_err, substrait_err,
 };
 use datafusion::execution::{FunctionRegistry, SessionState};
+use datafusion::logical_expr::TableType;
 use datafusion::logical_expr::{Expr, Extension, LogicalPlan};
+use datafusion::physical_plan::ExecutionPlan;
 use std::sync::{Arc, RwLock};
 use substrait::proto;
 use substrait::proto::expression as substrait_expression;
@@ -44,6 +47,26 @@ use substrait::proto::{
     FilterRel, JoinRel, ProjectRel, ReadRel, Rel, SetRel, SortRel, r#type,
 };
 
+/// Advisory hints extracted from a Substrait `RelCommon.hint.stats` message,
+/// passed to [`SubstraitConsumer::resolve_table_ref`] so that implementors can
+/// incorporate them into the returned [`TableProvider`].
+///
+/// The struct is `#[non_exhaustive]` so that new fields can be added in future
+/// versions without breaking existing implementations.
+#[non_exhaustive]
+#[derive(Debug, Clone, Default)]
+pub struct SubstraitHints {
+    /// Estimated number of rows, from `hint.stats.row_count`.
+    ///
+    /// `None` means the hint was absent or could not be reliably interpreted
+    /// (e.g. proto3 default-zero or a non-finite value).
+    pub row_count: Option<f64>,
+    /// Estimated average byte size per record, from `hint.stats.record_size`.
+    ///
+    /// `None` means the hint was absent or non-positive / non-finite.
+    pub record_size: Option<f64>,
+}
+
 #[async_trait]
 /// This trait is used to consume Substrait plans, converting them into DataFusion Logical Plans.
 /// It can be implemented by users to allow for custom handling of relations, expressions, etc.
@@ -67,7 +90,7 @@ use substrait::proto::{
 /// # use datafusion::logical_expr::expr::ScalarFunction;
 /// # use datafusion_substrait::extensions::Extensions;
 /// # use datafusion_substrait::logical_plan::consumer::{
-/// #     from_project_rel, from_substrait_rel, from_substrait_rex, SubstraitConsumer
+/// #     from_project_rel, from_substrait_rel, from_substrait_rex, SubstraitConsumer, SubstraitHints
 /// # };
 ///
 /// struct CustomSubstraitConsumer {
@@ -80,6 +103,7 @@ use substrait::proto::{
 ///     async fn resolve_table_ref(
 ///         &self,
 ///         table_ref: &TableReference,
+///         _hints: SubstraitHints,
 ///     ) -> Result<Option<Arc<dyn TableProvider>>> {
 ///         let table = table_ref.table().to_string();
 ///         let schema = self.state.schema_for_ref(table_ref.clone())?;
@@ -162,6 +186,7 @@ pub trait SubstraitConsumer: Send + Sync + Sized {
     async fn resolve_table_ref(
         &self,
         table_ref: &TableReference,
+        hints: SubstraitHints,
     ) -> datafusion::common::Result<Option<Arc<dyn TableProvider>>>;
 
     // TODO: Remove these two methods
@@ -471,6 +496,57 @@ pub trait SubstraitConsumer: Send + Sync + Sized {
     }
 }
 
+/// Wraps an inner [`TableProvider`] and overrides its `statistics()` return value.
+///
+/// Used by [`DefaultSubstraitConsumer`] to inject a row-count hint carried in a
+/// Substrait `RelCommon.hint.stats` when the resolved provider has no statistics.
+#[derive(Debug)]
+struct StatisticsOverrideTableProvider {
+    inner: Arc<dyn TableProvider>,
+    statistics: Statistics,
+}
+
+#[async_trait]
+impl TableProvider for StatisticsOverrideTableProvider {
+    fn as_any(&self) -> &dyn std::any::Any {
+        self
+    }
+
+    fn schema(&self) -> SchemaRef {
+        self.inner.schema()
+    }
+
+    fn constraints(&self) -> Option<&datafusion::common::Constraints> {
+        self.inner.constraints()
+    }
+
+    fn table_type(&self) -> TableType {
+        self.inner.table_type()
+    }
+
+    fn supports_filters_pushdown(
+        &self,
+        filters: &[&Expr],
+    ) -> datafusion::common::Result<Vec<datafusion::logical_expr::TableProviderFilterPushDown>>
+    {
+        self.inner.supports_filters_pushdown(filters)
+    }
+
+    fn statistics(&self) -> Option<Statistics> {
+        Some(self.statistics.clone())
+    }
+
+    async fn scan(
+        &self,
+        state: &dyn Session,
+        projection: Option<&Vec<usize>>,
+        filters: &[Expr],
+        limit: Option<usize>,
+    ) -> datafusion::common::Result<Arc<dyn ExecutionPlan>> {
+        self.inner.scan(state, projection, filters, limit).await
+    }
+}
+
 /// Default SubstraitConsumer for converting standard Substrait without user-defined extensions.
 ///
 /// Used as the consumer in [crate::logical_plan::consumer::from_substrait_plan]
@@ -495,11 +571,79 @@ impl SubstraitConsumer for DefaultSubstraitConsumer<'_> {
     async fn resolve_table_ref(
         &self,
         table_ref: &TableReference,
+        hints: SubstraitHints,
     ) -> datafusion::common::Result<Option<Arc<dyn TableProvider>>> {
         let table = table_ref.table().to_string();
         let schema = self.state.schema_for_ref(table_ref.clone())?;
-        let table_provider = schema.table(&table).await?;
-        Ok(table_provider)
+        let provider = schema.table(&table).await?;
+        // If the Substrait plan provides hints and the provider is missing the
+        // corresponding statistics fields, wrap it to expose those hints to
+        // DataFusion (e.g. for downstream optimizer rules).
+        // We check each field individually so that a provider returning
+        // Some(Statistics { num_rows: Absent, ... }) also gets hints injected,
+        // and any already-known values from the provider are preserved.
+        let has_hints = hints.row_count.is_some() || hints.record_size.is_some();
+        let provider = match provider {
+            Some(provider) if has_hints => {
+                let existing = provider.statistics();
+                let row_count_absent = existing
+                    .as_ref()
+                    .map_or(true, |s| matches!(s.num_rows, Precision::Absent));
+                let byte_size_absent = existing
+                    .as_ref()
+                    .map_or(true, |s| matches!(s.total_byte_size, Precision::Absent));
+                let inject_row_count = hints.row_count.is_some() && row_count_absent;
+                // total_byte_size = row_count * record_size, so both hints must be
+                // present to reconstruct it.  A record_size-only hint (row_count absent)
+                // is therefore silently ignored here, since we cannot compute a
+                // meaningful total_byte_size from record_size alone.
+                let inject_byte_size = hints.row_count.is_some()
+                    && hints.record_size.is_some()
+                    && byte_size_absent;
+                if inject_row_count || inject_byte_size {
+                    let num_rows = if inject_row_count {
+                        Precision::Inexact(hints.row_count.unwrap().round() as usize)
+                    } else {
+                        existing
+                            .as_ref()
+                            .map_or(Precision::Absent, |s| s.num_rows.clone())
+                    };
+                    let total_byte_size = if inject_byte_size {
+                        // Use the effective row count that was resolved above:
+                        // the provider's own value when present (keeping
+                        // num_rows and total_byte_size internally consistent),
+                        // or the hint's row_count when the provider had none.
+                        let effective_rows = match &num_rows {
+                            Precision::Exact(n) | Precision::Inexact(n) => *n as f64,
+                            Precision::Absent => hints.row_count.unwrap(),
+                        };
+                        Precision::Inexact(
+                            (effective_rows * hints.record_size.unwrap()).round() as usize,
+                        )
+                    } else {
+                        existing
+                            .as_ref()
+                            .map_or(Precision::Absent, |s| s.total_byte_size.clone())
+                    };
+                    let column_statistics = existing
+                        .map(|s| s.column_statistics)
+                        .unwrap_or_else(|| Statistics::unknown_column(&provider.schema()));
+                    let statistics = Statistics {
+                        num_rows,
+                        total_byte_size,
+                        column_statistics,
+                    };
+                    Some(Arc::new(StatisticsOverrideTableProvider {
+                        inner: provider,
+                        statistics,
+                    }) as Arc<dyn TableProvider>)
+                } else {
+                    Some(provider)
+                }
+            }
+            provider => provider,
+        };
+        Ok(provider)
     }
 
     fn get_extensions(&self) -> &Extensions {

datafusion/substrait/src/logical_plan/producer/rel/read_rel.rs

@@ -18,7 +18,9 @@
 use crate::logical_plan::producer::{
     SubstraitProducer, to_substrait_literal, to_substrait_named_struct,
 };
+use datafusion::common::stats::Precision;
 use datafusion::common::{DFSchema, ToDFSchema, substrait_datafusion_err};
+use datafusion::datasource::source_as_provider;
 use datafusion::logical_expr::utils::conjunction;
 use datafusion::logical_expr::{EmptyRelation, Expr, TableScan, Values};
 use datafusion::scalar::ScalarValue;
@@ -29,7 +31,7 @@ use substrait::proto::expression::mask_expression::{StructItem, StructSelect};
 use substrait::proto::expression::nested::Struct as NestedStruct;
 use substrait::proto::read_rel::{NamedTable, ReadType, VirtualTable};
 use substrait::proto::rel::RelType;
-use substrait::proto::{ReadRel, Rel};
+use substrait::proto::{ReadRel, Rel, RelCommon, rel_common};
 
 /// Converts rows of literal expressions into Substrait literal structs.
 ///
@@ -131,9 +133,53 @@ pub fn from_table_scan(
         Some(Box::new(filter_expr))
     };
 
+    // Both Exact and Inexact statistics are serialised as f64 hints.
+    // On the consumer side, hints are always read back as Precision::Inexact,
+    // so Exact values become Inexact after a round-trip — this is an accepted
+    // semantic limitation since RelCommon.hint.stats is advisory.
+    // Zero values write 0.0, which the consumer discards due to the proto3
+    // default-zero ambiguity; that loss is also accepted.
+    //
+    // source_as_provider succeeds for any TableScan (whose source is always a
+    // TableProviderWrapper), so .ok() here only suppresses the error type; it
+    // will never silently discard a real failure in this context.
+    let common = source_as_provider(&scan.source)
+        .ok()
+        .and_then(|provider| provider.statistics())
+        .and_then(|stats| {
+            let row_count = match stats.num_rows {
+                Precision::Exact(n) | Precision::Inexact(n) => Some(n as f64),
+                Precision::Absent => None,
+            };
+            // record_size = total_byte_size / num_rows (both must be present)
+            let record_size = match (stats.total_byte_size, row_count) {
+                (Precision::Exact(b) | Precision::Inexact(b), Some(rows))
+                    if rows > 0.0 =>
+                {
+                    Some(b as f64 / rows)
+                }
+                _ => None,
+            };
+            if row_count.is_none() && record_size.is_none() {
+                return None;
+            }
+            Some(RelCommon {
+                emit_kind: None,
+                hint: Some(rel_common::Hint {
+                    stats: Some(rel_common::hint::Stats {
+                        row_count: row_count.unwrap_or(0.0),
+                        record_size: record_size.unwrap_or(0.0),
+                        ..Default::default()
+                    }),
+                    ..Default::default()
+                }),
+                advanced_extension: None,
+            })
+        });
+
     Ok(Box::new(Rel {
         rel_type: Some(RelType::Read(Box::new(ReadRel {
-            common: None,
+            common,
             base_schema: Some(base_schema),
             filter: filter_option,
             best_effort_filter: None,