Create protobuf/ folder

Author: gabotechs

src/distributed_ext.rs

@@ -2,7 +2,7 @@ use crate::channel_manager_ext::set_distributed_channel_resolver;
 use crate::config_extension_ext::{
     set_distributed_option_extension, set_distributed_option_extension_from_headers,
 };
-use crate::user_codec_ext::set_distributed_user_codec;
+use crate::protobuf::set_distributed_user_codec;
 use crate::ChannelResolver;
 use datafusion::common::DataFusionError;
 use datafusion::config::ConfigExtension;

src/execution_plans/arrow_flight_read.rs

@@ -1,9 +1,9 @@
 use crate::channel_manager_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::DistributedCodec;
 use crate::flight_service::{DoGet, StageKey};
-use crate::stage::{proto_from_stage, ExecutionStage};
+use crate::protobuf::{proto_from_stage, DistributedCodec};
+use crate::stage::ExecutionStage;
 use crate::ChannelResolver;
 use arrow_flight::decode::FlightRecordBatchStream;
 use arrow_flight::error::FlightError;

src/execution_plans/mod.rs

@@ -1,7 +1,5 @@
 mod arrow_flight_read;
-mod codec;
 mod partition_isolator;
 
 pub use arrow_flight_read::ArrowFlightReadExec;
-pub use codec::DistributedCodec;
 pub use partition_isolator::{PartitionGroup, PartitionIsolatorExec};

src/flight_service/do_get.rs

@@ -1,10 +1,11 @@
 use super::service::StageKey;
 use crate::config_extension_ext::ContextGrpcMetadata;
 use crate::errors::datafusion_error_to_tonic_status;
-use crate::execution_plans::{DistributedCodec, PartitionGroup};
+use crate::execution_plans::PartitionGroup;
 use crate::flight_service::service::ArrowFlightEndpoint;
 use crate::flight_service::session_builder::DistributedSessionBuilderContext;
-use crate::stage::{stage_from_proto, ExecutionStage, ExecutionStageProto};
+use crate::protobuf::{stage_from_proto, DistributedCodec, ExecutionStageProto};
+use crate::stage::ExecutionStage;
 use arrow_flight::encode::FlightDataEncoderBuilder;
 use arrow_flight::error::FlightError;
 use arrow_flight::flight_service_server::FlightService;
@@ -169,16 +170,15 @@ impl ArrowFlightEndpoint {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::flight_service::session_builder::DefaultSessionBuilder;
+    use crate::stage::ExecutionTask;
+    use arrow_flight::Ticket;
+    use prost::{bytes::Bytes, Message};
+    use tonic::Request;
     use uuid::Uuid;
 
     #[tokio::test]
     async fn test_task_data_partition_counting() {
-        use crate::flight_service::session_builder::DefaultSessionBuilder;
-        use crate::task::ExecutionTask;
-        use arrow_flight::Ticket;
-        use prost::{bytes::Bytes, Message};
-        use tonic::Request;
-
         // Create ArrowFlightEndpoint with DefaultSessionBuilder
         let endpoint =
             ArrowFlightEndpoint::try_new(DefaultSessionBuilder).expect("Failed to create endpoint");

src/lib.rs

@@ -9,9 +9,8 @@ mod execution_plans;
 mod flight_service;
 mod physical_optimizer;
 mod stage;
-mod task;
-mod user_codec_ext;
 
+mod protobuf;
 #[cfg(any(feature = "integration", test))]
 pub mod test_utils;
 

src/execution_plans/codec.rs renamed to src/protobuf/distributed_codec.rs

@@ -1,7 +1,6 @@
-use super::PartitionIsolatorExec;
+use super::get_distributed_user_codec;
 use crate::common::ComposedPhysicalExtensionCodec;
-use crate::execution_plans::arrow_flight_read::ArrowFlightReadExec;
-use crate::user_codec_ext::get_distributed_user_codec;
+use crate::{ArrowFlightReadExec, PartitionIsolatorExec};
 use datafusion::arrow::datatypes::Schema;
 use datafusion::execution::FunctionRegistry;
 use datafusion::physical_plan::ExecutionPlan;

src/stage/proto.rs renamed to src/protobuf/execution_stage_proto.rs

@@ -1,5 +1,7 @@
 use std::sync::Arc;
 
+use crate::stage::ExecutionTask;
+use crate::ExecutionStage;
 use datafusion::{
     common::internal_datafusion_err,
     error::{DataFusionError, Result},
@@ -11,10 +13,6 @@ use datafusion_proto::{
     protobuf::PhysicalPlanNode,
 };
 
-use crate::task::ExecutionTask;
-
-use super::ExecutionStage;
-
 #[derive(Clone, PartialEq, ::prost::Message)]
 pub struct ExecutionStageProto {
     /// Our query id
@@ -98,6 +96,9 @@ pub fn stage_from_proto(
 mod tests {
     use std::sync::Arc;
 
+    use crate::protobuf::execution_stage_proto::ExecutionStageProto;
+    use crate::protobuf::{proto_from_stage, stage_from_proto};
+    use crate::ExecutionStage;
     use datafusion::{
         arrow::{
             array::{RecordBatch, StringArray, UInt8Array},
@@ -112,9 +113,6 @@ mod tests {
     use prost::Message;
     use uuid::Uuid;
 
-    use crate::stage::proto::proto_from_stage;
-    use crate::stage::{proto::stage_from_proto, ExecutionStage, ExecutionStageProto};
-
     // create a simple mem table
     fn create_mem_table() -> Arc<MemTable> {
         let fields = vec![

src/protobuf/mod.rs

@@ -0,0 +1,7 @@
+mod distributed_codec;
+mod execution_stage_proto;
+mod user_codec;
+
+pub(crate) use distributed_codec::DistributedCodec;
+pub(crate) use execution_stage_proto::{proto_from_stage, stage_from_proto, ExecutionStageProto};
+pub(crate) use user_codec::{get_distributed_user_codec, set_distributed_user_codec};

src/user_codec_ext.rs renamed to src/protobuf/user_codec.rs

@@ -1,9 +1,7 @@
-use super::ExecutionStage;
+use super::{ExecutionStage, ExecutionTask};
+use crate::stage::task::format_pg;
+use crate::ArrowFlightReadExec;
 use crate::PartitionIsolatorExec;
-use crate::{
-    task::{format_pg, ExecutionTask},
-    ArrowFlightReadExec,
-};
 use datafusion::physical_plan::{displayable, ExecutionPlan, ExecutionPlanProperties};
 use datafusion::{
     error::Result,