You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by ag...@apache.org on 2026/05/09 19:58:15 UTC
(datafusion-ballista) branch main updated: feat(scheduler): broadcast-style hash join for small-side joins (#1647)
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 2ca0c4b3f feat(scheduler): broadcast-style hash join for small-side joins (#1647)
2ca0c4b3f is described below
commit 2ca0c4b3f3f01e058d100821e8291760c84f9eae
Author: Andy Grove <ag...@apache.org>
AuthorDate: Sat May 9 13:58:10 2026 -0600
feat(scheduler): broadcast-style hash join for small-side joins (#1647)
---
ballista/core/proto/ballista.proto | 4 +
ballista/core/src/config.rs | 18 +
.../core/src/execution_plans/shuffle_reader.rs | 142 ++++-
.../core/src/execution_plans/unresolved_shuffle.rs | 60 ++-
ballista/core/src/extension.rs | 34 +-
ballista/core/src/serde/generated/ballista.rs | 8 +
ballista/core/src/serde/mod.rs | 104 +++-
ballista/scheduler/src/planner.rs | 579 ++++++++++++++++++++-
ballista/scheduler/src/state/aqe/mod.rs | 8 +
examples/Cargo.toml | 5 +
examples/examples/standalone-broadcast-join.rs | 185 +++++++
11 files changed, 1093 insertions(+), 54 deletions(-)
diff --git a/ballista/core/proto/ballista.proto b/ballista/core/proto/ballista.proto
index a4f8bf22f..bdd68ba88 100644
--- a/ballista/core/proto/ballista.proto
+++ b/ballista/core/proto/ballista.proto
@@ -76,6 +76,8 @@ message UnresolvedShuffleExecNode {
uint32 stage_id = 1;
datafusion_common.Schema schema = 2;
datafusion.Partitioning partitioning = 5;
+ bool broadcast = 6;
+ uint32 upstream_partition_count = 7;
}
message ShuffleReaderExecNode {
@@ -84,6 +86,8 @@ message ShuffleReaderExecNode {
// The stage to read from
uint32 stage_id = 3;
datafusion.Partitioning partitioning = 4;
+ bool broadcast = 5;
+ uint32 upstream_partition_count = 6;
}
message ShuffleReaderPartition {
diff --git a/ballista/core/src/config.rs b/ballista/core/src/config.rs
index 84d851e19..334d5b991 100644
--- a/ballista/core/src/config.rs
+++ b/ballista/core/src/config.rs
@@ -80,6 +80,11 @@ pub const BALLISTA_SHUFFLE_WRITER_CHANNEL_CAPACITY: &str =
/// sort shuffle writer spills its in-memory batches to disk.
pub const BALLISTA_SHUFFLE_SORT_BASED_MEMORY_LIMIT_PER_TASK_BYTES: &str =
"ballista.shuffle.sort_based.memory_limit_per_task_bytes";
+/// Configuration key for the byte-size threshold below which a hash join's
+/// smaller side is promoted to `CollectLeft` and lowered via the broadcast
+/// pattern in the distributed planner. Set to `0` to disable promotion.
+pub const BALLISTA_BROADCAST_JOIN_THRESHOLD_BYTES: &str =
+ "ballista.optimizer.broadcast_join_threshold_bytes";
/// Result type for configuration parsing operations.
pub type ParseResult<T> = result::Result<T, String>;
@@ -152,6 +157,12 @@ static CONFIG_ENTRIES: LazyLock<HashMap<String, ConfigEntry>> = LazyLock::new(||
memory per executor is approximately concurrent_tasks * this value.".to_string(),
DataType::UInt64,
Some((256 * 1024 * 1024).to_string())),
+ ConfigEntry::new(BALLISTA_BROADCAST_JOIN_THRESHOLD_BYTES.to_string(),
+ "Byte-size threshold below which a hash join's smaller side is \
+ promoted to CollectLeft and lowered via the broadcast pattern. \
+ Set to 0 to disable promotion.".to_string(),
+ DataType::UInt64,
+ Some((10 * 1024 * 1024).to_string())),
ConfigEntry::new(BALLISTA_CLIENT_PULL.to_string(),
"Should client employ pull or push job tracking. In pull mode client will make a request to server in the loop, until job finishes. Pull mode is kept for legacy clients.".to_string(),
DataType::Boolean,
@@ -365,6 +376,13 @@ impl BallistaConfig {
self.get_usize_setting(BALLISTA_SHUFFLE_SORT_BASED_MEMORY_LIMIT_PER_TASK_BYTES)
}
+ /// Returns the byte-size threshold below which a hash join's smaller side
+ /// is promoted to `CollectLeft` and lowered via the broadcast pattern.
+ /// `0` disables promotion.
+ pub fn broadcast_join_threshold_bytes(&self) -> usize {
+ self.get_usize_setting(BALLISTA_BROADCAST_JOIN_THRESHOLD_BYTES)
+ }
+
/// Should client employ pull or push job tracking strategy
pub fn client_pull(&self) -> bool {
self.get_bool_setting(BALLISTA_CLIENT_PULL)
diff --git a/ballista/core/src/execution_plans/shuffle_reader.rs b/ballista/core/src/execution_plans/shuffle_reader.rs
index 303fb6ba6..ee75917c5 100644
--- a/ballista/core/src/execution_plans/shuffle_reader.rs
+++ b/ballista/core/src/execution_plans/shuffle_reader.rs
@@ -69,6 +69,15 @@ pub struct ShuffleReaderExec {
pub(crate) schema: SchemaRef,
/// Each partition of a shuffle can read data from multiple locations
pub partition: Vec<Vec<PartitionLocation>>,
+ /// When true, every call to `execute(partition)` reads `partition[0]`
+ /// (which holds the flattened concatenation of all upstream partition
+ /// locations) regardless of the partition index. Used for the
+ /// distributed broadcast hash-join lowering.
+ pub broadcast: bool,
+ /// Number of shuffle output partitions on the upstream stage. Useful for
+ /// metrics and EXPLAIN output. For non-broadcast readers this equals
+ /// `partition.len()`.
+ pub upstream_partition_count: usize,
/// Execution metrics
metrics: ExecutionPlanMetricsSet,
properties: Arc<PlanProperties>,
@@ -77,13 +86,15 @@ pub struct ShuffleReaderExec {
}
impl ShuffleReaderExec {
- /// Create a new ShuffleReaderExec
+ /// Create a new ShuffleReaderExec for a standard one-to-one
+ /// per-partition read.
pub fn try_new(
stage_id: usize,
partition: Vec<Vec<PartitionLocation>>,
schema: SchemaRef,
partitioning: Partitioning,
) -> Result<Self> {
+ let upstream_partition_count = partition.len();
let properties = Arc::new(PlanProperties::new(
datafusion::physical_expr::EquivalenceProperties::new(schema.clone()),
partitioning,
@@ -94,6 +105,37 @@ impl ShuffleReaderExec {
stage_id,
schema,
partition,
+ broadcast: false,
+ upstream_partition_count,
+ metrics: ExecutionPlanMetricsSet::new(),
+ properties,
+ work_dir: None, // to be updated at the executor side
+ client_pool: None, // to be updated at the executor side
+ })
+ }
+
+ /// Create a broadcast ShuffleReaderExec. `all_locations` is the
+ /// flattened concatenation of every upstream partition's locations.
+ /// The reader has one logical output partition that fans in all of
+ /// them.
+ pub fn try_new_broadcast(
+ stage_id: usize,
+ all_locations: Vec<PartitionLocation>,
+ schema: SchemaRef,
+ upstream_partition_count: usize,
+ ) -> Result<Self> {
+ let properties = Arc::new(PlanProperties::new(
+ datafusion::physical_expr::EquivalenceProperties::new(schema.clone()),
+ Partitioning::UnknownPartitioning(1),
+ datafusion::physical_plan::execution_plan::EmissionType::Incremental,
+ datafusion::physical_plan::execution_plan::Boundedness::Bounded,
+ ));
+ Ok(Self {
+ stage_id,
+ schema,
+ partition: vec![all_locations],
+ broadcast: true,
+ upstream_partition_count,
metrics: ExecutionPlanMetricsSet::new(),
properties,
work_dir: None, // to be updated at the executor side
@@ -107,6 +149,8 @@ impl ShuffleReaderExec {
stage_id: self.stage_id,
schema: self.schema.clone(),
partition: self.partition.clone(),
+ broadcast: self.broadcast,
+ upstream_partition_count: self.upstream_partition_count,
metrics: self.metrics.clone(),
properties: self.properties.clone(),
work_dir: Some(work_dir),
@@ -119,6 +163,8 @@ impl ShuffleReaderExec {
stage_id: self.stage_id,
schema: self.schema.clone(),
partition: self.partition.clone(),
+ broadcast: self.broadcast,
+ upstream_partition_count: self.upstream_partition_count,
metrics: self.metrics.clone(),
properties: self.properties.clone(),
work_dir: self.work_dir.clone(),
@@ -135,11 +181,19 @@ impl DisplayAs for ShuffleReaderExec {
) -> std::fmt::Result {
match t {
DisplayFormatType::Default | DisplayFormatType::Verbose => {
- write!(
- f,
- "ShuffleReaderExec: partitioning: {}",
- self.properties.partitioning,
- )
+ if self.broadcast {
+ write!(
+ f,
+ "ShuffleReaderExec: broadcast=true, upstream_partitions: {}",
+ self.upstream_partition_count,
+ )
+ } else {
+ write!(
+ f,
+ "ShuffleReaderExec: partitioning: {}",
+ self.properties.partitioning,
+ )
+ }
}
DisplayFormatType::TreeRender => {
write!(f, "partitioning={}", self.properties.partitioning)
@@ -177,6 +231,8 @@ impl ExecutionPlan for ShuffleReaderExec {
stage_id: self.stage_id,
schema: self.schema.clone(),
partition: self.partition.clone(),
+ broadcast: self.broadcast,
+ upstream_partition_count: self.upstream_partition_count,
metrics: ExecutionPlanMetricsSet::new(),
properties: self.properties.clone(),
work_dir: self.work_dir.clone(),
@@ -196,6 +252,10 @@ impl ExecutionPlan for ShuffleReaderExec {
) -> Result<SendableRecordBatchStream> {
let task_id = context.task_id().unwrap_or_else(|| partition.to_string());
debug!("ShuffleReaderExec::execute({task_id})");
+ // Broadcast readers have a single logical output partition; always
+ // serve from partition[0] regardless of which physical partition the
+ // caller asked for.
+ let partition = if self.broadcast { 0 } else { partition };
let config = context.session_config();
let batch_size = config.batch_size();
@@ -262,6 +322,26 @@ impl ExecutionPlan for ShuffleReaderExec {
}
fn partition_statistics(&self, partition: Option<usize>) -> Result<Statistics> {
+ if self.broadcast {
+ if let Some(idx) = partition
+ && idx != 0
+ {
+ return datafusion::common::internal_err!(
+ "Broadcast ShuffleReaderExec: invalid partition index {idx}, only partition 0 exists"
+ );
+ }
+ let all_locations: &[PartitionLocation] =
+ self.partition.first().map(|v| v.as_slice()).unwrap_or(&[]);
+ let stats = stats_for_partitions(
+ self.schema.fields().len(),
+ all_locations.iter().map(|loc| loc.partition_stats),
+ );
+ trace!(
+ "broadcast shuffle reader at stage {} returned aggregated statistics: {:?}",
+ self.stage_id, stats
+ );
+ return Ok(stats);
+ }
if let Some(idx) = partition {
let partition_count = self.properties().partitioning.partition_count();
if idx >= partition_count {
@@ -1452,4 +1532,54 @@ mod tests {
Ok(())
}
+
+ #[test]
+ fn broadcast_reader_aggregates_stats_across_upstream_partitions() {
+ let schema = Arc::new(Schema::new(vec![Field::new("c", DataType::Int32, false)]));
+
+ let locs: Vec<PartitionLocation> = (0..3)
+ .map(|partition_id| PartitionLocation {
+ map_partition_id: 0,
+ partition_id: PartitionId {
+ job_id: "j".to_string(),
+ stage_id: 7,
+ partition_id,
+ },
+ executor_meta: ExecutorMetadata {
+ id: format!("exec-{partition_id}"),
+ host: "localhost".to_string(),
+ port: 50051,
+ grpc_port: 50052,
+ specification: ExecutorSpecification::default(),
+ os_info: ExecutorOperatingSystemSpecification::default(),
+ },
+ partition_stats: PartitionStats::new(Some(100), Some(1), Some(1024)),
+ file_id: None,
+ is_sort_shuffle: false,
+ })
+ .collect();
+
+ let reader = ShuffleReaderExec::try_new_broadcast(7, locs, schema, 3).unwrap();
+
+ assert!(reader.broadcast);
+ assert_eq!(reader.upstream_partition_count, 3);
+ assert_eq!(reader.properties().partitioning.partition_count(), 1);
+ assert_eq!(reader.partition[0].len(), 3);
+
+ let stats = reader.partition_statistics(Some(0)).unwrap();
+ assert_eq!(stats.num_rows.get_value().copied(), Some(300));
+ assert_eq!(stats.total_byte_size.get_value().copied(), Some(3072));
+ }
+
+ #[test]
+ fn broadcast_reader_rejects_out_of_range_partition_index() {
+ let schema = Arc::new(Schema::new(vec![Field::new("c", DataType::Int32, false)]));
+ let reader = ShuffleReaderExec::try_new_broadcast(7, vec![], schema, 3).unwrap();
+ let err = reader.partition_statistics(Some(1)).unwrap_err();
+ let msg = err.to_string();
+ assert!(
+ msg.contains("invalid partition index 1"),
+ "unexpected error message: {msg}"
+ );
+ }
}
diff --git a/ballista/core/src/execution_plans/unresolved_shuffle.rs b/ballista/core/src/execution_plans/unresolved_shuffle.rs
index a6a8b6623..06157b75a 100644
--- a/ballista/core/src/execution_plans/unresolved_shuffle.rs
+++ b/ballista/core/src/execution_plans/unresolved_shuffle.rs
@@ -41,12 +41,24 @@ pub struct UnresolvedShuffleExec {
/// The partition count this node will have once it is replaced with a ShuffleReaderExec.
pub output_partition_count: usize,
+ /// The number of shuffle output partitions on the upstream stage. For
+ /// non-broadcast readers this equals `output_partition_count`. For
+ /// broadcast readers this is M (one logical output partition fans in
+ /// all M upstream partition files).
+ pub upstream_partition_count: usize,
+
+ /// When true, the resolved `ShuffleReaderExec` reads *all* upstream
+ /// partition files into its single output partition (broadcast pattern).
+ pub broadcast: bool,
+
properties: Arc<PlanProperties>,
}
impl UnresolvedShuffleExec {
- /// Create a new UnresolvedShuffleExec
+ /// Create a new UnresolvedShuffleExec for a standard one-to-one
+ /// per-partition read.
pub fn new(stage_id: usize, schema: SchemaRef, partitioning: Partitioning) -> Self {
+ let partition_count = partitioning.partition_count();
let properties = Arc::new(PlanProperties::new(
datafusion::physical_expr::EquivalenceProperties::new(schema.clone()),
partitioning,
@@ -56,7 +68,33 @@ impl UnresolvedShuffleExec {
Self {
stage_id,
schema,
- output_partition_count: properties.partitioning.partition_count(),
+ output_partition_count: partition_count,
+ upstream_partition_count: partition_count,
+ broadcast: false,
+ properties,
+ }
+ }
+
+ /// Create a broadcast UnresolvedShuffleExec. The resolved reader has
+ /// one logical output partition that fans in all `upstream_partition_count`
+ /// shuffle files from the upstream stage.
+ pub fn new_broadcast(
+ stage_id: usize,
+ schema: SchemaRef,
+ upstream_partition_count: usize,
+ ) -> Self {
+ let properties = Arc::new(PlanProperties::new(
+ datafusion::physical_expr::EquivalenceProperties::new(schema.clone()),
+ Partitioning::UnknownPartitioning(1),
+ datafusion::physical_plan::execution_plan::EmissionType::Incremental,
+ datafusion::physical_plan::execution_plan::Boundedness::Bounded,
+ ));
+ Self {
+ stage_id,
+ schema,
+ output_partition_count: 1,
+ upstream_partition_count,
+ broadcast: true,
properties,
}
}
@@ -70,11 +108,19 @@ impl DisplayAs for UnresolvedShuffleExec {
) -> std::fmt::Result {
match t {
DisplayFormatType::Default | DisplayFormatType::Verbose => {
- write!(
- f,
- "UnresolvedShuffleExec: partitioning: {}",
- self.properties().output_partitioning()
- )
+ if self.broadcast {
+ write!(
+ f,
+ "UnresolvedShuffleExec: broadcast=true, upstream_partitions: {}",
+ self.upstream_partition_count,
+ )
+ } else {
+ write!(
+ f,
+ "UnresolvedShuffleExec: partitioning: {}",
+ self.properties().output_partitioning()
+ )
+ }
}
DisplayFormatType::TreeRender => {
write!(
diff --git a/ballista/core/src/extension.rs b/ballista/core/src/extension.rs
index 10dc8909c..c8745cec3 100644
--- a/ballista/core/src/extension.rs
+++ b/ballista/core/src/extension.rs
@@ -16,7 +16,8 @@
// under the License.
use crate::config::{
- BALLISTA_CLIENT_GRPC_MAX_MESSAGE_SIZE, BALLISTA_CLIENT_USE_TLS, BALLISTA_JOB_NAME,
+ BALLISTA_BROADCAST_JOIN_THRESHOLD_BYTES, BALLISTA_CLIENT_GRPC_MAX_MESSAGE_SIZE,
+ BALLISTA_CLIENT_USE_TLS, BALLISTA_JOB_NAME,
BALLISTA_SHUFFLE_READER_FORCE_REMOTE_READ, BALLISTA_SHUFFLE_READER_MAX_REQUESTS,
BALLISTA_SHUFFLE_READER_REMOTE_PREFER_FLIGHT, BALLISTA_STANDALONE_PARALLELISM,
BallistaConfig,
@@ -178,6 +179,17 @@ pub trait SessionConfigExt {
/// This option to be used to configure standalone session context
fn with_ballista_standalone_parallelism(self, parallelism: usize) -> Self;
+ /// Returns the byte-size threshold below which a hash join's smaller side
+ /// is promoted to `CollectLeft` and lowered via the broadcast pattern in
+ /// the distributed planner. `0` disables promotion.
+ fn ballista_broadcast_join_threshold_bytes(&self) -> usize;
+
+ /// Sets the byte-size threshold below which a hash join's smaller side
+ /// is promoted to `CollectLeft` and lowered via the broadcast pattern in
+ /// the distributed planner. Setting `0` disables promotion.
+ fn with_ballista_broadcast_join_threshold_bytes(self, threshold_bytes: usize)
+ -> Self;
+
/// retrieves grpc client max message size
fn ballista_grpc_client_max_message_size(&self) -> usize;
@@ -423,6 +435,26 @@ impl SessionConfigExt for SessionConfig {
}
}
+ fn ballista_broadcast_join_threshold_bytes(&self) -> usize {
+ self.options()
+ .extensions
+ .get::<BallistaConfig>()
+ .map(|c| c.broadcast_join_threshold_bytes())
+ .unwrap_or_else(|| BallistaConfig::default().broadcast_join_threshold_bytes())
+ }
+
+ fn with_ballista_broadcast_join_threshold_bytes(
+ self,
+ threshold_bytes: usize,
+ ) -> Self {
+ if self.options().extensions.get::<BallistaConfig>().is_some() {
+ self.set_usize(BALLISTA_BROADCAST_JOIN_THRESHOLD_BYTES, threshold_bytes)
+ } else {
+ self.with_option_extension(BallistaConfig::default())
+ .set_usize(BALLISTA_BROADCAST_JOIN_THRESHOLD_BYTES, threshold_bytes)
+ }
+ }
+
fn ballista_shuffle_reader_maximum_concurrent_requests(&self) -> usize {
self.options()
.extensions
diff --git a/ballista/core/src/serde/generated/ballista.rs b/ballista/core/src/serde/generated/ballista.rs
index d888280e7..897326fdd 100644
--- a/ballista/core/src/serde/generated/ballista.rs
+++ b/ballista/core/src/serde/generated/ballista.rs
@@ -91,6 +91,10 @@ pub struct UnresolvedShuffleExecNode {
pub schema: ::core::option::Option<::datafusion_proto_common::Schema>,
#[prost(message, optional, tag = "5")]
pub partitioning: ::core::option::Option<::datafusion_proto::protobuf::Partitioning>,
+ #[prost(bool, tag = "6")]
+ pub broadcast: bool,
+ #[prost(uint32, tag = "7")]
+ pub upstream_partition_count: u32,
}
#[derive(Clone, PartialEq, ::prost::Message)]
pub struct ShuffleReaderExecNode {
@@ -103,6 +107,10 @@ pub struct ShuffleReaderExecNode {
pub stage_id: u32,
#[prost(message, optional, tag = "4")]
pub partitioning: ::core::option::Option<::datafusion_proto::protobuf::Partitioning>,
+ #[prost(bool, tag = "5")]
+ pub broadcast: bool,
+ #[prost(uint32, tag = "6")]
+ pub upstream_partition_count: u32,
}
#[derive(Clone, PartialEq, ::prost::Message)]
pub struct ShuffleReaderPartition {
diff --git a/ballista/core/src/serde/mod.rs b/ballista/core/src/serde/mod.rs
index 031f93a4c..d24398201 100644
--- a/ballista/core/src/serde/mod.rs
+++ b/ballista/core/src/serde/mod.rs
@@ -425,13 +425,28 @@ impl PhysicalExtensionCodec for BallistaPhysicalExtensionCodec {
)?;
let partitioning = partitioning
.ok_or_else(|| proto_error("missing required partitioning field"))?;
- let shuffle_reader = ShuffleReaderExec::try_new(
- stage_id,
- partition_location,
- schema,
- partitioning,
- )?;
- Ok(Arc::new(shuffle_reader))
+ let exec = if shuffle_reader.broadcast {
+ let all_locations = partition_location
+ .into_iter()
+ .next()
+ .ok_or_else(|| proto_error(
+ "broadcast ShuffleReaderExec: expected exactly one partition in proto"
+ ))?;
+ ShuffleReaderExec::try_new_broadcast(
+ stage_id,
+ all_locations,
+ schema,
+ shuffle_reader.upstream_partition_count as usize,
+ )?
+ } else {
+ ShuffleReaderExec::try_new(
+ stage_id,
+ partition_location,
+ schema,
+ partitioning,
+ )?
+ };
+ Ok(Arc::new(exec))
}
PhysicalPlanType::UnresolvedShuffle(unresolved_shuffle) => {
let schema: SchemaRef =
@@ -445,11 +460,20 @@ impl PhysicalExtensionCodec for BallistaPhysicalExtensionCodec {
)?;
let partitioning = partitioning
.ok_or_else(|| proto_error("missing required partitioning field"))?;
- Ok(Arc::new(UnresolvedShuffleExec::new(
- unresolved_shuffle.stage_id as usize,
- schema,
- partitioning,
- )))
+ let exec = if unresolved_shuffle.broadcast {
+ UnresolvedShuffleExec::new_broadcast(
+ unresolved_shuffle.stage_id as usize,
+ schema,
+ unresolved_shuffle.upstream_partition_count as usize,
+ )
+ } else {
+ UnresolvedShuffleExec::new(
+ unresolved_shuffle.stage_id as usize,
+ schema,
+ partitioning,
+ )
+ };
+ Ok(Arc::new(exec))
}
}
}
@@ -571,6 +595,8 @@ impl PhysicalExtensionCodec for BallistaPhysicalExtensionCodec {
partition,
schema: Some(exec.schema().as_ref().try_into()?),
partitioning: Some(partitioning),
+ broadcast: exec.broadcast,
+ upstream_partition_count: exec.upstream_partition_count as u32,
},
)),
};
@@ -594,6 +620,8 @@ impl PhysicalExtensionCodec for BallistaPhysicalExtensionCodec {
stage_id: exec.stage_id as u32,
schema: Some(exec.schema().as_ref().try_into()?),
partitioning: Some(partitioning),
+ broadcast: exec.broadcast,
+ upstream_partition_count: exec.upstream_partition_count as u32,
},
)),
};
@@ -753,4 +781,56 @@ mod test {
assert_eq!(decoded_exec.schema().as_ref(), schema.as_ref());
assert_eq!(&decoded_exec.properties().partitioning, &partitioning);
}
+
+ #[tokio::test]
+ async fn test_broadcast_unresolved_shuffle_exec_roundtrip() {
+ let schema = create_test_schema();
+ let original_exec = UnresolvedShuffleExec::new_broadcast(7, schema.clone(), 4);
+
+ let codec = BallistaPhysicalExtensionCodec::default();
+ let mut buf: Vec<u8> = vec![];
+ codec
+ .try_encode(Arc::new(original_exec.clone()), &mut buf)
+ .unwrap();
+
+ let ctx = SessionContext::new().task_ctx();
+ let decoded_plan = codec.try_decode(&buf, &[], &ctx).unwrap();
+
+ let decoded_exec = decoded_plan
+ .as_any()
+ .downcast_ref::<UnresolvedShuffleExec>()
+ .expect("Expected UnresolvedShuffleExec");
+
+ assert_eq!(decoded_exec.stage_id, 7);
+ assert!(decoded_exec.broadcast);
+ assert_eq!(decoded_exec.upstream_partition_count, 4);
+ assert_eq!(decoded_exec.output_partition_count, 1);
+ }
+
+ #[tokio::test]
+ async fn test_broadcast_shuffle_reader_exec_roundtrip() {
+ let schema = create_test_schema();
+ let original_exec =
+ ShuffleReaderExec::try_new_broadcast(7, Vec::new(), schema.clone(), 4)
+ .unwrap();
+
+ let codec = BallistaPhysicalExtensionCodec::default();
+ let mut buf: Vec<u8> = vec![];
+ codec
+ .try_encode(Arc::new(original_exec.clone()), &mut buf)
+ .unwrap();
+
+ let ctx = SessionContext::new().task_ctx();
+ let decoded_plan = codec.try_decode(&buf, &[], &ctx).unwrap();
+
+ let decoded_exec = decoded_plan
+ .as_any()
+ .downcast_ref::<ShuffleReaderExec>()
+ .expect("Expected ShuffleReaderExec");
+
+ assert_eq!(decoded_exec.stage_id, 7);
+ assert!(decoded_exec.broadcast);
+ assert_eq!(decoded_exec.upstream_partition_count, 4);
+ assert_eq!(decoded_exec.partition.len(), 1);
+ }
}
diff --git a/ballista/scheduler/src/planner.rs b/ballista/scheduler/src/planner.rs
index 094c2ef7a..dc6c8b690 100644
--- a/ballista/scheduler/src/planner.rs
+++ b/ballista/scheduler/src/planner.rs
@@ -31,16 +31,20 @@ use ballista_core::{
},
serde::scheduler::PartitionLocation,
};
+use datafusion::arrow::datatypes::DataType;
use datafusion::config::ConfigOptions;
use datafusion::physical_optimizer::PhysicalOptimizerRule;
use datafusion::physical_optimizer::enforce_sorting::EnforceSorting;
use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec;
+use datafusion::physical_plan::joins::{HashJoinExec, PartitionMode};
use datafusion::physical_plan::repartition::RepartitionExec;
use datafusion::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
use datafusion::physical_plan::{
ExecutionPlan, Partitioning, with_new_children_if_necessary,
};
+use crate::physical_optimizer::join_selection::should_swap_join_order;
+
use log::info;
type PartialQueryStageResult = (Arc<dyn ExecutionPlan>, Vec<Arc<dyn ShuffleWriter>>);
@@ -126,11 +130,59 @@ impl DefaultDistributedPlanner {
execution_plan: Arc<dyn ExecutionPlan>,
config: &ConfigOptions,
) -> Result<PartialQueryStageResult> {
+ // Apply broadcast-join promotion before recursing.
+ let execution_plan = Self::maybe_promote_to_broadcast(execution_plan, config)?;
+
// recurse down and replace children
if execution_plan.children().is_empty() {
return Ok((execution_plan, vec![]));
}
+ // Broadcast-join lowering: HashJoinExec(CollectLeft) gets its own
+ // controlled recursion so the build side is written as a broadcast stage.
+ if let Some(hash_join) = execution_plan.as_any().downcast_ref::<HashJoinExec>()
+ && *hash_join.partition_mode() == PartitionMode::CollectLeft
+ {
+ // Build subtree: peel CoalescePartitionsExec if present, then
+ // recurse to lower its internal stages.
+ let mut build = hash_join.left().clone();
+ if let Some(coalesce) =
+ build.as_any().downcast_ref::<CoalescePartitionsExec>()
+ {
+ build = coalesce.children()[0].clone();
+ }
+ let (build, mut stages) =
+ self.plan_query_stages_internal(job_id, build, config)?;
+ let build_partitions =
+ build.properties().output_partitioning().partition_count();
+
+ let build_writer = create_shuffle_writer_with_config(
+ job_id,
+ self.next_stage_id(),
+ build,
+ None,
+ config,
+ )?;
+ let broadcast_left = Arc::new(UnresolvedShuffleExec::new_broadcast(
+ build_writer.stage_id(),
+ build_writer.schema(),
+ build_partitions,
+ ));
+ stages.push(build_writer);
+
+ // Probe subtree: recurse normally.
+ let (probe, mut probe_stages) = self.plan_query_stages_internal(
+ job_id,
+ hash_join.right().clone(),
+ config,
+ )?;
+ stages.append(&mut probe_stages);
+
+ let new_join =
+ execution_plan.with_new_children(vec![broadcast_left, probe])?;
+ return Ok((new_join, stages));
+ }
+
let mut stages = vec![];
let mut children = vec![];
for child in execution_plan.children() {
@@ -216,6 +268,150 @@ impl DefaultDistributedPlanner {
self.next_stage_id += 1;
self.next_stage_id
}
+
+ /// If `plan` is a `HashJoinExec(Partitioned)` whose smaller side fits
+ /// under the broadcast threshold, returns a rewritten
+ /// `HashJoinExec(CollectLeft)` (with a swap if the small side was on
+ /// the right) wrapped so the build subtree is a single-partition input.
+ /// Otherwise returns the input unchanged.
+ fn maybe_promote_to_broadcast(
+ plan: Arc<dyn ExecutionPlan>,
+ config: &ConfigOptions,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ let threshold_bytes = config
+ .extensions
+ .get::<BallistaConfig>()
+ .map(|c| c.broadcast_join_threshold_bytes())
+ .unwrap_or_else(|| {
+ BallistaConfig::default().broadcast_join_threshold_bytes()
+ });
+ if threshold_bytes == 0 {
+ info!("broadcast check: threshold is 0, broadcast disabled");
+ return Ok(plan);
+ }
+ let Some(hash_join) = plan.as_any().downcast_ref::<HashJoinExec>() else {
+ return Ok(plan);
+ };
+ info!(
+ "broadcast check: evaluating HashJoinExec mode={:?} join_type={:?} threshold={threshold_bytes}",
+ hash_join.partition_mode(),
+ hash_join.join_type(),
+ );
+ if *hash_join.partition_mode() != PartitionMode::Partitioned {
+ return Ok(plan);
+ }
+ if hash_join.null_aware {
+ return Ok(plan);
+ }
+
+ let left = hash_join.left();
+ let right = hash_join.right();
+
+ fn under(plan: &dyn ExecutionPlan, threshold: usize) -> bool {
+ let Ok(stats) = plan.partition_statistics(None) else {
+ info!(
+ "broadcast check: partition_statistics returned error for {}",
+ plan.name()
+ );
+ return false;
+ };
+ info!(
+ "broadcast check: {} total_byte_size={:?} num_rows={:?} threshold={}",
+ plan.name(),
+ stats.total_byte_size,
+ stats.num_rows,
+ threshold,
+ );
+ if let Some(bytes) = stats.total_byte_size.get_value() {
+ *bytes != 0 && *bytes < threshold
+ } else if let Some(rows) = stats.num_rows.get_value() {
+ let schema = plan.schema();
+ let bytes_per_row: usize = schema
+ .fields()
+ .iter()
+ .map(|f| {
+ match f.data_type() {
+ DataType::Boolean => 1,
+ DataType::Int8 | DataType::UInt8 => 1,
+ DataType::Int16 | DataType::UInt16 => 2,
+ DataType::Int32 | DataType::UInt32 | DataType::Float32 => 4,
+ DataType::Int64 | DataType::UInt64 | DataType::Float64 => 8,
+ DataType::Date32 => 4,
+ DataType::Date64 => 8,
+ DataType::Decimal128(_, _) => 16,
+ DataType::Decimal256(_, _) => 32,
+ _ => 32, // conservative estimate for variable-length types
+ }
+ })
+ .sum();
+ let estimated_bytes = *rows * bytes_per_row.max(8);
+ info!(
+ "broadcast check: estimated {estimated_bytes} bytes ({rows} rows * {bytes_per_row} bytes/row from {} columns)",
+ schema.fields().len(),
+ );
+ estimated_bytes != 0 && estimated_bytes < threshold
+ } else {
+ false
+ }
+ }
+
+ let left_under = under(&**left, threshold_bytes);
+ let right_under = under(&**right, threshold_bytes);
+ if !left_under && !right_under {
+ info!("broadcast check: neither side under threshold, skipping promotion");
+ return Ok(plan);
+ }
+
+ // Determine swap: put the smaller side on the left (build).
+ let swap = if left_under && right_under {
+ should_swap_join_order(&**left, &**right)?
+ } else {
+ right_under
+ };
+
+ info!(
+ "broadcast check: promoting to CollectLeft (left_under={left_under}, right_under={right_under}, swap={swap})"
+ );
+
+ let promoted: Arc<dyn ExecutionPlan> = if swap {
+ if !hash_join.join_type().supports_swap() {
+ info!(
+ "broadcast check: join type {:?} does not support swap, skipping",
+ hash_join.join_type()
+ );
+ return Ok(plan);
+ }
+ hash_join.swap_inputs(PartitionMode::CollectLeft)?
+ } else {
+ Arc::new(
+ hash_join
+ .builder()
+ .with_partition_mode(PartitionMode::CollectLeft)
+ .build()?,
+ )
+ };
+
+ let promoted_join = promoted
+ .as_any()
+ .downcast_ref::<HashJoinExec>()
+ .expect("promoted plan must still be a HashJoinExec");
+ let new_left: Arc<dyn ExecutionPlan> = if promoted_join
+ .left()
+ .properties()
+ .output_partitioning()
+ .partition_count()
+ > 1
+ {
+ Arc::new(CoalescePartitionsExec::new(promoted_join.left().clone()))
+ } else {
+ promoted_join.left().clone()
+ };
+ let new_right = promoted_join.right().clone();
+ Ok(with_new_children_if_necessary(
+ promoted,
+ vec![new_left, new_right],
+ )?)
+ }
}
fn create_unresolved_shuffle(
@@ -262,7 +458,6 @@ pub fn remove_unresolved_shuffles(
if let Some(unresolved_shuffle) =
child.as_any().downcast_ref::<UnresolvedShuffleExec>()
{
- let mut relevant_locations = vec![];
let p = partition_locations
.get(&unresolved_shuffle.stage_id)
.ok_or_else(|| {
@@ -273,23 +468,38 @@ pub fn remove_unresolved_shuffles(
})?
.clone();
- for i in 0..unresolved_shuffle.output_partition_count {
- if let Some(x) = p.get(&i) {
- relevant_locations.push(x.to_owned());
- } else {
- relevant_locations.push(vec![]);
+ if unresolved_shuffle.broadcast {
+ let mut all_locations = vec![];
+ for i in 0..unresolved_shuffle.upstream_partition_count {
+ if let Some(locs) = p.get(&i) {
+ all_locations.extend(locs.iter().cloned());
+ }
+ }
+ new_children.push(Arc::new(ShuffleReaderExec::try_new_broadcast(
+ unresolved_shuffle.stage_id,
+ all_locations,
+ unresolved_shuffle.schema().clone(),
+ unresolved_shuffle.upstream_partition_count,
+ )?));
+ } else {
+ let mut relevant_locations = vec![];
+ for i in 0..unresolved_shuffle.output_partition_count {
+ if let Some(x) = p.get(&i) {
+ relevant_locations.push(x.to_owned());
+ } else {
+ relevant_locations.push(vec![]);
+ }
}
+ new_children.push(Arc::new(ShuffleReaderExec::try_new(
+ unresolved_shuffle.stage_id,
+ relevant_locations,
+ unresolved_shuffle.schema().clone(),
+ unresolved_shuffle
+ .properties()
+ .output_partitioning()
+ .clone(),
+ )?));
}
-
- new_children.push(Arc::new(ShuffleReaderExec::try_new(
- unresolved_shuffle.stage_id,
- relevant_locations,
- unresolved_shuffle.schema().clone(),
- unresolved_shuffle
- .properties()
- .output_partitioning()
- .clone(),
- )?))
} else {
new_children.push(remove_unresolved_shuffles(
child.clone(),
@@ -310,13 +520,20 @@ pub fn rollback_resolved_shuffles(
for child in stage.children() {
if let Some(shuffle_reader) = child.as_any().downcast_ref::<ShuffleReaderExec>() {
let stage_id = shuffle_reader.stage_id;
-
- let unresolved_shuffle = Arc::new(UnresolvedShuffleExec::new(
- stage_id,
- shuffle_reader.schema(),
- shuffle_reader.properties().partitioning.clone(),
- ));
- new_children.push(unresolved_shuffle);
+ let unresolved = if shuffle_reader.broadcast {
+ Arc::new(UnresolvedShuffleExec::new_broadcast(
+ stage_id,
+ shuffle_reader.schema(),
+ shuffle_reader.upstream_partition_count,
+ ))
+ } else {
+ Arc::new(UnresolvedShuffleExec::new(
+ stage_id,
+ shuffle_reader.schema(),
+ shuffle_reader.properties().partitioning.clone(),
+ ))
+ };
+ new_children.push(unresolved);
} else {
new_children.push(rollback_resolved_shuffles(child.clone())?);
}
@@ -499,6 +716,10 @@ mod test {
let ctx = datafusion_test_context("testdata").await?;
let session_state = ctx.state();
+ // CSV files have no byte stats, so broadcast promotion never fires
+ // even with the default threshold — no opt-out needed.
+ let options = ctx.state().config().options().clone();
+
// simplified form of TPC-H query 12
let df = ctx
.sql(
@@ -542,11 +763,7 @@ order by
let mut planner = DefaultDistributedPlanner::new();
let job_uuid = Uuid::new_v4();
- let stages = planner.plan_query_stages(
- &job_uuid.to_string(),
- plan,
- ctx.state().config().options(),
- )?;
+ let stages = planner.plan_query_stages(&job_uuid.to_string(), plan, &options)?;
for (i, stage) in stages.iter().enumerate() {
println!("Stage {i}:\n{}", displayable(stage.as_ref()).indent(false));
}
@@ -674,6 +891,236 @@ order by
Ok(())
}
+ #[tokio::test]
+ async fn distributed_broadcast_join_plan() -> Result<(), BallistaError> {
+ use datafusion::physical_plan::joins::PartitionMode;
+
+ let (ctx, options) = make_broadcast_test_ctx(10 * 1024 * 1024, 1)?;
+
+ let df = ctx
+ .sql("select count(*) from big join small on big.k = small.k")
+ .await?;
+
+ let plan = df.into_optimized_plan()?;
+ let plan = ctx.state().create_physical_plan(&plan).await?;
+
+ let mut planner = DefaultDistributedPlanner::new();
+ let job_uuid = Uuid::new_v4();
+ let stages = planner.plan_query_stages(&job_uuid.to_string(), plan, &options)?;
+ for (i, stage) in stages.iter().enumerate() {
+ println!("Stage {i}:\n{}", displayable(stage.as_ref()).indent(false));
+ }
+
+ let mut found_broadcast_join = false;
+ for stage in &stages {
+ let mut walker: Vec<Arc<dyn ExecutionPlan>> =
+ vec![stage.clone() as Arc<dyn ExecutionPlan>];
+ while let Some(node) = walker.pop() {
+ if let Some(hj) = node.as_any().downcast_ref::<HashJoinExec>() {
+ assert_eq!(*hj.partition_mode(), PartitionMode::CollectLeft);
+ let left = hj.children()[0].clone();
+ let unresolved = left
+ .as_any()
+ .downcast_ref::<UnresolvedShuffleExec>()
+ .expect("left input should be UnresolvedShuffleExec");
+ assert!(unresolved.broadcast, "left input should be broadcast");
+ assert_eq!(unresolved.output_partition_count, 1);
+ found_broadcast_join = true;
+ }
+ walker.extend(node.children().iter().map(|c| (*c).clone()));
+ }
+ }
+ assert!(
+ found_broadcast_join,
+ "expected a broadcast HashJoinExec in stages"
+ );
+
+ Ok(())
+ }
+
+ #[tokio::test]
+ async fn distributed_join_plan_no_broadcast_when_threshold_zero()
+ -> Result<(), BallistaError> {
+ use datafusion::physical_plan::joins::PartitionMode;
+
+ let (ctx, options) = make_broadcast_test_ctx(0, 1)?;
+
+ let df = ctx
+ .sql("select count(*) from big join small on big.k = small.k")
+ .await?;
+
+ let plan = df.into_optimized_plan()?;
+ let plan = ctx.state().create_physical_plan(&plan).await?;
+
+ let mut planner = DefaultDistributedPlanner::new();
+ let job_uuid = Uuid::new_v4();
+ let stages = planner.plan_query_stages(&job_uuid.to_string(), plan, &options)?;
+
+ for stage in &stages {
+ let mut walker: Vec<Arc<dyn ExecutionPlan>> =
+ vec![stage.clone() as Arc<dyn ExecutionPlan>];
+ while let Some(node) = walker.pop() {
+ if let Some(unresolved) =
+ node.as_any().downcast_ref::<UnresolvedShuffleExec>()
+ {
+ assert!(
+ !unresolved.broadcast,
+ "no broadcast reader expected with threshold=0"
+ );
+ }
+ if let Some(hj) = node.as_any().downcast_ref::<HashJoinExec>() {
+ assert_ne!(
+ *hj.partition_mode(),
+ PartitionMode::CollectLeft,
+ "no CollectLeft promotion expected with threshold=0"
+ );
+ }
+ walker.extend(node.children().iter().map(|c| (*c).clone()));
+ }
+ }
+
+ Ok(())
+ }
+
+ #[tokio::test]
+ async fn distributed_broadcast_join_plan_multi_partition_build()
+ -> Result<(), BallistaError> {
+ // Build a MemTable with THREE partitions on the small side, so the
+ // broadcast build stage has 3 input partitions and writes 3 shuffle
+ // files. The broadcast UnresolvedShuffleExec must report
+ // upstream_partition_count = 3.
+ let (ctx, options) = make_broadcast_test_ctx(10 * 1024 * 1024 * 1024, 3)?;
+
+ let df = ctx
+ .sql("select count(*) from big join small on big.k = small.k")
+ .await?;
+
+ let plan = df.into_optimized_plan()?;
+ let plan = ctx.state().create_physical_plan(&plan).await?;
+
+ let mut planner = DefaultDistributedPlanner::new();
+ let job_uuid = Uuid::new_v4();
+ let stages = planner.plan_query_stages(&job_uuid.to_string(), plan, &options)?;
+ for (i, stage) in stages.iter().enumerate() {
+ println!("Stage {i}:\n{}", displayable(stage.as_ref()).indent(false));
+ }
+
+ let mut max_upstream = 0;
+ for stage in &stages {
+ let mut walker: Vec<Arc<dyn ExecutionPlan>> =
+ vec![stage.clone() as Arc<dyn ExecutionPlan>];
+ while let Some(node) = walker.pop() {
+ if let Some(unresolved) =
+ node.as_any().downcast_ref::<UnresolvedShuffleExec>()
+ && unresolved.broadcast
+ {
+ max_upstream = max_upstream.max(unresolved.upstream_partition_count);
+ }
+ walker.extend(node.children().iter().map(|c| (*c).clone()));
+ }
+ }
+ assert!(
+ max_upstream >= 2,
+ "expected broadcast UnresolvedShuffleExec with upstream_partition_count >= 2, got {max_upstream}"
+ );
+
+ Ok(())
+ }
+
+ #[tokio::test]
+ async fn remove_unresolved_shuffles_broadcasts_locations() -> Result<(), BallistaError>
+ {
+ use ballista_core::execution_plans::ShuffleReaderExec;
+ use ballista_core::serde::scheduler::{
+ ExecutorMetadata, ExecutorOperatingSystemSpecification,
+ ExecutorSpecification, PartitionId, PartitionLocation, PartitionStats,
+ };
+ use datafusion::arrow::datatypes::{DataType, Field, Schema};
+
+ let schema = Arc::new(Schema::new(vec![Field::new("c", DataType::Int32, false)]));
+ let unresolved =
+ Arc::new(UnresolvedShuffleExec::new_broadcast(42, schema.clone(), 3))
+ as Arc<dyn ExecutionPlan>;
+
+ let make_loc = |partition_id: usize| PartitionLocation {
+ map_partition_id: partition_id,
+ partition_id: PartitionId {
+ job_id: "job".to_string(),
+ stage_id: 42,
+ partition_id,
+ },
+ executor_meta: ExecutorMetadata {
+ id: format!("exec-{partition_id}"),
+ host: "localhost".to_string(),
+ port: 50050,
+ grpc_port: 50051,
+ specification: ExecutorSpecification::default().with_task_slots(1),
+ os_info: ExecutorOperatingSystemSpecification::default(),
+ },
+ partition_stats: PartitionStats::new(Some(10), None, Some(1)),
+ file_id: None,
+ is_sort_shuffle: false,
+ };
+
+ let mut by_partition: std::collections::HashMap<usize, Vec<PartitionLocation>> =
+ Default::default();
+ by_partition.insert(0, vec![make_loc(0)]);
+ by_partition.insert(1, vec![make_loc(1)]);
+ by_partition.insert(2, vec![make_loc(2)]);
+
+ let mut by_stage = std::collections::HashMap::new();
+ by_stage.insert(42usize, by_partition);
+
+ let parent: Arc<dyn ExecutionPlan> = Arc::new(
+ datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec::new(
+ unresolved,
+ ),
+ );
+
+ let resolved = crate::planner::remove_unresolved_shuffles(parent, &by_stage)?;
+
+ let resolved_child = resolved.children()[0].clone();
+ let reader = resolved_child
+ .as_any()
+ .downcast_ref::<ShuffleReaderExec>()
+ .expect("expected resolved ShuffleReaderExec");
+ assert!(reader.broadcast);
+ assert_eq!(reader.upstream_partition_count, 3);
+ assert_eq!(reader.partition.len(), 1);
+ assert_eq!(reader.partition[0].len(), 3);
+
+ Ok(())
+ }
+
+ #[tokio::test]
+ async fn rollback_resolved_shuffles_preserves_broadcast() -> Result<(), BallistaError>
+ {
+ use ballista_core::execution_plans::ShuffleReaderExec;
+ use datafusion::arrow::datatypes::{DataType, Field, Schema};
+
+ let schema = Arc::new(Schema::new(vec![Field::new("c", DataType::Int32, false)]));
+ let reader = Arc::new(
+ ShuffleReaderExec::try_new_broadcast(42, vec![], schema, 3).unwrap(),
+ ) as Arc<dyn ExecutionPlan>;
+ let parent: Arc<dyn ExecutionPlan> = Arc::new(
+ datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec::new(
+ reader,
+ ),
+ );
+
+ let rolled_back = crate::planner::rollback_resolved_shuffles(parent)?;
+ let child = rolled_back.children()[0].clone();
+ let unresolved = child
+ .as_any()
+ .downcast_ref::<UnresolvedShuffleExec>()
+ .expect("expected rolled-back UnresolvedShuffleExec");
+ assert!(unresolved.broadcast);
+ assert_eq!(unresolved.upstream_partition_count, 3);
+ assert_eq!(unresolved.output_partition_count, 1);
+
+ Ok(())
+ }
+
#[tokio::test]
async fn distributed_window_plan() -> Result<(), BallistaError> {
let ctx = datafusion_test_context("testdata").await?;
@@ -833,6 +1280,82 @@ order by
Ok(())
}
+ fn make_broadcast_test_ctx(
+ threshold_bytes: usize,
+ small_partitions: usize,
+ ) -> Result<
+ (
+ datafusion::prelude::SessionContext,
+ datafusion::config::ConfigOptions,
+ ),
+ BallistaError,
+ > {
+ use ballista_core::extension::SessionConfigExt;
+ use datafusion::arrow::array::Int32Array;
+ use datafusion::arrow::datatypes::{DataType, Field, Schema};
+ use datafusion::arrow::record_batch::RecordBatch;
+ use datafusion::datasource::MemTable;
+ use datafusion::prelude::SessionConfig;
+
+ let big_schema = Arc::new(Schema::new(vec![
+ Field::new("k", DataType::Int32, false),
+ Field::new("v", DataType::Int32, false),
+ ]));
+ let small_schema = Arc::new(Schema::new(vec![
+ Field::new("k", DataType::Int32, false),
+ Field::new("name", DataType::Int32, false),
+ ]));
+
+ let big_batch = RecordBatch::try_new(
+ big_schema.clone(),
+ vec![
+ Arc::new(Int32Array::from((0..1000).collect::<Vec<i32>>())),
+ Arc::new(Int32Array::from(
+ (0..1000).map(|i| i * 2).collect::<Vec<i32>>(),
+ )),
+ ],
+ )
+ .map_err(|e| BallistaError::General(e.to_string()))?;
+
+ let mk_small = || -> Result<RecordBatch, BallistaError> {
+ RecordBatch::try_new(
+ small_schema.clone(),
+ vec![
+ Arc::new(Int32Array::from(vec![1, 2])),
+ Arc::new(Int32Array::from(vec![10, 20])),
+ ],
+ )
+ .map_err(|e| BallistaError::General(e.to_string()))
+ };
+
+ let session_config = SessionConfig::new()
+ .with_target_partitions(2)
+ .set_usize(
+ "datafusion.optimizer.hash_join_single_partition_threshold",
+ 0,
+ )
+ .with_ballista_broadcast_join_threshold_bytes(threshold_bytes);
+ let ctx = datafusion::prelude::SessionContext::new_with_config(session_config);
+ ctx.register_batch("big", big_batch)?;
+
+ if small_partitions <= 1 {
+ let small_batch = mk_small()?;
+ ctx.register_batch("small", small_batch)?;
+ } else {
+ let mut partitions: Vec<Vec<RecordBatch>> =
+ Vec::with_capacity(small_partitions);
+ for _ in 0..small_partitions {
+ partitions.push(vec![mk_small()?]);
+ }
+ let small_table = MemTable::try_new(small_schema, partitions)
+ .map_err(|e| BallistaError::General(e.to_string()))?;
+ ctx.register_table("small", Arc::new(small_table))?;
+ }
+
+ let options = ctx.state().config().options().clone();
+ Ok((ctx, (*options).clone()))
+ }
+
fn roundtrip_operator(
ctx: &TaskContext,
plan: Arc<dyn ExecutionPlan>,
diff --git a/ballista/scheduler/src/state/aqe/mod.rs b/ballista/scheduler/src/state/aqe/mod.rs
index f83c15c3a..470a6e9ee 100644
--- a/ballista/scheduler/src/state/aqe/mod.rs
+++ b/ballista/scheduler/src/state/aqe/mod.rs
@@ -42,6 +42,14 @@ use std::sync::Arc;
use std::time::{SystemTime, UNIX_EPOCH};
use std::vec;
+// TODO: the AQE planner runs DataFusion's DefaultPhysicalPlanner with a
+// list of PhysicalOptimizerRules and never goes through
+// DefaultDistributedPlanner::plan_query_stages_internal, so neither
+// maybe_promote_to_broadcast nor the HashJoinExec(CollectLeft) shuffle
+// lowering fire here. Joins that would broadcast under the default
+// planner stay on the Partitioned shuffle path. Move the lowering into
+// an AQE optimizer rule in a follow-up PR.
+
mod adapter;
mod execution_plan;
pub mod optimizer_rule;
diff --git a/examples/Cargo.toml b/examples/Cargo.toml
index 18fac834f..c328261de 100644
--- a/examples/Cargo.toml
+++ b/examples/Cargo.toml
@@ -33,6 +33,11 @@ name = "standalone_sql"
path = "examples/standalone-sql.rs"
required-features = ["ballista/standalone"]
+[[example]]
+name = "standalone_broadcast_join"
+path = "examples/standalone-broadcast-join.rs"
+required-features = ["ballista/standalone"]
+
[[example]]
name = "mtls-cluster"
path = "examples/mtls-cluster.rs"
diff --git a/examples/examples/standalone-broadcast-join.rs b/examples/examples/standalone-broadcast-join.rs
new file mode 100644
index 000000000..f32c229d0
--- /dev/null
+++ b/examples/examples/standalone-broadcast-join.rs
@@ -0,0 +1,185 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Standalone broadcast-join correctness example.
+//!
+//! Builds two Parquet-backed tables, runs a hash join in a standalone Ballista
+//! cluster with the broadcast-join threshold enabled, and asserts the result
+//! matches a plain DataFusion reference run.
+//!
+//! Parquet files are used (not CSV) because `DataSourceExec` for Parquet
+//! reports `total_byte_size` from row-group metadata, which is required for
+//! `maybe_promote_to_broadcast` to fire. CSV sources report `Absent` for
+//! `total_byte_size` and would silently skip promotion.
+
+use std::sync::Arc;
+
+use ballista::datafusion::{
+ arrow::{
+ array::Int32Array,
+ datatypes::{DataType, Field, Schema},
+ record_batch::RecordBatch,
+ },
+ common::Result,
+ execution::{SessionStateBuilder, options::ParquetReadOptions},
+ prelude::{SessionConfig, SessionContext},
+};
+use ballista::prelude::{SessionConfigExt, SessionContextExt};
+
+#[tokio::main]
+async fn main() -> Result<()> {
+ // ── Build shared in-memory batches ─────────────────────────────────────
+
+ let big_schema = Arc::new(Schema::new(vec![
+ Field::new("k", DataType::Int32, false),
+ Field::new("v", DataType::Int32, false),
+ ]));
+ let small_schema = Arc::new(Schema::new(vec![
+ Field::new("k", DataType::Int32, false),
+ Field::new("name", DataType::Int32, false),
+ ]));
+
+ let big_batch = RecordBatch::try_new(
+ big_schema.clone(),
+ vec![
+ Arc::new(Int32Array::from((0..1000).collect::<Vec<i32>>())),
+ Arc::new(Int32Array::from(
+ (0..1000).map(|i: i32| i * 2).collect::<Vec<i32>>(),
+ )),
+ ],
+ )?;
+
+ let small_batch = RecordBatch::try_new(
+ small_schema,
+ vec![
+ Arc::new(Int32Array::from(vec![1i32, 2, 3, 4, 5])),
+ Arc::new(Int32Array::from(vec![10i32, 20, 30, 40, 50])),
+ ],
+ )?;
+
+ const QUERY: &str = "select count(*) from big join small on big.k = small.k";
+
+ // ── Reference run: plain DataFusion ───────────────────────────────────
+
+ let ref_ctx = SessionContext::new();
+ ref_ctx.register_batch("big", big_batch.clone())?;
+ ref_ctx.register_batch("small", small_batch.clone())?;
+
+ let ref_batches = ref_ctx.sql(QUERY).await?.collect().await?;
+ let ref_count: i64 = ref_batches[0]
+ .column(0)
+ .as_any()
+ .downcast_ref::<ballista::datafusion::arrow::array::Int64Array>()
+ .expect("count(*) should produce Int64Array")
+ .value(0);
+
+ // ── Write Parquet files so Ballista can serialize the table references ─
+ //
+ // Parquet is required here (not CSV) because DataFusion's Parquet reader
+ // populates `Statistics::total_byte_size` from row-group metadata.
+ // The Ballista broadcast-join optimizer (`supports_collect_by_thresholds`)
+ // checks `total_byte_size` first; when it is `Absent` (as with CSV) the
+ // check falls back to `num_rows`, and when that is also unavailable the
+ // promotion is silently skipped.
+
+ let tmp = tempfile::tempdir().expect("tmp dir");
+ let big_path = tmp.path().join("big.parquet");
+ let small_path = tmp.path().join("small.parquet");
+
+ // Use a plain (non-Ballista) context to write the Parquet files.
+ let write_ctx = SessionContext::new();
+ write_ctx
+ .read_batch(big_batch.clone())?
+ .write_parquet(
+ big_path.to_str().expect("valid path"),
+ Default::default(),
+ None,
+ )
+ .await?;
+ write_ctx
+ .read_batch(small_batch.clone())?
+ .write_parquet(
+ small_path.to_str().expect("valid path"),
+ Default::default(),
+ None,
+ )
+ .await?;
+
+ // ── Ballista standalone run ────────────────────────────────────────────
+
+ let config = SessionConfig::new_with_ballista()
+ .with_target_partitions(2)
+ .with_ballista_standalone_parallelism(2)
+ // Enable the Ballista broadcast-join promotion.
+ // new_with_ballista() already disables DataFusion's own single-partition
+ // threshold via ballista_restricted_configuration(), so we only need
+ // to ensure the Ballista threshold is set to a value large enough that
+ // the small table (5 rows, ~300 B as Parquet) is promoted.
+ .with_ballista_broadcast_join_threshold_bytes(10 * 1024 * 1024);
+
+ let state = SessionStateBuilder::new()
+ .with_config(config)
+ .with_default_features()
+ .build();
+
+ let ctx = SessionContext::standalone_with_state(state).await?;
+
+ ctx.register_parquet(
+ "big",
+ big_path.to_str().expect("valid path"),
+ ParquetReadOptions::default(),
+ )
+ .await?;
+ ctx.register_parquet(
+ "small",
+ small_path.to_str().expect("valid path"),
+ ParquetReadOptions::default(),
+ )
+ .await?;
+
+ // Print the physical plan so we can confirm broadcast lowering.
+ // Because DataFusion's hash_join_single_partition_threshold is set to 0
+ // by new_with_ballista(), the plan arrives at Ballista as Partitioned.
+ // The Ballista JoinSelection optimizer then sees the Parquet byte stats
+ // for the small table and promotes it to a broadcast (CollectLeft) join.
+ let physical_plan = ctx.sql(QUERY).await?.create_physical_plan().await?;
+ println!(
+ "Physical plan:\n{}",
+ ballista::datafusion::physical_plan::displayable(physical_plan.as_ref())
+ .indent(true)
+ );
+
+ let result_batches = ctx.sql(QUERY).await?.collect().await?;
+ let result_count: i64 = result_batches[0]
+ .column(0)
+ .as_any()
+ .downcast_ref::<ballista::datafusion::arrow::array::Int64Array>()
+ .expect("count(*) should produce Int64Array")
+ .value(0);
+
+ assert_eq!(
+ result_count, ref_count,
+ "Ballista result ({result_count}) does not match reference ({ref_count})"
+ );
+
+ println!(
+ "ok: broadcast join produced {result_count} rows (matches reference); \
+ threshold=10 MiB, small table written as Parquet so byte stats are present"
+ );
+
+ Ok(())
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org