execution_plans: add metrics collector and re-writer

This change introduces new structs / concepts:

1. MetricsCollector

Collects metrics from an instance of `StageExec` (which is a task) using a
pre-order traversal. It stops at any `ArrowFlightReadExec` nodes and collects
child task metrics from them, if there are any.

2. MetricsWrapperExec

A new `ExecutionPlan` node which cannot be executed. It wraps a "real" `ExecutionPlan`
node and stores metrics. This let's you "override" the `metrics()` method on
`ExecutionPlan` nodes. This override applies when displaying the plan as well.

This struct is private.

3. Task MetricsRewriter

Rewrites a task by wrapping each node in a `MetricsWrapperExec`.

Informs #123.

Author: jayshrivastava

src/execution_plans/arrow_flight_read.rs

@@ -3,13 +3,17 @@ use crate::config_extension_ext::ContextGrpcMetadata;
 use crate::errors::{map_flight_to_datafusion_error, map_status_to_datafusion_error};
 use crate::execution_plans::StageExec;
 use crate::flight_service::DoGet;
+use crate::metrics::proto::MetricsSetProto;
 use crate::protobuf::{proto_from_stage, DistributedCodec, StageKey};
 use crate::ChannelResolver;
 use arrow_flight::decode::FlightRecordBatchStream;
 use arrow_flight::error::FlightError;
 use arrow_flight::flight_service_client::FlightServiceClient;
 use arrow_flight::Ticket;
+use dashmap::DashMap;
 use datafusion::arrow::datatypes::SchemaRef;
+#[cfg(test)]
+use datafusion::arrow::record_batch::RecordBatch;
 use datafusion::common::{exec_err, internal_datafusion_err, internal_err, plan_err};
 use datafusion::error::DataFusionError;
 use datafusion::execution::{SendableRecordBatchStream, TaskContext};
@@ -56,6 +60,16 @@ pub struct ArrowFlightReadReadyExec {
     /// the properties we advertise for this execution plan
     properties: PlanProperties,
     pub(crate) stage_num: usize,
+    /// metrics_collection is used to collect metrics from child tasks. It is empty when an ArrowFlightReadReadyExec is instansiated
+    /// (deserialized, created via [ArrowFlightReadExec::new_ready] etc). Metrics are populated in this map via [ArrowFlightReadExec::execute].
+    ///
+    /// An instance may recieve metrics for 0 to N child tasks, where N is the number of tasks in the stage it is reading from.
+    /// This is because, by convention, the ArrowFlightEndpoint sends metrics for a task to the last ArrowFlightReadExec to read from it, which
+    /// may or may not be this instance.
+    pub(super) metrics_collection: Arc<DashMap<StageKey, Vec<MetricsSetProto>>>,
+    /// Test-only field to provide record batches directly without network calls
+    #[cfg(test)]
+    pub(crate) test_data: Option<Vec<RecordBatch>>,
 }
 
 impl ArrowFlightReadExec {
@@ -85,6 +99,9 @@ impl ArrowFlightReadExec {
         Self::Ready(ArrowFlightReadReadyExec {
             properties,
             stage_num,
+            metrics_collection: Arc::new(DashMap::new()),
+            #[cfg(test)]
+            test_data: None,
         })
     }
 
@@ -98,6 +115,20 @@ impl ArrowFlightReadExec {
             _ => internal_err!("ArrowFlightReadExec is already distributed"),
         }
     }
+
+    /// Test-only method to set record batches for direct execution
+    #[cfg(test)]
+    pub fn with_test_data(mut self, batches: Vec<RecordBatch>) -> Self {
+        match &mut self {
+            ArrowFlightReadExec::Ready(exec) => {
+                exec.test_data = Some(batches);
+            }
+            ArrowFlightReadExec::Pending(_) => {
+                panic!("Cannot set test data on pending ArrowFlightReadExec");
+            }
+        }
+        self
+    }
 }
 
 impl DisplayAs for ArrowFlightReadExec {
@@ -151,6 +182,16 @@ impl ExecutionPlan for ArrowFlightReadExec {
             return exec_err!("ArrowFlightReadExec is not ready, was the distributed optimization step performed?");
         };
 
+        // Check for test data first
+        #[cfg(test)]
+        if let Some(test_batches) = &self_ready.test_data {
+            use futures::stream;
+            return Ok(Box::pin(RecordBatchStreamAdapter::new(
+                self.schema(),
+                stream::iter(test_batches.clone().into_iter().map(Ok)),
+            )));
+        }
+
         // get the channel manager and current stage from our context
         let channel_resolver = get_distributed_channel_resolver(context.session_config())?;