Merge pull request #39 from datafusion-contrib/run-cargo-fmt

Run `cargo fmt`

Author: gabotechs

src/codec.rs

@@ -2,17 +2,15 @@ use std::sync::Arc;
 
 use arrow::datatypes::Schema;
 use datafusion::{
-    common::{Result, internal_datafusion_err, internal_err},
+    common::{internal_datafusion_err, internal_err, Result},
     datasource::source::DataSourceExec,
     execution::FunctionRegistry,
-    physical_plan::{ExecutionPlan, displayable},
+    physical_plan::{displayable, ExecutionPlan},
 };
 use datafusion_proto::{
     physical_plan::{
-        DefaultPhysicalExtensionCodec,
-        PhysicalExtensionCodec,
-        from_proto::parse_protobuf_partitioning,
-        to_proto::serialize_partitioning,
+        from_proto::parse_protobuf_partitioning, to_proto::serialize_partitioning,
+        DefaultPhysicalExtensionCodec, PhysicalExtensionCodec,
     },
     protobuf,
 };
@@ -23,11 +21,8 @@ use crate::{
     logging::trace,
     max_rows::MaxRowsExec,
     protobuf::{
-        DfRayExecNode,
-        DfRayStageReaderExecNode,
-        MaxRowsExecNode,
+        df_ray_exec_node::Payload, DfRayExecNode, DfRayStageReaderExecNode, MaxRowsExecNode,
         PartitionIsolatorExecNode,
-        df_ray_exec_node::Payload,
     },
     stage_reader::DFRayStageReaderExec,
 };
@@ -95,16 +90,12 @@ impl PhysicalExtensionCodec for DFRayCodec {
                         )))
                     }
                 }
-                Payload::NumpangExec(_) => {
-                    Err(internal_datafusion_err!(
-                        "NumpangExec not supported in open source version"
-                    ))
-                }
-                Payload::ContextExec(_) => {
-                    Err(internal_datafusion_err!(
-                        "ContextExec not supported in open source version"
-                    ))
-                }
+                Payload::NumpangExec(_) => Err(internal_datafusion_err!(
+                    "NumpangExec not supported in open source version"
+                )),
+                Payload::ContextExec(_) => Err(internal_datafusion_err!(
+                    "ContextExec not supported in open source version"
+                )),
             }
         } else {
             internal_err!("cannot decode proto extension in dfray codec")
@@ -169,16 +160,14 @@ mod test {
 
     use arrow::datatypes::DataType;
     use datafusion::{
-        physical_plan::{Partitioning, displayable},
+        physical_plan::{displayable, Partitioning},
         prelude::SessionContext,
     };
     use datafusion_proto::physical_plan::AsExecutionPlan;
 
     use super::*;
     use crate::{
-        isolator::PartitionIsolatorExec,
-        max_rows::MaxRowsExec,
-        stage_reader::DFRayStageReaderExec,
+        isolator::PartitionIsolatorExec, max_rows::MaxRowsExec, stage_reader::DFRayStageReaderExec,
     };
 
     fn create_test_schema() -> Arc<arrow::datatypes::Schema> {

src/explain.rs

@@ -15,27 +15,19 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use std::{
-    any::Any,
-    fmt::Formatter,
-    sync::Arc,
-};
+use std::{any::Any, fmt::Formatter, sync::Arc};
 
-use arrow::{
-    array::StringArray,
-    datatypes::SchemaRef,
-    record_batch::RecordBatch,
-};
+use arrow::{array::StringArray, datatypes::SchemaRef, record_batch::RecordBatch};
 use datafusion::{
     execution::TaskContext,
+    physical_expr::EquivalenceProperties,
     physical_plan::{
+        displayable,
         execution_plan::{Boundedness, EmissionType},
         memory::MemoryStream,
-        ExecutionPlan, Partitioning,
-        PlanProperties, DisplayAs, DisplayFormatType,
-        SendableRecordBatchStream, displayable,
+        DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties,
+        SendableRecordBatchStream,
     },
-    physical_expr::EquivalenceProperties,
 };
 
 /// Custom distributed EXPLAIN execution plan that also returns distributed plan and stages
@@ -96,7 +88,10 @@ impl DistributedExplainExec {
         let mut result = String::new();
         for (i, stage) in stages.iter().enumerate() {
             result.push_str(&format!("Stage {}:\n", stage.stage_id));
-            result.push_str(&format!("  Partition Groups: {:?}\n", stage.partition_groups));
+            result.push_str(&format!(
+                "  Partition Groups: {:?}\n",
+                stage.partition_groups
+            ));
             result.push_str(&format!("  Full Partitions: {}\n", stage.full_partitions));
             result.push_str("  Plan:\n");
             let plan_display = format!("{}", displayable(stage.plan.as_ref()).indent(true));
@@ -150,13 +145,13 @@ impl ExecutionPlan for DistributedExplainExec {
         _context: Arc<TaskContext>,
     ) -> datafusion::error::Result<SendableRecordBatchStream> {
         let schema = self.schema.clone();
-        
+
         // Create the result data with our 4 plan types
         let plan_types = StringArray::from(vec![
-            "logical_plan", 
-            "physical_plan", 
-            "distributed_plan", 
-            "distributed_stages"
+            "logical_plan",
+            "physical_plan",
+            "distributed_plan",
+            "distributed_stages",
         ]);
         let plans = StringArray::from(vec![
             self.logical_plan.as_str(),
@@ -165,14 +160,13 @@ impl ExecutionPlan for DistributedExplainExec {
             self.distributed_stages.as_str(),
         ]);
 
-        let batch = RecordBatch::try_new(
-            schema.clone(),
-            vec![Arc::new(plan_types), Arc::new(plans)],
-        ).map_err(|e| datafusion::error::DataFusionError::ArrowError(e, None))?;
+        let batch =
+            RecordBatch::try_new(schema.clone(), vec![Arc::new(plan_types), Arc::new(plans)])
+                .map_err(|e| datafusion::error::DataFusionError::ArrowError(e, None))?;
 
         // Use MemoryStream which is designed for DataFusion execution plans
         let stream = MemoryStream::try_new(vec![batch], schema, None)?;
-        
+
         Ok(Box::pin(stream))
     }
 
@@ -182,16 +176,20 @@ impl ExecutionPlan for DistributedExplainExec {
 }
 
 /// Check if this is an EXPLAIN query (but not EXPLAIN ANALYZE)
-/// 
+///
 /// This function distinguishes between:
 /// - EXPLAIN queries (returns true) - show plan information only
 /// - EXPLAIN ANALYZE queries (returns false) - execute and show runtime stats
 /// - Regular queries (returns false) - normal query execution
 pub fn is_explain_query(query: &str) -> bool {
     let query_upper = query.trim().to_uppercase();
     // Must start with "EXPLAIN" followed by whitespace or end of string
-    let is_explain = query_upper.starts_with("EXPLAIN") && 
-        (query_upper.len() == 7 || query_upper.chars().nth(7).is_some_and(|c| c.is_whitespace()));
+    let is_explain = query_upper.starts_with("EXPLAIN")
+        && (query_upper.len() == 7
+            || query_upper
+                .chars()
+                .nth(7)
+                .is_some_and(|c| c.is_whitespace()));
     let is_explain_analyze = query_upper.starts_with("EXPLAIN ANALYZE");
     is_explain && !is_explain_analyze
 }
@@ -223,7 +221,7 @@ mod tests {
         // Test edge cases
         assert!(!is_explain_query(""));
         assert!(!is_explain_query("   "));
-        assert!(!is_explain_query("EXPLAINSELECT"));  // No space
+        assert!(!is_explain_query("EXPLAINSELECT")); // No space
         assert!(is_explain_query("EXPLAIN")); // Just EXPLAIN
     }
 }

src/flight.rs

@@ -18,21 +18,9 @@
 use std::sync::Arc;
 
 use arrow_flight::{
-    Action,
-    ActionType,
-    Criteria,
-    Empty,
-    FlightData,
-    FlightDescriptor,
-    FlightInfo,
-    HandshakeRequest,
-    HandshakeResponse,
-    PollInfo,
-    PutResult,
-    SchemaResult,
-    Ticket,
-    flight_service_server::FlightService,
-    sql::server::FlightSqlService,
+    flight_service_server::FlightService, sql::server::FlightSqlService, Action, ActionType,
+    Criteria, Empty, FlightData, FlightDescriptor, FlightInfo, HandshakeRequest, HandshakeResponse,
+    PollInfo, PutResult, SchemaResult, Ticket,
 };
 use futures::stream::BoxStream;
 use tonic::{Request, Response, Status, Streaming};
@@ -46,7 +34,7 @@ pub type DoActionStream = BoxStream<'static, Result<arrow_flight::Result, Status
 pub trait FlightHandler: Send + Sync {
     async fn do_get(&self, request: Request<Ticket>) -> Result<Response<DoGetStream>, Status>;
     async fn do_action(&self, request: Request<Action>)
-    -> Result<Response<DoActionStream>, Status>;
+        -> Result<Response<DoActionStream>, Status>;
 }
 
 pub struct FlightServ {