execution debugging

Signed-off-by: Alex Qyoun-ae <[email protected]>

Author: MazterQyou

datafusion/core/src/physical_plan/hash_join.rs

@@ -298,6 +298,7 @@ impl ExecutionPlan for HashJoinExec {
         partition: usize,
         context: Arc<TaskContext>,
     ) -> Result<SendableRecordBatchStream> {
+        println!("HASH JOIN EXEC ENGAGED\n");
         // This is a hacky way to support type coercion for join expressions
         // Without this it would panic later, in build_join_indexes => equal_rows, when it would try to downcast both sides to same primitive type
         // TODO Remove this after rebasing on top of commit ac2e5d15 "Support type coercion for equijoin (#4666)". It was first released at DF 16.0
@@ -326,6 +327,7 @@ impl ExecutionPlan for HashJoinExec {
 
         // we only want to compute the build side once for PartitionMode::CollectLeft
         let left_data = {
+            println!("PARTITION MODE: {:#?}", self.mode);
             match self.mode {
                 PartitionMode::CollectLeft => {
                     let mut build_side = self.build_side.lock().await;
@@ -425,6 +427,7 @@ impl ExecutionPlan for HashJoinExec {
                         concat_batches(&self.left.schema(), &batches, num_rows)?;
 
                     let left_side = Arc::new((hashmap, single_batch));
+                    println!("LEFT SIDE HASH JOIN:\n{:#?}\n", left_side);
 
                     debug!(
                         "Built build-side {} of hash join containing {} rows in {} ms",
@@ -1086,7 +1089,12 @@ impl Stream for HashJoinStream {
                             JoinType::Inner | JoinType::Right => {}
                         }
                     }
-                    Some(result.map(|x| x.0))
+                    let result = Some(result.map(|x| x.0));
+                    println!(
+                        "HASH JOIN EXEC ISSUING BATCH 1:\n{:#?}\n",
+                        result.as_ref().unwrap()
+                    );
+                    result
                 }
                 other => {
                     let timer = self.join_metrics.join_time.timer();
@@ -1115,6 +1123,7 @@ impl Stream for HashJoinStream {
                             }
                             timer.done();
                             self.is_exhausted = true;
+                            println!("HASH JOIN EXEC ISSUING BATCH 2:\n{:#?}\n", result);
                             return Some(result);
                         }
                         JoinType::Left

datafusion/core/src/physical_plan/planner.rs

@@ -39,7 +39,6 @@ use crate::physical_optimizer::optimizer::PhysicalOptimizerRule;
 use crate::physical_plan::coalesce_partitions::CoalescePartitionsExec;
 use crate::physical_plan::cross_join::CrossJoinExec;
 use crate::physical_plan::explain::ExplainExec;
-use crate::physical_plan::expressions;
 use crate::physical_plan::expressions::{
     CaseExpr, Column, GetIndexedFieldExpr, Literal, PhysicalSortExpr,
 };
@@ -54,6 +53,7 @@ use crate::physical_plan::subquery::SubqueryExec;
 use crate::physical_plan::udf;
 use crate::physical_plan::udtf;
 use crate::physical_plan::windows::WindowAggExec;
+use crate::physical_plan::{expressions, DisplayFormatType};
 use crate::physical_plan::{join_utils, Partitioning};
 use crate::physical_plan::{AggregateExpr, ExecutionPlan, PhysicalExpr, WindowExpr};
 use crate::scalar::ScalarValue;
@@ -74,6 +74,7 @@ use datafusion_physical_expr::expressions::{any, OuterColumn};
 use futures::future::BoxFuture;
 use futures::{FutureExt, StreamExt, TryStreamExt};
 use log::{debug, trace};
+use std::fmt::Debug;
 use std::sync::Arc;
 
 fn create_function_physical_name(
@@ -434,15 +435,45 @@ impl PhysicalPlanner for DefaultPhysicalPlanner {
         logical_plan: &LogicalPlan,
         session_state: &SessionState,
     ) -> Result<Arc<dyn ExecutionPlan>> {
-        match self.handle_explain(logical_plan, session_state).await? {
+        let x = match self.handle_explain(logical_plan, session_state).await? {
             Some(plan) => Ok(plan),
             None => {
                 let plan = self
                     .create_initial_plan(logical_plan, session_state)
                     .await?;
                 self.optimize_internal(plan, session_state, |_, _| {})
             }
+        }?;
+
+        struct X {
+            plan: Arc<dyn ExecutionPlan>,
+            level: usize,
+        }
+
+        impl Debug for X {
+            fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+                for _ in 0..self.level {
+                    write!(f, "  ")?;
+                }
+                self.plan.fmt_as(DisplayFormatType::Default, f)?;
+                writeln!(f)?;
+                for child in self.plan.children() {
+                    let a = X {
+                        plan: Arc::clone(&child),
+                        level: self.level + 1,
+                    };
+                    write!(f, "{:?}", a)?;
+                }
+                Ok(())
+            }
         }
+
+        let a = X {
+            plan: Arc::clone(&x),
+            level: 0,
+        };
+        println!("PHYSICAL PLAN:\n{:?}\n", a);
+        Ok(x)
     }
 
     /// Create a physical expression from a logical expression
@@ -645,8 +676,10 @@ impl DefaultPhysicalPlanner {
                     aggr_expr,
                     ..
                 }) => {
+                    println!("\nAGGREGATE INITIAL PLAN");
                     // Initially need to perform the aggregate and then merge the partitions
                     let input_exec = self.create_initial_plan(input, session_state).await?;
+                    println!("INPUT EXEC: {:?}", input_exec);
                     let physical_input_schema = input_exec.schema();
                     let logical_input_schema = input.as_ref().schema();
 
@@ -664,6 +697,7 @@ impl DefaultPhysicalPlanner {
                             ))
                         })
                         .collect::<Result<Vec<_>>>()?;
+                    println!("GROUPS: {:?}", groups);
                     let aggregates = aggr_expr
                         .iter()
                         .map(|e| {
@@ -675,6 +709,7 @@ impl DefaultPhysicalPlanner {
                             )
                         })
                         .collect::<Result<Vec<_>>>()?;
+                    println!("AGGREGATES: {:?}", aggregates);
 
                     let initial_aggr = Arc::new(HashAggregateExec::try_new(
                         AggregateMode::Partial,
@@ -683,20 +718,27 @@ impl DefaultPhysicalPlanner {
                         input_exec,
                         physical_input_schema.clone(),
                     )?);
+                    println!("INITIAL AGGR: {:?}", initial_aggr);
 
                     // update group column indices based on partial aggregate plan evaluation
                     let final_group: Vec<Arc<dyn PhysicalExpr>> = initial_aggr.output_group_expr();
+                    println!("FINAL GROUP: {:?}", final_group);
 
                     // TODO: dictionary type not yet supported in Hash Repartition
                     let contains_dict = groups
                         .iter()
                         .flat_map(|x| x.0.data_type(physical_input_schema.as_ref()))
                         .any(|x| matches!(x, DataType::Dictionary(_, _)));
+                    println!("CONTAINS DICT: {:?}", contains_dict);
 
                     let can_repartition = !groups.is_empty()
                         && session_state.config.target_partitions > 1
                         && session_state.config.repartition_aggregations
                         && !contains_dict;
+                    println!("GROUPS IS EMPTY: {:?}", groups.is_empty());
+                    println!("TARGET PARTITIONS: {:?}", session_state.config.target_partitions);
+                    println!("REPARTITION AGGREGATIONS: {:?}", session_state.config.repartition_aggregations);
+                    println!("CAN REPARTITION: {:?}", can_repartition);
 
                     let (initial_aggr, next_partition_mode): (
                         Arc<dyn ExecutionPlan>,
@@ -717,6 +759,9 @@ impl DefaultPhysicalPlanner {
                         // first aggregation and the expressions corresponding to the respective aggregate
                         (initial_aggr, AggregateMode::Final)
                     };
+                    println!("INITIAL AGGR AFTER REPARTITION: {:?}", initial_aggr);
+                    println!("NEXT PARTITION MODE: {:?}", next_partition_mode);
+                    println!();
 
                     Ok(Arc::new(HashAggregateExec::try_new(
                         next_partition_mode,