address pr comments

Author: jayshrivastava

src/execution_plans/metrics.rs

@@ -38,7 +38,6 @@ pub struct TaskMetricsCollector {
 }
 
 /// MetricsCollectorResult is the result of collecting metrics from a task.
-#[allow(dead_code)]
 pub struct MetricsCollectorResult {
     // metrics is a collection of metrics for a task ordered using a pre-order traversal of the task's plan.
     pub(super) task_metrics: Vec<MetricsSet>,
@@ -108,7 +107,6 @@ impl TreeNodeRewriter for TaskMetricsCollector {
 }
 
 impl TaskMetricsCollector {
-    #[allow(dead_code)]
     pub fn new() -> Self {
         Self {
             task_metrics: Vec::new(),
@@ -120,7 +118,6 @@ impl TaskMetricsCollector {
     /// Returns
     /// - a vec representing the metrics for the current task (ordered using a pre-order traversal)
     /// - a map representing the metrics for some subset of child tasks collected from NetworkShuffleExec leaves
-    #[allow(dead_code)]
     pub fn collect(
         mut self,
         plan: Arc<dyn ExecutionPlan>,
@@ -305,10 +302,10 @@ pub fn collect_and_create_metrics_flight_data(
             df_metrics_set_to_proto(metrics)
                 .map_err(|err| FlightError::ProtocolError(err.to_string()))
         })
-        .collect::<Result<Vec<MetricsSetProto>, FlightError>>()?;
+        .collect::<Result<Vec<_>, _>>()?;
     result
         .child_task_metrics
-        .insert(stage_key.clone(), proto_task_metrics.clone());
+        .insert(stage_key, proto_task_metrics);
 
     // Serialize the metrics for all tasks.
     let mut task_metrics_set = vec![];

src/execution_plans/metrics_collecting_stream.rs

@@ -27,7 +27,6 @@ impl<S> MetricsCollectingStream<S>
 where
     S: Stream<Item = Result<FlightData, FlightError>> + Send,
 {
-    #[allow(dead_code)]
     pub fn new(
         stream: S,
         metrics_collection: Arc<DashMap<StageKey, Vec<MetricsSetProto>>>,

src/flight_service/do_get.rs

@@ -144,7 +144,6 @@ impl ArrowFlightEndpoint {
         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());
@@ -158,11 +157,10 @@ fn missing(field: &'static str) -> impl FnOnce() -> Status {
     move || Status::invalid_argument(format!("Missing field '{field}'"))
 }
 
-// Creates a tonic response from a stream of record batches. Handles
-// - RecordBatch to flight conversion partition tracking, stage eviction, and trailing metrics.
+/// 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,
@@ -183,12 +181,14 @@ fn flight_stream_from_record_batch_stream(
             {
                 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}"
-                        ))
-                    })?;
+                let metrics_stream =
+                    collect_and_create_metrics_flight_data(stage_key, stage_data.stage).map_err(
+                        |err| {
+                            Status::internal(format!(
+                                "error collecting metrics in arrow flight endpoint: {err}"
+                            ))
+                        },
+                    )?;
 
                 return Ok(Some(metrics_stream));
             }

src/flight_service/trailing_flight_data_stream.rs

@@ -27,8 +27,6 @@ where
     T: Stream<Item = Result<FlightData, FlightError>> + Send,
     F: FnOnce() -> Result<Option<T>, FlightError>,
 {
-    // TODO: remove
-    #[allow(dead_code)]
     pub fn new(on_complete: F, inner: S) -> Self {
         Self {
             inner,

src/metrics/proto.rs

@@ -137,7 +137,6 @@ pub struct ProtoLabel {
 /// df_metrics_set_to_proto converts a [datafusion::physical_plan::metrics::MetricsSet] to a [MetricsSetProto].
 /// Custom metrics are filtered out, but any other errors are returned.
 /// TODO(#140): Support custom metrics.
-#[allow(dead_code)]
 pub fn df_metrics_set_to_proto(
     metrics_set: &MetricsSet,
 ) -> Result<MetricsSetProto, DataFusionError> {
@@ -164,7 +163,6 @@ pub fn df_metrics_set_to_proto(
 }
 
 /// metrics_set_proto_to_df converts a [MetricsSetProto] to a [datafusion::physical_plan::metrics::MetricsSet].
-#[allow(dead_code)]
 pub fn metrics_set_proto_to_df(
     metrics_set_proto: &MetricsSetProto,
 ) -> Result<MetricsSet, DataFusionError> {
@@ -178,12 +176,10 @@ pub fn metrics_set_proto_to_df(
 }
 
 /// Custom metrics are not supported in proto conversion.
-#[allow(dead_code)]
 const CUSTOM_METRICS_NOT_SUPPORTED: &str =
     "custom metrics are not supported in metrics proto conversion";
 
 /// df_metric_to_proto converts a `datafusion::physical_plan::metrics::Metric` to a `MetricProto`. It does not consume the Arc<Metric>.
-#[allow(dead_code)]
 pub fn df_metric_to_proto(metric: Arc<Metric>) -> Result<MetricProto, DataFusionError> {
     let partition = metric.partition().map(|p| p as u64);
     let labels = metric