allow nested_loop_joins

Author: robtandy

src/common/util.rs

@@ -1,7 +1,10 @@
+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;
+use std::sync::Arc;
 
 pub fn display_plan_with_partition_in_out(plan: &dyn ExecutionPlan) -> Result<String> {
     let mut f = String::new();
@@ -34,3 +37,29 @@ pub fn display_plan_with_partition_in_out(plan: &dyn ExecutionPlan) -> Result<St
     visit(plan, 0, &mut f)?;
     Ok(f)
 }
+
+/// Returns a boolean indicating if this stage can be divided into more than one task.
+///
+/// Some Plan nodes need to materialize all partitions inorder to execute such as
+/// NestedLoopJoinExec.   Rewriting the plan to accommodate dividing it into tasks
+/// would result in redundant work.
+///
+/// The plans we cannot split are:
+/// - NestedLoopJoinExec
+pub fn can_be_divided(plan: &Arc<dyn ExecutionPlan>) -> Result<bool> {
+    // recursively check to see if this stages plan contains a NestedLoopJoinExec
+    let mut has_unsplittable_plan = false;
+    let search = |f: &Arc<dyn ExecutionPlan>| {
+        if f.as_any()
+            .downcast_ref::<datafusion::physical_plan::joins::NestedLoopJoinExec>()
+            .is_some()
+        {
+            has_unsplittable_plan = true;
+            return Ok(TreeNodeRecursion::Stop);
+        }
+
+        Ok(TreeNodeRecursion::Continue)
+    };
+    plan.apply(search)?;
+    Ok(!has_unsplittable_plan)
+}

src/test_utils/tpch.rs

@@ -160,13 +160,18 @@ where
 
 macro_rules! must_generate_tpch_table {
     ($generator:ident, $arrow:ident, $name:literal, $data_dir:expr) => {
-        generate_table(
-            // TODO: Consider adjusting the partitions and batch sizes.
-            $arrow::new($generator::new(SCALE_FACTOR, 1, 1)).with_batch_size(1000),
-            $name,
-            $data_dir,
-        )
-        .expect(concat!("Failed to generate ", $name, " table"));
+        let data_dir = $data_dir.join(format!("{}.parquet", $name));
+        fs::create_dir_all(data_dir.clone()).expect("Failed to create data directory");
+        // create three partitions for the table
+        (1..=3).for_each(|part| {
+            generate_table(
+                // TODO: Consider adjusting the partitions and batch sizes.
+                $arrow::new($generator::new(SCALE_FACTOR, part, 3)).with_batch_size(1000),
+                &format!("{}.parquet", part),
+                &data_dir.clone().into_boxed_path(),
+            )
+            .expect(concat!("Failed to generate ", $name, " table"));
+        });
     };
 }
 

tests/non_distributed_consistency_test.rs tests/tpch_validation_test.rstests/non_distributed_consistency_test.rs renamed to tests/tpch_validation_test.rs

@@ -6,9 +6,12 @@ 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::{DistributedPhysicalOptimizerRule, SessionBuilder};
+    use datafusion_distributed::{
+        display_stage_graphviz, DistributedPhysicalOptimizerRule, ExecutionStage, SessionBuilder,
+    };
     use futures::TryStreamExt;
     use std::error::Error;
     use std::sync::Arc;
@@ -119,8 +122,6 @@ mod tests {
     }
 
     #[tokio::test]
-    // TODO: Add support for NestedLoopJoinExec to support query 22.
-    #[ignore]
     async fn test_tpch_22() -> Result<(), Box<dyn Error>> {
         test_tpch_query(22).await
     }
@@ -155,7 +156,6 @@ mod tests {
 
             config.options_mut().optimizer.prefer_hash_join = true;
             // end critical options section
-
             let rule = DistributedPhysicalOptimizerRule::new().with_maximum_partitions_per_task(2);
             Ok(builder
                 .with_config(config)
@@ -174,7 +174,6 @@ mod tests {
     // and once in a non-distributed manner. For each query, it asserts that the results are identical.
     async fn run_tpch_query(ctx2: SessionContext, query_id: u8) -> Result<(), Box<dyn Error>> {
         ensure_tpch_data().await;
-
         let sql = get_test_tpch_query(query_id);
 
         // Context 1: Non-distributed execution.
@@ -205,19 +204,21 @@ mod tests {
             .await?;
         }
 
+        // Query 15 has three queries in it, one creating the view, the second
+        // executing, which we want to capture the output of, and the third
+        // tearing down the view
         let (stream1, stream2) = if query_id == 15 {
             let queries: Vec<&str> = sql
                 .split(';')
                 .map(str::trim)
                 .filter(|s| !s.is_empty())
                 .collect();
 
-            println!("queryies: {:?}", queries);
-
             ctx1.sql(queries[0]).await?.collect().await?;
             ctx2.sql(queries[0]).await?.collect().await?;
             let df1 = ctx1.sql(queries[1]).await?;
             let df2 = ctx2.sql(queries[1]).await?;
+
             let stream1 = df1.execute_stream().await?;
             let stream2 = df2.execute_stream().await?;
 
@@ -227,6 +228,7 @@ mod tests {
         } else {
             let stream1 = ctx1.sql(&sql).await?.execute_stream().await?;
             let stream2 = ctx2.sql(&sql).await?.execute_stream().await?;
+
             (stream1, stream2)
         };