feat: Single join for correlated scalar subqueries

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();
                 }

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"),
         }
     }

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()

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

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

datafusion/expr/src/logical_plan/builder.rs

@@ -1703,7 +1703,7 @@ pub fn build_join_schema(
                 .collect::<Vec<_>>();
             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)))

datafusion/expr/src/logical_plan/invariants.rs

@@ -169,31 +169,22 @@ 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.
+        //
+        // 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() {
-            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"
-                        )
-                    }
-                }
-            }?;
+            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 {
@@ -323,14 +314,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 +351,73 @@ fn check_no_outer_references(inner_plan: &LogicalPlan) -> Result<()> {
     }
 }
 
-fn check_aggregation_in_scalar_subquery(
+/// 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.
+///
+/// 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<bool> {
+    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::<Vec<_>>()))
-            .collect::<Result<Vec<_>>>()?
-            .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::<Vec<_>>()))
+        .collect::<Result<Vec<_>>>()?
+        .into_iter()
+        .flatten();
+
+    Ok(group_columns.all(|group| inner_subquery_cols.contains(&group)))
 }
 
 fn strip_inner_query(inner_plan: &LogicalPlan) -> &LogicalPlan {

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;