logging

Author: matthewgapp

datafusion/core/src/physical_planner.rs

@@ -1322,6 +1322,7 @@ impl DefaultPhysicalPlanner {
                 }
                 LogicalPlan::RecursiveQuery(RecursiveQuery { name, static_term, recursive_term, is_distinct }) => {
                     let static_term = self.create_initial_plan(static_term, session_state).await?;
+                    println!("recursiving term: {:?}", recursive_term);
                     let recursive_term = self.create_initial_plan(recursive_term, session_state).await?;
 
                     Ok(Arc::new(RecursiveQueryExec::new(name.clone(), static_term, recursive_term, *is_distinct)))
@@ -1989,6 +1990,7 @@ impl DefaultPhysicalPlanner {
         F: FnMut(&dyn ExecutionPlan, &dyn PhysicalOptimizerRule),
     {
         let optimizers = session_state.physical_optimizers();
+        println!("optimizers: {}", optimizers.len());
         debug!(
             "Input physical plan:\n{}\n",
             displayable(plan.as_ref()).indent(false)

datafusion/execution/src/task.rs

@@ -33,7 +33,7 @@ use crate::{
     runtime_env::{RuntimeConfig, RuntimeEnv},
 };
 
-use arrow::record_batch::RecordBatch;
+use arrow::{compute::kernels::partition, record_batch::RecordBatch};
 // use futures::channel::mpsc::Receiver as SingleChannelReceiver;
 use tokio::sync::mpsc::Receiver as SingleChannelReceiver;
 // use futures::lock::Mutex;
@@ -65,6 +65,8 @@ pub struct TaskContext {
     window_functions: HashMap<String, Arc<WindowUDF>>,
     /// Runtime environment associated with this task context
     runtime: Arc<RuntimeEnv>,
+    // TODO: to remove, this is only for testing
+    active_partition: Mutex<Option<usize>>,
     /// Registered relation handlers
     relation_handlers: Mutex<HashMap<String, RelationHandler>>,
 }
@@ -84,6 +86,7 @@ impl Default for TaskContext {
             window_functions: HashMap::new(),
             runtime: Arc::new(runtime),
             relation_handlers: Mutex::new(HashMap::new()),
+            active_partition: Mutex::new(None),
         }
     }
 }
@@ -112,6 +115,19 @@ impl TaskContext {
             window_functions,
             runtime,
             relation_handlers: Mutex::new(HashMap::new()),
+            active_partition: Mutex::new(None),
+        }
+    }
+
+    pub fn set_and_increment_partition(&self) -> usize {
+        let mut partition_guard = self.active_partition.lock();
+        if let Some(ref mut partition) = &mut *partition_guard {
+            let prev = *partition;
+            *partition += 1;
+            prev
+        } else {
+            *partition_guard = Some(1);
+            0
         }
     }
 

datafusion/physical-expr/src/expressions/binary.rs

@@ -38,6 +38,7 @@ use arrow::compute::kernels::comparison::regexp_is_match_utf8_scalar;
 use arrow::compute::kernels::concat_elements::concat_elements_utf8;
 use arrow::datatypes::*;
 use arrow::record_batch::RecordBatch;
+use arrow::util::pretty::print_batches;
 use arrow_array::Datum;
 use datafusion_common::cast::as_boolean_array;
 use datafusion_common::{internal_err, DataFusionError, Result, ScalarValue};
@@ -292,6 +293,13 @@ impl PhysicalExpr for BinaryExpr {
         let schema = batch.schema();
         let input_schema = schema.as_ref();
 
+        println!(
+            "evaluate binary on batch left {:?}, right {:?}, operator {:?}",
+            left_value, right_value, self.op
+        );
+
+        print_batches(&[batch.clone()]).unwrap();
+
         if self.is_datum_operator() {
             return match (&left_value, &right_value) {
                 (ColumnarValue::Array(left), ColumnarValue::Array(right)) => {

datafusion/physical-plan/src/coalesce_batches.rs

@@ -159,6 +159,7 @@ impl ExecutionPlan for CoalesceBatchesExec {
         partition: usize,
         context: Arc<TaskContext>,
     ) -> Result<SendableRecordBatchStream> {
+        println!("CoalesceBatchesExec::execute partition={}", partition);
         Ok(Box::pin(CoalesceBatchesStream {
             input: self.input.execute(partition, context)?,
             schema: self.input.schema(),

datafusion/physical-plan/src/continuance.rs

@@ -102,7 +102,7 @@ impl ExecutionPlan for ContinuanceExec {
     }
 
     fn output_partitioning(&self) -> Partitioning {
-        Partitioning::UnknownPartitioning(0)
+        Partitioning::UnknownPartitioning(1)
     }
 
     fn maintains_input_order(&self) -> Vec<bool> {

datafusion/physical-plan/src/recursive_query.rs

@@ -41,7 +41,7 @@ use super::{
     metrics::{ExecutionPlanMetricsSet, MetricsSet},
     SendableRecordBatchStream, Statistics,
 };
-use arrow::error::{ArrowError, Result as ArrowResult};
+use arrow::error::ArrowError;
 use tokio::sync::mpsc::{Receiver, Sender};
 
 // use crate::execution::context::TaskContext;
@@ -150,6 +150,7 @@ impl ExecutionPlan for RecursiveQueryExec {
         partition: usize,
         context: Arc<TaskContext>,
     ) -> Result<SendableRecordBatchStream> {
+        println!("recursive query exec execute");
         // All partitions must be coalesced before coming to RecursiveQueryExec.
         // TODO: we might be able to handle multiple partitions in the future.
         if partition != 0 {
@@ -306,8 +307,14 @@ impl RecursiveQueryStream {
             }
         }
 
-        self.recursive_stream =
-            Some(self.recursive_term.execute(0, self.task_context.clone())?);
+        println!("executing recursive term");
+
+        // let partition = self.task_context.set_and_increment_partition();
+        let partition = 0;
+        self.recursive_stream = Some(
+            self.recursive_term
+                .execute(partition, self.task_context.clone())?,
+        );
         self.poll_next(cx)
     }
 }

datafusion/physical-plan/src/repartition/mod.rs

@@ -483,6 +483,7 @@ impl ExecutionPlan for RepartitionExec {
 
         // if this is the first partition to be invoked then we need to set up initial state
         if state.channels.is_empty() {
+            println!("channels are empty");
             let (txs, rxs) = if self.preserve_order {
                 let (txs, rxs) =
                     partition_aware_channels(num_input_partitions, num_output_partitions);
@@ -555,6 +556,8 @@ impl ExecutionPlan for RepartitionExec {
         );
 
         println!("execute partition: {}", partition);
+        println!("channel number: {}", state.channels.len());
+        println!("preserve order {}", self.preserve_order);
 
         // now return stream for the specified *output* partition which will
         // read from the channel
@@ -597,11 +600,13 @@ impl ExecutionPlan for RepartitionExec {
                 merge_reservation,
             )
         } else {
+            let input = rx.swap_remove(0);
+            println!("input partition: {:?}", input);
             Ok(Box::pin(RepartitionStream {
                 num_input_partitions,
                 num_input_partitions_processed: 0,
                 schema: self.input.schema(),
-                input: rx.swap_remove(0),
+                input,
                 drop_helper: Arc::clone(&state.abort_helper),
                 reservation,
             }))

datafusion/sql/src/query.rs

@@ -120,6 +120,8 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
                             // it as a CTE.
                             planner_context.insert_cte(cte_name.clone(), named_rel);
 
+                            // this uses the named_relation we inserted above to resolve the
+                            // relation
                             let recursive_plan = self
                                 .set_expr_to_plan(*right, &mut planner_context.clone())?;