fix: clear error instead of panic on mismatched column statistics

datafusion/common/src/stats.rs

@@ -413,6 +413,10 @@ impl NdvFallback {
 impl Default for Statistics {
     /// Returns a new [`Statistics`] instance with all fields set to unknown
     /// and no columns.
+    ///
+    /// Unless the schema has no fields, this is not valid as the statistics of
+    /// a plan, which must hold one [`ColumnStatistics`] per field; use
+    /// [`Statistics::new_unknown`] instead.
     fn default() -> Self {
         Self {
             num_rows: Precision::Absent,

datafusion/core/tests/physical_optimizer/aggregate_statistics.rs

@@ -34,7 +34,7 @@ use datafusion_common::cast::as_int64_array;
 use datafusion_common::config::ConfigOptions;
 use datafusion_common::stats::Precision;
 use datafusion_common::{ColumnStatistics, Result, Statistics};
-use datafusion_common::{ScalarValue, assert_batches_eq};
+use datafusion_common::{ScalarValue, assert_batches_eq, assert_contains};
 use datafusion_datasource::file_scan_config::FileScanConfigBuilder;
 use datafusion_execution::TaskContext;
 use datafusion_execution::object_store::ObjectStoreUrl;
@@ -54,7 +54,7 @@ use datafusion_physical_plan::common;
 use datafusion_physical_plan::displayable;
 use datafusion_physical_plan::filter::FilterExec;
 use datafusion_physical_plan::operator_statistics::{
-    ClosureStatisticsProvider, StatisticsRegistry, StatisticsResult,
+    ClosureStatisticsProvider, ExtendedStatistics, StatisticsRegistry, StatisticsResult,
 };
 use datafusion_physical_plan::projection::ProjectionExec;
 use datafusion_physical_plan::statistics::StatisticsArgs;
@@ -905,3 +905,102 @@ async fn aggregate_statistics_consults_statistics_providers() -> Result<()> {
     );
     Ok(())
 }
+
+/// Returns a context with a table `t(a, b)` whose scans take their statistics
+/// from `scan_statistics`.
+fn context_with_scan_statistics<F>(scan_statistics: F) -> Result<SessionContext>
+where
+    F: Fn(&dyn ExecutionPlan, &[ExtendedStatistics]) -> Result<StatisticsResult>
+        + Send
+        + Sync
+        + 'static,
+{
+    let provider = ClosureStatisticsProvider::with_matches(
+        |plan| plan.is::<DataSourceExec>(),
+        scan_statistics,
+    );
+    let state = SessionStateBuilder::new()
+        .with_default_features()
+        .with_statistics_registry(StatisticsRegistry::with_providers(vec![Arc::new(
+            provider,
+        )]))
+        .build();
+    let ctx = SessionContext::new_with_state(state);
+
+    let schema = Arc::new(Schema::new(vec![
+        Field::new("a", DataType::Int32, true),
+        Field::new("b", DataType::Int32, true),
+    ]));
+    let batch = RecordBatch::try_new(
+        Arc::clone(&schema),
+        vec![
+            Arc::new(Int32Array::from(vec![1, 2])),
+            Arc::new(Int32Array::from(vec![10, 20])),
+        ],
+    )?;
+    ctx.register_batch("t", batch)?;
+    Ok(ctx)
+}
+
+/// A provider reporting fewer column statistics than the scan has columns
+/// breaks the invariant that consumers index by column, so planning fails with
+/// an error naming the provider instead of a consumer reading out of bounds.
+#[tokio::test]
+async fn short_column_statistics_are_reported() -> Result<()> {
+    let ctx = context_with_scan_statistics(|plan, child_stats| {
+        let child_stats = child_stats
+            .iter()
+            .map(|c| Arc::clone(c.base_arc()))
+            .collect::<Vec<_>>();
+        let mut stats = Arc::unwrap_or_clone(
+            plan.statistics_from_inputs(&child_stats, &StatisticsArgs::new())?,
+        );
+        // Drop the trailing column's statistics so every scan reports one
+        // fewer entry than its schema has columns.
+        stats.column_statistics.pop();
+        Ok(StatisticsResult::Computed(stats.into()))
+    })?;
+
+    // `count(b)` reads the statistics; a plain scan fails during planning too.
+    for (sql, entries, fields) in
+        [("SELECT count(b) FROM t", 0, 1), ("SELECT * FROM t", 1, 2)]
+    {
+        let err = ctx.sql(sql).await?.collect().await.expect_err(sql);
+        assert_contains!(
+            err.to_string(),
+            format!(
+                "Statistics provider ClosureStatisticsProvider returned {entries} column statistics for DataSourceExec, expected {fields} (one per field of its schema)"
+            )
+        );
+    }
+    Ok(())
+}
+
+/// A provider reporting the statistics of the whole table for a scan of only
+/// `b` would make `min(b)` and `max(b)` read the statistics of `a`, so planning
+/// fails with an error naming the provider instead of answering from them.
+#[tokio::test]
+async fn table_statistics_for_a_projected_scan_are_reported() -> Result<()> {
+    let column_statistics = |min: i32, max: i32| ColumnStatistics {
+        null_count: Precision::Exact(0),
+        min_value: Precision::Exact(ScalarValue::from(min)),
+        max_value: Precision::Exact(ScalarValue::from(max)),
+        ..ColumnStatistics::new_unknown()
+    };
+    let table_statistics = Statistics {
+        num_rows: Precision::Exact(2),
+        total_byte_size: Precision::Absent,
+        column_statistics: vec![column_statistics(1, 2), column_statistics(10, 20)],
+    };
+    let ctx = context_with_scan_statistics(move |_plan, _child_stats| {
+        Ok(StatisticsResult::Computed(table_statistics.clone().into()))
+    })?;
+
+    let sql = "SELECT min(b), max(b) FROM t";
+    let err = ctx.sql(sql).await?.collect().await.expect_err(sql);
+    assert_contains!(
+        err.to_string(),
+        "Statistics provider ClosureStatisticsProvider returned 2 column statistics for DataSourceExec, expected 1 (one per field of its schema)"
+    );
+    Ok(())
+}

datafusion/functions-aggregate/src/count.rs

@@ -414,7 +414,8 @@ impl AggregateUDFImpl for Count {
             // Only column references can be resolved from statistics;
             // expressions like casts or literals are not supported.
             let col_expr = expr.downcast_ref::<expressions::Column>()?;
-            if let Precision::Exact(dc) = col_stats[col_expr.index()].distinct_count {
+            if let Precision::Exact(dc) = col_stats.get(col_expr.index())?.distinct_count
+            {
                 let dc = i64::try_from(dc).ok()?;
                 return Some(ScalarValue::Int64(Some(dc)));
             }
@@ -427,7 +428,7 @@ impl AggregateUDFImpl for Count {
 
         // TODO optimize with exprs other than Column
         if let Some(col_expr) = expr.downcast_ref::<expressions::Column>() {
-            if let Precision::Exact(val) = col_stats[col_expr.index()].null_count {
+            if let Precision::Exact(val) = col_stats.get(col_expr.index())?.null_count {
                 let count = i64::try_from(num_rows - val).ok()?;
                 return Some(ScalarValue::Int64(Some(count)));
             }
@@ -1204,4 +1205,28 @@ mod tests {
         assert_eq!(merged.evaluate()?, ScalarValue::Int64(Some(3)));
         Ok(())
     }
+
+    /// `value_from_stats` is public and has no error channel, so a direct
+    /// caller passing fewer column statistics than the schema has fields gets
+    /// `None` instead of a panic, as with `sum`. The statistics walk rejects
+    /// such statistics during planning.
+    #[test]
+    fn count_value_from_stats_out_of_bounds_column_returns_none() {
+        let statistics = datafusion_common::Statistics {
+            num_rows: Precision::Exact(10),
+            total_byte_size: Precision::Absent,
+            column_statistics: vec![datafusion_common::ColumnStatistics::default()],
+        };
+        let return_type = DataType::Int64;
+        let exprs: Vec<Arc<dyn PhysicalExpr>> = vec![Arc::new(Column::new("b", 1))];
+        for is_distinct in [false, true] {
+            let statistics_args = StatisticsArgs {
+                statistics: &statistics,
+                return_type: &return_type,
+                is_distinct,
+                exprs: &exprs,
+            };
+            assert_eq!(Count::new().value_from_stats(&statistics_args), None);
+        }
+    }
 }

datafusion/functions-aggregate/src/min_max.rs

@@ -193,7 +193,7 @@ trait FromColumnStatistics {
                             statistics_args.exprs[0].downcast_ref::<expressions::Column>()
                         {
                             return self.value_from_column_statistics(
-                                &col_stats[col_expr.index()],
+                                col_stats.get(col_expr.index())?,
                             );
                         }
                     }
@@ -1722,4 +1722,26 @@ mod tests {
 
         Ok(())
     }
+
+    /// `value_from_stats` is public and has no error channel, so a direct
+    /// caller passing fewer column statistics than the schema has fields gets
+    /// `None` instead of a panic, as with `count` and `sum`.
+    #[test]
+    fn min_max_value_from_stats_out_of_bounds_column_returns_none() {
+        let statistics = datafusion_common::Statistics {
+            num_rows: Precision::Exact(10),
+            total_byte_size: Precision::Absent,
+            column_statistics: vec![ColumnStatistics::new_unknown()],
+        };
+        let exprs: Vec<Arc<dyn datafusion_physical_expr::PhysicalExpr>> =
+            vec![Arc::new(expressions::Column::new("b", 1))];
+        let statistics_args = StatisticsArgs {
+            statistics: &statistics,
+            return_type: &DataType::Int32,
+            is_distinct: false,
+            exprs: &exprs,
+        };
+        assert_eq!(Min::new().value_from_stats(&statistics_args), None);
+        assert_eq!(Max::new().value_from_stats(&statistics_args), None);
+    }
 }

datafusion/physical-plan/src/execution_plan.rs

@@ -740,11 +740,15 @@ pub trait ExecutionPlan: Any + Debug + DisplayAs + Send + Sync {
     /// (the default), not an error.
     /// If `args.partition()` is `None`, it returns statistics for all partitions.
     ///
+    /// The result must hold one [`ColumnStatistics`] per field of
+    /// [`Self::schema`]; [`StatisticsContext::compute`] returns an error otherwise.
+    ///
     /// Implementations should not call [`StatisticsContext::compute`] from within
     /// this method; child statistics are provided via `input_stats`.
     ///
     /// Use [`StatisticsContext::compute`] to initiate a full plan-tree walk.
     ///
+    /// [`ColumnStatistics`]: datafusion_common::ColumnStatistics
     /// [`StatisticsContext::compute`]: crate::statistics::StatisticsContext::compute
     fn statistics_from_inputs(
         &self,

datafusion/physical-plan/src/operator_statistics/mod.rs

@@ -263,6 +263,12 @@ pub enum StatisticsResult {
 /// `COUNT(*)`, `MIN` and `MAX` without scanning and remove limits. Return
 /// `Inexact` for anything that is an estimate.
 ///
+/// Computed statistics must also hold one [`ColumnStatistics`] per field of the
+/// node's schema, or the statistics walk returns an error. To override only
+/// some values, start from [`Statistics::new_unknown`].
+///
+/// [`ColumnStatistics`]: datafusion_common::ColumnStatistics
+///
 /// # Implementing a Custom Provider
 ///
 /// ```ignore
@@ -1224,7 +1230,9 @@ mod tests {
     };
     use arrow::datatypes::{DataType, Field, Schema};
     use datafusion_common::stats::Precision;
-    use datafusion_common::{ColumnStatistics, ScalarValue, internal_err};
+    use datafusion_common::{
+        ColumnStatistics, ScalarValue, assert_contains, internal_err,
+    };
     use datafusion_execution::TaskContext;
     use datafusion_expr::Operator;
     use datafusion_physical_expr::PhysicalExpr;
@@ -2777,11 +2785,8 @@ mod tests {
         let provider = ClosureStatisticsProvider::new(|plan, _child_stats| {
             if plan.downcast_ref::<FilterExec>().is_some() {
                 Ok(StatisticsResult::Computed(ExtendedStatistics::from(
-                    Statistics {
-                        num_rows: Precision::Inexact(42),
-                        total_byte_size: Precision::Absent,
-                        column_statistics: vec![],
-                    },
+                    Statistics::new_unknown(&plan.schema())
+                        .with_num_rows(Precision::Inexact(42)),
                 )))
             } else {
                 Ok(StatisticsResult::Delegate)
@@ -2810,18 +2815,12 @@ mod tests {
             }
             match child_stats[0].base.num_rows.get_value().copied() {
                 Some(500) => Ok(StatisticsResult::Computed(ExtendedStatistics::from(
-                    Statistics {
-                        num_rows: Precision::Inexact(100),
-                        total_byte_size: Precision::Absent,
-                        column_statistics: vec![],
-                    },
+                    Statistics::new_unknown(&plan.schema())
+                        .with_num_rows(Precision::Inexact(100)),
                 ))),
                 Some(200) => Ok(StatisticsResult::Computed(ExtendedStatistics::from(
-                    Statistics {
-                        num_rows: Precision::Inexact(50),
-                        total_byte_size: Precision::Absent,
-                        column_statistics: vec![],
-                    },
+                    Statistics::new_unknown(&plan.schema())
+                        .with_num_rows(Precision::Inexact(50)),
                 ))),
                 _ => Ok(StatisticsResult::Delegate),
             }
@@ -2841,4 +2840,49 @@ mod tests {
         assert_eq!(stats_b.base.num_rows, Precision::Inexact(50));
         Ok(())
     }
+
+    /// Consumers index `column_statistics` by column, so the walk rejects an
+    /// operator or a provider that does not report one entry per field.
+    #[test]
+    fn test_wrong_column_statistics_count_is_rejected() {
+        // Two fields: one entry too few, and one too many.
+        for num_columns in [1, 3] {
+            let source = MockSourceExec::with_column_stats(
+                make_schema(),
+                Precision::Exact(10),
+                vec![ColumnStatistics::new_unknown(); num_columns],
+            );
+            let err = compute(&StatisticsRegistry::new(), &source).unwrap_err();
+            assert_contains!(
+                err.to_string(),
+                format!(
+                    "MockSourceExec returned {num_columns} column statistics, expected 2 (one per field of its schema)"
+                )
+            );
+        }
+
+        // No entries, and one too many.
+        for num_columns in [0, 3] {
+            let provider = ClosureStatisticsProvider::new(move |_plan, _child_stats| {
+                Ok(StatisticsResult::Computed(ExtendedStatistics::from(
+                    Statistics {
+                        num_rows: Precision::Exact(10),
+                        total_byte_size: Precision::Absent,
+                        column_statistics: vec![
+                            ColumnStatistics::new_unknown();
+                            num_columns
+                        ],
+                    },
+                )))
+            });
+            let registry = StatisticsRegistry::with_providers(vec![Arc::new(provider)]);
+            let err = compute(&registry, make_source(10).as_ref()).unwrap_err();
+            assert_contains!(
+                err.to_string(),
+                format!(
+                    "Statistics provider ClosureStatisticsProvider returned {num_columns} column statistics for MockSourceExec, expected 2 (one per field of its schema)"
+                )
+            );
+        }
+    }
 }