perf: optimize shuffle writer with buffered I/O and fix file size bug

1. Add BufWriter for buffered file I/O in shuffle writer

   Wrapping File with BufWriter reduces syscalls when writing
   multiple small batches to shuffle files. This affects both
   the hash-partitioned shuffle path and the non-partitioned
   write_stream_to_disk utility.

2. Fix file size read before writer finish

   Previously, fs::metadata() was called before writer.finish(),
   which could report incorrect file sizes since data may not
   have been fully flushed to disk. This is especially important
   now that BufWriter is used, as buffered data would not be
   reflected in the file size until after finish() flushes it.

Co-Authored-By: Claude Opus 4.5 <[email protected]>

Author: andygrove

ballista/core/src/execution_plans/shuffle_writer.rs

@@ -29,6 +29,7 @@ use std::fmt::Debug;
 use std::fs;
 use std::fs::File;
 use std::future::Future;
+use std::io::BufWriter;
 use std::iter::Iterator;
 use std::path::PathBuf;
 use std::sync::Arc;
@@ -105,7 +106,7 @@ impl std::fmt::Display for ShuffleWriterExec {
 pub struct WriteTracker {
     pub num_batches: usize,
     pub num_rows: usize,
-    pub writer: StreamWriter<File>,
+    pub writer: StreamWriter<BufWriter<File>>,
     pub path: PathBuf,
 }
 
@@ -295,7 +296,8 @@ impl ShuffleWriterExec {
                                                 CompressionType::LZ4_FRAME,
                                             ))?;
 
-                                        let file = File::create(path.clone())?;
+                                        let file =
+                                            BufWriter::new(File::create(path.clone())?);
                                         let mut writer =
                                             StreamWriter::try_new_with_options(
                                                 file,
@@ -323,8 +325,8 @@ impl ShuffleWriterExec {
 
                     for (i, w) in writers.iter_mut().enumerate() {
                         if let Some(w) = w {
-                            let num_bytes = fs::metadata(&w.path)?.len();
                             w.writer.finish()?;
+                            let num_bytes = fs::metadata(&w.path)?.len();
                             debug!(
                                 "Finished writing shuffle partition {} at {:?}. Batches: {}. Rows: {}. Bytes: {}.",
                                 i, w.path, w.num_batches, w.num_rows, num_bytes

ballista/core/src/utils.rs

@@ -31,6 +31,7 @@ use datafusion::physical_plan::metrics::MetricsSet;
 use datafusion::physical_plan::{ExecutionPlan, RecordBatchStream, metrics};
 use futures::StreamExt;
 use log::error;
+use std::io::BufWriter;
 use std::sync::Arc;
 use std::time::{Duration, SystemTime, UNIX_EPOCH};
 use std::{fs::File, pin::Pin};
@@ -148,10 +149,10 @@ pub async fn write_stream_to_disk(
     path: &str,
     disk_write_metric: &metrics::Time,
 ) -> Result<PartitionStats> {
-    let file = File::create(path).map_err(|e| {
+    let file = BufWriter::new(File::create(path).map_err(|e| {
         error!("Failed to create partition file at {path}: {e:?}");
         BallistaError::IoError(e)
-    })?;
+    })?);
 
     let mut num_rows = 0;
     let mut num_batches = 0;