feat: let a standalone cluster run push-staged

ballista/client/tests/common/mod.rs

@@ -156,6 +156,43 @@ pub async fn setup_test_cluster_with_state(session_state: SessionState) -> (Stri
     (host, addr.port())
 }
 
+/// Starts a scheduler and `executors` executors of `vcores` vcores each, all
+/// in this process and all from `session_state`.
+#[allow(dead_code)]
+pub async fn setup_test_cluster_with_executors(
+    session_state: &SessionState,
+    executors: usize,
+    vcores: usize,
+) -> (String, u16) {
+    let addr = ballista_scheduler::standalone::new_standalone_scheduler_from_state(
+        session_state,
+    )
+    .await
+    .expect("scheduler to be created");
+
+    let host = "localhost".to_string();
+
+    for _ in 0..executors {
+        let scheduler =
+            connect_to_scheduler(format!("http://{}:{}", host, addr.port())).await;
+        ballista_executor::new_standalone_executor_from_state(
+            scheduler,
+            vcores,
+            session_state,
+        )
+        .await
+        .expect("executor to be created");
+    }
+
+    log::info!(
+        "test scheduler with {executors} executors created at: {}:{}",
+        host,
+        addr.port()
+    );
+
+    (host, addr.port())
+}
+
 #[allow(dead_code)]
 pub async fn setup_test_cluster_with_builders(
     config_producer: ConfigProducer,

ballista/client/tests/standalone_push_staged.rs

@@ -0,0 +1,159 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+mod common;
+
+/// End-to-end tests of a standalone cluster running push-staged, selected with
+/// `ballista.standalone.task_scheduling_policy`. Each query's results are
+/// compared with the same query run on a pull-staged standalone cluster.
+#[cfg(test)]
+#[cfg(feature = "standalone")]
+mod standalone_push_staged {
+    use crate::common::setup_test_cluster_with_executors;
+    use ballista::prelude::SessionContextExt;
+    use ballista_core::config::{
+        BALLISTA_STANDALONE_PARALLELISM, BALLISTA_STANDALONE_TASK_SCHEDULING_POLICY,
+    };
+    use ballista_core::extension::SessionConfigExt;
+    use ballista_core::serde::protobuf::PollWorkParams;
+    use ballista_core::serde::protobuf::scheduler_grpc_client::SchedulerGrpcClient;
+    use datafusion::arrow::util::pretty::pretty_format_batches;
+    use datafusion::execution::{SessionState, SessionStateBuilder};
+    use datafusion::prelude::{CsvReadOptions, SessionConfig, SessionContext};
+    use std::io::Write;
+    use std::path::Path;
+    use tempfile::TempDir;
+
+    /// Writes an `orders` table split over four files, so its scan has four
+    /// partitions, and a small `customers` table.
+    fn write_tables(dir: &Path) {
+        let orders = dir.join("orders");
+        std::fs::create_dir(&orders).unwrap();
+        for file in 0..4 {
+            let mut out =
+                std::fs::File::create(orders.join(format!("part-{file}.csv"))).unwrap();
+            writeln!(out, "o_id,o_cust,o_amount").unwrap();
+            for i in 0..250 {
+                let id = file * 250 + i;
+                writeln!(out, "{id},{},{}", id % 97, (id * 7) % 1000).unwrap();
+            }
+        }
+        let mut out = std::fs::File::create(dir.join("customers.csv")).unwrap();
+        writeln!(out, "c_id,c_name,c_region").unwrap();
+        for id in 0..100 {
+            writeln!(out, "{id},customer-{id},region-{}", id % 5).unwrap();
+        }
+    }
+
+    fn state(policy: &str) -> SessionState {
+        let config = SessionConfig::new_with_ballista()
+            .with_target_partitions(4)
+            .set_str(BALLISTA_STANDALONE_PARALLELISM, "4")
+            .set_str(BALLISTA_STANDALONE_TASK_SCHEDULING_POLICY, policy);
+        SessionStateBuilder::new()
+            .with_config(config)
+            .with_default_features()
+            .build()
+    }
+
+    async fn register_tables(ctx: &SessionContext, dir: &Path) {
+        ctx.register_csv(
+            "orders",
+            dir.join("orders").to_str().unwrap(),
+            CsvReadOptions::new(),
+        )
+        .await
+        .unwrap();
+        ctx.register_csv(
+            "customers",
+            dir.join("customers.csv").to_str().unwrap(),
+            CsvReadOptions::new(),
+        )
+        .await
+        .unwrap();
+    }
+
+    async fn rows(ctx: &SessionContext, sql: &str) -> String {
+        let batches = ctx.sql(sql).await.unwrap().collect().await.unwrap();
+        pretty_format_batches(&batches).unwrap().to_string()
+    }
+
+    const QUERIES: &[&str] = &[
+        "SELECT count(*), sum(o_amount) FROM orders",
+        "SELECT o_cust, count(*) AS n, sum(o_amount) AS total \
+         FROM orders GROUP BY o_cust ORDER BY o_cust",
+        "SELECT c_region, count(*) AS n, sum(o_amount) AS total \
+         FROM orders JOIN customers ON o_cust = c_id \
+         GROUP BY c_region ORDER BY c_region",
+    ];
+
+    async fn expected(dir: &Path) -> Vec<String> {
+        let ctx = SessionContext::standalone_with_state(state("pull-staged"))
+            .await
+            .unwrap();
+        register_tables(&ctx, dir).await;
+        let mut out = Vec::new();
+        for sql in QUERIES {
+            out.push(rows(&ctx, sql).await);
+        }
+        out
+    }
+
+    #[tokio::test]
+    async fn push_staged_standalone_matches_pull_staged() {
+        let dir = TempDir::new().unwrap();
+        write_tables(dir.path());
+        let expected = expected(dir.path()).await;
+
+        let ctx = SessionContext::standalone_with_state(state("push-staged"))
+            .await
+            .unwrap();
+        register_tables(&ctx, dir.path()).await;
+        for (sql, expected) in QUERIES.iter().zip(&expected) {
+            assert_eq!(&rows(&ctx, sql).await, expected, "{sql}");
+        }
+    }
+
+    #[tokio::test]
+    async fn push_staged_cluster_of_two_executors_matches_pull_staged() {
+        let dir = TempDir::new().unwrap();
+        write_tables(dir.path());
+        let expected = expected(dir.path()).await;
+
+        let state = state("push-staged");
+        let (host, port) = setup_test_cluster_with_executors(&state, 2, 2).await;
+
+        // The results match pull-staged by design, so check the scheduler
+        // really is push-staged: one of those rejects polling for work.
+        let mut scheduler = SchedulerGrpcClient::connect(format!("http://{host}:{port}"))
+            .await
+            .unwrap();
+        let status = scheduler
+            .poll_work(PollWorkParams::default())
+            .await
+            .unwrap_err();
+        assert_eq!(status.code(), tonic::Code::FailedPrecondition);
+
+        let ctx =
+            SessionContext::remote_with_state(&format!("df://{host}:{port}"), state)
+                .await
+                .unwrap();
+        register_tables(&ctx, dir.path()).await;
+        for (sql, expected) in QUERIES.iter().zip(&expected) {
+            assert_eq!(&rows(&ctx, sql).await, expected, "{sql}");
+        }
+    }
+}

ballista/core/src/config.rs

@@ -32,6 +32,9 @@ use std::{collections::HashMap, fmt::Display};
 pub const BALLISTA_JOB_NAME: &str = "ballista.job.name";
 /// Configuration key for standalone processing parallelism.
 pub const BALLISTA_STANDALONE_PARALLELISM: &str = "ballista.standalone.parallelism";
+/// Configuration key for the task scheduling policy of a standalone cluster.
+pub const BALLISTA_STANDALONE_TASK_SCHEDULING_POLICY: &str =
+    "ballista.standalone.task_scheduling_policy";
 /// Configuration key for disabling default cache extension node.
 pub const BALLISTA_CACHE_NOOP: &str = "ballista.cache.noop";
 /// Configuration key for maximum concurrent shuffle read requests.
@@ -252,6 +255,9 @@ static CONFIG_ENTRIES: LazyLock<HashMap<String, ConfigEntry>> = LazyLock::new(||
                          "Number of vcores advertised by the standalone in-process executor.".to_string(),
                          DataType::UInt16, Some(std::thread::available_parallelism().map(|v| v.get()).unwrap_or(1).to_string()))
             .with_doc_default("number of available CPU cores"),
+        ConfigEntry::new(BALLISTA_STANDALONE_TASK_SCHEDULING_POLICY.to_string(),
+                         "Task scheduling policy of a standalone in-process cluster: `pull-staged`, where the executor polls the scheduler for work, or `push-staged`, where the scheduler pushes tasks to the executor.".to_string(),
+                         DataType::Utf8, Some(TaskSchedulingPolicy::PullStaged.to_string())),
         ConfigEntry::new(BALLISTA_CACHE_NOOP.to_string(),
                          "When true (the default), DataFrame::cache() does nothing. When false, the first action on a cached DataFrame materializes it in the cluster, and later actions read it back instead of recomputing it".to_string(),
                          DataType::Boolean,
@@ -748,6 +754,25 @@ impl BallistaConfig {
         self.get_usize_setting(BALLISTA_STANDALONE_PARALLELISM)
     }
 
+    /// Returns the task scheduling policy of a standalone in-process cluster.
+    pub fn standalone_task_scheduling_policy(
+        &self,
+    ) -> std::result::Result<TaskSchedulingPolicy, DataFusionError> {
+        match self
+            .get_string_setting(BALLISTA_STANDALONE_TASK_SCHEDULING_POLICY)
+            .trim()
+            .to_lowercase()
+            .as_str()
+        {
+            "pull-staged" => Ok(TaskSchedulingPolicy::PullStaged),
+            "push-staged" => Ok(TaskSchedulingPolicy::PushStaged),
+            other => Err(DataFusionError::Configuration(format!(
+                "Got an invalid standalone task scheduling policy {other}: \
+                Valid options are: pull-staged and push-staged"
+            ))),
+        }
+    }
+
     /// Deprecated alias for [`Self::standalone_parallelism`].
     #[deprecated(since = "55.0.0", note = "renamed to `standalone_parallelism`")]
     pub fn default_standalone_parallelism(&self) -> usize {
@@ -1262,6 +1287,46 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn standalone_task_scheduling_policy_defaults_to_pull_staged() {
+        use crate::extension::SessionConfigExt;
+        use datafusion::prelude::SessionConfig;
+
+        let mut config = SessionConfig::new_with_ballista();
+        assert!(matches!(
+            config
+                .ballista_config()
+                .standalone_task_scheduling_policy()
+                .unwrap(),
+            TaskSchedulingPolicy::PullStaged
+        ));
+
+        config
+            .options_mut()
+            .set(BALLISTA_STANDALONE_TASK_SCHEDULING_POLICY, "Push-Staged")
+            .unwrap();
+        assert!(matches!(
+            config
+                .ballista_config()
+                .standalone_task_scheduling_policy()
+                .unwrap(),
+            TaskSchedulingPolicy::PushStaged
+        ));
+
+        config
+            .options_mut()
+            .set(BALLISTA_STANDALONE_TASK_SCHEDULING_POLICY, "eager")
+            .unwrap();
+        let err = config
+            .ballista_config()
+            .standalone_task_scheduling_policy()
+            .unwrap_err();
+        assert!(
+            err.to_string()
+                .contains("invalid standalone task scheduling policy eager")
+        );
+    }
+
     #[test]
     fn execution_mode_defaults_to_blocking_and_parses_pipelined() {
         use crate::extension::SessionConfigExt;

ballista/executor/src/executor_server.rs

@@ -31,6 +31,7 @@ use tokio::sync::mpsc;
 
 use log::{debug, error, info, warn};
 use tonic::transport::Channel;
+use tonic::transport::server::TcpIncoming;
 use tonic::{Request, Response, Status};
 
 use ballista_core::error::BallistaError;
@@ -106,13 +107,51 @@ struct CuratorTaskStatus {
 ///
 /// Returns a handle to the server task that can be awaited for completion.
 pub async fn startup<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan>(
+    scheduler: SchedulerGrpcClient<Channel>,
+    config: Arc<ExecutorProcessConfig>,
+    executor: Arc<Executor>,
+    codec: BallistaCodec<T, U>,
+    stop_send: mpsc::Sender<bool>,
+    shutdown_noti: &ShutdownNotifier,
+    health: ExecutorHealth,
+) -> Result<ServerHandle, BallistaError> {
+    // The listening socket is bound here rather than inside the spawned task,
+    // because registration with the scheduler makes the scheduler dial this
+    // port back to check connectivity. Binding lazily inside the server future
+    // let that callback lose the race and get ECONNREFUSED, which fails
+    // registration and takes the executor down with it.
+    let addr = format!("{}:{}", config.bind_host, executor.metadata.grpc_port);
+    let addr = addr.parse().unwrap();
+    let incoming = create_grpc_server_incoming(addr, &config.grpc_server_config)?;
+    info!("Ballista v{BALLISTA_VERSION} Rust Executor Grpc Server listening on {addr:?}");
+    startup_with_incoming(
+        scheduler,
+        config,
+        executor,
+        codec,
+        stop_send,
+        shutdown_noti,
+        health,
+        incoming,
+    )
+    .await
+}
+
+/// Like [`startup`], but serves the executor gRPC service on `incoming`, a
+/// socket the caller has already bound to `executor.metadata.grpc_port`.
+#[allow(clippy::too_many_arguments)]
+pub(crate) async fn startup_with_incoming<
+    T: 'static + AsLogicalPlan,
+    U: 'static + AsExecutionPlan,
+>(
     mut scheduler: SchedulerGrpcClient<Channel>,
     config: Arc<ExecutorProcessConfig>,
     executor: Arc<Executor>,
     codec: BallistaCodec<T, U>,
     stop_send: mpsc::Sender<bool>,
     shutdown_noti: &ShutdownNotifier,
     health: ExecutorHealth,
+    incoming: TcpIncoming,
 ) -> Result<ServerHandle, BallistaError> {
     let channel_buf_size = executor.vcores * 50;
     let (tx_task, rx_task) = mpsc::channel::<CuratorTaskDefinition>(channel_buf_size);
@@ -135,23 +174,9 @@ pub async fn startup<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan>(
         health,
     );
 
-    // 1. Start executor grpc service
-    //
-    // The listening socket is bound here rather than inside the spawned task,
-    // because step 2 registers with the scheduler and the scheduler dials this
-    // port back to check connectivity. Binding lazily inside the server future
-    // let that callback lose the race and get ECONNREFUSED, which fails
-    // registration and takes the executor down with it.
+    // 1. Start executor grpc service on the already-bound socket
     let server = {
-        let executor_meta = executor.metadata.clone();
-        let addr = format!("{}:{}", config.bind_host, executor_meta.grpc_port);
-        let addr = addr.parse().unwrap();
         let grpc_server_config = config.grpc_server_config.clone();
-
-        let incoming = create_grpc_server_incoming(addr, &grpc_server_config)?;
-        info!(
-            "Ballista v{BALLISTA_VERSION} Rust Executor Grpc Server listening on {addr:?}"
-        );
         let server = ExecutorGrpcServer::new(executor_server.clone())
             .max_encoding_message_size(config.grpc_max_encoding_message_size as usize)
             .max_decoding_message_size(config.grpc_max_decoding_message_size as usize);