Planner: support LATERAL subqueries

Author: aalexandrov

datafusion/sql/src/planner.rs

@@ -17,12 +17,13 @@
 
 //! [`SqlToRel`]: SQL Query Planner (produces [`LogicalPlan`] from SQL AST)
 use std::collections::HashMap;
+use std::ops::Deref;
 use std::sync::Arc;
 use std::vec;
 
 use arrow_schema::*;
 use datafusion_common::{
-    field_not_found, internal_err, plan_datafusion_err, SchemaError,
+    field_not_found, internal_err, plan_datafusion_err, DFSchemaRef, SchemaError,
 };
 use datafusion_expr::WindowUDF;
 use sqlparser::ast::TimezoneInfo;
@@ -149,6 +150,9 @@ pub struct PlannerContext {
     ctes: HashMap<String, Arc<LogicalPlan>>,
     /// The query schema of the outer query plan, used to resolve the columns in subquery
     outer_query_schema: Option<DFSchema>,
+    /// The joined schemas of all FROM clauses planned so far. When planning LATERAL
+    /// FROM clauses, this should become a suffix of the `outer_query_schema`.
+    outer_from_schema: Option<DFSchema>,
 }
 
 impl Default for PlannerContext {
@@ -164,6 +168,7 @@ impl PlannerContext {
             prepare_param_data_types: vec![],
             ctes: HashMap::new(),
             outer_query_schema: None,
+            outer_from_schema: None,
         }
     }
 
@@ -191,6 +196,25 @@ impl PlannerContext {
         schema
     }
 
+    /// sets the FROM schema, returning the existing one, if
+    /// any
+    pub fn set_outer_from_schema(
+        &mut self,
+        mut schema: Option<DFSchema>,
+    ) -> Option<DFSchema> {
+        std::mem::swap(&mut self.outer_from_schema, &mut schema);
+        schema
+    }
+
+    /// extends the FROM schema, returning the existing one, if any
+    pub fn extend_outer_from_schema(&mut self, schema: &DFSchemaRef) -> Result<()> {
+        self.outer_from_schema = match self.outer_from_schema.as_ref() {
+            Some(from_schema) => Some(from_schema.join(schema)?),
+            None => Some(schema.deref().clone()),
+        };
+        Ok(())
+    }
+
     /// Return the types of parameters (`$1`, `$2`, etc) if known
     pub fn prepare_param_data_types(&self) -> &[DataType] {
         &self.prepare_param_data_types

datafusion/sql/src/relation/join.rs

@@ -18,7 +18,7 @@
 use crate::planner::{ContextProvider, PlannerContext, SqlToRel};
 use datafusion_common::{not_impl_err, Column, Result};
 use datafusion_expr::{JoinType, LogicalPlan, LogicalPlanBuilder};
-use sqlparser::ast::{Join, JoinConstraint, JoinOperator, TableWithJoins};
+use sqlparser::ast::{Join, JoinConstraint, JoinOperator, TableFactor, TableWithJoins};
 use std::collections::HashSet;
 
 impl<'a, S: ContextProvider> SqlToRel<'a, S> {
@@ -33,28 +33,20 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
         // So always use original global CTEs to plan CTEs in from clause.
         // Btw, don't need to add CTEs in from to global CTEs.
         let origin_planner_context = planner_context.clone();
-        let left = self.create_relation(t.relation, planner_context)?;
-        match t.joins.len() {
-            0 => {
-                *planner_context = origin_planner_context;
-                Ok(left)
-            }
-            _ => {
-                let mut joins = t.joins.into_iter();
+        let mut left = if is_lateral_factor(&t.relation) {
+            self.create_relation_subquery(t.relation, planner_context)?
+        } else {
+            self.create_relation(t.relation, planner_context)?
+        };
+        if !t.joins.is_empty() {
+            for join in t.joins {
                 *planner_context = origin_planner_context.clone();
-                let mut left = self.parse_relation_join(
-                    left,
-                    joins.next().unwrap(), // length of joins > 0
-                    planner_context,
-                )?;
-                for join in joins {
-                    *planner_context = origin_planner_context.clone();
-                    left = self.parse_relation_join(left, join, planner_context)?;
-                }
-                *planner_context = origin_planner_context;
-                Ok(left)
+                planner_context.extend_outer_from_schema(left.schema())?;
+                left = self.parse_relation_join(left, join, planner_context)?;
             }
         }
+        *planner_context = origin_planner_context;
+        Ok(left)
     }
 
     fn parse_relation_join(
@@ -63,7 +55,11 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
         join: Join,
         planner_context: &mut PlannerContext,
     ) -> Result<LogicalPlan> {
-        let right = self.create_relation(join.relation, planner_context)?;
+        let right = if is_lateral_join(&join)? {
+            self.create_relation_subquery(join.relation, planner_context)?
+        } else {
+            self.create_relation(join.relation, planner_context)?
+        };
         match join.join_operator {
             JoinOperator::LeftOuter(constraint) => {
                 self.parse_join(left, right, constraint, JoinType::Left, planner_context)
@@ -167,3 +163,30 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
         }
     }
 }
+
+/// Return `true` iff the given [`TableFactor`] is lateral.
+pub(crate) fn is_lateral_factor(factor: &TableFactor) -> bool {
+    match factor {
+        TableFactor::Derived { lateral, .. } => *lateral,
+        TableFactor::Function { lateral, .. } => *lateral,
+        _ => false,
+    }
+}
+
+/// Return `true` iff the given [`Join`] is lateral.
+pub(crate) fn is_lateral_join(join: &Join) -> Result<bool> {
+    let is_lateral_syntax = is_lateral_factor(&join.relation);
+    let is_apply_syntax = match join.join_operator {
+        JoinOperator::FullOuter(..)
+        | JoinOperator::RightOuter(..)
+        | JoinOperator::RightAnti(..)
+        | JoinOperator::RightSemi(..)
+            if is_lateral_syntax =>
+        {
+            return not_impl_err!("NONE constraint is not supported");
+        }
+        JoinOperator::CrossApply | JoinOperator::OuterApply => true,
+        _ => false,
+    };
+    Ok(is_lateral_syntax || is_apply_syntax)
+}

datafusion/sql/src/relation/mod.rs

@@ -15,9 +15,13 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use std::sync::Arc;
+
 use crate::planner::{ContextProvider, PlannerContext, SqlToRel};
 use datafusion_common::{not_impl_err, plan_err, DFSchema, Result, TableReference};
+use datafusion_expr::builder::subquery_alias;
 use datafusion_expr::{expr::Unnest, Expr, LogicalPlan, LogicalPlanBuilder};
+use datafusion_expr::{Subquery, SubqueryAlias};
 use sqlparser::ast::{FunctionArg, FunctionArgExpr, TableFactor};
 
 mod join;
@@ -143,4 +147,39 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
             Ok(plan)
         }
     }
+
+    pub(crate) fn create_relation_subquery(
+        &self,
+        subquery: TableFactor,
+        planner_context: &mut PlannerContext,
+    ) -> Result<LogicalPlan> {
+        let old_from_sch = planner_context.set_outer_from_schema(None).unwrap(); // set in plan_table_with_joins
+        let new_outer_sch = match planner_context.outer_query_schema() {
+            Some(lhs) => Some(lhs.join(&old_from_sch)?),
+            None => Some(old_from_sch.clone()),
+        };
+        let old_outer_sch = planner_context.set_outer_query_schema(new_outer_sch);
+
+        let plan = self.create_relation(subquery, planner_context)?;
+        let outer_ref_columns = plan.all_out_ref_exprs();
+
+        planner_context.set_outer_query_schema(old_outer_sch);
+        planner_context.set_outer_from_schema(Some(old_from_sch));
+
+        match plan {
+            LogicalPlan::SubqueryAlias(SubqueryAlias { input, alias, .. }) => {
+                subquery_alias(
+                    LogicalPlan::Subquery(Subquery {
+                        subquery: input,
+                        outer_ref_columns,
+                    }),
+                    alias,
+                )
+            }
+            plan => Ok(LogicalPlan::Subquery(Subquery {
+                subquery: Arc::new(plan),
+                outer_ref_columns,
+            })),
+        }
+    }
 }

datafusion/sql/src/select.rs

@@ -25,7 +25,7 @@ use crate::utils::{
     resolve_positions_to_exprs,
 };
 
-use datafusion_common::{not_impl_err, plan_err, DataFusionError, Result};
+use datafusion_common::{not_impl_err, plan_err, DFSchema, DataFusionError, Result};
 use datafusion_common::{Column, UnnestOptions};
 use datafusion_expr::expr::Alias;
 use datafusion_expr::expr_rewriter::{
@@ -380,19 +380,30 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
         match from.len() {
             0 => Ok(LogicalPlanBuilder::empty(true).build()?),
             1 => {
-                let from = from.remove(0);
-                self.plan_table_with_joins(from, planner_context)
+                let input = from.remove(0);
+                self.plan_table_with_joins(input, planner_context)
             }
             _ => {
-                let mut plans = from
-                    .into_iter()
-                    .map(|t| self.plan_table_with_joins(t, planner_context));
-
-                let mut left = LogicalPlanBuilder::from(plans.next().unwrap()?);
-
-                for right in plans {
-                    left = left.cross_join(right?)?;
+                let mut from = from.into_iter();
+
+                let mut left = LogicalPlanBuilder::from({
+                    let input = from.next().unwrap();
+                    self.plan_table_with_joins(input, planner_context)?
+                });
+                let old_outer_from_schema = {
+                    let left_schema = Some(DFSchema::clone(left.schema()));
+                    planner_context.set_outer_from_schema(left_schema)
+                };
+                for input in from {
+                    // Join `input` with the current result (`left`).
+                    let right = self.plan_table_with_joins(input, planner_context)?;
+                    left = left.cross_join(right)?;
+                    // Update the outer FROM schema.
+                    let left_schema = Some(DFSchema::clone(left.schema()));
+                    planner_context.set_outer_from_schema(left_schema);
                 }
+                planner_context.set_outer_from_schema(old_outer_from_schema);
+
                 Ok(left.build()?)
             }
         }

datafusion/sql/tests/sql_integration.rs

@@ -3188,6 +3188,56 @@ fn join_on_complex_condition() {
     quick_test(sql, expected);
 }
 
+#[test]
+fn lateral_comma_join() {
+    let sql = "SELECT j1_string, j2_string FROM
+            j1, \
+            LATERAL (SELECT * FROM j2 WHERE j1_id < j2_id) AS j2";
+    let expected = "Projection: j1.j1_string, j2.j2_string\
+            \n  CrossJoin:\
+            \n    TableScan: j1\
+            \n    SubqueryAlias: j2\
+            \n      Subquery:\
+            \n        Projection: j2.j2_id, j2.j2_string\
+            \n          Filter: outer_ref(j1.j1_id) < j2.j2_id\
+            \n            TableScan: j2";
+    quick_test(sql, expected);
+}
+
+#[test]
+fn lateral_left_join() {
+    let sql = "SELECT j1_string, j2_string FROM 
+            j1 \
+            LEFT JOIN LATERAL (SELECT * FROM j2 WHERE j1_id < j2_id) AS j2 ON(true);";
+    let expected = "Projection: j1.j1_string, j2.j2_string\
+            \n  Left Join:  Filter: Boolean(true)\
+            \n    TableScan: j1\
+            \n    SubqueryAlias: j2\
+            \n      Subquery:\
+            \n        Projection: j2.j2_id, j2.j2_string\
+            \n          Filter: outer_ref(j1.j1_id) < j2.j2_id\
+            \n            TableScan: j2";
+    quick_test(sql, expected);
+}
+
+#[test]
+fn lateral_nested_left_join() {
+    let sql = "SELECT * FROM 
+            j1, \
+            (j2 LEFT JOIN LATERAL (SELECT * FROM j3 WHERE j1_id + j2_id = j3_id) AS j3 ON(true))";
+    let expected = "Projection: j1.j1_id, j1.j1_string, j2.j2_id, j2.j2_string, j3.j3_id, j3.j3_string\
+            \n  CrossJoin:\
+            \n    TableScan: j1\
+            \n    Left Join:  Filter: Boolean(true)\
+            \n      TableScan: j2\
+            \n      SubqueryAlias: j3\
+            \n        Subquery:\
+            \n          Projection: j3.j3_id, j3.j3_string\
+            \n            Filter: outer_ref(j1.j1_id) + outer_ref(j2.j2_id) = j3.j3_id\
+            \n              TableScan: j3";
+    quick_test(sql, expected);
+}
+
 #[test]
 fn hive_aggregate_with_filter() -> Result<()> {
     let dialect = &HiveDialect {};