From ae14e824d78ef8e658de6ed99fd318194f35f7b4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20Heres?= Date: Sun, 30 Aug 2026 01:19:29 +0200 Subject: [PATCH 1/3] feat: add single join and use it for correlated scalar subqueries Correlated scalar subqueries must return at most one row per set of outer values. DataFusion enforced that syntactically: the analyzer rejected any correlated scalar subquery that did not have an aggregate on top, with "Correlated scalar subquery must be aggregated to return at most one row". So a lookup like select o_orderkey, (select c_name from customer where c_custkey = o_custkey) from orders did not plan at all, and had to be written with a redundant aggregate (min/max/any_value) whose only job was to prove the row count -- paying a full hash aggregation over the subquery side. This adds the "single join" of Neumann and Kemper's unnesting paper as JoinType::LeftSingle / JoinType::RightSingle: a left/right outer join that emits exactly one row per row of its preserved side and raises "Scalar subquery returned more than one row" when a second row matches. ScalarSubqueryToJoin now emits it for the subqueries whose shape does not already guarantee the property, so those queries decorrelate into a plain join with no aggregate at all. Subqueries that do guarantee it -- an aggregate grouped only by columns the correlated predicate fixes -- keep their plain LEFT JOIN, so every plan that worked before is byte-identical, including TPC-H q2/q17/q20 and TPC-DS q1/q6/q30/q32/q41/q81/q92. Implemented in HashJoinExec (both build- and probe-driven, so the join can still be swapped for build-side selection) and NestedLoopJoinExec. The matched bitmaps double as duplicate detectors, so the check costs one already-cached bit test per matched row and covers matches spread across probe batches and partitions. SortMergeJoinExec and PiecewiseMergeJoin reject single joins so the planner falls back to a join that implements them. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01SgqwvctZdvR1ZCz2hbkEJC --- .../common/src/functional_dependencies.rs | 10 +- datafusion/common/src/join_type.rs | 44 ++++++++- datafusion/core/src/physical_planner.rs | 4 + datafusion/core/tests/dataframe/mod.rs | 3 + .../enforce_distribution.rs | 2 + datafusion/expr/src/logical_plan/builder.rs | 4 +- .../expr/src/logical_plan/invariants.rs | 97 +++++++++---------- datafusion/expr/src/logical_plan/mod.rs | 5 +- datafusion/expr/src/logical_plan/plan.rs | 31 ++++-- datafusion/optimizer/src/eliminate_join.rs | 15 ++- .../optimizer/src/optimize_projections/mod.rs | 7 +- datafusion/optimizer/src/push_down_filter.rs | 34 ++++--- .../optimizer/src/scalar_subquery_to_join.rs | 21 +++- .../physical-expr/src/equivalence/class.rs | 7 +- .../enforce_distribution.rs | 15 +-- .../enforce_sorting/sort_pushdown.rs | 6 +- .../src/projection_pushdown.rs | 7 +- .../physical-plan/src/joins/hash_join/exec.rs | 12 ++- .../src/joins/hash_join/stream.rs | 27 +++++- .../src/joins/nested_loop_join.rs | 48 +++++++-- .../src/joins/piecewise_merge_join/exec.rs | 11 ++- datafusion/physical-plan/src/joins/proto.rs | 4 + .../src/joins/sort_merge_join/exec.rs | 13 ++- .../src/joins/sort_merge_join/filter.rs | 6 ++ .../src/joins/symmetric_hash_join.rs | 4 + datafusion/physical-plan/src/joins/utils.rs | 85 ++++++++++++++-- .../src/operator_statistics/mod.rs | 3 + .../proto/datafusion_common.proto | 2 + .../proto-common/src/generated/pbjson.rs | 6 ++ .../proto-common/src/generated/prost.rs | 6 ++ datafusion/proto-models/src/from_proto.rs | 2 + .../src/generated/datafusion_proto_common.rs | 6 ++ datafusion/proto-models/src/to_proto.rs | 4 + datafusion/sql/src/unparser/plan.rs | 6 ++ .../src/logical_plan/producer/rel/join.rs | 15 ++- 35 files changed, 439 insertions(+), 133 deletions(-) diff --git a/datafusion/common/src/functional_dependencies.rs b/datafusion/common/src/functional_dependencies.rs index e8275aac2da4c..7be98fdb7155e 100644 --- a/datafusion/common/src/functional_dependencies.rs +++ b/datafusion/common/src/functional_dependencies.rs @@ -341,7 +341,11 @@ impl FunctionalDependencies { let mut left_func_dependencies = self.clone(); match join_type { - JoinType::Inner | JoinType::Left | JoinType::Right => { + JoinType::Inner + | JoinType::Left + | JoinType::Right + | JoinType::LeftSingle + | JoinType::RightSingle => { // Add offset to right schema: right_func_dependencies.add_offset(left_cols_len); @@ -351,10 +355,10 @@ impl FunctionalDependencies { right_func_dependencies = right_func_dependencies.with_dependency(Dependency::Multi); - if *join_type == JoinType::Left { + if matches!(join_type, JoinType::Left | JoinType::LeftSingle) { // Downgrade the right side, since it may have additional NULL values: right_func_dependencies.downgrade_dependencies(); - } else if *join_type == JoinType::Right { + } else if matches!(join_type, JoinType::Right | JoinType::RightSingle) { // Downgrade the left side, since it may have additional NULL values: left_func_dependencies.downgrade_dependencies(); } diff --git a/datafusion/common/src/join_type.rs b/datafusion/common/src/join_type.rs index c77a1475ed227..33bc64d45d03b 100644 --- a/datafusion/common/src/join_type.rs +++ b/datafusion/common/src/join_type.rs @@ -72,11 +72,42 @@ pub enum JoinType { /// Same logic as the LeftMark Join above, however it returns a record for each record from the /// right input. RightMark, + /// Left Single Join + /// + /// Returns one record for each record from the left input, padded with the columns of the + /// single matching record from the right input, or with NULLs when there is no match. If a + /// left record matches more than one right record the join fails with an error, because the + /// result would not be a scalar. + /// + /// This is the "single join" of [1] and is used to decorrelate scalar subqueries without + /// forcing an aggregate on top of the subquery to enforce its at-most-one-row property. + /// + /// [1]: http://btw2017.informatik.uni-stuttgart.de/slidesandpapers/F1-10-37/paper_web.pdf + LeftSingle, + /// Right Single Join + /// + /// Same logic as the LeftSingle Join above, however it returns a record for each record from + /// the right input, padded with the single matching record from the left input. + RightSingle, } impl JoinType { pub fn is_outer(self) -> bool { - self == JoinType::Left || self == JoinType::Right || self == JoinType::Full + matches!( + self, + JoinType::Left + | JoinType::Right + | JoinType::Full + | JoinType::LeftSingle + | JoinType::RightSingle + ) + } + + /// Returns true for the "single" join types, which emit at most one row from + /// the non-preserved side per row of the preserved side and error when more + /// than one row matches. + pub fn is_single(self) -> bool { + matches!(self, JoinType::LeftSingle | JoinType::RightSingle) } /// Returns the `JoinType` if the (2) inputs were swapped @@ -94,6 +125,8 @@ impl JoinType { JoinType::RightAnti => JoinType::LeftAnti, JoinType::LeftMark => JoinType::RightMark, JoinType::RightMark => JoinType::LeftMark, + JoinType::LeftSingle => JoinType::RightSingle, + JoinType::RightSingle => JoinType::LeftSingle, } } @@ -123,6 +156,8 @@ impl JoinType { JoinType::RightAnti => (true, false), JoinType::LeftMark => (false, true), JoinType::RightMark => (true, false), + JoinType::LeftSingle => (false, true), + JoinType::RightSingle => (true, false), } } @@ -140,6 +175,8 @@ impl JoinType { | JoinType::RightAnti | JoinType::LeftMark | JoinType::RightMark + | JoinType::LeftSingle + | JoinType::RightSingle ) } @@ -154,6 +191,7 @@ impl JoinType { | JoinType::LeftAnti | JoinType::LeftMark | JoinType::RightSemi + | JoinType::LeftSingle ) } @@ -189,6 +227,8 @@ impl Display for JoinType { JoinType::RightAnti => "RightAnti", JoinType::LeftMark => "LeftMark", JoinType::RightMark => "RightMark", + JoinType::LeftSingle => "LeftSingle", + JoinType::RightSingle => "RightSingle", }; write!(f, "{join_type}") } @@ -210,6 +250,8 @@ impl FromStr for JoinType { "RIGHTANTI" => Ok(JoinType::RightAnti), "LEFTMARK" => Ok(JoinType::LeftMark), "RIGHTMARK" => Ok(JoinType::RightMark), + "LEFTSINGLE" => Ok(JoinType::LeftSingle), + "RIGHTSINGLE" => Ok(JoinType::RightSingle), _ => _not_impl_err!("The join type {s} does not exist or is not implemented"), } } diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index aad392693723b..38c912d491abe 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -1610,6 +1610,10 @@ impl DefaultPhysicalPlanner { | JoinType::RightAnti | JoinType::LeftMark | JoinType::RightMark + // Single joins are only implemented by hash and + // nested loop joins. + | JoinType::LeftSingle + | JoinType::RightSingle ) && session_state .config_options() diff --git a/datafusion/core/tests/dataframe/mod.rs b/datafusion/core/tests/dataframe/mod.rs index 44ab14e6cc137..f6e6e38c1aad9 100644 --- a/datafusion/core/tests/dataframe/mod.rs +++ b/datafusion/core/tests/dataframe/mod.rs @@ -2601,6 +2601,9 @@ async fn verify_join_output_partitioning() -> Result<()> { let join_schema = physical_plan.schema(); match join_type { + JoinType::LeftSingle | JoinType::RightSingle => { + unreachable!("single joins are not part of this test's join types") + } JoinType::Left | JoinType::LeftSemi | JoinType::LeftAnti diff --git a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs index 489076331bbf5..398ad93f8fde6 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs +++ b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs @@ -1493,11 +1493,13 @@ fn multi_hash_joins() -> Result<()> { assert_plan!(plan_distrib, plan_sort); } JoinType::RightSemi | JoinType::RightAnti | JoinType::RightMark => {} + JoinType::LeftSingle | JoinType::RightSingle => {} } match join_type { + JoinType::LeftSingle | JoinType::RightSingle => {} JoinType::Inner | JoinType::Left | JoinType::Right diff --git a/datafusion/expr/src/logical_plan/builder.rs b/datafusion/expr/src/logical_plan/builder.rs index ef5e496b0d7de..ef466c72f4456 100644 --- a/datafusion/expr/src/logical_plan/builder.rs +++ b/datafusion/expr/src/logical_plan/builder.rs @@ -1703,7 +1703,7 @@ pub fn build_join_schema( .collect::>(); left_fields.into_iter().chain(right_fields).collect() } - JoinType::Left => { + JoinType::Left | JoinType::LeftSingle => { // left then right, right set to nullable in case of not matched scenario let left_fields = left_fields .map(|(q, f)| (q.cloned(), Arc::clone(f))) @@ -1713,7 +1713,7 @@ pub fn build_join_schema( .chain(nullify_fields(right_fields)) .collect() } - JoinType::Right => { + JoinType::Right | JoinType::RightSingle => { // left then right, left set to nullable in case of not matched scenario let right_fields = right_fields .map(|(q, f)| (q.cloned(), Arc::clone(f))) diff --git a/datafusion/expr/src/logical_plan/invariants.rs b/datafusion/expr/src/logical_plan/invariants.rs index f36653694c21d..cad7f171c619a 100644 --- a/datafusion/expr/src/logical_plan/invariants.rs +++ b/datafusion/expr/src/logical_plan/invariants.rs @@ -169,31 +169,11 @@ pub fn check_subquery_expr( subquery.subquery.schema().field_names().join(", ") ); } - // Correlated scalar subquery must be aggregated to return at most one row + // A correlated scalar subquery must return at most one row per set of + // outer values. Subqueries whose shape does not guarantee that are + // still valid: `ScalarSubqueryToJoin` decorrelates them with a single + // join, which raises an error at runtime if a second row shows up. if !subquery.outer_ref_columns.is_empty() { - match strip_inner_query(inner_plan) { - LogicalPlan::Aggregate(agg) => { - check_aggregation_in_scalar_subquery(inner_plan, agg) - } - LogicalPlan::Filter(Filter { input, .. }) - if matches!(input.as_ref(), LogicalPlan::Aggregate(_)) => - { - if let LogicalPlan::Aggregate(agg) = input.as_ref() { - check_aggregation_in_scalar_subquery(inner_plan, agg) - } else { - Ok(()) - } - } - _ => { - if inner_plan.max_rows().is_some_and(|max_row| max_row <= 1) { - Ok(()) - } else { - plan_err!( - "Correlated scalar subquery must be aggregated to return at most one row" - ) - } - } - }?; match outer_plan { LogicalPlan::Projection(_) | LogicalPlan::Filter(_) => Ok(()), LogicalPlan::Aggregate(Aggregate { @@ -323,14 +303,16 @@ fn check_inner_plan(inner_plan: &LogicalPlan) -> Result<()> { JoinType::Left | JoinType::LeftSemi | JoinType::LeftAnti - | JoinType::LeftMark => { + | JoinType::LeftMark + | JoinType::LeftSingle => { check_inner_plan(left)?; check_no_outer_references(right) } JoinType::Right | JoinType::RightSemi | JoinType::RightAnti - | JoinType::RightMark => { + | JoinType::RightMark + | JoinType::RightSingle => { check_no_outer_references(left)?; check_inner_plan(right) } @@ -358,35 +340,48 @@ fn check_no_outer_references(inner_plan: &LogicalPlan) -> Result<()> { } } -fn check_aggregation_in_scalar_subquery( +/// Returns true when the shape of a correlated scalar subquery already +/// guarantees it produces at most one row per set of outer values. +/// +/// That is the case when the subquery aggregates and groups only by columns the +/// correlated predicate already fixes, or when it cannot return more than one +/// row at all. Such a subquery can be decorrelated with a plain `LEFT JOIN`; +/// the rest need a single join, which enforces the property at runtime. +pub fn correlated_scalar_subquery_yields_single_row( inner_plan: &LogicalPlan, - agg: &Aggregate, -) -> Result<()> { +) -> Result { + let agg = match strip_inner_query(inner_plan) { + LogicalPlan::Aggregate(agg) => agg, + LogicalPlan::Filter(Filter { input, .. }) => match input.as_ref() { + LogicalPlan::Aggregate(agg) => agg, + _ => return Ok(inner_plan.max_rows().is_some_and(|rows| rows <= 1)), + }, + _ => return Ok(inner_plan.max_rows().is_some_and(|rows| rows <= 1)), + }; + + // A `GROUP BY` with no aggregate is a `DISTINCT`, which can still return + // one row per distinct group. if agg.aggr_expr.is_empty() { - return plan_err!( - "Correlated scalar subquery must be aggregated to return at most one row" - ); + return Ok(false); } - if !agg.group_expr.is_empty() { - let correlated_exprs = get_correlated_expressions(inner_plan)?; - let inner_subquery_cols = - collect_subquery_cols(&correlated_exprs, agg.input.schema())?; - let mut group_columns = agg - .group_expr - .iter() - .map(|group| Ok(group.column_refs().into_iter().cloned().collect::>())) - .collect::>>()? - .into_iter() - .flatten(); - - if !group_columns.all(|group| inner_subquery_cols.contains(&group)) { - // Group BY columns must be a subset of columns in the correlated expressions - return plan_err!( - "A GROUP BY clause in a scalar correlated subquery cannot contain non-correlated columns" - ); - } + if agg.group_expr.is_empty() { + return Ok(true); } - Ok(()) + + // Grouping by anything the correlated predicate does not fix can produce + // several groups -- and so several rows -- for one set of outer values. + let correlated_exprs = get_correlated_expressions(inner_plan)?; + let inner_subquery_cols = + collect_subquery_cols(&correlated_exprs, agg.input.schema())?; + let mut group_columns = agg + .group_expr + .iter() + .map(|group| Ok(group.column_refs().into_iter().cloned().collect::>())) + .collect::>>()? + .into_iter() + .flatten(); + + Ok(group_columns.all(|group| inner_subquery_cols.contains(&group))) } fn strip_inner_query(inner_plan: &LogicalPlan) -> &LogicalPlan { diff --git a/datafusion/expr/src/logical_plan/mod.rs b/datafusion/expr/src/logical_plan/mod.rs index 4766c3f33379f..ac87000fdb702 100644 --- a/datafusion/expr/src/logical_plan/mod.rs +++ b/datafusion/expr/src/logical_plan/mod.rs @@ -21,7 +21,10 @@ pub mod display; pub mod dml; mod extension; pub(crate) mod invariants; -pub use invariants::{InvariantLevel, assert_expected_schema, check_subquery_expr}; +pub use invariants::{ + InvariantLevel, assert_expected_schema, check_subquery_expr, + correlated_scalar_subquery_yields_single_row, +}; mod plan; mod statement; pub mod tree_node; diff --git a/datafusion/expr/src/logical_plan/plan.rs b/datafusion/expr/src/logical_plan/plan.rs index 1a8cd81aa74bd..ec9f992685b68 100644 --- a/datafusion/expr/src/logical_plan/plan.rs +++ b/datafusion/expr/src/logical_plan/plan.rs @@ -554,7 +554,12 @@ impl LogicalPlan { join_type, .. }) => match join_type { - JoinType::Inner | JoinType::Left | JoinType::Right | JoinType::Full => { + JoinType::Inner + | JoinType::Left + | JoinType::Right + | JoinType::Full + | JoinType::LeftSingle + | JoinType::RightSingle => { if left.schema().fields().is_empty() { right.head_output_expr() } else { @@ -1412,12 +1417,16 @@ impl LogicalPlan { (left_max, right_max, _) => Some(left_max * right_max), } } - JoinType::LeftSemi | JoinType::LeftAnti | JoinType::LeftMark => { - left.max_rows() - } - JoinType::RightSemi | JoinType::RightAnti | JoinType::RightMark => { - right.max_rows() - } + // A single join emits exactly one output row per row of the + // preserved side, so it preserves that side's row count. + JoinType::LeftSemi + | JoinType::LeftAnti + | JoinType::LeftMark + | JoinType::LeftSingle => left.max_rows(), + JoinType::RightSemi + | JoinType::RightAnti + | JoinType::RightMark + | JoinType::RightSingle => right.max_rows(), }, LogicalPlan::Repartition(Repartition { input, .. }) => input.max_rows(), LogicalPlan::Union(Union { inputs, .. }) => { @@ -1485,8 +1494,12 @@ impl LogicalPlan { JoinType::Inner if on.is_empty() && filter.is_none() => { left.min_rows().saturating_mul(right.min_rows()) } - JoinType::Left | JoinType::LeftMark => left.min_rows(), - JoinType::Right | JoinType::RightMark => right.min_rows(), + JoinType::Left | JoinType::LeftMark | JoinType::LeftSingle => { + left.min_rows() + } + JoinType::Right | JoinType::RightMark | JoinType::RightSingle => { + right.min_rows() + } JoinType::Full => left.min_rows().max(right.min_rows()), JoinType::Inner | JoinType::LeftSemi diff --git a/datafusion/optimizer/src/eliminate_join.rs b/datafusion/optimizer/src/eliminate_join.rs index 56aa8887065be..9a03f27493950 100644 --- a/datafusion/optimizer/src/eliminate_join.rs +++ b/datafusion/optimizer/src/eliminate_join.rs @@ -510,7 +510,11 @@ fn child_duplicate_insensitivity( JoinType::RightSemi | JoinType::RightAnti | JoinType::RightMark => { (true, duplicate_insensitive) } - JoinType::Left | JoinType::Right | JoinType::Full => (false, false), + JoinType::Left + | JoinType::Right + | JoinType::Full + | JoinType::LeftSingle + | JoinType::RightSingle => (false, false), } } @@ -601,9 +605,12 @@ fn split_join_output_columns( ) -> (LiveColumns, LiveColumns) { let left_len = join.left.schema().fields().len(); match join.join_type { - JoinType::Inner | JoinType::Left | JoinType::Right | JoinType::Full => { - live.split_at(left_len) - } + JoinType::Inner + | JoinType::Left + | JoinType::Right + | JoinType::Full + | JoinType::LeftSingle + | JoinType::RightSingle => live.split_at(left_len), // A semi/anti/mark join outputs only the surviving side's columns, with // the same index space, so `live` passes straight through to that side. JoinType::LeftSemi | JoinType::LeftAnti | JoinType::LeftMark => { diff --git a/datafusion/optimizer/src/optimize_projections/mod.rs b/datafusion/optimizer/src/optimize_projections/mod.rs index 413efd95588d6..9b8f4a68f5926 100644 --- a/datafusion/optimizer/src/optimize_projections/mod.rs +++ b/datafusion/optimizer/src/optimize_projections/mod.rs @@ -765,7 +765,12 @@ fn split_join_requirements( ) -> (RequiredIndices, RequiredIndices) { match join_type { // In these cases requirements are split between left/right children: - JoinType::Inner | JoinType::Left | JoinType::Right | JoinType::Full => { + JoinType::Inner + | JoinType::Left + | JoinType::Right + | JoinType::Full + | JoinType::LeftSingle + | JoinType::RightSingle => { // Decrease right side indices by `left_len` so that they point to valid // positions within the right child: indices.split_off(left_len) diff --git a/datafusion/optimizer/src/push_down_filter.rs b/datafusion/optimizer/src/push_down_filter.rs index 8f72c1a8c311f..01b49b0fe5f4e 100644 --- a/datafusion/optimizer/src/push_down_filter.rs +++ b/datafusion/optimizer/src/push_down_filter.rs @@ -169,8 +169,8 @@ pub struct PushDownFilter {} pub(crate) fn lr_is_preserved(join_type: JoinType) -> (bool, bool) { match join_type { JoinType::Inner => (true, true), - JoinType::Left => (true, false), - JoinType::Right => (false, true), + JoinType::Left | JoinType::LeftSingle => (true, false), + JoinType::Right | JoinType::RightSingle => (false, true), JoinType::Full => (false, false), // No columns from the right side of the join can be referenced in output // predicates for semi/anti joins, so whether we specify t/f doesn't matter. @@ -704,20 +704,22 @@ fn infer_join_predicates_from_on_filters( on_filters, inferred_predicates, ), - JoinType::Left | JoinType::LeftSemi | JoinType::LeftMark => { - infer_join_predicates_impl::( - join_col_keys, - on_filters, - inferred_predicates, - ) - } - JoinType::Right | JoinType::RightSemi | JoinType::RightMark => { - infer_join_predicates_impl::( - join_col_keys, - on_filters, - inferred_predicates, - ) - } + JoinType::Left + | JoinType::LeftSemi + | JoinType::LeftMark + | JoinType::LeftSingle => infer_join_predicates_impl::( + join_col_keys, + on_filters, + inferred_predicates, + ), + JoinType::Right + | JoinType::RightSemi + | JoinType::RightMark + | JoinType::RightSingle => infer_join_predicates_impl::( + join_col_keys, + on_filters, + inferred_predicates, + ), } } diff --git a/datafusion/optimizer/src/scalar_subquery_to_join.rs b/datafusion/optimizer/src/scalar_subquery_to_join.rs index 44011a125ba96..18349dc2f5ebb 100644 --- a/datafusion/optimizer/src/scalar_subquery_to_join.rs +++ b/datafusion/optimizer/src/scalar_subquery_to_join.rs @@ -34,7 +34,10 @@ use datafusion_common::{Column, Result, ScalarValue, assert_or_internal_err, pla use datafusion_expr::expr_rewriter::create_col_from_scalar_expr; use datafusion_expr::logical_plan::{JoinType, Subquery}; use datafusion_expr::utils::conjunction; -use datafusion_expr::{Expr, LogicalPlan, LogicalPlanBuilder, lit, not, when}; +use datafusion_expr::{ + Expr, LogicalPlan, LogicalPlanBuilder, correlated_scalar_subquery_yields_single_row, + lit, not, when, +}; /// Optimizer rule that rewrites scalar subquery filters to joins and places an /// additional projection on top of the filter to preserve the original schema. @@ -386,8 +389,22 @@ fn build_join( // columns. let join_filter = join_filter_opt.or_else(|| Some(lit(true))); + // A subquery whose shape already guarantees at most one row per set of + // outer values joins with a plain `LEFT JOIN`; the rest need a single + // join, which raises an error at runtime when a second row matches. Only + // correlated subqueries can be ambiguous here -- an uncorrelated one is + // checked by `ScalarSubqueryExec` or, when this rule rewrites it, joined + // on `Boolean(true)` against a plan that already returns a single row. + let join_type = if subquery.outer_ref_columns.is_empty() + || correlated_scalar_subquery_yields_single_row(subquery_plan)? + { + JoinType::Left + } else { + JoinType::LeftSingle + }; + let new_plan = LogicalPlanBuilder::from(outer_input.clone()) - .join_on(aliased_subquery, JoinType::Left, join_filter)? + .join_on(aliased_subquery, join_type, join_filter)? .build()?; // Add count-bug compensation for each of the subquery's projected diff --git a/datafusion/physical-expr/src/equivalence/class.rs b/datafusion/physical-expr/src/equivalence/class.rs index 06f384ac2db03..b17fdbceba37d 100644 --- a/datafusion/physical-expr/src/equivalence/class.rs +++ b/datafusion/physical-expr/src/equivalence/class.rs @@ -813,7 +813,12 @@ impl EquivalenceGroup { on: &[(PhysicalExprRef, PhysicalExprRef)], ) -> Result { let group = match join_type { - JoinType::Inner | JoinType::Left | JoinType::Full | JoinType::Right => { + JoinType::Inner + | JoinType::Left + | JoinType::Full + | JoinType::Right + | JoinType::LeftSingle + | JoinType::RightSingle => { let mut result = Self::new( self.iter().cloned().chain( right_equivalences diff --git a/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs index 0368577f9a24f..31bc9e6ab818f 100644 --- a/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs +++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs @@ -169,11 +169,13 @@ pub fn adjust_input_keys_ordering( PartitionMode::CollectLeft => { // Push down requirements to the right side requirements.children[1].data = match join_type { - JoinType::Inner | JoinType::Right => shift_right_required( - &requirements.data, - left.schema().fields().len(), - ) - .unwrap_or_default(), + JoinType::Inner | JoinType::Right | JoinType::RightSingle => { + shift_right_required( + &requirements.data, + left.schema().fields().len(), + ) + .unwrap_or_default() + } JoinType::RightSemi | JoinType::RightAnti | JoinType::RightMark => { requirements.data.clone() } @@ -181,7 +183,8 @@ pub fn adjust_input_keys_ordering( | JoinType::LeftSemi | JoinType::LeftAnti | JoinType::Full - | JoinType::LeftMark => vec![], + | JoinType::LeftMark + | JoinType::LeftSingle => vec![], }; } PartitionMode::Auto => { diff --git a/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs b/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs index 03a28e5f647fc..8a605c7b1c719 100644 --- a/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs +++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs @@ -781,7 +781,9 @@ fn expr_source_side( | JoinType::Right | JoinType::Full | JoinType::LeftMark - | JoinType::RightMark => { + | JoinType::RightMark + | JoinType::LeftSingle + | JoinType::RightSingle => { let eq_group = eqp.eq_group(); let mut right_ordering = ordering.clone(); let (mut valid_left, mut valid_right) = (true, true); @@ -1039,7 +1041,7 @@ fn build_join_column_index(plan: &HashJoinExec) -> Vec { }; match plan.join_type() { - JoinType::Inner | JoinType::Right => { + JoinType::Inner | JoinType::Right | JoinType::RightSingle => { map_fields(plan.left().schema(), JoinSide::Left) .into_iter() .chain(map_fields(plan.right().schema(), JoinSide::Right)) diff --git a/datafusion/physical-optimizer/src/projection_pushdown.rs b/datafusion/physical-optimizer/src/projection_pushdown.rs index fe71c211769c8..98d28edad5321 100644 --- a/datafusion/physical-optimizer/src/projection_pushdown.rs +++ b/datafusion/physical-optimizer/src/projection_pushdown.rs @@ -135,7 +135,12 @@ fn try_push_down_join_filter( let new_lhs_length = lhs_rewrite.data.0.schema().fields.len(); let projections = match projections.as_ref() { None => match join.join_type() { - JoinType::Inner | JoinType::Left | JoinType::Right | JoinType::Full => { + JoinType::Inner + | JoinType::Left + | JoinType::Right + | JoinType::Full + | JoinType::LeftSingle + | JoinType::RightSingle => { // Build projections that ignore the newly projected columns. let mut projections = Vec::new(); projections.extend(0..original_lhs_length); diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index 1b36dabebc564..b0d7a9dedc272 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -1056,6 +1056,7 @@ impl HashJoinExec { | JoinType::RightAnti | JoinType::RightSemi | JoinType::RightMark + | JoinType::RightSingle ), ] } @@ -1125,12 +1126,14 @@ impl HashJoinExec { | JoinType::RightSemi | JoinType::Right | JoinType::RightAnti - | JoinType::RightMark => EmissionType::Incremental, + | JoinType::RightMark + | JoinType::RightSingle => EmissionType::Incremental, // If we need to generate unmatched rows from the *build side*, // we need to emit them at the end. JoinType::Left | JoinType::LeftAnti | JoinType::LeftMark + | JoinType::LeftSingle | JoinType::Full => EmissionType::Both, } } else { @@ -2414,8 +2417,8 @@ mod proto_tests { fn lr_is_preserved(join_type: JoinType) -> (bool, bool) { match join_type { JoinType::Inner => (true, true), - JoinType::Left => (true, false), - JoinType::Right => (false, true), + JoinType::Left | JoinType::LeftSingle => (true, false), + JoinType::Right | JoinType::RightSingle => (false, true), JoinType::Full => (false, false), // Callers restrict the non-output side of semi joins to join-key columns. JoinType::LeftSemi | JoinType::RightSemi => (true, true), @@ -4225,6 +4228,9 @@ mod tests { let num_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); assert_eq!(num_rows, 0, "unexpected rows for {join_type}"); } + JoinType::LeftSingle | JoinType::RightSingle => { + unreachable!("single joins are not part of this test's join types") + } JoinType::Left => { allow_duplicates! { assert_snapshot!(batches_to_sort_string(&batches), @r" diff --git a/datafusion/physical-plan/src/joins/hash_join/stream.rs b/datafusion/physical-plan/src/joins/hash_join/stream.rs index fe6eaff9e53a2..8fccdad0451ed 100644 --- a/datafusion/physical-plan/src/joins/hash_join/stream.rs +++ b/datafusion/physical-plan/src/joins/hash_join/stream.rs @@ -43,7 +43,8 @@ use crate::{ BuildProbeJoinMetrics, ColumnIndex, JoinFilter, JoinHashMapType, StatefulStreamResult, adjust_indices_by_join_type, apply_join_filter_to_indices, build_batch_empty_build_side, build_batch_from_indices, - need_produce_result_in_final, + check_single_join_probe_indices, need_produce_result_in_final, + single_join_too_many_rows_err, }, }; @@ -888,12 +889,30 @@ impl HashJoinStream { (left_indices, right_indices) }; + // A probe-driven single join must not match a probe row twice. + if self.join_type == JoinType::RightSingle { + check_single_join_probe_indices(&right_indices, state.joined_probe_idx)?; + } + // mark joined left-side indices as visited, if required by join type if need_produce_result_in_final(self.join_type) { let mut bitmap = build_side.left_data.visited_indices_bitmap().lock(); - left_indices.iter().flatten().for_each(|x| { - bitmap.set_bit(x as usize, true); - }); + if self.join_type == JoinType::LeftSingle { + // A build-driven single join must not match a build row twice. + // The visited bitmap doubles as the duplicate detector, so this + // also catches matches spread over several probe batches and, + // under `PartitionMode::CollectLeft`, over several partitions. + for index in left_indices.iter().flatten() { + if bitmap.get_bit(index as usize) { + return single_join_too_many_rows_err(); + } + bitmap.set_bit(index as usize, true); + } + } else { + left_indices.iter().flatten().for_each(|x| { + bitmap.set_bit(x as usize, true); + }); + } } // The goals of index alignment for different join types are: diff --git a/datafusion/physical-plan/src/joins/nested_loop_join.rs b/datafusion/physical-plan/src/joins/nested_loop_join.rs index bb91735369b9d..db689188ec2b4 100644 --- a/datafusion/physical-plan/src/joins/nested_loop_join.rs +++ b/datafusion/physical-plan/src/joins/nested_loop_join.rs @@ -18,14 +18,14 @@ //! [`NestedLoopJoinExec`]: joins without equijoin (equality predicates). use std::fmt::Formatter; -use std::ops::{BitOr, ControlFlow}; +use std::ops::{BitAnd, BitOr, ControlFlow}; use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; use std::task::Poll; use super::utils::{ asymmetric_join_output_partitioning, need_produce_result_in_final, - reorder_output_after_swap, swap_join_projection, + reorder_output_after_swap, single_join_too_many_rows_err, swap_join_projection, }; use crate::common::can_project; use crate::execution_plan::{EmissionType, boundedness_from_children}; @@ -386,12 +386,14 @@ impl NestedLoopJoinExec { | JoinType::RightSemi | JoinType::Right | JoinType::RightAnti - | JoinType::RightMark => EmissionType::Incremental, + | JoinType::RightMark + | JoinType::RightSingle => EmissionType::Incremental, // If we need to generate unmatched rows from the *build side*, // we need to emit them at the end. JoinType::Left | JoinType::LeftAnti | JoinType::LeftMark + | JoinType::LeftSingle | JoinType::Full => EmissionType::Both, } } else { @@ -2590,6 +2592,17 @@ impl NestedLoopJoinStream { None }; + // A single join must not match a row of its preserved side twice. The + // matched bitmaps double as the duplicate detectors: a bit that is + // already set means an earlier pair matched the same row, whether in + // this left range, an earlier one, or -- for the left bitmap, which is + // shared -- another partition. + let single_side = match self.join_type { + JoinType::LeftSingle => Some(JoinSide::Left), + JoinType::RightSingle => Some(JoinSide::Right), + _ => None, + }; + // Set the matched bit for left and right side bitmap for (i, is_matched) in bitmap_combined.iter().enumerate() { let is_matched = is_matched.ok_or_else(|| { @@ -2602,6 +2615,9 @@ impl NestedLoopJoinStream { if let Some(bitmap) = left_bitmap.as_mut() && is_matched { + if single_side == Some(JoinSide::Left) && bitmap.get_bit(l_index) { + return single_join_too_many_rows_err(); + } // Map local index back to absolute left index within the batch bitmap.set_bit(l_index, true); } @@ -2609,6 +2625,9 @@ impl NestedLoopJoinStream { if let Some(bitmap) = local_right_bitmap.as_mut() && is_matched { + if single_side == Some(JoinSide::Right) && bitmap.get_bit(r_index) { + return single_join_too_many_rows_err(); + } bitmap.set_bit(r_index, true); } } @@ -2630,6 +2649,13 @@ impl NestedLoopJoinStream { ) })? .finish(); + // A right row already matched by an earlier left range must not + // match again in this one. + if single_side == Some(JoinSide::Right) + && buf.bitand(¤t_right_bitmap).count_set_bits() > 0 + { + return single_join_too_many_rows_err(); + } let updated_global_right_bitmap = buf.bitor(¤t_right_bitmap); self.current_right_batch_matched = @@ -3153,10 +3179,12 @@ fn build_unmatched_batch_empty_schema( | JoinType::Right | JoinType::Full | JoinType::LeftAnti - | JoinType::RightAnti => batch_bitmap.false_count(), + | JoinType::RightAnti + | JoinType::LeftSingle + | JoinType::RightSingle => batch_bitmap.false_count(), JoinType::LeftSemi | JoinType::RightSemi => batch_bitmap.true_count(), JoinType::LeftMark | JoinType::RightMark => batch_bitmap.len(), - _ => unreachable!(), + JoinType::Inner => unreachable!("inner joins have no unmatched rows"), }; if output_schema.fields().is_empty() { @@ -3242,11 +3270,15 @@ fn build_unmatched_batch( } match join_type { - JoinType::Full | JoinType::Right | JoinType::Left => { - if join_type == JoinType::Right { + JoinType::Full + | JoinType::Right + | JoinType::Left + | JoinType::LeftSingle + | JoinType::RightSingle => { + if matches!(join_type, JoinType::Right | JoinType::RightSingle) { debug_assert_eq!(batch_side, JoinSide::Right); } - if join_type == JoinType::Left { + if matches!(join_type, JoinType::Left | JoinType::LeftSingle) { debug_assert_eq!(batch_side, JoinSide::Left); } diff --git a/datafusion/physical-plan/src/joins/piecewise_merge_join/exec.rs b/datafusion/physical-plan/src/joins/piecewise_merge_join/exec.rs index efaf051f7a169..529ebdf02eb11 100644 --- a/datafusion/physical-plan/src/joins/piecewise_merge_join/exec.rs +++ b/datafusion/physical-plan/src/joins/piecewise_merge_join/exec.rs @@ -293,6 +293,11 @@ impl PiecewiseMergeJoinExec { // Left Semi/Anti are handled by `ExistencePWMJStream` (the marked side is // already the buffered side, so no input swap is needed). Right existence joins // and Mark joins are not yet supported. + if join_type.is_single() { + return not_impl_err!( + "Join type {join_type} is currently not supported for PiecewiseMergeJoin" + ); + } if is_existence_join(join_type) && !is_supported_existence_join(join_type) { return not_impl_err!( "Existence join {join_type} is currently not supported for PiecewiseMergeJoin" @@ -405,11 +410,13 @@ impl PiecewiseMergeJoinExec { | JoinType::Full | JoinType::RightSemi | JoinType::RightAnti - | JoinType::RightMark => JoinSide::Right, + | JoinType::RightMark + | JoinType::RightSingle => JoinSide::Right, JoinType::Left | JoinType::LeftAnti | JoinType::LeftSemi - | JoinType::LeftMark => JoinSide::Left, + | JoinType::LeftMark + | JoinType::LeftSingle => JoinSide::Left, } } diff --git a/datafusion/physical-plan/src/joins/proto.rs b/datafusion/physical-plan/src/joins/proto.rs index e91e00f89737a..d0715a38418c6 100644 --- a/datafusion/physical-plan/src/joins/proto.rs +++ b/datafusion/physical-plan/src/joins/proto.rs @@ -45,6 +45,8 @@ pub(crate) fn join_type_to_proto(join_type: JoinType) -> protobuf::JoinType { JoinType::RightAnti => protobuf::JoinType::Rightanti, JoinType::LeftMark => protobuf::JoinType::Leftmark, JoinType::RightMark => protobuf::JoinType::Rightmark, + JoinType::LeftSingle => protobuf::JoinType::Leftsingle, + JoinType::RightSingle => protobuf::JoinType::Rightsingle, } } @@ -62,6 +64,8 @@ pub(crate) fn join_type_from_proto(value: i32, plan_name: &str) -> Result JoinType::RightAnti, protobuf::JoinType::Leftmark => JoinType::LeftMark, protobuf::JoinType::Rightmark => JoinType::RightMark, + protobuf::JoinType::Leftsingle => JoinType::LeftSingle, + protobuf::JoinType::Rightsingle => JoinType::RightSingle, }) } diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs index 911eca0a97928..a91eca0f4bbe7 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs @@ -53,7 +53,7 @@ use datafusion_common::project_schema; use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::{ JoinSide, JoinType, NullEquality, Result, assert_eq_or_internal_err, internal_err, - plan_err, + not_impl_err, plan_err, }; use datafusion_execution::TaskContext; use datafusion_execution::memory_pool::MemoryConsumer; @@ -154,6 +154,11 @@ impl SortMergeJoinExec { sort_options: Vec, null_equality: NullEquality, ) -> Result { + if join_type.is_single() { + return not_impl_err!( + "Join type {join_type} is currently not supported for SortMergeJoinExec" + ); + } let left_schema = left.schema(); let right_schema = right.schema(); @@ -243,13 +248,15 @@ impl SortMergeJoinExec { JoinType::Right | JoinType::RightSemi | JoinType::RightAnti - | JoinType::RightMark => JoinSide::Right, + | JoinType::RightMark + | JoinType::RightSingle => JoinSide::Right, JoinType::Inner | JoinType::Left | JoinType::Full | JoinType::LeftAnti | JoinType::LeftSemi - | JoinType::LeftMark => JoinSide::Left, + | JoinType::LeftMark + | JoinType::LeftSingle => JoinSide::Left, } } diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/filter.rs b/datafusion/physical-plan/src/joins/sort_merge_join/filter.rs index 4fc6cccaa8838..9d0041020032c 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/filter.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/filter.rs @@ -290,6 +290,9 @@ pub fn get_corrected_filter_mask( | JoinType::RightAnti => { unreachable!("Semi/anti/mark joins are handled by BitwiseSortMergeJoinStream") } + JoinType::LeftSingle | JoinType::RightSingle => { + unreachable!("Single joins are rejected by SortMergeJoinExec::try_new") + } JoinType::Inner => None, } } @@ -383,6 +386,9 @@ pub fn filter_record_batch_by_join_type( | JoinType::RightMark => unreachable!( "Semi/anti/mark joins are handled by SemiAntiMarkSortMergeJoinStream" ), + JoinType::LeftSingle | JoinType::RightSingle => { + unreachable!("Single joins are rejected by SortMergeJoinExec::try_new") + } JoinType::Inner => Ok(filter_record_batch(record_batch, corrected_mask)?), } } diff --git a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs index aec693e839718..6daa750a20e3e 100644 --- a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs +++ b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs @@ -713,6 +713,8 @@ impl ExecutionPlan for SymmetricHashJoinExec { JoinType::RightAnti => protobuf::JoinType::Rightanti, JoinType::LeftMark => protobuf::JoinType::Leftmark, JoinType::RightMark => protobuf::JoinType::Rightmark, + JoinType::LeftSingle => protobuf::JoinType::Leftsingle, + JoinType::RightSingle => protobuf::JoinType::Rightsingle, }; let null_equality = match null_equality { NullEquality::NullEqualsNothing => protobuf::NullEquality::NullEqualsNothing, @@ -860,6 +862,8 @@ impl SymmetricHashJoinExec { protobuf::JoinType::Rightanti => JoinType::RightAnti, protobuf::JoinType::Leftmark => JoinType::LeftMark, protobuf::JoinType::Rightmark => JoinType::RightMark, + protobuf::JoinType::Leftsingle => JoinType::LeftSingle, + protobuf::JoinType::Rightsingle => JoinType::RightSingle, }; let null_equality = match protobuf::NullEquality::try_from(*null_equality) .map_err(|_| { diff --git a/datafusion/physical-plan/src/joins/utils.rs b/datafusion/physical-plan/src/joins/utils.rs index 5e8f3c5de929b..abb3595ea6324 100644 --- a/datafusion/physical-plan/src/joins/utils.rs +++ b/datafusion/physical-plan/src/joins/utils.rs @@ -67,7 +67,7 @@ use datafusion_common::stats::Precision; use datafusion_common::utils::memory::RecordBatchMemoryCounter; use datafusion_common::utils::normalize_float_zero; use datafusion_common::{ - DataFusionError, JoinSide, JoinType, NullEquality, Result, SharedResult, + DataFusionError, JoinSide, JoinType, NullEquality, Result, SharedResult, exec_err, internal_datafusion_err, not_impl_err, plan_err, }; use datafusion_expr::interval_arithmetic::Interval; @@ -253,6 +253,8 @@ fn output_join_field(old_field: &Field, join_type: &JoinType, is_left: bool) -> JoinType::RightAnti => false, // doesn't introduce nulls (or can it??) JoinType::LeftMark => false, JoinType::RightMark => false, + JoinType::LeftSingle => !is_left, // right input is padded with nulls + JoinType::RightSingle => is_left, // left input is padded with nulls }; if force_nullable { @@ -303,7 +305,12 @@ pub fn build_join_schema( }; let (fields, column_indices): (SchemaBuilder, Vec) = match join_type { - JoinType::Inner | JoinType::Left | JoinType::Full | JoinType::Right => { + JoinType::Inner + | JoinType::Left + | JoinType::Full + | JoinType::Right + | JoinType::LeftSingle + | JoinType::RightSingle => { // left then right left_fields().chain(right_fields()).unzip() } @@ -516,7 +523,12 @@ fn estimate_join_cardinality( .unzip::<_, _, Vec<_>, Vec<_>>(); match join_type { - JoinType::Inner | JoinType::Left | JoinType::Right | JoinType::Full => { + JoinType::Inner + | JoinType::Left + | JoinType::Right + | JoinType::Full + | JoinType::LeftSingle + | JoinType::RightSingle => { let ij_cardinality = estimate_inner_join_cardinality( Statistics { num_rows: left_stats.num_rows, @@ -538,6 +550,9 @@ fn estimate_join_cardinality( JoinType::Inner => ij_cardinality, JoinType::Left => ij_cardinality.max(&left_stats.num_rows), JoinType::Right => ij_cardinality.max(&right_stats.num_rows), + // A single join emits exactly one row per preserved-side row. + JoinType::LeftSingle => left_stats.num_rows, + JoinType::RightSingle => right_stats.num_rows, JoinType::Full => ij_cardinality .max(&left_stats.num_rows) .add(&ij_cardinality.max(&right_stats.num_rows)) @@ -1171,6 +1186,7 @@ pub(crate) fn need_produce_right_in_final(join_type: JoinType) -> bool { | JoinType::RightAnti | JoinType::RightMark | JoinType::RightSemi + | JoinType::RightSingle ) } @@ -1179,6 +1195,46 @@ pub(crate) fn need_produce_right_in_final(join_type: JoinType) -> bool { /// /// For example of the `Left` join, in each iteration of right side, can get the matched result, but need /// to maintain the matched indices bit map to get the unmatched row for the left side. +/// The error a "single" join raises when a row of its preserved side matches more +/// than one row of the other side. +/// +/// Single joins exist to evaluate scalar subqueries, so this matches the error +/// [`ScalarSubqueryExec`] raises for the uncorrelated case. +/// +/// [`ScalarSubqueryExec`]: crate::scalar_subquery::ScalarSubqueryExec +pub(crate) fn single_join_too_many_rows_err() -> Result { + exec_err!("Scalar subquery returned more than one row") +} + +/// Verifies that a probe-driven single join (`RightSingle`) matched each probe +/// row at most once. +/// +/// `probe_indices` holds the matched probe-side indices of one chunk of a probe +/// batch, in ascending order, so a second match for a probe row is either +/// adjacent to the first or -- when a probe row's matches straddle a chunk +/// boundary -- equal to `last_joined_probe_idx`, the last probe index joined by +/// the preceding chunk of the same batch. +pub(crate) fn check_single_join_probe_indices( + probe_indices: &UInt32Array, + last_joined_probe_idx: Option, +) -> Result<()> { + debug_assert_eq!( + probe_indices.null_count(), + 0, + "expected matched probe indices to have no nulls" + ); + let indices = probe_indices.values(); + if let (Some(first), Some(last)) = (indices.first(), last_joined_probe_idx) + && *first as usize == last + { + return single_join_too_many_rows_err(); + } + if indices.windows(2).any(|pair| pair[0] == pair[1]) { + return single_join_too_many_rows_err(); + } + Ok(()) +} + pub(crate) fn need_produce_result_in_final(join_type: JoinType) -> bool { matches!( join_type, @@ -1186,6 +1242,7 @@ pub(crate) fn need_produce_result_in_final(join_type: JoinType) -> bool { | JoinType::LeftAnti | JoinType::LeftSemi | JoinType::LeftMark + | JoinType::LeftSingle | JoinType::Full ) } @@ -1442,12 +1499,12 @@ pub(crate) fn adjust_indices_by_join_type( // matched Ok((left_indices, right_indices)) } - JoinType::Left => { + JoinType::Left | JoinType::LeftSingle => { // matched Ok((left_indices, right_indices)) // unmatched left row will be produced in the end of loop, and it has been set in the left visited bitmap } - JoinType::Right => { + JoinType::Right | JoinType::RightSingle => { // combine the matched and unmatched right result together append_right_indices( left_indices, @@ -1904,12 +1961,17 @@ pub(crate) fn symmetric_join_output_partitioning( let left_partitioning = left.output_partitioning(); let right_partitioning = right.output_partitioning(); let result = match join_type { - JoinType::Left | JoinType::LeftSemi | JoinType::LeftAnti | JoinType::LeftMark => { - left_partitioning.clone() - } + JoinType::Left + | JoinType::LeftSemi + | JoinType::LeftAnti + | JoinType::LeftMark + | JoinType::LeftSingle => left_partitioning.clone(), JoinType::RightSemi | JoinType::RightAnti | JoinType::RightMark => { right_partitioning.clone() } + JoinType::RightSingle => { + adjust_right_output_partitioning(right_partitioning, left_columns_len)? + } JoinType::Inner | JoinType::Right => { adjust_right_output_partitioning(right_partitioning, left_columns_len)? } @@ -1938,9 +2000,14 @@ pub(crate) fn asymmetric_join_output_partitioning( | JoinType::LeftSemi | JoinType::LeftAnti | JoinType::Full - | JoinType::LeftMark => Partitioning::UnknownPartitioning( + | JoinType::LeftMark + | JoinType::LeftSingle => Partitioning::UnknownPartitioning( right.output_partitioning().partition_count(), ), + JoinType::RightSingle => adjust_right_output_partitioning( + right.output_partitioning(), + left.schema().fields().len(), + )?, }; Ok(result) } diff --git a/datafusion/physical-plan/src/operator_statistics/mod.rs b/datafusion/physical-plan/src/operator_statistics/mod.rs index b19f4e5fd4693..0ca2e272305b8 100644 --- a/datafusion/physical-plan/src/operator_statistics/mod.rs +++ b/datafusion/physical-plan/src/operator_statistics/mod.rs @@ -936,6 +936,9 @@ impl StatisticsProvider for JoinStatisticsProvider { } JoinType::LeftMark => left_rows, JoinType::RightMark => right_rows, + // A single join emits exactly one row per preserved-side row. + JoinType::LeftSingle => left_rows, + JoinType::RightSingle => right_rows, }; // NL join inner with exact inputs is an exact Cartesian product; diff --git a/datafusion/proto-common/proto/datafusion_common.proto b/datafusion/proto-common/proto/datafusion_common.proto index 27d1101036d9b..b8c6746cad3c1 100644 --- a/datafusion/proto-common/proto/datafusion_common.proto +++ b/datafusion/proto-common/proto/datafusion_common.proto @@ -88,6 +88,8 @@ enum JoinType { RIGHTANTI = 7; LEFTMARK = 8; RIGHTMARK = 9; + LEFTSINGLE = 10; + RIGHTSINGLE = 11; } enum JoinConstraint { diff --git a/datafusion/proto-common/src/generated/pbjson.rs b/datafusion/proto-common/src/generated/pbjson.rs index c222cd1cb8687..2ad9776a445d0 100644 --- a/datafusion/proto-common/src/generated/pbjson.rs +++ b/datafusion/proto-common/src/generated/pbjson.rs @@ -4922,6 +4922,8 @@ impl serde::Serialize for JoinType { Self::Rightanti => "RIGHTANTI", Self::Leftmark => "LEFTMARK", Self::Rightmark => "RIGHTMARK", + Self::Leftsingle => "LEFTSINGLE", + Self::Rightsingle => "RIGHTSINGLE", }; serializer.serialize_str(variant) } @@ -4943,6 +4945,8 @@ impl<'de> serde::Deserialize<'de> for JoinType { "RIGHTANTI", "LEFTMARK", "RIGHTMARK", + "LEFTSINGLE", + "RIGHTSINGLE", ]; struct GeneratedVisitor; @@ -4993,6 +4997,8 @@ impl<'de> serde::Deserialize<'de> for JoinType { "RIGHTANTI" => Ok(JoinType::Rightanti), "LEFTMARK" => Ok(JoinType::Leftmark), "RIGHTMARK" => Ok(JoinType::Rightmark), + "LEFTSINGLE" => Ok(JoinType::Leftsingle), + "RIGHTSINGLE" => Ok(JoinType::Rightsingle), _ => Err(serde::de::Error::unknown_variant(value, FIELDS)), } } diff --git a/datafusion/proto-common/src/generated/prost.rs b/datafusion/proto-common/src/generated/prost.rs index bdbe38538e1d7..e275ddb99e233 100644 --- a/datafusion/proto-common/src/generated/prost.rs +++ b/datafusion/proto-common/src/generated/prost.rs @@ -1053,6 +1053,8 @@ pub enum JoinType { Rightanti = 7, Leftmark = 8, Rightmark = 9, + Leftsingle = 10, + Rightsingle = 11, } impl JoinType { /// String value of the enum field names used in the ProtoBuf definition. @@ -1071,6 +1073,8 @@ impl JoinType { Self::Rightanti => "RIGHTANTI", Self::Leftmark => "LEFTMARK", Self::Rightmark => "RIGHTMARK", + Self::Leftsingle => "LEFTSINGLE", + Self::Rightsingle => "RIGHTSINGLE", } } /// Creates an enum from field names used in the ProtoBuf definition. @@ -1086,6 +1090,8 @@ impl JoinType { "RIGHTANTI" => Some(Self::Rightanti), "LEFTMARK" => Some(Self::Leftmark), "RIGHTMARK" => Some(Self::Rightmark), + "LEFTSINGLE" => Some(Self::Leftsingle), + "RIGHTSINGLE" => Some(Self::Rightsingle), _ => None, } } diff --git a/datafusion/proto-models/src/from_proto.rs b/datafusion/proto-models/src/from_proto.rs index 74ead8c52049b..a77a38f21b262 100644 --- a/datafusion/proto-models/src/from_proto.rs +++ b/datafusion/proto-models/src/from_proto.rs @@ -166,6 +166,8 @@ impl From for JoinType { protobuf::JoinType::Rightanti => JoinType::RightAnti, protobuf::JoinType::Leftmark => JoinType::LeftMark, protobuf::JoinType::Rightmark => JoinType::RightMark, + protobuf::JoinType::Leftsingle => JoinType::LeftSingle, + protobuf::JoinType::Rightsingle => JoinType::RightSingle, } } } diff --git a/datafusion/proto-models/src/generated/datafusion_proto_common.rs b/datafusion/proto-models/src/generated/datafusion_proto_common.rs index bdbe38538e1d7..e275ddb99e233 100644 --- a/datafusion/proto-models/src/generated/datafusion_proto_common.rs +++ b/datafusion/proto-models/src/generated/datafusion_proto_common.rs @@ -1053,6 +1053,8 @@ pub enum JoinType { Rightanti = 7, Leftmark = 8, Rightmark = 9, + Leftsingle = 10, + Rightsingle = 11, } impl JoinType { /// String value of the enum field names used in the ProtoBuf definition. @@ -1071,6 +1073,8 @@ impl JoinType { Self::Rightanti => "RIGHTANTI", Self::Leftmark => "LEFTMARK", Self::Rightmark => "RIGHTMARK", + Self::Leftsingle => "LEFTSINGLE", + Self::Rightsingle => "RIGHTSINGLE", } } /// Creates an enum from field names used in the ProtoBuf definition. @@ -1086,6 +1090,8 @@ impl JoinType { "RIGHTANTI" => Some(Self::Rightanti), "LEFTMARK" => Some(Self::Leftmark), "RIGHTMARK" => Some(Self::Rightmark), + "LEFTSINGLE" => Some(Self::Leftsingle), + "RIGHTSINGLE" => Some(Self::Rightsingle), _ => None, } } diff --git a/datafusion/proto-models/src/to_proto.rs b/datafusion/proto-models/src/to_proto.rs index d1c857a7c5cba..9d5cf355ffe6a 100644 --- a/datafusion/proto-models/src/to_proto.rs +++ b/datafusion/proto-models/src/to_proto.rs @@ -172,6 +172,8 @@ impl From for protobuf::JoinType { JoinType::RightAnti => protobuf::JoinType::Rightanti, JoinType::LeftMark => protobuf::JoinType::Leftmark, JoinType::RightMark => protobuf::JoinType::Rightmark, + JoinType::LeftSingle => protobuf::JoinType::Leftsingle, + JoinType::RightSingle => protobuf::JoinType::Rightsingle, } } } @@ -266,6 +268,8 @@ mod tests { JoinType::RightAnti, JoinType::LeftMark, JoinType::RightMark, + JoinType::LeftSingle, + JoinType::RightSingle, ] { assert_eq!( JoinType::from(protobuf::JoinType::from(join_type)), diff --git a/datafusion/sql/src/unparser/plan.rs b/datafusion/sql/src/unparser/plan.rs index 9922509a0e609..ec4c0bf9bea7c 100644 --- a/datafusion/sql/src/unparser/plan.rs +++ b/datafusion/sql/src/unparser/plan.rs @@ -1441,6 +1441,9 @@ impl Unparser<'_> { select.projection(projection); } } + JoinType::LeftSingle | JoinType::RightSingle => { + return not_impl_err!("Unparsing of Single join type"); + } JoinType::Inner | JoinType::Left | JoinType::Right @@ -2413,6 +2416,9 @@ impl Unparser<'_> { JoinType::LeftMark | JoinType::RightMark => { unimplemented!("Unparsing of Mark join type") } + JoinType::LeftSingle | JoinType::RightSingle => { + return not_impl_err!("Unparsing of Single join type"); + } }) } diff --git a/datafusion/substrait/src/logical_plan/producer/rel/join.rs b/datafusion/substrait/src/logical_plan/producer/rel/join.rs index 9094774780e10..0f053c14423ea 100644 --- a/datafusion/substrait/src/logical_plan/producer/rel/join.rs +++ b/datafusion/substrait/src/logical_plan/producer/rel/join.rs @@ -36,7 +36,7 @@ pub fn from_join( let left = producer.handle_plan(join.left.as_ref())?; let right = producer.handle_plan(join.right.as_ref())?; - let join_type = to_substrait_jointype(join.join_type); + let join_type = to_substrait_jointype(join.join_type)?; let join_expr = to_substrait_join_expr(join.on.clone(), join.null_equality, join.filter.clone()); @@ -79,8 +79,10 @@ fn to_substrait_join_expr( conjunction(all_conditions) } -fn to_substrait_jointype(join_type: JoinType) -> join_rel::JoinType { - match join_type { +fn to_substrait_jointype( + join_type: JoinType, +) -> datafusion::common::Result { + Ok(match join_type { JoinType::Inner => join_rel::JoinType::Inner, JoinType::Left => join_rel::JoinType::Left, JoinType::Right => join_rel::JoinType::Right, @@ -91,7 +93,12 @@ fn to_substrait_jointype(join_type: JoinType) -> join_rel::JoinType { JoinType::RightMark => join_rel::JoinType::RightMark, JoinType::RightAnti => join_rel::JoinType::RightAnti, JoinType::RightSemi => join_rel::JoinType::RightSemi, - } + // Substrait has no equivalent of the single join, which fails at + // runtime when more than one row matches. + JoinType::LeftSingle | JoinType::RightSingle => { + return not_impl_err!("join type: `{join_type}`"); + } + }) } #[cfg(test)] From 0e14ea3f76ef219cc0b56de14f7d3c58d6ccca66 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20Heres?= Date: Sun, 30 Aug 2026 02:27:37 +0200 Subject: [PATCH 2/3] test: cover the single join, and keep the LIMIT case's plan error Adds the single-join section of `subquery.slt` -- attribute lookup with and without a match, the more-than-one-row error, the same in a filter rather than a projection, the nested-loop path for a correlation with no equijoin key, and a plan assertion that a provably-single-row subquery still gets a plain left join -- plus `HashJoinExec` and `NestedLoopJoinExec` unit tests covering both directions across partition modes and batch sizes. Writing them found two gaps, both fixed here: - `projection_pushdown` and `NestedLoopJoinExec::build_unmatched_batch` hit catch-all arms for the new join types, panicking on a correlation with no equijoin key. - `NestedLoopJoinStream::update_matched_bitmap`, the path taken when the right batch is large relative to the batch size, set the matched bitmaps without checking them, so a second match went unreported. Decorrelation still cannot pull a correlated predicate through a `LIMIT`, and those subqueries no longer hit the "must be aggregated" check that used to report them. They now say so directly instead of reaching the physical planner as a leftover `ScalarSubquery` expression. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01SgqwvctZdvR1ZCz2hbkEJC --- .../expr/src/logical_plan/invariants.rs | 36 ++ .../optimizer/src/scalar_subquery_to_join.rs | 5 + .../physical-plan/src/joins/hash_join/exec.rs | 148 +++++- .../src/joins/nested_loop_join.rs | 91 ++++ datafusion/physical-plan/src/reuse.rs | 478 ++++++++++++++++++ .../sqllogictest/test_files/subquery.slt | 155 +++++- 6 files changed, 906 insertions(+), 7 deletions(-) create mode 100644 datafusion/physical-plan/src/reuse.rs diff --git a/datafusion/expr/src/logical_plan/invariants.rs b/datafusion/expr/src/logical_plan/invariants.rs index cad7f171c619a..64ca3e53e986a 100644 --- a/datafusion/expr/src/logical_plan/invariants.rs +++ b/datafusion/expr/src/logical_plan/invariants.rs @@ -173,7 +173,18 @@ pub fn check_subquery_expr( // outer values. Subqueries whose shape does not guarantee that are // still valid: `ScalarSubqueryToJoin` decorrelates them with a single // join, which raises an error at runtime if a second row shows up. + // + // Decorrelation works by pulling the correlated predicate above the + // subquery, which a row-limiting operator in between makes unsound. + // Such a subquery is only usable when it returns a single row anyway. if !subquery.outer_ref_columns.is_empty() { + if correlation_is_below_limit(inner_plan) + && inner_plan.max_rows().is_none_or(|rows| rows > 1) + { + return plan_err!( + "Correlated scalar subquery with a LIMIT must be limited to a single row" + ); + } match outer_plan { LogicalPlan::Projection(_) | LogicalPlan::Filter(_) => Ok(()), LogicalPlan::Aggregate(Aggregate { @@ -340,6 +351,31 @@ fn check_no_outer_references(inner_plan: &LogicalPlan) -> Result<()> { } } +/// Returns true when a `LIMIT` sits above an outer reference in the subquery, +/// which stops [`PullUpCorrelatedExpr`] from pulling the correlated predicate +/// above the subquery. +/// +/// [`PullUpCorrelatedExpr`]: https://docs.rs/datafusion-optimizer/latest/datafusion_optimizer/decorrelate/struct.PullUpCorrelatedExpr.html +fn correlation_is_below_limit(inner_plan: &LogicalPlan) -> bool { + let mut found = false; + inner_plan + .apply(|plan| { + Ok( + if matches!(plan, LogicalPlan::Limit(_)) + && !plan.all_out_ref_exprs().is_empty() + { + found = true; + TreeNodeRecursion::Stop + } else { + TreeNodeRecursion::Continue + }, + ) + }) + // the closure always returns Ok + .expect("infallible"); + found +} + /// Returns true when the shape of a correlated scalar subquery already /// guarantees it produces at most one row per set of outer values. /// diff --git a/datafusion/optimizer/src/scalar_subquery_to_join.rs b/datafusion/optimizer/src/scalar_subquery_to_join.rs index 18349dc2f5ebb..499d5a6a64942 100644 --- a/datafusion/optimizer/src/scalar_subquery_to_join.rs +++ b/datafusion/optimizer/src/scalar_subquery_to_join.rs @@ -395,6 +395,11 @@ fn build_join( // correlated subqueries can be ambiguous here -- an uncorrelated one is // checked by `ScalarSubqueryExec` or, when this rule rewrites it, joined // on `Boolean(true)` against a plan that already returns a single row. + // + // The single join only sees the outer rows that reach it, so an outer row + // a pushed-down filter removed cannot raise the error. Whether a + // more-than-one-row subquery is reported therefore depends on the plan, + // the same way it does in other engines. let join_type = if subquery.outer_ref_columns.is_empty() || correlated_scalar_subquery_yields_single_row(subquery_plan)? { diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index b0d7a9dedc272..d6a10d54a1f81 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -4159,6 +4159,124 @@ mod tests { Ok(()) } + /// A single join emits one row per row of its preserved side, padded with + /// nulls where nothing matched -- the same output a left/right join gives + /// when the other side is unique on the join keys. + #[rstest] + #[tokio::test] + async fn join_single_at_most_one_match( + #[values(JoinType::LeftSingle, JoinType::RightSingle)] join_type: JoinType, + #[values(PartitionMode::CollectLeft, PartitionMode::Partitioned)] + partition_mode: PartitionMode, + #[values(1, 8192)] batch_size: usize, + ) -> Result<()> { + let task_ctx = prepare_task_ctx(batch_size, false); + // 7 has no match on the right, 6 has no match on the left. + let left = build_table( + ("a1", &vec![1, 2, 3]), + ("b1", &vec![4, 5, 7]), + ("c1", &vec![7, 8, 9]), + ); + let right = build_table( + ("a2", &vec![10, 20, 30]), + ("b2", &vec![4, 5, 6]), + ("c2", &vec![70, 80, 90]), + ); + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b2", &right.schema())?) as _, + )]; + + let (columns, batches, _) = join_collect_with_partition_mode( + left, + right, + on, + &join_type, + partition_mode, + NullEquality::NullEqualsNothing, + task_ctx, + ) + .await?; + + assert_eq!(columns, vec!["a1", "b1", "c1", "a2", "b2", "c2"]); + + // LeftSingle keeps every left row, RightSingle every right row. + if join_type == JoinType::LeftSingle { + allow_duplicates! { + assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+----+----+----+----+----+ + | a1 | b1 | c1 | a2 | b2 | c2 | + +----+----+----+----+----+----+ + | 1 | 4 | 7 | 10 | 4 | 70 | + | 2 | 5 | 8 | 20 | 5 | 80 | + | 3 | 7 | 9 | | | | + +----+----+----+----+----+----+ + "); + } + } else { + allow_duplicates! { + assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+----+----+----+----+----+ + | a1 | b1 | c1 | a2 | b2 | c2 | + +----+----+----+----+----+----+ + | | | | 30 | 6 | 90 | + | 1 | 4 | 7 | 10 | 4 | 70 | + | 2 | 5 | 8 | 20 | 5 | 80 | + +----+----+----+----+----+----+ + "); + } + } + + Ok(()) + } + + /// A second match for a row of the preserved side makes the join fail, + /// however the matches are spread over batches and partitions. + #[rstest] + #[tokio::test] + async fn join_single_rejects_second_match( + #[values(JoinType::LeftSingle, JoinType::RightSingle)] join_type: JoinType, + #[values(PartitionMode::CollectLeft, PartitionMode::Partitioned)] + partition_mode: PartitionMode, + #[values(1, 8192)] batch_size: usize, + ) -> Result<()> { + let task_ctx = prepare_task_ctx(batch_size, false); + // Key 4 matches twice, whichever side is preserved. + let left = build_table( + ("a1", &vec![1, 2]), + ("b1", &vec![4, 4]), + ("c1", &vec![7, 8]), + ); + let right = build_table( + ("a2", &vec![10, 20]), + ("b2", &vec![4, 4]), + ("c2", &vec![70, 80]), + ); + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b2", &right.schema())?) as _, + )]; + + let err = join_collect_with_partition_mode( + left, + right, + on, + &join_type, + partition_mode, + NullEquality::NullEqualsNothing, + task_ctx, + ) + .await + .unwrap_err(); + + assert_contains!( + err.to_string(), + "Scalar subquery returned more than one row" + ); + + Ok(()) + } + /// Under NullEqualsNothing, NULL join keys are not inserted into the hash /// map, so a build side whose keys are all NULL produces an empty map even /// though it contains rows. Join types that emit unmatched build rows must @@ -4193,6 +4311,8 @@ mod tests { JoinType::RightAnti, JoinType::LeftMark, JoinType::RightMark, + JoinType::LeftSingle, + JoinType::RightSingle, ] { let (_, batches, metrics) = join_collect_with_partition_mode( Arc::clone(&left), @@ -4228,8 +4348,32 @@ mod tests { let num_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); assert_eq!(num_rows, 0, "unexpected rows for {join_type}"); } - JoinType::LeftSingle | JoinType::RightSingle => { - unreachable!("single joins are not part of this test's join types") + // Nothing matches, so the single joins degenerate to their + // outer counterparts. + JoinType::LeftSingle => { + allow_duplicates! { + assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+----+----+----+ + | a1 | b1 | a2 | b1 | + +----+----+----+----+ + | 1 | | | | + | 2 | | | | + +----+----+----+----+ + "); + } + } + JoinType::RightSingle => { + allow_duplicates! { + assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+----+----+----+ + | a1 | b1 | a2 | b1 | + +----+----+----+----+ + | | | 10 | 4 | + | | | 20 | | + | | | 30 | 6 | + +----+----+----+----+ + "); + } } JoinType::Left => { allow_duplicates! { diff --git a/datafusion/physical-plan/src/joins/nested_loop_join.rs b/datafusion/physical-plan/src/joins/nested_loop_join.rs index db689188ec2b4..6413d875e6a6d 100644 --- a/datafusion/physical-plan/src/joins/nested_loop_join.rs +++ b/datafusion/physical-plan/src/joins/nested_loop_join.rs @@ -2946,6 +2946,14 @@ impl NestedLoopJoinStream { // 1. Maybe update the left bitmap if need_produce_result_in_final(self.join_type) && r_matched_bitmap.has_true() { let mut bitmap = left_data.bitmap().lock(); + // A single join must not match a left row twice, whether inside this + // right batch or against an earlier one -- and, since the bitmap is + // shared, whether in this partition or another. + if self.join_type == JoinType::LeftSingle + && (r_matched_bitmap.true_count() > 1 || bitmap.get_bit(l_index)) + { + return single_join_too_many_rows_err(); + } bitmap.set_bit(l_index, true); } @@ -2959,6 +2967,14 @@ impl NestedLoopJoinStream { })?; let (buf, nulls) = right_bitmap.into_parts(); debug_assert!(nulls.is_none()); + // Likewise, a bit already set means an earlier left row matched the + // same right row. + if self.join_type == JoinType::RightSingle + && buf.bitand(r_matched_bitmap.values()).count_set_bits() > 0 + { + self.current_right_batch_matched = Some(BooleanArray::new(buf, None)); + return single_join_too_many_rows_err(); + } let updated_right_bitmap = buf.bitor(r_matched_bitmap.values()); self.current_right_batch_matched = @@ -3743,6 +3759,81 @@ pub(crate) mod tests { Ok(()) } + /// A single join emits one row per row of its preserved side; the filter + /// here leaves at most one match for each, so nothing is rejected. + #[rstest] + #[tokio::test] + async fn join_single_with_filter( + #[values(JoinType::LeftSingle, JoinType::RightSingle)] join_type: JoinType, + #[values(1, 2, 16)] batch_size: usize, + ) -> Result<()> { + let task_ctx = new_task_ctx(batch_size); + let left = build_left_table(); + let right = build_right_table(); + + let filter = prepare_join_filter(); + let (columns, batches, _) = multi_partitioned_join_collect( + left, + right, + &join_type, + Some(filter), + task_ctx, + ) + .await?; + assert_eq!(columns, vec!["a1", "b1", "c1", "a2", "b2", "c2"]); + + if join_type == JoinType::LeftSingle { + allow_duplicates!(assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+----+-----+----+----+----+ + | a1 | b1 | c1 | a2 | b2 | c2 | + +----+----+-----+----+----+----+ + | 11 | 8 | 110 | | | | + | 5 | 5 | 50 | 2 | 2 | 80 | + | 9 | 8 | 90 | | | | + +----+----+-----+----+----+----+ + ")); + } else { + allow_duplicates!(assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+----+----+----+----+-----+ + | a1 | b1 | c1 | a2 | b2 | c2 | + +----+----+----+----+----+-----+ + | | | | 10 | 10 | 100 | + | | | | 12 | 10 | 40 | + | 5 | 5 | 50 | 2 | 2 | 80 | + +----+----+----+----+----+-----+ + ")); + } + + Ok(()) + } + + /// Without a filter every row of one side matches every row of the other, + /// so a single join rejects the second match. + #[rstest] + #[tokio::test] + async fn join_single_rejects_second_match( + #[values(JoinType::LeftSingle, JoinType::RightSingle)] join_type: JoinType, + #[values(1, 2, 16)] batch_size: usize, + ) -> Result<()> { + let task_ctx = new_task_ctx(batch_size); + let err = multi_partitioned_join_collect( + build_left_table(), + build_right_table(), + &join_type, + None, + task_ctx, + ) + .await + .unwrap_err(); + + assert_contains!( + err.to_string(), + "Scalar subquery returned more than one row" + ); + + Ok(()) + } + #[rstest] #[tokio::test] async fn join_right_with_filter(#[values(1, 2, 16)] batch_size: usize) -> Result<()> { diff --git a/datafusion/physical-plan/src/reuse.rs b/datafusion/physical-plan/src/reuse.rs new file mode 100644 index 0000000000000..096e2eac5062b --- /dev/null +++ b/datafusion/physical-plan/src/reuse.rs @@ -0,0 +1,478 @@ +// 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. + +//! [`ReuseExec`]: execute a subplan once and distribute it to several consumers. + +use std::pin::Pin; +use std::sync::Arc; +use std::task::{Context, Poll, Waker}; + +use crate::execution_plan::{CardinalityEffect, EvaluationType, SchedulingType}; +use crate::metrics::{ExecutionPlanMetricsSet, MetricsSet}; +use crate::statistics::{ChildStats, StatisticsArgs}; +use crate::{ + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, + ExecutionPlanProperties, PlanProperties, RecordBatchStream, ReplaceChildrenOptions, + SendableRecordBatchStream, Statistics, validate_child_count, +}; + +use arrow::datatypes::SchemaRef; +use arrow::record_batch::RecordBatch; +use datafusion_common::tree_node::TreeNodeRecursion; +use datafusion_common::{DataFusionError, Result, internal_err}; +use datafusion_common_runtime::SpawnedTask; +use datafusion_execution::TaskContext; +use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; +use datafusion_physical_expr::PhysicalExpr; + +use futures::{Stream, StreamExt}; +use parking_lot::Mutex; + +/// Executes its input **once** and distributes the result to every consumer +/// that shares this operator. +/// +/// A plan is a tree, so a subplan appearing in two places is executed twice. +/// When the same `Arc` is installed at both places, the first +/// consumer to call [`ExecutionPlan::execute`] starts the input, and each batch +/// it produces is handed to every consumer. +/// +/// # Retention +/// +/// Batches are not cached wholesale. Each batch is released as soon as all +/// consumers have read it, so consumers that keep pace with each other cost +/// roughly one batch of retention apiece. +/// +/// A consumer that attaches late is the case that costs memory: everything +/// produced before it attaches must be held for it. That is what happens under +/// [`ScalarSubqueryExec`], which runs the subquery to completion before +/// executing the main input, so the whole subplan output is retained. Bounding +/// the buffer instead would deadlock there — the producer would block on a +/// consumer that cannot start until the producer has finished. +/// +/// # Sharing +/// +/// Sharing is by `Arc` identity: two separately constructed `ReuseExec`s over +/// equal inputs share nothing. Rewriting a plan through +/// [`ExecutionPlan::with_new_children`] rebuilds the operator and drops the +/// sharing — the result stays correct, it just recomputes. +/// +/// [`ScalarSubqueryExec`]: crate::scalar_subquery::ScalarSubqueryExec +#[derive(Debug)] +pub struct ReuseExec { + /// The subplan to execute once. + input: Arc, + /// How many plan sites share this operator. Used to know when a batch has + /// been seen by everyone and can be dropped. + consumers: usize, + /// Created by whichever consumer executes first. + state: Mutex>>, + cache: Arc, + metrics: ExecutionPlanMetricsSet, +} + +impl ReuseExec { + /// Create a [`ReuseExec`] over `input` shared by `consumers` plan sites. + pub fn new(input: Arc, consumers: usize) -> Self { + let cache = Self::compute_properties(&input); + Self { + input, + consumers, + state: Mutex::new(None), + cache: Arc::new(cache), + metrics: ExecutionPlanMetricsSet::new(), + } + } + + /// The subplan being reused. + pub fn input(&self) -> &Arc { + &self.input + } + + /// Number of plan sites sharing this operator. + pub fn consumers(&self) -> usize { + self.consumers + } + + /// Partitioning, ordering and emission all pass through: batches are + /// forwarded as they are produced, in order, per partition. + fn compute_properties(input: &Arc) -> PlanProperties { + PlanProperties::new( + input.equivalence_properties().clone(), + input.output_partitioning().clone(), + input.pipeline_behavior(), + input.boundedness(), + ) + .with_evaluation_type(EvaluationType::Eager) + .with_scheduling_type(SchedulingType::Cooperative) + } + + /// Start the input on first use; later callers join the running execution. + fn shared_state(&self, context: &Arc) -> Result> { + let mut guard = self.state.lock(); + if let Some(state) = guard.as_ref() { + return Ok(Arc::clone(state)); + } + + let partition_count = self.input.output_partitioning().partition_count(); + let mut logs = Vec::with_capacity(partition_count); + let mut tasks = Vec::with_capacity(partition_count); + + for partition in 0..partition_count { + let reservation = + MemoryConsumer::new(format!("ReuseExec[{partition}]")) + .register(context.memory_pool()); + let log = Arc::new(PartitionLog::new(self.consumers, reservation)); + let stream = self.input.execute(partition, Arc::clone(context))?; + tasks.push(SpawnedTask::spawn(pull_from_input( + Arc::clone(&log), + stream, + ))); + logs.push(log); + } + + let state = Arc::new(ReuseState { + logs, + _tasks: tasks, + }); + *guard = Some(Arc::clone(&state)); + Ok(state) + } +} + +/// The running execution of the input, shared by all consumers. +#[derive(Debug)] +struct ReuseState { + logs: Vec>, + /// Producer tasks; aborted when the last consumer drops the state. + _tasks: Vec>, +} + +/// An append-only log of one partition's batches, read concurrently by every +/// consumer at its own pace. +#[derive(Debug)] +struct PartitionLog { + inner: Mutex, + reservation: MemoryReservation, +} + +#[derive(Debug)] +struct LogState { + /// Produced batches. A slot becomes `None` once every consumer has read it. + batches: Vec>, + /// How many consumers have yet to read each slot. + unread: Vec, + /// Consumers still reading. New batches start with this many readers. + live: usize, + finished: bool, + error: Option>, + /// Consumers parked waiting for the producer. + wakers: Vec, +} + +impl PartitionLog { + fn new(consumers: usize, reservation: MemoryReservation) -> Self { + Self { + inner: Mutex::new(LogState { + batches: Vec::new(), + unread: Vec::new(), + live: consumers, + finished: false, + error: None, + wakers: Vec::new(), + }), + reservation, + } + } + + /// Append a batch for all live consumers. Returns `Err` if the buffer could + /// not be accounted for, which stops the producer. + fn push(&self, batch: RecordBatch) -> Result<()> { + let size = batch.get_array_memory_size(); + if let Err(e) = self.reservation.try_grow(size) { + self.fail(e); + return internal_err!("ReuseExec: memory reservation failed"); + } + let mut state = self.inner.lock(); + let live = state.live; + state.batches.push(Some(batch)); + state.unread.push(live); + // Nobody left to read it; release straight away. + if live == 0 { + let last = state.batches.len() - 1; + state.batches[last] = None; + self.reservation.shrink(size); + } + state.wake_all(); + Ok(()) + } + + fn fail(&self, error: DataFusionError) { + let mut state = self.inner.lock(); + if state.error.is_none() { + state.error = Some(Arc::new(error)); + } + state.finished = true; + state.wake_all(); + } + + fn finish(&self) { + let mut state = self.inner.lock(); + state.finished = true; + state.wake_all(); + } +} + +impl LogState { + fn wake_all(&mut self) { + for waker in self.wakers.drain(..) { + waker.wake(); + } + } +} + +/// Drive one input partition into its log. +async fn pull_from_input(log: Arc, mut stream: SendableRecordBatchStream) { + while let Some(batch) = stream.next().await { + match batch { + Ok(batch) => { + if log.push(batch).is_err() { + return; + } + } + Err(e) => { + log.fail(e); + return; + } + } + } + log.finish(); +} + +/// One consumer's view of a partition log. +struct ReuseStream { + log: Arc, + /// Keeps the producer tasks alive while any consumer is reading. + _state: Arc, + schema: SchemaRef, + cursor: usize, + done: bool, +} + +impl Stream for ReuseStream { + type Item = Result; + + fn poll_next( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll> { + if self.done { + return Poll::Ready(None); + } + let mut state = self.log.inner.lock(); + + if let Some(error) = &state.error { + let error = Arc::clone(error); + drop(state); + self.done = true; + return Poll::Ready(Some(Err(DataFusionError::Shared(error)))); + } + + if self.cursor < state.batches.len() { + let index = self.cursor; + let batch = state.batches[index] + .clone() + .expect("batch released while a consumer still needed it"); + state.unread[index] = state.unread[index].saturating_sub(1); + if state.unread[index] == 0 { + state.batches[index] = None; + self.log.reservation.shrink(batch.get_array_memory_size()); + } + drop(state); + self.cursor += 1; + return Poll::Ready(Some(Ok(batch))); + } + + if state.finished { + drop(state); + self.done = true; + return Poll::Ready(None); + } + + state.wakers.push(cx.waker().clone()); + Poll::Ready(Some(Ok(RecordBatch::new_empty(Arc::clone(&self.schema))))) + } +} + +impl RecordBatchStream for ReuseStream { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +} + +impl Drop for ReuseStream { + fn drop(&mut self) { + // Give up this consumer's claim so retained batches can be released + // even when the stream is abandoned early (a LIMIT upstream, say). + let mut state = self.log.inner.lock(); + state.live = state.live.saturating_sub(1); + let mut freed = 0; + for index in self.cursor..state.batches.len() { + state.unread[index] = state.unread[index].saturating_sub(1); + if state.unread[index] == 0 { + if let Some(batch) = state.batches[index].take() { + freed += batch.get_array_memory_size(); + } + } + } + drop(state); + if freed > 0 { + self.log.reservation.shrink(freed); + } + } +} + +impl DisplayAs for ReuseExec { + fn fmt_as( + &self, + t: DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!(f, "ReuseExec: consumers={}", self.consumers) + } + DisplayFormatType::TreeRender => write!(f, "ReuseExec"), + } + } +} + +impl ExecutionPlan for ReuseExec { + fn name(&self) -> &'static str { + "ReuseExec" + } + + fn properties(&self) -> &Arc { + &self.cache + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + + fn maintains_input_order(&self) -> Vec { + vec![true] + } + + fn benefits_from_input_partitioning(&self) -> Vec { + vec![false] + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + + fn replace_children( + self: Arc, + mut children: Vec>, + options: ReplaceChildrenOptions, + ) -> Result> { + validate_child_count!(self, children); + let input = children.swap_remove(0); + // A rebuilt operator starts a fresh execution, so the sharing the + // optimizer established is lost here. That costs a recomputation, not + // correctness. + match options.children_properties { + ChildrenPropertiesMode::Keep => Ok(Arc::new(Self { + input, + consumers: self.consumers, + state: Mutex::new(None), + cache: Arc::clone(&self.cache), + metrics: ExecutionPlanMetricsSet::new(), + })), + ChildrenPropertiesMode::Recompute => { + Ok(Arc::new(Self::new(input, self.consumers))) + } + } + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn with_new_children_and_same_properties( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + let partition_count = self.input.output_partitioning().partition_count(); + if partition >= partition_count { + return internal_err!( + "ReuseExec invalid partition {partition} (expected less than {partition_count})" + ); + } + + let state = self.shared_state(&context)?; + let log = Arc::clone(&state.logs[partition]); + Ok(Box::pin(ReuseStream { + log, + _state: state, + schema: self.schema(), + cursor: 0, + done: false, + })) + } + + fn metrics(&self) -> Option { + Some(self.metrics.clone_inner()) + } + + fn child_stats_requests(&self, partition: Option) -> Vec { + vec![ChildStats::At(partition)] + } + + /// Distributing changes when rows appear, not which rows or how many. + fn statistics_from_inputs( + &self, + input_stats: &[Arc], + _args: &StatisticsArgs, + ) -> Result> { + Ok(Arc::clone(&input_stats[0])) + } + + fn cardinality_effect(&self) -> CardinalityEffect { + CardinalityEffect::Equal + } +} diff --git a/datafusion/sqllogictest/test_files/subquery.slt b/datafusion/sqllogictest/test_files/subquery.slt index 2a022b947d0f8..03823dfa0bd6f 100644 --- a/datafusion/sqllogictest/test_files/subquery.slt +++ b/datafusion/sqllogictest/test_files/subquery.slt @@ -432,8 +432,20 @@ SELECT t1_id, t1_name, t1_int, (select t2_id, t2_name FROM t2 WHERE t2.t2_id = t statement error DataFusion error: Invalid \(non-executable\) plan after Analyzer\ncaused by\nError during planning: In/Exist/SetComparison subquery can only be used in Projection, Filter, TableScan, Window functions, Aggregate and Join plan nodes, but was used in \[Sort: t1.t1_int IN \(\) ASC NULLS LAST\] SELECT t1_id, t1_name, t1_int FROM t1 order by t1_int in (SELECT t2_int FROM t2 WHERE t1.t1_id > t1.t1_int) +# A correlated scalar subquery with no aggregate decorrelates into a single +# join, which fails when a second row matches. #non_aggregated_correlated_scalar_subquery -statement error DataFusion error: Invalid \(non-executable\) plan after Analyzer\ncaused by\nError during planning: Correlated scalar subquery must be aggregated to return at most one row +query TT +explain SELECT t1_id, (SELECT t2_int FROM t2 WHERE t2.t2_int = t1.t1_int) as t2_int from t1 +---- +logical_plan +01)Projection: t1.t1_id, __scalar_sq_1.t2_int AS t2_int +02)--LeftSingle Join: t1.t1_int = __scalar_sq_1.t2_int +03)----TableScan: t1 projection=[t1_id, t1_int] +04)----SubqueryAlias: __scalar_sq_1 +05)------TableScan: t2 projection=[t2_int] + +statement error DataFusion error: Execution error: Scalar subquery returned more than one row SELECT t1_id, (SELECT t2_int FROM t2 WHERE t2.t2_int = t1.t1_int) as t2_int from t1 #non_aggregated_correlated_scalar_subquery_unique @@ -446,12 +458,21 @@ SELECT t1_id, (SELECT t3_int FROM t3 WHERE t3.t3_id = t1.t1_id) as t3_int from t 44 3 -#non_aggregated_correlated_scalar_subquery -statement error DataFusion error: Invalid \(non-executable\) plan after Analyzer\ncaused by\nError during planning: Correlated scalar subquery must be aggregated to return at most one row +# `GROUP BY` alone does not prove at most one row per outer row, so this also +# takes a single join -- and here the grouping happens to make it hold. +#non_aggregated_correlated_scalar_subquery_with_group_by +query II rowsort SELECT t1_id, (SELECT t2_int FROM t2 WHERE t2.t2_int = t1_int group by t2_int) as t2_int from t1 +---- +11 1 +22 NULL +33 3 +44 NULL +# A `LIMIT` above the correlated predicate blocks decorrelation, so it is only +# allowed when the subquery is limited to a single row. #non_aggregated_correlated_scalar_subquery_with_limit -statement error DataFusion error: Invalid \(non-executable\) plan after Analyzer\ncaused by\nError during planning: Correlated scalar subquery must be aggregated to return at most one row +statement error DataFusion error: Invalid \(non-executable\) plan after Analyzer\ncaused by\nError during planning: Correlated scalar subquery with a LIMIT must be limited to a single row SELECT t1_id, (SELECT t2_int FROM t2 WHERE t2.t2_int = t1.t1_int limit 2) as t2_int from t1 #non_aggregated_correlated_scalar_subquery_with_single_row @@ -513,9 +534,17 @@ logical_plan 06)----------TableScan: t2 07)--TableScan: t1 projection=[t1_id] +# Grouping by a column the correlated predicate does not fix can produce more +# than one row per outer row, so this takes a single join instead of the plain +# left join used when the grouping is covered by the correlation. #aggregated_correlated_scalar_subquery_with_extra_group_by_columns -statement error DataFusion error: Invalid \(non-executable\) plan after Analyzer\ncaused by\nError during planning: A GROUP BY clause in a scalar correlated subquery cannot contain non-correlated columns +query II rowsort SELECT t1_id, (SELECT sum(t2_int) FROM t2 WHERE t2.t2_id = t1.t1_id group by t2_name) as t2_sum from t1 +---- +11 3 +22 1 +33 NULL +44 3 #support_agg_correlated_columns query TT @@ -2737,3 +2766,119 @@ b 400 statement ok DROP TABLE metrics; + +########## +## Single join +########## + +# A correlated scalar subquery that does not prove its own at-most-one-row +# property decorrelates into a single join, which checks it at runtime. These +# queries have no aggregate to pay for at all. + +statement ok +CREATE TABLE sj_fact(id INT, dim_id INT) AS VALUES +(1, 10), +(2, 20), +(3, 30), +(4, NULL); + +statement ok +CREATE TABLE sj_dim(dim_id INT, name TEXT) AS VALUES +(10, 'ten'), +(20, 'twenty'); + +statement ok +CREATE TABLE sj_dup(dim_id INT, name TEXT) AS VALUES +(10, 'ten'), +(10, 'ten again'); + +# An attribute lookup needs no aggregate: outer rows with no match are padded +# with NULL, exactly like the left join it lowers to. +query IT rowsort +SELECT id, (SELECT name FROM sj_dim WHERE sj_dim.dim_id = sj_fact.dim_id) AS name +FROM sj_fact +---- +1 ten +2 twenty +3 NULL +4 NULL + +query TT +EXPLAIN SELECT id, (SELECT name FROM sj_dim WHERE sj_dim.dim_id = sj_fact.dim_id) AS name +FROM sj_fact +---- +logical_plan +01)Projection: sj_fact.id, __scalar_sq_1.name AS name +02)--LeftSingle Join: sj_fact.dim_id = __scalar_sq_1.dim_id +03)----TableScan: sj_fact projection=[id, dim_id] +04)----SubqueryAlias: __scalar_sq_1 +05)------Projection: sj_dim.name, sj_dim.dim_id +06)--------TableScan: sj_dim projection=[dim_id, name] +physical_plan +01)HashJoinExec: mode=CollectLeft, join_type=LeftSingle, on=[(dim_id@1, dim_id@1)], projection=[id@0, name@2] +02)--DataSourceExec: partitions=1, partition_sizes=[1] +03)--DataSourceExec: partitions=1, partition_sizes=[1] + +# A second matching row makes the subquery non-scalar, which is an error -- +# the same one an uncorrelated scalar subquery raises. +statement error DataFusion error: Execution error: Scalar subquery returned more than one row +SELECT id, (SELECT name FROM sj_dup WHERE sj_dup.dim_id = sj_fact.dim_id) AS name +FROM sj_fact + +# The same in a filter rather than a projection. +query I rowsort +SELECT id FROM sj_fact +WHERE (SELECT name FROM sj_dim WHERE sj_dim.dim_id = sj_fact.dim_id) = 'ten' +---- +1 + +statement error DataFusion error: Execution error: Scalar subquery returned more than one row +SELECT id FROM sj_fact +WHERE (SELECT name FROM sj_dup WHERE sj_dup.dim_id = sj_fact.dim_id) = 'ten' + +# A correlation with no equijoin key runs as a nested loop single join. +query IT rowsort +SELECT id, (SELECT name FROM sj_dim WHERE sj_dim.dim_id > sj_fact.dim_id + 5) AS name +FROM sj_fact +---- +1 twenty +2 NULL +3 NULL +4 NULL + +statement error DataFusion error: Execution error: Scalar subquery returned more than one row +SELECT id, (SELECT name FROM sj_dim WHERE sj_dim.dim_id >= sj_fact.dim_id) AS name +FROM sj_fact + +# A subquery that does prove the property keeps its plain left join, so the +# runtime check is only paid where it is needed. +query TT +EXPLAIN SELECT id, (SELECT max(name) FROM sj_dup WHERE sj_dup.dim_id = sj_fact.dim_id) AS name +FROM sj_fact +---- +logical_plan +01)Projection: sj_fact.id, __scalar_sq_1.max(sj_dup.name) AS name +02)--Left Join: sj_fact.dim_id = __scalar_sq_1.dim_id +03)----TableScan: sj_fact projection=[id, dim_id] +04)----SubqueryAlias: __scalar_sq_1 +05)------Projection: max(sj_dup.name), sj_dup.dim_id +06)--------Aggregate: groupBy=[[sj_dup.dim_id]], aggr=[[max(sj_dup.name)]] +07)----------TableScan: sj_dup projection=[dim_id, name] +physical_plan +01)ProjectionExec: expr=[id@0 as id, max(sj_dup.name)@1 as name] +02)--HashJoinExec: mode=CollectLeft, join_type=Left, on=[(dim_id@1, dim_id@1)], projection=[id@0, max(sj_dup.name)@2] +03)----DataSourceExec: partitions=1, partition_sizes=[1] +04)----ProjectionExec: expr=[max(sj_dup.name)@1 as max(sj_dup.name), dim_id@0 as dim_id] +05)------AggregateExec: mode=FinalPartitioned, gby=[dim_id@0 as dim_id], aggr=[max(sj_dup.name)] +06)--------RepartitionExec: partitioning=Hash([dim_id@0], 4), input_partitions=1 +07)----------AggregateExec: mode=Partial, gby=[dim_id@0 as dim_id], aggr=[max(sj_dup.name)] +08)------------DataSourceExec: partitions=1, partition_sizes=[1] + +statement ok +DROP TABLE sj_fact; + +statement ok +DROP TABLE sj_dim; + +statement ok +DROP TABLE sj_dup; From 8b90f76ac9685f52da9a7bf6e08274286692d2d2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20Heres?= Date: Sun, 30 Aug 2026 06:45:00 +0200 Subject: [PATCH 3/3] perf: keep the plain left join when the subquery is already unique on the keys Preferring `LEFT JOIN` over a single join matters for more than the runtime check. A left join lets the optimizer turn the join inner when a predicate above it rejects nulls, fold a comparison against an outer column into a second join key, and finish with a semi join. None of that is sound for a single join, because all three change how many rows match -- which is the one thing a single join has to observe. So a correlated scalar subquery in a `WHERE` clause was about 1.6x slower than the aggregated form it replaces. `correlated_scalar_subquery_yields_single_row` only looks for an aggregate the rule can see from the top of the subquery. This also asks the decorrelated subquery's functional dependencies whether it is already unique on the columns the join equates with the outer plan, which covers uniqueness it inherits rather than declares: - a declared key, when the correlated predicate carries a further condition and so does not reduce to the `Filter::is_scalar` shape `max_rows` looks for; - a `GROUP BY` on the correlated column with nothing to aggregate; - uniqueness carried up through the subquery's own joins and projections. Only equality conjuncts count as join keys: uniqueness on `(a, b)` says nothing about how many rows `sub.a = outer.a AND sub.b > outer.b` matches. On TPC-H SF10 with `c_custkey` declared a primary key, `where (select c_mktsegment from customer where c_custkey = o_custkey) = 'BUILDING'` goes from 1.5x slower than the `min()` form to 1.6x faster (87 ms -> 55 ms), and the non-pushable variant from 1.55x slower to 1.15x faster. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01SgqwvctZdvR1ZCz2hbkEJC --- .../optimizer/src/scalar_subquery_to_join.rs | 98 ++++++++++++++++--- .../sqllogictest/test_files/subquery.slt | 79 +++++++++++++++ 2 files changed, 164 insertions(+), 13 deletions(-) diff --git a/datafusion/optimizer/src/scalar_subquery_to_join.rs b/datafusion/optimizer/src/scalar_subquery_to_join.rs index 499d5a6a64942..7bf9f7897736a 100644 --- a/datafusion/optimizer/src/scalar_subquery_to_join.rs +++ b/datafusion/optimizer/src/scalar_subquery_to_join.rs @@ -30,13 +30,15 @@ use datafusion_common::alias::AliasGenerator; use datafusion_common::tree_node::{ Transformed, TransformedResult, TreeNode, TreeNodeRecursion, TreeNodeRewriter, }; -use datafusion_common::{Column, Result, ScalarValue, assert_or_internal_err, plan_err}; +use datafusion_common::{ + Column, Dependency, Result, ScalarValue, assert_or_internal_err, plan_err, +}; use datafusion_expr::expr_rewriter::create_col_from_scalar_expr; use datafusion_expr::logical_plan::{JoinType, Subquery}; -use datafusion_expr::utils::conjunction; +use datafusion_expr::utils::{conjunction, split_conjunction}; use datafusion_expr::{ - Expr, LogicalPlan, LogicalPlanBuilder, correlated_scalar_subquery_yields_single_row, - lit, not, when, + BinaryExpr, Expr, LogicalPlan, LogicalPlanBuilder, Operator, + correlated_scalar_subquery_yields_single_row, lit, not, when, }; /// Optimizer rule that rewrites scalar subquery filters to joins and places an @@ -345,6 +347,69 @@ impl TreeNodeRewriter for ExtractScalarSubQuery<'_> { /// column to its `CASE WHEN __always_true IS NULL THEN ... END` compensation /// expression, which the caller must substitute into any expression that /// references those columns. +/// Returns true when the decorrelated subquery cannot match more than one row +/// for one set of outer values, because it is already unique on the columns the +/// join equates with the outer plan. +/// +/// This complements the shape test in +/// [`correlated_scalar_subquery_yields_single_row`]: it also covers a subquery +/// that inherits uniqueness from a declared constraint or from an aggregate +/// further down, rather than from an aggregate the caller can see. +/// +/// Only equality conjuncts count. Uniqueness on `(a, b)` says nothing about how +/// many rows `sub.a = outer.a AND sub.b > outer.b` matches. +fn subquery_unique_on_join_keys( + subquery: &LogicalPlan, + join_filter: &Expr, + subquery_alias: &str, +) -> bool { + let schema = subquery.schema(); + let mut join_key_indices = Vec::new(); + + for conjunct in split_conjunction(join_filter) { + let Expr::BinaryExpr(BinaryExpr { + left, + op: Operator::Eq, + right, + }) = conjunct + else { + continue; + }; + for (side, other) in [(left, right), (right, left)] { + let Some(column) = side.try_as_col() else { + continue; + }; + // The other side has to come from the outer plan, or this equality + // relates two subquery columns and is not a join key. + let is_join_key = qualified_by(column, subquery_alias) + && !other + .column_refs() + .iter() + .any(|column| qualified_by(column, subquery_alias)); + if is_join_key && let Some(index) = schema.maybe_index_of_column(column) { + join_key_indices.push(index); + } + } + } + + // A nullable determinant is still enough: the join uses + // `NullEquality::NullEqualsNothing`, so a NULL key matches nothing at all. + schema.functional_dependencies().iter().any(|dependency| { + dependency.mode == Dependency::Single + && dependency + .source_indices + .iter() + .all(|index| join_key_indices.contains(index)) + }) +} + +fn qualified_by(column: &Column, alias: &str) -> bool { + column + .relation + .as_ref() + .is_some_and(|relation| relation.table() == alias) +} + fn build_join( subquery: &Subquery, outer_input: &LogicalPlan, @@ -387,14 +452,20 @@ fn build_join( // the decorrelated subquery still yields at most one row per outer row // because its aggregate is grouped by the (empty) set of correlated inner // columns. - let join_filter = join_filter_opt.or_else(|| Some(lit(true))); - - // A subquery whose shape already guarantees at most one row per set of - // outer values joins with a plain `LEFT JOIN`; the rest need a single - // join, which raises an error at runtime when a second row matches. Only - // correlated subqueries can be ambiguous here -- an uncorrelated one is - // checked by `ScalarSubqueryExec` or, when this rule rewrites it, joined - // on `Boolean(true)` against a plan that already returns a single row. + let join_filter = join_filter_opt.unwrap_or_else(|| lit(true)); + + // A subquery that cannot produce more than one row per set of outer values + // joins with a plain `LEFT JOIN`; the rest need a single join, which raises + // an error at runtime when a second row matches. Only correlated + // subqueries can be ambiguous here -- an uncorrelated one is checked by + // `ScalarSubqueryExec` or, when this rule rewrites it, joined on + // `Boolean(true)` against a plan that already returns a single row. + // + // Preferring the plain join matters beyond saving the runtime check: the + // optimizer may turn a `LEFT JOIN` inner when a predicate above it rejects + // nulls, fold a comparison into a second join key, and finish with a semi + // join. None of that is sound for a single join, because all three change + // how many rows match, which is what the single join has to observe. // // The single join only sees the outer rows that reach it, so an outer row // a pushed-down filter removed cannot raise the error. Whether a @@ -402,6 +473,7 @@ fn build_join( // the same way it does in other engines. let join_type = if subquery.outer_ref_columns.is_empty() || correlated_scalar_subquery_yields_single_row(subquery_plan)? + || subquery_unique_on_join_keys(&aliased_subquery, &join_filter, subquery_alias) { JoinType::Left } else { @@ -409,7 +481,7 @@ fn build_join( }; let new_plan = LogicalPlanBuilder::from(outer_input.clone()) - .join_on(aliased_subquery, join_type, join_filter)? + .join_on(aliased_subquery, join_type, Some(join_filter))? .build()?; // Add count-bug compensation for each of the subquery's projected diff --git a/datafusion/sqllogictest/test_files/subquery.slt b/datafusion/sqllogictest/test_files/subquery.slt index 03823dfa0bd6f..5c9322e5f2c02 100644 --- a/datafusion/sqllogictest/test_files/subquery.slt +++ b/datafusion/sqllogictest/test_files/subquery.slt @@ -2874,6 +2874,82 @@ physical_plan 07)----------AggregateExec: mode=Partial, gby=[dim_id@0 as dim_id], aggr=[max(sj_dup.name)] 08)------------DataSourceExec: partitions=1, partition_sizes=[1] +# Uniqueness the subquery inherits counts too, not just an aggregate the rule +# can see: a declared key still proves the property when the correlated +# predicate carries an extra condition, so this keeps the plain left join and +# with it every rewrite a single join would block. +statement ok +CREATE TABLE sj_dim_pk(dim_id INT PRIMARY KEY, name TEXT, ok BOOLEAN) AS VALUES +(10, 'ten', true), +(20, 'twenty', false); + +query TT +EXPLAIN SELECT id, (SELECT name FROM sj_dim_pk WHERE sj_dim_pk.dim_id = sj_fact.dim_id AND ok) AS name +FROM sj_fact +---- +logical_plan +01)Projection: sj_fact.id, __scalar_sq_1.name AS name +02)--Left Join: sj_fact.dim_id = __scalar_sq_1.dim_id +03)----TableScan: sj_fact projection=[id, dim_id] +04)----SubqueryAlias: __scalar_sq_1 +05)------Projection: sj_dim_pk.name, sj_dim_pk.dim_id +06)--------Filter: sj_dim_pk.ok +07)----------TableScan: sj_dim_pk projection=[dim_id, name, ok] +physical_plan +01)HashJoinExec: mode=CollectLeft, join_type=Right, on=[(dim_id@1, dim_id@1)], projection=[id@2, name@0] +02)--FilterExec: ok@2, projection=[name@1, dim_id@0] +03)----DataSourceExec: partitions=1, partition_sizes=[1] +04)--DataSourceExec: partitions=1, partition_sizes=[1] + +query IT rowsort +SELECT id, (SELECT name FROM sj_dim_pk WHERE sj_dim_pk.dim_id = sj_fact.dim_id AND ok) AS name +FROM sj_fact +---- +1 ten +2 NULL +3 NULL +4 NULL + +# Grouping by the correlated column makes the subquery unique on the join key +# even with no aggregate to group for. +query TT +EXPLAIN SELECT id, (SELECT dim_id FROM sj_dup WHERE sj_dup.dim_id = sj_fact.dim_id GROUP BY dim_id) AS dim_id +FROM sj_fact +---- +logical_plan +01)Projection: sj_fact.id, __scalar_sq_1.dim_id AS dim_id +02)--Left Join: sj_fact.dim_id = __scalar_sq_1.dim_id +03)----TableScan: sj_fact projection=[id, dim_id] +04)----SubqueryAlias: __scalar_sq_1 +05)------Aggregate: groupBy=[[sj_dup.dim_id]], aggr=[[]] +06)--------TableScan: sj_dup projection=[dim_id] +physical_plan +01)HashJoinExec: mode=CollectLeft, join_type=Right, on=[(dim_id@0, dim_id@1)], projection=[id@1, dim_id@0] +02)--CoalescePartitionsExec +03)----AggregateExec: mode=FinalPartitioned, gby=[dim_id@0 as dim_id], aggr=[] +04)------RepartitionExec: partitioning=Hash([dim_id@0], 4), input_partitions=1 +05)--------AggregateExec: mode=Partial, gby=[dim_id@0 as dim_id], aggr=[] +06)----------DataSourceExec: partitions=1, partition_sizes=[1] +07)--DataSourceExec: partitions=1, partition_sizes=[1] + +# A unique key proves nothing when the correlated predicate is an inequality: +# many rows can still match one outer row. +query TT +EXPLAIN SELECT id, (SELECT name FROM sj_dim_pk WHERE sj_dim_pk.dim_id > sj_fact.dim_id) AS name +FROM sj_fact +---- +logical_plan +01)Projection: sj_fact.id, __scalar_sq_1.name AS name +02)--LeftSingle Join: Filter: __scalar_sq_1.dim_id > sj_fact.dim_id +03)----TableScan: sj_fact projection=[id, dim_id] +04)----SubqueryAlias: __scalar_sq_1 +05)------Projection: sj_dim_pk.name, sj_dim_pk.dim_id +06)--------TableScan: sj_dim_pk projection=[dim_id, name] +physical_plan +01)NestedLoopJoinExec: join_type=LeftSingle, filter=dim_id@1 > dim_id@0, projection=[id@0, name@2] +02)--DataSourceExec: partitions=1, partition_sizes=[1] +03)--DataSourceExec: partitions=1, partition_sizes=[1] + statement ok DROP TABLE sj_fact; @@ -2882,3 +2958,6 @@ DROP TABLE sj_dim; statement ok DROP TABLE sj_dup; + +statement ok +DROP TABLE sj_dim_pk;