From ce1ab9c297da72adacb0a621133fbb6509858373 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 16 Jul 2026 10:22:04 -0600 Subject: [PATCH] fix: restore input ordering a SortMergeJoinExec requires after lowering Lowering a plan into stages rewrites nodes, and a rewrite can drop an ordering its parent depends on. Promoting a `SortMergeJoinExec` to a broadcast `HashJoinExec(CollectLeft)` produces a join that does not preserve input order, so a parent `SortMergeJoinExec` reading that output merged unsorted input. A merge join given unsorted input does not error: it stops matching once the cursors diverge, silently dropping rows and returning a wrong answer. Re-establish any hard input ordering a rewritten child no longer carries. Sorts are only added, never removed, so an input that already satisfies its requirement is left untouched. Also validate the planned stages: every `SortMergeJoinExec` input must satisfy the join's required ordering. An unsorted merge-join input is always a bug, and failing loudly is better than returning wrong results. The check reports no false positives across the full SF1 TPC-DS gate under both `prefer_hash_join=true` and `false`. TPC-DS q4 and q78 now match single-process DataFusion and join the correctness gate, taking it from 86 to 88 queries. q78 was previously filed as "distributed LIMIT drops rows": the wrong answer does not in fact depend on `ORDER BY` or `LIMIT` and reproduces with neither. --- ballista/scheduler/src/planner.rs | 239 +++++++++++++++++++++++++++++- benchmarks/src/bin/tpcds.rs | 14 +- 2 files changed, 241 insertions(+), 12 deletions(-) diff --git a/ballista/scheduler/src/planner.rs b/ballista/scheduler/src/planner.rs index 4448a8eb12..daf84b1320 100644 --- a/ballista/scheduler/src/planner.rs +++ b/ballista/scheduler/src/planner.rs @@ -34,6 +34,9 @@ use ballista_core::{ }; use datafusion::arrow::datatypes::DataType; use datafusion::config::ConfigOptions; +use datafusion::physical_expr::{ + LexOrdering, LexRequirement, OrderingRequirements, PhysicalSortRequirement, +}; use datafusion::physical_optimizer::PhysicalOptimizerRule; use datafusion::physical_optimizer::enforce_sorting::EnforceSorting; use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec; @@ -55,6 +58,120 @@ use crate::physical_optimizer::join_selection::{ type PartialQueryStageResult = (Arc, Vec>); +/// Returns the `Hard` ordering requirements `parent` places on the child at +/// `index`, or `None` when the child's ordering is unconstrained. `Soft` +/// requirements are opportunistic and are deliberately ignored. +fn hard_ordering_requirement( + parent: &Arc, + index: usize, +) -> Option> { + match parent.required_input_ordering().into_iter().nth(index) { + Some(Some(OrderingRequirements::Hard(reqs))) => Some(reqs), + _ => None, + } +} + +/// Whether `child` already carries one of the orderings in `reqs`. +fn ordering_is_satisfied( + child: &Arc, + reqs: &[LexRequirement], +) -> Result { + for req in reqs { + if child + .properties() + .equivalence_properties() + .ordering_satisfy_requirement(req.clone())? + { + return Ok(true); + } + } + Ok(false) +} + +/// Re-establishes any input ordering a rewritten child no longer carries. +/// +/// Lowering a plan into stages rewrites nodes, and a rewrite can drop an +/// ordering its parent depends on: promoting a `SortMergeJoinExec` to a +/// broadcast `HashJoinExec(CollectLeft)` yields a join that does not preserve +/// input order, so a parent `SortMergeJoinExec` reading it would merge unsorted +/// input and silently drop rows. Sorts are only added, never removed, so a +/// child that already satisfies its requirement is returned untouched. +fn restore_required_input_ordering( + parent: &Arc, + children: Vec>, +) -> Result>> { + let mut restored = Vec::with_capacity(children.len()); + for (index, child) in children.into_iter().enumerate() { + let Some(reqs) = hard_ordering_requirement(parent, index) else { + restored.push(child); + continue; + }; + if ordering_is_satisfied(&child, &reqs)? { + restored.push(child); + continue; + } + // Enforce the first alternative, dropping requirements on expressions + // that are constant for this child and so already trivially ordered. + let Some(req) = reqs.into_iter().next() else { + restored.push(child); + continue; + }; + let mut sort_reqs: Vec = req.into(); + sort_reqs.retain(|sort_expr| { + child + .properties() + .equivalence_properties() + .is_expr_constant(&sort_expr.expr) + .is_none() + }); + let Some(ordering) = + LexOrdering::new(sort_reqs.into_iter().map(Into::into).collect::>()) + else { + restored.push(child); + continue; + }; + debug!( + "restoring ordering {ordering} required by {} on input {index}", + parent.name(), + ); + let mut sort = SortExec::new(ordering, Arc::clone(&child)); + if child.properties().output_partitioning().partition_count() > 1 { + sort = sort.with_preserve_partitioning(true); + } + restored.push(Arc::new(sort)); + } + Ok(restored) +} + +/// Verifies that every `SortMergeJoinExec` in `plan` has inputs carrying the +/// ordering the join requires. +/// +/// A merge join reads both inputs as sorted streams; given unsorted input it +/// does not error, it simply stops matching once the cursors diverge and +/// returns a wrong answer. Checking the planned stages turns that silent +/// corruption into a loud failure. +pub fn validate_sort_merge_join_orderings(plan: &Arc) -> Result<()> { + if plan.downcast_ref::().is_some() { + for (index, child) in plan.children().into_iter().enumerate() { + let Some(reqs) = hard_ordering_requirement(plan, index) else { + continue; + }; + if !ordering_is_satisfied(child, &reqs)? { + return Err(BallistaError::Internal(format!( + "SortMergeJoinExec input {index} does not satisfy its required \ + ordering; merging unsorted input would silently drop rows. \ + Required one of {reqs:?}, input plan:\n{}", + datafusion::physical_plan::displayable(child.as_ref()).indent(false) + ))); + } + } + } + for child in plan.children() { + validate_sort_merge_join_orderings(child)?; + } + Ok(()) +} + /// Trait for breaking an execution plan into distributed query stages. /// /// The planner creates a DAG of stages where each stage can be executed @@ -122,6 +239,11 @@ impl DistributedPlanner for DefaultDistributedPlanner { None, config, )?); + for stage in &stages { + for child in stage.children() { + validate_sort_merge_join_orderings(child)?; + } + } Ok(stages) } } @@ -195,6 +317,7 @@ impl DefaultDistributedPlanner { children.push(new_child); stages.append(&mut child_stages); } + let children = restore_required_input_ordering(&execution_plan, children)?; if let Some(_coalesce) = execution_plan.downcast_ref::() { let input = children[0].clone(); @@ -721,7 +844,10 @@ pub(crate) fn create_shuffle_writer_with_config( #[cfg(test)] mod test { use crate::assert_plan; - use crate::planner::{DefaultDistributedPlanner, DistributedPlanner}; + use crate::planner::{ + DefaultDistributedPlanner, DistributedPlanner, restore_required_input_ordering, + validate_sort_merge_join_orderings, + }; use crate::test_utils::datafusion_test_context; use ballista_core::error::BallistaError; use ballista_core::execution_plans::{SortShuffleWriterExec, UnresolvedShuffleExec}; @@ -1855,4 +1981,115 @@ order by (proto).try_into_physical_plan(ctx, codec.physical_extension_codec())?; Ok(result_exec_plan) } + + /// Lowering can replace a `SortMergeJoinExec` with a broadcast + /// `HashJoinExec`, which does not preserve input order. A parent + /// `SortMergeJoinExec` reading that output must have its ordering restored, + /// or it merges unsorted input and silently drops rows. + #[tokio::test] + async fn restores_ordering_for_sort_merge_join_input() -> Result<(), BallistaError> { + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::common::{JoinType, NullEquality}; + use datafusion::physical_plan::empty::EmptyExec; + use datafusion::physical_plan::joins::SortMergeJoinExec; + + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let unsorted: Arc = + Arc::new(EmptyExec::new(Arc::clone(&schema))); + let on = vec![( + Arc::new(Column::new("a", 0)) as _, + Arc::new(Column::new("a", 0)) as _, + )]; + let smj: Arc = Arc::new(SortMergeJoinExec::try_new( + Arc::clone(&unsorted), + Arc::clone(&unsorted), + on, + None, + JoinType::Left, + vec![SortOptions::default()], + NullEquality::NullEqualsNothing, + )?); + + // Both inputs lack the join's required ordering, so both gain a sort. + let restored = restore_required_input_ordering( + &smj, + vec![Arc::clone(&unsorted), Arc::clone(&unsorted)], + )?; + let rebuilt = smj.clone().with_new_children(restored.clone())?; + assert_plan!( + rebuilt.as_ref(), + @r" + SortMergeJoinExec: join_type=Left, on=[(a@0, a@0)] + SortExec: expr=[a@0 ASC], preserve_partitioning=[false] + EmptyExec + SortExec: expr=[a@0 ASC], preserve_partitioning=[false] + EmptyExec + " + ); + + // An input that already carries the ordering is left untouched rather + // than being wrapped in a second, redundant sort. + let sorted = restored[0].clone(); + let already_ok = restore_required_input_ordering( + &smj, + vec![Arc::clone(&sorted), Arc::clone(&sorted)], + )?; + let rebuilt = smj.with_new_children(already_ok)?; + assert_plan!( + rebuilt.as_ref(), + @r" + SortMergeJoinExec: join_type=Left, on=[(a@0, a@0)] + SortExec: expr=[a@0 ASC], preserve_partitioning=[false] + EmptyExec + SortExec: expr=[a@0 ASC], preserve_partitioning=[false] + EmptyExec + " + ); + Ok(()) + } + + /// The validation rejects a `SortMergeJoinExec` whose inputs are unsorted, + /// so the wrong-results failure mode surfaces as an error. + #[tokio::test] + async fn validate_rejects_unsorted_sort_merge_join_input() -> Result<(), BallistaError> + { + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::common::{JoinType, NullEquality}; + use datafusion::physical_plan::empty::EmptyExec; + use datafusion::physical_plan::joins::SortMergeJoinExec; + + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let unsorted: Arc = + Arc::new(EmptyExec::new(Arc::clone(&schema))); + let on = vec![( + Arc::new(Column::new("a", 0)) as _, + Arc::new(Column::new("a", 0)) as _, + )]; + let smj: Arc = Arc::new(SortMergeJoinExec::try_new( + Arc::clone(&unsorted), + Arc::clone(&unsorted), + on, + None, + JoinType::Left, + vec![SortOptions::default()], + NullEquality::NullEqualsNothing, + )?); + + let err = validate_sort_merge_join_orderings(&smj) + .expect_err("unsorted SortMergeJoinExec input must be rejected"); + assert!( + err.to_string() + .contains("does not satisfy its required ordering"), + "unexpected error: {err}" + ); + + let restored = restore_required_input_ordering( + &smj, + vec![Arc::clone(&unsorted), Arc::clone(&unsorted)], + )?; + let repaired = smj.with_new_children(restored)?; + validate_sort_merge_join_orderings(&repaired) + .expect("restored plan must satisfy the join's ordering requirement"); + Ok(()) + } } diff --git a/benchmarks/src/bin/tpcds.rs b/benchmarks/src/bin/tpcds.rs index ca95b61483..214f2320bb 100644 --- a/benchmarks/src/bin/tpcds.rs +++ b/benchmarks/src/bin/tpcds.rs @@ -62,21 +62,13 @@ const TABLES: &[&str] = &[ /// under the default (static) planner; this list reflects that configuration. /// Remove an entry as the underlying cause is fixed. const SKIP: &[(usize, &str)] = &[ - // Distributed execution diverges from single-process DataFusion on the same - // data (confirmed reproducible; DataFusion is correct under both join modes). - // See https://github.com/apache/datafusion-ballista/issues/2046 - ( - 4, - "distributed result diverges from DataFusion (issue #2046)", - ), + // Distributed set operations diverge from single-process DataFusion on the + // same data (confirmed reproducible; DataFusion is correct under both join + // modes). See https://github.com/apache/datafusion-ballista/issues/2046 ( 38, "distributed INTERSECT diverges from DataFusion (issue #2046)", ), - ( - 78, - "distributed LIMIT drops rows vs DataFusion (issue #2046)", - ), ( 87, "distributed EXCEPT diverges from DataFusion (issue #2046)",