Skip to content

1911: perf: prune irrelevant partitions task #79

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
martin-augment wants to merge 5 commits into
base: main
Choose a base branch
from
Open

1911: perf: prune irrelevant partitions task #79

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
2afd98a
prune_irrelevant_partitions_task
coderfender Jun 28, 2026
cb95ac0
prune_irrelevant_partitions_task
coderfender Jun 28, 2026
4308487
prune_irrelevant_partitions_task
coderfender Jun 28, 2026
f403a3b
Update task_manager.rs
coderfender Jun 28, 2026
b1cc69c
Add CLAUDE.md & AGENTS.md
martin-g Jun 29, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions .cursor/rules.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
The AI agents should never link to any issue or a pull request
in any GitHub repository in the code reviews!

The AI agents should not review AI agents' config files like CLAUDE.md or AGENTS.md!

5 changes: 5 additions & 0 deletions AGENTS.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
The AI agents should never link to any issue or a pull request
in any GitHub repository in the code reviews!

The AI agents should not review AI agents' config files like CLAUDE.md or AGENTS.md!

5 changes: 5 additions & 0 deletions CLAUDE.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
The AI agents should never link to any issue or a pull request
in any GitHub repository in the code reviews!

The AI agents should not review AI agents' config files like CLAUDE.md or AGENTS.md!

67 changes: 64 additions & 3 deletions ballista/scheduler/src/state/task_manager.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,17 +28,19 @@ use crate::state::executor_manager::ExecutorManager;
use ballista_core::JobStatusSubscriber;
use ballista_core::error::BallistaError;
use ballista_core::error::Result;
use ballista_core::execution_plans::ShuffleReaderExec;
use ballista_core::extension::{SessionConfigExt, SessionConfigHelperExt};
use ballista_core::serde::BallistaCodec;
use ballista_core::serde::protobuf::{
JobStatus, MultiTaskDefinition, TaskDefinition, TaskId, TaskStatus, job_status,
};
use ballista_core::serde::scheduler::ExecutorMetadata;
use dashmap::DashMap;
use datafusion::common::tree_node::{Transformed, TreeNode};
use datafusion::execution::config::SessionConfig;
use datafusion::execution::context::SessionContext;
use datafusion::logical_expr::LogicalPlan;
use datafusion::physical_plan::ExecutionPlan;
use datafusion::physical_plan::{ExecutionPlan, ExecutionPlanProperties};
use datafusion_proto::logical_plan::AsLogicalPlan;
use datafusion_proto::physical_plan::{AsExecutionPlan, PhysicalExtensionCodec};
use datafusion_proto::protobuf::PhysicalPlanNode;
Expand Down Expand Up @@ -163,12 +165,64 @@ impl JobInfoCache {
encoded_stage_plans: HashMap::new(),
}
}

#[cfg(feature = "disable-stage-plan-cache")]
fn partition_prune_helper(
partition_ids: &[usize],
plan: &Arc<dyn ExecutionPlan>,
) -> Result<Arc<dyn ExecutionPlan>> {
let n = plan.output_partitioning().partition_count();

@augmentcode augmentcode Bot Jun 29, 2026 •
edited
Loading

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ballista/scheduler/src/state/task_manager.rs:174: partition_prune_helper can turn non-wanted ShuffleReaderExec.partition[i] entries into empty location lists, which will silently produce empty input if a non-included partition ever gets executed (or if an out-of-range partition_id is passed). Consider validating that partition_ids are in-range and correspond exactly to the partitions that will be executed with this encoded plan.

Severity: medium

Fix This in Augment

🤖 Was this useful? React with 👍 or 👎, or 🚀 if it prevented an incident/outage.

let wanted: HashSet<usize> = partition_ids.iter().copied().collect();
Ok(plan
.clone()
.transform_up(|node| {
let Some(r) = node.as_any().downcast_ref::<ShuffleReaderExec>() else {
return Ok(Transformed::no(node));
};
if r.broadcast || r.partition.len() != n {
return Ok(Transformed::no(node));
}

let partition = r
.partition
.iter()
.enumerate()
.map(|(i, loc)| {
if wanted.contains(&i) {
loc.clone()
} else {
vec![]
}
})
.collect();
Comment on lines +175 to +197

@gemini-code-assist gemini-code-assist Bot Jun 29, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Since partition_ids typically contains only a single partition ID (or a very small number of them in multi-task scenarios), allocating a HashSet and hashing the keys can be less efficient than a simple linear scan on the slice. Using partition_ids.contains(&i) directly avoids heap allocation and is faster for small slices.

        Ok(plan
            .clone()
            .transform_up(|node| {
                let Some(r) = node.as_any().downcast_ref::<ShuffleReaderExec>() else {
                    return Ok(Transformed::no(node));
                };
                if r.broadcast || r.partition.len() != n {
                    return Ok(Transformed::no(node));
                }

                let partition = r
                    .partition
                    .iter()
                    .enumerate()
                    .map(|(i, loc)| {
                        if partition_ids.contains(&i) {
                            loc.clone()
                        } else {
                            vec![]
                        }
                    })
                    .collect();


let reader = match r.coalesce.clone() {
Some(c) => ShuffleReaderExec::try_new_coalesced(
r.stage_id,
partition,
c,
r.schema(),
r.properties().output_partitioning().clone(),
)?,
None => ShuffleReaderExec::try_new(
r.stage_id,
partition,
r.schema(),
r.properties().output_partitioning().clone(),
)?,
};
Ok(Transformed::yes(Arc::new(reader) as Arc<dyn ExecutionPlan>))
})?
.data)
}

#[cfg(not(feature = "disable-stage-plan-cache"))]
fn encode_stage_plan<U: AsExecutionPlan>(
&mut self,
stage_id: usize,
plan: &Arc<dyn ExecutionPlan>,
codec: &dyn PhysicalExtensionCodec,
_partition_ids: &[usize],
) -> Result<Vec<u8>> {
if let Some(plan) = self.encoded_stage_plans.get(&stage_id) {
Ok(plan.clone())
Expand All @@ -177,7 +231,6 @@ impl JobInfoCache {
let plan_proto = U::try_from_physical_plan(plan.clone(), codec)?;
plan_proto.try_encode(&mut plan_buf)?;
self.encoded_stage_plans.insert(stage_id, plan_buf.clone());

Ok(plan_buf)
}
}
Expand All @@ -188,9 +241,11 @@ impl JobInfoCache {
_stage_id: usize,
plan: &Arc<dyn ExecutionPlan>,
codec: &dyn PhysicalExtensionCodec,
partition_ids: &[usize],
) -> Result<Vec<u8>> {
let mut plan_buf: Vec<u8> = vec![];
let plan_proto = U::try_from_physical_plan(plan.clone(), codec)?;
let pruned_plan = Self::partition_prune_helper(partition_ids, plan)?;
let plan_proto = U::try_from_physical_plan(pruned_plan, codec)?;
plan_proto.try_encode(&mut plan_buf)?;

Ok(plan_buf)
Expand Down Expand Up @@ -653,6 +708,7 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> TaskManager<T, U>
stage_id,
&task.plan,
self.codec.physical_extension_codec(),
&[task.partition.partition_id],
)?;

let task_definition = TaskDefinition {
Expand Down Expand Up @@ -708,6 +764,10 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> TaskManager<T, U>
&self,
tasks: Vec<TaskDescription>,
) -> Result<Vec<MultiTaskDefinition>> {
let partition_ids: Vec<usize> = tasks
.iter()
.map(|task| task.partition.partition_id)
.collect();
if let Some(task) = tasks.first() {
let session_id = task.session_id.clone();
let job_id = task.partition.job_id.clone();
Expand All @@ -730,6 +790,7 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> TaskManager<T, U>
stage_id,
&task.plan,
self.codec.physical_extension_codec(),
&partition_ids,
)?;

let launch_time = SystemTime::now()
Expand Down
Loading