This is an automated email from the ASF dual-hosted git repository.
andygrove pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-ballista.git
The following commit(s) were added to refs/heads/main by this push:
new 1cd1affc4 feat: improve state creation (#1963)
1cd1affc4 is described below
commit 1cd1affc40e1cb9b2f84e0397ff4f1f2003947c5
Author: Marko Milenković <[email protected]>
AuthorDate: Sat Jul 11 15:02:13 2026 +0200
feat: improve state creation (#1963)
---
ballista/scheduler/src/state/aqe/planner.rs | 47 ++++++++++++-----------------
1 file changed, 20 insertions(+), 27 deletions(-)
diff --git a/ballista/scheduler/src/state/aqe/planner.rs
b/ballista/scheduler/src/state/aqe/planner.rs
index 498191227..67d303761 100644
--- a/ballista/scheduler/src/state/aqe/planner.rs
+++ b/ballista/scheduler/src/state/aqe/planner.rs
@@ -30,14 +30,14 @@ use ballista_core::serde::scheduler::PartitionLocation;
use datafusion::common;
use datafusion::common::{HashMap, exec_err};
use datafusion::error::DataFusionError;
+#[cfg(test)]
+use datafusion::execution::config::SessionConfig;
use datafusion::execution::context::SessionContext;
-use datafusion::execution::runtime_env::RuntimeEnv;
use datafusion::execution::{SessionState, SessionStateBuilder};
use datafusion::logical_expr::LogicalPlan;
use datafusion::physical_optimizer::PhysicalOptimizerRule;
use datafusion::physical_plan::{ExecutionPlan, ExecutionPlanProperties,
displayable};
use datafusion::physical_planner::DefaultPhysicalPlanner;
-use datafusion::prelude::SessionConfig;
use log::debug;
use std::collections::HashSet;
use std::fmt::{Debug, Formatter};
@@ -86,9 +86,7 @@ impl AdaptivePlanner {
/// Creates a new `AdaptivePlanner` with the specified physical optimizer
rules.
///
/// # Arguments:
- ///
- /// * `session_config` - The session configuration for the job.
- /// * `runtime_env` - runtime environment
+ /// * `state_builder` - Session state builder,
/// * `plan` - The physical execution plan for the job.
/// * `job_name` - The name of the job.
/// * `physical_optimizer_rules` - A list of physical optimizer rules to
apply.
@@ -96,17 +94,13 @@ impl AdaptivePlanner {
/// # Returns
/// A new instance of `AdaptivePlanner` or an error if the initialization
fails.
pub fn try_new_with_optimizers(
- session_config: &SessionConfig,
- runtime_env: Arc<RuntimeEnv>,
+ state_builder: SessionStateBuilder,
plan: Arc<dyn ExecutionPlan>,
job_name: String,
physical_optimizer_rules: Vec<PhysicalOptimizerRuleRef>,
) -> common::Result<Self> {
- let session_state = Self::create_session_state(
- session_config,
- runtime_env,
- physical_optimizer_rules,
- );
+ let session_state =
+ Self::create_session_state(state_builder,
physical_optimizer_rules);
let planner = DefaultPhysicalPlanner::default();
let plan = planner.optimize_physical_plan(plan, &session_state, |_, _|
{})?;
@@ -137,9 +131,10 @@ impl AdaptivePlanner {
job_name: String,
) -> common::Result<Self> {
let plan_id_generator = Arc::new(AtomicUsize::new(0));
+ let state_builder = SessionStateBuilder::new_with_default_features()
+ .with_config(session_config.clone());
Self::try_new_with_optimizers(
- session_config,
- RuntimeEnv::default().into(),
+ state_builder,
plan,
job_name,
Self::default_optimizers(plan_id_generator),
@@ -166,14 +161,14 @@ impl AdaptivePlanner {
// running standard set of optimizers, which will
// after each stage.
let plan_id_generator = Arc::new(AtomicUsize::new(0));
- let runtime_env = ctx.runtime_env();
- let plan_preparation_stage = Self::create_session_state(
- ctx.state().config(),
- ctx.runtime_env(),
+
+ let plan_preparation_state_builder =
SessionStateBuilder::from(ctx.state());
+ let plan_preparation_state = Self::create_session_state(
+ plan_preparation_state_builder,
Self::plan_preparation_optimizers(plan_id_generator.clone()),
);
- let plan = plan_preparation_stage
+ let plan = plan_preparation_state
.create_physical_plan(logical_plan)
.await?;
@@ -184,9 +179,9 @@ impl AdaptivePlanner {
.await
.map_err(|e| DataFusionError::Execution(e.to_string()))?;
+ let state_builder = SessionStateBuilder::from(ctx.state());
Self::try_new_with_optimizers(
- ctx.state().config(),
- runtime_env,
+ state_builder,
plan,
job_name,
Self::default_optimizers(plan_id_generator),
@@ -551,22 +546,20 @@ impl AdaptivePlanner {
/// Creates a session state with the given configuration and optimizer
rules.
///
/// # Arguments
- /// * `session_config` - The session configuration.
+ /// * `session_builder` - The session builder.
/// * `physical_optimizers` - A list of physical optimizer rules.
///
/// # Returns
/// A new `SessionState` instance.
fn create_session_state(
- session_config: &SessionConfig,
- runtime_env: Arc<RuntimeEnv>,
+ builder: SessionStateBuilder,
physical_optimizers: Vec<PhysicalOptimizerRuleRef>,
) -> SessionState {
- SessionStateBuilder::new_with_default_features()
+ builder
.with_physical_optimizer_rules(physical_optimizers)
- .with_config(session_config.clone())
- .with_runtime_env(runtime_env)
.build()
}
+
/// Recursively finds runnable exchanges in the execution plan.
///
/// # Arguments
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
