Skip to content

Remove the hash-partitioning shuffle writer #2106

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

Merged
andygrove merged 14 commits into from
Jul 21, 2026
Merged

Remove the hash-partitioning shuffle writer #2106

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
14 commits
Select commit Hold shift + click to select a range
22aa1ed
refactor: drop hash-partitioning branch from ShuffleWriterExec
andygrove Jul 20, 2026
0ee427a
test: update EXPLAIN ANALYZE golden after repart_time removal
andygrove Jul 20, 2026
8126688
refactor: always use sort-based shuffle for hash-repartition stages
andygrove Jul 20, 2026
854ff1e
refactor: drop hash writer arm from shuffle_bench
andygrove Jul 20, 2026
db172cd
test: drop hash-vs-sort shuffle comparison tests
andygrove Jul 20, 2026
340ce0b
test: drop sort-shuffle enable toggle from AQE tests
andygrove Jul 20, 2026
26cf662
refactor: gate sort-shuffle reads on on-disk index only
andygrove Jul 20, 2026
78d51a1
test: fix local shuffle read test for single-partition writer
andygrove Jul 20, 2026
d17bef5
feat: remove ballista.shuffle.sort_based.enabled config key
andygrove Jul 20, 2026
da8a035
docs: remove hash-based shuffle documentation
andygrove Jul 20, 2026
c1d3be4
style: apply rustfmt to shuffle_bench
andygrove Jul 20, 2026
85e5ad7
refactor: clean up stale references to removed sort_based.enabled key
andygrove Jul 20, 2026
bec1af9
docs: apply prettier to shuffle tuning table
andygrove Jul 20, 2026
beee639
refactor: drop the unused output partitioning from ShuffleWriterExec
andygrove Jul 21, 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
2 changes: 1 addition & 1 deletion ballista/client/tests/context_checks.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1192,7 +1192,7 @@ mod supported {
"| | PlaceholderRowExec, metrics=[...] |",
"| | |",
"| | =========SuccessfulStage[stage_id=2, partitions=16]========= |",
"| | ShuffleWriterExec: partitioning: None, metrics=[output_rows=..., input_rows=..., repart_time=..., write_time=...] |",
"| | ShuffleWriterExec: partitioning: None, metrics=[output_rows=..., input_rows=..., write_time=...] |",
"| | ProjectionExec: expr=[count(Int64(1))@1 as count(*), id@0 as id], metrics=[output_rows=..., elapsed_compute=..., output_bytes=..., output_batches=..., expr_0_eval_time=..., expr_1_eval_time=...] |",
"| | AggregateExec: mode=FinalPartitioned, gby=[id@0 as id], aggr=[count(Int64(1))], metrics=[output_rows=..., elapsed_compute=..., output_bytes=..., output_batches=..., spill_count=..., spilled_bytes=..., spilled_rows=..., peak_mem_used=..., aggregate_arguments_time=..., aggregation_time=..., emitting_time=..., time_calculating_group_ids=...] |",
"| | ShuffleReaderExec: upstream_stage: 1, partitioning: Hash([id@0], 16), metrics=[output_rows=..., elapsed_compute=..., output_bytes=..., output_batches=..., decoded_bytes=..., fetch_requests=..., fetch_retries=..., local_partitions=..., remote_partitions=..., fetch_time=..., local_read_time=..., permit_wait_time=...] |",
Expand Down
86 changes: 0 additions & 86 deletions ballista/client/tests/sort_shuffle.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,14 +37,12 @@ mod sort_shuffle_tests {
BALLISTA_SHUFFLE_READER_MAX_BLOCKS_PER_ADDRESS,
BALLISTA_SHUFFLE_READER_MAX_BYTES_IN_FLIGHT,
BALLISTA_SHUFFLE_READER_REMOTE_PREFER_FLIGHT,
BALLISTA_SHUFFLE_SORT_BASED_ENABLED,
};
use datafusion::arrow::util::pretty::pretty_format_batches;
use datafusion::common::Result;
use datafusion::execution::SessionStateBuilder;
use datafusion::prelude::{ParquetReadOptions, SessionConfig, SessionContext};
use rstest::rstest;
use std::collections::HashSet;

/// Read mode for shuffle data
#[derive(Debug, Clone, Copy)]
Expand All @@ -69,7 +67,6 @@ mod sort_shuffle_tests {
aqe_enabled: bool,
) -> SessionContext {
let mut config = SessionConfig::new_with_ballista()
.set_str(BALLISTA_SHUFFLE_SORT_BASED_ENABLED, "true")
.set_bool(BALLISTA_ADAPTIVE_PLANNER_ENABLED, aqe_enabled);

// Configure read mode
Expand Down Expand Up @@ -103,7 +100,6 @@ mod sort_shuffle_tests {
/// complete and return correct results.
async fn create_tiny_budget_remote_context() -> SessionContext {
let config = SessionConfig::new_with_ballista()
.set_str(BALLISTA_SHUFFLE_SORT_BASED_ENABLED, "true")
.set_str(BALLISTA_SHUFFLE_READER_FORCE_REMOTE_READ, "true")
.set_str(BALLISTA_SHUFFLE_READER_REMOTE_PREFER_FLIGHT, "true")
.set_str(BALLISTA_SHUFFLE_READER_MAX_BYTES_IN_FLIGHT, "65536")
Expand All @@ -117,17 +113,6 @@ mod sort_shuffle_tests {
SessionContext::standalone_with_state(state).await.unwrap()
}

/// Creates a standalone session context with hash-based shuffle.
async fn create_hash_shuffle_context() -> SessionContext {
let config = SessionConfig::new_with_ballista()
.set_str(BALLISTA_SHUFFLE_SORT_BASED_ENABLED, "false");
let state = SessionStateBuilder::new()
.with_config(config)
.with_default_features()
.build();
SessionContext::standalone_with_state(state).await.unwrap()
}

/// Registers test data in the context.
async fn register_test_data(ctx: &SessionContext) {
ctx.register_parquet(
Expand All @@ -154,21 +139,6 @@ mod sort_shuffle_tests {
);
}

/// Extracts values from a result set, ignoring order.
fn extract_values_unordered(
results: &[datafusion::arrow::record_batch::RecordBatch],
) -> HashSet<String> {
pretty_format_batches(results)
.unwrap()
.to_string()
.trim()
.lines()
.skip(3) // Skip header lines
.filter(|line| !line.starts_with('+'))
.map(|s| s.to_string())
.collect()
}

// ==================== Basic Aggregation Tests ====================

#[rstest]
Expand Down Expand Up @@ -327,62 +297,6 @@ mod sort_shuffle_tests {
Ok(())
}

// ==================== Comparison with Hash Shuffle ====================

#[tokio::test]
async fn test_sort_vs_hash_shuffle_group_by() -> Result<()> {
// Test with sort shuffle (local read is sufficient for comparison)
let sort_ctx = create_sort_shuffle_context(ReadMode::Local).await;
register_test_data(&sort_ctx).await;
let sort_results = sort_ctx
.sql("SELECT bool_col, SUM(id) as total FROM test GROUP BY bool_col")
.await?
.collect()
.await?;

// Test with hash shuffle
let hash_ctx = create_hash_shuffle_context().await;
register_test_data(&hash_ctx).await;
let hash_results = hash_ctx
.sql("SELECT bool_col, SUM(id) as total FROM test GROUP BY bool_col")
.await?
.collect()
.await?;

// Results should be equivalent (order may differ)
let sort_values = extract_values_unordered(&sort_results);
let hash_values = extract_values_unordered(&hash_results);
assert_eq!(sort_values, hash_values);
Ok(())
}

#[tokio::test]
async fn test_sort_vs_hash_shuffle_distinct() -> Result<()> {
// Test with sort shuffle (local read is sufficient for comparison)
let sort_ctx = create_sort_shuffle_context(ReadMode::Local).await;
register_test_data(&sort_ctx).await;
let sort_results = sort_ctx
.sql("SELECT DISTINCT bool_col FROM test")
.await?
.collect()
.await?;

// Test with hash shuffle
let hash_ctx = create_hash_shuffle_context().await;
register_test_data(&hash_ctx).await;
let hash_results = hash_ctx
.sql("SELECT DISTINCT bool_col FROM test")
.await?
.collect()
.await?;

// Results should be equivalent
let sort_values = extract_values_unordered(&sort_results);
let hash_values = extract_values_unordered(&hash_results);
assert_eq!(sort_values, hash_values);
Ok(())
}

// ==================== Edge Cases ====================

#[rstest]
Expand Down
15 changes: 0 additions & 15 deletions ballista/core/src/config.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -84,9 +84,6 @@ pub const BALLISTA_CLIENT_IO_RETRY_WAIT_TIME_MS: &str =
"ballista.client.io_retry_wait_time_ms";
/// Enables adaptive query planning
pub const BALLISTA_ADAPTIVE_PLANNER_ENABLED: &str = "ballista.planner.adaptive.enabled";
/// Configuration key for enabling sort-based shuffle.
pub const BALLISTA_SHUFFLE_SORT_BASED_ENABLED: &str =
"ballista.shuffle.sort_based.enabled";
/// Configuration key for sort shuffle target batch size in rows.
pub const BALLISTA_SHUFFLE_SORT_BASED_BATCH_SIZE: &str =
"ballista.shuffle.sort_based.batch_size";
Expand Down Expand Up @@ -226,10 +223,6 @@ static CONFIG_ENTRIES: LazyLock<HashMap<String, ConfigEntry>> = LazyLock::new(||
"Enables Adaptive Query Planning (EXPERIMENTAL)".to_string(),
DataType::Boolean,
Some(false.to_string())),
ConfigEntry::new(BALLISTA_SHUFFLE_SORT_BASED_ENABLED.to_string(),
"Enable sort-based shuffle which writes consolidated files with index".to_string(),
DataType::Boolean,
Some(true.to_string())),
ConfigEntry::new(BALLISTA_SHUFFLE_SORT_BASED_BATCH_SIZE.to_string(),
"Target batch size in rows for coalescing small batches in sort shuffle".to_string(),
DataType::UInt64,
Expand Down Expand Up @@ -576,14 +569,6 @@ impl BallistaConfig {
self.get_bool_setting(BALLISTA_ADAPTIVE_PLANNER_ENABLED)
}

/// Returns whether sort-based shuffle is enabled.
///
/// When enabled, shuffle writes produce a single consolidated file per input
/// partition with an index file, rather than one file per output partition.
pub fn shuffle_sort_based_enabled(&self) -> bool {
self.get_bool_setting(BALLISTA_SHUFFLE_SORT_BASED_ENABLED)
}

/// Returns the target batch size for sort-based shuffle.
pub fn shuffle_sort_based_batch_size(&self) -> usize {
self.get_usize_setting(BALLISTA_SHUFFLE_SORT_BASED_BATCH_SIZE)
Expand Down
23 changes: 9 additions & 14 deletions ballista/core/src/execution_plans/shuffle_reader.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -772,7 +772,6 @@ fn send_fetch_partitions(
let max_blocks_per_addr =
ballista_config.shuffle_reader_max_blocks_in_flight_per_address();
let default_block_size = ballista_config.shuffle_reader_default_block_size_bytes();
let sort_shuffle_enabled = config.ballista_sort_shuffle_enabled();

let (response_sender, response_receiver) = mpsc::channel(max_reqs.max(1));

Expand Down Expand Up @@ -816,7 +815,7 @@ fn send_fetch_partitions(
for p in local_locations {
let r = {
let _timer = local_read_time.timer();
fetch_partition_local(&work_dir, &p, sort_shuffle_enabled)
fetch_partition_local(&work_dir, &p)
};
if let Err(e) = response_sender_c.blocking_send(r) {
error!("Fail to send response event to the channel due to {e}");
Expand Down Expand Up @@ -1112,7 +1111,6 @@ async fn fetch_partition_remote(
fn fetch_partition_local(
work_dir: &str,
location: &PartitionLocation,
sort_shuffle_enabled: bool,
) -> result::Result<SendableRecordBatchStream, BallistaError> {
let path = &location.path(work_dir)?;
let metadata = &location.executor_meta;
Expand All @@ -1123,10 +1121,10 @@ fn fetch_partition_local(
// replace this check with open, and check for error
//
// Check if this is a sort-based shuffle output (has index file)
if sort_shuffle_enabled && is_sort_shuffle_output(data_path) {
// note: in some cases sort shuffle is not going to be used
// even its enabled. thus we need to check if there is
// sort shuffle file index
if is_sort_shuffle_output(data_path) {
// A stage's on-disk layout is authoritative: sort-shuffle outputs have a
// companion index file. Standard single-partition outputs do not, so a
// missing index means this is a plain Arrow IPC file.
debug!(
"Reading sort-based shuffle for partition {} from {:?}",
partition_id.partition_id, data_path
Expand All @@ -1147,7 +1145,7 @@ fn fetch_partition_local(
});
}
debug!("fetch local partition file: {data_path:?} ");
// Standard hash-based shuffle - read the file directly
// Standard single-file shuffle output - read the file directly
let reader = fetch_partition_local_inner(path).map_err(|e| {
// return BallistaError::FetchFailed may let scheduler retry this task.
BallistaError::FetchFailed(
Expand Down Expand Up @@ -1290,7 +1288,6 @@ mod tests {
use datafusion::common::DataFusionError;
use datafusion::datasource::memory::MemorySourceConfig;
use datafusion::datasource::source::DataSourceExec;
use datafusion::physical_expr::expressions::Column;
use datafusion::physical_plan::common;
use datafusion::prelude::SessionContext;
use tempfile::{TempDir, tempdir};
Expand Down Expand Up @@ -1744,7 +1741,6 @@ mod tests {
1,
create_test_data_plan().unwrap(),
work_dir.path().to_str().unwrap().to_owned(),
Some(Partitioning::Hash(vec![Arc::new(Column::new("a", 0))], 1)),
)
.unwrap();

Expand Down Expand Up @@ -1772,10 +1768,9 @@ mod tests {
.map_err(|e| DataFusionError::Execution(format!("{e:?}")))
.unwrap();

// Writer's K=1 hash routes every row to partition 0. Its coordinator
// reads its full input slice (2 partitions × 2 batches = 4 batches),
// all routed to the single output file.
assert_eq!(result.len(), 4);
// With single-partition (None) output, executing input partition 0
// writes just that partition's 2 batches to a single output file.
assert_eq!(result.len(), 2);
for b in result {
assert_eq!(b, create_test_batch())
}
Expand Down
Loading
Loading