first round of otpimizations

Author: mattcuento

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

@@ -80,20 +80,19 @@ impl ScratchSpace {
 
         let arrays = evaluate_expressions_to_arrays(exprs, batch)?;
 
-        self.hash_buffer.clear();
+        // Resize preserves capacity, only grows when needed
         self.hash_buffer.resize(num_rows, 0);
         create_hashes(
             &arrays,
             REPARTITION_RANDOM_STATE.random_state(),
             &mut self.hash_buffer,
         )?;
 
-        self.partition_ids.clear();
-        self.partition_ids.extend(
-            self.hash_buffer
-                .iter()
-                .map(|h| (*h % num_partitions as u64) as usize),
-        );
+        // Resize and populate partition IDs directly
+        self.partition_ids.resize(num_rows, 0);
+        for (i, hash) in self.hash_buffer.iter().enumerate() {
+            self.partition_ids[i] = (*hash % num_partitions as u64) as usize;
+        }
 
         self.map_partition_ids_to_starts_and_indices(num_partitions, num_rows);
         Ok(())
@@ -105,9 +104,10 @@ impl ScratchSpace {
         num_partitions: usize,
         num_rows: usize,
     ) {
-        self.partition_starts.clear();
+        // Truncate then resize for partition_starts (needs zeroing)
+        self.partition_starts.truncate(0);
         self.partition_starts.resize(num_partitions + 1, 0);
-        self.partition_row_indices.clear();
+        // Resize preserves capacity for partition_row_indices
         self.partition_row_indices.resize(num_rows, 0);
 
         // Count rows per partition
@@ -141,10 +141,13 @@ 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.
 pub fn materialize_partition(
     partition_indices: &[(u32, u32)],
     input_batches: &InputBatchStore,
     target_batch_size: usize,
+    scratch_indices: &mut Vec<u64>,
 ) -> Result<Vec<RecordBatch>> {
     if partition_indices.is_empty() {
         return Ok(Vec::new());
@@ -166,12 +169,16 @@ pub fn materialize_partition(
         }
 
         let batch = input_batches.get_batch(current_batch_idx);
-        let idx_array = UInt64Array::from(
+
+        // Reuse scratch buffer to reduce allocations
+        scratch_indices.clear();
+        scratch_indices.extend(
             partition_indices[start..end]
                 .iter()
-                .map(|(_, r)| *r as u64)
-                .collect::<Vec<_>>(),
+                .map(|(_, r)| *r as u64),
         );
+        let idx_array = UInt64Array::from(scratch_indices.clone());
+
         coalescer.push_batch_with_indices(batch.clone(), &idx_array)?;
         while let Some(completed) = coalescer.next_completed_batch() {
             result.push(completed);
@@ -315,7 +322,8 @@ mod tests {
     fn test_materialize_partition_empty() {
         let schema = create_test_schema();
         let store = InputBatchStore::new(schema);
-        let result = materialize_partition(&[], &store, 8192).unwrap();
+        let mut scratch = Vec::new();
+        let result = materialize_partition(&[], &store, 8192, &mut scratch).unwrap();
         assert!(result.is_empty());
     }
 
@@ -327,7 +335,8 @@ mod tests {
 
         // Select rows 0, 2, 4 from batch 0
         let indices = vec![(0u32, 0u32), (0, 2), (0, 4)];
-        let result = materialize_partition(&indices, &store, 8192).unwrap();
+        let mut scratch = Vec::new();
+        let result = materialize_partition(&indices, &store, 8192, &mut scratch).unwrap();
 
         let total_rows: usize = result.iter().map(|b| b.num_rows()).sum();
         assert_eq!(total_rows, 3);
@@ -352,7 +361,8 @@ 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 result = materialize_partition(&indices, &store, 8192).unwrap();
+        let mut scratch = Vec::new();
+        let result = materialize_partition(&indices, &store, 8192, &mut scratch).unwrap();
 
         let total_rows: usize = result.iter().map(|b| b.num_rows()).sum();
         assert_eq!(total_rows, 3);
@@ -371,7 +381,8 @@ mod tests {
         let indices: Vec<(u32, u32)> = (0..20000).map(|i| (0u32, i as u32)).collect();
 
         // Use target_batch_size of 8192
-        let result = materialize_partition(&indices, &store, 8192).unwrap();
+        let mut scratch = Vec::new();
+        let result = materialize_partition(&indices, &store, 8192, &mut scratch).unwrap();
 
         // Should produce multiple batches
         assert!(result.len() >= 2, "Expected multiple output batches");
@@ -438,7 +449,8 @@ mod tests {
         }
 
         // Use target_batch_size of 8192
-        let result = materialize_partition(&indices, &store, 8192).unwrap();
+        let mut scratch = Vec::new();
+        let result = materialize_partition(&indices, &store, 8192, &mut scratch).unwrap();
 
         // Should coalesce into far fewer output batches than input batches
         assert!(

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

@@ -211,8 +211,12 @@ impl SortShuffleWriterExec {
             };
 
             // Index-based partition tracking: store (batch_idx, row_idx) pairs per partition
-            let mut partition_indices: Vec<Vec<(u32, u32)>> =
-                (0..num_output_partitions).map(|_| Vec::new()).collect();
+            // Pre-allocate capacity to reduce reallocations during index insertion
+            let rows_per_partition_estimate =
+                (config.batch_size / num_output_partitions).max(64);
+            let mut partition_indices: Vec<Vec<(u32, u32)>> = (0..num_output_partitions)
+                .map(|_| Vec::with_capacity(rows_per_partition_estimate))
+                .collect();
             let mut input_store = InputBatchStore::new(schema.clone());
             let mut scratch = ScratchSpace::new(num_output_partitions);
 
@@ -244,8 +248,11 @@ impl SortShuffleWriterExec {
                 let batch_idx = input_store.push_batch(input_batch) as u32;
                 for (partition_id, indices) in partition_indices.iter_mut().enumerate() {
                     let row_indices = scratch.partition_indices(partition_id);
-                    for &row_idx in row_indices {
-                        indices.push((batch_idx, row_idx));
+                    // Skip empty partitions to reduce iterator overhead
+                    if !row_indices.is_empty() {
+                        for &row_idx in row_indices {
+                            indices.push((batch_idx, row_idx));
+                        }
                     }
                 }
 
@@ -343,11 +350,15 @@ fn spill_all_partitions(
     schema: &SchemaRef,
     batch_size: usize,
 ) -> Result<()> {
+    // Reusable scratch buffer for materialize_partition
+    let mut scratch_indices = Vec::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)?;
+        let batches =
+            materialize_partition(indices, input_store, batch_size, &mut scratch_indices)?;
         if !batches.is_empty() {
             spill_manager
                 .spill(partition_id, batches, schema)
@@ -403,6 +414,9 @@ 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();
+
     // Write partitions in order
     for (partition_id, indices) in partition_indices.iter().enumerate() {
         // Set the starting batch index for this partition
@@ -427,8 +441,12 @@ fn finalize_output(
         }
 
         // Then materialize remaining buffered data from indices
-        let materialized_batches =
-            materialize_partition(indices, input_store, config.batch_size)?;
+        let materialized_batches = materialize_partition(
+            indices,
+            input_store,
+            config.batch_size,
+            &mut scratch_indices,
+        )?;
         for batch in materialized_batches {
             partition_rows += batch.num_rows() as u64;
             partition_bytes += batch.get_array_memory_size() as u64;