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

Author: phillipleblanc

ballista/core/src/execution_plans/shuffle_reader.rs

@@ -29,7 +29,9 @@ use std::sync::Arc;
 use std::task::{Context, Poll};
 
 use crate::client::BallistaClient;
-use crate::extension::{BallistaConfigGrpcEndpoint, SessionConfigExt};
+use crate::extension::{
+    BallistaConfigGrpcEndpoint, SessionConfigExt, ShuffleReadMetricsCallback,
+};
 use crate::serde::scheduler::{PartitionLocation, PartitionStats};
 
 use datafusion::arrow::datatypes::SchemaRef;
@@ -164,6 +166,7 @@ impl ExecutionPlan for ShuffleReaderExec {
         let prefer_flight = config.ballista_shuffle_reader_remote_prefer_flight();
         let customize_endpoint = config.ballista_override_create_grpc_client_endpoint();
         let use_tls = config.ballista_use_tls();
+        let metrics_callback = config.ballista_shuffle_read_metrics_callback();
 
         if force_remote_read {
             debug!(
@@ -199,6 +202,7 @@ impl ExecutionPlan for ShuffleReaderExec {
             prefer_flight,
             customize_endpoint,
             use_tls,
+            metrics_callback,
         );
 
         let result = RecordBatchStreamAdapter::new(
@@ -396,6 +400,7 @@ fn send_fetch_partitions(
     flight_transport: bool,
     customize_endpoint: Option<Arc<BallistaConfigGrpcEndpoint>>,
     use_tls: bool,
+    metrics_callback: Option<Arc<dyn ShuffleReadMetricsCallback>>,
 ) -> AbortableReceiverStream {
     let (response_sender, response_receiver) = mpsc::channel(max_request_num);
     let semaphore = Arc::new(Semaphore::new(max_request_num));
@@ -413,8 +418,10 @@ fn send_fetch_partitions(
     // keep local shuffle files reading in serial order for memory control.
     let response_sender_c = response_sender.clone();
     let customize_endpoint_c = customize_endpoint.clone();
+    let metrics_callback_c = metrics_callback.clone();
     spawned_tasks.push(SpawnedTask::spawn(async move {
         for p in local_locations {
+            let start_time = std::time::Instant::now();
             let r = PartitionReaderEnum::Local
                 .fetch_partition(
                     &p,
@@ -424,6 +431,25 @@ fn send_fetch_partitions(
                     use_tls,
                 )
                 .await;
+
+            // Record local read metrics if callback is set and read succeeded
+            if r.is_ok() {
+                if let Some(ref callback) = metrics_callback_c {
+                    let duration_ms = start_time.elapsed().as_millis() as u64;
+                    let bytes = p.partition_stats.num_bytes().unwrap_or(0);
+                    let rows = p.partition_stats.num_rows().unwrap_or(0);
+                    callback.record_local_read(
+                        &p.partition_id.job_id,
+                        p.partition_id.stage_id,
+                        p.partition_id.partition_id,
+                        &p.executor_meta.id,
+                        bytes,
+                        rows,
+                        duration_ms,
+                    );
+                }
+            }
+
             if let Err(e) = response_sender_c.send(r).await {
                 error!("Fail to send response event to the channel due to {e}");
             }
@@ -434,9 +460,11 @@ fn send_fetch_partitions(
         let semaphore = semaphore.clone();
         let response_sender = response_sender.clone();
         let customize_endpoint_c = customize_endpoint.clone();
+        let metrics_callback_c = metrics_callback.clone();
         spawned_tasks.push(SpawnedTask::spawn(async move {
             // Block if exceeds max request number.
             let permit = semaphore.acquire_owned().await.unwrap();
+            let start_time = std::time::Instant::now();
             let r = PartitionReaderEnum::FlightRemote
                 .fetch_partition(
                     &p,
@@ -446,6 +474,25 @@ fn send_fetch_partitions(
                     use_tls,
                 )
                 .await;
+
+            // Record remote read metrics if callback is set and read succeeded
+            if r.is_ok() {
+                if let Some(ref callback) = metrics_callback_c {
+                    let duration_ms = start_time.elapsed().as_millis() as u64;
+                    let bytes = p.partition_stats.num_bytes().unwrap_or(0);
+                    let rows = p.partition_stats.num_rows().unwrap_or(0);
+                    callback.record_remote_read(
+                        &p.partition_id.job_id,
+                        p.partition_id.stage_id,
+                        p.partition_id.partition_id,
+                        &p.executor_meta.id,
+                        bytes,
+                        rows,
+                        duration_ms,
+                    );
+                }
+            }
+
             // Block if the channel buffer is full.
             if let Err(e) = response_sender.send(r).await {
                 error!("Fail to send response event to the channel due to {e}");
@@ -992,6 +1039,7 @@ mod tests {
             true,
             None,
             false,
+            None, // No metrics callback in tests
         );
 
         let stream = RecordBatchStreamAdapter::new(

ballista/core/src/extension.rs

@@ -175,6 +175,20 @@ pub trait SessionConfigExt {
 
     /// Get whether to use TLS for executor connections
     fn ballista_use_tls(&self) -> bool;
+
+    /// Set a callback for recording shuffle read metrics (local vs remote).
+    ///
+    /// This callback will be invoked by the shuffle reader during execution
+    /// to record detailed metrics about local and remote shuffle reads.
+    fn with_ballista_shuffle_read_metrics_callback(
+        self,
+        callback: Arc<dyn ShuffleReadMetricsCallback>,
+    ) -> Self;
+
+    /// Get the shuffle read metrics callback if one has been set.
+    fn ballista_shuffle_read_metrics_callback(
+        &self,
+    ) -> Option<Arc<dyn ShuffleReadMetricsCallback>>;
 }
 
 /// [SessionConfigHelperExt] is set of [SessionConfig] extension methods
@@ -459,6 +473,21 @@ impl SessionConfigExt for SessionConfig {
             .map(|ext| ext.0)
             .unwrap_or(false)
     }
+
+    fn with_ballista_shuffle_read_metrics_callback(
+        self,
+        callback: Arc<dyn ShuffleReadMetricsCallback>,
+    ) -> Self {
+        let extension = ShuffleReadMetricsCallbackExtension::new(callback);
+        self.with_extension(Arc::new(extension))
+    }
+
+    fn ballista_shuffle_read_metrics_callback(
+        &self,
+    ) -> Option<Arc<dyn ShuffleReadMetricsCallback>> {
+        self.get_extension::<ShuffleReadMetricsCallbackExtension>()
+            .map(|ext| ext.callback())
+    }
 }
 
 impl SessionConfigHelperExt for SessionConfig {
@@ -658,6 +687,79 @@ impl BallistaConfigGrpcEndpoint {
 #[derive(Clone, Copy)]
 pub struct BallistaUseTls(pub bool);
 
+/// Callback trait for recording shuffle read metrics from the shuffle reader.
+///
+/// This trait is designed to be passed via session config extension to the
+/// shuffle reader, allowing external systems (like Spice) to capture detailed
+/// shuffle read locality metrics without creating circular dependencies.
+pub trait ShuffleReadMetricsCallback: Send + Sync {
+    /// Record a local shuffle read operation.
+    ///
+    /// Called when the shuffle reader successfully reads data from a local file
+    /// (i.e., the partition was produced by this executor).
+    ///
+    /// # Arguments
+    /// * `job_id` - The job identifier
+    /// * `stage_id` - The stage that is reading the shuffle data
+    /// * `partition` - The partition being read
+    /// * `source_executor_id` - The executor that produced the shuffle data (same as current executor for local reads)
+    /// * `bytes` - Number of bytes read
+    /// * `rows` - Number of rows read
+    /// * `duration_ms` - Time taken to read the partition
+    fn record_local_read(
+        &self,
+        job_id: &str,
+        stage_id: usize,
+        partition: usize,
+        source_executor_id: &str,
+        bytes: u64,
+        rows: u64,
+        duration_ms: u64,
+    );
+
+    /// Record a remote shuffle read operation.
+    ///
+    /// Called when the shuffle reader fetches data from a remote executor
+    /// via Arrow Flight.
+    ///
+    /// # Arguments
+    /// * `job_id` - The job identifier
+    /// * `stage_id` - The stage that is reading the shuffle data
+    /// * `partition` - The partition being read
+    /// * `source_executor_id` - The executor that produced the shuffle data
+    /// * `bytes` - Number of bytes read
+    /// * `rows` - Number of rows read
+    /// * `duration_ms` - Time taken to fetch the partition
+    fn record_remote_read(
+        &self,
+        job_id: &str,
+        stage_id: usize,
+        partition: usize,
+        source_executor_id: &str,
+        bytes: u64,
+        rows: u64,
+        duration_ms: u64,
+    );
+}
+
+/// Session config extension wrapper for the shuffle read metrics callback.
+#[derive(Clone)]
+pub struct ShuffleReadMetricsCallbackExtension {
+    callback: Arc<dyn ShuffleReadMetricsCallback>,
+}
+
+impl ShuffleReadMetricsCallbackExtension {
+    /// Create a new extension wrapping the provided callback.
+    pub fn new(callback: Arc<dyn ShuffleReadMetricsCallback>) -> Self {
+        Self { callback }
+    }
+
+    /// Get the callback.
+    pub fn callback(&self) -> Arc<dyn ShuffleReadMetricsCallback> {
+        Arc::clone(&self.callback)
+    }
+}
+
 #[cfg(test)]
 mod test {
     use datafusion::{

ballista/executor/src/metrics/mod.rs

@@ -86,6 +86,36 @@ pub trait ExecutorMetricsCollector: Send + Sync {
         duration_ms: u64,
     );
 
+    /// Record local shuffle read metrics (data read from local disk).
+    ///
+    /// Called when shuffle data is read from a local file. This means the partition
+    /// was written by this same executor in a previous stage, avoiding network transfer.
+    fn record_shuffle_read_local(
+        &self,
+        job_id: &str,
+        stage_id: usize,
+        partition: usize,
+        bytes: u64,
+        rows: u64,
+        duration_ms: u64,
+    );
+
+    /// Record remote shuffle read metrics (data fetched from another executor).
+    ///
+    /// Called when shuffle data must be fetched over the network from another
+    /// executor that produced the partition. The `source_executor_id` identifies
+    /// the executor that holds the shuffle data.
+    fn record_shuffle_read_remote(
+        &self,
+        job_id: &str,
+        stage_id: usize,
+        partition: usize,
+        source_executor_id: &str,
+        bytes: u64,
+        rows: u64,
+        duration_ms: u64,
+    );
+
     /// Record executor memory availability.
     ///
     /// Called periodically (e.g., during heartbeat) to report the executor's
@@ -154,6 +184,35 @@ impl ExecutorMetricsCollector for LoggingMetricsCollector {
         );
     }
 
+    fn record_shuffle_read_local(
+        &self,
+        job_id: &str,
+        stage_id: usize,
+        partition: usize,
+        bytes: u64,
+        rows: u64,
+        duration_ms: u64,
+    ) {
+        info!(
+            "=== [{job_id}/{stage_id}/{partition}] Local shuffle read: {bytes} bytes, {rows} rows in {duration_ms}ms ==="
+        );
+    }
+
+    fn record_shuffle_read_remote(
+        &self,
+        job_id: &str,
+        stage_id: usize,
+        partition: usize,
+        source_executor_id: &str,
+        bytes: u64,
+        rows: u64,
+        duration_ms: u64,
+    ) {
+        info!(
+            "=== [{job_id}/{stage_id}/{partition}] Remote shuffle read from {source_executor_id}: {bytes} bytes, {rows} rows in {duration_ms}ms ==="
+        );
+    }
+
     fn record_memory_available(&self, available_bytes: u64) {
         info!("=== Executor memory available: {available_bytes} bytes ===");
     }

ballista/scheduler/src/metrics/mod.rs

@@ -134,6 +134,32 @@ pub trait SchedulerMetricsCollector: Send + Sync {
     /// stage-level retries and tracks individual task retry attempts.
     fn record_task_retry(&self, job_id: &str, stage_id: usize);
 
+    /// Record a shuffle affinity hit - task was assigned to an executor that has
+    /// local shuffle data from a parent stage.
+    ///
+    /// Called when the scheduler assigns a task to an executor that already has
+    /// the required shuffle partitions from upstream stages stored locally.
+    /// This indicates the task can read shuffle data without network transfer.
+    fn record_task_shuffle_affinity_hit(
+        &self,
+        job_id: &str,
+        stage_id: usize,
+        executor_id: &str,
+    );
+
+    /// Record a shuffle affinity miss - task was assigned to an executor that does
+    /// NOT have local shuffle data from a parent stage.
+    ///
+    /// Called when the scheduler assigns a task to an executor that does not have
+    /// the required shuffle partitions locally. This indicates the task will need
+    /// to fetch shuffle data over the network from other executors.
+    fn record_task_shuffle_affinity_miss(
+        &self,
+        job_id: &str,
+        stage_id: usize,
+        executor_id: &str,
+    );
+
     // =========================================================================
     // Executor management events (new)
     // =========================================================================
@@ -207,6 +233,20 @@ impl SchedulerMetricsCollector for NoopMetricsCollector {
     ) {
     }
     fn record_task_retry(&self, _job_id: &str, _stage_id: usize) {}
+    fn record_task_shuffle_affinity_hit(
+        &self,
+        _job_id: &str,
+        _stage_id: usize,
+        _executor_id: &str,
+    ) {
+    }
+    fn record_task_shuffle_affinity_miss(
+        &self,
+        _job_id: &str,
+        _stage_id: usize,
+        _executor_id: &str,
+    ) {
+    }
 
     // Executor management
     fn set_active_executor_count(&self, _count: usize) {}