Clean up in memory shuffles

Author: lukekim

ballista/core/src/config.rs

@@ -46,6 +46,9 @@ pub const BALLISTA_SHUFFLE_READER_REMOTE_PREFER_FLIGHT: &str =
     "ballista.shuffle.remote_read_prefer_flight";
 /// Configuration key for shuffle storage mode (disk or memory).
 pub const BALLISTA_SHUFFLE_MEMORY_MODE: &str = "ballista.shuffle.memory_mode";
+/// Configuration key indicating if this is the final output stage.
+/// When true, shuffle data is always written to disk regardless of memory_mode setting.
+pub const BALLISTA_IS_FINAL_STAGE: &str = "ballista.shuffle.is_final_stage";
 /// Shuffle format configuration: "arrow_ipc" or "vortex"
 pub const BALLISTA_SHUFFLE_FORMAT: &str = "ballista.shuffle.format";
 
@@ -94,6 +97,10 @@ static CONFIG_ENTRIES: LazyLock<HashMap<String, ConfigEntry>> = LazyLock::new(||
                          "When enabled, shuffle data is kept in memory on executors instead of being written to disk. This can improve performance for workloads with sufficient memory.".to_string(),
                          DataType::Boolean,
                          Some((false).to_string())),
+        ConfigEntry::new(BALLISTA_IS_FINAL_STAGE.to_string(),
+                         "When true, indicates this is the final output stage. Final stages always write to disk regardless of memory_mode setting to ensure proper cleanup.".to_string(),
+                         DataType::Boolean,
+                         Some((false).to_string())),
         ConfigEntry::new(BALLISTA_GRPC_CLIENT_CONNECT_TIMEOUT_SECONDS.to_string(),
                          "Connection timeout for gRPC client in seconds".to_string(),
                          DataType::UInt64,
@@ -322,6 +329,12 @@ impl BallistaConfig {
         self.get_bool_setting(BALLISTA_SHUFFLE_MEMORY_MODE)
     }
 
+    /// Returns whether this is the final output stage.
+    /// Final stages always write to disk regardless of memory_mode setting.
+    pub fn is_final_stage(&self) -> bool {
+        self.get_bool_setting(BALLISTA_IS_FINAL_STAGE)
+    }
+
     /// Returns the configured shuffle format (ArrowIpc or Vortex)
     ///
     /// Note: Vortex format requires the 'vortex' feature to be enabled.
@@ -492,4 +505,34 @@ mod tests {
         assert_eq!(16777216, config.default_grpc_client_max_message_size());
         Ok(())
     }
+
+    #[test]
+    fn test_is_final_stage_default() {
+        let config = BallistaConfig::default();
+        // Default should be false
+        assert!(!config.is_final_stage());
+    }
+
+    #[test]
+    fn test_shuffle_memory_mode_default() {
+        let config = BallistaConfig::default();
+        // Default should be false (disk-based shuffles)
+        assert!(!config.shuffle_memory_mode());
+    }
+
+    #[test]
+    fn test_shuffle_format_default() {
+        let config = BallistaConfig::default();
+        // Default should be ArrowIpc
+        assert_eq!(config.shuffle_format(), ShuffleFormat::ArrowIpc);
+    }
+
+    #[test]
+    fn test_shuffle_format_parsing() {
+        assert_eq!("arrow_ipc".parse::<ShuffleFormat>().unwrap(), ShuffleFormat::ArrowIpc);
+        assert_eq!("arrow-ipc".parse::<ShuffleFormat>().unwrap(), ShuffleFormat::ArrowIpc);
+        assert_eq!("ipc".parse::<ShuffleFormat>().unwrap(), ShuffleFormat::ArrowIpc);
+        assert_eq!("vortex".parse::<ShuffleFormat>().unwrap(), ShuffleFormat::Vortex);
+        assert!("invalid".parse::<ShuffleFormat>().is_err());
+    }
 }

ballista/core/src/execution_plans/shuffle_manager.rs

@@ -212,6 +212,31 @@ impl InMemoryShuffleManager {
         log::debug!("Removed all shuffle partitions for job: {job_id}");
     }
 
+    /// Removes all partitions for a given stage within a job.
+    ///
+    /// This is called when a stage's output has been fully consumed by the next stage,
+    /// allowing the memory to be reclaimed immediately rather than waiting for job completion.
+    ///
+    /// # Arguments
+    /// * `job_id` - The job identifier
+    /// * `stage_id` - The stage identifier
+    ///
+    /// # Returns
+    /// The number of partitions that were removed
+    pub fn remove_stage_partitions(&self, job_id: &str, stage_id: usize) -> usize {
+        let prefix = format!("{job_id}/{stage_id}/");
+        let initial_count = self.partitions.len();
+        self.partitions.retain(|k, _| !k.starts_with(&prefix));
+        let removed = initial_count - self.partitions.len();
+        log::debug!(
+            "Removed {} shuffle partitions for stage: {}/{}",
+            removed,
+            job_id,
+            stage_id
+        );
+        removed
+    }
+
     /// Returns the total number of partitions stored in memory.
     pub fn partition_count(&self) -> usize {
         self.partitions.len()
@@ -334,4 +359,133 @@ mod tests {
         let key = InMemoryShuffleManager::hash_partition_key("job1", 1, 2, 3);
         assert_eq!(key, "job1/1/2/data-3");
     }
+
+    #[test]
+    fn test_remove_stage_partitions() {
+        let manager = InMemoryShuffleManager::new();
+        let batch = create_test_batch();
+        let schema = batch.schema();
+
+        // Store partitions for multiple stages in the same job
+        for stage in 0..3 {
+            for partition in 0..4 {
+                let key =
+                    InMemoryShuffleManager::partition_key("job1", stage, partition);
+                let data =
+                    ShufflePartitionData::new(schema.clone(), vec![batch.clone()]);
+                manager.store_partition(key, data);
+            }
+        }
+
+        assert_eq!(manager.partition_count(), 12);
+
+        // Remove stage 1 partitions
+        let removed = manager.remove_stage_partitions("job1", 1);
+        assert_eq!(removed, 4);
+        assert_eq!(manager.partition_count(), 8);
+
+        // Verify stage 0 and 2 partitions still exist
+        let key0 = InMemoryShuffleManager::partition_key("job1", 0, 0);
+        let key2 = InMemoryShuffleManager::partition_key("job1", 2, 0);
+        assert!(manager.contains_partition(&key0));
+        assert!(manager.contains_partition(&key2));
+
+        // Verify stage 1 partitions are gone
+        let key1 = InMemoryShuffleManager::partition_key("job1", 1, 0);
+        assert!(!manager.contains_partition(&key1));
+    }
+
+    #[test]
+    fn test_remove_stage_partitions_different_jobs() {
+        let manager = InMemoryShuffleManager::new();
+        let batch = create_test_batch();
+        let schema = batch.schema();
+
+        // Store partitions for stage 1 in two different jobs
+        for job in ["job1", "job2"] {
+            for partition in 0..3 {
+                let key = InMemoryShuffleManager::partition_key(job, 1, partition);
+                let data =
+                    ShufflePartitionData::new(schema.clone(), vec![batch.clone()]);
+                manager.store_partition(key, data);
+            }
+        }
+
+        assert_eq!(manager.partition_count(), 6);
+
+        // Remove stage 1 from job1 only
+        let removed = manager.remove_stage_partitions("job1", 1);
+        assert_eq!(removed, 3);
+        assert_eq!(manager.partition_count(), 3);
+
+        // Verify job2 stage 1 partitions still exist
+        let key = InMemoryShuffleManager::partition_key("job2", 1, 0);
+        assert!(manager.contains_partition(&key));
+    }
+
+    #[test]
+    fn test_remove_partition_returns_data() {
+        let manager = InMemoryShuffleManager::new();
+        let batch = create_test_batch();
+        let schema = batch.schema();
+        let data = ShufflePartitionData::new(schema.clone(), vec![batch]);
+
+        let key = InMemoryShuffleManager::partition_key("job1", 1, 0);
+        manager.store_partition(key.clone(), data);
+
+        assert!(manager.contains_partition(&key));
+
+        // Remove should return the data
+        let removed = manager.remove_partition(&key);
+        assert!(removed.is_some());
+        let removed_data = removed.unwrap();
+        assert_eq!(removed_data.num_rows, 3);
+        assert_eq!(removed_data.num_batches, 1);
+
+        // Partition should no longer exist
+        assert!(!manager.contains_partition(&key));
+
+        // Second remove should return None
+        let removed_again = manager.remove_partition(&key);
+        assert!(removed_again.is_none());
+    }
+
+    #[test]
+    fn test_total_memory_usage() {
+        let manager = InMemoryShuffleManager::new();
+        let batch = create_test_batch();
+        let schema = batch.schema();
+
+        // Store multiple partitions
+        for i in 0..3 {
+            let key = InMemoryShuffleManager::partition_key("job1", 1, i);
+            let data = ShufflePartitionData::new(schema.clone(), vec![batch.clone()]);
+            manager.store_partition(key, data);
+        }
+
+        // Memory usage should be > 0
+        let usage = manager.total_memory_usage();
+        assert!(usage > 0);
+
+        // Remove partitions and verify usage decreases
+        manager.remove_job_partitions("job1");
+        assert_eq!(manager.total_memory_usage(), 0);
+    }
+
+    #[test]
+    fn test_clear() {
+        let manager = InMemoryShuffleManager::new();
+        let batch = create_test_batch();
+        let schema = batch.schema();
+
+        for i in 0..5 {
+            let key = InMemoryShuffleManager::partition_key("job1", 1, i);
+            let data = ShufflePartitionData::new(schema.clone(), vec![batch.clone()]);
+            manager.store_partition(key, data);
+        }
+
+        assert_eq!(manager.partition_count(), 5);
+        manager.clear();
+        assert_eq!(manager.partition_count(), 0);
+    }
 }

ballista/core/src/execution_plans/shuffle_reader.rs

@@ -746,6 +746,10 @@ async fn fetch_partition_object_store(
 }
 
 /// Fetch partition data from in-memory shuffle storage.
+///
+/// After successfully fetching the data, the partition is removed from memory
+/// to allow for immediate memory reclamation. This is safe because each shuffle
+/// partition is typically read only once by the consuming stage.
 async fn fetch_partition_memory(
     location: &PartitionLocation,
 ) -> result::Result<SendableRecordBatchStream, BallistaError> {
@@ -759,17 +763,23 @@ async fn fetch_partition_memory(
     })?;
 
     let shuffle_manager = global_shuffle_manager();
-    let data = shuffle_manager.get_partition(key).map_err(|e| {
-        BallistaError::FetchFailed(
-            metadata.id.clone(),
-            partition_id.stage_id,
-            partition_id.partition_id,
-            e.to_string(),
-        )
-    })?;
+
+    // Remove and retrieve the partition data in one atomic operation
+    // This ensures the memory is reclaimed as soon as the data is read
+    let data = shuffle_manager.remove_partition(key).ok_or_else(|| {
+        // If remove fails, try a regular get (for retry scenarios)
+        shuffle_manager.get_partition(key).map_err(|e| {
+            BallistaError::FetchFailed(
+                metadata.id.clone(),
+                partition_id.stage_id,
+                partition_id.partition_id,
+                e.to_string(),
+            )
+        })
+    }).or_else(|result| result)?;
 
     debug!(
-        "Fetched partition {} from memory: {} batches, {} rows",
+        "Fetched and removed partition {} from memory: {} batches, {} rows",
         key, data.num_batches, data.num_rows
     );
 

ballista/core/src/execution_plans/shuffle_writer.rs

@@ -309,8 +309,13 @@ impl ShuffleWriterExec {
         let job_id = self.job_id.clone();
         let stage_id = self.stage_id;
 
-        // Check if memory mode is enabled
+        // Check if memory mode is enabled and this is not the final stage
+        // Final stages always write to disk to ensure proper cleanup via existing mechanisms
         let memory_mode = context.session_config().ballista_shuffle_memory_mode();
+        let is_final_stage = context.session_config().ballista_is_final_stage();
+
+        // Use memory mode only for intermediate stages, not for the final output stage
+        let use_memory = memory_mode && !is_final_stage;
 
         // Get shuffle format from session config
         let shuffle_format = context.session_config().ballista_shuffle_format();
@@ -320,7 +325,7 @@ impl ShuffleWriterExec {
             let now = Instant::now();
             let mut stream = plan.execute(input_partition, context)?;
 
-            if memory_mode {
+            if use_memory {
                 // Use in-memory shuffle storage with configurable format
                 Self::execute_shuffle_write_memory(
                     &job_id,
@@ -335,6 +340,9 @@ impl ShuffleWriterExec {
                 .await
             } else {
                 // Use disk-based shuffle storage with configurable format
+                // This is used for:
+                // 1. When memory_mode is disabled
+                // 2. For final stages (even if memory_mode is enabled)
                 Self::execute_shuffle_write_disk(
                     path,
                     input_partition,