second round

Author: mattcuento

ballista/core/src/execution_plans/sort_shuffle/buffer.rs

@@ -26,7 +26,7 @@
 
 use std::sync::Arc;
 
-use datafusion::arrow::array::UInt64Array;
+use datafusion::arrow::array::UInt64Builder;
 use datafusion::arrow::compute::BatchCoalescer;
 use datafusion::arrow::datatypes::SchemaRef;
 use datafusion::arrow::record_batch::RecordBatch;
@@ -142,12 +142,12 @@ impl ScratchSpace {
 /// Materializes a partition's data from `(batch_idx, row_idx)` pairs into
 /// coalesced `RecordBatch`es using `BatchCoalescer::push_batch_with_indices`.
 ///
-/// Uses `scratch_indices` as a reusable buffer to avoid allocations.
+/// Uses `scratch_builder` as a reusable builder to avoid allocations.
 pub fn materialize_partition(
     partition_indices: &[(u32, u32)],
     input_batches: &InputBatchStore,
     target_batch_size: usize,
-    scratch_indices: &mut Vec<u64>,
+    scratch_builder: &mut UInt64Builder,
 ) -> Result<Vec<RecordBatch>> {
     if partition_indices.is_empty() {
         return Ok(Vec::new());
@@ -170,14 +170,11 @@ pub fn materialize_partition(
 
         let batch = input_batches.get_batch(current_batch_idx);
 
-        // Reuse scratch buffer to reduce allocations
-        scratch_indices.clear();
-        scratch_indices.extend(
-            partition_indices[start..end]
-                .iter()
-                .map(|(_, r)| *r as u64),
-        );
-        let idx_array = UInt64Array::from(scratch_indices.clone());
+        // Use builder pattern to avoid Vec allocation and clone
+        for (_, r) in &partition_indices[start..end] {
+            scratch_builder.append_value(*r as u64);
+        }
+        let idx_array = scratch_builder.finish();
 
         coalescer.push_batch_with_indices(batch.clone(), &idx_array)?;
         while let Some(completed) = coalescer.next_completed_batch() {
@@ -322,7 +319,7 @@ mod tests {
     fn test_materialize_partition_empty() {
         let schema = create_test_schema();
         let store = InputBatchStore::new(schema);
-        let mut scratch = Vec::new();
+        let mut scratch = UInt64Builder::new();
         let result = materialize_partition(&[], &store, 8192, &mut scratch).unwrap();
         assert!(result.is_empty());
     }
@@ -335,7 +332,7 @@ mod tests {
 
         // Select rows 0, 2, 4 from batch 0
         let indices = vec![(0u32, 0u32), (0, 2), (0, 4)];
-        let mut scratch = Vec::new();
+        let mut scratch = UInt64Builder::new();
         let result = materialize_partition(&indices, &store, 8192, &mut scratch).unwrap();
 
         let total_rows: usize = result.iter().map(|b| b.num_rows()).sum();
@@ -361,7 +358,7 @@ mod tests {
 
         // Select row 1 from batch 0, rows 0 and 2 from batch 1
         let indices = vec![(0u32, 1u32), (1, 0), (1, 2)];
-        let mut scratch = Vec::new();
+        let mut scratch = UInt64Builder::new();
         let result = materialize_partition(&indices, &store, 8192, &mut scratch).unwrap();
 
         let total_rows: usize = result.iter().map(|b| b.num_rows()).sum();
@@ -381,7 +378,7 @@ mod tests {
         let indices: Vec<(u32, u32)> = (0..20000).map(|i| (0u32, i as u32)).collect();
 
         // Use target_batch_size of 8192
-        let mut scratch = Vec::new();
+        let mut scratch = UInt64Builder::new();
         let result = materialize_partition(&indices, &store, 8192, &mut scratch).unwrap();
 
         // Should produce multiple batches
@@ -449,7 +446,7 @@ mod tests {
         }
 
         // Use target_batch_size of 8192
-        let mut scratch = Vec::new();
+        let mut scratch = UInt64Builder::new();
         let result = materialize_partition(&indices, &store, 8192, &mut scratch).unwrap();
 
         // Should coalesce into far fewer output batches than input batches

ballista/core/src/execution_plans/sort_shuffle/writer.rs

@@ -350,15 +350,15 @@ fn spill_all_partitions(
     schema: &SchemaRef,
     batch_size: usize,
 ) -> Result<()> {
-    // Reusable scratch buffer for materialize_partition
-    let mut scratch_indices = Vec::new();
+    // Reusable scratch builder for materialize_partition
+    let mut scratch_builder = UInt64Builder::new();
 
     for (partition_id, indices) in partition_indices.iter_mut().enumerate() {
         if indices.is_empty() {
             continue;
         }
         let batches =
-            materialize_partition(indices, input_store, batch_size, &mut scratch_indices)?;
+            materialize_partition(indices, input_store, batch_size, &mut scratch_builder)?;
         if !batches.is_empty() {
             spill_manager
                 .spill(partition_id, batches, schema)
@@ -414,8 +414,8 @@ fn finalize_output(
     // FileReader supports random access to batches by index
     let mut cumulative_batch_count: i64 = 0;
 
-    // Reusable scratch buffer for materialize_partition
-    let mut scratch_indices = Vec::new();
+    // Reusable scratch builder for materialize_partition
+    let mut scratch_builder = UInt64Builder::new();
 
     // Write partitions in order
     for (partition_id, indices) in partition_indices.iter().enumerate() {
@@ -445,7 +445,7 @@ fn finalize_output(
             indices,
             input_store,
             config.batch_size,
-            &mut scratch_indices,
+            &mut scratch_builder,
         )?;
         for batch in materialized_batches {
             partition_rows += batch.num_rows() as u64;

benchmarks/src/bin/shuffle_bench.rs

@@ -333,6 +333,14 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
         let mut hash_file_count = 0;
         let mut hash_total_size = 0u64;
 
+        // Warmup iteration (discard result to avoid JIT/allocator initialization overhead)
+        {
+            let temp_dir = TempDir::new()?;
+            let work_dir = temp_dir.path().to_str().unwrap();
+            let _ = benchmark_hash_shuffle(&data, schema.clone(), opt.partitions, work_dir).await?;
+            println!("  Warmup iteration completed (not timed)");
+        }
+
         for i in 0..opt.iterations {
             let temp_dir = TempDir::new()?;
             let work_dir = temp_dir.path().to_str().unwrap();
@@ -380,6 +388,22 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
         let mut sort_file_count = 0;
         let mut sort_total_size = 0u64;
 
+        // Warmup iteration (discard result to avoid JIT/allocator initialization overhead)
+        {
+            let temp_dir = TempDir::new()?;
+            let work_dir = temp_dir.path().to_str().unwrap();
+            let _ = benchmark_sort_shuffle(
+                &data,
+                schema.clone(),
+                opt.partitions,
+                work_dir,
+                buffer_size,
+                memory_limit,
+            )
+            .await?;
+            println!("  Warmup iteration completed (not timed)");
+        }
+
         for i in 0..opt.iterations {
             let temp_dir = TempDir::new()?;
             let work_dir = temp_dir.path().to_str().unwrap();