Merge remote-tracking branch 'origin/spiceai-51' into lukim/object-store

Author: lukekim

Cargo.lock

@@ -50,6 +50,7 @@ aws-credential-types = { version = "1.2.0", optional = true }
 bytes = { workspace = true }
 chrono = { version = "0.4", default-features = false }
 clap = { workspace = true, optional = true }
+dashmap = "6"
 datafusion = { workspace = true }
 datafusion-proto = { workspace = true }
 datafusion-proto-common = { workspace = true }

ballista/core/Cargo.toml

@@ -48,6 +48,11 @@ pub const BALLISTA_SHUFFLE_READER_REMOTE_PREFER_FLIGHT: &str =
 pub const BALLISTA_SHUFFLE_STORAGE_TYPE: &str = "ballista.shuffle.storage_type";
 /// Configuration key for shuffle storage base URL/path.
 pub const BALLISTA_SHUFFLE_STORAGE_URL: &str = "ballista.shuffle.storage_url";
+/// 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";
 
@@ -100,6 +105,14 @@ static CONFIG_ENTRIES: LazyLock<HashMap<String, ConfigEntry>> = LazyLock::new(||
                          "Base URL/path for shuffle storage. For local: file path; For S3: s3://bucket/prefix; For Azure: abfs://container@account.dfs.core.windows.net/prefix".to_string(),
                          DataType::Utf8,
                          None),
+        ConfigEntry::new(BALLISTA_SHUFFLE_MEMORY_MODE.to_string(),
+                         "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,
@@ -329,6 +342,21 @@ impl BallistaConfig {
         self.settings.get(BALLISTA_SHUFFLE_STORAGE_URL).cloned()
     }
 
+    /// Returns whether in-memory shuffle mode is enabled.
+    ///
+    /// 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.
+    pub fn shuffle_memory_mode(&self) -> bool {
+        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.
@@ -499,4 +527,46 @@ 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/config.rs

@@ -19,6 +19,7 @@
 //! several Ballista executors.
 
 mod distributed_query;
+mod shuffle_manager;
 mod shuffle_reader;
 mod shuffle_writer;
 mod unresolved_shuffle;
@@ -27,6 +28,10 @@ mod unresolved_shuffle;
 pub mod vortex_shuffle;
 
 pub use distributed_query::DistributedQueryExec;
+pub use shuffle_manager::{
+    InMemoryShuffleManager, ShufflePartitionData, ShufflePartitionKey,
+    global_shuffle_manager,
+};
 pub use shuffle_reader::ShuffleReaderExec;
 pub use shuffle_writer::ShuffleWriterExec;
 pub use unresolved_shuffle::UnresolvedShuffleExec;