Add comprehensive metrics instrumentation for scheduler and executor (#10)

* Add shuffle read metrics extraction and QueryStageExecutor::plan() method

- Add public getter methods to PartitionStats (num_rows, num_batches, num_bytes)
- Extend QueryStageExecutor trait with plan() method to access underlying ExecutionPlan
- Add extract_shuffle_read_metrics() to walk plan tree and sum ShuffleReaderExec partition stats
- Record shuffle read metrics (bytes, rows, duration) after successful task execution in executor

* Add shuffle locality metrics to ExecutorMetricsCollector, SchedulerMetricsCollector, and ShuffleReaderExec

- Add record_shuffle_read_local/remote methods to ExecutorMetricsCollector trait
- Add record_task_shuffle_affinity_hit/miss methods to SchedulerMetricsCollector trait
- Add ShuffleReadMetricsCallback trait in ballista-core for tracking local vs remote reads
- Instrument shuffle_reader.rs to call metrics callback during partition fetches
- Add SessionConfigExt methods to pass metrics callback via session config

* Add metrics collector to SchedulerState and instrument executor and planning metrics

- Add metrics_collector field to SchedulerState struct
- Instrument record_planning_duration in submit_job
- Instrument record_executor_registered/deregistered and set_active_executor_count
- Update all SchedulerState constructors and call sites

* Add stage and task lifecycle metrics instrumentation to update_task_status flow

* Add shuffle affinity metrics to scheduler task binding

* Add actual task scheduling latency tracking

- Add schedulable_time_millis field to TaskDescription to track when a task became schedulable (when its stage transitioned to running state)
- Update all TaskDescription creation sites to pass RunningStage.stage_running_time
- Calculate actual scheduling latency in record_task_scheduled calls by computing the difference between current time and schedulable_time_millis
- This enables accurate scheduler_task_scheduling_latency_ms metrics instead of the previous placeholder value of 0

* fix lint

Author: phillipleblanc

ballista/core/src/client.rs

@@ -47,8 +47,6 @@ use datafusion::error::Result;
 use crate::extension::BallistaConfigGrpcEndpoint;
 use crate::serde::protobuf;
 
-use crate::utils::{GrpcClientConfig, create_grpc_client_connection};
-
 use crate::utils::create_grpc_client_endpoint;
 
 use datafusion::physical_plan::{RecordBatchStream, SendableRecordBatchStream};

ballista/core/src/execution_plans/distributed_query.rs

@@ -18,7 +18,8 @@
 use crate::client::BallistaClient;
 use crate::config::BallistaConfig;
 use crate::extension::{
-    BallistaConfigGrpcEndpoint, BallistaGrpcMetadataInterceptor, SessionConfigExt,
+    BallistaConfigGrpcEndpoint, BallistaGrpcMetadataInterceptor,
+    ResultFetchMetricsCallback, SessionConfigExt,
 };
 use crate::serde::protobuf::SuccessfulJob;
 use crate::serde::protobuf::{
@@ -27,8 +28,6 @@ use crate::serde::protobuf::{
     scheduler_grpc_client::SchedulerGrpcClient,
 };
 
-use crate::utils::{GrpcClientConfig, create_grpc_client_connection};
-
 use crate::utils::create_grpc_client_endpoint;
 
 use datafusion::arrow::datatypes::SchemaRef;
@@ -248,8 +247,6 @@ impl<T: 'static + AsLogicalPlan> ExecutionPlan for DistributedQueryExec<T> {
         let metric_total_bytes =
             MetricBuilder::new(&self.metrics).counter("transferred_bytes", partition);
 
-
-
         let interceptor = context.session_config().ballista_grpc_interceptor();
 
         let customize_endpoint = context
@@ -258,23 +255,22 @@ impl<T: 'static + AsLogicalPlan> ExecutionPlan for DistributedQueryExec<T> {
 
         let use_tls = context.session_config().ballista_use_tls();
 
+        let result_fetch_callback = context
+            .session_config()
+            .ballista_result_fetch_metrics_callback();
 
         let stream = futures::stream::once(
             execute_query(
                 self.scheduler_url.clone(),
                 self.session_id.clone(),
                 query,
-
-                self.config.default_grpc_client_max_message_size(),
-                GrpcClientConfig::from(&self.config),
                 Arc::new(self.metrics.clone()),
                 partition,
-
                 self.config.clone(),
                 interceptor,
                 customize_endpoint,
                 use_tls,
-
+                result_fetch_callback,
             )
             .map_err(|e| ArrowError::ExternalError(Box::new(e))),
         )
@@ -306,21 +302,18 @@ impl<T: 'static + AsLogicalPlan> ExecutionPlan for DistributedQueryExec<T> {
     }
 }
 
+#[allow(clippy::too_many_arguments)]
 async fn execute_query(
     scheduler_url: String,
     session_id: String,
     query: ExecuteQueryParams,
-
-    max_message_size: usize,
-    grpc_config: GrpcClientConfig,
     metrics: Arc<ExecutionPlanMetricsSet>,
     partition: usize,
-
     config: BallistaConfig,
     grpc_interceptor: Arc<BallistaGrpcMetadataInterceptor>,
     customize_endpoint: Option<Arc<BallistaConfigGrpcEndpoint>>,
     use_tls: bool,
-
+    result_fetch_callback: Option<Arc<dyn ResultFetchMetricsCallback>>,
 ) -> Result<impl Stream<Item = Result<RecordBatch>> + Send> {
     // Capture query submission time for total_query_time_ms
     let query_start_time = std::time::Instant::now();
@@ -450,12 +443,14 @@ async fn execute_query(
                 // This could be added in a future enhancement by wrapping the stream.
 
                 let streams = partition_location.into_iter().map(move |partition| {
+                    let callback = result_fetch_callback.clone();
                     let f = fetch_partition(
                         partition,
                         max_message_size,
                         true,
                         customize_endpoint.clone(),
                         use_tls,
+                        callback,
                     )
                     .map_err(|e| ArrowError::ExternalError(Box::new(e)));
 
@@ -474,13 +469,29 @@ async fn fetch_partition(
     flight_transport: bool,
     customize_endpoint: Option<Arc<BallistaConfigGrpcEndpoint>>,
     use_tls: bool,
+    metrics_callback: Option<Arc<dyn ResultFetchMetricsCallback>>,
 ) -> Result<SendableRecordBatchStream> {
+    let start_time = std::time::Instant::now();
+
     let metadata = location.executor_meta.ok_or_else(|| {
         DataFusionError::Internal("Received empty executor metadata".to_owned())
     })?;
     let partition_id = location.partition_id.ok_or_else(|| {
         DataFusionError::Internal("Received empty partition id".to_owned())
     })?;
+
+    // Extract stats before consuming location
+    let stats = location.partition_stats.as_ref();
+    #[expect(clippy::cast_sign_loss)]
+    let expected_bytes = stats.map(|s| s.num_bytes as u64).unwrap_or(0);
+    #[expect(clippy::cast_sign_loss)]
+    let expected_rows = stats.map(|s| s.num_rows as u64).unwrap_or(0);
+
+    let job_id = partition_id.job_id.clone();
+    let stage_id = partition_id.stage_id as usize;
+    let partition = partition_id.partition_id as usize;
+    let executor_id = metadata.id.clone();
+
     let host = metadata.host.as_str();
     let port = metadata.port as u16;
     let mut ballista_client = BallistaClient::try_new(
@@ -492,7 +503,8 @@ async fn fetch_partition(
     )
     .await
     .map_err(|e| DataFusionError::Execution(format!("{e:?}")))?;
-    ballista_client
+
+    let stream = ballista_client
         .fetch_partition(
             &metadata.id,
             &partition_id.into(),
@@ -502,5 +514,21 @@ async fn fetch_partition(
             flight_transport,
         )
         .await
-        .map_err(|e| DataFusionError::External(Box::new(e)))
+        .map_err(|e| DataFusionError::External(Box::new(e)))?;
+
+    // Record metrics after successful fetch
+    if let Some(callback) = metrics_callback {
+        let duration_ms = start_time.elapsed().as_millis() as u64;
+        callback.record_result_fetch(
+            &job_id,
+            stage_id,
+            partition,
+            &executor_id,
+            expected_bytes,
+            expected_rows,
+            duration_ms,
+        );
+    }
+
+    Ok(stream)
 }