Add debug print statement to DataFrame.collect for improved traceability

kosiew · Jan 13, 2025 · 42f4e7d · 42f4e7d
1 parent a8d4077
commit 42f4e7d
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 0 deletions.
diff --git a/datafusion/core/src/dataframe/mod.rs b/datafusion/core/src/dataframe/mod.rs
@@ -1208,6 +1208,7 @@ impl DataFrame {
     /// # }
     /// ```
     pub async fn collect(self) -> Result<Vec<RecordBatch>> {
+        println!("==> Datafrrame.collect");
         let task_ctx = Arc::new(self.task_ctx());
         let plan = self.create_physical_plan().await?;
         collect(plan, task_ctx).await

diff --git a/datafusion/physical-plan/src/execution_plan.rs b/datafusion/physical-plan/src/execution_plan.rs
@@ -814,6 +814,8 @@ pub fn execute_stream(
     plan: Arc<dyn ExecutionPlan>,
     context: Arc<TaskContext>,
 ) -> Result<SendableRecordBatchStream> {
+    let num_partitions = plan.output_partitioning().partition_count();
+    println!("==> num_partitions: {}", num_partitions);
     match plan.output_partitioning().partition_count() {
         0 => Ok(Box::pin(EmptyRecordBatchStream::new(plan.schema()))),
         1 => plan.execute(0, context),