feat: Column statistics support v1

ballista/core/proto/ballista.proto

@@ -643,6 +643,12 @@ message FailedTask {
   }
 }
 
+message TaskColumnStats {
+  uint32 column      = 1;
+  uint64 null_count  = 2;
+  bytes  hll_sketch  = 3;
+}
+
 message SuccessfulTask {
   string executor_id = 1;
   // TODO tasks are currently always shuffle writes but this will not always be the case
@@ -656,6 +662,7 @@ message SuccessfulTask {
   // executed `RuntimeStatsExec`; the scheduler groups by `order_by` tag
   // to combine reports across tasks/executors.
   repeated RuntimeStatsReport runtime_stats = 3;
+  repeated TaskColumnStats taskColumnStats = 4;
 }
 
 // One report per `RuntimeStatsExec` in the executed plan.

ballista/core/src/execution_plans/mod.rs

@@ -67,6 +67,19 @@ pub use unordered_range_repartition::UnorderedRangeRepartitionExec;
 pub use unresolved_shuffle::UnresolvedShuffleExec;
 
 use crate::JobId;
+use crate::serde::protobuf::{ShuffleWritePartition, TaskColumnStats};
+
+/// Result of a completed shuffle-write task, returned from the executor to
+/// the scheduler: per-partition file summaries plus per-column statistics
+/// folded across the task's output.
+#[derive(Debug, Clone)]
+pub struct ShuffleWriteResult {
+    /// Per-output-partition file summaries (location, row/batch/byte counts).
+    pub partitions: Vec<ShuffleWritePartition>,
+    /// Per-column statistics aggregated over this task's output. Empty until
+    /// collection is wired into the shuffle writers.
+    pub column_stats: Vec<TaskColumnStats>,
+}
 
 /// Creates the file path for a shuffle output partition.
 ///

ballista/core/src/serde/generated/ballista.rs

@@ -943,6 +943,15 @@ pub mod failed_task {
         TaskKilled(super::TaskKilled),
     }
 }
+#[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)]
+pub struct TaskColumnStats {
+    #[prost(uint32, tag = "1")]
+    pub column: u32,
+    #[prost(uint64, tag = "2")]
+    pub null_count: u64,
+    #[prost(bytes = "vec", tag = "3")]
+    pub hll_sketch: ::prost::alloc::vec::Vec<u8>,
+}
 #[derive(Clone, PartialEq, ::prost::Message)]
 pub struct SuccessfulTask {
     #[prost(string, tag = "1")]
@@ -960,6 +969,8 @@ pub struct SuccessfulTask {
     /// to combine reports across tasks/executors.
     #[prost(message, repeated, tag = "3")]
     pub runtime_stats: ::prost::alloc::vec::Vec<RuntimeStatsReport>,
+    #[prost(message, repeated, tag = "4")]
+    pub task_column_stats: ::prost::alloc::vec::Vec<TaskColumnStats>,
 }
 /// One report per `RuntimeStatsExec` in the executed plan.
 #[derive(Clone, PartialEq, ::prost::Message)]

ballista/executor/src/execution_engine.rs

@@ -24,7 +24,7 @@
 use ballista_core::client_pool::BallistaClientPool;
 use ballista_core::execution_plans::sort_shuffle::SortShuffleWriterExec;
 use ballista_core::execution_plans::{
-    RangeShuffleReaderExec, ShuffleReaderExec, ShuffleWriterExec,
+    RangeShuffleReaderExec, ShuffleReaderExec, ShuffleWriteResult, ShuffleWriterExec,
 };
 use ballista_core::serde::protobuf::ShuffleWritePartition;
 use ballista_core::serde::scheduler::PartitionStats;
@@ -84,7 +84,7 @@ pub trait QueryStageExecutor: Sync + Send + Debug + Display {
         &self,
         task_id: usize,
         context: Arc<TaskContext>,
-    ) -> Result<Vec<ShuffleWritePartition>>;
+    ) -> Result<ShuffleWriteResult>;
 
     /// Collects execution metrics from all operators in the plan.
     fn collect_plan_metrics(&self) -> Vec<MetricsSet>;
@@ -279,7 +279,7 @@ impl QueryStageExecutor for DefaultQueryStageExec {
         &self,
         task_id: usize,
         context: Arc<TaskContext>,
-    ) -> Result<Vec<ShuffleWritePartition>> {
+    ) -> Result<ShuffleWriteResult> {
         let (plan_arc, is_sort_shuffle): (Arc<dyn ExecutionPlan>, bool) =
             match &self.shuffle_writer {
                 ShuffleWriterVariant::Passthrough(writer) => {
@@ -304,7 +304,10 @@ impl QueryStageExecutor for DefaultQueryStageExec {
             result.is_ok(),
             DisplayableExecutionPlan::with_metrics(plan_arc.as_ref()).indent(true)
         );
-        result
+        result.map(|partitions| ShuffleWriteResult {
+            partitions,
+            column_stats: vec![],
+        })
     }
 
     fn collect_plan_metrics(&self) -> Vec<MetricsSet> {

ballista/executor/src/executor.rs

@@ -28,8 +28,8 @@ use ballista_core::ConfigProducer;
 use ballista_core::JobId;
 use ballista_core::RuntimeProducer;
 use ballista_core::error::BallistaError;
+use ballista_core::execution_plans::ShuffleWriteResult;
 use ballista_core::registry::BallistaFunctionRegistry;
-use ballista_core::serde::protobuf;
 use ballista_core::serde::protobuf::ExecutorRegistration;
 use ballista_core::serde::scheduler::TaskKey;
 use dashmap::DashMap;
@@ -228,7 +228,7 @@ impl Executor {
         key: TaskKey,
         query_stage_exec: Arc<dyn QueryStageExecutor>,
         task_ctx: Arc<TaskContext>,
-    ) -> Result<Vec<protobuf::ShuffleWritePartition>, BallistaError> {
+    ) -> Result<ShuffleWriteResult, BallistaError> {
         let (task, abort_handle) = futures::future::abortable(
             query_stage_exec.execute_query_stage(key.task_id, task_ctx),
         );

ballista/executor/src/lib.rs

@@ -59,15 +59,15 @@ pub use standalone::new_standalone_executor;
 pub use standalone::new_standalone_executor_from_builder;
 pub use standalone::new_standalone_executor_from_state;
 
-use log::info;
-
 use crate::shutdown::Shutdown;
+use ballista_core::execution_plans::ShuffleWriteResult;
 use ballista_core::serde::protobuf::{
-    FailedTask, OperatorMetricsSet, RuntimeStatsReport, ShuffleWritePartition,
-    SuccessfulTask, TaskStatus, task_status,
+    FailedTask, OperatorMetricsSet, RuntimeStatsReport, SuccessfulTask, TaskStatus,
+    task_status,
 };
 use ballista_core::serde::scheduler::TaskKey;
 use ballista_core::utils::GrpcServerConfig;
+use log::info;
 
 /// [ArrowFlightServerProvider] provides a function which creates a new Arrow Flight server.
 ///
@@ -119,7 +119,7 @@ pub struct TaskCompletionExtras {
 /// along with timing and metrics information into a status message that
 /// can be sent back to the scheduler.
 pub fn as_task_status(
-    execution_result: ballista_core::error::Result<Vec<ShuffleWritePartition>>,
+    execution_result: Result<ShuffleWriteResult, BallistaError>,
     executor_id: String,
     stage_attempt_num: usize,
     key: TaskKey,
@@ -133,13 +133,15 @@ pub fn as_task_status(
     let metrics = operator_metrics.unwrap_or_default();
     let task_id = key.task_id;
     match execution_result {
-        Ok(partitions) => {
+        Ok(shuffle_write_result) => {
             debug!(
                 "Task {task_id} finished with operator_metrics array size {} \
                  and {} runtime-stats report(s)",
                 metrics.len(),
                 runtime_stats.len(),
             );
+            let partition = shuffle_write_result.partitions;
+            let col_stats = shuffle_write_result.column_stats;
             TaskStatus {
                 task_id: task_id as u32,
                 job_id: key.job_id.clone().into(),
@@ -151,8 +153,9 @@ pub fn as_task_status(
                 metrics,
                 status: Some(task_status::Status::Successful(SuccessfulTask {
                     executor_id,
-                    partitions,
+                    partitions: partition,
                     runtime_stats,
+                    task_column_stats: col_stats,
                 })),
             }
         }

ballista/scheduler/src/scheduler_server/mod.rs

@@ -686,6 +686,7 @@ mod test {
                         executor_id: "executor-1".to_owned(),
                         partitions,
                         runtime_stats: vec![],
+                        task_column_stats: vec![],
                     })),
                 };
 

ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs

@@ -867,6 +867,7 @@ mod tests {
                         })
                         .collect(),
                     runtime_stats: vec![],
+                    task_column_stats: vec![],
                 })),
             })
             .collect();

ballista/scheduler/src/state/execution_stage.rs

@@ -15,14 +15,15 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use datafusion::common::ColumnStatistics;
+use datafusion::common::stats::Precision;
+use datafusion::config::ConfigOptions;
+use datafusion::physical_optimizer::aggregate_statistics::AggregateStatistics;
 use std::collections::{HashMap, HashSet, VecDeque};
 use std::convert::TryInto;
 use std::fmt::{Debug, Formatter};
 use std::sync::Arc;
 use std::time::{SystemTime, UNIX_EPOCH};
-
-use datafusion::config::ConfigOptions;
-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;
@@ -254,6 +255,8 @@ pub struct SuccessfulStage {
     pub stage_metrics: Vec<MetricsSet>,
     /// [SessionConfig] used for this stage
     pub session_config: Arc<SessionConfig>,
+    /// container to store col stats per stage
+    pub output_column_stats: Vec<ColumnStatistics>,
 }
 
 /// If a stage fails, it will be with an error message
@@ -651,6 +654,24 @@ impl RunningStage {
             warn!("The metrics for stage {} should not be none", self.stage_id);
             vec![]
         });
+        let mut output_column_stats: Vec<ColumnStatistics> = vec![];
+        for info in self.task_infos.iter() {
+            if let task_status::Status::Successful(task_status) = &info.task_status {
+                if output_column_stats.is_empty() {
+                    output_column_stats = vec![
+                        ColumnStatistics::new_unknown();
+                        task_status.task_column_stats.len()
+                    ];
+                }
+                for col_stats in &task_status.task_column_stats {
+                    let slot = &mut output_column_stats[col_stats.column as usize];
+                    slot.null_count = slot
+                        .null_count
+                        .add(&Precision::Exact(col_stats.null_count as usize));
+                }
+            }
+        }
+
         SuccessfulStage {
             stage_id: self.stage_id,
             stage_attempt_num: self.stage_attempt_num,
@@ -661,6 +682,7 @@ impl RunningStage {
             task_infos,
             stage_metrics,
             session_config: self.session_config.clone(),
+            output_column_stats,
         }
     }
 
@@ -1405,6 +1427,7 @@ mod tests {
                 executor_id: "executor-1".to_string(),
                 partitions: vec![],
                 runtime_stats: vec![],
+                task_column_stats: vec![],
             })),
             metrics: vec![],
         }
@@ -1804,6 +1827,7 @@ mod tests {
                     executor_id: executor.to_string(),
                     partitions: vec![],
                     runtime_stats: vec![],
+                    task_column_stats: vec![],
                 });
             stage.append_runtime_stats_reports(
                 task_id,

ballista/scheduler/src/test_utils.rs

@@ -309,6 +309,7 @@ pub fn default_task_runner() -> impl TaskRunner {
                     executor_id: executor_id.clone(),
                     partitions: partitions.clone(),
                     runtime_stats: vec![],
+                    task_column_stats: vec![],
                 })),
             });
         }
@@ -1265,6 +1266,7 @@ pub fn mock_completed_task(task: TaskDescription, executor_id: &str) -> TaskStat
             executor_id: executor_id.to_owned(),
             partitions,
             runtime_stats: vec![],
+            task_column_stats: vec![],
         })),
     }
 }