flight_service: emit metrics from ArrowFlightEndpoint

This change updates the ArrowFlightEndpoint to collect metrics and emit them. When the last partition
in a task is finished, the ArrowFlightEndpoint collects metrics and emits them via the
TrailingFlightDataStream.

Previously, we would determine if a partition is finished when the request first hit the endpoint. Now,
we do it on stream completition. This is crutial for metrics collection because we need to know that
the stream is exhausted, meaning that there's no data flowing in the plan and metrics are not actively
being updated.

Since the ArrowFlightEndpoint now emits metrics and NetworkBoundary plan nodes collect metrics,
all coordinating StageExecs will now have the full collection of metrics for all tasks. This commit
adds integration style tests that assert that the coordinator is recieving the full set of metrics.

Follow up work
- Only collect metrics if a configuration is set in the SessionContext, removing extra overhead
- Display metrics in the plan using EXPLAIN (ANALYZE) - consider using sqllogictest or similar
  to test the output

Author: jayshrivastava

src/execution_plans/metrics.rs

@@ -67,8 +67,8 @@ where
             };
             metrics_collection.insert(stage_key, task_metrics.metrics);
         }
-
         flight_data.app_metadata.clear();
+
         Ok(())
     }
 }
@@ -256,7 +256,8 @@ mod tests {
         // Create a stream that emits an error - should be propagated through
         let stream_error = FlightError::ProtocolError("stream error from inner stream".to_string());
         let error_stream = stream::iter(vec![Err(stream_error)]);
-        let mut collecting_stream = MetricsCollectingStream::new(error_stream, metrics_collection);
+        let mut collecting_stream =
+            MetricsCollectingStream::new(error_stream, metrics_collection.clone());
 
         let result = collecting_stream.next().await.unwrap();
         assert_protocol_error(result, "stream error from inner stream");

src/execution_plans/metrics_collecting_stream.rs

@@ -5,6 +5,7 @@ mod network_shuffle;
 mod partition_isolator;
 mod stage;
 
+pub use metrics::collect_and_create_metrics_flight_data;
 pub use network_coalesce::{NetworkCoalesceExec, NetworkCoalesceReady};
 pub use network_shuffle::{NetworkShuffleExec, NetworkShuffleReadyExec};
 pub use partition_isolator::PartitionIsolatorExec;

src/execution_plans/mod.rs

@@ -3,6 +3,7 @@ use crate::channel_resolver_ext::get_distributed_channel_resolver;
 use crate::common::scale_partitioning_props;
 use crate::config_extension_ext::ContextGrpcMetadata;
 use crate::distributed_physical_optimizer_rule::{NetworkBoundary, limit_tasks_err};
+use crate::execution_plans::metrics_collecting_stream::MetricsCollectingStream;
 use crate::execution_plans::{DistributedTaskContext, StageExec};
 use crate::flight_service::DoGet;
 use crate::metrics::proto::MetricsSetProto;
@@ -297,6 +298,7 @@ impl ExecutionPlan for NetworkCoalesceExec {
             return internal_err!("NetworkCoalesceExec: task is unassigned, cannot proceed");
         };
 
+        let metrics_collection_capture = self_ready.metrics_collection.clone();
         let stream = async move {
             let channel = channel_resolver.get_channel_for_url(&url).await?;
             let stream = FlightServiceClient::new(channel)
@@ -306,8 +308,13 @@ impl ExecutionPlan for NetworkCoalesceExec {
                 .into_inner()
                 .map_err(|err| FlightError::Tonic(Box::new(err)));
 
-            Ok(FlightRecordBatchStream::new_from_flight_data(stream)
-                .map_err(map_flight_to_datafusion_error))
+            let metrics_collecting_stream =
+                MetricsCollectingStream::new(stream, metrics_collection_capture);
+
+            Ok(
+                FlightRecordBatchStream::new_from_flight_data(metrics_collecting_stream)
+                    .map_err(map_flight_to_datafusion_error),
+            )
         }
         .try_flatten_stream();
 

src/execution_plans/network_coalesce.rs

@@ -3,6 +3,7 @@ use crate::channel_resolver_ext::get_distributed_channel_resolver;
 use crate::common::scale_partitioning;
 use crate::config_extension_ext::ContextGrpcMetadata;
 use crate::distributed_physical_optimizer_rule::NetworkBoundary;
+use crate::execution_plans::metrics_collecting_stream::MetricsCollectingStream;
 use crate::execution_plans::{DistributedTaskContext, StageExec};
 use crate::flight_service::DoGet;
 use crate::metrics::proto::MetricsSetProto;
@@ -330,6 +331,7 @@ impl ExecutionPlan for NetworkShuffleExec {
                 },
             );
 
+            let metrics_collection_capture = self_ready.metrics_collection.clone();
             async move {
                 let url = task.url.ok_or(internal_datafusion_err!(
                     "NetworkShuffleExec: task is unassigned, cannot proceed"
@@ -343,8 +345,13 @@ impl ExecutionPlan for NetworkShuffleExec {
                     .into_inner()
                     .map_err(|err| FlightError::Tonic(Box::new(err)));
 
-                Ok(FlightRecordBatchStream::new_from_flight_data(stream)
-                    .map_err(map_flight_to_datafusion_error))
+                let metrics_collecting_stream =
+                    MetricsCollectingStream::new(stream, metrics_collection_capture);
+
+                Ok(
+                    FlightRecordBatchStream::new_from_flight_data(metrics_collecting_stream)
+                        .map_err(map_flight_to_datafusion_error),
+                )
             }
             .try_flatten_stream()
             .boxed()

src/execution_plans/network_shuffle.rs

@@ -1,7 +1,10 @@
 use crate::config_extension_ext::ContextGrpcMetadata;
-use crate::execution_plans::{DistributedTaskContext, StageExec};
+use crate::execution_plans::{
+    DistributedTaskContext, StageExec, collect_and_create_metrics_flight_data,
+};
 use crate::flight_service::service::ArrowFlightEndpoint;
 use crate::flight_service::session_builder::DistributedSessionBuilderContext;
+use crate::flight_service::trailing_flight_data_stream::TrailingFlightDataStream;
 use crate::protobuf::{
     DistributedCodec, StageExecProto, StageKey, datafusion_error_to_tonic_status, stage_from_proto,
 };
@@ -94,12 +97,6 @@ impl ArrowFlightEndpoint {
             })
             .await?;
         let stage = Arc::clone(&stage_data.stage);
-        let num_partitions_remaining = Arc::clone(&stage_data.num_partitions_remaining);
-
-        // If all the partitions are done, remove the stage from the cache.
-        if num_partitions_remaining.fetch_sub(1, Ordering::SeqCst) <= 1 {
-            self.task_data_entries.remove(key);
-        }
 
         // Find out which partition group we are executing
         let cfg = session_state.config_mut();
@@ -126,15 +123,30 @@ impl ArrowFlightEndpoint {
             .execute(doget.target_partition as usize, session_state.task_ctx())
             .map_err(|err| Status::internal(format!("Error executing stage plan: {err:#?}")))?;
 
-        Ok(record_batch_stream_to_response(stream))
+        let task_data_capture = self.task_data_entries.clone();
+        Ok(flight_stream_from_record_batch_stream(
+            key.clone(),
+            stage,
+            stage_data.clone(),
+            move || {
+                task_data_capture.remove(key.clone());
+            },
+            stream,
+        ))
     }
 }
 
 fn missing(field: &'static str) -> impl FnOnce() -> Status {
     move || Status::invalid_argument(format!("Missing field '{field}'"))
 }
 
-fn record_batch_stream_to_response(
+// Creates a tonic response from a stream of record batches. Handles
+// - RecordBatch to flight conversion partition tracking, stage eviction, and trailing metrics.
+fn flight_stream_from_record_batch_stream(
+    stage_key: StageKey,
+    stage: Arc<StageExec>,
+    stage_data: TaskData,
+    evict_stage: impl FnOnce() + Send + 'static,
     stream: SendableRecordBatchStream,
 ) -> Response<<ArrowFlightEndpoint as FlightService>::DoGetStream> {
     let flight_data_stream =
@@ -144,7 +156,31 @@ fn record_batch_stream_to_response(
                 FlightError::Tonic(Box::new(datafusion_error_to_tonic_status(&err)))
             }));
 
-    Response::new(Box::pin(flight_data_stream.map_err(|err| match err {
+    let trailing_metrics_stream = TrailingFlightDataStream::new(
+        move || {
+            if stage_data
+                .num_partitions_remaining
+                .fetch_sub(1, Ordering::SeqCst)
+                == 1
+            {
+                evict_stage();
+
+                let metrics_stream = collect_and_create_metrics_flight_data(stage_key, stage)
+                    .map_err(|err| {
+                        Status::internal(format!(
+                            "error collecting metrics in arrow flight endpoint: {err}"
+                        ))
+                    })?;
+
+                return Ok(Some(metrics_stream));
+            }
+
+            Ok(None)
+        },
+        flight_data_stream,
+    );
+
+    Response::new(Box::pin(trailing_metrics_stream.map_err(|err| match err {
         FlightError::Tonic(status) => *status,
         _ => Status::internal(format!("Error during flight stream: {err}")),
     })))
@@ -215,24 +251,27 @@ mod tests {
         let stage_proto = proto_from_stage(&stage, &DefaultPhysicalExtensionCodec {}).unwrap();
         let stage_proto_for_closure = stage_proto.clone();
         let endpoint_ref = &endpoint;
+
         let do_get = async move |partition: u64, task_number: u64, stage_key: StageKey| {
             let stage_proto = stage_proto_for_closure.clone();
-            // Create DoGet message
             let doget = DoGet {
                 stage_proto: Some(stage_proto),
                 target_task_index: task_number,
                 target_partition: partition,
                 stage_key: Some(stage_key),
             };
 
-            // Create Flight ticket
             let ticket = Ticket {
                 ticket: Bytes::from(doget.encode_to_vec()),
             };
 
-            // Call the actual get() method
             let request = Request::new(ticket);
-            endpoint_ref.get(request).await
+            let response = endpoint_ref.get(request).await?;
+            let mut stream = response.into_inner();
+
+            // Consume the stream.
+            while let Some(_flight_data) = stream.try_next().await? {}
+            Ok::<(), Status>(())
         };
 
         // For each task, call do_get() for each partition except the last.
@@ -248,22 +287,22 @@ mod tests {
 
         // Run the last partition of task 0. Any partition number works. Verify that the task state
         // is evicted because all partitions have been processed.
-        let result = do_get(1, 0, task_keys[0].clone()).await;
+        let result = do_get(2, 0, task_keys[0].clone()).await;
         assert!(result.is_ok());
         let stored_stage_keys = endpoint.task_data_entries.keys().collect::<Vec<StageKey>>();
         assert_eq!(stored_stage_keys.len(), 2);
         assert!(stored_stage_keys.contains(&task_keys[1]));
         assert!(stored_stage_keys.contains(&task_keys[2]));
 
         // Run the last partition of task 1.
-        let result = do_get(1, 1, task_keys[1].clone()).await;
+        let result = do_get(2, 1, task_keys[1].clone()).await;
         assert!(result.is_ok());
         let stored_stage_keys = endpoint.task_data_entries.keys().collect::<Vec<StageKey>>();
         assert_eq!(stored_stage_keys.len(), 1);
         assert!(stored_stage_keys.contains(&task_keys[2]));
 
         // Run the last partition of the last task.
-        let result = do_get(1, 2, task_keys[2].clone()).await;
+        let result = do_get(2, 2, task_keys[2].clone()).await;
         assert!(result.is_ok());
         let stored_stage_keys = endpoint.task_data_entries.keys().collect::<Vec<StageKey>>();
         assert_eq!(stored_stage_keys.len(), 0);

src/flight_service/do_get.rs

@@ -1,7 +1,7 @@
 mod do_get;
 mod service;
 mod session_builder;
-mod trailing_flight_data_stream;
+pub(super) mod trailing_flight_data_stream;
 pub(crate) use do_get::DoGet;
 
 pub use service::ArrowFlightEndpoint;

src/flight_service/mod.rs

@@ -17,7 +17,7 @@ use tonic::{Request, Response, Status, Streaming};
 
 pub struct ArrowFlightEndpoint {
     pub(super) runtime: Arc<RuntimeEnv>,
-    pub(super) task_data_entries: TTLMap<StageKey, Arc<OnceCell<TaskData>>>,
+    pub(super) task_data_entries: Arc<TTLMap<StageKey, Arc<OnceCell<TaskData>>>>,
     pub(super) session_builder: Arc<dyn DistributedSessionBuilder + Send + Sync>,
 }
 
@@ -28,7 +28,7 @@ impl ArrowFlightEndpoint {
         let ttl_map = TTLMap::try_new(TTLMapConfig::default())?;
         Ok(Self {
             runtime: Arc::new(RuntimeEnv::default()),
-            task_data_entries: ttl_map,
+            task_data_entries: Arc::new(ttl_map),
             session_builder: Arc::new(session_builder),
         })
     }

src/flight_service/service.rs

@@ -8,22 +8,24 @@ use tokio::pin;
 /// TrailingFlightDataStream - wraps a FlightData stream. It calls the `on_complete` closure when the stream is finished.
 /// If the closure returns a new stream, it will be appended to the original stream and consumed.
 #[pin_project]
-pub struct TrailingFlightDataStream<S, F>
+pub struct TrailingFlightDataStream<S, T, F>
 where
     S: Stream<Item = Result<FlightData, FlightError>> + Send,
-    F: FnOnce() -> Result<Option<S>, FlightError>,
+    T: Stream<Item = Result<FlightData, FlightError>> + Send,
+    F: FnOnce() -> Result<Option<T>, FlightError>,
 {
     #[pin]
     inner: S,
     on_complete: Option<F>,
     #[pin]
-    trailing_stream: Option<S>,
+    trailing_stream: Option<T>,
 }
 
-impl<S, F> TrailingFlightDataStream<S, F>
+impl<S, T, F> TrailingFlightDataStream<S, T, F>
 where
     S: Stream<Item = Result<FlightData, FlightError>> + Send,
-    F: FnOnce() -> Result<Option<S>, FlightError>,
+    T: Stream<Item = Result<FlightData, FlightError>> + Send,
+    F: FnOnce() -> Result<Option<T>, FlightError>,
 {
     // TODO: remove
     #[allow(dead_code)]
@@ -36,10 +38,11 @@ where
     }
 }
 
-impl<S, F> Stream for TrailingFlightDataStream<S, F>
+impl<S, T, F> Stream for TrailingFlightDataStream<S, T, F>
 where
     S: Stream<Item = Result<FlightData, FlightError>> + Send,
-    F: FnOnce() -> Result<Option<S>, FlightError>,
+    T: Stream<Item = Result<FlightData, FlightError>> + Send,
+    F: FnOnce() -> Result<Option<T>, FlightError>,
 {
     type Item = Result<FlightData, FlightError>;
 
@@ -74,7 +77,7 @@ mod tests {
     use arrow::record_batch::RecordBatch;
     use arrow_flight::FlightData;
     use arrow_flight::decode::FlightRecordBatchStream;
-    use arrow_flight::encode::FlightDataEncoderBuilder;
+    use arrow_flight::encode::{FlightDataEncoder, FlightDataEncoderBuilder};
     use futures::stream::{self, StreamExt};
     use std::sync::Arc;
 
@@ -186,7 +189,7 @@ mod tests {
             )))),
         ];
         let inner_stream = stream::iter(data);
-        let on_complete = || Ok(None);
+        let on_complete = || Ok(None::<FlightDataEncoder>);
         let trailing_stream = TrailingFlightDataStream::new(on_complete, inner_stream);
         let record_batches = FlightRecordBatchStream::new_from_flight_data(trailing_stream)
             .collect::<Vec<Result<RecordBatch, FlightError>>>()
@@ -202,8 +205,7 @@ mod tests {
         let name_array = StringArray::from(vec!["item1"]);
         let value_array = Int32Array::from(vec![1]);
         let inner_stream = create_flight_data_stream(name_array, value_array);
-
-        let on_complete = || -> Result<Option<_>, FlightError> {
+        let on_complete = || -> Result<Option<FlightDataEncoder>, FlightError> {
             Err(FlightError::ExternalError(Box::new(std::io::Error::new(
                 std::io::ErrorKind::Other,
                 "callback error",
@@ -225,7 +227,7 @@ mod tests {
             StringArray::from(vec!["item1"] as Vec<&str>),
             Int32Array::from(vec![1] as Vec<i32>),
         );
-        let on_complete = || Ok(None);
+        let on_complete = || Ok(None::<FlightDataEncoder>);
         let trailing_stream = TrailingFlightDataStream::new(on_complete, inner_stream);
         let record_batches = FlightRecordBatchStream::new_from_flight_data(trailing_stream)
             .collect::<Vec<Result<RecordBatch, FlightError>>>()