protobuf: move StageKey to protobuf module

The StageKey will be used by various modules due to metrics collection for EXPLAIN ANALYZE. It will be more organized to import from the protobuf module than flight_service, which is its old module.

Author: jayshrivastava

src/execution_plans/arrow_flight_read.rs

@@ -2,8 +2,8 @@ use crate::channel_resolver_ext::get_distributed_channel_resolver;
 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, StageKey};
-use crate::protobuf::{proto_from_stage, DistributedCodec};
+use crate::flight_service::DoGet;
+use crate::protobuf::{proto_from_stage, DistributedCodec, StageKey};
 use crate::ChannelResolver;
 use arrow_flight::decode::FlightRecordBatchStream;
 use arrow_flight::error::FlightError;

src/flight_service/do_get.rs

@@ -1,10 +1,9 @@
-use super::service::StageKey;
 use crate::config_extension_ext::ContextGrpcMetadata;
 use crate::errors::datafusion_error_to_tonic_status;
 use crate::execution_plans::{PartitionGroup, StageExec};
 use crate::flight_service::service::ArrowFlightEndpoint;
 use crate::flight_service::session_builder::DistributedSessionBuilderContext;
-use crate::protobuf::{stage_from_proto, DistributedCodec, StageExecProto};
+use crate::protobuf::{stage_from_proto, DistributedCodec, StageExecProto, StageKey};
 use arrow_flight::encode::FlightDataEncoderBuilder;
 use arrow_flight::error::FlightError;
 use arrow_flight::flight_service_server::FlightService;

src/flight_service/mod.rs

@@ -4,7 +4,7 @@ mod session_builder;
 
 pub(crate) use do_get::DoGet;
 
-pub use service::{ArrowFlightEndpoint, StageKey};
+pub use service::ArrowFlightEndpoint;
 pub use session_builder::{
     DefaultSessionBuilder, DistributedSessionBuilder, DistributedSessionBuilderContext,
     MappedDistributedSessionBuilder, MappedDistributedSessionBuilderExt,

src/flight_service/service.rs

@@ -1,6 +1,7 @@
 use crate::common::ttl_map::{TTLMap, TTLMapConfig};
 use crate::flight_service::do_get::TaskData;
 use crate::flight_service::DistributedSessionBuilder;
+use crate::protobuf::StageKey;
 use arrow_flight::flight_service_server::FlightService;
 use arrow_flight::{
     Action, ActionType, Criteria, Empty, FlightData, FlightDescriptor, FlightInfo,
@@ -14,20 +15,6 @@ use std::sync::Arc;
 use tokio::sync::OnceCell;
 use tonic::{Request, Response, Status, Streaming};
 
-/// A key that uniquely identifies a stage in a query
-#[derive(Clone, Hash, Eq, PartialEq, ::prost::Message)]
-pub struct StageKey {
-    /// Our query id
-    #[prost(string, tag = "1")]
-    pub query_id: String,
-    /// Our stage id
-    #[prost(uint64, tag = "2")]
-    pub stage_id: u64,
-    /// The task number within the stage
-    #[prost(uint64, tag = "3")]
-    pub task_number: u64,
-}
-
 pub struct ArrowFlightEndpoint {
     pub(super) runtime: Arc<RuntimeEnv>,
     pub(super) task_data_entries: TTLMap<StageKey, Arc<OnceCell<TaskData>>>,

src/protobuf/mod.rs

@@ -3,5 +3,5 @@ mod stage_proto;
 mod user_codec;
 
 pub(crate) use distributed_codec::DistributedCodec;
-pub(crate) use stage_proto::{proto_from_stage, stage_from_proto, StageExecProto};
+pub(crate) use stage_proto::{proto_from_stage, stage_from_proto, StageExecProto, StageKey};
 pub(crate) use user_codec::{get_distributed_user_codec, set_distributed_user_codec};

src/protobuf/stage_proto.rs

@@ -9,9 +9,34 @@ use datafusion_proto::{
     physical_plan::{AsExecutionPlan, PhysicalExtensionCodec},
     protobuf::PhysicalPlanNode,
 };
+use std::fmt::Display;
 use std::sync::Arc;
 use url::Url;
 
+/// A key that uniquely identifies a stage in a query
+#[derive(Clone, Hash, Eq, PartialEq, ::prost::Message)]
+pub struct StageKey {
+    /// Our query id
+    #[prost(string, tag = "1")]
+    pub query_id: String,
+    /// Our stage id
+    #[prost(uint64, tag = "2")]
+    pub stage_id: u64,
+    /// The task number within the stage
+    #[prost(uint64, tag = "3")]
+    pub task_number: u64,
+}
+
+impl Display for StageKey {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(
+            f,
+            "StageKey_QueryID_{}_StageID_{}_TaskNumber_{}",
+            self.query_id, self.stage_id, self.task_number
+        )
+    }
+}
+
 #[derive(Clone, PartialEq, ::prost::Message)]
 pub struct StageExecProto {
     /// Our query id