add shuffle_write.output_io_time metric

Author: zhangli20

native-engine/datafusion-ext-plans/src/common/ipc_compression.rs

@@ -82,6 +82,10 @@ impl<W: Write> IpcCompressionWriter<W> {
     pub fn inner(&self) -> &W {
         &self.output
     }
+
+    pub fn inner_mut(&mut self) -> &mut W {
+        &mut self.output
+    }
 }
 
 pub struct IpcCompressionReader<R: Read + 'static> {

native-engine/datafusion-ext-plans/src/common/timer_helper.rs

@@ -14,6 +14,7 @@
 
 use std::{
     future::Future,
+    io::Write,
     sync::{
         atomic::{AtomicUsize, Ordering::Relaxed},
         Arc,
@@ -37,6 +38,8 @@ pub trait TimerHelper {
 
     fn duration(&self) -> Duration;
     fn sub_duration(&self, duration: Duration);
+
+    fn wrap_writer<W: Write>(&self, w: W) -> TimedWriter<W>;
 }
 
 impl TimerHelper for Time {
@@ -85,6 +88,22 @@ impl TimerHelper for Time {
         };
         xtime.nanos.fetch_sub(duration.as_nanos() as usize, Relaxed);
     }
+
+    fn wrap_writer<W: Write>(&self, w: W) -> TimedWriter<W> {
+        TimedWriter(w, self.clone())
+    }
+}
+
+pub struct TimedWriter<W: Write>(pub W, Time);
+
+impl<W: Write> Write for TimedWriter<W> {
+    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
+        self.1.with_timer(|| self.0.write(buf))
+    }
+
+    fn flush(&mut self) -> std::io::Result<()> {
+        self.1.with_timer(|| self.0.flush())
+    }
 }
 
 pub trait RegisterTimer {

native-engine/datafusion-ext-plans/src/shuffle/single_repartitioner.rs

@@ -20,43 +20,57 @@ use std::{
 
 use arrow::record_batch::RecordBatch;
 use async_trait::async_trait;
-use datafusion::{common::Result, physical_plan::metrics::BaselineMetrics};
+use datafusion::{
+    common::Result,
+    physical_plan::metrics::{BaselineMetrics, Time},
+};
 use tokio::sync::Mutex;
 
-use crate::{common::ipc_compression::IpcCompressionWriter, shuffle::ShuffleRepartitioner};
+use crate::{
+    common::{
+        ipc_compression::IpcCompressionWriter,
+        timer_helper::{TimedWriter, TimerHelper},
+    },
+    shuffle::ShuffleRepartitioner,
+};
 
 pub struct SingleShuffleRepartitioner {
     output_data_file: String,
     output_index_file: String,
-    output_data: Arc<Mutex<Option<IpcCompressionWriter<File>>>>,
+    output_data: Arc<Mutex<Option<IpcCompressionWriter<TimedWriter<File>>>>>,
+    output_io_time: Time,
     metrics: BaselineMetrics,
 }
 
 impl SingleShuffleRepartitioner {
     pub fn new(
         output_data_file: String,
         output_index_file: String,
+        output_io_time: Time,
         metrics: BaselineMetrics,
     ) -> Self {
         Self {
             output_data_file,
             output_index_file,
             output_data: Arc::new(Mutex::default()),
+            output_io_time,
             metrics,
         }
     }
 
     fn get_output_writer<'a>(
         &self,
-        output_data: &'a mut Option<IpcCompressionWriter<File>>,
-    ) -> Result<&'a mut IpcCompressionWriter<File>> {
+        output_data: &'a mut Option<IpcCompressionWriter<TimedWriter<File>>>,
+    ) -> Result<&'a mut IpcCompressionWriter<TimedWriter<File>>> {
         if output_data.is_none() {
             *output_data = Some(IpcCompressionWriter::new(
-                OpenOptions::new()
-                    .write(true)
-                    .create(true)
-                    .truncate(true)
-                    .open(&self.output_data_file)?,
+                self.output_io_time.wrap_writer(
+                    OpenOptions::new()
+                        .write(true)
+                        .create(true)
+                        .truncate(true)
+                        .open(&self.output_data_file)?,
+                ),
             ));
         }
         Ok(output_data.as_mut().unwrap())
@@ -79,16 +93,33 @@ impl ShuffleRepartitioner for SingleShuffleRepartitioner {
 
         // write index file
         if let Some(output_writer) = output_data.as_mut() {
+            let mut output_index = self.output_io_time.wrap_writer(
+                OpenOptions::new()
+                    .write(true)
+                    .create(true)
+                    .truncate(true)
+                    .open(&self.output_index_file)?,
+            );
             output_writer.finish_current_buf()?;
-            let offset = output_writer.inner().stream_position()?;
-            let mut output_index = File::create(&self.output_index_file)?;
+            let offset = output_writer.inner_mut().0.stream_position()?;
             output_index.write_all(&[0u8; 8])?;
             output_index.write_all(&(offset as i64).to_le_bytes()[..])?;
         } else {
             // write empty data file and index file
-            let output_data = File::create(&self.output_data_file)?;
-            output_data.set_len(0)?;
-            let mut output_index = File::create(&self.output_index_file)?;
+            let _output_data = self.output_io_time.wrap_writer(
+                OpenOptions::new()
+                    .write(true)
+                    .create(true)
+                    .truncate(true)
+                    .open(&self.output_data_file)?,
+            );
+            let mut output_index = self.output_io_time.wrap_writer(
+                OpenOptions::new()
+                    .write(true)
+                    .create(true)
+                    .truncate(true)
+                    .open(&self.output_index_file)?,
+            );
             output_index.write_all(&[0u8; 16])?;
         }
         Ok(())

native-engine/datafusion-ext-plans/src/shuffle/sort_repartitioner.rs

@@ -13,7 +13,7 @@
 // limitations under the License.
 
 use std::{
-    fs::{File, OpenOptions},
+    fs::OpenOptions,
     io::{BufReader, Read, Seek, Write},
     sync::Weak,
 };
@@ -23,7 +23,10 @@ use async_trait::async_trait;
 use bytesize::ByteSize;
 use datafusion::{
     common::{DataFusionError, Result},
-    physical_plan::{metrics::ExecutionPlanMetricsSet, Partitioning},
+    physical_plan::{
+        metrics::{ExecutionPlanMetricsSet, Time},
+        Partitioning,
+    },
 };
 use datafusion_ext_commons::{
     df_execution_err,
@@ -32,6 +35,7 @@ use datafusion_ext_commons::{
 use futures::lock::Mutex;
 
 use crate::{
+    common::timer_helper::TimerHelper,
     memmgr::{
         metrics::SpillMetrics,
         spill::{try_new_spill, Spill},
@@ -49,6 +53,7 @@ pub struct SortShuffleRepartitioner {
     spills: Mutex<Vec<ShuffleSpill>>,
     partitioning: Partitioning,
     num_output_partitions: usize,
+    output_io_time: Time,
     spill_metrics: SpillMetrics,
 }
 
@@ -58,6 +63,7 @@ impl SortShuffleRepartitioner {
         output_data_file: String,
         output_index_file: String,
         partitioning: Partitioning,
+        output_io_time: Time,
         metrics: &ExecutionPlanMetricsSet,
     ) -> Self {
         let num_output_partitions = partitioning.partition_count();
@@ -70,6 +76,7 @@ impl SortShuffleRepartitioner {
             spills: Mutex::default(),
             partitioning,
             num_output_partitions,
+            output_io_time,
             spill_metrics: SpillMetrics::new(metrics, partition_id),
         }
     }
@@ -158,22 +165,33 @@ impl ShuffleRepartitioner for SortShuffleRepartitioner {
         // no spills - directly write current batches into final file
         if spills.is_empty() {
             let partitioning = self.partitioning.clone();
+            let output_io_time = self.output_io_time.clone();
             tokio::task::spawn_blocking(move || {
-                let mut output_data = OpenOptions::new()
-                    .write(true)
-                    .create(true)
-                    .truncate(true)
-                    .open(&data_file)?;
-                let mut output_index = OpenOptions::new()
-                    .write(true)
-                    .create(true)
-                    .truncate(true)
-                    .open(&index_file)?;
+                let mut output_data = output_io_time.wrap_writer(
+                    OpenOptions::new()
+                        .write(true)
+                        .create(true)
+                        .truncate(true)
+                        .open(&data_file)?,
+                );
+                let mut output_index = output_io_time.wrap_writer(
+                    OpenOptions::new()
+                        .write(true)
+                        .create(true)
+                        .truncate(true)
+                        .open(&index_file)?,
+                );
 
+                // write data file
                 let offsets = data.write(&mut output_data, &partitioning)?;
+
+                // write index file
+                let mut offsets_data = vec![];
                 for offset in offsets {
-                    output_index.write_all(&(offset as i64).to_le_bytes()[..])?;
+                    offsets_data.extend_from_slice(&(offset as i64).to_le_bytes()[..]);
                 }
+                output_index.write_all(&offsets_data)?;
+
                 Ok::<(), DataFusionError>(())
             })
             .await
@@ -216,13 +234,22 @@ impl ShuffleRepartitioner for SortShuffleRepartitioner {
         let mut offsets = vec![0];
 
         // append partition in each spills
+        let output_io_time = self.output_io_time.clone();
         tokio::task::spawn_blocking(move || {
-            let mut output_data = OpenOptions::new()
-                .write(true)
-                .create(true)
-                .truncate(true)
-                .open(data_file)?;
-            let mut cur_partition_id = 0;
+            let mut output_data = output_io_time.wrap_writer(
+                OpenOptions::new()
+                    .write(true)
+                    .create(true)
+                    .truncate(true)
+                    .open(&data_file)?,
+            );
+            let mut output_index = output_io_time.wrap_writer(
+                OpenOptions::new()
+                    .write(true)
+                    .create(true)
+                    .truncate(true)
+                    .open(&index_file)?,
+            );
 
             if !spills.is_empty() {
                 // select partitions from spills
@@ -243,14 +270,15 @@ impl ShuffleRepartitioner for SortShuffleRepartitioner {
                     num_output_partitions,
                 );
 
+                let mut cur_partition_id = 0;
                 loop {
                     let mut min_spill = cursors.peek_mut();
                     if min_spill.cur + 1 >= min_spill.offsets.len() {
                         break;
                     }
 
                     while cur_partition_id < min_spill.cur {
-                        offsets.push(output_data.stream_position()?);
+                        offsets.push(output_data.0.stream_position()?);
                         cur_partition_id += 1;
                     }
                     let (spill_offset_start, spill_offset_end) = (
@@ -269,12 +297,15 @@ impl ShuffleRepartitioner for SortShuffleRepartitioner {
             }
 
             // add one extra offset at last to ease partition length computation
-            offsets.resize(num_output_partitions + 1, output_data.stream_position()?);
+            offsets.resize(num_output_partitions + 1, output_data.0.stream_position()?);
 
-            let mut output_index = File::create(index_file)?;
+            // write index file
+            let mut offsets_data = vec![];
             for offset in offsets {
-                output_index.write_all(&(offset as i64).to_le_bytes()[..])?;
+                offsets_data.extend_from_slice(&(offset as i64).to_le_bytes()[..]);
             }
+            output_index.write_all(&offsets_data)?;
+
             Ok::<(), DataFusionError>(())
         })
         .await

native-engine/datafusion-ext-plans/src/shuffle_writer_exec.rs

@@ -33,7 +33,10 @@ use datafusion_ext_commons::df_execution_err;
 use futures::{stream::once, TryStreamExt};
 
 use crate::{
-    common::batch_statisitcs::{stat_input, InputBatchStatistics},
+    common::{
+        batch_statisitcs::{stat_input, InputBatchStatistics},
+        timer_helper::RegisterTimer,
+    },
     memmgr::MemManager,
     shuffle::{
         single_repartitioner::SingleShuffleRepartitioner,
@@ -110,11 +113,13 @@ impl ExecutionPlan for ShuffleWriterExec {
     ) -> Result<SendableRecordBatchStream> {
         // record uncompressed data size
         let data_size_metric = MetricBuilder::new(&self.metrics).counter("data_size", partition);
+        let output_time = self.metrics.register_timer("output_io_time", partition);
 
         let repartitioner: Arc<dyn ShuffleRepartitioner> = match &self.partitioning {
             p if p.partition_count() == 1 => Arc::new(SingleShuffleRepartitioner::new(
                 self.output_data_file.clone(),
                 self.output_index_file.clone(),
+                output_time,
                 BaselineMetrics::new(&self.metrics, partition),
             )),
             Partitioning::Hash(..) => {
@@ -123,6 +128,7 @@ impl ExecutionPlan for ShuffleWriterExec {
                     self.output_data_file.clone(),
                     self.output_index_file.clone(),
                     self.partitioning.clone(),
+                    output_time,
                     &self.metrics,
                 ));
                 MemManager::register_consumer(partitioner.clone(), true);

spark-extension-shims-spark3/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeShuffleExchangeExec.scala

@@ -57,7 +57,8 @@ case class NativeShuffleExchangeExec(
             "mem_spill_size",
             "mem_spill_iotime",
             "disk_spill_size",
-            "disk_spill_iotime"))
+            "disk_spill_iotime",
+            "output_io_time"))
         .toSeq: _*)).toMap
 
   lazy val readMetrics: Map[String, SQLMetric] =

spark-extension/src/main/scala/org/apache/spark/sql/blaze/NativeHelper.scala

@@ -108,7 +108,8 @@ object NativeHelper extends Logging {
       "mem_spill_size" -> sizeMetric("Native.mem_spill_size"),
       "mem_spill_iotime" -> nanoTimingMetric("Native.mem_spill_iotime"),
       "disk_spill_size" -> sizeMetric("Native.disk_spill_size"),
-      "disk_spill_iotime" -> nanoTimingMetric("Native.disk_spill_iotime"))
+      "disk_spill_iotime" -> nanoTimingMetric("Native.disk_spill_iotime"),
+      "output_io_time" -> nanoTimingMetric("Native.output_io_time"))
 
     if (BlazeConf.INPUT_BATCH_STATISTICS_ENABLE.booleanConf()) {
       metrics ++= TreeMap(

spark-extension/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeShuffleExchangeBase.scala

@@ -167,22 +167,15 @@ abstract class NativeShuffleExchangeBase(
     val nativeInputRDD = rdd.asInstanceOf[NativeRDD]
     val numPartitions = outputPartitioning.numPartitions
     val nativeMetrics = MetricNode(
-      Map(),
+      metrics,
       nativeInputRDD.metrics :: Nil,
       Some({
-        case ("stage_id", v) => metrics("stage_id") += v
-        case ("data_size", v) => metrics("dataSize") += v
         case ("output_rows", v) =>
           val shuffleWriteMetrics = TaskContext.get.taskMetrics().shuffleWriteMetrics
           new SQLShuffleWriteMetricsReporter(shuffleWriteMetrics, metrics).incRecordsWritten(v)
         case ("elapsed_compute", v) =>
           val shuffleWriteMetrics = TaskContext.get.taskMetrics().shuffleWriteMetrics
           new SQLShuffleWriteMetricsReporter(shuffleWriteMetrics, metrics).incWriteTime(v)
-        case ("mem_spill_count", v) if v > 0 => metrics("mem_spill_count").add(v)
-        case ("mem_spill_size", v) if v > 0 => metrics("mem_spill_size").add(v)
-        case ("mem_spill_iotime", v) if v > 0 => metrics("mem_spill_iotime").add(v)
-        case ("disk_spill_size", v) if v > 0 => metrics("disk_spill_size").add(v)
-        case ("disk_spill_iotime", v) if v > 0 => metrics("disk_spill_iotime").add(v)
         case _ =>
       }))
     val nativeHashExprs = this.nativeHashExprs