add support for collect left hash joins

Author: robtandy

benchmarks/src/tpch/run.rs

@@ -36,14 +36,16 @@ use datafusion::datasource::listing::{
 };
 use datafusion::datasource::{MemTable, TableProvider};
 use datafusion::error::{DataFusionError, Result};
-use datafusion::execution::SessionStateBuilder;
+use datafusion::execution::{SessionState, SessionStateBuilder};
 use datafusion::physical_plan::display::DisplayableExecutionPlan;
 use datafusion::physical_plan::{collect, displayable};
 use datafusion::prelude::*;
 
 use crate::util::{print_memory_stats, BenchmarkRun, CommonOpt, QueryResult};
 use datafusion_distributed::test_utils::localhost::start_localhost_context;
-use datafusion_distributed::{DistributedPhysicalOptimizerRule, SessionBuilder};
+use datafusion_distributed::{
+    DistributedPhysicalOptimizerRule, DistributedSessionBuilder, DistributedSessionBuilderContext,
+};
 use log::info;
 use structopt::StructOpt;
 
@@ -110,31 +112,19 @@ pub struct RunOpt {
 }
 
 #[async_trait]
-impl SessionBuilder for RunOpt {
-    fn session_state_builder(
+impl DistributedSessionBuilder for RunOpt {
+    async fn build_session_state(
         &self,
-        mut builder: SessionStateBuilder,
-    ) -> Result<SessionStateBuilder, DataFusionError> {
-        let mut config = self
+        _ctx: DistributedSessionBuilderContext,
+    ) -> Result<SessionState, DataFusionError> {
+        let mut builder = SessionStateBuilder::new().with_default_features();
+
+        let config = self
             .common
             .config()?
             .with_collect_statistics(!self.disable_statistics)
             .with_target_partitions(self.partitions());
 
-        // // FIXME: these three options are critical for the correct function of the library
-        // // but we are not enforcing that the user sets them.  They are here at the moment
-        // // but we should figure out a way to do this better.
-        // config
-        //     .options_mut()
-        //     .optimizer
-        //     .hash_join_single_partition_threshold = 0;
-        // config
-        //     .options_mut()
-        //     .optimizer
-        //     .hash_join_single_partition_threshold_rows = 0;
-        //
-        // config.options_mut().optimizer.prefer_hash_join = self.prefer_hash_join;
-        // end critical options section
         let rt_builder = self.common.runtime_env_builder()?;
 
         if self.distributed {
@@ -145,17 +135,14 @@ impl SessionBuilder for RunOpt {
             builder = builder.with_physical_optimizer_rule(Arc::new(rule));
         }
 
-        Ok(builder
+        let state = builder
             .with_config(config)
-            .with_runtime_env(rt_builder.build_arc()?))
-    }
+            .with_runtime_env(rt_builder.build_arc()?)
+            .build();
 
-    async fn session_context(
-        &self,
-        ctx: SessionContext,
-    ) -> std::result::Result<SessionContext, DataFusionError> {
+        let ctx = SessionContext::from(state);
         self.register_tables(&ctx).await?;
-        Ok(ctx)
+        Ok(ctx.state())
     }
 }
 

src/common/mod.rs

@@ -1,2 +1,3 @@
+#[allow(unused)]
 pub mod ttl_map;
 pub mod util;

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::take(&mut shard);
                     for key in keys_to_delete {
                         data.remove(&key);
                     }
@@ -253,14 +252,14 @@ where
 
     /// run_gc_loop will continuously clear expired entries from the map, checking every `period`. The
     /// function terminates if `shutdown` is signalled.
-    async fn run_gc_loop(time: Arc<AtomicU64>, period: Duration, buckets: &Vec<Bucket<K>>) {
+    async fn run_gc_loop(time: Arc<AtomicU64>, period: Duration, buckets: &[Bucket<K>]) {
         loop {
             tokio::time::sleep(period).await;
             Self::gc(time.clone(), buckets);
         }
     }
 
-    fn gc(time: Arc<AtomicU64>, buckets: &Vec<Bucket<K>>) {
+    fn gc(time: Arc<AtomicU64>, buckets: &[Bucket<K>]) {
         let index = time.load(std::sync::atomic::Ordering::SeqCst) % buckets.len() as u64;
         buckets[index as usize].clear();
         time.fetch_add(1, std::sync::atomic::Ordering::SeqCst);

src/common/util.rs

@@ -37,44 +37,3 @@ 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
-<<<<<<< HEAD
-/// - HashJoinExec when not in Partitioned mode
-=======
->>>>>>> robtandy/nested_loop_joins
-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);
-<<<<<<< HEAD
-        } else if let Some(hash_join) = f
-            .as_any()
-            .downcast_ref::<datafusion::physical_plan::joins::HashJoinExec>()
-        {
-            if hash_join.partition_mode() != &PartitionMode::Partitioned {
-                has_unsplittable_plan = true;
-                return Ok(TreeNodeRecursion::Stop);
-            }
-=======
->>>>>>> robtandy/nested_loop_joins
-        }
-
-        Ok(TreeNodeRecursion::Continue)
-    };
-    plan.apply(search)?;
-    Ok(!has_unsplittable_plan)
-}