big fmt

Author: adriangb

src/common/ttl_map.rs

@@ -27,10 +27,10 @@ use dashmap::{DashMap, Entry};
 use datafusion::error::DataFusionError;
 use std::collections::HashSet;
 use std::hash::Hash;
+use std::sync::Arc;
 use std::sync::atomic::AtomicU64;
 #[cfg(test)]
 use std::sync::atomic::{AtomicUsize, Ordering::Relaxed};
-use std::sync::Arc;
 use std::time::Duration;
 use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender};
 
@@ -289,7 +289,7 @@ where
 mod tests {
     use super::*;
     use std::sync::atomic::Ordering;
-    use tokio::time::{sleep, Duration};
+    use tokio::time::{Duration, sleep};
 
     #[tokio::test]
     async fn test_basic_insert_and_get() {

src/config_extension_ext.rs

@@ -1,4 +1,4 @@
-use datafusion::common::{internal_datafusion_err, DataFusionError};
+use datafusion::common::{DataFusionError, internal_datafusion_err};
 use datafusion::config::ConfigExtension;
 use datafusion::prelude::SessionConfig;
 use http::{HeaderMap, HeaderName};
@@ -84,8 +84,8 @@ impl ContextGrpcMetadata {
 #[cfg(test)]
 mod tests {
     use crate::config_extension_ext::{
-        set_distributed_option_extension, set_distributed_option_extension_from_headers,
-        ContextGrpcMetadata,
+        ContextGrpcMetadata, set_distributed_option_extension,
+        set_distributed_option_extension_from_headers,
     };
     use datafusion::common::extensions_options;
     use datafusion::config::ConfigExtension;

src/distributed_ext.rs

@@ -1,9 +1,9 @@
+use crate::ChannelResolver;
 use crate::channel_resolver_ext::set_distributed_channel_resolver;
 use crate::config_extension_ext::{
     set_distributed_option_extension, set_distributed_option_extension_from_headers,
 };
 use crate::protobuf::set_distributed_user_codec;
-use crate::ChannelResolver;
 use datafusion::common::DataFusionError;
 use datafusion::config::ConfigExtension;
 use datafusion::execution::{SessionState, SessionStateBuilder};

src/distributed_physical_optimizer_rule.rs

@@ -12,7 +12,7 @@ use datafusion::{
     config::ConfigOptions,
     error::Result,
     physical_optimizer::PhysicalOptimizerRule,
-    physical_plan::{repartition::RepartitionExec, ExecutionPlan},
+    physical_plan::{ExecutionPlan, repartition::RepartitionExec},
 };
 use uuid::Uuid;
 
@@ -83,13 +83,14 @@ impl DistributedPhysicalOptimizerRule {
                     internal_datafusion_err!("Expected RepartitionExec to have a child"),
                 )?);
 
-                let maybe_isolated_plan =
-                    if let Some(ppt) = self.partitions_per_task && can_be_divided(&plan)? {
-                        let isolated = Arc::new(PartitionIsolatorExec::new(child, ppt));
-                        plan.with_new_children(vec![isolated])?
-                    } else {
-                        plan
-                    };
+                let maybe_isolated_plan = if let Some(ppt) = self.partitions_per_task
+                    && can_be_divided(&plan)?
+                {
+                    let isolated = Arc::new(PartitionIsolatorExec::new(child, ppt));
+                    plan.with_new_children(vec![isolated])?
+                } else {
+                    plan
+                };
 
                 return Ok(Transformed::yes(Arc::new(
                     ArrowFlightReadExec::new_pending(Arc::clone(&maybe_isolated_plan)),

src/errors/arrow_error.rs

@@ -130,7 +130,10 @@ impl ArrowErrorProto {
                 ctx: ctx.cloned(),
             },
             ArrowError::OffsetOverflowError(size) => ArrowErrorProto {
-                inner: Some(ArrowErrorInnerProto::ExternalError(format!("Offset overflow error: {}", size))),
+                inner: Some(ArrowErrorInnerProto::ExternalError(format!(
+                    "Offset overflow error: {}",
+                    size
+                ))),
                 ctx: ctx.cloned(),
             },
         }

src/execution_plans/arrow_flight_read.rs

@@ -1,15 +1,15 @@
+use crate::ChannelResolver;
 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;
 use crate::metrics::proto::MetricsSetProto;
-use crate::protobuf::{proto_from_stage, DistributedCodec, StageKey};
-use crate::ChannelResolver;
+use crate::protobuf::{DistributedCodec, StageKey, proto_from_stage};
+use arrow_flight::Ticket;
 use arrow_flight::decode::FlightRecordBatchStream;
 use arrow_flight::error::FlightError;
 use arrow_flight::flight_service_client::FlightServiceClient;
-use arrow_flight::Ticket;
 use dashmap::DashMap;
 use datafusion::arrow::datatypes::SchemaRef;
 use datafusion::common::{exec_err, internal_datafusion_err, internal_err, plan_err};
@@ -27,8 +27,8 @@ use prost::Message;
 use std::any::Any;
 use std::fmt::Formatter;
 use std::sync::Arc;
-use tonic::metadata::MetadataMap;
 use tonic::Request;
+use tonic::metadata::MetadataMap;
 
 /// This node has two variants.
 /// 1. Pending: it acts as a placeholder for the distributed optimization step to mark it as ready.
@@ -158,7 +158,9 @@ impl ExecutionPlan for ArrowFlightReadExec {
         context: Arc<TaskContext>,
     ) -> Result<SendableRecordBatchStream, DataFusionError> {
         let ArrowFlightReadExec::Ready(self_ready) = self else {
-            return exec_err!("ArrowFlightReadExec is not ready, was the distributed optimization step performed?");
+            return exec_err!(
+                "ArrowFlightReadExec is not ready, was the distributed optimization step performed?"
+            );
         };
 
         // get the channel manager and current stage from our context

src/execution_plans/metrics.rs

@@ -4,7 +4,7 @@ use std::collections::HashMap;
 use std::sync::Arc;
 
 use crate::execution_plans::{ArrowFlightReadExec, StageExec};
-use crate::metrics::proto::{metrics_set_proto_to_df, MetricsSetProto};
+use crate::metrics::proto::{MetricsSetProto, metrics_set_proto_to_df};
 use crate::protobuf::StageKey;
 use datafusion::common::internal_err;
 use datafusion::common::tree_node::{Transformed, TreeNode, TreeNodeRecursion, TreeNodeRewriter};
@@ -143,7 +143,11 @@ impl TaskMetricsRewriter {
     ) -> Result<Arc<dyn ExecutionPlan>> {
         let transformed = plan.rewrite(&mut self)?;
         if self.idx != self.metrics.len() {
-            return internal_err!("too many metrics sets provided to rewrite task: {} metrics sets provided, {} nodes in plan", self.metrics.len(), self.idx);
+            return internal_err!(
+                "too many metrics sets provided to rewrite task: {} metrics sets provided, {} nodes in plan",
+                self.metrics.len(),
+                self.idx
+            );
         }
         Ok(transformed.data)
     }
@@ -269,11 +273,11 @@ mod tests {
     use datafusion::arrow::array::{Int32Array, StringArray};
     use datafusion::arrow::record_batch::RecordBatch;
 
-    use crate::test_utils::in_memory_channel_resolver::InMemoryChannelResolver;
-    use crate::test_utils::session_context::register_temp_parquet_table;
     use crate::DistributedExt;
     use crate::DistributedPhysicalOptimizerRule;
-    use datafusion::execution::{context::SessionContext, SessionStateBuilder};
+    use crate::test_utils::in_memory_channel_resolver::InMemoryChannelResolver;
+    use crate::test_utils::session_context::register_temp_parquet_table;
+    use datafusion::execution::{SessionStateBuilder, context::SessionContext};
     use datafusion::physical_plan::metrics::MetricValue;
     use datafusion::prelude::SessionConfig;
     use datafusion::{

src/execution_plans/mod.rs

@@ -5,4 +5,4 @@ mod stage;
 
 pub use arrow_flight_read::ArrowFlightReadExec;
 pub use partition_isolator::{PartitionGroup, PartitionIsolatorExec};
-pub use stage::{display_plan_graphviz, ExecutionTask, StageExec};
+pub use stage::{ExecutionTask, StageExec, display_plan_graphviz};

src/execution_plans/stage.rs

@@ -7,7 +7,7 @@ use datafusion::common::{
 use datafusion::error::{DataFusionError, Result};
 use datafusion::execution::TaskContext;
 use datafusion::physical_plan::{
-    displayable, DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties,
+    DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, displayable,
 };
 use datafusion::prelude::SessionContext;
 use itertools::Itertools;

src/flight_service/do_get.rs

@@ -3,18 +3,18 @@ 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, StageKey};
+use crate::protobuf::{DistributedCodec, StageExecProto, StageKey, stage_from_proto};
+use arrow_flight::Ticket;
 use arrow_flight::encode::FlightDataEncoderBuilder;
 use arrow_flight::error::FlightError;
 use arrow_flight::flight_service_server::FlightService;
-use arrow_flight::Ticket;
 use datafusion::execution::{SendableRecordBatchStream, SessionState};
 use futures::TryStreamExt;
 use http::HeaderMap;
 use prost::Message;
 use std::fmt::Display;
-use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::Arc;
+use std::sync::atomic::{AtomicUsize, Ordering};
 use tokio::sync::OnceCell;
 use tonic::{Request, Response, Status};
 
@@ -172,17 +172,17 @@ fn record_batch_stream_to_response(
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::ExecutionTask;
     use crate::flight_service::session_builder::DefaultSessionBuilder;
     use crate::protobuf::proto_from_stage;
-    use crate::ExecutionTask;
     use arrow::datatypes::{Schema, SchemaRef};
     use arrow_flight::Ticket;
     use datafusion::physical_expr::Partitioning;
+    use datafusion::physical_plan::ExecutionPlan;
     use datafusion::physical_plan::empty::EmptyExec;
     use datafusion::physical_plan::repartition::RepartitionExec;
-    use datafusion::physical_plan::ExecutionPlan;
     use datafusion_proto::physical_plan::DefaultPhysicalExtensionCodec;
-    use prost::{bytes::Bytes, Message};
+    use prost::{Message, bytes::Bytes};
     use tonic::Request;
     use uuid::Uuid;