Move CBOs and Statistics to physical plan

ballista/rust/core/proto/ballista.proto

@@ -275,6 +275,7 @@ message Statistics {
   int64 num_rows = 1;
   int64 total_byte_size = 2;
   repeated ColumnStats column_stats = 3;
+  bool is_exact = 4;
 }
 
 message PartitionedFile {

ballista/rust/core/src/datasource.rs

@@ -20,7 +20,7 @@ use std::{any::Any, sync::Arc};
 use datafusion::arrow::datatypes::SchemaRef;
 use datafusion::error::Result as DFResult;
 use datafusion::{
-    datasource::{datasource::Statistics, TableProvider},
+    datasource::TableProvider,
     logical_plan::{Expr, LogicalPlan},
     physical_plan::ExecutionPlan,
 };
@@ -61,12 +61,4 @@ impl TableProvider for DfTableAdapter {
     ) -> DFResult<Arc<dyn ExecutionPlan>> {
         Ok(self.plan.clone())
     }
-
-    fn statistics(&self) -> Statistics {
-        Statistics {
-            num_rows: None,
-            total_byte_size: None,
-            column_statistics: None,
-        }
-    }
 }

ballista/rust/core/src/execution_plans/distributed_query.rs

@@ -35,7 +35,7 @@ use datafusion::error::{DataFusionError, Result};
 use datafusion::logical_plan::LogicalPlan;
 use datafusion::physical_plan::{
     DisplayFormatType, ExecutionPlan, Partitioning, RecordBatchStream,
-    SendableRecordBatchStream,
+    SendableRecordBatchStream, Statistics,
 };
 
 use async_trait::async_trait;
@@ -203,6 +203,12 @@ impl ExecutionPlan for DistributedQueryExec {
             }
         }
     }
+
+    fn statistics(&self) -> Statistics {
+        // We cannot infer the statistics until the logical plan
+        // is converted to a physical plan.
+        Statistics::default()
+    }
 }
 
 async fn fetch_partition(

ballista/rust/core/src/execution_plans/shuffle_reader.rs

@@ -21,7 +21,7 @@ use std::{any::Any, pin::Pin};
 
 use crate::client::BallistaClient;
 use crate::memory_stream::MemoryStream;
-use crate::serde::scheduler::PartitionLocation;
+use crate::serde::scheduler::{PartitionLocation, PartitionStats};
 
 use crate::utils::WrappedStream;
 use async_trait::async_trait;
@@ -31,7 +31,9 @@ use datafusion::arrow::record_batch::RecordBatch;
 use datafusion::physical_plan::metrics::{
     ExecutionPlanMetricsSet, MetricBuilder, MetricsSet,
 };
-use datafusion::physical_plan::{DisplayFormatType, ExecutionPlan, Metric, Partitioning};
+use datafusion::physical_plan::{
+    DisplayFormatType, ExecutionPlan, Metric, Partitioning, Statistics,
+};
 use datafusion::{
     error::{DataFusionError, Result},
     physical_plan::RecordBatchStream,
@@ -156,6 +158,38 @@ impl ExecutionPlan for ShuffleReaderExec {
     fn metrics(&self) -> Option<MetricsSet> {
         Some(self.metrics.clone_inner())
     }
+
+    fn statistics(&self) -> Statistics {
+        stats_for_partitions(
+            self.partition
+                .iter()
+                .flatten()
+                .map(|loc| loc.partition_stats),
+        )
+    }
+}
+
+fn stats_for_partitions(
+    partition_stats: impl Iterator<Item = PartitionStats>,
+) -> Statistics {
+    // TODO stats: add column statistics to PartitionStats
+    partition_stats.fold(
+        Statistics {
+            is_exact: true,
+            num_rows: Some(0),
+            total_byte_size: Some(0),
+            column_statistics: None,
+        },
+        |mut acc, part| {
+            // if any statistic is unkown it makes the entire statistic unkown
+            acc.num_rows = acc.num_rows.zip(part.num_rows).map(|(a, b)| a + b as usize);
+            acc.total_byte_size = acc
+                .total_byte_size
+                .zip(part.num_bytes)
+                .map(|(a, b)| a + b as usize);
+            acc
+        },
+    )
 }
 
 async fn fetch_partition(
@@ -177,3 +211,76 @@ async fn fetch_partition(
         .await
         .map_err(|e| DataFusionError::Execution(format!("{:?}", e)))?)
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[tokio::test]
+    async fn test_stats_for_partitions_empty() {
+        let result = stats_for_partitions(std::iter::empty());
+
+        let exptected = Statistics {
+            is_exact: true,
+            num_rows: Some(0),
+            total_byte_size: Some(0),
+            column_statistics: None,
+        };
+
+        assert_eq!(result, exptected);
+    }
+
+    #[tokio::test]
+    async fn test_stats_for_partitions_full() {
+        let part_stats = vec![
+            PartitionStats {
+                num_rows: Some(10),
+                num_bytes: Some(84),
+                num_batches: Some(1),
+            },
+            PartitionStats {
+                num_rows: Some(4),
+                num_bytes: Some(65),
+                num_batches: None,
+            },
+        ];
+
+        let result = stats_for_partitions(part_stats.into_iter());
+
+        let exptected = Statistics {
+            is_exact: true,
+            num_rows: Some(14),
+            total_byte_size: Some(149),
+            column_statistics: None,
+        };
+
+        assert_eq!(result, exptected);
+    }
+
+    #[tokio::test]
+    async fn test_stats_for_partitions_missing() {
+        let part_stats = vec![
+            PartitionStats {
+                num_rows: Some(10),
+                num_bytes: Some(84),
+                num_batches: Some(1),
+            },
+            PartitionStats {
+                num_rows: None,
+                num_bytes: None,
+                num_batches: None,
+            },
+        ];
+
+        let result = stats_for_partitions(part_stats.into_iter());
+
+        let exptected = Statistics {
+            is_exact: true,
+            num_rows: None,
+            total_byte_size: None,
+            column_statistics: None,
+        };
+
+        assert_eq!(result, exptected);
+    }
+}

ballista/rust/core/src/execution_plans/shuffle_writer.rs

@@ -51,7 +51,7 @@ use datafusion::physical_plan::metrics::{
 use datafusion::physical_plan::repartition::RepartitionExec;
 use datafusion::physical_plan::Partitioning::RoundRobinBatch;
 use datafusion::physical_plan::{
-    DisplayFormatType, ExecutionPlan, Metric, Partitioning, RecordBatchStream,
+    DisplayFormatType, ExecutionPlan, Metric, Partitioning, RecordBatchStream, Statistics,
 };
 use futures::StreamExt;
 use hashbrown::HashMap;
@@ -417,6 +417,10 @@ impl ExecutionPlan for ShuffleWriterExec {
             }
         }
     }
+
+    fn statistics(&self) -> Statistics {
+        self.plan.statistics()
+    }
 }
 
 fn result_schema() -> SchemaRef {

ballista/rust/core/src/execution_plans/unresolved_shuffle.rs

@@ -23,7 +23,9 @@ use crate::serde::scheduler::PartitionLocation;
 
 use async_trait::async_trait;
 use datafusion::arrow::datatypes::SchemaRef;
-use datafusion::physical_plan::{DisplayFormatType, ExecutionPlan, Partitioning};
+use datafusion::physical_plan::{
+    DisplayFormatType, ExecutionPlan, Partitioning, Statistics,
+};
 use datafusion::{
     error::{DataFusionError, Result},
     physical_plan::RecordBatchStream,
@@ -117,4 +119,12 @@ impl ExecutionPlan for UnresolvedShuffleExec {
             }
         }
     }
+
+    fn statistics(&self) -> Statistics {
+        // We could try to fetch the statistics here from the shuffle writer,
+        // but it is much more valuable to optimize the plan once this
+        // nodes has been replaced by the actual ShuffleReaderExec which will
+        // have more accurate statistics from its input partitions.
+        Statistics::default()
+    }
 }

ballista/rust/core/src/serde/logical_plan/from_proto.rs

@@ -359,6 +359,7 @@ impl TryInto<Statistics> for &protobuf::Statistics {
             num_rows: Some(self.num_rows as usize),
             total_byte_size: Some(self.total_byte_size as usize),
             column_statistics: Some(column_statistics),
+            is_exact: self.is_exact,
         })
     }
 }
@@ -1177,8 +1178,7 @@ impl TryInto<Field> for &protobuf::Field {
 }
 
 use crate::serde::protobuf::ColumnStats;
-use datafusion::datasource::datasource::{ColumnStatistics, Statistics};
-use datafusion::physical_plan::{aggregates, windows};
+use datafusion::physical_plan::{aggregates, windows, ColumnStatistics, Statistics};
 use datafusion::prelude::{
     array, date_part, date_trunc, length, lower, ltrim, md5, rtrim, sha224, sha256,
     sha384, sha512, trim, upper,

ballista/rust/core/src/serde/logical_plan/to_proto.rs

@@ -25,7 +25,6 @@ use crate::serde::{protobuf, BallistaError};
 use datafusion::arrow::datatypes::{
     DataType, Field, IntervalUnit, Schema, SchemaRef, TimeUnit,
 };
-use datafusion::datasource::datasource::{ColumnStatistics, Statistics};
 use datafusion::datasource::{CsvFile, PartitionedFile, TableDescriptor};
 use datafusion::logical_plan::{
     window_frames::{WindowFrame, WindowFrameBound, WindowFrameUnits},
@@ -36,6 +35,7 @@ use datafusion::physical_plan::functions::BuiltinScalarFunction;
 use datafusion::physical_plan::window_functions::{
     BuiltInWindowFunction, WindowFunction,
 };
+use datafusion::physical_plan::{ColumnStatistics, Statistics};
 use datafusion::{datasource::parquet::ParquetTable, logical_plan::exprlist_to_fields};
 use protobuf::{
     arrow_type, logical_expr_node::ExprType, scalar_type, DateUnit, PrimitiveScalarType,
@@ -278,6 +278,7 @@ impl From<&Statistics> for protobuf::Statistics {
             num_rows: s.num_rows.map(|n| n as i64).unwrap_or(none_value),
             total_byte_size: s.total_byte_size.map(|n| n as i64).unwrap_or(none_value),
             column_stats,
+            is_exact: s.is_exact,
         }
     }
 }

ballista/rust/core/src/serde/physical_plan/from_proto.rs

@@ -34,7 +34,6 @@ use datafusion::arrow::datatypes::{DataType, Schema, SchemaRef};
 use datafusion::catalog::catalog::{
     CatalogList, CatalogProvider, MemoryCatalogList, MemoryCatalogProvider,
 };
-use datafusion::datasource::datasource::Statistics;
 use datafusion::datasource::object_store::ObjectStoreRegistry;
 use datafusion::datasource::FilePartition;
 use datafusion::execution::context::{
@@ -74,7 +73,9 @@ use datafusion::physical_plan::{
     sort::{SortExec, SortOptions},
     Partitioning,
 };
-use datafusion::physical_plan::{AggregateExpr, ExecutionPlan, PhysicalExpr, WindowExpr};
+use datafusion::physical_plan::{
+    AggregateExpr, ExecutionPlan, PhysicalExpr, Statistics, WindowExpr,
+};
 use datafusion::prelude::CsvReadOptions;
 use log::debug;
 use protobuf::physical_expr_node::ExprType;

ballista/rust/executor/src/collect.rs

@@ -28,7 +28,7 @@ use datafusion::arrow::{
 };
 use datafusion::error::DataFusionError;
 use datafusion::physical_plan::{
-    DisplayFormatType, ExecutionPlan, Partitioning, SendableRecordBatchStream,
+    DisplayFormatType, ExecutionPlan, Partitioning, SendableRecordBatchStream, Statistics,
 };
 use datafusion::{error::Result, physical_plan::RecordBatchStream};
 use futures::stream::SelectAll;
@@ -116,6 +116,10 @@ impl ExecutionPlan for CollectExec {
             }
         }
     }
+
+    fn statistics(&self) -> Statistics {
+        self.plan.statistics()
+    }
 }
 
 struct MergedRecordBatchStream {

datafusion/src/datasource/csv.rs

@@ -39,7 +39,6 @@ use std::io::{Read, Seek};
 use std::string::String;
 use std::sync::{Arc, Mutex};
 
-use crate::datasource::datasource::Statistics;
 use crate::datasource::{Source, TableProvider};
 use crate::error::{DataFusionError, Result};
 use crate::logical_plan::Expr;
@@ -54,7 +53,6 @@ pub struct CsvFile {
     has_header: bool,
     delimiter: u8,
     file_extension: String,
-    statistics: Statistics,
 }
 
 impl CsvFile {
@@ -82,7 +80,6 @@ impl CsvFile {
             has_header: options.has_header,
             delimiter: options.delimiter,
             file_extension: String::from(options.file_extension),
-            statistics: Statistics::default(),
         })
     }
 
@@ -105,7 +102,6 @@ impl CsvFile {
             schema,
             has_header: options.has_header,
             delimiter: options.delimiter,
-            statistics: Statistics::default(),
             file_extension: String::new(),
         })
     }
@@ -133,7 +129,6 @@ impl CsvFile {
             schema,
             has_header: options.has_header,
             delimiter: options.delimiter,
-            statistics: Statistics::default(),
             file_extension: String::new(),
         })
     }
@@ -210,10 +205,6 @@ impl TableProvider for CsvFile {
         };
         Ok(Arc::new(exec))
     }
-
-    fn statistics(&self) -> Statistics {
-        self.statistics.clone()
-    }
 }
 
 #[cfg(test)]