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 6540d0851 Remove the hash-partitioning shuffle writer (#2106)
6540d0851 is described below
commit 6540d0851d7604fd759b41594d9652e5a4432489
Author: Andy Grove <[email protected]>
AuthorDate: Tue Jul 21 11:58:49 2026 -0600
Remove the hash-partitioning shuffle writer (#2106)
* refactor: drop hash-partitioning branch from ShuffleWriterExec
Guard try_new() to reject Some(Partitioning::Hash(..)) and delete the
now-unreachable Hash match arm from execute_shuffle_write, along with
WriteTracker and the repart_time metric it used. Hash-repartition
stages are handled by SortShuffleWriterExec instead; ShuffleWriterExec
now only ever writes the single-file (None) passthrough path.
Update the module's existing tests that exercised the Hash arm to
exercise the surviving None path instead (test, test_partitioned
folded into test, display_renders_child_operator_metrics,
test_hash_repart_write_failure_propagates renamed to
test_create_dir_failure_propagates), and add
try_new_rejects_hash_partitioning to cover the new guard.
* test: update EXPLAIN ANALYZE golden after repart_time removal
* refactor: always use sort-based shuffle for hash-repartition stages
* refactor: drop hash writer arm from shuffle_bench
* test: drop hash-vs-sort shuffle comparison tests
* test: drop sort-shuffle enable toggle from AQE tests
* refactor: gate sort-shuffle reads on on-disk index only
* test: fix local shuffle read test for single-partition writer
* feat: remove ballista.shuffle.sort_based.enabled config key
* docs: remove hash-based shuffle documentation
* style: apply rustfmt to shuffle_bench
* refactor: clean up stale references to removed sort_based.enabled key
* docs: apply prettier to shuffle tuning table
* refactor: drop the unused output partitioning from ShuffleWriterExec
`ShuffleWriterExec` no longer repartitions, so its `Option<Partitioning>`
was dead in every real path: the planner passes `None` at all call sites,
the AQE adapter passes `None` or `Hash` (which routes to
`SortShuffleWriterExec`), and proto decode can only produce `None` or
`Hash`. Rather than reject `Some(Partitioning::Hash(..))` at runtime,
remove the parameter and field entirely so the writer is unambiguously
the passthrough writer.
- `ShuffleWriterExec::try_new` drops its `shuffle_output_partitioning`
argument; the struct drops the field and derives its properties from
the child plan.
- The `ShuffleWriter` trait impl returns `None` unconditionally.
- Proto decode errors on a `ShuffleWriterExecNode` that still carries an
`output_partitioning`; encode always writes `None`.
- `create_shuffle_writer_with_config` errors on a non-hash `Some(..)`
instead of silently building a writer that could not be encoded.
- `ShuffleWriterVariant::Hash` is renamed to `Passthrough`, which is what
the variant has actually meant since the hash write path was removed.
Plan display keeps `partitioning: None` so TPC-H plan-stability golden
files are unchanged.
---
ballista/client/tests/context_checks.rs | 2 +-
ballista/client/tests/sort_shuffle.rs | 86 ----
ballista/core/src/config.rs | 15 -
.../core/src/execution_plans/shuffle_reader.rs | 23 +-
.../core/src/execution_plans/shuffle_writer.rs | 570 +++++----------------
ballista/core/src/extension.rs | 11 -
ballista/core/src/serde/mod.rs | 41 +-
ballista/executor/src/execution_engine.rs | 22 +-
ballista/executor/src/executor.rs | 3 +-
ballista/executor/src/flight_service.rs | 2 +-
ballista/scheduler/src/planner.rs | 33 +-
ballista/scheduler/src/state/aqe/test/mod.rs | 16 -
.../scheduler/src/state/aqe/test/plan_to_stages.rs | 33 +-
benchmarks/src/bin/shuffle_bench.rs | 82 +--
benchmarks/src/bin/tpch.rs | 2 +-
docs/source/user-guide/configs.md | 3 +-
docs/source/user-guide/tuning-guide.md | 38 +-
python/python/tests/test_context.py | 4 +-
18 files changed, 220 insertions(+), 766 deletions(-)
diff --git a/ballista/client/tests/context_checks.rs
b/ballista/client/tests/context_checks.rs
index dd67cc103..5d9cf83ea 100644
--- a/ballista/client/tests/context_checks.rs
+++ b/ballista/client/tests/context_checks.rs
@@ -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=...]
|",
diff --git a/ballista/client/tests/sort_shuffle.rs
b/ballista/client/tests/sort_shuffle.rs
index 0a9cc4bf7..fbeb7f052 100644
--- a/ballista/client/tests/sort_shuffle.rs
+++ b/ballista/client/tests/sort_shuffle.rs
@@ -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)]
@@ -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
@@ -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")
@@ -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(
@@ -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]
@@ -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]
diff --git a/ballista/core/src/config.rs b/ballista/core/src/config.rs
index 1ec5cadbd..10f64aaaa 100644
--- a/ballista/core/src/config.rs
+++ b/ballista/core/src/config.rs
@@ -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";
@@ -233,10 +230,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,
@@ -589,14 +582,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)
diff --git a/ballista/core/src/execution_plans/shuffle_reader.rs
b/ballista/core/src/execution_plans/shuffle_reader.rs
index c4d7b511b..d85f1978f 100644
--- a/ballista/core/src/execution_plans/shuffle_reader.rs
+++ b/ballista/core/src/execution_plans/shuffle_reader.rs
@@ -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));
@@ -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}");
@@ -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;
@@ -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
@@ -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(
@@ -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};
@@ -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();
@@ -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())
}
diff --git a/ballista/core/src/execution_plans/shuffle_writer.rs
b/ballista/core/src/execution_plans/shuffle_writer.rs
index 56d378740..0742dbf21 100644
--- a/ballista/core/src/execution_plans/shuffle_writer.rs
+++ b/ballista/core/src/execution_plans/shuffle_writer.rs
@@ -20,14 +20,9 @@
//! partition is re-partitioned and streamed to disk in Arrow IPC format.
Future stages of the query
//! will use the ShuffleReaderExec to read these results.
-use datafusion::arrow::ipc::writer::StreamWriter;
use std::fmt::Debug;
-use std::fs;
-use std::fs::File;
use std::future::Future;
-use std::io::BufWriter;
use std::iter::Iterator;
-use std::path::PathBuf;
use std::sync::Arc;
use std::time::Instant;
@@ -39,7 +34,6 @@ use crate::utils;
use crate::serde::protobuf::ShuffleWritePartition;
use crate::serde::scheduler::PartitionStats;
-use crate::utils::create_write_options;
use datafusion::arrow::array::{
ArrayBuilder, ArrayRef, StringBuilder, StructBuilder, UInt32Builder,
UInt64Builder,
};
@@ -57,14 +51,14 @@ use datafusion::physical_plan::{
DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties,
SendableRecordBatchStream, Statistics,
};
-use futures::{StreamExt, TryStreamExt};
+use futures::TryStreamExt;
use datafusion::arrow::error::ArrowError;
use datafusion::execution::context::TaskContext;
-use datafusion::physical_plan::repartition::{BatchPartitioner,
RepartitionExec};
+use datafusion::physical_plan::repartition::RepartitionExec;
use
datafusion::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
-use log::{debug, error};
+use log::debug;
use std::sync::Mutex;
use tokio::sync::oneshot;
@@ -148,12 +142,11 @@ fn walk_child_partition_mapping(
/// the wire. The executor-side writer then uses these ids directly rather
/// than re-walking the plan.
///
-/// Three cases:
+/// Two cases:
///
/// - `SortShuffleWriter(Hash(K))` — HashSpace by construction; the K-space
/// `[0..K-1]` is intrinsic. Input ids are irrelevant.
-/// - `ShuffleWriter(Hash(K)|RoundRobin(K))` — same K-space semantics.
-/// - `ShuffleWriter(None)` — passthrough; the child's plan shape decides:
+/// - `ShuffleWriter` — always passthrough; the child's plan shape decides:
/// - `SortPreservingMergeExec` in the child chain → `[0]` (collapse).
/// - `RepartitionExec(Hash|RoundRobin)` in the child chain → `[0..K-1]`.
/// - Otherwise (leaf-preserved partitioning) → the input ids as-is.
@@ -171,36 +164,19 @@ pub fn compute_global_output_partition_ids(
};
return (0..*k).collect();
}
- if let Some(w) = stage_plan.downcast_ref::<ShuffleWriterExec>() {
- match w.shuffle_output_partitioning() {
- Some(Partitioning::Hash(_, k)) |
Some(Partitioning::RoundRobinBatch(k)) => {
- return (0..*k).collect();
- }
- None => {
- let children = stage_plan.children();
- let [child] = children.as_slice() else {
- unreachable!("ShuffleWriterExec always has exactly one
child");
- };
- return match walk_child_partition_mapping(
- child,
- global_input_partition_ids,
- ) {
- GlobalPartitionMap::Collapsed => vec![0],
- GlobalPartitionMap::HashSpace => {
- let k =
-
child.properties().output_partitioning().partition_count();
- (0..k).collect()
- }
- GlobalPartitionMap::PassThrough(ids) => ids,
- };
- }
- Some(Partitioning::UnknownPartitioning(_)) => {
- panic!(
- "ShuffleWriterExec::UnknownPartitioning has no meaningful
output-id \
- space; upstream planner should never emit this shape"
- );
+ if stage_plan.downcast_ref::<ShuffleWriterExec>().is_some() {
+ let children = stage_plan.children();
+ let [child] = children.as_slice() else {
+ unreachable!("ShuffleWriterExec always has exactly one child");
+ };
+ return match walk_child_partition_mapping(child,
global_input_partition_ids) {
+ GlobalPartitionMap::Collapsed => vec![0],
+ GlobalPartitionMap::HashSpace => {
+ let k =
child.properties().output_partitioning().partition_count();
+ (0..k).collect()
}
- }
+ GlobalPartitionMap::PassThrough(ids) => ids,
+ };
}
global_input_partition_ids.to_vec()
}
@@ -240,13 +216,13 @@ impl Debug for WriterState {
/// partition is re-partitioned and streamed to disk in Arrow IPC format.
Future stages of the query
/// will use the ShuffleReaderExec to read these results.
///
-/// # Threading (passthrough / `None` branch)
+/// # Threading
///
/// One Ballista task holds one `ShuffleWriterExec`. The first `execute(k)`
/// call from the executor initializes K oneshot handoffs and spawns
-/// `run_coordinator`, which invokes `execute_shuffle_write`. In the
-/// passthrough branch (`shuffle_output_partitioning == None`, i.e. the
-/// child already has the target K partitions), `execute_shuffle_write`
+/// `run_coordinator`, which invokes `execute_shuffle_write`. The writer only
+/// ever passes its child's partitioning through (the child already has the
+/// target K partitions), so `execute_shuffle_write`
/// spawns K concurrent tokio tasks — one per output partition — each
/// pulling `child.execute(k)` and streaming directly to
/// `data-{task_id}.arrow` for partition k. Each drain emits one summary;
@@ -263,8 +239,8 @@ impl Debug for WriterState {
/// deadlock the scatter side.
///
/// Data flows bottom → top (child produces, executor consumes), matching
-/// DataFusion's convention. Under passthrough M = K (child preserves
-/// partition count). Example: K=3.
+/// DataFusion's convention. M = K (the child preserves partition count).
+/// Example: K=3.
///
/// ```text
/// K=3 output partitions (pulled by executor)
@@ -315,18 +291,16 @@ impl Debug for WriterState {
/// what differs, and downstream `ShuffleReaderExec` uses the summaries to
/// open whichever set of files is right.
///
-/// The `Hash` branch (`shuffle_output_partitioning == Some(Hash)`) is a
-/// legacy shape kept for correctness while the planner transitions to
-/// `RepartitionExec(Hash) → ShuffleWriterExec(None)`; not diagrammed here.
+/// This writer never repartitions: hash-repartition stages use
+/// [`SortShuffleWriterExec`] instead, so the only scheme left here is
+/// passthrough and the writer carries no output-partitioning of its own.
///
/// The coordinator + oneshot plumbing is a shared idiom with
-/// [`SortShuffleWriterExec`] and
-/// with the `Hash` branch; passthrough alone doesn't structurally need it —
-/// only the K concurrent drains, which are the real deadlock guard. Once
-/// `Hash` is retired in favor of DataFusion's `RepartitionExec(Hash)`, this
-/// branch could collapse to per-`execute(k)` eager K-spawn with
-/// `JoinHandle` handoff, keeping the coordinator idiom only where it does
-/// real work (SortShuffle's M×K summary re-bucketing).
+/// [`SortShuffleWriterExec`]; passthrough alone doesn't structurally need
+/// it — only the K concurrent drains, which are the real deadlock guard.
+/// This could collapse to per-`execute(k)` eager K-spawn with `JoinHandle`
+/// handoff, keeping the coordinator idiom only where it does real work
+/// (SortShuffle's M×K summary re-bucketing).
#[derive(Debug)]
pub struct ShuffleWriterExec {
/// Unique ID for the job (query) that this stage is a part of
@@ -337,9 +311,6 @@ pub struct ShuffleWriterExec {
plan: Arc<dyn ExecutionPlan>,
/// Path to write output streams to
work_dir: String,
- /// Optional shuffle output partitioning.
- /// If it's none, it means there's no need to do repartitioning.
- shuffle_output_partitioning: Option<Partitioning>,
/// Task id (the task's append-order slot in `RunningStage.task_infos`)
/// used as `file_id` in shuffle paths so files from different tasks
/// (including retries) don't collide. Seeded by the executor's
@@ -351,11 +322,10 @@ pub struct ShuffleWriterExec {
/// `global_output_partition_ids[i]` globally.
///
/// Consumed by the passthrough (`None`) branch when the plan is a straight
- /// pass-through of the slice. The `Hash` branch ignores it — the hash
- /// K-space is its own global identity — and even for `None`, if the child
- /// plan contains a partitioning-resetting operator (SPM → 1 output,
- /// RepartitionExec::Hash → 0..K) the writer detects that at path-build
- /// time and uses `local` directly instead of
`global_output_partition_ids[local]`.
+ /// pass-through of the slice. If the child plan contains a
+ /// partitioning-resetting operator (SPM → 1 output, RepartitionExec::Hash
+ /// → 0..K) the writer detects that at path-build time and uses `local`
+ /// directly instead of `global_output_partition_ids[local]`.
global_output_partition_ids: Vec<usize>,
/// Execution metrics
metrics: ExecutionPlanMetricsSet,
@@ -374,7 +344,6 @@ impl Clone for ShuffleWriterExec {
stage_id: self.stage_id,
plan: self.plan.clone(),
work_dir: self.work_dir.clone(),
- shuffle_output_partitioning:
self.shuffle_output_partitioning.clone(),
task_id: self.task_id,
global_output_partition_ids:
self.global_output_partition_ids.clone(),
metrics: self.metrics.clone(),
@@ -391,28 +360,16 @@ impl std::fmt::Display for ShuffleWriterExec {
.indent(false);
write!(
f,
- "ShuffleWriterExec: job={} stage={} work_dir={} partitioning={:?}
plan: \n {}",
- self.job_id,
- self.stage_id,
- self.work_dir,
- self.shuffle_output_partitioning,
- printable_plan
+ "ShuffleWriterExec: job={} stage={} work_dir={} plan: \n {}",
+ self.job_id, self.stage_id, self.work_dir, printable_plan
)
}
}
-pub struct WriteTracker {
- pub num_batches: usize,
- pub num_rows: usize,
- pub writer: StreamWriter<BufWriter<File>>,
- pub path: PathBuf,
-}
-
#[derive(Debug, Clone)]
struct ShuffleWriteMetrics {
/// Time spend writing batches to shuffle files
write_time: metrics::Time,
- repart_time: metrics::Time,
input_rows: metrics::Count,
output_rows: metrics::Count,
}
@@ -428,8 +385,6 @@ impl ShuffleWriteMetrics {
fn new(input_partition: usize, metrics: &ExecutionPlanMetricsSet) -> Self {
let write_time =
MetricBuilder::new(metrics).subset_time("write_time",
input_partition);
- let repart_time =
- MetricBuilder::new(metrics).subset_time("repart_time",
input_partition);
let input_rows =
MetricBuilder::new(metrics).counter("input_rows", input_partition);
@@ -438,7 +393,6 @@ impl ShuffleWriteMetrics {
Self {
write_time,
- repart_time,
input_rows,
output_rows,
}
@@ -454,13 +408,10 @@ impl ShuffleWriterExec {
stage_id: usize,
plan: Arc<dyn ExecutionPlan>,
work_dir: String,
- shuffle_output_partitioning: Option<Partitioning>,
) -> Result<Self> {
- // If [`shuffle_output_partitioning`] is none, then there's no need to
do repartitioning.
- // Therefore, the partition is the same as its input plan's.
- let partitioning = shuffle_output_partitioning
- .clone()
- .unwrap_or_else(||
plan.properties().output_partitioning().clone());
+ // This writer never repartitions, so its output partitioning is
+ // exactly its input plan's.
+ let partitioning = plan.properties().output_partitioning().clone();
let output_partition_count = partitioning.partition_count();
let properties = Arc::new(PlanProperties::new(
datafusion::physical_expr::EquivalenceProperties::new(plan.schema()),
@@ -473,15 +424,12 @@ impl ShuffleWriterExec {
// unit tests that don't set a slice. Once every construction path
// calls `with_global_output_partition_ids`, drop the default and fold
it into
// `try_new`'s signature.
- let child_partition_count =
- plan.properties().output_partitioning().partition_count();
- let default_partition_slice: Vec<usize> =
(0..child_partition_count).collect();
+ let default_partition_slice: Vec<usize> =
(0..output_partition_count).collect();
Ok(Self {
job_id,
stage_id,
plan,
work_dir,
- shuffle_output_partitioning,
task_id: 0,
global_output_partition_ids: default_partition_slice,
metrics: ExecutionPlanMetricsSet::new(),
@@ -539,28 +487,20 @@ impl ShuffleWriterExec {
.partition_count()
}
- /// Get the true output partitioning
- pub fn shuffle_output_partitioning(&self) -> Option<&Partitioning> {
- self.shuffle_output_partitioning.as_ref()
- }
-
/// Executes the shuffle write operation for this task.
///
/// Returns `(handoff_idx, summary)` pairs. `handoff_idx` indexes into the
/// coordinator's per-output-partition oneshot slots (0..K in this
/// operator's output_partitioning). `summary.partition_id` is the
- /// **global** output partition id downstream will address:
- ///
- /// - Passthrough (`None`): computed via `walk_child_partition_mapping`
- /// over the child plan — either `global_output_partition_ids[local]`,
`local` (hash
- /// K-space), or `0` (collapsed / SPM).
- /// - Hash: hash bucket space is global; global = local.
+ /// **global** output partition id downstream will address, computed via
+ /// `walk_child_partition_mapping` over the child plan — either
+ /// `global_output_partition_ids[local]`, `local` (hash K-space), or `0`
+ /// (collapsed / SPM).
pub fn execute_shuffle_write(
self,
context: Arc<TaskContext>,
) -> impl Future<Output = Result<Vec<(usize, ShuffleWritePartition)>>> {
let task_id = self.task_id;
- let output_partitioning = self.shuffle_output_partitioning.clone();
let plan = self.plan.clone();
let partition_map =
walk_child_partition_mapping(&plan,
&self.global_output_partition_ids);
@@ -572,272 +512,83 @@ impl ShuffleWriterExec {
let compression_type = config.shuffle_compression_codec()?;
let channel_capacity = config.shuffle_writer_channel_capacity();
- match output_partitioning {
- None => {
- // Passthrough shuffle: drain each of the child's output
- // partitions into its own file. All K must drain
- // CONCURRENTLY, not sequentially — coordinating operators
- // below (DynamicRangeRepartitionExec) push to all K
- // senders from shared scatter tasks; draining one to EOF
- // before the next starts fills up the undrained channel
- // and deadlocks the scatter side.
- let num_partitions =
-
plan.properties().output_partitioning().partition_count();
- let mut handles = Vec::with_capacity(num_partitions);
- for local_input_partition in 0..num_partitions {
- // Each drain owns its own metric bucket, keyed by the
- // operator-local input partition it drains.
Passthrough
- // is 1:1 so local input == local output here.
- let write_metrics =
- ShuffleWriteMetrics::new(local_input_partition,
&metrics);
- let global_partition =
- partition_map.resolve(local_input_partition) as
usize;
- let path = create_shuffle_path(
- &self.work_dir,
- &self.job_id,
- self.stage_id,
- global_partition,
- Some(task_id as u64),
- false,
- )?;
-
- if let Some(parent) = path.parent() {
- std::fs::create_dir_all(parent)?;
- }
-
- debug!("Writing results to {path:?}");
-
- let mut stream =
- plan.execute(local_input_partition,
context.clone())?;
- handles.push(tokio::spawn(async move {
- let stats = utils::write_stream_to_disk(
- &mut stream,
- path.as_path(),
- &write_metrics.write_time,
- channel_capacity,
- compression_type,
- )
- .await
- .map_err(|e|
DataFusionError::Execution(format!("{e:?}")))?;
- let rows = stats.num_rows.unwrap_or(0) as usize;
- write_metrics.input_rows.add(rows);
- write_metrics.output_rows.add(rows);
- Ok::<_, DataFusionError>((local_input_partition,
stats))
- }));
- }
-
- let mut results = Vec::with_capacity(num_partitions);
- for handle in handles {
- let (local_input_partition, stats) =
- handle.await.map_err(|e| {
- DataFusionError::Execution(format!(
- "shuffle-write drain task panicked: {e}"
- ))
- })??;
- results.push((
- local_input_partition,
- ShuffleWritePartition {
- partition_id: partition_map
- .resolve(local_input_partition),
- num_batches: stats.num_batches.unwrap_or(0),
- num_rows: stats.num_rows.unwrap_or(0),
- num_bytes: stats.num_bytes.unwrap_or(0),
- file_id: Some(task_id as u64),
- is_sort_shuffle: false,
- },
- ));
- }
- debug!(
- "task_id {} drained {} partitions in {}s",
- task_id,
- num_partitions,
- now.elapsed().as_secs()
- );
- Ok(results)
+ // Passthrough shuffle: drain each of the child's output
+ // partitions into its own file. All K must drain
+ // CONCURRENTLY, not sequentially — coordinating operators
+ // below (DynamicRangeRepartitionExec) push to all K
+ // senders from shared scatter tasks; draining one to EOF
+ // before the next starts fills up the undrained channel
+ // and deadlocks the scatter side.
+ let num_partitions =
+ plan.properties().output_partitioning().partition_count();
+ let mut handles = Vec::with_capacity(num_partitions);
+ for local_input_partition in 0..num_partitions {
+ // Each drain owns its own metric bucket, keyed by the
+ // operator-local input partition it drains. Passthrough
+ // is 1:1 so local input == local output here.
+ let write_metrics =
+ ShuffleWriteMetrics::new(local_input_partition, &metrics);
+ let global_partition =
+ partition_map.resolve(local_input_partition) as usize;
+ let path = create_shuffle_path(
+ &self.work_dir,
+ &self.job_id,
+ self.stage_id,
+ global_partition,
+ Some(task_id as u64),
+ false,
+ )?;
+
+ if let Some(parent) = path.parent() {
+ std::fs::create_dir_all(parent)?;
}
- // TODO: just use Datafusion's hash partition operator and
remove this branch
- Some(Partitioning::Hash(exprs, num_output_partitions)) => {
- // Task drains ALL of the child's (already
slice-restricted)
- // input partitions and routes rows into K writers by hash.
- // file_id = task_id so files from different tasks
- // (including retries) don't collide. Hash buckets 0..K are
- // a global K-space and aren't further sliced.
- //
- // One metric bucket per operator-local input partition,
- // matching the pre-K-drain model (before: N tasks × 1
- // bucket; now: 1 task × N buckets). One `BatchPartitioner`
- // per input partition so `repart_time` disaggregates the
- // same way — total partitioner count across the stage is
- // unchanged.
- let num_input_partitions =
-
plan.properties().output_partitioning().partition_count();
- let write_metrics_per_input: Vec<ShuffleWriteMetrics> = (0
- ..num_input_partitions)
- .map(|i| ShuffleWriteMetrics::new(i, &metrics))
- .collect();
- let schema = plan.schema();
- let (tx, mut rx) = tokio::sync::mpsc::channel::<(usize,
RecordBatch)>(
+
+ debug!("Writing results to {path:?}");
+
+ let mut stream = plan.execute(local_input_partition,
context.clone())?;
+ handles.push(tokio::spawn(async move {
+ let stats = utils::write_stream_to_disk(
+ &mut stream,
+ path.as_path(),
+ &write_metrics.write_time,
channel_capacity,
- );
- let write_times: Vec<metrics::Time> =
write_metrics_per_input
- .iter()
- .map(|m| m.write_time.clone())
- .collect();
- let output_rows_per_input: Vec<metrics::Count> =
- write_metrics_per_input
- .iter()
- .map(|m| m.output_rows.clone())
- .collect();
- let partitioners: Vec<BatchPartitioner> =
write_metrics_per_input
- .iter()
- .map(|m| {
- BatchPartitioner::new_hash_partitioner(
- exprs.clone(),
- num_output_partitions,
- m.repart_time.clone(),
- )
- })
- .collect::<Result<Vec<_>, _>>()?;
- let work_dir = self.work_dir.clone();
- let job_id = self.job_id.clone();
- let stage_id = self.stage_id;
-
- let handle = tokio::task::spawn_blocking(move || {
- let mut writers: Vec<Option<WriteTracker>> =
- (0..num_output_partitions).map(|_| None).collect();
- let mut partitioners = partitioners;
-
- while let Some((local_input_partition, input_batch)) =
- rx.blocking_recv()
- {
- let write_time =
&write_times[local_input_partition];
- let output_rows =
- &output_rows_per_input[local_input_partition];
- partitioners[local_input_partition].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 p = create_shuffle_path(
- &work_dir,
- &job_id,
- stage_id,
- output_partition,
- Some(task_id as u64),
- false,
- )?;
-
- if let Some(parent) = p.parent() {
-
std::fs::create_dir_all(parent)?;
- }
-
- debug!("Writing results to {p:?}");
-
- let options =
-
create_write_options(compression_type)?;
-
- 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,
- });
- }
- }
- 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((
- i,
- ShuffleWritePartition {
- partition_id: i as u64,
- num_batches: w.num_batches as u64,
- num_rows: w.num_rows as u64,
- num_bytes,
- file_id: Some(task_id as u64),
- is_sort_shuffle: false,
- },
- ));
- }
- }
- Ok(part_locs)
- });
-
- let mut stream_err = None;
- 'outer: for (local_input_partition, per_input_metrics) in
- write_metrics_per_input.iter().enumerate()
- {
- let mut stream =
- plan.execute(local_input_partition,
context.clone())?;
- loop {
- match stream.next().await {
- Some(Ok(batch)) => {
-
per_input_metrics.input_rows.add(batch.num_rows());
- if tx
- .send((local_input_partition, batch))
- .await
- .is_err()
- {
- break 'outer;
- }
- }
- Some(Err(e)) => {
- stream_err = Some(e);
- break 'outer;
- }
- None => break,
- }
- }
- }
- drop(tx);
-
- let write_result = handle.await.map_err(|e| {
- DataFusionError::Execution(format!(
- "Shuffle writer task failed: {e}"
- ))
- })?;
- if let Some(e) = stream_err {
- if let Err(write_err) = &write_result {
- error!("Shuffle writer also failed: {write_err}");
- }
- return Err(e);
- }
- write_result
- }
+ compression_type,
+ )
+ .await
+ .map_err(|e|
DataFusionError::Execution(format!("{e:?}")))?;
+ let rows = stats.num_rows.unwrap_or(0) as usize;
+ write_metrics.input_rows.add(rows);
+ write_metrics.output_rows.add(rows);
+ Ok::<_, DataFusionError>((local_input_partition, stats))
+ }));
+ }
- _ => Err(DataFusionError::Execution(
- "Invalid shuffle partitioning scheme".to_owned(),
- )),
+ let mut results = Vec::with_capacity(num_partitions);
+ for handle in handles {
+ let (local_input_partition, stats) = handle.await.map_err(|e| {
+ DataFusionError::Execution(format!(
+ "shuffle-write drain task panicked: {e}"
+ ))
+ })??;
+ results.push((
+ local_input_partition,
+ ShuffleWritePartition {
+ partition_id:
partition_map.resolve(local_input_partition),
+ num_batches: stats.num_batches.unwrap_or(0),
+ num_rows: stats.num_rows.unwrap_or(0),
+ num_bytes: stats.num_bytes.unwrap_or(0),
+ file_id: Some(task_id as u64),
+ is_sort_shuffle: false,
+ },
+ ));
}
+ debug!(
+ "task_id {} drained {} partitions in {}s",
+ task_id,
+ num_partitions,
+ now.elapsed().as_secs()
+ );
+ Ok(results)
}
}
}
@@ -849,25 +600,13 @@ impl DisplayAs for ShuffleWriterExec {
f: &mut std::fmt::Formatter,
) -> std::fmt::Result {
match t {
+ // "None" is retained for plan-shape stability: this writer never
+ // repartitions, so the value can only ever be None.
DisplayFormatType::Default | DisplayFormatType::Verbose => {
- write!(
- f,
- "ShuffleWriterExec: partitioning: {}",
- self.shuffle_output_partitioning
- .as_ref()
- .map(|p| p.to_string())
- .unwrap_or("None".to_string())
- )
+ write!(f, "ShuffleWriterExec: partitioning: None")
}
DisplayFormatType::TreeRender => {
- write!(
- f,
- "partitioning={}",
- self.shuffle_output_partitioning
- .as_ref()
- .map(|p| p.to_string())
- .unwrap_or("None".to_string())
- )
+ write!(f, "partitioning=None")
}
}
}
@@ -907,7 +646,6 @@ impl ExecutionPlan for ShuffleWriterExec {
self.stage_id,
input,
self.work_dir.clone(),
- self.shuffle_output_partitioning.clone(),
)?
.with_task_id(self.task_id)
.with_global_output_partition_ids(
@@ -1020,8 +758,9 @@ impl ShuffleWriter for ShuffleWriterExec {
self.stage_id
}
+ /// Always `None`: this writer preserves its input partitioning.
fn shuffle_output_partitioning(&self) -> Option<&Partitioning> {
- self.shuffle_output_partitioning.as_ref()
+ None
}
fn input_partition_count(&self) -> usize {
@@ -1217,25 +956,23 @@ mod tests {
}
#[tokio::test]
- // number of rows in each partition is a function of the hash output, so
don't test here
- #[cfg(not(feature = "force_hash_collisions"))]
async fn test() -> Result<()> {
let session_ctx = SessionContext::new();
let task_ctx = session_ctx.task_ctx();
- let input_plan =
Arc::new(CoalescePartitionsExec::new(create_input_plan()?));
+ // No output partitioning: passthrough writer, one file per one of
+ // the input plan's 2 partitions.
+ let input_plan = create_input_plan()?;
let work_dir = TempDir::new()?;
let query_stage = Arc::new(ShuffleWriterExec::try_new(
JobId::new("jobOne"),
1,
input_plan,
work_dir.path().to_str().unwrap().to_owned(),
- Some(Partitioning::Hash(vec![Arc::new(Column::new("a", 0))], 2)),
)?);
let batches = drive_all_partitions(query_stage, task_ctx).await?;
// K=2 output partitions -> one metadata batch per execute(N) call
- // (drive_all_partitions collects one per K). Empty output slots
- // now emit 0-row batches (no summary), non-empty ones emit 1-row.
+ // (drive_all_partitions collects one per K).
assert_eq!(2, batches.len());
for batch in &batches {
assert_eq!(4, batch.num_columns());
@@ -1275,7 +1012,6 @@ mod tests {
}
#[tokio::test]
- #[cfg(not(feature = "force_hash_collisions"))]
async fn display_renders_child_operator_metrics() -> Result<()> {
let session_ctx = SessionContext::new();
let task_ctx = session_ctx.task_ctx();
@@ -1287,7 +1023,6 @@ mod tests {
1,
input_plan,
work_dir.path().to_str().unwrap().to_owned(),
- Some(Partitioning::Hash(vec![Arc::new(Column::new("a", 0))], 2)),
)?;
let mut stream = query_stage.execute(0, task_ctx)?;
let _ = utils::collect_stream(&mut stream)
@@ -1306,45 +1041,6 @@ mod tests {
Ok(())
}
- #[tokio::test]
- // number of rows in each partition is a function of the hash output, so
don't test here
- #[cfg(not(feature = "force_hash_collisions"))]
- async fn test_partitioned() -> Result<()> {
- let session_ctx = SessionContext::new();
- let task_ctx = session_ctx.task_ctx();
-
- let input_plan = create_input_plan()?;
- let work_dir = TempDir::new()?;
- let query_stage = Arc::new(ShuffleWriterExec::try_new(
- JobId::new("jobOne"),
- 1,
- input_plan,
- work_dir.path().to_str().unwrap().to_owned(),
- Some(Partitioning::Hash(vec![Arc::new(Column::new("a", 0))], 2)),
- )?);
- let batches = drive_all_partitions(query_stage, task_ctx).await?;
- assert_eq!(2, batches.len());
- let total: u64 = batches
- .iter()
- .flat_map(|b| {
- let stats =
b.column(3).as_any().downcast_ref::<StructArray>().unwrap();
- let num_rows = stats
- .column_by_name("num_rows")
- .unwrap()
- .as_any()
- .downcast_ref::<UInt64Array>()
- .unwrap()
- .clone();
- (0..b.num_rows()).map(move |i| num_rows.value(i))
- })
- .sum();
- // Row conservation across the K output partitions: writer reads its
- // full input slice (2 partitions × 2 batches × 2 rows = 8).
- assert_eq!(8, total);
-
- Ok(())
- }
-
#[tokio::test]
async fn test_no_repart_write_failure_propagates() -> Result<()> {
let session_ctx = SessionContext::new();
@@ -1370,7 +1066,6 @@ mod tests {
1,
input_plan,
work_dir,
- None,
)?);
let result = drive_all_partitions(query_stage, task_ctx).await;
assert!(
@@ -1381,13 +1076,13 @@ mod tests {
}
#[tokio::test]
- async fn test_hash_repart_write_failure_propagates() -> Result<()> {
+ async fn test_create_dir_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.
+ // create_dir_all fails when the writer tries to create the
+ // output partition subdirectory underneath it.
// Path structure: work_dir / job_id / stage_id / ...
let tmp = tempfile::TempDir::new().unwrap();
let job_dir = tmp.path().join("jobOne");
@@ -1402,12 +1097,11 @@ mod tests {
1,
input_plan,
work_dir,
- Some(Partitioning::Hash(vec![Arc::new(Column::new("a", 0))], 2)),
)?);
let result = drive_all_partitions(query_stage, task_ctx).await;
assert!(
result.is_err(),
- "expected create_dir_all failure in hash writer to propagate"
+ "expected create_dir_all failure in writer to propagate"
);
Ok(())
}
diff --git a/ballista/core/src/extension.rs b/ballista/core/src/extension.rs
index 28fd77e65..31a9e12af 100644
--- a/ballista/core/src/extension.rs
+++ b/ballista/core/src/extension.rs
@@ -261,9 +261,6 @@ pub trait SessionConfigExt {
/// Get whether to use TLS for executor connections
fn ballista_use_tls(&self) -> bool;
- /// Is short shuffle used
- fn ballista_sort_shuffle_enabled(&self) -> bool;
-
/// Returns whether the AQE coalesce-shuffle-partitions rule is enabled.
fn ballista_coalesce_enabled(&self) -> bool;
/// Sets whether the AQE coalesce-shuffle-partitions rule is enabled.
@@ -537,14 +534,6 @@ impl SessionConfigExt for SessionConfig {
})
}
- fn ballista_sort_shuffle_enabled(&self) -> bool {
- self.options()
- .extensions
- .get::<BallistaConfig>()
- .map(|c| c.shuffle_sort_based_enabled())
- .unwrap_or_else(||
BallistaConfig::default().shuffle_sort_based_enabled())
- }
-
fn with_ballista_shuffle_reader_maximum_concurrent_requests(
self,
max_requests: usize,
diff --git a/ballista/core/src/serde/mod.rs b/ballista/core/src/serde/mod.rs
index a22bb977d..40df51bd4 100644
--- a/ballista/core/src/serde/mod.rs
+++ b/ballista/core/src/serde/mod.rs
@@ -389,19 +389,22 @@ impl PhysicalExtensionCodec for
BallistaPhysicalExtensionCodec {
PhysicalPlanType::ShuffleWriter(shuffle_writer) => {
let input = inputs[0].clone();
- let shuffle_output_partitioning =
parse_protobuf_hash_partitioning(
- shuffle_writer.output_partitioning.as_ref(),
- &decode_ctx,
- input.schema().as_ref(),
- &converter,
- )?;
+ // ShuffleWriterExec never repartitions. A plan that still
+ // carries an output partitioning is a legacy hash-shuffle
+ // plan, which this writer no longer implements.
+ if shuffle_writer.output_partitioning.is_some() {
+ return Err(DataFusionError::Internal(
+ "hash-partitioned ShuffleWriterExec is no longer
supported; \
+ hash-repartition stages use SortShuffleWriterExec"
+ .to_string(),
+ ));
+ }
Ok(Arc::new(ShuffleWriterExec::try_new(
shuffle_writer.job_id.clone().into(),
shuffle_writer.stage_id as usize,
input,
"".to_string(), // this is intentional but hacky - the
executor will fill this in
- shuffle_output_partitioning,
)?))
}
PhysicalPlanType::SortShuffleWriter(sort_shuffle_writer) => {
@@ -603,33 +606,15 @@ impl PhysicalExtensionCodec for
BallistaPhysicalExtensionCodec {
buf: &mut Vec<u8>,
) -> Result<(), DataFusionError> {
if let Some(exec) = node.downcast_ref::<ShuffleWriterExec>() {
- // note that we use shuffle_output_partitioning() rather than
output_partitioning()
- // to get the true output partitioning
- let output_partitioning = match exec.shuffle_output_partitioning()
{
- Some(Partitioning::Hash(exprs, partition_count)) => {
- Some(datafusion_proto::protobuf::PhysicalHashRepartition {
- hash_expr: exprs
- .iter()
-
.map(|expr|datafusion_proto::physical_plan::to_proto::serialize_physical_expr(&expr.clone(),
self.default_codec.as_ref()))
- .collect::<Result<Vec<_>, DataFusionError>>()?,
- partition_count: *partition_count as u64,
- })
- }
- None => None,
- other => {
- return Err(DataFusionError::Internal(format!(
- "physical_plan::to_proto() invalid partitioning for
ShuffleWriterExec: {other:?}"
- )));
- }
- };
-
let proto = protobuf::BallistaPhysicalPlanNode {
physical_plan_type: Some(PhysicalPlanType::ShuffleWriter(
protobuf::ShuffleWriterExecNode {
job_id: exec.job_id().to_string(),
stage_id: exec.stage_id() as u32,
input: None,
- output_partitioning,
+ // This writer preserves its input partitioning, so
+ // there is never a shuffle partitioning to encode.
+ output_partitioning: None,
},
)),
};
diff --git a/ballista/executor/src/execution_engine.rs
b/ballista/executor/src/execution_engine.rs
index 2a83e762b..7f779d2c8 100644
--- a/ballista/executor/src/execution_engine.rs
+++ b/ballista/executor/src/execution_engine.rs
@@ -149,18 +149,17 @@ impl ExecutionEngine for DefaultExecutionEngine {
// the query plan created by the scheduler always starts with a
shuffle writer
// (either ShuffleWriterExec or SortShuffleWriterExec)
- if let Some(shuffle_writer) = plan.downcast_ref::<ShuffleWriterExec>()
{
+ if plan.downcast_ref::<ShuffleWriterExec>().is_some() {
let exec = ShuffleWriterExec::try_new(
job_id,
stage_id,
plan.children()[0].clone(),
work_dir.to_string(),
- shuffle_writer.shuffle_output_partitioning().cloned(),
)?
.with_task_id(task_id)
.with_global_output_partition_ids(global_output_partition_ids);
Ok(Arc::new(DefaultQueryStageExec::new(
- ShuffleWriterVariant::Hash(exec),
+ ShuffleWriterVariant::Passthrough(exec),
)))
} else if let Some(sort_shuffle_writer) =
plan.downcast_ref::<SortShuffleWriterExec>()
@@ -190,8 +189,9 @@ impl ExecutionEngine for DefaultExecutionEngine {
/// Enum representing the different shuffle writer implementations.
#[derive(Debug, Clone)]
pub enum ShuffleWriterVariant {
- /// Hash-based shuffle writer (original implementation).
- Hash(ShuffleWriterExec),
+ /// Passthrough shuffle writer: preserves its input partitioning,
+ /// one file per output partition.
+ Passthrough(ShuffleWriterExec),
/// Sort-based shuffle writer.
Sort(SortShuffleWriterExec),
}
@@ -216,7 +216,7 @@ impl DefaultQueryStageExec {
impl Display for DefaultQueryStageExec {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match &self.shuffle_writer {
- ShuffleWriterVariant::Hash(writer) => {
+ ShuffleWriterVariant::Passthrough(writer) => {
let stage_metrics: Vec<String> = writer
.metrics()
.unwrap_or_default()
@@ -225,7 +225,7 @@ impl Display for DefaultQueryStageExec {
.collect();
write!(
f,
- "DefaultQueryStageExec(Hash): ({})\n{}",
+ "DefaultQueryStageExec(Passthrough): ({})\n{}",
stage_metrics.join(", "),
writer
)
@@ -257,7 +257,9 @@ impl QueryStageExecutor for DefaultQueryStageExec {
) -> Result<Vec<ShuffleWritePartition>> {
let (plan_arc, is_sort_shuffle): (Arc<dyn ExecutionPlan>, bool) =
match &self.shuffle_writer {
- ShuffleWriterVariant::Hash(writer) =>
(Arc::new(writer.clone()), false),
+ ShuffleWriterVariant::Passthrough(writer) => {
+ (Arc::new(writer.clone()), false)
+ }
ShuffleWriterVariant::Sort(writer) =>
(Arc::new(writer.clone()), true),
};
debug!(
@@ -282,7 +284,9 @@ impl QueryStageExecutor for DefaultQueryStageExec {
fn collect_plan_metrics(&self) -> Vec<MetricsSet> {
match &self.shuffle_writer {
- ShuffleWriterVariant::Hash(writer) =>
utils::collect_plan_metrics(writer),
+ ShuffleWriterVariant::Passthrough(writer) => {
+ utils::collect_plan_metrics(writer)
+ }
ShuffleWriterVariant::Sort(writer) =>
utils::collect_plan_metrics(writer),
}
}
diff --git a/ballista/executor/src/executor.rs
b/ballista/executor/src/executor.rs
index 3a1c103be..dce01ae7f 100644
--- a/ballista/executor/src/executor.rs
+++ b/ballista/executor/src/executor.rs
@@ -424,12 +424,11 @@ mod test {
1,
Arc::new(NeverendingOperator::new()),
work_dir.clone(),
- None,
)
.expect("creating shuffle writer");
let query_stage_exec =
-
DefaultQueryStageExec::new(ShuffleWriterVariant::Hash(shuffle_write));
+
DefaultQueryStageExec::new(ShuffleWriterVariant::Passthrough(shuffle_write));
let executor_registration = ExecutorRegistration {
id: "executor".to_string(),
diff --git a/ballista/executor/src/flight_service.rs
b/ballista/executor/src/flight_service.rs
index 32a0db404..3e3c8b079 100644
--- a/ballista/executor/src/flight_service.rs
+++ b/ballista/executor/src/flight_service.rs
@@ -143,7 +143,7 @@ impl FlightService for BallistaFlightService {
));
}
- // Standard hash-based shuffle - read the entire file
+ // Standard single-file shuffle output - read the entire file
let file = File::open(&path)
.map_err(|e| {
BallistaError::General(format!(
diff --git a/ballista/scheduler/src/planner.rs
b/ballista/scheduler/src/planner.rs
index 0128540c5..728355899 100644
--- a/ballista/scheduler/src/planner.rs
+++ b/ballista/scheduler/src/planner.rs
@@ -700,16 +700,17 @@ pub(crate) fn create_shuffle_writer_with_config(
) -> Result<Arc<dyn ShuffleWriter>> {
let plan = make_empty_exec_serde_safe(plan)?;
- // Check if sort-based shuffle is enabled
+ // Sort-based shuffle is the only shuffle writer for hash-repartition
+ // stages. Its tuning values still come from the session config.
let ballista_config = config
.extensions
.get::<BallistaConfig>()
.cloned()
.unwrap_or_default();
- if ballista_config.shuffle_sort_based_enabled() {
- // Sort shuffle requires hash partitioning
- if let Some(Partitioning::Hash(exprs, partition_count)) = partitioning
{
+ // Sort shuffle requires hash partitioning.
+ match partitioning {
+ Some(Partitioning::Hash(exprs, partition_count)) => {
let sort_config = SortShuffleConfig::new(
true,
ballista_config.shuffle_sort_based_batch_size(),
@@ -718,25 +719,27 @@ pub(crate) fn create_shuffle_writer_with_config(
ballista_config.shuffle_sort_based_memory_limit_per_task_bytes(),
);
- return Ok(Arc::new(SortShuffleWriterExec::try_new(
+ Ok(Arc::new(SortShuffleWriterExec::try_new(
job_id.to_owned(),
stage_id,
plan,
"".to_owned(),
Partitioning::Hash(exprs, partition_count),
sort_config,
- )?));
+ )?))
}
+ // Stages that don't repartition write their input partitioning
+ // through: one file per output partition.
+ None => Ok(Arc::new(ShuffleWriterExec::try_new(
+ job_id.to_owned(),
+ stage_id,
+ plan,
+ "".to_owned(),
+ )?)),
+ Some(other) => Err(BallistaError::General(format!(
+ "unsupported shuffle output partitioning: {other}"
+ ))),
}
-
- // Fall back to standard shuffle writer
- Ok(Arc::new(ShuffleWriterExec::try_new(
- job_id.to_owned(),
- stage_id,
- plan,
- "".to_owned(),
- partitioning,
- )?))
}
#[cfg(test)]
diff --git a/ballista/scheduler/src/state/aqe/test/mod.rs
b/ballista/scheduler/src/state/aqe/test/mod.rs
index e4bffbbb7..52f173d19 100644
--- a/ballista/scheduler/src/state/aqe/test/mod.rs
+++ b/ballista/scheduler/src/state/aqe/test/mod.rs
@@ -30,8 +30,6 @@ mod plan_to_stages;
/// A table whose statistics are declared rather than measured
mod stats_table;
-use ballista_core::config::BALLISTA_SHUFFLE_SORT_BASED_ENABLED;
-use ballista_core::extension::SessionConfigExt;
use ballista_core::serde::scheduler::{
ExecutorMetadata, ExecutorOperatingSystemSpecification,
ExecutorSpecification,
PartitionId, PartitionLocation, PartitionStats,
@@ -135,17 +133,3 @@ pub(crate) fn mock_context() -> SessionContext {
SessionContext::new_with_state(state)
}
-
-pub(crate) fn mock_context_sort_shuffle() -> SessionContext {
- let config = SessionConfig::new_with_ballista()
- .set_str(BALLISTA_SHUFFLE_SORT_BASED_ENABLED, "true")
- .with_target_partitions(2)
- .with_round_robin_repartition(false);
-
- let state = SessionStateBuilder::new()
- .with_config(config)
- .with_default_features()
- .build();
-
- SessionContext::new_with_state(state)
-}
diff --git a/ballista/scheduler/src/state/aqe/test/plan_to_stages.rs
b/ballista/scheduler/src/state/aqe/test/plan_to_stages.rs
index b115f273c..993996af4 100644
--- a/ballista/scheduler/src/state/aqe/test/plan_to_stages.rs
+++ b/ballista/scheduler/src/state/aqe/test/plan_to_stages.rs
@@ -19,8 +19,7 @@ use crate::assert_plan;
use crate::state::aqe::execution_plan::ExchangeExec;
use crate::state::aqe::planner::AdaptivePlanner;
use crate::state::aqe::test::{
- mock_batch, mock_context, mock_context_sort_shuffle, mock_memory_table,
- mock_partitions_with_statistics,
+ mock_batch, mock_context, mock_memory_table,
mock_partitions_with_statistics,
};
use ballista_core::execution_plans::SortShuffleWriterExec;
use datafusion::arrow::datatypes::{DataType, Field, Schema};
@@ -455,36 +454,6 @@ async fn should_ignore_inactive_stages() ->
datafusion::error::Result<()> {
Ok(())
}
-#[tokio::test]
-async fn should_use_sort_shuffle_when_enabled() ->
datafusion::error::Result<()> {
- let ctx = mock_context_sort_shuffle();
- ctx.register_batch("t", mock_batch()?)?;
-
- let q = r#"
- select min(a) as c0, max(b) as c1, c as c2 from t group by c
- "#;
-
- let plan = ctx.sql(q).await?.create_physical_plan().await?;
- let mut planner = AdaptivePlanner::try_from_plan(
- ctx.state().config(),
- plan,
- "test_job".to_owned(),
- )?;
-
- let stages = planner.runnable_stages()?.unwrap();
- assert_eq!(1, stages.len());
-
- let plan = stages.first().unwrap().plan.as_ref();
- assert!(
- (plan as &dyn ExecutionPlan)
- .downcast_ref::<SortShuffleWriterExec>()
- .is_some(),
- "expected SortShuffleWriterExec when sort shuffle is enabled, got
plan: {plan:?}"
- );
-
- Ok(())
-}
-
#[tokio::test]
async fn should_use_sort_shuffle_by_default() -> datafusion::error::Result<()>
{
let ctx = mock_context();
diff --git a/benchmarks/src/bin/shuffle_bench.rs
b/benchmarks/src/bin/shuffle_bench.rs
index 8fd9184d2..e2bdbcbf3 100644
--- a/benchmarks/src/bin/shuffle_bench.rs
+++ b/benchmarks/src/bin/shuffle_bench.rs
@@ -17,14 +17,13 @@
//! Standalone shuffle benchmark for profiling Ballista shuffle write
//! performance outside of a cluster. Streams input from Parquet files and
-//! drives either the hash-based or sort-based shuffle writer end-to-end.
+//! drives the sort-based shuffle writer end-to-end.
//!
//! # Usage
//!
//! ```sh
//! cargo run --release --bin shuffle_bench -- \
//! --input /data/tpch-sf100/lineitem/ \
-//! --writer sort \
//! --partitions 200 \
//! --hash-columns 0,3
//! ```
@@ -33,10 +32,9 @@
//! ```sh
//! cargo flamegraph --release --bin shuffle_bench -- \
//! --input /data/tpch-sf100/lineitem/ \
-//! --writer sort --partitions 200
+//! --partitions 200
//! ```
-use ballista_core::execution_plans::ShuffleWriterExec;
use ballista_core::execution_plans::sort_shuffle::{
SortShuffleConfig, SortShuffleWriterExec,
};
@@ -70,12 +68,8 @@ struct Args {
#[arg(long)]
input: PathBuf,
- /// Shuffle writer to drive: `hash` (default) or `sort`.
- #[arg(long, default_value = "hash")]
- writer: String,
-
/// Partitioning scheme: `hash`, `single`, or `round-robin`. Currently
- /// both writers only support `hash`; other values are rejected.
+ /// the writer only supports `hash`; other values are rejected.
#[arg(long, default_value = "hash")]
partitioning: String,
@@ -117,27 +111,11 @@ struct Args {
vcores: usize,
}
-#[derive(Clone, Copy, Debug)]
-enum WriterKind {
- Hash,
- Sort,
-}
-
#[derive(Clone, Copy, Debug)]
enum PartitioningKind {
Hash,
}
-fn parse_writer(s: &str) -> Result<WriterKind, String> {
- match s.to_lowercase().as_str() {
- "hash" => Ok(WriterKind::Hash),
- "sort" => Ok(WriterKind::Sort),
- other => Err(format!(
- "unknown writer: {other} (expected 'hash' or 'sort')"
- )),
- }
-}
-
fn parse_partitioning(s: &str) -> Result<PartitioningKind, String> {
match s.to_lowercase().as_str() {
"hash" => Ok(PartitioningKind::Hash),
@@ -227,7 +205,6 @@ fn build_partitioning(
async fn execute_shuffle_write(
args: &Args,
- writer_kind: WriterKind,
partitioning_kind: PartitioningKind,
hash_col_indices: &[usize],
work_dir: PathBuf,
@@ -270,43 +247,25 @@ async fn execute_shuffle_write(
let work_dir_str = work_dir.to_str().unwrap().to_string();
fs::create_dir_all(&work_dir).expect("create work dir");
- let metrics: MetricsSet = match writer_kind {
- WriterKind::Hash => {
- let exec = ShuffleWriterExec::try_new(
- format!("bench_job_{task_id}").into(),
- 1,
- input,
- work_dir_str,
- Some(partitioning),
- )?;
- let task_ctx = ctx.task_ctx();
- let mut stream = exec.execute(0, task_ctx)?;
- let _ = utils::collect_stream(&mut stream).await;
- exec.metrics().unwrap_or_default()
- }
- WriterKind::Sort => {
- let cfg = SortShuffleConfig::new(true, args.batch_size);
- let exec = SortShuffleWriterExec::try_new(
- format!("bench_job_{task_id}").into(),
- 1,
- input,
- work_dir_str,
- partitioning,
- cfg,
- )?;
- let task_ctx = ctx.task_ctx();
- let mut stream = exec.execute(0, task_ctx)?;
- let _ = utils::collect_stream(&mut stream).await;
- exec.metrics().unwrap_or_default()
- }
- };
+ let cfg = SortShuffleConfig::new(true, args.batch_size);
+ let exec = SortShuffleWriterExec::try_new(
+ format!("bench_job_{task_id}").into(),
+ 1,
+ input,
+ work_dir_str,
+ partitioning,
+ cfg,
+ )?;
+ let task_ctx = ctx.task_ctx();
+ let mut stream = exec.execute(0, task_ctx)?;
+ let _ = utils::collect_stream(&mut stream).await;
+ let metrics: MetricsSet = exec.metrics().unwrap_or_default();
Ok(metrics)
}
fn run_iteration(
args: &Args,
- writer_kind: WriterKind,
partitioning_kind: PartitioningKind,
hash_col_indices: &[usize],
) -> (f64, Option<MetricsSet>) {
@@ -317,7 +276,6 @@ fn run_iteration(
let work_dir = args.output_dir.join("task_0");
let metrics = execute_shuffle_write(
args,
- writer_kind,
partitioning_kind,
hash_col_indices,
work_dir.clone(),
@@ -337,7 +295,6 @@ fn run_iteration(
handles.push(tokio::spawn(async move {
let m = execute_shuffle_write(
&args,
- writer_kind,
partitioning_kind,
&hash_col_indices,
work_dir.clone(),
@@ -433,10 +390,6 @@ fn main() {
.init();
let args = Args::parse();
- let writer_kind = parse_writer(&args.writer).unwrap_or_else(|e| {
- eprintln!("error: {e}");
- std::process::exit(2);
- });
let partitioning_kind =
parse_partitioning(&args.partitioning).unwrap_or_else(|e| {
eprintln!("error: {e}");
std::process::exit(2);
@@ -448,7 +401,6 @@ fn main() {
let (schema, total_rows) = read_parquet_metadata(&args.input, args.limit);
println!("=== Ballista Shuffle Benchmark ===");
- println!("Writer: {writer_kind:?}");
println!("Partitioning: {partitioning_kind:?}");
println!("Input: {}", args.input.display());
println!(
@@ -483,7 +435,7 @@ fn main() {
format!("iter {}/{}", i - args.warmup + 1, args.iterations)
};
let (elapsed, metrics) =
- run_iteration(&args, writer_kind, partitioning_kind,
&hash_col_indices);
+ run_iteration(&args, partitioning_kind, &hash_col_indices);
if !is_warmup {
times.push(elapsed);
if metrics.is_some() {
diff --git a/benchmarks/src/bin/tpch.rs b/benchmarks/src/bin/tpch.rs
index 2650f3401..53e1fa0f0 100644
--- a/benchmarks/src/bin/tpch.rs
+++ b/benchmarks/src/bin/tpch.rs
@@ -127,7 +127,7 @@ struct BallistaBenchmarkOpt {
/// Configuration overrides in key=value format.
/// Can be specified multiple times, e.g.
- /// -c ballista.shuffle.sort_based.enabled=true
+ /// -c ballista.shuffle.sort_based.batch_size=8192
/// -c datafusion.execution.target_partitions=16
#[structopt(short = "c", long = "config", number_of_values = 1)]
config_overrides: Vec<String>,
diff --git a/docs/source/user-guide/configs.md
b/docs/source/user-guide/configs.md
index 83f0351db..8bc7276f4 100644
--- a/docs/source/user-guide/configs.md
+++ b/docs/source/user-guide/configs.md
@@ -92,14 +92,13 @@ let expected = [
The following session-level keys control Ballista's shuffle behavior. See
the [tuning guide](tuning-guide.md#shuffle-implementation) for an
-explanation of the sort-based (default) and hash-based shuffle writers.
+explanation of the sort-based shuffle writer.
| key | type | default
| description
|
| ------------------------------------------------------- | ------- |
--------- |
------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
|
| ballista.shuffle.max_concurrent_read_requests | UInt64 | 64
| Maximum number of concurrent fetch requests the shuffle reader will issue.
|
| ballista.shuffle.force_remote_read | Boolean | false
| Forces the shuffle reader to fetch every partition through Arrow Flight,
even when the data is local. Intended for testing.
|
| ballista.shuffle.remote_read_prefer_flight | Boolean | false
| For remote reads, prefer the Arrow Flight reader over the block reader. The
block reader is generally faster.
|
-| ballista.shuffle.sort_based.enabled | Boolean | true
| Enables the sort-based shuffle writer (consolidated data file per input
partition with an index, instead of one file per (input partition, output
partition) pair).
|
| ballista.shuffle.sort_based.batch_size | UInt64 | 8192
| Target row count when coalescing buffered batches before they are written or
spilled.
|
| ballista.shuffle.sort_based.memory_limit_per_task_bytes | UInt64 |
268435456 | Per-task buffered-bytes budget at which the sort-based writer
spills to disk (256 MiB default), counted independently of the runtime memory
pool. Set to `0` to disable the per-task budget and spill only under
memory-pool pressure — safe only with a bounded memory pool, otherwise the
writer never spills and may run out of memory. |
diff --git a/docs/source/user-guide/tuning-guide.md
b/docs/source/user-guide/tuning-guide.md
index b9f3d5d7f..5391c399d 100644
--- a/docs/source/user-guide/tuning-guide.md
+++ b/docs/source/user-guide/tuning-guide.md
@@ -127,11 +127,12 @@ scheduler or executors.
Ballista exchanges data between query stages by writing the output of each
upstream task to local files, which downstream tasks read either from disk
-(when co-located) or over Arrow Flight. Two shuffle implementations are
-available, with different trade-offs around file count, memory use, and
-write latency.
+(when co-located) or over Arrow Flight. Ballista uses a sort-based shuffle
+writer that bounds file count and memory use. Single-partition stages (a
+query's final output, coalesce, and broadcast-build stages) are written
+directly to one file.
-### Sort-based shuffle (default)
+### Sort-based shuffle
The sort-based writer accumulates incoming batches in memory, tracking each
row's output partition. It spills the buffered batches to disk when either of
@@ -160,7 +161,7 @@ output partition.
This produces `2 × N` files instead of `N × M`, coalesces small batches
to a target size before writing, and bounds shuffle memory use via
-spilling — at the cost of higher write latency than the hash writer.
+spilling.
Worst-case sort-shuffle memory per executor is approximately
`vcores × memory_limit_per_task_bytes`, since one writer task can run per
@@ -169,31 +170,12 @@ sooner, or raise it to keep more data in memory and
reduce spill I/O. Setting it
to `0` removes the budget entirely and is safe only with a bounded memory pool
(see the warning above).
-### Hash-based shuffle (opt-in)
-
-The hash-based writer hashes each incoming `RecordBatch` and immediately
-encodes the per-partition slices to Arrow IPC, streaming them into one
-file per `(input_partition, output_partition)` pair. Nothing is buffered
-in memory across batches.
-
-This is simple and low latency, but for `N` input partitions and `M`
-output partitions it produces `N × M` files. Wide shuffles can therefore
-generate a large number of small files. Consider switching to the
-hash-based writer for narrow shuffles where the additional buffering
-and merging of the sort-based writer is unnecessary overhead:
-
-```rust
-let session_config = SessionConfig::new_with_ballista()
- .set_bool("ballista.shuffle.sort_based.enabled", false);
-```
-
The following session-level keys tune its behavior:
-| key | type | default
| description
|
-| ------------------------------------------------------- | ------- |
--------- |
------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
|
-| ballista.shuffle.sort_based.enabled | Boolean | true
| Enables the sort-based shuffle writer.
|
-| ballista.shuffle.sort_based.batch_size | UInt64 | 8192
| Target row count when coalescing buffered batches before they are written or
spilled.
|
-| ballista.shuffle.sort_based.memory_limit_per_task_bytes | UInt64 |
268435456 | Per-task buffered-bytes budget at which the writer spills to disk
(256 MiB default). Counted independently of the runtime memory pool. Set to `0`
to spill only under memory pressure — safe only with a bounded memory pool,
otherwise the writer never spills and may run out of memory. |
+| key | type | default
| description
|
+| ------------------------------------------------------- | ------ | ---------
|
------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
|
+| ballista.shuffle.sort_based.batch_size | UInt64 | 8192
| Target row count when coalescing buffered batches before they are written or
spilled.
|
+| ballista.shuffle.sort_based.memory_limit_per_task_bytes | UInt64 | 268435456
| Per-task buffered-bytes budget at which the writer spills to disk (256 MiB
default). Counted independently of the runtime memory pool. Set to `0` to spill
only under memory pressure — safe only with a bounded memory pool, otherwise
the writer never spills and may run out of memory. |
## Adaptive Query Execution (Experimental)
diff --git a/python/python/tests/test_context.py
b/python/python/tests/test_context.py
index 047c8986e..6e5180378 100644
--- a/python/python/tests/test_context.py
+++ b/python/python/tests/test_context.py
@@ -95,7 +95,7 @@ def test_cluster_config_propagates_to_distributed_dataframe():
def test_cluster_config_accepts_ballista_namespaced_keys():
- """Ballista-namespaced keys (e.g. ``ballista.shuffle.sort_based.enabled``)
+ """Ballista-namespaced keys (e.g.
``ballista.shuffle.sort_based.batch_size``)
are not understood by the local DataFusion ``SessionConfig`` and used to
panic when applied to it. They are forwarded to the scheduler only and
must be ignored locally rather than crashing context construction.
@@ -103,7 +103,7 @@ def test_cluster_config_accepts_ballista_namespaced_keys():
(address, port) = setup_test_cluster()
overrides = {
"datafusion.execution.target_partitions": "8",
- "ballista.shuffle.sort_based.enabled": "true",
+ "ballista.shuffle.sort_based.batch_size": "8192",
}
ctx = BallistaSessionContext(
address=f"df://{address}:{port}",
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
