minor: split repartition time and send time metrics (#11440)

apache · Jul 16, 2024 · 2837e02 · 2837e02
1 parent f204869
commit 2837e02
Showing 1 changed file with 9 additions and 3 deletions.
diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs
@@ -261,6 +261,7 @@ impl BatchPartitioner {
                     num_partitions: partitions,
                     hash_buffer,
                 } => {
+                    // Tracking time required for distributing indexes across output partitions
                     let timer = self.timer.timer();
 
                     let arrays = exprs
@@ -282,6 +283,11 @@ impl BatchPartitioner {
                             .append_value(index as u64);
                     }
 
+                    // Finished building index-arrays for output partitions
+                    timer.done();
+
+                    // Borrowing partitioner timer to prevent moving `self` to closure
+                    let partitioner_timer = &self.timer;
                     let it = indices
                         .into_iter()
                         .enumerate()
@@ -290,6 +296,9 @@ impl BatchPartitioner {
                             (!indices.is_empty()).then_some((partition, indices))
                         })
                         .map(move |(partition, indices)| {
+                            // Tracking time required for repartitioned batches construction
+                            let _timer = partitioner_timer.timer();
+
                             // Produce batches based on indices
                             let columns = batch
                                 .columns()
@@ -303,9 +312,6 @@ impl BatchPartitioner {
                             let batch =
                                 RecordBatch::try_new(batch.schema(), columns).unwrap();
 
-                            // bind timer so it drops w/ this iterator
-                            let _ = &timer;
-
                             Ok((partition, batch))
                         });