Dynamically optimize aggregate based on shuffle stats (#919)

ballista/scheduler/src/state/execution_graph/execution_stage.rs

@@ -22,6 +22,7 @@ use std::iter::FromIterator;
 use std::sync::Arc;
 use std::time::{SystemTime, UNIX_EPOCH};
 
+use datafusion::physical_optimizer::aggregate_statistics::AggregateStatistics;
 use datafusion::physical_optimizer::join_selection::JoinSelection;
 use datafusion::physical_optimizer::PhysicalOptimizerRule;
 use datafusion::physical_plan::display::DisplayableExecutionPlan;
@@ -348,9 +349,13 @@ impl UnresolvedStage {
             &input_locations,
         )?;
 
-        // Optimize join order based on new resolved statistics
+        // Optimize join order and statistics based on new resolved statistics
         let optimize_join = JoinSelection::new();
-        let plan = optimize_join.optimize(plan, SessionConfig::default().options())?;
+        let config = SessionConfig::default();
+        let plan = optimize_join.optimize(plan, config.options())?;
+        let optimize_aggregate = AggregateStatistics::new();
+        let plan =
+            optimize_aggregate.optimize(plan, SessionConfig::default().options())?;
 
         Ok(ResolvedStage::new(
             self.stage_id,