Update spilled_bytes to report spill file size (#16535)

Author: ding-young

datafusion/execution/src/disk_manager.rs

@@ -366,6 +366,10 @@ impl RefCountedTempFile {
 
         Ok(())
     }
+
+    pub fn current_disk_usage(&self) -> u64 {
+        self.current_file_disk_usage
+    }
 }
 
 /// When the temporary file is dropped, subtract its disk usage from the disk manager's total

datafusion/physical-plan/src/metrics/baseline.rs

@@ -151,7 +151,7 @@ pub struct SpillMetrics {
     /// count of spills during the execution of the operator
     pub spill_file_count: Count,
 
-    /// total spilled bytes during the execution of the operator
+    /// total bytes actually written to disk during the execution of the operator
     pub spilled_bytes: Count,
 
     /// total spilled rows during the execution of the operator

datafusion/physical-plan/src/sorts/sort.rs

@@ -1514,7 +1514,7 @@ mod tests {
         // bytes. We leave a little wiggle room for the actual numbers.
         assert!((3..=10).contains(&spill_count));
         assert!((9000..=10000).contains(&spilled_rows));
-        assert!((38000..=42000).contains(&spilled_bytes));
+        assert!((38000..=44000).contains(&spilled_bytes));
 
         let columns = result[0].columns();
 

datafusion/physical-plan/src/spill/in_progress_spill_file.rs

@@ -75,15 +75,14 @@ impl InProgressSpillFile {
             }
         }
         if let Some(writer) = &mut self.writer {
-            let (spilled_rows, spilled_bytes) = writer.write(batch)?;
+            let (spilled_rows, _) = writer.write(batch)?;
             if let Some(in_progress_file) = &mut self.in_progress_file {
                 in_progress_file.update_disk_usage()?;
             } else {
                 unreachable!() // Already checked inside current function
             }
 
             // Update metrics
-            self.spill_writer.metrics.spilled_bytes.add(spilled_bytes);
             self.spill_writer.metrics.spilled_rows.add(spilled_rows);
         }
         Ok(())
@@ -98,6 +97,14 @@ impl InProgressSpillFile {
             return Ok(None);
         }
 
+        // Since spill files are append-only, add the file size to spilled_bytes
+        if let Some(in_progress_file) = &mut self.in_progress_file {
+            // Since writer.finish() writes continuation marker and message length at the end
+            in_progress_file.update_disk_usage()?;
+            let size = in_progress_file.current_disk_usage();
+            self.spill_writer.metrics.spilled_bytes.add(size as usize);
+        }
+
         Ok(self.in_progress_file.take())
     }
 }

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

@@ -809,12 +809,13 @@ mod tests {
                 Arc::new(StringArray::from(vec!["d", "e", "f"])),
             ],
         )?;
-
+        // After appending each batch, spilled_rows should increase, while spill_file_count and
+        // spilled_bytes remain the same (spilled_bytes is updated only after finish() is called)
         in_progress_file.append_batch(&batch1)?;
-        verify_metrics(&in_progress_file, 1, 356, 3)?;
+        verify_metrics(&in_progress_file, 1, 0, 3)?;
 
         in_progress_file.append_batch(&batch2)?;
-        verify_metrics(&in_progress_file, 1, 712, 6)?;
+        verify_metrics(&in_progress_file, 1, 0, 6)?;
 
         let completed_file = in_progress_file.finish()?;
         assert!(completed_file.is_some());