Skip to content

1537: feat: move shuffle writer disk I/O off tokio worker threads #10

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 3 commits into
base: main
Choose a base branch
from
Open

1537: feat: move shuffle writer disk I/O off tokio worker threads #10

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
517bce1
feat: move shuffle writer disk I/O off tokio worker threads
hcrosse Apr 2, 2026
fdaf1b7
feat: add concurrent partition execution to shuffle_bench
hcrosse Apr 2, 2026
472f2f3
Add CLAUDE.md & AGENTS.md
martin-g Apr 2, 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!

2 changes: 1 addition & 1 deletion ballista/core/Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -64,7 +64,7 @@ prost = { workspace = true }
prost-types = { workspace = true }
rand = { workspace = true }
serde = { workspace = true, features = ["derive"] }
tokio = { workspace = true, features = ["rt-multi-thread"] }
tokio = { workspace = true, features = ["fs", "rt-multi-thread"] }
tokio-stream = { workspace = true, features = ["net"] }
tonic = { workspace = true }
tonic-prost = { workspace = true }
Expand Down
263 changes: 173 additions & 90 deletions ballista/core/src/execution_plans/shuffle_writer.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -214,14 +214,12 @@ impl ShuffleWriterExec {

match output_partitioning {
None => {
let timer = write_metrics.write_time.timer();
path.push(format!("{input_partition}"));
std::fs::create_dir_all(&path)?;
tokio::fs::create_dir_all(&path).await?;
path.push("data.arrow");
let path = path.to_str().unwrap();
debug!("Writing results to {path}");

// stream results to disk
let stats = utils::write_stream_to_disk(
&mut stream,
path,
Expand All @@ -236,7 +234,6 @@ impl ShuffleWriterExec {
write_metrics
.output_rows
.add(stats.num_rows.unwrap_or(0) as usize);
timer.done();

info!(
"Executed partition {} in {} seconds. Statistics: {}",
Expand All @@ -255,96 +252,114 @@ impl ShuffleWriterExec {
}

Some(Partitioning::Hash(exprs, num_output_partitions)) => {
// we won't necessary produce output for every possible partition, so we
// create writers on demand
let mut writers: Vec<Option<WriteTracker>> = vec![];
for _ in 0..num_output_partitions {
writers.push(None);
}

let mut partitioner = BatchPartitioner::new_hash_partitioner(
exprs,
num_output_partitions,
write_metrics.repart_time.clone(),
);

while let Some(result) = stream.next().await {
let input_batch = result?;

write_metrics.input_rows.add(input_batch.num_rows());

partitioner.partition(
input_batch,
|output_partition, output_batch| {
// partition func in datafusion make sure not write empty output_batch.
let timer = write_metrics.write_time.timer();
match &mut writers[output_partition] {
Some(w) => {
w.num_batches += 1;
w.num_rows += output_batch.num_rows();
w.writer.write(&output_batch)?;
}
None => {
let mut path = path.clone();
path.push(format!("{output_partition}"));
std::fs::create_dir_all(&path)?;

path.push(format!(
"data-{input_partition}.arrow"
));
debug!("Writing results to {path:?}");

let options = IpcWriteOptions::default()
.try_with_compression(Some(
CompressionType::LZ4_FRAME,
))?;

let file =
BufWriter::new(File::create(path.clone())?);
let mut writer =
StreamWriter::try_new_with_options(
file,
stream.schema().as_ref(),
options,
)?;

writer.write(&output_batch)?;
writers[output_partition] = Some(WriteTracker {
num_batches: 1,
num_rows: output_batch.num_rows(),
writer,
path,
});
let schema = stream.schema();
let (tx, mut rx) = tokio::sync::mpsc::channel::<RecordBatch>(2);

@gemini-code-assist gemini-code-assist Bot Apr 2, 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

The channel capacity of 2 is quite low for a shuffle operation. While it provides backpressure, such a small buffer might lead to excessive context switching between the async stream reader and the blocking writer task, potentially reducing throughput. Consider increasing this to a larger value like 16 or 32 to allow for better pipelining of record batches.

Suggested change
let (tx, mut rx) = tokio::sync::mpsc::channel::<RecordBatch>(2);
let (tx, mut rx) = tokio::sync::mpsc::channel::<RecordBatch>(16);

@martin-augment martin-augment Apr 2, 2026

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

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

value:good-to-have; category:bug; feedback: The Gemini AI reviewer is correct! The channel used for passing back the data from the blocking task may contain at most two items. If the receiver is slow for any reason then the sender will block and the thread context will switch. A channel with more slots will allow to reduce the thread context switching by using more memory.

let write_time = write_metrics.write_time.clone();
let repart_time = write_metrics.repart_time.clone();
let output_rows = write_metrics.output_rows.clone();

let handle = tokio::task::spawn_blocking(move || {
let mut writers: Vec<Option<WriteTracker>> =
(0..num_output_partitions).map(|_| None).collect();
let mut partitioner = BatchPartitioner::new_hash_partitioner(
exprs,
num_output_partitions,
repart_time,
);

while let Some(input_batch) = rx.blocking_recv() {
partitioner.partition(
input_batch,
|output_partition, output_batch| {
let timer = write_time.timer();
match &mut writers[output_partition] {
Some(w) => {
w.num_batches += 1;
w.num_rows += output_batch.num_rows();
w.writer.write(&output_batch)?;
}
None => {
let mut p = path.clone();
p.push(format!("{output_partition}"));
std::fs::create_dir_all(&p)?;
p.push(format!(
"data-{input_partition}.arrow"
));
debug!("Writing results to {p:?}");

let options = IpcWriteOptions::default()
.try_with_compression(Some(
CompressionType::LZ4_FRAME,
))?;
let file =
BufWriter::new(File::create(p.clone())?);
let mut writer =
StreamWriter::try_new_with_options(
file,
schema.as_ref(),
options,
)?;
writer.write(&output_batch)?;
writers[output_partition] =
Some(WriteTracker {
num_batches: 1,
num_rows: output_batch.num_rows(),
writer,
path: p,
});
}
}
}
write_metrics.output_rows.add(output_batch.num_rows());
timer.done();
Ok(())
},
)?;
}
output_rows.add(output_batch.num_rows());
timer.done();
Ok(())
},
)?;
}

let mut part_locs = vec![];

for (i, w) in writers.iter_mut().enumerate() {
if let Some(w) = w {
w.writer.finish()?;
let num_bytes = fs::metadata(&w.path)?.len();
debug!(
"Finished writing shuffle partition {} at {:?}. Batches: {}. Rows: {}. Bytes: {}.",
i, w.path, w.num_batches, w.num_rows, num_bytes
);

part_locs.push(ShuffleWritePartition {
partition_id: i as u64,
path: w.path.to_string_lossy().to_string(),
num_batches: w.num_batches as u64,
num_rows: w.num_rows as u64,
num_bytes,
});
let mut part_locs = vec![];
for (i, w) in writers.iter_mut().enumerate() {
if let Some(w) = w {
w.writer.finish()?;
let num_bytes = fs::metadata(&w.path)?.len();
debug!(
"Finished writing shuffle partition {} at {:?}. Batches: {}. Rows: {}. Bytes: {}.",
i, w.path, w.num_batches, w.num_rows, num_bytes
);
part_locs.push(ShuffleWritePartition {
partition_id: i as u64,
path: w.path.to_string_lossy().to_string(),
num_batches: w.num_batches as u64,
num_rows: w.num_rows as u64,
num_bytes,
});
}
}
Ok(part_locs)
});

let stream_err = loop {
match stream.next().await {
Some(Ok(batch)) => {
write_metrics.input_rows.add(batch.num_rows());
if tx.send(batch).await.is_err() {
break None;
}
}
Some(Err(e)) => break Some(e),
None => break None,
}
};
drop(tx);

let write_result = handle.await.map_err(|e| {
DataFusionError::Execution(format!(
"Shuffle writer task failed: {e}"
))
})?;
if let Some(e) = stream_err {
return Err(e);
}
Ok(part_locs)
write_result
}

_ => Err(DataFusionError::Execution(
Expand Down Expand Up @@ -648,6 +663,74 @@ mod tests {
Ok(())
}

#[tokio::test]
async fn test_no_repart_write_failure_propagates() -> Result<()> {
let session_ctx = SessionContext::new();
let task_ctx = session_ctx.task_ctx();

// Place a directory at the data.arrow path so File::create fails
// inside write_stream_to_disk, not at create_dir_all.
// Path structure: work_dir / job_id / stage_id / partition / data.arrow
let tmp = tempfile::TempDir::new().unwrap();
let data_arrow_dir = tmp
.path()
.join("jobOne")
.join("1")
.join("0")
.join("data.arrow");
std::fs::create_dir_all(&data_arrow_dir).unwrap();
let work_dir = tmp.path().to_str().unwrap().to_owned();

let input_plan = Arc::new(CoalescePartitionsExec::new(create_input_plan()?));
let query_stage = ShuffleWriterExec::try_new(
"jobOne".to_owned(),
1,
input_plan,
work_dir,
None,
)?;
let mut stream = query_stage.execute(0, task_ctx)?;
let result = utils::collect_stream(&mut stream).await;
assert!(
result.is_err(),
"expected File::create failure in write_stream_to_disk to propagate"
);
Ok(())
}

#[tokio::test]
async fn test_hash_repart_write_failure_propagates() -> Result<()> {
let session_ctx = SessionContext::new();
let task_ctx = session_ctx.task_ctx();

// Place a regular file at the stage_id path component so
// create_dir_all fails when the writer tries to create
// output partition subdirectories underneath it.
// Path structure: work_dir / job_id / stage_id / ...
let tmp = tempfile::TempDir::new().unwrap();
let job_dir = tmp.path().join("jobOne");
std::fs::create_dir_all(&job_dir).unwrap();
let stage_id_as_file = job_dir.join("1");
std::fs::File::create(&stage_id_as_file).unwrap();
let work_dir = tmp.path().to_str().unwrap().to_owned();

let input_plan = Arc::new(CoalescePartitionsExec::new(create_input_plan()?));
let query_stage = ShuffleWriterExec::try_new(
"jobOne".to_owned(),
1,
input_plan,
work_dir,
Some(Partitioning::Hash(vec![Arc::new(Column::new("a", 0))], 2)),
)?;
let mut stream = query_stage.execute(0, task_ctx)?;
let result = utils::collect_stream(&mut stream).await;
assert!(
result.is_err(),
"expected create_dir_all failure in hash writer to propagate"
);
Ok(())
}

fn create_input_plan() -> Result<Arc<dyn ExecutionPlan>> {
let schema = Arc::new(Schema::new(vec![
Field::new("a", DataType::UInt32, true),
Expand Down
Loading
Loading