refactor modules and files

Author: jayshrivastava

src/execution_plans/metrics.rs

@@ -1,11 +1,10 @@
 mod metrics;
-mod metrics_collecting_stream;
 mod network_coalesce;
 mod network_shuffle;
 mod partition_isolator;
 mod stage;
 
-pub use metrics::collect_and_create_metrics_flight_data;
+pub use metrics::MetricsWrapperExec;
 pub use network_coalesce::{NetworkCoalesceExec, NetworkCoalesceReady};
 pub use network_shuffle::{NetworkShuffleExec, NetworkShuffleReadyExec};
 pub use partition_isolator::PartitionIsolatorExec;

src/execution_plans/mod.rs

@@ -3,9 +3,9 @@ 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::MetricsCollectingStream;
 use crate::metrics::proto::MetricsSetProto;
 use crate::protobuf::{DistributedCodec, StageKey, proto_from_input_stage};
 use crate::protobuf::{map_flight_to_datafusion_error, map_status_to_datafusion_error};

src/execution_plans/network_coalesce.rs

@@ -3,9 +3,9 @@ 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::MetricsCollectingStream;
 use crate::metrics::proto::MetricsSetProto;
 use crate::protobuf::{DistributedCodec, StageKey, proto_from_input_stage};
 use crate::protobuf::{map_flight_to_datafusion_error, map_status_to_datafusion_error};

src/execution_plans/network_shuffle.rs

@@ -1,14 +1,18 @@
 use crate::common::with_callback;
 use crate::config_extension_ext::ContextGrpcMetadata;
-use crate::execution_plans::{
-    DistributedTaskContext, StageExec, collect_and_create_metrics_flight_data,
-};
+use crate::execution_plans::{DistributedTaskContext, StageExec};
 use crate::flight_service::service::ArrowFlightEndpoint;
 use crate::flight_service::session_builder::DistributedSessionBuilderContext;
 use crate::flight_service::trailing_flight_data_stream::TrailingFlightDataStream;
+use crate::metrics::TaskMetricsCollector;
+use crate::metrics::proto::df_metrics_set_to_proto;
 use crate::protobuf::{
-    DistributedCodec, StageKey, datafusion_error_to_tonic_status, stage_from_proto,
+    datafusion_error_to_tonic_status, stage_from_proto, AppMetadata, DistributedCodec, FlightAppMetadata, MetricsCollection, StageKey, TaskMetrics
 };
+use arrow::array::RecordBatch;
+use arrow::datatypes::SchemaRef;
+use arrow::ipc::writer::{DictionaryTracker, IpcDataGenerator, IpcWriteOptions};
+use arrow_flight::FlightData;
 use arrow_flight::Ticket;
 use arrow_flight::encode::FlightDataEncoderBuilder;
 use arrow_flight::error::FlightError;
@@ -18,6 +22,7 @@ use datafusion::common::exec_datafusion_err;
 use datafusion::execution::SendableRecordBatchStream;
 use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
 use futures::TryStreamExt;
+use futures::{Stream, stream};
 use prost::Message;
 use std::sync::Arc;
 use std::sync::atomic::{AtomicUsize, Ordering};
@@ -204,6 +209,77 @@ fn flight_stream_from_record_batch_stream(
     })))
 }
 
+// Collects metrics from the provided stage and encodes it into a stream of flight data using
+// the schema of the stage.
+pub fn collect_and_create_metrics_flight_data(
+    stage_key: StageKey,
+    stage: Arc<StageExec>,
+) -> Result<impl Stream<Item = Result<FlightData, FlightError>> + Send + 'static, FlightError> {
+    // Get the metrics for the task executed on this worker. Separately, collect metrics for child tasks.
+    let mut result = TaskMetricsCollector::new()
+        .collect(stage.plan.clone())
+        .map_err(|err| FlightError::ProtocolError(err.to_string()))?;
+
+    // Add the metrics for this task into the collection of task metrics.
+    // Skip any metrics that can't be converted to proto (unsupported types)
+    let proto_task_metrics = result
+        .task_metrics
+        .iter()
+        .map(|metrics| {
+            df_metrics_set_to_proto(metrics)
+                .map_err(|err| FlightError::ProtocolError(err.to_string()))
+        })
+        .collect::<Result<Vec<_>, _>>()?;
+    result
+        .child_task_metrics
+        .insert(stage_key, proto_task_metrics);
+
+    // Serialize the metrics for all tasks.
+    let mut task_metrics_set = vec![];
+    for (stage_key, metrics) in result.child_task_metrics.into_iter() {
+        task_metrics_set.push(TaskMetrics {
+            stage_key: Some(stage_key),
+            metrics,
+        });
+    }
+
+    let flight_app_metadata = FlightAppMetadata {
+        content: Some(AppMetadata::MetricsCollection(MetricsCollection {
+            tasks: task_metrics_set,
+        })),
+    };
+
+    let metrics_flight_data =
+        empty_flight_data_with_app_metadata(flight_app_metadata, stage.plan.schema())?;
+    Ok(Box::pin(stream::once(
+        async move { Ok(metrics_flight_data) },
+    )))
+}
+
+/// Creates a FlightData with the given app_metadata and empty RecordBatch using the provided schema.
+/// We don't use [arrow_flight::encode::FlightDataEncoder] (and by extension, the [arrow_flight::encode::FlightDataEncoderBuilder])
+/// since they skip messages with empty RecordBatch data.
+pub fn empty_flight_data_with_app_metadata(
+    metadata: FlightAppMetadata,
+    schema: SchemaRef,
+) -> Result<FlightData, FlightError> {
+    let mut buf = vec![];
+    metadata
+        .encode(&mut buf)
+        .map_err(|err| FlightError::ProtocolError(err.to_string()))?;
+
+    let empty_batch = RecordBatch::new_empty(schema);
+    let options = IpcWriteOptions::default();
+    let data_gen = IpcDataGenerator::default();
+    let mut dictionary_tracker = DictionaryTracker::new(true);
+    let (_, encoded_data) = data_gen
+        .encoded_batch(&empty_batch, &mut dictionary_tracker, &options)
+        .map_err(|e| {
+            FlightError::ProtocolError(format!("Failed to create empty batch FlightData: {e}"))
+        })?;
+    Ok(FlightData::from(encoded_data).with_app_metadata(buf))
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

src/flight_service/do_get.rs

@@ -1 +1,6 @@
+mod metrics_collecting_stream;
 pub(crate) mod proto;
+mod task_metrics_collector;
+mod task_metrics_rewriter;
+pub(crate) use metrics_collecting_stream::MetricsCollectingStream;
+pub(crate) use task_metrics_collector::TaskMetricsCollector;