feat: rewrite subquery into dependent join logical plan

datafusion/core/src/physical_planner.rs

@@ -1250,6 +1250,11 @@ impl DefaultPhysicalPlanner {
                     "Unsupported logical plan: Analyze must be root of the plan"
                 )
             }
+            LogicalPlan::DependentJoin(_) => {
+                return internal_err!(
+                    "Optimizors have not completely remove dependent join"
+                )
+            }
         };
         Ok(exec_node)
     }

datafusion/expr/src/logical_plan/builder.rs

@@ -31,10 +31,10 @@ use crate::expr_rewriter::{
     rewrite_sort_cols_by_aggs,
 };
 use crate::logical_plan::{
-    Aggregate, Analyze, Distinct, DistinctOn, EmptyRelation, Explain, Filter, Join,
-    JoinConstraint, JoinType, Limit, LogicalPlan, Partitioning, PlanType, Prepare,
-    Projection, Repartition, Sort, SubqueryAlias, TableScan, Union, Unnest, Values,
-    Window,
+    Aggregate, Analyze, DependentJoin, Distinct, DistinctOn, EmptyRelation, Explain,
+    Filter, Join, JoinConstraint, JoinType, Limit, LogicalPlan, Partitioning, PlanType,
+    Prepare, Projection, Repartition, Sort, SubqueryAlias, TableScan, Union, Unnest,
+    Values, Window,
 };
 use crate::select_expr::SelectExpr;
 use crate::utils::{
@@ -883,6 +883,47 @@ impl LogicalPlanBuilder {
         ))))
     }
 
+    /// Build a dependent join provided a subquery plan
+    /// this function should only be used by the optimizor
+    /// a dependent join node will provides all columns belonging to the LHS
+    /// and one additional column as the result of evaluating the subquery on the RHS
+    /// under the name "subquery_name.output"
+    pub fn dependent_join(
+        self,
+        right: LogicalPlan,
+        correlated_columns: Vec<(usize, Column, DataType)>,
+        subquery_expr: Option<Expr>,
+        subquery_depth: usize,
+        subquery_name: String,
+        lateral_join_condition: Option<(JoinType, Expr)>,
+    ) -> Result<Self> {
+        let left = self.build()?;
+        let schema = left.schema();
+        // TODO: for lateral join, output schema is similar to a normal join
+        let qualified_fields = schema
+            .iter()
+            .map(|(q, f)| (q.cloned(), Arc::clone(f)))
+            .chain(
+                subquery_expr
+                    .iter()
+                    .map(|expr| subquery_output_field(&subquery_name, expr)),
+            )
+            .collect();
+        let metadata = schema.metadata().clone();
+        let dfschema = DFSchema::new_with_metadata(qualified_fields, metadata)?;
+
+        Ok(Self::new(LogicalPlan::DependentJoin(DependentJoin {
+            schema: DFSchemaRef::new(dfschema),
+            left: Arc::new(left),
+            right: Arc::new(right),
+            correlated_columns,
+            subquery_expr,
+            subquery_name,
+            subquery_depth,
+            lateral_join_condition,
+        })))
+    }
+
     /// Apply a join to `right` using explicitly specified columns and an
     /// optional filter expression.
     ///
@@ -1547,6 +1588,26 @@ fn mark_field(schema: &DFSchema) -> (Option<TableReference>, Arc<Field>) {
     )
 }
 
+fn subquery_output_field(
+    subquery_alias: &str,
+    subquery_expr: &Expr,
+) -> (Option<TableReference>, Arc<Field>) {
+    // TODO: check nullability
+    let field = match subquery_expr {
+        Expr::InSubquery(_) => Arc::new(Field::new("output", DataType::Boolean, false)),
+        Expr::Exists(_) => Arc::new(Field::new("output", DataType::Boolean, false)),
+        Expr::ScalarSubquery(sq) => {
+            let data_type = sq.subquery.schema().field(0).data_type().clone();
+            Arc::new(Field::new("output", data_type, false))
+        }
+        _ => {
+            unreachable!()
+        }
+    };
+
+    (Some(TableReference::bare(subquery_alias)), field)
+}
+
 /// Creates a schema for a join operation.
 /// The fields from the left side are first
 pub fn build_join_schema(

datafusion/expr/src/logical_plan/display.rs

@@ -485,6 +485,7 @@ impl<'a, 'b> PgJsonVisitor<'a, 'b> {
 
                 object
             }
+            LogicalPlan::DependentJoin(..) => json!({}),
             LogicalPlan::Join(Join {
                 on: ref keys,
                 filter,

datafusion/expr/src/logical_plan/invariants.rs

@@ -201,20 +201,27 @@ pub fn check_subquery_expr(
             }?;
             match outer_plan {
                 LogicalPlan::Projection(_)
-                | LogicalPlan::Filter(_) => Ok(()),
-                LogicalPlan::Aggregate(Aggregate { group_expr, aggr_expr, .. }) => {
+                | LogicalPlan::Filter(_)
+                | LogicalPlan::DependentJoin(_) => Ok(()),
+                LogicalPlan::Aggregate(Aggregate {
+                    group_expr,
+                    aggr_expr,
+                    ..
+                }) => {
                     if group_expr.contains(expr) && !aggr_expr.contains(expr) {
                         // TODO revisit this validation logic
                         plan_err!(
-                            "Correlated scalar subquery in the GROUP BY clause must also be in the aggregate expressions"
+                            "Correlated scalar subquery in the GROUP BY clause must \
+                            also be in the aggregate expressions"
                         )
                     } else {
                         Ok(())
                     }
                 }
                 _ => plan_err!(
-                    "Correlated scalar subquery can only be used in Projection, Filter, Aggregate plan nodes"
-                )
+                    "Correlated scalar subquery can only be used in Projection, Filter, \
+                    Aggregate, DependentJoin plan nodes"
+                ),
             }?;
         }
         check_correlations_in_subquery(inner_plan)
@@ -235,11 +242,12 @@ pub fn check_subquery_expr(
             | LogicalPlan::TableScan(_)
             | LogicalPlan::Window(_)
             | LogicalPlan::Aggregate(_)
-            | LogicalPlan::Join(_) => Ok(()),
+            | LogicalPlan::Join(_)
+            | LogicalPlan::DependentJoin(_) => Ok(()),
             _ => plan_err!(
                 "In/Exist subquery can only be used in \
-                Projection, Filter, TableScan, Window functions, Aggregate and Join plan nodes, \
-                but was used in [{}]",
+                Projection, Filter, TableScan, Window functions, Aggregate, Join and \
+                Dependent Join plan nodes, but was used in [{}]",
                 outer_plan.display()
             ),
         }?;
@@ -323,6 +331,7 @@ fn check_inner_plan(inner_plan: &LogicalPlan) -> Result<()> {
             }
         },
         LogicalPlan::Extension(_) => Ok(()),
+        LogicalPlan::DependentJoin(_) => Ok(()),
         plan => check_no_outer_references(plan),
     }
 }

datafusion/expr/src/logical_plan/mod.rs

@@ -37,11 +37,12 @@ pub use ddl::{
 };
 pub use dml::{DmlStatement, WriteOp};
 pub use plan::{
-    projection_schema, Aggregate, Analyze, ColumnUnnestList, DescribeTable, Distinct,
-    DistinctOn, EmptyRelation, Explain, ExplainFormat, Extension, FetchType, Filter,
-    Join, JoinConstraint, JoinType, Limit, LogicalPlan, Partitioning, PlanType,
-    Projection, RecursiveQuery, Repartition, SkipType, Sort, StringifiedPlan, Subquery,
-    SubqueryAlias, TableScan, ToStringifiedPlan, Union, Unnest, Values, Window,
+    projection_schema, Aggregate, Analyze, ColumnUnnestList, DependentJoin,
+    DescribeTable, Distinct, DistinctOn, EmptyRelation, Explain, ExplainFormat,
+    Extension, FetchType, Filter, Join, JoinConstraint, JoinType, Limit, LogicalPlan,
+    Partitioning, PlanType, Projection, RecursiveQuery, Repartition, SkipType, Sort,
+    StringifiedPlan, Subquery, SubqueryAlias, TableScan, ToStringifiedPlan, Union,
+    Unnest, Values, Window,
 };
 pub use statement::{
     Deallocate, Execute, Prepare, SetVariable, Statement, TransactionAccessMode,

datafusion/expr/src/logical_plan/plan.rs

@@ -287,6 +287,99 @@ pub enum LogicalPlan {
     Unnest(Unnest),
     /// A variadic query (e.g. "Recursive CTEs")
     RecursiveQuery(RecursiveQuery),
+    /// A node type that only exist during subquery decorrelation
+    /// TODO: maybe we can avoid creating new type of LogicalPlan for this usecase
+    DependentJoin(DependentJoin),
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+pub struct DependentJoin {
+    pub schema: DFSchemaRef,
+    // All combinations of (subquery depth,Column and its DataType) on the RHS (and its descendant)
+    // which points to a column on the LHS of this dependent join
+    // Note that not all outer_refs from the RHS are mentioned in this vectors
+    // because RHS may reference columns provided somewhere from the above parent dependent join.
+    // Depths of each correlated_columns should always be gte current dependent join
+    // subquery_depth
+    pub correlated_columns: Vec<(usize, Column, DataType)>,
+    // the upper expr that containing the subquery expr
+    // i.e for predicates: where outer = scalar_sq + 1
+    // correlated exprs are `scalar_sq + 1`
+    pub subquery_expr: Option<Expr>,
+    // begins with depth = 1
+    pub subquery_depth: usize,
+    pub left: Arc<LogicalPlan>,
+    // dependent side accessing columns from left hand side (and maybe columns)
+    // belong to the parent dependent join node in case of recursion)
+    pub right: Arc<LogicalPlan>,
+    pub subquery_name: String,
+
+    pub lateral_join_condition: Option<(JoinType, Expr)>,
+}
+
+impl Display for DependentJoin {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        let correlated_str = self
+            .correlated_columns
+            .iter()
+            .map(|(level, col, _)| format!("{col} lvl {level}"))
+            .collect::<Vec<String>>()
+            .join(", ");
+        let lateral_join_info =
+            if let Some((join_type, join_expr)) = &self.lateral_join_condition {
+                format!(" lateral {join_type} join with {join_expr}")
+            } else {
+                "".to_string()
+            };
+        let subquery_expr_str = if let Some(expr) = &self.subquery_expr {
+            format!(" with expr {expr}")
+        } else {
+            "".to_string()
+        };
+        write!(
+            f,
+            "DependentJoin on [{correlated_str}]{subquery_expr_str}\
+                        {lateral_join_info} depth {0}",
+            self.subquery_depth,
+        )
+    }
+}
+
+impl PartialOrd for DependentJoin {
+    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+        #[derive(PartialEq, PartialOrd)]
+        struct ComparableJoin<'a> {
+            correlated_columns: &'a Vec<(usize, Column, DataType)>,
+            // the upper expr that containing the subquery expr
+            // i.e for predicates: where outer = scalar_sq + 1
+            // correlated exprs are `scalar_sq + 1`
+            subquery_expr: &'a Option<Expr>,
+
+            depth: &'a usize,
+            left: &'a Arc<LogicalPlan>,
+            // dependent side accessing columns from left hand side (and maybe columns)
+            // belong to the parent dependent join node in case of recursion)
+            right: &'a Arc<LogicalPlan>,
+            lateral_join_condition: &'a Option<(JoinType, Expr)>,
+        }
+        let comparable_self = ComparableJoin {
+            left: &self.left,
+            right: &self.right,
+            correlated_columns: &self.correlated_columns,
+            subquery_expr: &self.subquery_expr,
+            depth: &self.subquery_depth,
+            lateral_join_condition: &self.lateral_join_condition,
+        };
+        let comparable_other = ComparableJoin {
+            left: &other.left,
+            right: &other.right,
+            correlated_columns: &other.correlated_columns,
+            subquery_expr: &other.subquery_expr,
+            depth: &other.subquery_depth,
+            lateral_join_condition: &other.lateral_join_condition,
+        };
+        comparable_self.partial_cmp(&comparable_other)
+    }
 }
 
 impl Default for LogicalPlan {
@@ -318,6 +411,7 @@ impl LogicalPlan {
     /// Get a reference to the logical plan's schema
     pub fn schema(&self) -> &DFSchemaRef {
         match self {
+            LogicalPlan::DependentJoin(DependentJoin { schema, .. }) => schema,
             LogicalPlan::EmptyRelation(EmptyRelation { schema, .. }) => schema,
             LogicalPlan::Values(Values { schema, .. }) => schema,
             LogicalPlan::TableScan(TableScan {
@@ -452,6 +546,9 @@ impl LogicalPlan {
             LogicalPlan::Aggregate(Aggregate { input, .. }) => vec![input],
             LogicalPlan::Sort(Sort { input, .. }) => vec![input],
             LogicalPlan::Join(Join { left, right, .. }) => vec![left, right],
+            LogicalPlan::DependentJoin(DependentJoin { left, right, .. }) => {
+                vec![left, right]
+            }
             LogicalPlan::Limit(Limit { input, .. }) => vec![input],
             LogicalPlan::Subquery(Subquery { subquery, .. }) => vec![subquery],
             LogicalPlan::SubqueryAlias(SubqueryAlias { input, .. }) => vec![input],
@@ -540,6 +637,7 @@ impl LogicalPlan {
             | LogicalPlan::Limit(Limit { input, .. })
             | LogicalPlan::Repartition(Repartition { input, .. })
             | LogicalPlan::Window(Window { input, .. }) => input.head_output_expr(),
+            LogicalPlan::DependentJoin(..) => todo!(),
             LogicalPlan::Join(Join {
                 left,
                 right,
@@ -647,6 +745,7 @@ impl LogicalPlan {
             }) => Aggregate::try_new(input, group_expr, aggr_expr)
                 .map(LogicalPlan::Aggregate),
             LogicalPlan::Sort(_) => Ok(self),
+            LogicalPlan::DependentJoin(_) => todo!(),
             LogicalPlan::Join(Join {
                 left,
                 right,
@@ -1138,6 +1237,7 @@ impl LogicalPlan {
                     unnest_with_options(input, columns.clone(), options.clone())?;
                 Ok(new_plan)
             }
+            LogicalPlan::DependentJoin(_) => todo!(),
         }
     }
 
@@ -1290,6 +1390,7 @@ impl LogicalPlan {
     /// If `Some(n)` then the plan can return at most `n` rows but may return fewer.
     pub fn max_rows(self: &LogicalPlan) -> Option<usize> {
         match self {
+            LogicalPlan::DependentJoin(DependentJoin { left, .. }) => left.max_rows(),
             LogicalPlan::Projection(Projection { input, .. }) => input.max_rows(),
             LogicalPlan::Filter(filter) => {
                 if filter.is_scalar() {
@@ -1882,6 +1983,10 @@ impl LogicalPlan {
 
                         Ok(())
                     }
+
+                    LogicalPlan::DependentJoin(dependent_join) => {
+                        Display::fmt(dependent_join, f)
+                    },
                     LogicalPlan::Join(Join {
                         on: ref keys,
                         filter,