feat: add scatter_time metric to shuffle write profiling

Author: andygrove

native/core/src/execution/shuffle/metrics.rs

@@ -26,6 +26,9 @@ pub(super) struct ShufflePartitionerMetrics {
     /// Time to perform repartitioning
     pub(super) repart_time: Time,
 
+    /// Time scattering values to per-partition buffers
+    pub(super) scatter_time: Time,
+
     /// Time encoding batches to IPC format
     pub(super) encode_time: Time,
 
@@ -50,6 +53,7 @@ impl ShufflePartitionerMetrics {
         Self {
             baseline: BaselineMetrics::new(metrics, partition),
             repart_time: MetricBuilder::new(metrics).subset_time("repart_time", partition),
+            scatter_time: MetricBuilder::new(metrics).subset_time("scatter_time", partition),
             encode_time: MetricBuilder::new(metrics).subset_time("encode_time", partition),
             write_time: MetricBuilder::new(metrics).subset_time("write_time", partition),
             input_batches: MetricBuilder::new(metrics).counter("input_batches", partition),

native/core/src/execution/shuffle/partitioners/multi_partition.rs

@@ -415,6 +415,8 @@ impl MultiPartitionShuffleRepartitioner {
         // Track memory before scatter
         let mem_before: usize = self.partition_buffers.iter().map(|b| b.memory_size()).sum();
 
+        let scatter_start = Instant::now();
+
         // Column-oriented scatter: for each column, iterate by partition then by
         // rows within that partition. This keeps writes to the same partition buffer
         // sequential for better cache locality.
@@ -535,6 +537,8 @@ impl MultiPartitionShuffleRepartitioner {
             }
         }
 
+        self.metrics.scatter_time.add_duration(scatter_start.elapsed());
+
         // Update row counts from partition_starts (O(num_partitions), not O(num_rows))
         for p in 0..num_partitions {
             let count = (partition_starts[p + 1] - partition_starts[p]) as usize;

spark/src/main/scala/org/apache/spark/sql/comet/CometMetricNode.scala

@@ -248,6 +248,9 @@ object CometMetricNode {
     Map(
       "elapsed_compute" -> SQLMetrics.createNanoTimingMetric(sc, "native shuffle writer time"),
       "repart_time" -> SQLMetrics.createNanoTimingMetric(sc, "repartition time"),
+      "scatter_time" -> SQLMetrics.createNanoTimingMetric(
+        sc,
+        "scatter to partition buffers time"),
       "encode_time" -> SQLMetrics.createNanoTimingMetric(sc, "encoding and compression time"),
       "decode_time" -> SQLMetrics.createNanoTimingMetric(sc, "decoding and decompression time"),
       "spill_count" -> SQLMetrics.createMetric(sc, "number of spills"),

spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala

@@ -78,6 +78,7 @@ class CometNativeShuffleWriter[K, V](
     val detailedMetrics = Seq(
       "elapsed_compute",
       "repart_time",
+      "scatter_time",
       "encode_time",
       "input_batches",
       "spill_count",