check if plan can be divided into tasks

Author: robtandy

src/common/ttl_map.rs

@@ -27,7 +27,6 @@ use dashmap::{DashMap, Entry};
 use datafusion::error::DataFusionError;
 use std::collections::HashSet;
 use std::hash::Hash;
-use std::mem;
 use std::sync::atomic::AtomicU64;
 #[cfg(test)]
 use std::sync::atomic::{AtomicUsize, Ordering::Relaxed};
@@ -94,7 +93,7 @@ where
                     shard.insert(key);
                 }
                 BucketOp::Clear => {
-                    let keys_to_delete = mem::replace(&mut shard, HashSet::new());
+                    let keys_to_delete = std::mem::replace(&mut shard, HashSet::new());
                     for key in keys_to_delete {
                         data.remove(&key);
                     }

src/common/util.rs

@@ -1,6 +1,5 @@
 use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion};
 use datafusion::error::Result;
-use datafusion::physical_plan::joins::PartitionMode;
 use datafusion::physical_plan::{displayable, ExecutionPlan, ExecutionPlanProperties};
 
 use std::fmt::Write;

src/physical_optimizer.rs

@@ -1,6 +1,7 @@
 use std::sync::Arc;
 
 use super::stage::ExecutionStage;
+use crate::common::util::can_be_divided;
 use crate::{plan::PartitionIsolatorExec, ArrowFlightReadExec};
 use datafusion::common::tree_node::TreeNodeRecursion;
 use datafusion::error::DataFusionError;
@@ -83,12 +84,14 @@ impl DistributedPhysicalOptimizerRule {
                     internal_datafusion_err!("Expected RepartitionExec to have a child"),
                 )?);
 
-                let maybe_isolated_plan = if let Some(ppt) = self.partitions_per_task {
-                    let isolated = Arc::new(PartitionIsolatorExec::new(child, ppt));
-                    plan.with_new_children(vec![isolated])?
-                } else {
-                    plan
-                };
+                let maybe_isolated_plan =
+                    if can_be_divided(&plan)? && self.partitions_per_task.is_some() {
+                        let ppt = self.partitions_per_task.unwrap();
+                        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(
@@ -120,7 +123,7 @@ impl DistributedPhysicalOptimizerRule {
     ) -> Result<ExecutionStage, DataFusionError> {
         let mut inputs = vec![];
 
-        let distributed = plan.transform_down(|plan| {
+        let distributed = plan.clone().transform_down(|plan| {
             let Some(node) = plan.as_any().downcast_ref::<ArrowFlightReadExec>() else {
                 return Ok(Transformed::no(plan));
             };
@@ -137,9 +140,13 @@ impl DistributedPhysicalOptimizerRule {
         let mut stage = ExecutionStage::new(query_id, *num, distributed.data, inputs);
         *num += 1;
 
-        if let Some(partitions_per_task) = self.partitions_per_task {
-            stage = stage.with_maximum_partitions_per_task(partitions_per_task);
-        }
+        stage = match (self.partitions_per_task, can_be_divided(&plan)?) {
+            (Some(partitions_per_task), true) => {
+                stage.with_maximum_partitions_per_task(partitions_per_task)
+            }
+            (_, _) => stage,
+        };
+
         stage.depth = depth;
 
         Ok(stage)

tests/tpch_validation_test.rs

@@ -6,12 +6,10 @@ mod tests {
     use async_trait::async_trait;
     use datafusion::error::DataFusionError;
     use datafusion::execution::SessionStateBuilder;
-    use datafusion::physical_plan::displayable;
+
     use datafusion::prelude::{SessionConfig, SessionContext};
     use datafusion_distributed::test_utils::localhost::start_localhost_context;
-    use datafusion_distributed::{
-        display_stage_graphviz, DistributedPhysicalOptimizerRule, ExecutionStage, SessionBuilder,
-    };
+    use datafusion_distributed::{DistributedPhysicalOptimizerRule, SessionBuilder};
     use futures::TryStreamExt;
     use std::error::Error;
     use std::sync::Arc;