Remove compute_statistics free function, let callers manage StatisticsArgs

Callers now create StatisticsArgs directly and call
plan.statistics_with_args(). The cache is created in StatisticsArgs::new()
and shared through compute_child_statistics calls.

Author: asolimando

datafusion/core/src/datasource/file_format/csv.rs

@@ -45,7 +45,8 @@ mod tests {
     use datafusion_datasource::file_format::FileFormat;
     use datafusion_datasource::write::BatchSerializer;
     use datafusion_expr::{col, lit};
-    use datafusion_physical_plan::{ExecutionPlan, collect, compute_statistics};
+    use datafusion_physical_plan::statistics::StatisticsArgs;
+    use datafusion_physical_plan::{ExecutionPlan, collect};
 
     use arrow::array::{
         Array, BooleanArray, Float64Array, Int32Array, RecordBatch, StringArray,
@@ -216,11 +217,13 @@ mod tests {
 
         // test metadata
         assert_eq!(
-            compute_statistics(exec.as_ref(), None)?.num_rows,
+            exec.statistics_with_args(&StatisticsArgs::new(None))?
+                .num_rows,
             Precision::Absent
         );
         assert_eq!(
-            compute_statistics(exec.as_ref(), None)?.total_byte_size,
+            exec.statistics_with_args(&StatisticsArgs::new(None))?
+                .total_byte_size,
             Precision::Absent
         );
 

datafusion/core/src/datasource/file_format/json.rs

@@ -36,7 +36,8 @@ mod tests {
         BatchDeserializer, DecoderDeserializer, DeserializerOutput,
     };
     use datafusion_datasource::file_format::FileFormat;
-    use datafusion_physical_plan::{ExecutionPlan, collect, compute_statistics};
+    use datafusion_physical_plan::statistics::StatisticsArgs;
+    use datafusion_physical_plan::{ExecutionPlan, collect};
 
     use arrow::compute::concat_batches;
     use arrow::datatypes::{DataType, Field};
@@ -118,11 +119,13 @@ mod tests {
 
         // test metadata
         assert_eq!(
-            compute_statistics(exec.as_ref(), None)?.num_rows,
+            exec.statistics_with_args(&StatisticsArgs::new(None))?
+                .num_rows,
             Precision::Absent
         );
         assert_eq!(
-            compute_statistics(exec.as_ref(), None)?.total_byte_size,
+            exec.statistics_with_args(&StatisticsArgs::new(None))?
+                .total_byte_size,
             Precision::Absent
         );
 

datafusion/core/src/datasource/file_format/parquet.rs

@@ -141,8 +141,9 @@ mod tests {
     use datafusion_execution::object_store::ObjectStoreUrl;
     use datafusion_execution::runtime_env::RuntimeEnv;
     use datafusion_expr::dml::InsertOp;
+    use datafusion_physical_plan::statistics::StatisticsArgs;
     use datafusion_physical_plan::stream::RecordBatchStreamAdapter;
-    use datafusion_physical_plan::{ExecutionPlan, collect, compute_statistics};
+    use datafusion_physical_plan::{ExecutionPlan, collect};
 
     use crate::test_util::bounded_stream;
     use arrow::array::{
@@ -715,12 +716,14 @@ mod tests {
 
         // test metadata
         assert_eq!(
-            compute_statistics(exec.as_ref(), None)?.num_rows,
+            exec.statistics_with_args(&StatisticsArgs::new(None))?
+                .num_rows,
             Precision::Exact(8)
         );
         // TODO correct byte size: https://github.com/apache/datafusion/issues/14936
         assert_eq!(
-            compute_statistics(exec.as_ref(), None)?.total_byte_size,
+            exec.statistics_with_args(&StatisticsArgs::new(None))?
+                .total_byte_size,
             Precision::Absent,
         );
 
@@ -764,11 +767,13 @@ mod tests {
 
         // note: even if the limit is set, the executor rounds up to the batch size
         assert_eq!(
-            compute_statistics(exec.as_ref(), None)?.num_rows,
+            exec.statistics_with_args(&StatisticsArgs::new(None))?
+                .num_rows,
             Precision::Exact(8)
         );
         assert_eq!(
-            compute_statistics(exec.as_ref(), None)?.total_byte_size,
+            exec.statistics_with_args(&StatisticsArgs::new(None))?
+                .total_byte_size,
             Precision::Absent,
         );
         let batches = collect(exec, task_ctx).await?;

datafusion/core/src/datasource/listing/table.rs

@@ -144,9 +144,8 @@ mod tests {
     use datafusion_physical_expr::expressions::binary;
     use datafusion_physical_expr_common::sort_expr::LexOrdering;
     use datafusion_physical_plan::empty::EmptyExec;
-    use datafusion_physical_plan::{
-        ExecutionPlanProperties, collect, compute_statistics,
-    };
+    use datafusion_physical_plan::statistics::StatisticsArgs;
+    use datafusion_physical_plan::{ExecutionPlanProperties, collect};
     use std::collections::HashMap;
     use std::io::Write;
     use std::sync::Arc;
@@ -249,11 +248,13 @@ mod tests {
 
         // test metadata
         assert_eq!(
-            compute_statistics(exec.as_ref(), None)?.num_rows,
+            exec.statistics_with_args(&StatisticsArgs::new(None))?
+                .num_rows,
             Precision::Exact(8)
         );
         assert_eq!(
-            compute_statistics(exec.as_ref(), None)?.total_byte_size,
+            exec.statistics_with_args(&StatisticsArgs::new(None))?
+                .total_byte_size,
             Precision::Absent,
         );
 
@@ -1357,13 +1358,17 @@ mod tests {
 
         let exec_default = table_default.scan(&state, None, &[], None).await?;
         assert_eq!(
-            compute_statistics(exec_default.as_ref(), None)?.num_rows,
+            exec_default
+                .statistics_with_args(&StatisticsArgs::new(None))?
+                .num_rows,
             Precision::Absent
         );
 
         // TODO correct byte size: https://github.com/apache/datafusion/issues/14936
         assert_eq!(
-            compute_statistics(exec_default.as_ref(), None)?.total_byte_size,
+            exec_default
+                .statistics_with_args(&StatisticsArgs::new(None))?
+                .total_byte_size,
             Precision::Absent
         );
 
@@ -1378,11 +1383,15 @@ mod tests {
 
         let exec_disabled = table_disabled.scan(&state, None, &[], None).await?;
         assert_eq!(
-            compute_statistics(exec_disabled.as_ref(), None)?.num_rows,
+            exec_disabled
+                .statistics_with_args(&StatisticsArgs::new(None))?
+                .num_rows,
             Precision::Absent
         );
         assert_eq!(
-            compute_statistics(exec_disabled.as_ref(), None)?.total_byte_size,
+            exec_disabled
+                .statistics_with_args(&StatisticsArgs::new(None))?
+                .total_byte_size,
             Precision::Absent
         );
 
@@ -1397,12 +1406,16 @@ mod tests {
 
         let exec_enabled = table_enabled.scan(&state, None, &[], None).await?;
         assert_eq!(
-            compute_statistics(exec_enabled.as_ref(), None)?.num_rows,
+            exec_enabled
+                .statistics_with_args(&StatisticsArgs::new(None))?
+                .num_rows,
             Precision::Exact(8)
         );
         // TODO correct byte size: https://github.com/apache/datafusion/issues/14936
         assert_eq!(
-            compute_statistics(exec_enabled.as_ref(), None)?.total_byte_size,
+            exec_enabled
+                .statistics_with_args(&StatisticsArgs::new(None))?
+                .total_byte_size,
             Precision::Absent,
         );
 

datafusion/core/tests/custom_sources_cases/statistics.rs

@@ -36,7 +36,6 @@ use datafusion_catalog::Session;
 use datafusion_common::{project_schema, stats::Precision};
 use datafusion_physical_expr::EquivalenceProperties;
 use datafusion_physical_plan::StatisticsArgs;
-use datafusion_physical_plan::compute_statistics;
 use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType};
 
 use async_trait::async_trait;
@@ -232,7 +231,10 @@ async fn sql_basic() -> Result<()> {
     let physical_plan = df.create_physical_plan().await.unwrap();
 
     // the statistics should be those of the source
-    assert_eq!(stats, *compute_statistics(physical_plan.as_ref(), None)?);
+    assert_eq!(
+        stats,
+        *physical_plan.statistics_with_args(&StatisticsArgs::new(None))?
+    );
 
     Ok(())
 }
@@ -248,7 +250,7 @@ async fn sql_filter() -> Result<()> {
         .unwrap();
 
     let physical_plan = df.create_physical_plan().await.unwrap();
-    let stats = compute_statistics(physical_plan.as_ref(), None)?;
+    let stats = physical_plan.statistics_with_args(&StatisticsArgs::new(None))?;
     assert_eq!(stats.num_rows, Precision::Inexact(7));
 
     Ok(())
@@ -263,7 +265,7 @@ async fn sql_limit() -> Result<()> {
     let physical_plan = df.create_physical_plan().await.unwrap();
     // when the limit is smaller than the original number of lines we mark the statistics as inexact
     // and cap NDV at the new row count
-    let limit_stats = compute_statistics(physical_plan.as_ref(), None)?;
+    let limit_stats = physical_plan.statistics_with_args(&StatisticsArgs::new(None))?;
     assert_eq!(limit_stats.num_rows, Precision::Exact(5));
     // c1: NDV=2 stays at 2 (already below limit of 5)
     assert_eq!(
@@ -282,7 +284,10 @@ async fn sql_limit() -> Result<()> {
         .unwrap();
     let physical_plan = df.create_physical_plan().await.unwrap();
     // when the limit is larger than the original number of lines, statistics remain unchanged
-    assert_eq!(stats, *compute_statistics(physical_plan.as_ref(), None)?);
+    assert_eq!(
+        stats,
+        *physical_plan.statistics_with_args(&StatisticsArgs::new(None))?
+    );
 
     Ok(())
 }
@@ -299,7 +304,7 @@ async fn sql_window() -> Result<()> {
 
     let physical_plan = df.create_physical_plan().await.unwrap();
 
-    let result = compute_statistics(physical_plan.as_ref(), None)?;
+    let result = physical_plan.statistics_with_args(&StatisticsArgs::new(None))?;
 
     assert_eq!(stats.num_rows, result.num_rows);
     let col_stats = &result.column_statistics;

datafusion/core/tests/parquet/file_statistics.rs

@@ -43,8 +43,8 @@ use datafusion_common::config::ConfigOptions;
 use datafusion_physical_optimizer::PhysicalOptimizerRule;
 use datafusion_physical_optimizer::filter_pushdown::FilterPushdown;
 use datafusion_physical_plan::ExecutionPlan;
-use datafusion_physical_plan::compute_statistics;
 use datafusion_physical_plan::filter::FilterExec;
+use datafusion_physical_plan::statistics::StatisticsArgs;
 use tempfile::tempdir;
 
 #[tokio::test]
@@ -64,7 +64,9 @@ async fn check_stats_precision_with_filter_pushdown() {
     // Scan without filter, stats are exact
     let exec = table.scan(&state, None, &[], None).await.unwrap();
     assert_eq!(
-        compute_statistics(exec.as_ref(), None).unwrap().num_rows,
+        exec.statistics_with_args(&StatisticsArgs::new(None))
+            .unwrap()
+            .num_rows,
         Precision::Exact(8),
         "Stats without filter should be exact"
     );
@@ -96,7 +98,8 @@ async fn check_stats_precision_with_filter_pushdown() {
     );
     // Scan with filter pushdown, stats are inexact
     assert_eq!(
-        compute_statistics(optimized_exec.as_ref(), None)
+        optimized_exec
+            .statistics_with_args(&StatisticsArgs::new(None))
             .unwrap()
             .num_rows,
         Precision::Inexact(8),
@@ -126,11 +129,15 @@ async fn load_table_stats_with_session_level_cache() {
     let exec1 = table1.scan(&state1, None, &[], None).await.unwrap();
 
     assert_eq!(
-        compute_statistics(exec1.as_ref(), None).unwrap().num_rows,
+        exec1
+            .statistics_with_args(&StatisticsArgs::new(None))
+            .unwrap()
+            .num_rows,
         Precision::Exact(8)
     );
     assert_eq!(
-        compute_statistics(exec1.as_ref(), None)
+        exec1
+            .statistics_with_args(&StatisticsArgs::new(None))
             .unwrap()
             .total_byte_size,
         // Byte size is absent because we cannot estimate the output size
@@ -144,11 +151,15 @@ async fn load_table_stats_with_session_level_cache() {
     assert_eq!(get_static_cache_size(&state2), 0);
     let exec2 = table2.scan(&state2, None, &[], None).await.unwrap();
     assert_eq!(
-        compute_statistics(exec2.as_ref(), None).unwrap().num_rows,
+        exec2
+            .statistics_with_args(&StatisticsArgs::new(None))
+            .unwrap()
+            .num_rows,
         Precision::Exact(8)
     );
     assert_eq!(
-        compute_statistics(exec2.as_ref(), None)
+        exec2
+            .statistics_with_args(&StatisticsArgs::new(None))
             .unwrap()
             .total_byte_size,
         // Absent because the data contains variable length columns
@@ -161,11 +172,15 @@ async fn load_table_stats_with_session_level_cache() {
     assert_eq!(get_static_cache_size(&state1), 1);
     let exec3 = table1.scan(&state1, None, &[], None).await.unwrap();
     assert_eq!(
-        compute_statistics(exec3.as_ref(), None).unwrap().num_rows,
+        exec3
+            .statistics_with_args(&StatisticsArgs::new(None))
+            .unwrap()
+            .num_rows,
         Precision::Exact(8)
     );
     assert_eq!(
-        compute_statistics(exec3.as_ref(), None)
+        exec3
+            .statistics_with_args(&StatisticsArgs::new(None))
             .unwrap()
             .total_byte_size,
         // Absent because the data contains variable length columns