Don't expose final stage

Author: lukekim

ballista/core/src/config.rs

@@ -50,7 +50,8 @@ pub const BALLISTA_SHUFFLE_STORAGE_TYPE: &str = "ballista.shuffle.storage_type";
 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.
+/// Internal configuration key indicating if this is the final output stage.
+/// This is set by the scheduler based on stage topology, NOT user-configurable.
 /// 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"
@@ -109,10 +110,9 @@ 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())),
+        // Note: BALLISTA_IS_FINAL_STAGE is intentionally NOT in CONFIG_ENTRIES.
+        // It's an internal flag set by the scheduler based on stage topology,
+        // not a user-configurable setting.
         ConfigEntry::new(BALLISTA_GRPC_CLIENT_CONNECT_TIMEOUT_SECONDS.to_string(),
                          "Connection timeout for gRPC client in seconds".to_string(),
                          DataType::UInt64,
@@ -352,9 +352,21 @@ impl BallistaConfig {
     }
 
     /// Returns whether this is the final output stage.
+    /// This is an internal flag set by the scheduler, not user-configurable.
     /// 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)
+        self.settings
+            .get(BALLISTA_IS_FINAL_STAGE)
+            .and_then(|v| v.parse::<bool>().ok())
+            .unwrap_or(false)
+    }
+
+    /// Sets the internal is_final_stage flag.
+    /// This should only be called by the scheduler when creating task configurations.
+    pub fn with_is_final_stage(mut self, is_final: bool) -> Self {
+        self.settings
+            .insert(BALLISTA_IS_FINAL_STAGE.to_string(), is_final.to_string());
+        self
     }
 
     /// Returns the configured shuffle format (ArrowIpc or Vortex)

ballista/core/src/extension.rs

@@ -16,12 +16,11 @@
 // under the License.
 
 use crate::config::{
-    BALLISTA_GRPC_CLIENT_MAX_MESSAGE_SIZE, BALLISTA_IS_FINAL_STAGE, BALLISTA_JOB_NAME,
-    BALLISTA_SHUFFLE_FORMAT, BALLISTA_SHUFFLE_MEMORY_MODE,
-    BALLISTA_SHUFFLE_READER_FORCE_REMOTE_READ, BALLISTA_SHUFFLE_READER_MAX_REQUESTS,
-    BALLISTA_SHUFFLE_READER_REMOTE_PREFER_FLIGHT, BALLISTA_SHUFFLE_STORAGE_TYPE,
-    BALLISTA_SHUFFLE_STORAGE_URL, BALLISTA_STANDALONE_PARALLELISM, BallistaConfig,
-    ShuffleFormat,
+    BALLISTA_GRPC_CLIENT_MAX_MESSAGE_SIZE, BALLISTA_JOB_NAME, BALLISTA_SHUFFLE_FORMAT,
+    BALLISTA_SHUFFLE_MEMORY_MODE, BALLISTA_SHUFFLE_READER_FORCE_REMOTE_READ,
+    BALLISTA_SHUFFLE_READER_MAX_REQUESTS, BALLISTA_SHUFFLE_READER_REMOTE_PREFER_FLIGHT,
+    BALLISTA_SHUFFLE_STORAGE_TYPE, BALLISTA_SHUFFLE_STORAGE_URL,
+    BALLISTA_STANDALONE_PARALLELISM, BallistaConfig, ShuffleFormat,
 };
 use crate::planner::BallistaQueryPlanner;
 use crate::serde::protobuf::KeyValuePair;
@@ -513,16 +512,20 @@ impl SessionConfigExt for SessionConfig {
             .extensions
             .get::<BallistaConfig>()
             .map(|c| c.is_final_stage())
-            .unwrap_or_else(|| BallistaConfig::default().is_final_stage())
+            .unwrap_or(false)
     }
 
     fn with_ballista_is_final_stage(self, is_final: bool) -> Self {
-        if self.options().extensions.get::<BallistaConfig>().is_some() {
-            self.set_bool(BALLISTA_IS_FINAL_STAGE, is_final)
-        } else {
-            self.with_option_extension(BallistaConfig::default())
-                .set_bool(BALLISTA_IS_FINAL_STAGE, is_final)
-        }
+        // is_final_stage is an internal flag, not a user-configurable setting,
+        // so we modify the BallistaConfig directly instead of using set_bool
+        let ballista_config = self
+            .options()
+            .extensions
+            .get::<BallistaConfig>()
+            .cloned()
+            .unwrap_or_default()
+            .with_is_final_stage(is_final);
+        self.with_option_extension(ballista_config)
     }
 
     fn with_ballista_grpc_metadata(self, metadata: HashMap<String, String>) -> Self {
@@ -980,7 +983,7 @@ mod test {
     };
 
     use crate::{
-        config::BALLISTA_JOB_NAME,
+        config::{BALLISTA_JOB_NAME, BallistaConfig},
         extension::{SessionConfigExt, SessionConfigHelperExt, SessionStateExt},
     };
 
@@ -1065,17 +1068,21 @@ mod test {
     }
 
     #[test]
-    fn test_is_final_stage_serialization() {
-        use crate::config::BALLISTA_IS_FINAL_STAGE;
-
-        // Test that is_final_stage is included in key-value pairs
+    fn test_is_final_stage_internal_setting() {
+        // Test that is_final_stage is properly stored in BallistaConfig
         let config =
             SessionConfig::new_with_ballista().with_ballista_is_final_stage(true);
-        let pairs = config.to_key_value_pairs();
 
-        let is_final_pair = pairs.iter().find(|p| p.key == BALLISTA_IS_FINAL_STAGE);
-        assert!(is_final_pair.is_some());
-        assert_eq!(is_final_pair.unwrap().value, Some("true".to_string()));
+        // Verify via the getter
+        assert!(config.ballista_is_final_stage());
+
+        // Verify the internal BallistaConfig has the setting
+        let ballista_config = config
+            .options()
+            .extensions
+            .get::<BallistaConfig>()
+            .expect("BallistaConfig should exist");
+        assert!(ballista_config.is_final_stage());
     }
 
     #[test]